scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_key_soft_budget_edit

Browse source
This commit is contained in:
yassin 2026-09-07 19:23:42 +00:00
commit 4728d7f609
1848 changed files with 90160 additions and 20864 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
.circleci/config.yml
View file

@ -575,7 +575,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"langfuse\""
@ -630,7 +630,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -737,7 +737,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -782,7 +782,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"assistants\""
@ -909,7 +909,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
@ -999,7 +999,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1054,7 +1054,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 8 \
@ -1090,7 +1090,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1134,7 +1134,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1178,7 +1178,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1222,7 +1222,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1267,7 +1267,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1312,7 +1312,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
@ -1391,7 +1391,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
@ -1444,7 +1444,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 -n 2 \
@ -1705,7 +1705,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -1794,7 +1794,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-s -v -x \
-s -v \
--junitxml=test-results/junit.xml \
-n 4 \
--durations=5"
@ -2012,7 +2012,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -2092,7 +2092,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2195,7 +2195,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2266,7 +2266,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2350,7 +2350,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
-vv \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
@ -2446,7 +2446,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
-v \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -2516,7 +2516,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
-vv -s \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m

14
.github/actions/setup-uv-with-retries/action.yml vendored
View file

@ -1,11 +1,7 @@
name: "Set up uv with retries"
description: >-
Install uv via astral-sh/setup-uv, retrying on transient failures. Even with
an exact pinned version, the action resolves the artifact URL by fetching
https://raw.githubusercontent.com/astral-sh/versions/main/v1/uv.ndjson in a
single request with no retry, timeout, or fallback, so one connection-level
network error ("fetch failed") fails the whole job before any test runs.
Retrying the full step covers the manifest fetch and the binary download.
Install uv via astral-sh/setup-uv, retrying the full setup step so manifest
resolution and binary downloads get fresh attempts after transient failures.
inputs:
version:
@ -18,7 +14,7 @@ runs:
- name: Set up uv (attempt 1)
id: attempt-1
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}
@ -31,7 +27,7 @@ runs:
id: attempt-2
if: steps.attempt-1.outcome == 'failure'
continue-on-error: true
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}
@ -42,6 +38,6 @@ runs:
- name: Set up uv (attempt 3)
if: steps.attempt-2.outcome == 'failure'
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}

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

@ -79,6 +79,11 @@ test_paths:
- tests/load_tests/test_otel_load_test.py
- tests/load_tests/test_vertex_embeddings_load_test.py
- tests/load_tests/test_vertex_load_tests.py
- reason: >-
Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
locally rather than in pull-request jobs
paths:
- tests/load_tests/test_granian_admission_saturation.py
- reason: >-
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a

38
.github/e2e-stack/assert_tests_ran.py vendored Normal file
View file

@ -0,0 +1,38 @@
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
def main() -> int:
selected: Final = tuple(sys.argv[2:])
try:
report: Final = ET.parse(Path(sys.argv[1])).getroot()
except (ET.ParseError, OSError):
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
missing: Final = tuple(path for path in selected if path not in passed)
for path in selected:
collected: Final = sum(case.get("file") == path for case in cases)
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
for case in cases:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")):
continue
_ = sys.stdout.write(f" failed: {case.get('classname', '')}::{case.get('name', '')}\n")
if (
selected
and not missing
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
):
return 0
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
return 1
if __name__ == "__main__":
sys.exit(main())

17
.github/e2e-stack/down.sh vendored Executable file
View file

@ -0,0 +1,17 @@
#!/usr/bin/env bash
set -uo pipefail
STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}"
for pid_file in "${STACK_DIR}"/pids/*.pid; do
[[ -f "${pid_file}" ]] || continue
pkill -TERM -P "$(cat "${pid_file}")" 2>/dev/null
kill -TERM "$(cat "${pid_file}")" 2>/dev/null
rm -f "${pid_file}"
done
for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do
docker rm -f "${container}" >/dev/null 2>&1
done
exit 0

49
.github/e2e-stack/secrets_to_env.py vendored Normal file
View file

@ -0,0 +1,49 @@
import os
import re
import sys
from pathlib import Path
from typing import Final
from pydantic import TypeAdapter, ValidationError
secrets_adapter: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
ENV_NAME: Final = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
MIN_MASKED_LENGTH: Final = 8
def main() -> int:
env_path: Final = Path(sys.argv[1])
try:
secrets: Final = {
key: value.rstrip("\r\n") for key, value in secrets_adapter.validate_json(sys.stdin.read()).items()
}
except (ValidationError, UnicodeError):
_ = sys.stderr.write("expected a JSON object containing string environment values\n")
return 1
unusable: Final = tuple(
key
for key, value in secrets.items()
if ENV_NAME.fullmatch(key) is None or any(char in value for char in "'\n\r\0")
)
if unusable:
_ = sys.stderr.write(
f"these names or values cannot be represented in both bash and dotenv: {' '.join(sorted(unusable))}\n"
)
return 1
for value in secrets.values():
if len(value) >= MIN_MASKED_LENGTH:
_ = sys.stdout.write(f"::add-mask::{value.replace('%', '%25')}\n")
sys.stdout.flush()
lines: Final = tuple(f"{key}='{value}'" for key, value in secrets.items() if value)
try:
with os.fdopen(os.open(env_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT | os.O_NOFOLLOW, 0o600), "w") as handle:
os.fchmod(handle.fileno(), 0o600)
_ = handle.write("\n".join(lines) + "\n")
except OSError:
_ = sys.stderr.write("could not write the environment file\n")
return 1
return 0
if __name__ == "__main__":
sys.exit(main())

44
.github/e2e-stack/select_tests.py vendored Normal file
View file

@ -0,0 +1,44 @@
import re
import sys
from typing import Final
SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
UNSUPPORTED: Final = re.compile(
r"^tests/e2e/(ui|claude_code|load)/"
r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
)
HARNESS: Final = re.compile(
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
r"|^tests/e2e/gateway/"
r"|^\.github/e2e-stack/"
r"|^\.github/workflows/test-e2e-changed\.yml$"
)
UNEXPANDED: Final = re.compile(r"[*?\[]")
def is_selectable(path: str) -> bool:
return SELECTABLE.match(path) is not None and UNSUPPORTED.match(path) is None
def select(changed: tuple[str, ...], canary: tuple[str, ...]) -> tuple[str, ...]:
direct: Final = frozenset(path for path in changed if is_selectable(path))
harness_changed: Final = any(HARNESS.match(path) for path in changed)
canary_tests: Final = frozenset(path for path in canary if harness_changed and is_selectable(path))
return tuple(sorted(direct | canary_tests))
def main() -> int:
canary: Final = tuple(sys.argv[1:])
unexpanded: Final = tuple(path for path in canary if UNEXPANDED.search(path))
if unexpanded:
_ = sys.stderr.write(f"the canary paths reached the selector unexpanded: {' '.join(unexpanded)}\n")
return 1
changed: Final = tuple(line.strip() for line in sys.stdin if line.strip())
_ = sys.stdout.write(" ".join(select(changed, canary)) + "\n")
return 0
if __name__ == "__main__":
sys.exit(main())

207
.github/e2e-stack/up.sh vendored Executable file
View file

@ -0,0 +1,207 @@
#!/usr/bin/env bash
set -euo pipefail
umask 077
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}"
CERTS_DIR="${STACK_DIR}/certs"
LOGS_DIR="${STACK_DIR}/logs"
PIDS_DIR="${STACK_DIR}/pids"
POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}"
VALKEY_IMAGE="${E2E_VALKEY_IMAGE:-valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa}"
JAEGER_IMAGE="${E2E_JAEGER_IMAGE:-jaegertracing/jaeger:2.10.0}"
NGINX_IMAGE="${E2E_NGINX_IMAGE:-nginx:1.29.1-alpine@sha256:42a516af16b852e33b7682d5ef8acbd5d13fe08fecadc7ed98605ba5e3b26ab8}"
LB_PORT="${E2E_LB_PORT:-4000}"
GATEWAY_PORT_1="${E2E_GATEWAY_PORT_1:-4010}"
GATEWAY_PORT_2="${E2E_GATEWAY_PORT_2:-4011}"
BACKEND_PORT="${E2E_BACKEND_PORT:-4001}"
REDIS_PORT="${E2E_REDIS_PORT:-6379}"
DATABASE_HOST="${E2E_DATABASE_HOST:-127.0.0.1}"
DATABASE_PORT="${E2E_DATABASE_PORT:-5432}"
DATABASE_USER="${E2E_DATABASE_USER:-litellm}"
DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}"
DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}"
JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}"
JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}"
MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}"
mkdir -p "${CERTS_DIR}" "${LOGS_DIR}" "${PIDS_DIR}"
chmod 700 "${STACK_DIR}" "${LOGS_DIR}" "${PIDS_DIR}"
chmod 755 "${CERTS_DIR}"
log() { printf 'e2e-stack: %s\n' "$*"; }
port_open() { (exec 3<>"/dev/tcp/127.0.0.1/$1") 2>/dev/null; }
wait_for() {
local label="$1" check="$2" deadline=$((SECONDS + ${3:-120}))
until eval "${check}"; do
if ((SECONDS >= deadline)); then
log "timed out waiting for ${label}"
exit 1
fi
sleep 2
done
log "${label} is up"
}
if [[ -f "${REPO_ROOT}/tests/e2e/.env" ]]; then
set -a
source "${REPO_ROOT}/tests/e2e/.env"
set +a
fi
if [[ -z "${DD_API_KEY:-}" ]]; then
log "DD_API_KEY is empty; the gateway config enables the datadog callback, so put a Datadog API key in tests/e2e/.env"
exit 1
fi
export DD_SITE="${DD_SITE:-datadoghq.com}"
if ! port_open "${DATABASE_PORT}"; then
docker run -d --name e2e-postgres -p "${DATABASE_PORT}:5432" \
-e "POSTGRES_USER=${DATABASE_USER}" -e "POSTGRES_PASSWORD=${DATABASE_PASSWORD}" -e "POSTGRES_DB=${DATABASE_NAME}" \
"${POSTGRES_IMAGE}" >/dev/null
fi
wait_for "postgres" "port_open ${DATABASE_PORT}"
if ! port_open "${JAEGER_QUERY_PORT}"; then
docker run -d --name e2e-jaeger -p "${JAEGER_OTLP_PORT}:4318" -p "${JAEGER_QUERY_PORT}:16686" \
"${JAEGER_IMAGE}" >/dev/null
fi
wait_for "jaeger" "curl -fs http://127.0.0.1:${JAEGER_QUERY_PORT}/api/services >/dev/null"
openssl genrsa -out "${CERTS_DIR}/ca.key" 2048 2>/dev/null
openssl req -x509 -new -nodes -key "${CERTS_DIR}/ca.key" -sha256 -days 7 \
-subj "/CN=litellm-e2e-ca" \
-addext "basicConstraints=critical,CA:TRUE" -addext "keyUsage=critical,keyCertSign,cRLSign" \
-out "${CERTS_DIR}/ca.crt" 2>/dev/null
openssl genrsa -out "${CERTS_DIR}/server.key" 2048 2>/dev/null
openssl req -new -key "${CERTS_DIR}/server.key" -subj "/CN=localhost" -out "${CERTS_DIR}/server.csr" 2>/dev/null
openssl x509 -req -in "${CERTS_DIR}/server.csr" -CA "${CERTS_DIR}/ca.crt" -CAkey "${CERTS_DIR}/ca.key" \
-CAcreateserial -days 7 -sha256 \
-extfile <(printf 'basicConstraints=CA:FALSE\nkeyUsage=critical,digitalSignature,keyEncipherment\nextendedKeyUsage=serverAuth\nsubjectAltName=DNS:localhost,IP:127.0.0.1\n') \
-out "${CERTS_DIR}/server.crt" 2>/dev/null
chmod 644 "${CERTS_DIR}"/*.key "${CERTS_DIR}"/*.crt
CERTIFI_BUNDLE="$(cd "${REPO_ROOT}" && uv run --no-sync python -c 'import certifi; print(certifi.where())')"
cat "${CERTIFI_BUNDLE}" "${CERTS_DIR}/ca.crt" > "${CERTS_DIR}/ca-bundle.pem"
docker rm -f e2e-valkey >/dev/null 2>&1 || true
docker run -d --name e2e-valkey -p "${REDIS_PORT}:${REDIS_PORT}" -v "${CERTS_DIR}:/certs:ro" \
"${VALKEY_IMAGE}" valkey-server \
--cluster-enabled yes --port 0 --tls-port "${REDIS_PORT}" \
--tls-cert-file /certs/server.crt --tls-key-file /certs/server.key --tls-ca-cert-file /certs/ca.crt \
--tls-auth-clients no --cluster-announce-ip 127.0.0.1 >/dev/null
VALKEY_CLI="docker exec e2e-valkey valkey-cli --tls --cacert /certs/ca.crt -h 127.0.0.1 -p ${REDIS_PORT}"
wait_for "valkey" "${VALKEY_CLI} ping 2>/dev/null | grep -q PONG"
${VALKEY_CLI} cluster addslotsrange 0 16383 >/dev/null
wait_for "valkey cluster" "${VALKEY_CLI} cluster info 2>/dev/null | grep -q cluster_state:ok"
CONFIG_SOURCE="${REPO_ROOT}/tests/e2e/gateway/stage_mirror_ci_config.yml"
CONFIG_PATH="${CONFIG_SOURCE}"
if [[ "${REDIS_PORT}" != "6379" ]]; then
CONFIG_PATH="${STACK_DIR}/litellm-config.yml"
sed "s/port: 6379/port: ${REDIS_PORT}/" "${CONFIG_SOURCE}" > "${CONFIG_PATH}"
fi
SERVER_ENV=(
"LITELLM_MASTER_KEY=${MASTER_KEY}"
"DATABASE_HOST=${DATABASE_HOST}"
"DATABASE_PORT=${DATABASE_PORT}"
"DATABASE_USER=${DATABASE_USER}"
"DATABASE_PASSWORD=${DATABASE_PASSWORD}"
"DATABASE_NAME=${DATABASE_NAME}"
"DISABLE_SCHEMA_UPDATE=true"
"REDIS_HOST=127.0.0.1"
"REDIS_PORT=${REDIS_PORT}"
"REDIS_CLUSTER_NODES=[{\"host\":\"127.0.0.1\",\"port\":${REDIS_PORT}}]"
"CONFIG_FILE_PATH=${CONFIG_PATH}"
"STORE_MODEL_IN_DB=True"
"OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf"
"OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}"
"SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem"
"PYTHONPATH=${REPO_ROOT}"
)
if [[ -n "${VERTEXAI_CREDENTIALS:-}" ]]; then
printf '%s' "${VERTEXAI_CREDENTIALS}" > "${STACK_DIR}/vertex-adc.json"
SERVER_ENV+=("GOOGLE_APPLICATION_CREDENTIALS=${STACK_DIR}/vertex-adc.json")
fi
cd "${REPO_ROOT}"
log "running migrations"
env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/migrations.log" 2>&1
start_server() {
local name="$1"; shift
env "${SERVER_ENV[@]}" "$@" >"${LOGS_DIR}/${name}.log" 2>&1 &
echo $! > "${PIDS_DIR}/${name}.pid"
}
start_server backend uv run --no-sync uvicorn backend.main:app --host 0.0.0.0 --port "${BACKEND_PORT}"
start_server gateway-1 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_1}"
start_server gateway-2 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_2}"
if [[ "$(uname)" == "Linux" ]]; then
NGINX_UPSTREAM_HOST=127.0.0.1
NGINX_DOCKER_ARGS=(--network host)
else
NGINX_UPSTREAM_HOST=host.docker.internal
NGINX_DOCKER_ARGS=(-p "${LB_PORT}:${LB_PORT}")
fi
cat > "${STACK_DIR}/nginx.conf" <<EOF
events {}
http {
map \$http_upgrade \$connection_upgrade {
default upgrade;
'' close;
}
upstream litellm_gateways {
server ${NGINX_UPSTREAM_HOST}:${GATEWAY_PORT_1};
server ${NGINX_UPSTREAM_HOST}:${GATEWAY_PORT_2};
}
server {
listen ${LB_PORT};
client_max_body_size 100m;
location / {
proxy_pass http://litellm_gateways;
proxy_http_version 1.1;
proxy_set_header Host \$host;
proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
proxy_set_header X-Forwarded-Proto \$scheme;
proxy_set_header Upgrade \$http_upgrade;
proxy_set_header Connection \$connection_upgrade;
proxy_buffering off;
proxy_read_timeout 600s;
proxy_send_timeout 600s;
}
}
}
EOF
docker rm -f e2e-nginx >/dev/null 2>&1 || true
docker run -d --name e2e-nginx "${NGINX_DOCKER_ARGS[@]}" \
-v "${STACK_DIR}/nginx.conf:/etc/nginx/nginx.conf:ro" "${NGINX_IMAGE}" >/dev/null
wait_for "backend" "curl -fs http://127.0.0.1:${BACKEND_PORT}/health/liveliness >/dev/null" 300
wait_for "gateway-1" "curl -fs http://127.0.0.1:${GATEWAY_PORT_1}/health/liveliness >/dev/null" 300
wait_for "gateway-2" "curl -fs http://127.0.0.1:${GATEWAY_PORT_2}/health/liveliness >/dev/null" 300
wait_for "load balancer" "curl -fs http://127.0.0.1:${LB_PORT}/health/liveliness >/dev/null" 60
cat > "${STACK_DIR}/stack.env" <<EOF
LITELLM_PROXY_URL=http://127.0.0.1:${LB_PORT}
LITELLM_CONTROL_PLANE_URL=http://127.0.0.1:${BACKEND_PORT}
LITELLM_PROXY_REPLICA_URLS=http://127.0.0.1:${GATEWAY_PORT_1},http://127.0.0.1:${GATEWAY_PORT_2}
LITELLM_MASTER_KEY=${MASTER_KEY}
REDIS_HOST=127.0.0.1
REDIS_PORT=${REDIS_PORT}
E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem
DATABASE_URL=postgresql://${DATABASE_USER}:${DATABASE_PASSWORD}@${DATABASE_HOST}:${DATABASE_PORT}/${DATABASE_NAME}
EOF
log "stack is up; pytest env written to ${STACK_DIR}/stack.env"

22
.github/workflows/_test-unit-base.yml vendored
View file

@ -55,6 +55,9 @@ on:
permissions:
contents: read
env:
UV_PYTHON: "3.12"
jobs:
run:
name: Run tests
@ -82,7 +85,7 @@ jobs:
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
python-version: ${{ env.UV_PYTHON }}
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
@ -96,12 +99,10 @@ jobs:
timeout-minutes: 5
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
path: ${{ env.UV_CACHE_DIR }}
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
@ -112,7 +113,8 @@ jobs:
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra mongodb
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -134,12 +136,6 @@ jobs:
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DIST: ${{ inputs.dist }}
# coverage.py's sys.monitoring backend (PEP 669), the cheapest core it has.
# It is only the default from Python 3.14, and these shards run 3.12, so it
# has to be asked for. Coverage refuses it when branch measurement is on
# (`branch_right_left` needs > 3.14.0a5) and falls back to the slow core with
# a `no-sysmon` warning, so turning on `branch = true` here means giving this
# back until the runners move to 3.14.
COVERAGE_CORE: sysmon
run: |
if [ "${WORKERS}" = "0" ]; then

2
.github/workflows/auto-close-duplicates.yml vendored
View file

@ -22,7 +22,7 @@ on:
permissions: {}
jobs:
test:
sweep-tests:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5

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

@ -0,0 +1,45 @@
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:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
cancel-in-progress: true
jobs:
cost-map-guard:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Fetch the pull request head and its merge base
id: revisions
env:
GH_TOKEN: ${{ github.token }}
BASE_SHA: ${{ github.event.pull_request.base.sha }}
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"
echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Run the guard
env:
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
HEAD_REF: ${{ github.event.pull_request.head.ref }}
run: |
uv run --frozen python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"

130
.github/workflows/report-rust-release-wheel.yml vendored
View file

@ -1,130 +0,0 @@
name: Report LiteLLM Rust release wheel
on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs
workflow_run:
workflows:
- LiteLLM Rust
types:
- completed
permissions: {}
concurrency:
group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }}
cancel-in-progress: false
jobs:
report-release-wheel:
name: report release wheel
if: >-
github.event.workflow_run.event == 'pull_request' &&
github.event.workflow_run.path == '.github/workflows/test-rust.yml' &&
github.event.workflow_run.head_repository.full_name == github.repository &&
github.event.workflow_run.pull_requests[0].number != null
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
issues: write # PR comments use the issues API
pull-requests: read # Current-head validation rejects stale workflow runs
steps:
- name: Link release wheel report on PR
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
env:
COMMENT_MARKER: "<!-- litellm-release-wheel-size -->"
with:
script: |
const marker = process.env.COMMENT_MARKER;
const workflowRun = context.payload.workflow_run;
const allowedConclusions = new Set([
"action_required",
"cancelled",
"failure",
"neutral",
"skipped",
"stale",
"startup_failure",
"success",
"timed_out",
]);
if (
!allowedConclusions.has(workflowRun.conclusion) ||
workflowRun.event !== "pull_request" ||
workflowRun.path !== ".github/workflows/test-rust.yml" ||
workflowRun.head_repository?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
workflowRun.pull_requests?.length !== 1
) {
throw new Error("unexpected source workflow");
}
const pullRequest = workflowRun.pull_requests[0];
const pullRequestNumber = pullRequest.number;
const headSha = workflowRun.head_sha;
const runId = workflowRun.id;
if (
!Number.isSafeInteger(pullRequestNumber) ||
pullRequestNumber <= 0 ||
!Number.isSafeInteger(runId) ||
runId <= 0 ||
!/^[0-9a-f]{40}$/.test(headSha) ||
pullRequest.head?.sha !== headSha
) {
throw new Error("invalid source workflow metadata");
}
const runUrl =
`${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` +
`/actions/runs/${runId}`;
const result =
workflowRun.conclusion === "success"
? "successfully"
: `with \`${workflowRun.conclusion}\``;
const body = [
marker,
"## LiteLLM Rust workflow",
"",
`Workflow completed ${result} for \`${headSha}\``,
"",
`[View workflow run](${runUrl})`,
].join("\n");
const comments = await github.paginate(github.rest.issues.listComments, {
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
per_page: 100,
});
const existing = comments.find(
(comment) =>
comment.user?.login === "github-actions[bot]" &&
comment.body?.startsWith(marker),
);
const currentPullRequest = (
await github.rest.pulls.get({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: pullRequestNumber,
})
).data;
if (
currentPullRequest.state !== "open" ||
currentPullRequest.head.repo?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
currentPullRequest.head.sha !== headSha
) {
core.info("source workflow no longer matches the current pull request head");
return;
}
if (existing) {
await github.rest.issues.updateComment({
owner: context.repo.owner,
repo: context.repo.repo,
comment_id: existing.id,
body,
});
} else {
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
body,
});
}

9
.github/workflows/test-code-quality.yml vendored
View file

@ -74,6 +74,15 @@ jobs:
- name: check_workflow_startup_safety
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py
- name: check_workflow_job_name_collisions
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py
- name: test_workflow_job_name_collisions
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py

240
.github/workflows/test-e2e-changed.yml vendored Normal file
View file

@ -0,0 +1,240 @@
name: e2e-changed-tests
on:
pull_request:
concurrency:
group: e2e-changed-${{ github.event.pull_request.number }}
cancel-in-progress: true
permissions: {}
jobs:
detect:
name: Detect changed e2e tests
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
pull-requests: read
outputs:
tests: ${{ steps.changed.outputs.tests }}
any: ${{ steps.changed.outputs.any }}
steps:
- name: Checkout the selector and the canary suite
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: |
.github/e2e-stack
tests/e2e/access_control
persist-credentials: false
ref: ${{ github.sha }}
- name: List the e2e test files this PR added or modified
id: changed
env:
GH_TOKEN: ${{ github.token }}
REPO: ${{ github.repository }}
PR_NUMBER: ${{ github.event.pull_request.number }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
gh api "repos/${REPO}/pulls/${PR_NUMBER}" \
--jq 'select(.head.sha == env.HEAD_SHA and .changed_files < 3000) | .head.sha' \
| grep -Fxq "${HEAD_SHA}"
files="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}/files" --paginate \
--jq '.[] | select(.status != "removed") | .filename')"
gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}"
tests="$(printf '%s\n' "${files}" \
| python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py)"
echo "tests=${tests}" >> "${GITHUB_OUTPUT}"
if [ -n "${tests}" ]; then
echo "any=true" >> "${GITHUB_OUTPUT}"
echo "selected e2e tests: ${tests}"
else
echo "any=false" >> "${GITHUB_OUTPUT}"
echo "no changed e2e test files supported by this stack; nothing to run"
fi
run:
name: Run changed e2e tests against the stage-mirror stack
needs: detect
if: needs.detect.outputs.any == 'true' && github.event.pull_request.head.repo.full_name == github.repository
runs-on: ubuntu-latest
timeout-minutes: 90
environment: e2e-changed
permissions:
contents: read
id-token: write
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U litellm"
--health-interval 5s
--health-timeout 5s
--health-retries 10
jaeger:
image: jaegertracing/jaeger:2.10.0
ports:
- 4318:4318
- 16686:16686
steps:
- name: Validate configuration
env:
ROLE: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
run: test -n "${ROLE}" || { echo "::error::Set repo variable E2E_AWS_ROLE_TO_ASSUME to an OIDC role with read access to the e2e secrets"; exit 1; }
- name: Checkout
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
ref: ${{ github.sha }}
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.13"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen \
--extra proxy --extra proxy-runtime --extra extra_proxy \
--extra semantic-router --extra bedrock-realtime \
--group ci --group proxy-dev --group e2e-dev
uv pip install "pipecat-ai[openai]==1.4.0"
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Install Playwright chromium
run: uv run --no-sync playwright install --with-deps chromium
- name: Configure AWS credentials
id: aws
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
with:
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
aws-region: us-east-1
role-session-name: litellm-e2e-changed-${{ github.run_id }}
role-duration-seconds: 900
output-env-credentials: false
output-credentials: true
- name: Fetch provider credentials from AWS Secrets Manager
env:
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
AWS_DEFAULT_REGION: us-east-1
run: |
umask 077
aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-provider-keys \
--query SecretString --output text \
| uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env
aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license \
--query SecretString --output text \
| jq -R -s '{"LITELLM_LICENSE": .}' \
| uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env
- name: Boot the stage-mirror stack
id: boot
run: |
umask 077
if ! bash .github/e2e-stack/up.sh > "${RUNNER_TEMP}/e2e-boot.log" 2>&1; then
echo "::error::stage-mirror stack failed to boot; raw logs are not published"
exit 1
fi
- name: Export stack environment
run: |
master_key="$(grep '^LITELLM_MASTER_KEY=' "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" | cut -d= -f2-)"
echo "::add-mask::${master_key}"
cat "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" >> "${GITHUB_ENV}"
- name: Run the selected tests three times
env:
TESTS: ${{ needs.detect.outputs.tests }}
E2E_FIXTURE_MODE: live
run: |
umask 077
read -r -a test_files <<< "${TESTS}"
for pass in 1 2 3; do
report="${RUNNER_TEMP}/e2e-pass-${pass}.xml"
log="${RUNNER_TEMP}/e2e-pass-${pass}.log"
echo "::group::pass ${pass} of 3"
set +e
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \
-o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1
status=$?
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"
verified=$?
set -e
grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1
echo "::endgroup::"
if [ "${status}" = "5" ]; then
echo "::error::the selected files collected no runnable tests, so nothing was verified"
exit 1
fi
if [ "${status}" != "0" ]; then
echo "::error::pass ${pass} of 3 failed with exit code ${status}"
exit "${status}"
fi
if [ "${verified}" != "0" ]; then
echo "::error::pass ${pass} of 3 did not verify every selected file"
exit 1
fi
echo "pass ${pass} of 3 passed"
done
- name: Stop the stack
if: always() && steps.boot.outcome != 'skipped'
run: bash .github/e2e-stack/down.sh
- name: Remove credentials and raw output
if: always()
run: |
rm -f tests/e2e/.env "${RUNNER_TEMP}/e2e-boot.log" "${RUNNER_TEMP}"/e2e-pass-*.log "${RUNNER_TEMP}"/e2e-pass-*.xml
rm -rf "${RUNNER_TEMP}/litellm-e2e-stack"
gate:
name: e2e-changed-tests
needs: [detect, run]
if: always()
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- name: Require three successful passes when tests changed
env:
DETECT_RESULT: ${{ needs.detect.result }}
ANY_TESTS: ${{ needs.detect.outputs.any }}
RUN_RESULT: ${{ needs.run.result }}
run: |
if [ "${DETECT_RESULT}" != "success" ]; then
echo "::error::changed-test detection did not succeed"
exit 1
fi
if [ "${ANY_TESTS}" = "false" ]; then
echo "no changed e2e test files supported by this stack; nothing to run"
exit 0
fi
if [ "${ANY_TESTS}" != "true" ] || [ "${RUN_RESULT}" != "success" ]; then
echo "::error::selected e2e tests require an approved, successful run; fork PRs must run from a reviewed same-repository branch"
exit 1
fi

37
.github/workflows/test-model-map.yml vendored
View file

@ -1,37 +0,0 @@
name: Validate model_prices_and_context_window.json
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
validate-model-prices-json:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Validate model_prices_and_context_window.json
run: |
jq empty model_prices_and_context_window.json
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Check model_prices_and_context_window.schema.json is in sync
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py --check

111
.github/workflows/test-rust.yml vendored
View file

@ -7,6 +7,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -22,6 +23,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -34,102 +36,89 @@ concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
env:
CARGO_TERM_COLOR: always
jobs:
rust-checks:
name: rustfmt, clippy, test
rust-lint:
runs-on: ubuntu-latest
timeout-minutes: 10
defaults:
run:
working-directory: litellm-rust
env:
CARGO_TERM_COLOR: always
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Cache Cargo registry and target
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- run: cargo fmt --check
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }}
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-
${{ runner.os }}-cargo-${{ github.job }}-
- name: Check Rust formatting
run: cargo fmt --check
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
- name: Run Clippy
run: cargo clippy --workspace --all-targets --locked -- -D warnings
- run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- name: Run Clippy with Bedrock auth
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Clippy with all gateway features
run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Rust tests
run: cargo test --workspace --locked
- name: Run core tests with Bedrock auth
run: cargo test -p litellm-core --features bedrock-auth --locked
# Not --all-features: python-config links libpython, which this job does not install.
- name: Run gateway tests with the server feature
run: cargo test -p litellm-ai-gateway --features server --locked
release-wheel:
name: release wheel
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 20
permissions:
contents: read
env:
CARGO_TERM_COLOR: always
timeout-minutes: 30
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
- uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Build release wheel
run: uv build --wheel --out-dir dist
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
- name: Build panic contract wheel
run: >-
- run: cargo test --workspace --locked
working-directory: litellm-rust
- run: cargo test -p litellm-core --features bedrock-auth --locked
working-directory: litellm-rust
- run: cargo test -p litellm-ai-gateway --features server --locked
working-directory: litellm-rust
- run: uv build --wheel --out-dir dist
- run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: >-
uv build --wheel --out-dir panic-dist
--config-setting "maturin.build-args=--features panic-test,extension-module"
- name: Smoke-test native panic unwinding
run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl
- name: Verify stripped native extension
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
- name: Test native route wheel
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl

31
.github/workflows/test-terraform-modules.yml vendored
View file

@ -4,6 +4,7 @@ on:
push:
paths:
- "terraform/litellm/aws/**"
- "terraform/litellm/gcp/**"
- ".github/workflows/test-terraform-modules.yml"
pull_request:
branches:
@ -13,6 +14,7 @@ on:
- "litellm_**"
paths:
- "terraform/litellm/aws/**"
- "terraform/litellm/gcp/**"
- ".github/workflows/test-terraform-modules.yml"
permissions:
@ -52,3 +54,32 @@ jobs:
# Plan-only, mock_provider-backed: no AWS credentials, no API calls.
- name: test
run: terraform test
gcp-module:
name: fmt, validate, test (gcp)
runs-on: ubuntu-latest
timeout-minutes: 15
defaults:
run:
working-directory: terraform/litellm/gcp
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2
with:
terraform_version: 1.13.3
terraform_wrapper: false
- name: fmt
run: terraform fmt -recursive -check -diff
- name: init
run: terraform init -backend=false -input=false
- name: validate
run: terraform validate
- name: test
run: terraform test

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

@ -96,7 +96,6 @@ jobs:
- shard: misc
artifact-name: misc
test-path: >-
tests/sdk_function_trace
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol

4
CLAUDE.md
View file

@ -37,7 +37,7 @@ 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 leave the section blank
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, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
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
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
@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on `litellm_internal_staging` in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice

2
Dockerfile
View file

@ -67,6 +67,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Copy full source tree
@ -89,6 +90,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

36
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 14074
"limit": 13429
},
"reportArgumentType": {
"limit": 2206
"limit": 2198
},
"reportAssignmentType": {
"limit": 319
@ -18,13 +18,13 @@
"limit": 40
},
"reportDeprecated": {
"limit": 211
"limit": 209
},
"reportDuplicateImport": {
"limit": 19
},
"reportExplicitAny": {
"limit": 4124
"limit": 3369
},
"reportFunctionMemberAccess": {
"limit": 7
@ -45,19 +45,19 @@
"limit": 24
},
"reportInvalidTypeForm": {
"limit": 34
"limit": 30
},
"reportInvalidTypeVarUse": {
"limit": 2
"limit": 1
},
"reportMatchNotExhaustive": {
"limit": 0
},
"reportMissingParameterType": {
"limit": 5601
"limit": 5570
},
"reportMissingTypeArgument": {
"limit": 15285
"limit": 15281
},
"reportMissingTypeStubs": {
"limit": 40
@ -84,37 +84,37 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1808
"limit": 1804
},
"reportRedeclaration": {
"limit": 8
},
"reportReturnType": {
"limit": 181
"limit": 180
},
"reportTypedDictNotRequiredAccess": {
"limit": 24
"limit": 22
},
"reportUndefinedVariable": {
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44360
"limit": 44358
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38311
"limit": 38271
},
"reportUnknownParameterType": {
"limit": 19624
"limit": 19584
},
"reportUnknownVariableType": {
"limit": 29847
"limit": 29814
},
"reportUnnecessaryCast": {
"limit": 111
"limit": 110
},
"reportUnnecessaryComparison": {
"limit": 687
@ -123,7 +123,7 @@
"limit": 4
},
"reportUnnecessaryIsInstance": {
"limit": 819
"limit": 816
},
"reportUntypedBaseClass": {
"limit": 0
@ -135,7 +135,7 @@
"limit": 21
},
"reportUnusedFunction": {
"limit": 138
"limit": 136
},
"reportUnusedImport": {
"limit": 542

146
ci_cd/cost_map_guard.py Normal file
View file

@ -0,0 +1,146 @@
"""Guard the cost map on pull requests.
Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file,
and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named
litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models.
"""
from __future__ import annotations
import argparse
import json
import subprocess
import sys
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
CostMap = dict[str, object]
@dataclass(frozen=True, slots=True)
class Snapshot:
cost_map: str
backup: str
schema: str
def _parse_object(text: str, path: str) -> CostMap | str:
try:
parsed: Final = json.loads(text)
except json.JSONDecodeError as error:
return f"{path} is not valid JSON: {error}"
return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
def _rendered_schema(cost_map: CostMap) -> str:
try:
return render(build_schema(cost_map))
except SystemExit as error:
return str(error)
def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]:
schema_text: Final = _rendered_schema(head_map)
if not schema_text.startswith("{"):
return (schema_text,)
backup_failure: Final = (
()
if head.backup == head.cost_map
else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
)
schema_failure: Final = (
()
if head.schema == schema_text
else (
f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
"run `python ci_cd/generate_model_prices_schema.py` and commit the result",
)
)
return (
*backup_failure,
*schema_failure,
*(
f"{COST_MAP_PATH} does not validate against its schema: {error}"
for error in validation_errors(head_map, json.loads(schema_text))[:20]
),
)
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 _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):
return (f"merge base: {base_map}",)
base_entries: Final = _entries(base_map)
head_entries: Final = _entries(head_map)
removed_fields: Final = tuple(
f"{key}.{field}"
for key, entry in base_entries.items()
if key in head_entries
for field in entry
if field not in head_entries[key]
)
return (
*(
f"bot PRs may only change the cost map files, not {path}"
for path in changed_files
if path not in GUARDED_PATHS
),
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),
*(
f"bot PRs may not change {key}"
for key in sorted(SPECIAL_ROOT_KEYS)
if base_map.get(key) != head_map.get(key)
),
)
def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]:
head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH)
if isinstance(head_map, str):
return (head_map,)
return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ()))
def _git(*args: str) -> str:
result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True)
return result.stdout if result.returncode == 0 else ""
def snapshot(revision: str) -> Snapshot:
return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS))
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")
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)
changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines())
failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot)
contract: Final = "bot contract enforced" if bot else "human PR, file checks only"
if failures:
print(f"cost map guard failed ({contract}):")
print("\n".join(f"- {failure}" for failure in failures))
return 1
print(f"cost map guard passed ({contract})")
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))

9
ci_cd/generate_model_prices_schema.py
View file

@ -58,6 +58,11 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
}
ARRAY_KEYS: dict[str, JsonSchema] = {
"supported_audio_formats": {
"type": "array",
"description": "Audio container formats the model can return.",
"items": {"type": "string", "enum": ["mp3", "wav"]},
},
"supported_endpoints": {
"type": "array",
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
@ -231,6 +236,10 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
},
"comment": STRING,
"audio_transcription_config": STRING,
"vertex_ai_audio_api": {
"type": "string",
"enum": ["lyria_predict", "lyria_interactions"],
},
}

2
docker/Dockerfile.database
View file

@ -65,6 +65,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Copy full source tree
@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

3
docker/Dockerfile.non_root
View file

@ -71,6 +71,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Copy full source tree
@ -99,6 +100,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13 \
--no-sources-package litellm-proxy-extras; \
else \
@ -109,6 +111,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13; \
fi

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

@ -433,9 +433,9 @@ _default_detect_secrets_config = {
"name": "ZendeskSecretKeyDetector",
"path": _custom_plugins_path + "/zendesk_secret_key.py",
},
{"name": "Base64HighEntropyString", "limit": 3.0},
{"name": "Base64HighEntropyString", "limit": 4.5},
{"name": "HexHighEntropyString", "limit": 3.0},
]
],
}
@ -466,16 +466,19 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
os.remove(temp_file.name)
detected_secrets = []
for file in secrets.files:
for found_secret in secrets[file]:
if found_secret.secret_value is None:
continue
detected_secrets.append(
{"type": found_secret.type, "value": found_secret.secret_value}
)
return detected_secrets
return [
{"type": found_secret.type, "value": found_secret.secret_value}
for file in sorted(secrets.files)
for found_secret in sorted(
secrets[file],
key=lambda secret: (
-len(secret.secret_value or ""),
secret.type,
secret.secret_value or "",
),
)
if found_secret.secret_value is not None
]
def redact_text(self, text: str, source: str = "message") -> str:
"""Replace every detected secret in ``text`` with ``[REDACTED]`` and

15
enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py
View file

@ -3,6 +3,7 @@ This plugin searches for OpenAI API Keys.
"""
import re
from collections.abc import Generator
from detect_secrets.plugins.base import RegexBasedDetector
@ -16,4 +17,16 @@ class OpenAIApiKeyDetector(RegexBasedDetector):
@property
def denylist(self) -> list[re.Pattern]:
return [re.compile(r"""(sk-[a-zA-Z0-9]{5,})""")]
return [
re.compile(
r"((?:(?<![a-zA-Z0-9])|(?<=%[0-9A-Fa-f]{2}))"
r"sk[-_]"
r"[a-zA-Z0-9_-]{5,}"
r"(?![a-zA-Z0-9_-]))"
)
]
def analyze_string(self, string: str) -> Generator[str, None, None]:
# the digit check lives outside the regex: a lookahead re-scans the token
# from every `sk` inside it, which is quadratic on `-sk-sk-sk-...` input
yield from (match for match in super().analyze_string(string) if re.search(r"[0-9]", match))

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

@ -161,6 +161,7 @@ class CheckBatchCost:
metadata: dict[str, object] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_hash": api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}

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

@ -46,6 +46,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
BATCH_CREATE_HIDDEN_PARAM,
FILE_LIST_CONTINUATION_CHUNK_SIZE,
MAX_FILE_LIST_LIMIT,
_is_base64_encoded_unified_file_id,
@ -1321,7 +1322,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
## Check if unified_file_id is in the response
unified_file_id = response._hidden_params.get("unified_file_id") # managed file id
unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id
is_batch_create: Final = unified_file_id is not None
is_batch_create: Final = response._hidden_params.get(BATCH_CREATE_HIDDEN_PARAM) is True
model_id = cast(Optional[str], response._hidden_params.get("model_id"))
model_name = cast(Optional[str], response._hidden_params.get("model_name"))
@ -1410,10 +1411,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
# Only record batch creation metric on actual create (not retrieve/cancel).
# unified_file_id in _hidden_params is only set by the create_batch endpoint.
original_unified_file_id = response._hidden_params.get("unified_file_id")
if original_unified_file_id:
if is_batch_create:
prom_logger = self._get_prometheus_logger()
if prom_logger:
batch_provider = ""

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

@ -16,6 +16,7 @@ from litellm.llms.base_llm.managed_resources.utils import (
is_base64_encoded_unified_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import LLMResponseTypes
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreCreateResponse,
@ -24,6 +25,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.caching.caching import DualCache
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
@ -156,7 +158,7 @@ class _PROXY_LiteLLMManagedVectorStores(
# Create vector store for each model
# Convert TypedDict to Dict[str, Any] for base class compatibility
request_data_dict: Dict[str, Any] = dict(create_request)
request_data_dict: Dict[str, object] = dict(create_request)
responses = await self.create_resource_for_each_model(
llm_router=llm_router,
request_data=request_data_dict,
@ -209,7 +211,7 @@ class _PROXY_LiteLLMManagedVectorStores(
limit: Optional[int] = None,
after: Optional[str] = None,
order: Optional[str] = None,
) -> Dict[str, Any]:
) -> Dict[str, object]:
"""
List vector stores created by a user.
@ -301,7 +303,7 @@ class _PROXY_LiteLLMManagedVectorStores(
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: Any,
cache: "DualCache",
data: Dict,
call_type: str,
) -> Union[Exception, str, Dict, None]:
@ -403,8 +405,8 @@ class _PROXY_LiteLLMManagedVectorStores(
self,
data: Dict,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
) -> Any:
response: LLMResponseTypes,
) -> LLMResponseTypes:
"""
Post-call hook to transform responses.

4
enterprise/pyproject.toml
View file

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

2
gateway/Dockerfile
View file

@ -47,6 +47,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
@ -59,6 +60,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--extra mongodb \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

12
helm/litellm/templates/migrations-job.yaml
View file

@ -77,4 +77,16 @@ spec:
volumes:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.migrationJob.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.migrationJob.affinity }}
affinity:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.migrationJob.tolerations }}
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

68
helm/litellm/tests/migration_job_tests.yaml
View file

@ -1,4 +1,4 @@
suite: test migrations Job ServiceAccount resolution and pod hardening
suite: test migrations Job ServiceAccount resolution, pod hardening, and scheduling
templates:
- migrations-job.yaml
values:
@ -188,3 +188,69 @@ tests:
asserts:
- notExists:
path: spec.activeDeadlineSeconds
- it: renders no scheduling fields by default
asserts:
- isNull:
path: spec.template.spec.nodeSelector
- isNull:
path: spec.template.spec.tolerations
- isNull:
path: spec.template.spec.affinity
- it: renders nodeSelector, tolerations, and affinity from the migrationJob values
set:
migrationJob.nodeSelector:
intent: no-csi-nodes
migrationJob.tolerations:
- key: intent
operator: Equal
value: no-csi-nodes
effect: NoSchedule
migrationJob.affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: intent
operator: In
values:
- no-csi-nodes
asserts:
- equal:
path: spec.template.spec.nodeSelector
value:
intent: no-csi-nodes
- equal:
path: spec.template.spec.tolerations
value:
- key: intent
operator: Equal
value: no-csi-nodes
effect: NoSchedule
- equal:
path: spec.template.spec.affinity
value:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: intent
operator: In
values:
- no-csi-nodes
- it: does not inherit the gateway's scheduling values
set:
gateway.nodeSelector:
intent: no-csi-nodes
gateway.tolerations:
- key: intent
operator: Equal
value: no-csi-nodes
effect: NoSchedule
asserts:
- isNull:
path: spec.template.spec.nodeSelector
- isNull:
path: spec.template.spec.tolerations

7
helm/litellm/values.yaml
View file

@ -152,6 +152,13 @@ migrationJob:
# the writable scratch space a read-only root filesystem needs.
volumes: []
volumeMounts: []
# Scheduling for the Job pod, same shape as gateway.nodeSelector /
# gateway.tolerations / gateway.affinity. The Job does not inherit the other
# components' scheduling values: a migration usually needs a larger node
# than the gateway, so pin it here explicitly.
nodeSelector: {}
tolerations: []
affinity: {}
image:
repository: ghcr.io/berriai/litellm-migrations
tag: "" # defaults to .Chart.AppVersion

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000001_add_guardrail_usage_units_cost/migration.sql Normal file
View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "cost" DOUBLE PRECISION;
ALTER TABLE "LiteLLM_DailyGuardrailUsageUnits" ADD COLUMN IF NOT EXISTS "untracked_units" BIGINT NOT NULL DEFAULT 0;

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260904000000_shadow_eval_model_scope/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "models" TEXT[] NOT NULL DEFAULT ARRAY[]::TEXT[];

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false;

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

@ -5,6 +5,7 @@ datasource client {
generator client {
provider = "prisma-client-py"
recursive_type_depth = -1
binaryTargets = ["native", "debian-openssl-1.1.x", "debian-openssl-3.0.x", "linux-musl", "linux-musl-openssl-3.0.x"]
}
@ -342,6 +343,7 @@ model LiteLLM_MCPServerTable {
delegate_auth_to_upstream Boolean @default(false)
oauth_passthrough Boolean @default(false)
dcr_bridge Boolean?
per_server_oauth_discovery Boolean @default(false)
is_byok Boolean @default(false)
byok_description String[] @default([])
byok_api_key_help_url String?
@ -1123,6 +1125,8 @@ model LiteLLM_DailyGuardrailUsageUnits {
api_key String // hashed virtual key; empty string when unknown
usage_unit String // provider counter name, e.g. Bedrock's contentPolicyUnits
units BigInt @default(0)
cost Float? // USD for the priced share of units; null only on rows written before this column existed
untracked_units BigInt @default(0) // units recorded with no known price, the share cost leaves out
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@ -1533,6 +1537,7 @@ model LiteLLM_ShadowEvalJob {
target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows
router_name String // first (often only) auto-router under evaluation; router_names is the full set
router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name)
models String[] @default([]) // model groups the sampled traffic is narrowed to; empty samples every model
direction String @default("forward") // forward | reverse
baseline_model String? // reverse only: the fixed model the router is judged against
judge_model String

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

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

5
litellm-rust/AGENTS.md
View file

@ -1,17 +1,18 @@
# AGENTS.md
litellm-rust has four crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
litellm-rust has five crates. A crate is a layer or shared foundation, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are modules inside the layers.
## Crates
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust. One public entrypoint per top-level call (`messages::messages()`), owning types, transforms, provider resolution, auth, and the provider HTTP call. Call it, get a typed response. |
| litellm-config | Config-loading boundary. Returns resolved core deployment data and optionally delegates loading to Python. |
| litellm-ai-gateway | The axum server (behind the `server` feature) plus the WebSocket hosts. Translates HTTP/WS to core entrypoints; owns no provider logic and no handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
Dependency direction is acyclic: `litellm-config` depends on `litellm-core`, the gateway depends on both, and `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`. The interop foundation depends on no LiteLLM domain crate.
## Where a route lives

1
litellm-rust/CLAUDE.md
View file

@ -24,6 +24,7 @@ the base when behavior is genuinely different, and say so explicitly in the PR.
## Crates (see AGENTS.md)
`litellm-core` **is** the LiteLLM SDK in Rust: it makes the LLM call.
`litellm-config` is the config-loading boundary and returns resolved core types.
`litellm-ai-gateway` is an HTTP/WebSocket server in front of it, and
`litellm-python-bridge` exposes it to the Python SDK. `litellm-python-interop`
holds domain-neutral PyO3 primitives shared by Python-facing Rust code. A crate

35
litellm-rust/Cargo.lock generated
View file

@ -1412,9 +1412,11 @@ dependencies = [
"base64",
"futures-channel",
"futures-util",
"litellm-config",
"litellm-core",
"pyo3",
"reqwest",
"rustls 0.23.42",
"rustls-native-certs",
"serde",
"serde_json",
"sha2 0.10.9",
@ -1425,6 +1427,16 @@ dependencies = [
"tracing",
]
[[package]]
name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-core",
"pyo3",
"serde_json",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-core"
version = "0.1.0"
@ -1435,14 +1447,17 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"base64",
"rand 0.8.7",
"reqwest",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -1461,7 +1476,6 @@ dependencies = [
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
]
[[package]]
@ -1511,6 +1525,16 @@ version = "0.3.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
[[package]]
name = "mime_guess"
version = "2.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
dependencies = [
"mime",
"unicase",
]
[[package]]
name = "mio"
version = "1.2.2"
@ -1969,6 +1993,7 @@ dependencies = [
"hyper-util",
"js-sys",
"log",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
@ -2736,6 +2761,12 @@ version = "1.20.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20"
[[package]]
name = "unicase"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
[[package]]
name = "unicode-ident"
version = "1.0.24"

6
litellm-rust/Cargo.toml
View file

@ -1,6 +1,7 @@
[workspace]
members = [
"crates/core",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
"crates/python-bridge",
@ -17,6 +18,7 @@ repository = "https://github.com/BerriAI/litellm"
tracing = "0.1"
tracing-subscriber = { version = "0.3", default-features = false, features = ["registry", "std"] }
litellm-core = { path = "crates/core" }
litellm-config = { path = "crates/config" }
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
litellm-python-interop = { path = "crates/python-interop" }
axum = "0.7"
@ -24,8 +26,10 @@ pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
sha2 = "0.10"

4
litellm-rust/README.md
View file

@ -25,11 +25,12 @@ coverage and production evidence.
| Crate | Role |
|-------|------|
| litellm-core | The SDK. Per-route entrypoints (`messages::messages()`), types, provider transforms (modules under `providers/`), provider resolution, auth, the provider HTTP call, and the router. |
| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
| litellm-ai-gateway | The axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. Owns API registration, domain wiring, and Python exception mapping. |
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
## Layout
@ -38,6 +39,7 @@ crates/
core/ The SDK: route modules + provider transforms.
src/messages/ mod.rs (entrypoint), types, transformation, prepare, handler, client
src/providers/anthropic/messages/transformation.rs
config/ Config loading and resolved deployments.
ai-gateway/ Axum server + WebSocket hosts; calls core entrypoints.
python-interop/ Domain-neutral PyO3 conversion and GIL primitives.
python-bridge/ PyO3 API adapter for Python LiteLLM.

2
litellm-rust/crates/CODING_STANDARDS/PROVIDER_CODING_STANDARDS.md
View file

@ -45,7 +45,7 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. `messages`
22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
24. Do not add new feature flags unless explicitly requested. Reuse the existing LiteLLM Rust rollout mechanism (`litellm.rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
## Checks before push

9
litellm-rust/crates/ai-gateway/AGENTS.md
View file

@ -9,19 +9,15 @@ such as `litellm_core::messages::messages`. No provider handler lives here.
src/
main.rs # entrypoint: build AppState (router + master key), bind, serve
state.rs # AppState — shared Arc<Router> + master_key
gil.rs # GIL-activity tracker (records Python acquisitions)
auth/ # authentication as an axum extractor — added to handler args
mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY)
routes/ # one module per route, all matching the same template
AGENTS.md # ← the route template (read this before adding a route)
mod.rs # app(): merges every module's router()
health.rs # simple route (one file): router() + liveness/readiness
gil.rs # simple route (one file): router() + GET /health/gil
realtime/ # route with logic → axum surface + a no-axum service:
mod.rs # router() + handler + WS<->events adapter (the axum surface)
service.rs # business logic (select deployment, call provider) — no axum, testable
python/ # Python interop (feature: python-config) — load-time only
mod.rs, config.rs, AGENTS.md
```
## Rules
@ -53,5 +49,6 @@ proxy in a later phase. Health routes don't add the extractor (unauthenticated).
## Python interop
Anything that calls into Python lives in `python/` and is **load-time only** — see
`python/AGENTS.md`. The realtime data path never takes the GIL.
Python-backed loading lives in `litellm-config` and is **load-time only**. The
gateway's `python-config` feature forwards to that crate. The realtime data path
never takes the GIL.

2
litellm-rust/crates/ai-gateway/ARCHITECTURE.md
View file

@ -9,4 +9,6 @@ flowchart LR
C[client] <--> G[Rust ai-gateway<br/>LLM inference]
G <--> O[OpenAI realtime]
G -. spend tracking callback .-> P[litellm proxy]
F[litellm-config<br/>load-time only] --> G
F -. Python backend .-> P
```

10
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -16,9 +16,14 @@ required-features = ["server"]
[dependencies]
tracing.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-config.workspace = true
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
# rustls and its root store are direct dependencies so `io::tls` can build the
# one TLS config the outbound dials use; see that module for why it has to.
rustls.workspace = true
rustls-native-certs.workspace = true
# `sync` powers the bounded mpsc channel the realtime logger drains.
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
tokio-tungstenite.workspace = true
@ -31,14 +36,15 @@ subtle = { workspace = true, optional = true }
# sha2 hashes the master key into user_api_key_hash (matches the proxy's
# SHA-256 hash_token) so the plaintext credential never enters a log payload.
sha2 = { workspace = true, optional = true }
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
tower = { version = "0.5.3", features = ["util"], optional = true }
[features]
default = []
server = ["dep:axum", "dep:subtle", "dep:sha2"]
# Build the gateway's config from the proxy YAML via an embedded Python
# interpreter (links libpython; requires `litellm` importable at runtime).
python-config = ["dep:pyo3"]
python-config = ["litellm-config/python"]
trace-parity = ["server", "dep:tower", "litellm-core/observability"]
[dev-dependencies]
futures-channel = "0.3"

22
litellm-rust/crates/ai-gateway/README.md
View file

@ -6,25 +6,30 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
## Crates
`litellm-rust` has four crates. A crate is a layer or shared foundation, not a route:
`litellm-rust` has five crates. A crate is a layer or shared foundation, not a route:
| Crate | Role |
|-------|------|
| litellm-core | The LiteLLM SDK in Rust — per-route entrypoints (`messages::messages()`) that resolve the provider, transform, and make the call; plus types, provider transforms, and the router. |
| litellm-config | Config-loading boundary. Returns resolved deployments and optionally delegates loading to Python. |
| litellm-ai-gateway | The Axum server (behind the `server` feature) and WebSocket hosts. Translates HTTP/WS to core entrypoints; no provider handlers. |
| litellm-python-interop | Domain-neutral PyO3 foundation for GIL handling and typed Python/Serde conversion. |
| litellm-python-bridge | PyO3 cdylib exposing LiteLLM Rust APIs to the Python SDK. |
Dependency direction is acyclic: `litellm-python-bridge` depends on the domain layers and `litellm-python-interop`; the interop foundation depends on no LiteLLM domain crate.
Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on the domain layers and Python interop.
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
- **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil`
- **Health:** `GET /health/readiness`, `GET /health/liveness`
- **Request logs:** POSTed to a LiteLLM proxy at `/v1/rust_control_plane/logs` (see [Request logging](#request-logging))
> **Realtime serving is pure Rust.** Python is used at **load time only** — to
> read the config once at boot. The realtime hot path never touches Python.
The former `/health/gil` route and its acquisition counter were removed. They
only observed the single startup config load and did not prove that every GIL
acquisition was instrumented
## Configuration (config.yaml)
The gateway loads its `model_list` from a **config.yaml**, the same as the
@ -43,9 +48,10 @@ model_list:
LITELLM_CONFIG_PATH=./config.yaml ./litellm-ai-gateway
```
At boot the gateway calls into `litellm.proxy.read_model_list`, which reuses the
**real proxy config reader** (`ProxyConfig.get_config`). That means everything
the proxy supports in config.yaml works here too:
At boot `litellm-config` calls into `litellm.proxy.read_model_list` and returns
resolved deployments to the gateway, which constructs the router. The Python
backend still reuses the **real proxy config reader** (`ProxyConfig.get_config`),
so everything the proxy supports in config.yaml works here too:
- `include:` to merge in other config files,
- `os.environ/VAR` secret references (resolved via the secret manager, never
@ -82,8 +88,8 @@ stand-in built from the environment:
|---|---|---|
| `OPENAI_REALTIME_MODEL` | `gpt-realtime` | The single deployment's model name (also the `?model=` clients pass). |
This mode links no libpython and needs no config file, but it only supports one
hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
The default workspace build links no libpython and needs no config file. This
fallback mode only supports one hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
stand-in only for the leanest possible build.
## Request logging

6
litellm-rust/crates/ai-gateway/config.yaml
View file

@ -1,8 +1,8 @@
# Sample realtime config for the LiteLLM Rust AI Gateway.
#
# The gateway loads this model_list at boot via the embedded python config
# reader (litellm.proxy.read_model_list), which reuses the proxy's own reader —
# so include:, os.environ/ secrets, and DB-stored models all work here too.
# litellm-config resolves this model_list at boot through the Python config
# reader (litellm.proxy.read_model_list), then the gateway builds its router.
# Includes, environment secrets, and database-stored models still work.
#
# Secrets are referenced (never inlined) via os.environ/. A real deploy can
# override this file (e.g. mount a Render secret file at LITELLM_CONFIG_PATH).

58
litellm-rust/crates/ai-gateway/src/gil.rs
View file

@ -1,58 +0,0 @@
//! GIL-activity tracking.
//!
//! Every acquisition of the Python GIL is recorded here so the `/health/gil`
//! endpoint can report whether Python was touched recently. The design goal is
//! that the GIL is acquired **only at load time** (config read) and never on the
//! realtime hot path — polling this endpoint during traffic should show the
//! count holding steady and `acquired_last_30s` falling to `false`.
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
/// Window (seconds) for the "recently acquired" signal.
pub const RECENT_WINDOW_SECS: u64 = 30;
static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0);
/// Unix seconds of the last acquisition; `0` means "never".
static LAST_GIL_UNIX_SECS: AtomicU64 = AtomicU64::new(0);
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
/// Record that the GIL was just acquired. Call immediately before taking the GIL.
///
/// Only invoked under the `python-config` feature; without it the gateway never
/// touches Python, so the recorder is unused (and the endpoint reports zero).
#[cfg_attr(not(feature = "python-config"), allow(dead_code))]
pub fn record_acquisition() {
GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed);
LAST_GIL_UNIX_SECS.store(now_unix_secs(), Ordering::Relaxed);
}
/// Point-in-time view of GIL activity.
pub struct GilSnapshot {
pub total_acquisitions: u64,
pub seconds_since_last: Option<u64>,
pub acquired_last_30s: bool,
}
/// Read the current GIL-activity snapshot.
pub fn snapshot() -> GilSnapshot {
let total = GIL_ACQUISITIONS.load(Ordering::Relaxed);
let last = LAST_GIL_UNIX_SECS.load(Ordering::Relaxed);
let seconds_since_last = if last == 0 {
None
} else {
Some(now_unix_secs().saturating_sub(last))
};
let acquired_last_30s = seconds_since_last.is_some_and(|secs| secs <= RECENT_WINDOW_SECS);
GilSnapshot {
total_acquisitions: total,
seconds_since_last,
acquired_last_30s,
}
}

1
litellm-rust/crates/ai-gateway/src/io/mod.rs
View file

@ -3,3 +3,4 @@ pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;
pub(crate) mod tls;

33
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream(
.map_err(|err| Error::Auth(err.to_string()))?,
);
let (upstream, _response) = connect_async(request)
let (upstream, _response) = connect_upstream(request)
.await
.map_err(|err| Error::Network(err.to_string()))?;
Ok(upstream)
@ -284,6 +286,33 @@ mod tests {
serde_json::from_str(raw).expect("valid event json")
}
/// The realtime dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result = dial_upstream(
"gpt-realtime",
"sk-test",
Some(&format!("wss://127.0.0.1:{port}")),
)
.await;
assert!(matches!(result, Err(Error::Network(_))));
}
#[test]
fn resolve_api_key_prefers_param_then_blank_falls_through() {
assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");

35
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use crate::io::tls::connect_upstream;
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection {
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match error {
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -138,13 +140,13 @@ async fn dial_upstream(
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_async(request),
connect_upstream(request),
)
.await
.map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match error {
.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -324,6 +326,29 @@ mod tests {
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
/// The Responses dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result =
dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await;
assert!(matches!(result, Err(Error::Network(_))));
}
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("local address");

80
litellm-rust/crates/ai-gateway/src/io/tls.rs Normal file
View file

@ -0,0 +1,80 @@
//! Outbound WebSocket dials over a TLS config this crate builds once and owns.
//!
//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth`
//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that
//! `tokio-tungstenite` uses when handed no connector panics rather than guess
//! between them. Naming ring on a connector of our own settles that for these
//! dials without touching the process-wide default, and building the config
//! once keeps the platform trust store, which `tokio-tungstenite` would
//! otherwise re-read on every dial, off the dial path.
use std::io;
use std::sync::{Arc, OnceLock};
use rustls::{ClientConfig, RootCertStore};
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::error::TlsError;
use tokio_tungstenite::tungstenite::handshake::client::Response;
use tokio_tungstenite::{
Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
};
static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
fn build_config() -> Result<ClientConfig, Box<Error>> {
let native = rustls_native_certs::load_native_certs();
let roots = {
let mut store = RootCertStore::empty();
let (added, _ignored) = store.add_parsable_certificates(native.certs);
if added == 0 {
return Err(Box::new(Error::Io(io::Error::other(format!(
"no usable native root certificates: {:?}",
native.errors
)))));
}
store
};
ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
.with_safe_default_protocol_versions()
.map(|builder| builder.with_root_certificates(roots).with_no_client_auth())
.map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error))))
}
fn tls_config() -> Result<Arc<ClientConfig>, Box<Error>> {
if let Some(config) = TLS_CONFIG.get() {
return Ok(Arc::clone(config));
}
let built = Arc::new(build_config()?);
Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
}
pub(crate) async fn connect_upstream<R>(
request: R,
) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
where
R: IntoClientRequest + Unpin,
{
let request = request.into_client_request().map_err(Box::new)?;
let connector = match request.uri().scheme_str() {
Some("wss") => Some(Connector::Rustls(tls_config()?)),
_ => None,
};
connect_async_tls_with_config(request, None, false, connector)
.await
.map_err(Box::new)
}
#[cfg(test)]
mod tests {
use super::build_config;
#[test]
fn builds_a_usable_config_with_both_provider_features_enabled() {
let config = build_config().expect("a client config");
assert!(!config.crypto_provider().cipher_suites.is_empty());
}
}

12
litellm-rust/crates/ai-gateway/src/lib.rs
View file

@ -10,29 +10,23 @@
//! - [`io`]: compatibility exports and realtime WebSocket splice helpers.
//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
//! binary turns on. The `python-config` feature additionally pulls in [`python`]
//! for the load-time config reader.
//! binary turns on.
pub mod audio_transcription;
mod client;
pub mod io;
pub mod ocr;
/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
/// the `python-config` reader, so it is available without either feature.
pub mod gil;
#[cfg(feature = "server")]
pub mod auth;
#[cfg(feature = "server")]
pub mod routes;
#[cfg(feature = "server")]
pub mod state;
#[cfg(feature = "trace-parity")]
pub mod trace_parity;
mod constants;
pub mod integrations;
#[cfg(feature = "server")]
mod realtime;
#[cfg(feature = "python-config")]
pub mod python;

10
litellm-rust/crates/ai-gateway/src/main.rs
View file

@ -14,12 +14,12 @@ use std::sync::Arc;
use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key};
use litellm_ai_gateway::routes;
use litellm_ai_gateway::state::AppState;
#[cfg(feature = "python-config")]
use litellm_config::load_model_list;
use litellm_core::router::{Deployment, LiteLLMParams, Router};
use litellm_ai_gateway::integrations::custom_logger::CustomLogger;
use litellm_ai_gateway::integrations::litellm_python_proxy_api::LiteLLMPythonProxyAPILogger;
#[cfg(feature = "python-config")]
use litellm_ai_gateway::python;
/// Bind to localhost by default so the gateway is not a public, unauthenticated
/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`).
@ -124,10 +124,10 @@ fn resolve_port() -> u16 {
fn build_router() -> Router {
#[cfg(feature = "python-config")]
if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") {
match python::config::load_router_from_config(&config_path) {
Ok(router) => {
match load_model_list(std::path::Path::new(&config_path)) {
Ok(deployments) => {
eprintln!("loaded model_list from {config_path} via python config reader");
return router;
return Router::new(deployments);
}
Err(err) => {
eprintln!("config load failed ({err}); falling back to env deployment");

90
litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs
View file

@ -12,6 +12,7 @@ use litellm_core::providers::azure_ai::ocr::transformation::{
AZURE_AI_OCR_CONFIG, AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG,
};
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
use litellm_core::providers::reducto::ocr::transformation as reducto;
use litellm_core::providers::vertex_ai::ocr::transformation as vertex_ai;
use litellm_core::providers::vertex_ai::ocr::transformation::{
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
@ -39,6 +40,7 @@ pub(super) fn ocr_provider_config(
) -> Option<&'static dyn OcrProviderConfig> {
match provider {
"mistral" => Some(&MISTRAL_OCR_CONFIG),
"reducto" => reducto::config_for_model(model),
"azure_ai" if is_azure_document_intelligence_model(model) => {
Some(&AZURE_DOCUMENT_INTELLIGENCE_OCR_CONFIG)
}
@ -334,6 +336,7 @@ fn operation_status(response_json: &Value) -> Result<&str, Error> {
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn poll_document_intelligence(
operation_url: &str,
original_url: &str,
@ -392,9 +395,11 @@ pub(super) async fn poll_document_intelligence(
#[cfg(test)]
mod tests {
use super::*;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::json;
use super::*;
#[test]
fn blocks_private_and_metadata_ips() {
assert!(is_blocked_ip("127.0.0.1".parse().unwrap()));
@ -438,4 +443,87 @@ mod tests {
assert_eq!(transformed, document);
}
#[test]
fn truncate_error_body_passes_short_strings_through() {
let body = "Unauthorized";
assert_eq!(truncate_error_body(body), "Unauthorized");
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(306);
let truncated = truncate_error_body(&body);
assert!(truncated.ends_with("... (truncated)"));
let prefix_chars = truncated
.strip_suffix("... (truncated)")
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, 256);
}
#[test]
fn truncate_error_body_does_not_split_multibyte_chars() {
let body = "é".repeat(266);
let truncated = truncate_error_body(&body);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(
ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document()
);
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}
#[test]
fn string_headers_accepts_string_values() {
let headers = json!({
"x-trace-id": "trace-1"
})
.as_object()
.unwrap()
.clone();
assert_eq!(
string_headers(Some(headers)).expect("string headers accepted"),
vec![("x-trace-id".to_string(), "trace-1".to_string())]
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({
"x-retry-count": 3
})
.as_object()
.unwrap()
.clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::InvalidRequest(
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
)
);
}
}

12
litellm-rust/crates/ai-gateway/src/ocr/handler.rs
View file

@ -50,7 +50,11 @@ pub(crate) async fn execute_ocr_provider_call(
.await?;
return Ok(request
.config
.transform_ocr_response(&request.model, response_json)?
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json());
}
@ -71,6 +75,10 @@ pub(crate) async fn execute_ocr_provider_call(
Ok(request
.config
.transform_ocr_response(&request.model, response_json)?
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json())
}

102
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -1,11 +1,15 @@
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use litellm_core::providers::reducto::ocr::transformation::{
build_upload_request, extract_document_source, extract_upload_file_id,
};
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use super::common_utils::{convert_document_url_to_data_uri, string_headers};
use super::common_utils::{convert_document_url_to_data_uri, string_headers, truncate_error_body};
use super::types::{PreparedOcrRequest, ProviderOcrRequest};
use crate::client::http_client;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
@ -89,22 +93,39 @@ impl OcrLifecycleHooks {
)?;
let model = request.model.clone();
let custom_llm_provider = request.custom_llm_provider.clone();
let document = if config.requires_data_uri_document() {
let is_reducto = custom_llm_provider == "reducto";
let document = if is_reducto {
let guarded_document = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, request.document)
.await?;
upload_reducto_document(
&guarded_document,
request.api_base.as_deref(),
request.timeout,
&upstream_headers,
)
.await?
} else if config.requires_data_uri_document() {
convert_document_url_to_data_uri(request.document).await?
} else {
request.document
};
let optional_params = request.optional_params;
let body = config
.transform_ocr_request(&request.model, document, request.optional_params)?
.transform_ocr_request(&request.model, document, optional_params.clone())?
.data;
let body = self
.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?;
let body = if is_reducto {
body
} else {
self.run_during_call_guardrails(&model, &custom_llm_provider, &url, body)
.await?
};
Ok(ProviderOcrRequest {
model,
config,
url,
body,
optional_params,
upstream_headers,
timeout: request.timeout,
})
@ -165,6 +186,63 @@ impl OcrLifecycleHooks {
}
}
async fn upload_reducto_document(
document: &Value,
api_base: Option<&str>,
timeout: Option<std::time::Duration>,
upstream_headers: &[(String, String)],
) -> Result<Value, Error> {
let source = extract_document_source(document)?;
let Some(authorization) = upstream_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
else {
return Err(Error::Auth(
"Reducto upload requires an Authorization header".to_string(),
));
};
let Some(upload) = build_upload_request(source, authorization, api_base) else {
return Ok(document.clone());
};
let part = reqwest::multipart::Part::bytes(upload.bytes)
.file_name(upload.file_name)
.mime_str(&upload.mime_type)
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let form = reqwest::multipart::Form::new().part("file", part);
let mut request_builder = http_client().post(upload.url).multipart(form);
for (name, value) in upstream_headers {
if !name.eq_ignore_ascii_case("content-type")
&& !name.eq_ignore_ascii_case("content-length")
{
request_builder = request_builder.header(name, value);
}
}
if let Some(timeout) = timeout {
request_builder = request_builder.timeout(timeout);
}
let response = request_builder
.send()
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let body = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
if !status.is_success() {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
}
let response_json: Value = serde_json::from_str(&body).map_err(|error| {
Error::InvalidResponse(format!("invalid Reducto upload response JSON: {error}"))
})?;
let file_id = extract_upload_file_id(&response_json)?;
Ok(json!({"type": "document_url", "document_url": file_id}))
}
impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLifecycleHooks {
type PreCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
@ -187,6 +265,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
@ -210,6 +294,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
})
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,

149
litellm-rust/crates/ai-gateway/src/ocr/mod.rs
View file

@ -24,4 +24,151 @@ pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
}
#[cfg(test)]
mod tests;
mod tests {
use serde_json::{Map, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::{OcrRequest, ocr};
use crate::integrations::types::RequestMetadata;
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
#[tokio::test]
async fn reducto_file_upload_then_parse_maps_response() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut upload_socket, _) = listener.accept().await.expect("accepts upload request");
let upload_request = read_http_request(&mut upload_socket).await;
let upload_body = r#"{"file_id":"reducto://uploaded.pdf"}"#;
let upload_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
upload_body.len(),
upload_body
);
upload_socket
.write_all(upload_response.as_bytes())
.await
.expect("writes upload response");
let (mut parse_socket, _) = listener.accept().await.expect("accepts parse request");
let parse_request = read_http_request(&mut parse_socket).await;
let parse_body = r#"{"job_id":"job_123","usage":{"num_pages":3,"credits":3},"result":{"chunks":[{"content":"Page 1 block A","blocks":[{"content":"Page 1 block A","bbox":{"page":1},"kind":"text"}]},{"content":"Page 2 block A","blocks":[{"content":"Page 2 block A","bbox":{"page":2},"kind":"table"}]},{"content":"Page 1 block B","blocks":[{"content":"Page 1 block B","bbox":{"page":1},"kind":"text"}]},{"content":"Page 3 block A","blocks":[{"content":"Page 3 block A","bbox":{"page":3},"kind":"figure"}]}]}}"#;
let parse_response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
parse_body.len(),
parse_body
);
parse_socket
.write_all(parse_response.as_bytes())
.await
.expect("writes parse response");
(upload_request, parse_request)
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.api_key = None;
request.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer test-key")),
("x-trace-id".to_string(), json!("trace-1")),
]));
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.optional_params = Map::from_iter([
(
"formatting".to_string(),
json!({"table_output_format": "html"}),
),
("retrieval".to_string(), json!({"chunk_mode": "section"})),
("settings".to_string(), json!({"ocr_system": "standard"})),
]);
let response = ocr(request).await.expect("Reducto OCR succeeds");
assert_eq!(response["pages"].as_array().map(Vec::len), Some(3));
assert_eq!(
response["pages"][0]["markdown"],
"Page 1 block A\n\nPage 1 block B"
);
assert_eq!(response["pages"][1]["markdown"], "Page 2 block A");
assert_eq!(response["pages"][2]["markdown"], "Page 3 block A");
assert_eq!(response["usage_info"]["pages_processed"], 3);
assert_eq!(response["usage_info"]["credits"], 3);
assert_eq!(response["provider_native_response"]["job_id"], "job_123");
let (upload_request, parse_request) = server.await.expect("server task completes");
assert!(
upload_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(upload_request.contains("application/pdf"));
assert!(upload_request.contains("%PDF-1.4"));
assert!(upload_request.contains("x-trace-id: trace-1"));
assert!(
parse_request
.to_ascii_lowercase()
.contains("authorization: bearer test-key")
);
assert!(parse_request.contains(r#""input":"reducto://uploaded.pdf""#));
assert!(parse_request.contains(r#""table_output_format":"html""#));
assert!(parse_request.contains(r#""chunk_mode":"section""#));
assert!(parse_request.contains(r#""ocr_system":"standard""#));
}
}

96
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
View file

@ -2,6 +2,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use serde_json::{Map, Value};
use super::common_utils::ocr_provider_config;
use super::hooks::OcrLifecycleHooks;
@ -28,17 +29,33 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
let model = provider_info.model.to_string();
let custom_llm_provider = provider_info.custom_llm_provider.to_string();
let config = ocr_provider_config(&custom_llm_provider, &model)
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()));
.ok_or_else(|| litellm_core::Error::InvalidProvider(custom_llm_provider.clone()))
.and_then(|config| {
validate_request_format(config, &request.optional_params, &custom_llm_provider)?;
Ok(config)
});
let optional_params = match &config {
Ok(config) => {
let supported = config.supported_ocr_params();
config.map_ocr_params(
let mut mapped = config.map_ocr_params(
&request
.optional_params
.into_iter()
.iter()
.filter(|(name, _)| supported.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
)
);
for name in [
"vertex_project",
"vertex_ai_project",
"vertex_location",
"vertex_ai_location",
] {
if let Some(value) = request.optional_params.get(name) {
mapped.insert(name.to_string(), value.clone());
}
}
mapped
}
Err(_) => request.optional_params,
};
@ -64,6 +81,26 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
}
}
fn validate_request_format(
config: &'static dyn litellm_core::ocr::transformation::OcrProviderConfig,
optional_params: &Map<String, Value>,
provider: &str,
) -> Result<(), litellm_core::Error> {
let Some(format) = optional_params.get("req_format") else {
return Ok(());
};
match format.as_str() {
Some("litellm") => Ok(()),
Some("native") if config.supported_ocr_params().contains(&"req_format") => Ok(()),
Some("native") => Err(litellm_core::Error::InvalidRequest(format!(
"`req_format=native` is not supported for provider {provider}"
))),
_ => Err(litellm_core::Error::InvalidRequest(format!(
"Invalid `req_format`: {format}. Expected `litellm` or `native`"
))),
}
}
fn new_ocr_call_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
@ -73,3 +110,54 @@ fn new_ocr_call_id() -> String {
.unwrap_or(0);
format!("ocr-{timestamp}-{sequence}")
}
#[cfg(test)]
mod tests {
use litellm_core::error::Error;
use serde_json::{Map, json};
use super::{OcrRequest, prepare_ocr_call};
use crate::integrations::types::RequestMetadata;
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
fn request_with_format(format: &str) -> OcrRequest<'_> {
let mut request = base_ocr_request("mistral/mistral-ocr-latest");
request.optional_params = Map::from_iter([("req_format".to_string(), json!(format))]);
request
}
#[test]
fn native_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("native"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("not supported for provider"))
);
}
#[test]
fn unknown_format_rejected_for_provider_without_support_as_bad_request() {
let prepared = prepare_ocr_call(request_with_format("raw"));
assert!(
matches!(prepared.request.config, Err(Error::InvalidRequest(message)) if message.contains("Invalid `req_format`"))
);
}
}

1
litellm-rust/crates/ai-gateway/src/ocr/types.rs
View file

@ -53,6 +53,7 @@ pub(crate) struct ProviderOcrRequest {
pub(crate) config: &'static dyn OcrProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) optional_params: Map<String, Value>,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
View file

@ -1,27 +0,0 @@
# ai-gateway/src/python — Python interop (load-time only)
Functions here embed the Python interpreter (pyo3) and take the GIL to call into
`litellm` (e.g. read the proxy `model_list`). Compiled only under the
`python-config` feature.
## Hard rule: non-hot-path functions only
Everything in this folder MUST run **at most once per process lifetime — at
startup / load time** (config read, warm-up). NEVER call into Python on the
request path:
- No GIL acquisition per request, per connection, or per realtime event.
- No Python call inside a route handler, the router's hot path, or any loop that
scales with traffic.
**Why:** the GIL serializes execution and would cap throughput; the realtime data
path must stay pure Rust. Every acquisition is recorded by `crate::gil` — poll
`GET /health/gil`, and `total_acquisitions` MUST stay flat under load.
## How to add one
Resolve whatever Python-derived data you need **once at boot** and hand the rest
of the gateway an owned, plain-Rust value (e.g. build a `Router` from the
resolved `model_list`). Record the acquisition via `crate::gil::record_acquisition()`
immediately before taking the GIL. If a function would need to run per request,
it does not belong here — move the work to Rust, or pre-resolve it at startup.

37
litellm-rust/crates/ai-gateway/src/python/config.rs
View file

@ -1,37 +0,0 @@
//! Build the router by calling the Python proxy config reader (load time only).
//!
//! Embeds the interpreter via pyo3 and calls
//! `litellm.proxy.read_model_list.read_model_list`, which reuses the proxy's
//! `os.environ/` + secret-manager resolution. The GIL is taken **once at boot**
//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python.
//!
//! Compiled only under the `python-config` feature.
use litellm_core::error::Error;
use litellm_core::router::{Deployment, Router};
use pyo3::prelude::*;
use crate::gil;
/// Load the router's `model_list` from `config_path` via the Python reader.
pub fn load_router_from_config(config_path: &str) -> Result<Router, Error> {
gil::record_acquisition();
Python::attach(|py| {
let model_list = py
.import("litellm.proxy.read_model_list")
.and_then(|module| module.getattr("read_model_list"))
.and_then(|reader| reader.call1((config_path,)))
.map_err(|err| Error::Routing(format!("read_model_list failed: {err}")))?;
let model_list_json: String = py
.import("json")
.and_then(|json| json.getattr("dumps"))
.and_then(|dumps| dumps.call1((model_list,)))
.and_then(|encoded| encoded.extract())
.map_err(|err| Error::Routing(format!("serializing model_list failed: {err}")))?;
let deployments: Vec<Deployment> = serde_json::from_str(&model_list_json)
.map_err(|err| Error::Routing(format!("parsing model_list failed: {err}")))?;
Ok(Router::new(deployments))
})
}

4
litellm-rust/crates/ai-gateway/src/python/mod.rs
View file

@ -1,4 +0,0 @@
//! Python interop for the gateway. See `AGENTS.md`: **load-time / non-hot-path
//! only.** Compiled only under the `python-config` feature.
pub mod config;

2
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
View file

@ -13,7 +13,7 @@ private). This is the norm — don't split until it hurts.
pub fn router() -> Router<AppState> { Router::new().route(PATH, get(handle)) }
async fn handle(...) -> impl IntoResponse { ... }
```
`health.rs` and `gil.rs` are examples.
`health.rs` is the example.
## Split out `service` when there's real logic
When a route has business logic worth testing without axum, put it in a sibling

30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
View file

@ -1,30 +0,0 @@
//! `GET /health/gil` — poll to confirm Python is only touched at load time.
//! Simple-route template: a `router()` plus its handler, in one file.
use axum::routing::get;
use axum::{Json, Router};
use serde::Serialize;
use crate::gil;
use crate::state::AppState;
/// This route's contribution to the app router.
pub fn router() -> Router<AppState> {
Router::new().route("/health/gil", get(status))
}
#[derive(Debug, Serialize)]
struct GilStatusResponse {
gil_acquired_last_30s: bool,
total_acquisitions: u64,
seconds_since_last: Option<u64>,
}
async fn status() -> Json<GilStatusResponse> {
let snapshot = gil::snapshot();
Json(GilStatusResponse {
gil_acquired_last_30s: snapshot.acquired_last_30s,
total_acquisitions: snapshot.total_acquisitions,
seconds_since_last: snapshot.seconds_since_last,
})
}

6
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -21,6 +21,12 @@ pub fn router() -> Router<AppState> {
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
}
#[tracing::instrument(
name = "messages_gateway_route",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
async fn handle(
_auth: RequireMasterKey,
State(state): State<AppState>,

6
litellm-rust/crates/ai-gateway/src/routes/messages/service.rs
View file

@ -12,6 +12,12 @@ pub(crate) enum MessagesResponse {
Stream(reqwest::Response),
}
#[tracing::instrument(
name = "messages_gateway_service",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
pub async fn run(
router: &Arc<Router>,
body: Value,

4
litellm-rust/crates/ai-gateway/src/routes/mod.rs
View file

@ -2,10 +2,9 @@
//!
//! **Template:** every route module exposes `pub fn router() -> Router<AppState>`
//! that mounts its own paths; [`app`] merges them. A trivial route is a single
//! file (`health.rs`, `gil.rs`); a non-trivial one is a folder (`realtime/`) with
//! file (`health.rs`); a non-trivial one is a folder (`realtime/`) with
//! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md.
pub mod gil;
pub mod health;
pub mod messages;
pub mod realtime;
@ -19,7 +18,6 @@ use crate::state::AppState;
pub fn app(state: AppState) -> Router {
Router::new()
.merge(health::router())
.merge(gil::router())
.merge(messages::router())
.merge(realtime::router())
.merge(responses::router())

65
litellm-rust/crates/ai-gateway/src/trace_parity.rs Normal file
View file

@ -0,0 +1,65 @@
//! Harness-only in-process adapters. Never mounted as production routes.
use std::sync::Arc;
use axum::body::{Body, to_bytes};
use axum::http::header::{AUTHORIZATION, CONTENT_TYPE};
use axum::http::{Request, StatusCode};
use litellm_core::Error;
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde::Serialize;
use serde_json::Value;
use tower::ServiceExt;
use crate::io::realtime_pool::RealtimePool;
use crate::routes;
use crate::state::AppState;
#[derive(Debug, Serialize)]
pub struct GatewayResponse {
pub status: u16,
pub body: Value,
}
pub async fn messages_request(
model_alias: String,
provider_model: String,
api_base: String,
body: Value,
) -> Result<GatewayResponse, Error> {
let state = AppState {
router: Arc::new(ModelRouter::new(vec![Deployment {
model_name: model_alias,
litellm_params: LiteLLMParams {
model: provider_model,
api_key: Some("trace-provider-key".to_string()),
api_base: Some(api_base),
},
}])),
master_key: Some(Arc::from("trace-master-key")),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
};
let request = Request::builder()
.method("POST")
.uri("/v1/messages")
.header(AUTHORIZATION, "Bearer trace-master-key")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(body.to_string()))
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let response = routes::app(state)
.oneshot(request)
.await
.map_err(|error| match error {})?;
let status: StatusCode = response.status();
let bytes = to_bytes(response.into_body(), usize::MAX)
.await
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
let body = serde_json::from_slice(&bytes).map_err(|error| {
Error::InvalidResponse(format!("gateway returned invalid JSON: {error}"))
})?;
Ok(GatewayResponse {
status: status.as_u16(),
body,
})
}

48
litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs Normal file
View file

@ -0,0 +1,48 @@
//! Guards the wiring, not just the helper: a `wss://` dial through the public
//! API has to resolve its own crypto provider, in a test binary where nothing
//! has installed a process-wide one, and has to leave it uninstalled.
use std::collections::HashMap;
use std::time::Duration;
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
use tokio::net::TcpListener;
async fn dead_tls_server() -> u16 {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
port
}
#[tokio::test]
async fn dialing_wss_returns_an_error_instead_of_panicking() {
let port = dead_tls_server().await;
let result = ResponsesWebSocketConnection::connect_url(
&format!("wss://127.0.0.1:{port}/"),
&HashMap::new(),
Some(Duration::from_secs(10)),
)
.await;
assert!(
result.is_err(),
"a plain TCP server cannot finish a TLS handshake"
);
assert!(
rustls::crypto::CryptoProvider::get_default().is_none(),
"the dial settles its provider on its own connector, not process-wide"
);
}

248
litellm-rust/crates/ai-gateway/src/ocr/tests.rs → litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -1,23 +1,23 @@
use std::sync::{Arc, Mutex};
use std::time::Duration;
use litellm_core::error::Error;
use litellm_core::http_utils::has_header;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::common_utils::{ocr_provider_config, string_headers, truncate_error_body};
use super::{OcrRequest, ocr};
use crate::integrations::custom_guardrail::{
use litellm_ai_gateway::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
};
use crate::integrations::custom_logger::{
use litellm_ai_gateway::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails,
};
use crate::integrations::types::RequestMetadata;
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
#[cfg(feature = "trace-parity")]
use litellm_core::observability::FunctionTrace;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "trace-parity")]
use tracing::instrument::WithSubscriber;
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -136,6 +136,7 @@ struct RecordingOcrGuardrail {
hooks: Vec<GuardrailEventHook>,
events: Mutex<Vec<&'static str>>,
block_pre_call: bool,
block_during_call: bool,
}
impl RecordingOcrGuardrail {
@ -144,6 +145,7 @@ impl RecordingOcrGuardrail {
hooks,
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: false,
}
}
@ -152,6 +154,16 @@ impl RecordingOcrGuardrail {
hooks: vec![GuardrailEventHook::PreCall],
events: Mutex::new(Vec::new()),
block_pre_call: true,
block_during_call: false,
}
}
fn blocking_during_call() -> Self {
Self {
hooks: vec![GuardrailEventHook::DuringCall],
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: true,
}
}
@ -193,91 +205,95 @@ impl CustomGuardrail for RecordingOcrGuardrail {
) -> GuardrailFuture<'a> {
Box::pin(async move {
self.events.lock().unwrap().push("async_moderation_hook");
if self.block_during_call {
return Ok(GuardrailDecision::Block(GuardrailError::blocked(
"blocked before provider",
)));
}
request.data["body"]["guarded_during"] = json!(true);
Ok(GuardrailDecision::Mask(request))
})
}
}
#[test]
fn truncate_error_body_passes_short_strings_through() {
let body = "Unauthorized";
assert_eq!(truncate_error_body(body), "Unauthorized");
fn base_ocr_request(model: &str) -> OcrRequest<'_> {
OcrRequest {
model,
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Map::new(),
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: None,
}
}
#[test]
fn truncate_error_body_caps_long_payloads() {
let body = "x".repeat(306);
let truncated = truncate_error_body(&body);
#[tokio::test]
async fn reducto_during_call_guardrail_blocks_before_upload() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let api_base = format!("http://{address}");
let guardrail = Arc::new(RecordingOcrGuardrail::blocking_during_call());
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
request.guardrails = vec![guardrail.clone()];
assert!(truncated.ends_with("... (truncated)"));
let prefix_chars = truncated
.strip_suffix("... (truncated)")
.expect("truncated marker present")
.chars()
.count();
assert_eq!(prefix_chars, 256);
let error = ocr(request).await.expect_err("guardrail blocks upload");
assert!(matches!(error, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_moderation_hook"]);
let accepted = tokio::time::timeout(Duration::from_millis(100), listener.accept()).await;
assert!(accepted.is_err(), "upload socket should not be touched");
}
#[test]
fn truncate_error_body_does_not_split_multibyte_chars() {
let body = "é".repeat(266);
let truncated = truncate_error_body(&body);
assert!(truncated.is_char_boundary(truncated.len()));
}
#[tokio::test]
async fn reducto_upload_error_body_is_truncated() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("test listener binds");
let address = listener.local_addr().expect("listener has local address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts upload request");
let _request = read_http_request(&mut socket).await;
let body = "x".repeat(300);
let response = format!(
"HTTP/1.1 500 Internal Server Error\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes upload response");
});
let api_base = format!("http://{address}");
let mut request = base_ocr_request("reducto/parse-v3");
request.api_base = Some(&api_base);
request.document = json!({
"type": "document_url",
"document_url": "data:application/pdf;base64,JVBERi0xLjQ="
});
let error = ocr(request).await.expect_err("upload should fail");
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(
ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document()
matches!(error, Error::Http { status: 500, body } if body.chars().count() < 300 && body.ends_with("... (truncated)"))
);
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}
#[test]
fn string_headers_accepts_string_values() {
let headers = json!({
"x-trace-id": "trace-1"
})
.as_object()
.unwrap()
.clone();
assert_eq!(
string_headers(Some(headers)).expect("string headers accepted"),
vec![("x-trace-id".to_string(), "trace-1".to_string())]
);
}
#[test]
fn auth_header_detection_is_case_insensitive() {
let headers = vec![
("x-trace-id".to_string(), "trace-1".to_string()),
("authorization".to_string(), "Bearer sk-test".to_string()),
];
assert!(has_header(&headers, "authorization"));
let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
assert!(has_header(&headers, "authorization"));
let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
assert!(!has_header(&headers, "authorization"));
server.await.expect("server task completes");
}
#[tokio::test]
@ -308,14 +324,17 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
let response = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -327,9 +346,10 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
..Default::default()
},
litellm_call_id: Some("ocr-call-1"),
})
.await
.expect("ocr request succeeds");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let response = call.await.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
assert_eq!(
@ -347,6 +367,16 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
error_kind: None,
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["success_callback"]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
@ -376,14 +406,17 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
});
let logger = Arc::new(RecordingOcrLogger::default());
let err = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -392,9 +425,10 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-2"),
})
.await
.expect_err("provider error propagates");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let err = call.await.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
@ -409,6 +443,16 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
error_kind: Some("HttpError".to_string()),
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["failure_callback"]
);
}
#[tokio::test]
@ -595,21 +639,3 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() {
"{poll_request}"
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = json!({
"x-retry-count": 3
})
.as_object()
.unwrap()
.clone();
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::InvalidRequest(
"OCR extra_headers.x-retry-count must be a string, got number".to_string()
)
);
}

16
litellm-rust/crates/config/Cargo.toml Normal file
View file

@ -0,0 +1,16 @@
[package]
name = "litellm-config"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-core.workspace = true
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
serde_json.workspace = true
thiserror.workspace = true
[features]
default = []
python = ["dep:pyo3"]

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

@ -0,0 +1,11 @@
use thiserror::Error as ThisError;
#[derive(Debug, ThisError)]
pub enum Error {
#[error("read_model_list failed: {0}")]
PythonLoading(String),
#[error("serializing model_list failed: {0}")]
Serialization(String),
#[error("parsing model_list failed: {0}")]
ModelListParsing(#[source] serde_json::Error),
}

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

@ -0,0 +1,7 @@
mod error;
#[cfg(feature = "python")]
mod python;
pub use error::Error;
#[cfg(feature = "python")]
pub use python::load_model_list;

76
litellm-rust/crates/config/src/python.rs Normal file
View file

@ -0,0 +1,76 @@
use std::path::Path;
use litellm_core::router::Deployment;
use pyo3::prelude::*;
use crate::Error;
pub fn load_model_list(config_path: &Path) -> Result<Vec<Deployment>, Error> {
Python::attach(|python| {
let model_list = python
.import("litellm.proxy.read_model_list")
.and_then(|module| module.getattr("read_model_list"))
.and_then(|reader| reader.call1((config_path.to_string_lossy().as_ref(),)))
.map_err(|error| Error::PythonLoading(error.to_string()))?;
let model_list_json = python
.import("json")
.and_then(|json| json.getattr("dumps"))
.and_then(|dumps| dumps.call1((model_list,)))
.and_then(|encoded| encoded.extract::<String>())
.map_err(|error| Error::Serialization(error.to_string()))?;
parse_model_list(&model_list_json)
})
}
fn parse_model_list(model_list_json: &str) -> Result<Vec<Deployment>, Error> {
serde_json::from_str(model_list_json).map_err(Error::ModelListParsing)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_resolved_model_list() {
let deployments = parse_model_list(
r#"[
{
"model_name": "realtime",
"litellm_params": {
"model": "openai/gpt-realtime",
"api_key": "resolved-secret",
"api_base": "https://api.example.test/v1"
}
},
{
"model_name": "without-optional-values",
"litellm_params": {"model": "openai/gpt-4.1"}
}
]"#,
)
.expect("resolved model list should parse");
assert_eq!(deployments.len(), 2);
assert_eq!(deployments[0].model_name, "realtime");
assert_eq!(
deployments[0].litellm_params.api_key.as_deref(),
Some("resolved-secret")
);
assert_eq!(
deployments[0].litellm_params.api_base.as_deref(),
Some("https://api.example.test/v1")
);
assert_eq!(deployments[1].litellm_params.api_key, None);
assert_eq!(deployments[1].litellm_params.api_base, None);
}
#[test]
fn malformed_model_list_returns_parsing_error() {
let error = parse_model_list(r#"[{"model_name":"missing-params"}]"#)
.expect_err("missing litellm_params should fail");
assert!(matches!(error, Error::ModelListParsing(_)));
}
}

5
litellm-rust/crates/core/Cargo.toml
View file

@ -6,12 +6,14 @@ license.workspace = true
repository.workspace = true
[dependencies]
base64.workspace = true
rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tracing.workspace = true
tracing-subscriber = { workspace = true, optional = true }
sha2.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
@ -30,6 +32,9 @@ bedrock-auth = [
"dep:aws-types",
"dep:aws-smithy-runtime-api",
]
observability = ["dep:tracing-subscriber"]
[dev-dependencies]
rstest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
tracing-subscriber.workspace = true

2
litellm-rust/crates/core/src/constants.rs
View file

@ -41,3 +41,5 @@ pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";
pub const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

17
litellm-rust/crates/core/src/http_utils.rs
View file

@ -101,6 +101,23 @@ mod tests {
assert!(!has_header(&headers, "authorization"));
}
#[test]
fn auth_header_detection_is_case_insensitive() {
let headers = vec![
("x-trace-id".to_string(), "trace-1".to_string()),
("authorization".to_string(), "Bearer sk-test".to_string()),
];
assert!(has_header(&headers, "authorization"));
let headers = vec![("Authorization".to_string(), "Bearer sk-test".to_string())];
assert!(has_header(&headers, "authorization"));
let headers = vec![("x-trace-id".to_string(), "trace-1".to_string())];
assert!(!has_header(&headers, "authorization"));
}
#[test]
fn bearer_detection_requires_a_non_empty_token() {
assert!(has_bearer_auth(&[(

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

@ -6,6 +6,8 @@ pub mod constants;
pub mod error;
pub mod http_utils;
pub mod messages;
#[cfg(any(feature = "observability", test))]
pub mod observability;
pub mod ocr;
pub mod providers;
pub mod realtime;

215
litellm-rust/crates/core/src/observability/function_trace.rs Normal file
View file

@ -0,0 +1,215 @@
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use serde::Serialize;
use tracing::span::{Attributes, Id};
use tracing::{Dispatch, Subscriber};
use tracing_subscriber::layer::Context;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::{Layer, Registry};
use super::function_trace_filter;
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionTraceEvent {
pub id: usize,
pub parent_id: Option<usize>,
pub function: &'static str,
pub module_path: Option<&'static str>,
pub file: Option<&'static str>,
pub line: Option<u32>,
}
#[derive(Clone, Default)]
pub struct FunctionTrace {
events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
span_events: Arc<Mutex<HashMap<Id, usize>>>,
}
impl FunctionTrace {
pub fn dispatcher(&self) -> Dispatch {
Dispatch::new(
Registry::default().with(
FunctionTraceLayer {
trace: self.clone(),
}
.with_filter(function_trace_filter()),
),
)
}
pub fn events(&self) -> Vec<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let parent_id = context.span(id).and_then(|span| {
let span_events = self
.trace
.span_events
.lock()
.unwrap_or_else(|error| error.into_inner());
span.scope()
.skip(1)
.find_map(|ancestor| span_events.get(&ancestor.id()).copied())
});
let mut events = self
.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner());
let event_id = events.len();
events.push(FunctionTraceEvent {
id: event_id,
parent_id,
function: attributes.metadata().name(),
module_path: attributes.metadata().module_path(),
file: attributes.metadata().file(),
line: attributes.metadata().line(),
});
self.trace
.span_events
.lock()
.unwrap_or_else(|error| error.into_inner())
.insert(id.clone(), event_id);
}
}
#[cfg(test)]
mod tests {
use crate::constants::FUNCTION_TRACE_TARGET;
use super::*;
fn event(
id: usize,
parent_id: Option<usize>,
function: &'static str,
) -> (usize, Option<usize>, &'static str) {
(id, parent_id, function)
}
fn structural_events(
events: &[FunctionTraceEvent],
) -> Vec<(usize, Option<usize>, &'static str)> {
events
.iter()
.map(|event| (event.id, event.parent_id, event.function))
.collect()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn outer() {
tokio::task::yield_now().await;
inner().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn inner() {
tokio::task::yield_now().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn concurrent_parent() {
tokio::join!(inner(), inner());
}
#[tokio::test]
async fn concurrent_futures_keep_separate_traces_across_yields() {
use tracing::instrument::WithSubscriber;
let first = FunctionTrace::default();
let second = FunctionTrace::default();
let outside = FunctionTrace::default();
async {
tokio::join!(
outer().with_subscriber(first.dispatcher()),
inner().with_subscriber(second.dispatcher()),
);
inner().await;
}
.with_subscriber(outside.dispatcher())
.await;
assert_eq!(
structural_events(&first.events()),
vec![event(0, None, "outer"), event(1, Some(0), "inner")],
);
assert_eq!(
structural_events(&second.events()),
vec![event(0, None, "inner")],
);
assert_eq!(
structural_events(&outside.events()),
vec![event(0, None, "inner")],
);
}
#[tokio::test]
async fn concurrent_siblings_keep_the_same_parent() {
use tracing::instrument::WithSubscriber;
let trace = FunctionTrace::default();
concurrent_parent()
.with_subscriber(trace.dispatcher())
.await;
assert_eq!(
structural_events(&trace.events()),
vec![
event(0, None, "concurrent_parent"),
event(1, Some(0), "inner"),
event(2, Some(0), "inner"),
]
);
}
#[test]
fn records_matching_spans_in_creation_order() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let _ignored = tracing::trace_span!(target: "other", "ignored");
let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
});
assert_eq!(
structural_events(&trace.events()),
vec![event(0, None, "same_name"), event(1, None, "same_name")]
);
}
#[test]
fn records_matching_span_nesting_depth() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
let _outer_guard = outer.enter();
let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
});
assert_eq!(
structural_events(&trace.events()),
vec![event(0, None, "outer"), event(1, Some(0), "inner")]
);
}
}

59
litellm-rust/crates/core/src/observability/mod.rs Normal file
View file

@ -0,0 +1,59 @@
use tracing::span::Id;
use tracing::{Level, Metadata, Subscriber};
use tracing_subscriber::filter::{FilterFn, LevelFilter, filter_fn};
use tracing_subscriber::layer::Context;
use tracing_subscriber::registry::LookupSpan;
use crate::constants::FUNCTION_TRACE_TARGET;
pub mod function_trace;
pub use function_trace::{FunctionTrace, FunctionTraceEvent};
pub fn function_trace_filter() -> FilterFn<impl Fn(&Metadata<'_>) -> bool> {
filter_fn(|metadata| {
metadata.is_span()
&& metadata.target() == FUNCTION_TRACE_TARGET
&& *metadata.level() == Level::TRACE
})
.with_max_level_hint(LevelFilter::TRACE)
}
pub fn span_depth<S>(context: &Context<'_, S>, id: &Id) -> usize
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use tracing::instrument::WithSubscriber;
use super::*;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn instrumented_with_literal_target() {}
#[tokio::test]
async fn literal_instrument_target_matches_filter_constant() {
assert_eq!(FUNCTION_TRACE_TARGET, "litellm::function_trace");
let trace = FunctionTrace::default();
instrumented_with_literal_target()
.with_subscriber(trace.dispatcher())
.await;
let events = trace.events();
assert_eq!(events.len(), 1);
assert_eq!(events[0].id, 0);
assert_eq!(events[0].parent_id, None);
assert_eq!(events[0].function, "instrumented_with_literal_target");
assert_eq!(events[0].module_path, Some(module_path!()));
assert_eq!(events[0].file, Some(file!()));
assert!(events[0].line.is_some());
}
}

9
litellm-rust/crates/core/src/ocr/transformation.rs
View file

@ -51,6 +51,15 @@ pub trait OcrProviderConfig: Sync {
response_json: Value,
) -> Result<OcrResponseData, Error>;
fn transform_ocr_response_with_params(
&self,
model: &str,
response_json: Value,
_optional_params: &Map<String, Value>,
) -> Result<OcrResponseData, Error> {
self.transform_ocr_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,

15
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -1,5 +1,5 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{Map, Value};
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct OcrRequestData {
@ -14,16 +14,25 @@ pub struct OcrResponseData {
pub document_annotation: Option<Value>,
pub usage_info: Option<Value>,
pub object: String,
pub extra_fields: Map<String, Value>,
pub provider_native_response: Option<Value>,
}
impl OcrResponseData {
pub fn into_json(self) -> Value {
serde_json::json!({
let mut response = serde_json::json!({
"pages": self.pages,
"model": self.model,
"document_annotation": self.document_annotation,
"usage_info": self.usage_info,
"object": self.object,
})
});
if let Value::Object(object) = &mut response {
object.extend(self.extra_fields);
if let Some(native_response) = self.provider_native_response {
object.insert("provider_native_response".to_string(), native_response);
}
}
response
}
}

887
litellm-rust/crates/core/src/providers/azure_ai/ocr/transformation.rs
View file

File diff suppressed because it is too large Load diff

2
litellm-rust/crates/core/src/providers/mistral/ocr/transformation.rs
View file

@ -134,6 +134,8 @@ impl OcrProviderConfig for MistralOcrConfig {
document_annotation,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}

1
litellm-rust/crates/core/src/providers/mod.rs
View file

@ -4,4 +4,5 @@ pub mod azure_ai;
pub mod bedrock;
pub mod mistral;
pub mod openai;
pub mod reducto;
pub mod vertex_ai;

1
litellm-rust/crates/core/src/providers/reducto/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod ocr;

4
litellm-rust/crates/core/src/providers/reducto/ocr/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub mod transformation;
#[cfg(test)]
mod tests;

202
litellm-rust/crates/core/src/providers/reducto/ocr/tests.rs Normal file
View file

@ -0,0 +1,202 @@
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::transformation::*;
use crate::ocr::transformation::OcrProviderConfig;
#[fixture]
fn parse_response() -> Value {
json!({
"job_id": "job_123",
"usage": {"num_pages": 3, "credits": 3},
"result": {
"chunks": [
{
"content": "Page 1 block A",
"blocks": [{
"content": "Page 1 block A",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 2 block A",
"blocks": [{
"content": "Page 2 block A",
"bbox": {"page": 2},
"kind": "table",
}],
},
{
"content": "Page 1 block B",
"blocks": [{
"content": "Page 1 block B",
"bbox": {"page": 1},
"kind": "text",
}],
},
{
"content": "Page 3 block A",
"blocks": [{
"content": "Page 3 block A",
"bbox": {"page": 3},
"kind": "figure",
}],
},
],
},
})
}
#[rstest]
fn test_parse_v3_file_upload_and_response_mapping(parse_response: Value) {
let source = classify_document_source("data:application/pdf;base64,JVBERi0xLjQ=")
.expect("PDF data URI should be valid");
let upload = build_upload_request(
source,
"Bearer test-key",
Some("https://platform.reducto.ai"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://platform.reducto.ai/upload");
assert_eq!(upload.authorization, "Bearer test-key");
assert_eq!(upload.file_name, "document");
assert_eq!(upload.mime_type, "application/pdf");
assert_eq!(upload.bytes, b"%PDF-1.4");
let optional_params = json!({
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
.as_object()
.expect("params should be an object")
.clone();
let request = build_parse_v3_request("reducto://uploaded.pdf", optional_params);
assert_eq!(
request.data,
json!({
"input": "reducto://uploaded.pdf",
"formatting": {"table_output_format": "html"},
"retrieval": {"chunk_mode": "section"},
"settings": {"ocr_system": "standard"},
})
);
let transformed = transform_reducto_response("parse-v3", parse_response.clone())
.expect("response should transform");
assert_eq!(
transformed.usage_info,
Some(json!({"pages_processed": 3, "credits": 3}))
);
assert_eq!(transformed.pages.len(), 3);
assert_eq!(
transformed.pages[0],
json!({
"index": 0,
"markdown": "Page 1 block A\n\nPage 1 block B",
"blocks": [
{"content": "Page 1 block A", "bbox": {"page": 1}, "kind": "text"},
{"content": "Page 1 block B", "bbox": {"page": 1}, "kind": "text"},
],
})
);
assert_eq!(transformed.pages[1]["markdown"], "Page 2 block A");
assert_eq!(transformed.pages[2]["markdown"], "Page 3 block A");
assert_eq!(transformed.provider_native_response, Some(parse_response));
}
#[rstest]
fn test_parse_v3_reducto_id_passthrough_skips_upload(parse_response: Value) {
let document = json!({
"type": "document_url",
"document_url": "reducto://already-uploaded.pdf",
});
let source = extract_document_source(&document).expect("Reducto ID should be valid");
assert!(build_upload_request(source.clone(), "Bearer test-key", None).is_none());
assert_eq!(
source,
ReductoDocumentSource::FileId("reducto://already-uploaded.pdf".to_string())
);
let request = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_request(
"parse-v3",
document,
json!({"retrieval": {"chunk_mode": "section"}})
.as_object()
.expect("params should be object")
.clone(),
)
.expect("direct ID should transform");
assert_eq!(request.data["input"], "reducto://already-uploaded.pdf");
assert_eq!(request.data["retrieval"]["chunk_mode"], "section");
let response = REDUCTO_PARSE_V3_CONFIG
.transform_ocr_response("parse-v3", parse_response)
.expect("response should transform");
assert!(
response.pages[0]["markdown"]
.as_str()
.expect("markdown should be string")
.starts_with("Page 1 block A")
);
}
#[rstest]
fn test_parse_legacy_wraps_enhance_under_options() {
let request = build_parse_legacy_request(
"reducto://legacy.pdf",
json!({"enhance": {"agentic": [{"type": "table"}]}})
.as_object()
.expect("params should be object"),
);
assert_eq!(
request.data,
json!({
"document_url": "reducto://legacy.pdf",
"options": {"enhance": {"agentic": [{"type": "table"}]}},
})
);
}
#[rstest]
fn test_parse_v3_image_data_uri_upload_uses_image_mime() {
let source = classify_document_source("data:image/png;base64,iVBORw0KGgo=")
.expect("PNG data URI should be valid");
let upload = build_upload_request(
source,
"Bearer programmatic-key",
Some("https://custom.reducto.test/"),
)
.expect("data URI should require upload");
assert_eq!(upload.url, "https://custom.reducto.test/upload");
assert_eq!(upload.authorization, "Bearer programmatic-key");
assert_eq!(upload.mime_type, "image/png");
assert_eq!(upload.bytes, b"\x89PNG\r\n\x1a\n");
}
#[rstest]
#[case::http("http://example.com/document.pdf")]
#[case::https("https://example.com/document.pdf")]
fn test_parse_v3_rejects_plain_http_urls(#[case] source: &str) {
let error = classify_document_source(source).expect_err("plain URL should be rejected");
assert!(error.to_string().contains("upload the file first"));
}
#[rstest]
fn test_parse_v3_uses_programmatic_api_key_over_env() {
let key = resolve_api_key(Some("passed-key"), &|_| Some("env-reducto-key".to_string()))
.expect("explicit key should resolve");
assert_eq!(key, "passed-key");
let headers = REDUCTO_PARSE_V3_CONFIG
.validate_environment(Vec::new(), Some("passed-key"), &|_| {
Some("env-reducto-key".to_string())
})
.expect("headers should validate");
assert_eq!(
headers,
vec![("Authorization".to_string(), "Bearer passed-key".to_string())]
);
}

407
litellm-rust/crates/core/src/providers/reducto/ocr/transformation.rs Normal file
View file

@ -0,0 +1,407 @@
use std::collections::BTreeMap;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use serde_json::{Map, Value, json};
use crate::error::{Error, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
pub const REDUCTO_API_BASE: &str = "https://platform.reducto.ai";
pub const REDUCTO_API_KEY_ENV: &str = "REDUCTO_API_KEY";
pub const REDUCTO_ID_PREFIX: &str = "reducto://";
const PARSE_V3_SUPPORTED_OCR_PARAMS: &[&str] = &["formatting", "retrieval", "settings"];
const PARSE_LEGACY_SUPPORTED_OCR_PARAMS: &[&str] = &["enhance"];
const MISSING_KEY_MESSAGE: &str = "Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()";
const DATA_URI_UPLOAD_REQUIRED: &str =
"Reducto data URI upload must complete before OCR request transformation";
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReductoDocumentSource {
FileId(String),
Upload { bytes: Vec<u8>, mime_type: String },
}
#[derive(Clone, PartialEq, Eq)]
pub struct ReductoUploadRequest {
pub url: String,
pub authorization: String,
pub file_name: &'static str,
pub bytes: Vec<u8>,
pub mime_type: String,
}
pub struct ReductoParseV3Config;
pub struct ReductoParseLegacyConfig;
pub const REDUCTO_PARSE_V3_CONFIG: ReductoParseV3Config = ReductoParseV3Config;
pub const REDUCTO_PARSE_LEGACY_CONFIG: ReductoParseLegacyConfig = ReductoParseLegacyConfig;
pub fn config_for_model(model: &str) -> Option<&'static dyn OcrProviderConfig> {
match model {
"parse-v3" => Some(&REDUCTO_PARSE_V3_CONFIG),
"parse-legacy" => Some(&REDUCTO_PARSE_LEGACY_CONFIG),
_ => None,
}
}
pub fn normalize_api_base(api_base: Option<&str>) -> String {
api_base
.map(str::trim)
.filter(|base| !base.is_empty())
.unwrap_or(REDUCTO_API_BASE)
.trim_end_matches('/')
.to_string()
}
pub fn parse_url(api_base: Option<&str>) -> String {
format!("{}/parse", normalize_api_base(api_base))
}
pub fn upload_url(api_base: Option<&str>) -> String {
format!("{}/upload", normalize_api_base(api_base))
}
pub fn resolve_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
.map(str::to_string)
.or_else(|| {
env_lookup(REDUCTO_API_KEY_ENV)
.map(|key| key.trim().to_string())
.filter(|key| !key.is_empty())
})
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
}
pub fn extract_document_source(document: &Value) -> Result<ReductoDocumentSource, Error> {
let document = document.as_object().ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(document),
})?;
let source = document
.get("document_url")
.and_then(Value::as_str)
.filter(|source| !source.is_empty())
.or_else(|| document.get("image_url").and_then(Value::as_str))
.ok_or_else(|| {
Error::InvalidRequest(
"Reducto expected OCR preprocessing to produce document_url or image_url"
.to_string(),
)
})?;
classify_document_source(source)
}
pub fn classify_document_source(source: &str) -> Result<ReductoDocumentSource, Error> {
if source.starts_with(REDUCTO_ID_PREFIX) {
return Ok(ReductoDocumentSource::FileId(source.to_string()));
}
if source.starts_with("http://") || source.starts_with("https://") {
return Err(Error::InvalidRequest(
"Reducto requires type='file' (auto-uploaded) or a reducto:// id. Plain http(s) URLs are not supported; upload the file first."
.to_string(),
));
}
if !source.starts_with("data:") {
return Err(Error::InvalidRequest(
"Reducto requires a reducto:// id or a base64 data URI after OCR preprocessing."
.to_string(),
));
}
let (header, encoded) = source
.split_once(',')
.ok_or_else(|| Error::InvalidRequest("Invalid Reducto data URI provided.".to_string()))?;
if !header.split(';').any(|part| part == "base64") {
return Err(Error::InvalidRequest(
"Reducto only supports base64-encoded data URIs.".to_string(),
));
}
let mime_type = header
.strip_prefix("data:")
.and_then(|header| header.split(';').next())
.filter(|mime| !mime.is_empty())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = BASE64_STANDARD.decode(encoded).map_err(|_| {
Error::InvalidRequest("Invalid Reducto base64 payload provided.".to_string())
})?;
Ok(ReductoDocumentSource::Upload { bytes, mime_type })
}
pub fn build_upload_request(
source: ReductoDocumentSource,
authorization: &str,
api_base: Option<&str>,
) -> Option<ReductoUploadRequest> {
let ReductoDocumentSource::Upload { bytes, mime_type } = source else {
return None;
};
Some(ReductoUploadRequest {
url: upload_url(api_base),
authorization: authorization.to_string(),
file_name: "document",
bytes,
mime_type,
})
}
pub fn extract_upload_file_id(response_json: &Value) -> Result<&str, Error> {
response_json
.as_object()
.and_then(|response| response.get("file_id"))
.and_then(Value::as_str)
.filter(|file_id| !file_id.is_empty())
.ok_or_else(|| {
Error::InvalidResponse(format!(
"Reducto /upload returned 200 without a file_id; got payload={response_json}"
))
})
}
pub fn build_parse_v3_request(
file_id: &str,
optional_params: Map<String, Value>,
) -> OcrRequestData {
let data = std::iter::once(("input".to_string(), Value::String(file_id.to_string())))
.chain(optional_params)
.collect();
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
pub fn build_parse_legacy_request(
file_id: &str,
optional_params: &Map<String, Value>,
) -> OcrRequestData {
let options = optional_params
.get("enhance")
.filter(|enhance| !enhance.is_null())
.map(|enhance| json!({"options": {"enhance": enhance}}));
let data = match options {
Some(Value::Object(options)) => std::iter::once((
"document_url".to_string(),
Value::String(file_id.to_string()),
))
.chain(options)
.collect(),
_ => Map::from_iter([(
"document_url".to_string(),
Value::String(file_id.to_string()),
)]),
};
OcrRequestData {
data: Value::Object(data),
files: None,
}
}
fn source_file_id(document: &Value) -> Result<String, Error> {
match extract_document_source(document)? {
ReductoDocumentSource::FileId(file_id) => Ok(file_id),
ReductoDocumentSource::Upload { .. } => Err(Error::Unsupported(DATA_URI_UPLOAD_REQUIRED)),
}
}
fn page_number(block: &Map<String, Value>) -> Option<i64> {
let page = block.get("bbox")?.as_object()?.get("page")?;
page.as_i64()
.or_else(|| page.as_u64().and_then(|page| i64::try_from(page).ok()))
.or_else(|| page.as_str().and_then(|page| page.parse().ok()))
}
fn chunks(result: &Map<String, Value>) -> &[Value] {
result
.get("chunks")
.and_then(Value::as_array)
.map(Vec::as_slice)
.unwrap_or_default()
}
fn build_pages(result: &Map<String, Value>) -> Vec<Value> {
let blocks_by_page = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("blocks").and_then(Value::as_array))
.flatten()
.filter_map(|block| block.as_object().map(|object| (block, object)))
.filter_map(|(block, object)| page_number(object).map(|page| (page, block.clone())))
.fold(
BTreeMap::<i64, Vec<Value>>::new(),
|mut pages, (page, block)| {
pages.entry(page).or_default().push(block);
pages
},
);
if blocks_by_page.is_empty() {
let markdown = chunks(result)
.iter()
.filter_map(Value::as_object)
.filter_map(|chunk| chunk.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
return if markdown.is_empty() {
Vec::new()
} else {
vec![json!({"index": 0, "markdown": markdown})]
};
}
blocks_by_page
.into_iter()
.map(|(page, blocks)| {
let markdown = blocks
.iter()
.filter_map(Value::as_object)
.filter_map(|block| block.get("content").and_then(Value::as_str))
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
json!({
"index": page.saturating_sub(1).max(0),
"markdown": markdown,
"blocks": blocks,
})
})
.collect()
}
pub fn transform_reducto_response(
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
let response = response_json
.as_object()
.ok_or_else(|| Error::InvalidType {
expected: "object",
actual: json_type_name(&response_json),
})?;
let empty_result = Map::new();
let result = match response.get("result") {
Some(Value::Object(result)) => result,
Some(Value::Null) => &empty_result,
Some(_) => {
return Err(Error::InvalidResponse(
"Reducto result must be an object".to_string(),
));
}
None => response,
};
let usage = response
.get("usage")
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let usage_info = Some(json!({
"pages_processed": usage.get("num_pages").cloned().unwrap_or(Value::Null),
"credits": usage.get("credits").cloned().unwrap_or(Value::Null),
}));
Ok(OcrResponseData {
pages: build_pages(result),
model: model.to_string(),
document_annotation: None,
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: Some(response_json),
})
}
impl OcrProviderConfig for ReductoParseV3Config {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_V3_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_v3_request(&file_id, optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}
impl OcrProviderConfig for ReductoParseLegacyConfig {
fn supported_ocr_params(&self) -> &'static [&'static str] {
PARSE_LEGACY_SUPPORTED_OCR_PARAMS
}
fn transform_ocr_request(
&self,
_model: &str,
document: Value,
optional_params: Map<String, Value>,
) -> Result<OcrRequestData, Error> {
let file_id = source_file_id(&document)?;
Ok(build_parse_legacy_request(&file_id, &optional_params))
}
fn transform_ocr_response(
&self,
model: &str,
response_json: Value,
) -> Result<OcrResponseData, Error> {
transform_reducto_response(model, response_json)
}
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(parse_url(api_base))
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_api_key(api_key, env_lookup)
}
}

34
litellm-rust/crates/core/src/providers/vertex_ai/ocr/transformation.rs
View file

@ -212,6 +212,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
MISTRAL_OCR_CONFIG.supported_ocr_params()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
@ -229,6 +230,7 @@ impl OcrProviderConfig for VertexAiOcrConfig {
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
@ -253,10 +255,21 @@ impl OcrProviderConfig for VertexAiOcrConfig {
}
impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn supported_ocr_params(&self) -> &'static [&'static str] {
DEEPSEEK_SUPPORTED_OCR_PARAMS
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
non_default_params
.iter()
.filter(|(name, _)| DEEPSEEK_SUPPORTED_OCR_PARAMS.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_request(
&self,
model: &str,
@ -283,6 +296,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn transform_ocr_response(
&self,
model: &str,
@ -335,9 +349,12 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
document_annotation: object.get("document_annotation").cloned(),
usage_info,
object: "ocr".to_string(),
extra_fields: Map::new(),
provider_native_response: None,
})
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn complete_url(
&self,
api_base: Option<&str>,
@ -360,6 +377,7 @@ impl OcrProviderConfig for VertexAiDeepSeekOcrConfig {
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[test]
fn vertex_mistral_url_uses_project_location_and_model() {
@ -411,6 +429,22 @@ mod tests {
);
}
#[rstest]
#[case::bare_model("deepseek-ocr-maas")]
#[case::namespaced_model("deepseek-ai/deepseek-ocr-maas")]
fn vertex_deepseek_request_uses_single_provider_namespace(#[case] model: &str) {
let body = VERTEX_AI_DEEPSEEK_OCR_CONFIG
.transform_ocr_request(
model,
json!({"type": "image_url", "image_url": "data:image/png;base64,AA=="}),
Map::new(),
)
.expect("request transforms")
.data;
assert_eq!(body["model"], "deepseek-ai/deepseek-ocr-maas");
}
#[test]
fn vertex_deepseek_response_wraps_markdown_content() {
let response = VERTEX_AI_DEEPSEEK_OCR_CONFIG

14
litellm-rust/crates/core/tests/workspace_crate_allowlist.rs
View file

@ -1,6 +1,7 @@
//! Enforcement: the litellm-rust workspace has exactly four crates.
//! Enforcement: the litellm-rust workspace has exactly five crates.
//!
//! `core` (the Rust SDK), `ai-gateway` (the HTTP/WebSocket host),
//! `core` (the Rust SDK), `config` (the config-loading boundary),
//! `ai-gateway` (the HTTP/WebSocket host),
//! `python-interop` (domain-neutral PyO3 primitives), and `python-bridge` (the
//! PyO3 cdylib). Adding or removing a crate must be a
//! deliberate act: this test fails until the allowlist here is updated, forcing
@ -19,13 +20,20 @@ use std::path::{Path, PathBuf};
/// workspace legitimately gains or loses a crate.
const EXPECTED_MEMBERS: &[&str] = &[
"crates/core",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
"crates/python-bridge",
];
/// The crate subdirectory names that must exist under `crates/`.
const EXPECTED_CRATE_DIRS: &[&str] = &["core", "ai-gateway", "python-interop", "python-bridge"];
const EXPECTED_CRATE_DIRS: &[&str] = &[
"core",
"config",
"ai-gateway",
"python-interop",
"python-bridge",
];
const MISMATCH: &str = "litellm-rust crate set changed — update this allowlist AND litellm-rust/AGENTS.md, and justify the crate per the rule (crate = layer needing independent compilation / its own deps / a separate artifact).";

9
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -14,11 +14,15 @@ default = ["abi3"]
abi3 = ["pyo3/abi3-py310"]
extension-module = ["pyo3/extension-module"]
panic-test = []
trace-parity = [
"dep:tracing",
"litellm-core/observability",
"litellm-ai-gateway/trace-parity",
]
[dependencies]
futures-util.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
tracing = { workspace = true, optional = true }
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
litellm-python-interop.workspace = true
@ -31,6 +35,7 @@ tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
tokio-tungstenite.workspace = true
tracing.workspace = true
[[bench]]
name = "serialization"

1
litellm-rust/crates/python-bridge/src/constants.rs
View file

@ -1 +0,0 @@
pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace";

230
litellm-rust/crates/python-bridge/src/function_trace.rs
View file

@ -1,216 +1,38 @@
use std::fmt::Display;
use std::future::Future;
use std::sync::{Arc, Mutex};
use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
use serde::Serialize;
use tracing::instrument::WithSubscriber;
use tracing::span::{Attributes, Id};
use tracing::{Dispatch, Level, Subscriber};
use tracing_subscriber::filter::{LevelFilter, filter_fn};
use tracing_subscriber::layer::Context;
use tracing_subscriber::prelude::*;
use tracing_subscriber::registry::LookupSpan;
use tracing_subscriber::{Layer, Registry};
use crate::constants::FUNCTION_TRACE_TARGET;
#[derive(Serialize)]
#[serde(untagged)]
pub(crate) enum TraceResponse<T> {
Plain(T),
Traced {
response: T,
trace: Vec<FunctionTraceEvent>,
},
pub(crate) struct TracedResponse<T> {
#[serde(skip_serializing_if = "Option::is_none")]
response: Option<T>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
trace: Vec<FunctionTraceEvent>,
}
pub(crate) async fn trace_call<T, E>(
pub(crate) async fn capture<T, E>(
future: impl Future<Output = Result<T, E>>,
enabled: bool,
) -> Result<TraceResponse<T>, E> {
if !enabled {
return future.await.map(TraceResponse::Plain);
}
) -> Result<TracedResponse<T>, E>
where
E: Display,
{
let trace = FunctionTrace::default();
let response = future.with_subscriber(trace.dispatcher()).await?;
Ok(TraceResponse::Traced {
response,
trace: trace.events(),
let result = future.with_subscriber(trace.dispatcher()).await;
let events = trace.events();
Ok(match result {
Ok(response) => TracedResponse {
response: Some(response),
error: None,
trace: events,
},
Err(error) => TracedResponse {
response: None,
error: Some(error.to_string()),
trace: events,
},
})
}
#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct FunctionTraceEvent {
pub function: &'static str,
pub depth: usize,
}
#[derive(Clone, Default)]
pub struct FunctionTrace {
events: Arc<Mutex<Vec<FunctionTraceEvent>>>,
}
impl FunctionTrace {
pub fn dispatcher(&self) -> Dispatch {
let filter = filter_fn(|metadata| {
metadata.is_span()
&& metadata.target() == FUNCTION_TRACE_TARGET
&& *metadata.level() == Level::TRACE
})
.with_max_level_hint(LevelFilter::TRACE);
Dispatch::new(
Registry::default().with(
FunctionTraceLayer {
trace: self.clone(),
}
.with_filter(filter),
),
)
}
pub fn events(&self) -> Vec<FunctionTraceEvent> {
self.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.clone()
}
}
struct FunctionTraceLayer {
trace: FunctionTrace,
}
impl<S> Layer<S> for FunctionTraceLayer
where
S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, S>) {
let depth = context
.span(id)
.map(|span| span.scope().skip(1).count())
.unwrap_or_default();
self.trace
.events
.lock()
.unwrap_or_else(|error| error.into_inner())
.push(FunctionTraceEvent {
function: attributes.metadata().name(),
depth,
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn outer() {
tokio::task::yield_now().await;
inner().await;
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
async fn inner() {
tokio::task::yield_now().await;
}
#[tokio::test]
async fn concurrent_futures_keep_separate_traces_across_yields() {
use tracing::instrument::WithSubscriber;
let first = FunctionTrace::default();
let second = FunctionTrace::default();
let outside = FunctionTrace::default();
async {
tokio::join!(
outer().with_subscriber(first.dispatcher()),
inner().with_subscriber(second.dispatcher()),
);
inner().await;
}
.with_subscriber(outside.dispatcher())
.await;
assert_eq!(
first.events(),
vec![
FunctionTraceEvent {
function: "outer",
depth: 0
},
FunctionTraceEvent {
function: "inner",
depth: 1
},
],
);
assert_eq!(
second.events(),
vec![FunctionTraceEvent {
function: "inner",
depth: 0
}],
);
assert_eq!(
outside.events(),
vec![FunctionTraceEvent {
function: "inner",
depth: 0
}],
);
}
#[test]
fn records_matching_spans_in_creation_order() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let _ignored = tracing::trace_span!(target: "other", "ignored");
let _first = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
let _wrong_level = tracing::debug_span!(target: FUNCTION_TRACE_TARGET, "wrong_level");
let _second = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "same_name");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
FunctionTraceEvent {
function: "same_name",
depth: 0,
},
]
);
}
#[test]
fn records_matching_span_nesting_depth() {
let trace = FunctionTrace::default();
let dispatch = trace.dispatcher();
tracing::dispatcher::with_default(&dispatch, || {
let outer = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "outer");
let _outer_guard = outer.enter();
let _inner = tracing::trace_span!(target: FUNCTION_TRACE_TARGET, "inner");
});
assert_eq!(
trace.events(),
vec![
FunctionTraceEvent {
function: "outer",
depth: 0,
},
FunctionTraceEvent {
function: "inner",
depth: 1,
},
]
);
}
}

40
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -1,8 +1,8 @@
mod constants;
mod diagnostics;
mod errors;
mod execution;
pub mod function_trace;
#[cfg(feature = "trace-parity")]
mod function_trace;
mod marshal;
mod routes;
@ -115,9 +115,43 @@ mod tests {
.extract::<Vec<String>>()
.expect("module names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.filter(|name| !name.starts_with('_'))
.collect();
assert_eq!(public_names, expected);
#[cfg(not(feature = "trace-parity"))]
assert!(!module.hasattr("_trace").expect("module lookup should work"));
#[cfg(feature = "trace-parity")]
{
let trace = module
.getattr("_trace")
.expect("trace build should expose its diagnostic namespace");
let trace_names: Vec<String> = trace
.cast::<PyModule>()
.expect("trace namespace should be a module")
.dict()
.keys()
.extract::<Vec<String>>()
.expect("trace names should be strings")
.into_iter()
.filter(|name| !name.starts_with("__"))
.collect();
assert_eq!(
trace_names,
[
"ocr",
"aocr",
"transcription",
"atranscription",
"messages",
"amessages",
"chat_completions",
"achat_completions",
"gateway_messages",
]
);
}
});
}

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