scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 178

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

Browse source 
# Conflicts:
#	type-discipline-budget.json
This commit is contained in:
ryan-crabbe-berri 2026-09-04 14:41:55 -07:00
commit 9bd34adb6d
1288 changed files with 151922 additions and 11227 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
.circleci/config.yml
View file

@ -112,10 +112,10 @@ commands:
node --version
npm --version
install_rust:
description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself."
description: "Install pinned rustup (1.28.2) and Rust toolchain (1.98.0) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself."
steps:
- run:
name: Install Rust (rustup 1.28.2, toolchain 1.97.1)
name: Install Rust (rustup 1.28.2, toolchain 1.98.0)
command: |
case "$(uname -m)" in
x86_64)
@ -135,7 +135,7 @@ commands:
"https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init"
echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c -
chmod +x /tmp/rustup-init
/tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1
/tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.98.0
rm -f /tmp/rustup-init
echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV"
export PATH="$HOME/.cargo/bin:$PATH"
@ -300,7 +300,7 @@ jobs:
if ($rustupActual -ne $rustupExpected) {
throw "rustup installer hash mismatch: expected $rustupExpected got $rustupActual"
}
& $rustupInit -y --profile minimal --default-toolchain stable
& $rustupInit -y --profile minimal --default-toolchain 1.98.0
if ($LASTEXITCODE -ne 0) {
exit $LASTEXITCODE
}
@ -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 }}

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

@ -4,6 +4,11 @@ description: >-
by a job nor listed here, so every entry below is a decision on the record.
test_paths:
- reason: >-
The Rust/Python parity harness is run manually through its local CLI. Recorded replay,
fixture generation, and harness checks are intentionally outside pull request CI
paths:
- tests/rust-python-harness
- reason: >-
What is left of the caching suite in tests/local_testing that runs nowhere. Every job that
globs that directory either deselects it (local_testing_part1 and part2 carry `-k "... and
@ -74,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

68
.github/scripts/smoke_test_native_wheel.py vendored Normal file
View file

@ -0,0 +1,68 @@
from __future__ import annotations
import subprocess
import sys
import tempfile
import zipfile
from pathlib import Path
from typing import Final
CHILD_SCRIPT: Final = """
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
import sys
native_path = Path(sys.argv[1])
spec = spec_from_file_location("litellm.rust_bridge._native", native_path)
if spec is None or spec.loader is None:
raise RuntimeError("cannot create native extension import specification")
module = module_from_spec(spec)
spec.loader.exec_module(module)
before = module.gil_stats()
if not isinstance(before.get("releases"), int):
raise AssertionError(f"unexpected gil_stats result: {before!r}")
try:
module._panic_for_test()
except BaseException as error:
if type(error).__name__ != "PanicException":
raise AssertionError(f"expected PanicException, got {type(error).__name__}") from error
else:
raise AssertionError("Rust panic returned without raising")
after = module.gil_stats()
if not isinstance(after.get("releases"), int):
raise AssertionError(f"native module unusable after panic: {after!r}")
"""
def main() -> int:
if len(sys.argv) != 2:
sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n")
return 2
wheel: Final = Path(sys.argv[1])
with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive:
native_members: Final = tuple(
member
for member in archive.infolist()
if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
)
if len(native_members) != 1:
sys.stderr.write(f"expected one native extension, found {len(native_members)}\n")
return 1
native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name
native_path.write_bytes(archive.read(native_members[0]))
result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False)
if result.returncode != 0:
sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n")
return 1
return 0
if __name__ == "__main__":
sys.exit(main())

282
.github/scripts/verify_linux_native_wheel.py vendored Normal file
View file

@ -0,0 +1,282 @@
from __future__ import annotations
import importlib.util
import os
import re
import subprocess
import sys
import zipfile
from collections.abc import Callable, Mapping, Sequence
from itertools import product
from pathlib import Path, PurePosixPath
from types import MappingProxyType, ModuleType
from typing import Final, Protocol
EXPECTED_PYTHON_TAG: Final = "cp310"
EXPECTED_ABI_TAG: Final = "abi3"
EXPECTED_PLATFORM_TAG: Final = "linux_x86_64"
class CommandRunner(Protocol):
def __call__(
self,
command: tuple[str, ...],
*,
check: bool,
capture_output: bool,
text: bool,
) -> subprocess.CompletedProcess[str]: ...
def _run_command(
command: tuple[str, ...],
*,
check: bool,
capture_output: bool,
text: bool,
) -> subprocess.CompletedProcess[str]:
return subprocess.run(command, check=check, capture_output=capture_output, text=text)
def _dist_info_directory(member: zipfile.ZipInfo) -> str | None:
parts: Final = PurePosixPath(member.filename).parts
if not parts or not parts[0].endswith(".dist-info"):
return None
return parts[0]
def _wheel_metadata_tags(archive: zipfile.ZipFile, members: tuple[zipfile.ZipInfo, ...]) -> tuple[str, ...]:
if len(members) != 1:
return ()
lines: Final = archive.read(members[0]).splitlines()
return tuple(line.removeprefix(b"Tag:").strip().decode("ascii") for line in lines if line.startswith(b"Tag:"))
def _load_native_module(native_path: Path) -> ModuleType | None:
module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path)
if module_spec is None or module_spec.loader is None:
return None
try:
native_module: Final = importlib.util.module_from_spec(module_spec)
module_spec.loader.exec_module(native_module)
except Exception as error: # noqa: BLE001 # native module initialization can raise arbitrary exceptions
sys.stderr.write(f"native module load failed: {error}\n")
return None
return native_module
def main(
argv: Sequence[str] | None = None,
environment: Mapping[str, str] | None = None,
load_native_module: Callable[[Path], ModuleType | None] = _load_native_module,
run_command: CommandRunner = _run_command,
) -> int:
arguments: Final = tuple(sys.argv if argv is None else argv)
resolved_environment: Final = os.environ if environment is None else environment
if len(arguments) != 2:
sys.stderr.write(f"usage: {Path(arguments[0]).name} WHEEL\n")
return 2
wheel: Final = Path(arguments[1])
wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3)
if len(wheel_tags) != 4:
sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n")
return 1
wheel_identity: Final = wheel_tags[0].split("-")
if len(wheel_identity) != 2 or wheel_identity[0] != "litellm" or not wheel_identity[1]:
sys.stderr.write(f"unexpected wheel identity: {wheel_tags[0]}\n")
return 1
expected_dist_info_directory: Final = f"{wheel_tags[0]}.dist-info"
expected_dist_info_directories: Final = frozenset((expected_dist_info_directory,))
python_tag: Final = wheel_tags[1]
abi_tag: Final = wheel_tags[2]
platform_tag: Final = wheel_tags[3]
expanded_filename_tags: Final = frozenset(
"-".join(tag) for tag in product(python_tag.split("."), abi_tag.split("."), platform_tag.split("."))
)
with zipfile.ZipFile(wheel) as archive:
wheel_members: Final = archive.infolist()
dist_info_directories: Final = frozenset(
directory for member in wheel_members if (directory := _dist_info_directory(member)) is not None
)
required_dist_info_files: Final = ("METADATA", "RECORD", "WHEEL")
dist_info_file_counts: Final = MappingProxyType(
{
filename: sum(
member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members
)
for filename in required_dist_info_files
}
)
wheel_metadata_members: Final = tuple(
member for member in wheel_members if member.filename == f"{expected_dist_info_directory}/WHEEL"
)
wheel_metadata_tags: Final = _wheel_metadata_tags(archive, wheel_metadata_members)
native_members: Final = tuple(
member
for member in wheel_members
if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
)
if len(native_members) != 1:
sys.stderr.write(f"expected one native extension, found {len(native_members)}\n")
return 1
unexpected_members: Final = tuple(
member.filename
for member in wheel_members
if member.filename.endswith((".pdb", ".dwp", ".rlib", ".rmeta", "Cargo.toml", "Cargo.lock"))
or any(part.endswith(".dSYM") for part in PurePosixPath(member.filename).parts)
)
native_member: Final = native_members[0]
uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members)
native_path: Final = wheel.parent / "native" / Path(native_member.filename).name
native_path.parent.mkdir(parents=True, exist_ok=True)
native_path.write_bytes(archive.read(native_member))
wheel_metadata_tags_match: Final = (
len(wheel_metadata_tags) == len(expanded_filename_tags)
and frozenset(wheel_metadata_tags) == expanded_filename_tags
)
commit_sha: Final = resolved_environment.get(
"RELEASE_WHEEL_COMMIT_SHA", resolved_environment.get("GITHUB_SHA", "unknown")
)
rustc_version: Final = run_command(
("rustc", "--version"),
check=True,
capture_output=True,
text=True,
).stdout.strip()
pyproject: Final = (Path(__file__).parents[2] / "pyproject.toml").read_text()
maturin_match: Final = re.search(r'"maturin==([^";]+)', pyproject)
if maturin_match is None:
sys.stderr.write("build-system does not pin an exact Maturin version\n")
return 1
maturin_version: Final = maturin_match.group(1)
native_percentage: Final = native_member.file_size / uncompressed_wheel_size * 100
size_report: Final = "\n".join(
(
"## Native wheel build report",
"",
"| Build | Value |",
"| --- | --- |",
f"| Commit | `{commit_sha}` |",
f"| Platform | `{platform_tag}` |",
f"| Python ABI | `{python_tag}-{abi_tag}` |",
f"| Rust compiler | `{rustc_version}` |",
f"| Maturin | `{maturin_version}` |",
"| Cargo profile | `release` |",
"",
"| Artifact | Size |",
"| --- | ---: |",
f"| Compressed wheel | {wheel.stat().st_size / 1_000_000:.2f} MB |",
f"| Uncompressed wheel | {uncompressed_wheel_size / 1_000_000:.2f} MB |",
f"| Native extension | {native_member.file_size / 1_000_000:.2f} MB |",
f"| Native share | {native_percentage:.2f}% |",
"",
)
)
summary_path: Final = resolved_environment.get("GITHUB_STEP_SUMMARY")
if summary_path is None:
sys.stdout.write(size_report)
else:
Path(summary_path).write_text(size_report)
sections: Final = run_command(
("readelf", "--sections", "--wide", str(native_path)),
check=True,
capture_output=True,
text=True,
).stdout
debug_sections: Final = tuple(section for section in (".debug_", ".zdebug_") if section in sections)
debug_sections_absent: Final = not debug_sections
static_symbol_table_absent: Final = ".symtab" not in sections
dynamic_symbols: Final = run_command(
("readelf", "--dyn-syms", "--wide", str(native_path)),
check=True,
capture_output=True,
text=True,
).stdout
extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols
native_module: Final = load_native_module(native_path)
native_module_loads: Final = native_module is not None
panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
native_size_limit: Final = 20_000_000
native_size_within_limit: Final = native_member.file_size <= native_size_limit
validations: Final = (
(f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG),
(f"ABI tag is {EXPECTED_ABI_TAG}", abi_tag == EXPECTED_ABI_TAG),
(f"Platform tag is {EXPECTED_PLATFORM_TAG}", platform_tag == EXPECTED_PLATFORM_TAG),
("Wheel dist-info directory matches the filename", dist_info_directories == expected_dist_info_directories),
(
"Required dist-info files are present exactly once",
all(count == 1 for count in dist_info_file_counts.values()),
),
("Wheel metadata tags match the filename", wheel_metadata_tags_match),
("Debug sections are absent", debug_sections_absent),
("Static symbol table is absent", static_symbol_table_absent),
("Python extension entry point is present", extension_entry_point_present),
("Native module loads", native_module_loads),
("Production module omits the panic test hook", panic_test_hook_absent),
("Native extension does not exceed 20 MB", native_size_within_limit),
("Wheel contents are valid", not unexpected_members),
)
verified_report: Final = size_report + "\n".join(
("", "| Validation | Expected | Result |", "| --- | --- | :---: |")
+ tuple(f"| {label} | Yes | {'O' if passed else 'X'} |" for label, passed in validations)
+ ("",)
)
if summary_path is not None:
Path(summary_path).write_text(verified_report)
invalid_dist_info_files: Final = any(count != 1 for count in dist_info_file_counts.values())
validation_errors: Final = tuple(
message
for failed, message in (
(bool(debug_sections), f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}"),
(not static_symbol_table_absent, f"{native_member.filename} contains a static symbol table"),
(not extension_entry_point_present, "native extension does not export PyInit__native"),
(
python_tag != EXPECTED_PYTHON_TAG,
f"unexpected Python tag: expected {EXPECTED_PYTHON_TAG}, found {python_tag}",
),
(abi_tag != EXPECTED_ABI_TAG, f"unexpected ABI tag: expected {EXPECTED_ABI_TAG}, found {abi_tag}"),
(
platform_tag != EXPECTED_PLATFORM_TAG,
f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}",
),
(
dist_info_directories != expected_dist_info_directories,
f"unexpected dist-info directories: expected {expected_dist_info_directory}, "
f"found {', '.join(sorted(dist_info_directories))}",
),
(invalid_dist_info_files, f"required dist-info file counts are invalid: {dist_info_file_counts}"),
(
not invalid_dist_info_files and not wheel_metadata_tags_match,
f"WHEEL tags do not match filename: expected {', '.join(sorted(expanded_filename_tags))}, "
f"found {', '.join(sorted(wheel_metadata_tags))}",
),
(
native_module is not None and not panic_test_hook_absent,
"production native module exposes _panic_for_test",
),
(
not native_size_within_limit,
f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB",
),
(bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"),
)
if failed
)
sys.stderr.write("".join(f"{message}\n" for message in validation_errors))
return 0 if all(passed for _, passed in validations) else 1
if __name__ == "__main__":
sys.exit(main())

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

@ -57,9 +57,15 @@ permissions:
jobs:
run:
name: Run tests
name: ${{ matrix.python-version == '3.12' && 'Run tests' || format('Run tests (Python {0})', matrix.python-version) }}
runs-on: ubuntu-latest
timeout-minutes: ${{ inputs.job-timeout-minutes }}
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
env:
UV_PYTHON: ${{ matrix.python-version }}
permissions:
contents: read
pull-requests: read
@ -82,7 +88,7 @@ jobs:
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
python-version: ${{ matrix.python-version }}
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
@ -96,12 +102,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${{ matrix.python-version }}-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}-
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
@ -113,6 +117,7 @@ jobs:
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
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,13 +139,7 @@ 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
COVERAGE_CORE: ${{ contains(fromJSON('["3.10", "3.11"]'), matrix.python-version) && 'ctrace' || 'sysmon' }}
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \
@ -167,7 +166,7 @@ jobs:
fi
- name: Save coverage report
if: always() && steps.changes.outputs.decision != 'skip'
if: always() && matrix.python-version == '3.12' && steps.changes.outputs.decision != 'skip'
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}

6
.github/workflows/image-scan.yml vendored
View file

@ -80,7 +80,7 @@ jobs:
LITELLM_IMAGE: litellm-image-scan:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
# Scans the whole shipped artifact: OS/apk plus every language package
# baked into the image, including ones no lockfile declares (e.g. prisma's
@ -124,7 +124,7 @@ jobs:
LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
migrations-image:
name: migrations-image
@ -185,7 +185,7 @@ jobs:
LITELLM_COMPONENT_PORT: "4000"
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
ui-image:
name: ui-image

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

@ -0,0 +1,130 @@
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,
});
}

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

@ -128,6 +128,9 @@ jobs:
- name: check_fastuuid_usage
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
- name: check_py310_typing_imports
run: uv run --no-sync python ./tests/code_coverage_tests/check_py310_typing_imports.py
- name: check_e2e_no_raw_requests
run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
@ -145,3 +148,33 @@ jobs:
- name: documentation_test_api_docs
run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py
python-310-import-smoke:
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.10"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
run: uv sync --frozen --extra proxy --python 3.10
- run: uv run --no-sync python --version
- name: Import litellm
run: uv run --no-sync python -c "import litellm"
- name: Check litellm CLI
run: uv run --no-sync litellm --version

24
.github/workflows/test-litellm-ui-build.yml vendored
View file

@ -19,9 +19,6 @@ jobs:
build-ui:
runs-on: ubuntu-latest
timeout-minutes: 10
defaults:
run:
working-directory: ui/litellm-dashboard
steps:
- name: Checkout repository
@ -35,18 +32,11 @@ jobs:
with:
category: ui
- name: Setup Node.js
# Built through the image stage rather than the checkout, because the
# stage copies ui/litellm-dashboard/ alone: an import reaching above the
# dashboard root resolves in a checkout and fails in every image we ship.
# Dockerfile, docker/Dockerfile.non_root and ui/Dockerfile share this
# stage verbatim, so building one covers all three.
- name: Build the dashboard as the shipped images build it
if: steps.changes.outputs.decision != 'skip'
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: npm ci
- name: Build
if: steps.changes.outputs.decision != 'skip'
run: npm run build
run: docker build --target ui-builder -f Dockerfile .

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

@ -4,6 +4,12 @@ on:
push:
paths:
- "litellm-rust/**"
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
pull_request:
branches:
@ -13,6 +19,12 @@ on:
- "litellm_**"
paths:
- "litellm-rust/**"
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
- ".github/workflows/test-rust.yml"
permissions:
@ -40,9 +52,7 @@ jobs:
persist-credentials: false
- name: Set up Rust
run: |
rustup toolchain install stable --profile minimal --component clippy,rustfmt
rustup default stable
run: rustup toolchain install
- name: Cache Cargo registry and target
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
@ -51,7 +61,7 @@ jobs:
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }}
key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-
@ -64,8 +74,62 @@ jobs:
- name: Run Clippy with Bedrock auth
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --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
runs-on: ubuntu-latest
timeout-minutes: 20
permissions:
contents: read
env:
CARGO_TERM_COLOR: always
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Set up Rust
run: rustup toolchain install
- name: Build release wheel
run: uv build --wheel --out-dir dist
- name: Build panic contract wheel
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

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

@ -150,6 +150,7 @@ jobs:
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/container_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints

8
Dockerfile
View file

@ -66,6 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
# Copy full source tree
@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
@ -101,6 +103,12 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
# The base image only configures Chainguard's authenticated apk repo, which
# requires an enterprise subscription. Add the public Wolfi repo so `apk add`
# also works for anyone installing extra packages into a running container.
# https://github.com/BerriAI/litellm/issues/33518
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile

2
backend/Dockerfile
View file

@ -46,6 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
@ -57,6 +58,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \

30
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 14076
"limit": 14074
},
"reportArgumentType": {
"limit": 2216
"limit": 2206
},
"reportAssignmentType": {
"limit": 319
@ -18,13 +18,13 @@
"limit": 40
},
"reportDeprecated": {
"limit": 211
"limit": 209
},
"reportDuplicateImport": {
"limit": 19
},
"reportExplicitAny": {
"limit": 4128
"limit": 4124
},
"reportFunctionMemberAccess": {
"limit": 7
@ -42,10 +42,10 @@
"limit": 12
},
"reportIndexIssue": {
"limit": 25
"limit": 24
},
"reportInvalidTypeForm": {
"limit": 34
"limit": 30
},
"reportInvalidTypeVarUse": {
"limit": 2
@ -57,7 +57,7 @@
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15306
"limit": 15285
},
"reportMissingTypeStubs": {
"limit": 40
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44364
"limit": 44360
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38350
"limit": 38309
},
"reportUnknownParameterType": {
"limit": 19626
"limit": 19622
},
"reportUnknownVariableType": {
"limit": 29890
"limit": 29846
},
"reportUnnecessaryCast": {
"limit": 111
},
"reportUnnecessaryComparison": {
"limit": 695
"limit": 687
},
"reportUnnecessaryContains": {
"limit": 5
"limit": 4
},
"reportUnnecessaryIsInstance": {
"limit": 826
"limit": 819
},
"reportUntypedBaseClass": {
"limit": 0
@ -138,7 +138,7 @@
"limit": 138
},
"reportUnusedImport": {
"limit": 543
"limit": 542
},
"reportUnusedVariable": {
"limit": 137

2
docker/Dockerfile.database
View file

@ -64,6 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
# Copy full source tree
@ -85,6 +86,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--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

@ -70,6 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
# Copy full source tree
@ -97,6 +98,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13 \
--no-sources-package litellm-proxy-extras; \
else \
@ -106,6 +108,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13; \
fi

5
docker/build_from_pip/Dockerfile.build_from_pip
View file

@ -55,7 +55,10 @@ RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/
ENV PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
COPY docker/prod_entrypoint.sh /app/docker/prod_entrypoint.sh
RUN sed -i 's/\r$//' /app/docker/prod_entrypoint.sh && chmod +x /app/docker/prod_entrypoint.sh
EXPOSE 4000/tcp
ENTRYPOINT ["litellm"]
ENTRYPOINT ["/app/docker/prod_entrypoint.sh"]
CMD ["--port", "4000"]

10
docker/component_entrypoint.sh
View file

@ -1,8 +1,10 @@
#!/bin/sh
if [ "$USE_DDTRACE" = "true" ]; then
export DD_TRACE_OPENAI_ENABLED="False"
exec ddtrace-run "$@"
fi
case "$USE_DDTRACE" in
[Tt][Rr][Uu][Ee])
export DD_TRACE_OPENAI_ENABLED="False"
exec ddtrace-run "$@"
;;
esac
exec "$@"

14
docker/prod_entrypoint.sh
View file

@ -1,8 +1,10 @@
#!/bin/sh
if [ "$USE_DDTRACE" = "true" ]; then
export DD_TRACE_OPENAI_ENABLED="False"
exec ddtrace-run litellm "$@"
else
exec litellm "$@"
fi
case "$USE_DDTRACE" in
[Tt][Rr][Uu][Ee])
export DD_TRACE_OPENAI_ENABLED="False"
exec ddtrace-run litellm "$@"
;;
esac
exec litellm "$@"

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

@ -11,17 +11,35 @@ import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import functools
import tempfile
from typing import Optional
from contextvars import ContextVar
from typing import TYPE_CHECKING, ClassVar, Literal, Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import walk_user_text
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
GUARDRAIL_NAME = "hide_secrets"
GUARDRAIL_PROVIDER = "hide-secrets"
# Per-invocation tally of redacted secrets by detect-secrets plugin type; None
# means the guardrail did not run, so _process_response records nothing.
_masked_entity_count: ContextVar[Optional[dict]] = ContextVar(
"hide_secrets_masked_entity_count", default=None
)
_custom_plugins_path = "file://" + os.path.join(
os.path.dirname(os.path.abspath(__file__)), "secrets_plugins"
)
@ -422,6 +440,10 @@ _default_detect_secrets_config = {
class _ENTERPRISE_SecretDetection(CustomGuardrail):
# Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
# path skips should_run_check and never sees data["prompt"]).
use_native_lifecycle_hooks: ClassVar[bool] = True
def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs):
self.user_defined_detect_secrets_config = detect_secrets_config
super().__init__(**kwargs)
@ -455,6 +477,26 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
return detected_secrets
def redact_text(self, text: str, source: str = "message") -> str:
"""Replace every detected secret in ``text`` with ``[REDACTED]`` and
tally the detected types into the per-invocation masked-entity count."""
detected_secrets = self.scan_message_for_secrets(text)
if not detected_secrets:
return text
counts = _masked_entity_count.get()
if counts is not None:
for secret in detected_secrets:
counts[secret["type"]] = counts.get(secret["type"], 0) + 1
secret_types = [secret["type"] for secret in detected_secrets]
verbose_proxy_logger.warning(
f"Detected and redacted secrets in {source}: {secret_types}"
)
return functools.reduce(
lambda redacted, secret: redacted.replace(secret["value"], "[REDACTED]"),
detected_secrets,
text,
)
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
if user_api_key_dict.permissions is not None:
if GUARDRAIL_NAME in user_api_key_dict.permissions:
@ -463,7 +505,45 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
return True
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
"""Unified-interface entrypoint, used by /guardrails/apply_guardrail
(the UI test playground). Proxied traffic keeps using
``async_pre_call_hook``, see ``use_native_lifecycle_hooks``."""
texts = inputs.get("texts")
if not texts or not any(texts):
return inputs
_masked_entity_count.set({})
return {**inputs, "texts": [self.redact_text(text) for text in texts]}
def _redact_prompt(self, data: dict) -> int:
"""Redact ``data["prompt"]`` (the text-completion shape, which
``walk_user_text`` does not cover) and return how many non-empty
strings were inspected."""
prompt = data.get("prompt")
if isinstance(prompt, str):
if not prompt:
return 0
data["prompt"] = self.redact_text(prompt, source="prompt")
return 1
if isinstance(prompt, list):
data["prompt"] = [ # mutable-ok: data["prompt"] is a list on the wire
self.redact_text(item, source="prompt")
if isinstance(item, str) and item
else item
for item in prompt
]
return sum(1 for item in prompt if isinstance(item, str) and item)
return 0
#### CALL HOOKS - proxy only ####
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -471,53 +551,84 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
data: dict,
call_type: str, # "completion", "embeddings", "image_generation", "moderation"
):
_masked_entity_count.set(None)
if await self.should_run_check(user_api_key_dict) is False:
return
_masked_entity_count.set({})
# Covers multimodal list content + Responses-API input.
def _redact_message_text(text: str) -> str:
detected_secrets = self.scan_message_for_secrets(text)
for secret in detected_secrets:
text = text.replace(secret["value"], "[REDACTED]")
if detected_secrets:
secret_types = [secret["type"] for secret in detected_secrets]
verbose_proxy_logger.warning(
f"Detected and redacted secrets in message: {secret_types}"
)
return text
inspected = walk_user_text(data, self.redact_text) + self._redact_prompt(data)
walk_user_text(data, _redact_message_text)
if inspected == 0:
# Image-only, empty-text, and unsupported payloads inspected
# nothing, so recording "allow" would count a run that never
# looked at any content.
_masked_entity_count.set(None)
if "prompt" in data:
if isinstance(data["prompt"], str):
detected_secrets = self.scan_message_for_secrets(data["prompt"])
for secret in detected_secrets:
data["prompt"] = data["prompt"].replace(
secret["value"], "[REDACTED]"
)
if len(detected_secrets) > 0:
secret_types = [secret["type"] for secret in detected_secrets]
verbose_proxy_logger.warning(
f"Detected and redacted secrets in prompt: {secret_types}"
)
elif isinstance(data["prompt"], list):
# Index back into the list — assigning to ``item`` would only
# rebind the loop variable and leave ``data["prompt"]``
# carrying the unredacted secret.
for idx, item in enumerate(data["prompt"]):
if isinstance(item, str):
detected_secrets = self.scan_message_for_secrets(item)
for secret in detected_secrets:
item = item.replace(secret["value"], "[REDACTED]")
data["prompt"][idx] = item
if len(detected_secrets) > 0:
secret_types = [
secret["type"] for secret in detected_secrets
]
verbose_proxy_logger.warning(
f"Detected and redacted secrets in prompt: {secret_types}"
)
# ``data["input"]`` (Responses API and embeddings/moderation) is
# already covered by ``walk_user_text`` above.
return
def _process_response(
self,
response: Optional[dict],
request_data: dict,
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
event_type: Optional[GuardrailEventHooks] = None,
original_inputs: Optional[dict] = None,
):
"""Record allow/mask plus the masked-entity tally for a completed run.
Records nothing when the guardrail inspected nothing (opted-out key,
empty inputs) or when the instance has no guardrail_name (legacy
``litellm_settings.callbacks`` deployments, which predate guardrail
telemetry and stay without it).
"""
counts = _masked_entity_count.get()
_masked_entity_count.set(None)
if counts is None or self.guardrail_name is None:
return response
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response="mask" if counts else "allow",
request_data=request_data,
guardrail_status="success",
duration=duration,
start_time=start_time,
end_time=end_time,
event_type=event_type,
guardrail_provider=GUARDRAIL_PROVIDER,
masked_entity_count=counts,
)
return response
def _process_error(
self,
e: Exception,
request_data: dict,
start_time: Optional[float] = None,
end_time: Optional[float] = None,
duration: Optional[float] = None,
event_type: Optional[GuardrailEventHooks] = None,
):
"""Label the failed run with this guardrail's provider so error rows
group with the successful ones in the monitor. Nameless legacy
instances record nothing, matching ``_process_response``."""
_masked_entity_count.set(None)
if self.guardrail_name is None:
raise e
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=e,
request_data=request_data,
guardrail_status=(
"guardrail_intervened"
if self._is_guardrail_intervention(e)
else "guardrail_failed_to_respond"
),
duration=duration,
start_time=start_time,
end_time=end_time,
event_type=event_type,
guardrail_provider=GUARDRAIL_PROVIDER,
)
raise e

27
enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
View file

@ -48,6 +48,18 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec
}
def _build_search_condition(search: str) -> dict[str, object]:
"""Match a row whose id, changed_by, object_id, or changed_by_api_key equals the search value."""
return {
"OR": (
{"id": search},
{"changed_by": search},
{"object_id": search},
{"changed_by_api_key": search},
)
}
@router.get(
"/audit",
tags=["Audit Logging"],
@ -83,6 +95,10 @@ async def get_audit_logs(
None,
description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
),
search: str | None = Query(
None,
description="Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value",
),
# Sorting parameters
sort_by: str | None = Query(
None,
@ -118,6 +134,11 @@ async def get_audit_logs(
*([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []),
]
and_conditions: Final[tuple[dict[str, object], ...]] = (
*json_field_conditions,
*((_build_search_condition(search),) if search else ()),
)
# Build filter conditions
where_conditions: Final[dict[str, object]] = {
**({"changed_by": changed_by} if changed_by else {}),
@ -126,14 +147,14 @@ async def get_audit_logs(
**({"table_name": table_name} if table_name else {}),
**({"object_id": object_id} if object_id else {}),
**({"updated_at": date_filter} if start_date or end_date else {}),
**({"AND": json_field_conditions} if json_field_conditions else {}),
**({"AND": and_conditions} if and_conditions else {}),
}
order_by: Final[dict[str, str]] = (
{sort_by: sort_order} if sort_by and isinstance(sort_by, str) else {"updated_at": sort_order}
)
audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table
# Get paginated results
audit_logs: Final = await audit_log_table.find_many(
@ -195,7 +216,7 @@ async def get_audit_log_by_id(
detail={"message": CommonProxyErrors.db_not_connected_error.value},
)
audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table
# Get the audit log by ID
audit_log: Final = await audit_log_table.find_unique(where={"id": id})

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

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

@ -50,10 +50,12 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
MAX_FILE_LIST_LIMIT,
_is_base64_encoded_unified_file_id,
apply_unified_file_ids,
decode_model_from_file_id,
ensure_batch_response_managed_file_ids,
get_batch_id_from_unified_batch_id,
get_content_type_from_file_object,
get_model_id_from_unified_batch_id,
get_original_file_id,
map_raw_file_ids_to_unified,
normalize_mime_type_for_provider,
resolve_managed_output_file_model_name,
@ -427,6 +429,103 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
detail=f"Object not found: {unified_object_id}",
)
async def enforce_batch_object_access(
self, object_id: str, user_api_key_dict: UserAPIKeyAuth
) -> None:
"""Deny access to a provider-format batch id owned by another caller.
Ids with no ownership row (batches created before ownership tracking,
or directly on the provider account) stay accessible so pass-through
reads keep working.
"""
if self.prisma_client is None:
return
managed_object = (
await self.prisma_client.db.litellm_managedobjecttable.find_first(
where={"OR": [{"unified_object_id": object_id}, {"model_object_id": object_id}]}
)
)
if managed_object is None:
return
if not can_access_resource(
user_api_key_dict=user_api_key_dict,
created_by=managed_object.created_by,
resource_team_id=managed_object.team_id,
):
raise HTTPException(
status_code=403,
detail=f"User {user_api_key_dict.user_id} does not have access to the object {object_id}",
)
async def enforce_provider_file_access(
self, file_id: str, user_api_key_dict: UserAPIKeyAuth
) -> None:
"""Deny access to a provider-format file id owned by another caller.
Ownership rows for provider-format ids are written when a managed
batch's output/error files are first synced; ids with no row stay
accessible so pass-through reads keep working.
"""
if self.prisma_client is None:
return
managed_file = (
await self.prisma_client.db.litellm_managedfiletable.find_first(
where={"OR": [{"unified_file_id": file_id}, {"flat_model_file_ids": {"has": file_id}}]}
)
)
if managed_file is None:
return
if not can_access_resource(
user_api_key_dict=user_api_key_dict,
created_by=managed_file.created_by,
resource_team_id=managed_file.team_id,
):
raise HTTPException(
status_code=403,
detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}",
)
async def store_batch_output_file_ownership(
self, response: LiteLLMBatch, litellm_parent_otel_span: Optional[Span]
) -> None:
"""Record ownership rows for a batch's provider-format output/error
file ids, inherited from the owning batch row (never the caller), so
file reads can be isolation-checked."""
provider_file_ids = tuple(
file_id
for file_id in (
getattr(response, "output_file_id", None),
getattr(response, "error_file_id", None),
)
if file_id and not _is_base64_encoded_unified_file_id(file_id)
)
if not provider_file_ids:
return
if self.prisma_client is None:
return
batch_row = (
await self.prisma_client.db.litellm_managedobjecttable.find_first(
where={"unified_object_id": response.id}
)
)
if batch_row is None or (
batch_row.created_by is None and batch_row.team_id is None
):
return
owner_identity = UserAPIKeyAuth(
user_id=batch_row.created_by, team_id=batch_row.team_id
)
for file_id in provider_file_ids:
model_name = decode_model_from_file_id(file_id)
raw_file_id = get_original_file_id(file_id)
await self.store_unified_file_id(
file_id=file_id,
file_object=None,
litellm_parent_otel_span=litellm_parent_otel_span,
model_mappings={model_name: raw_file_id} if model_name else {},
user_api_key_dict=owner_identity,
)
async def list_user_batches(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -613,6 +712,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
status_code=403,
detail=f"User {user_api_key_dict.user_id} does not have access to the file {retrieve_file_id}",
)
if retrieve_file_id:
await self.enforce_provider_file_access(
retrieve_file_id, user_api_key_dict
)
return False
async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None:
@ -765,6 +868,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
data["model"] = potential_model_id
data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id)
elif retrieve_object_id and accessor_key == "batch_id":
await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
elif call_type == CallTypes.acreate_fine_tuning_job.value:
input_file_id = cast(Optional[str], data.get("training_file"))
if input_file_id:
@ -1297,7 +1402,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict=user_api_key_dict,
request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}),
persist_attribution=is_batch_create,
create_if_missing=is_batch_create,
)
if not is_batch_create:
await self.store_batch_output_file_ownership(
response=response,
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.

4
enterprise/litellm_enterprise/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -39,9 +39,9 @@ async def available_enterprise_users(
if not premium_user:
# check if SSO is enabled - show 5 user limit
from litellm.proxy.auth.auth_utils import _has_user_setup_sso
from litellm.proxy.auth.auth_utils import has_user_setup_sso
if _has_user_setup_sso():
if has_user_setup_sso():
premium_user_data = EnterpriseLicenseData(
max_users=5,
)

4
enterprise/pyproject.toml
View file

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

2
helm/litellm-helm/Chart.yaml
View file

@ -18,7 +18,7 @@ type: application
# This is the chart version. This version number should be incremented each time you make changes
# to the chart and its templates, including the app version.
# Versions are expected to follow Semantic Versioning (https://semver.org/)
version: 1.1.2
version: 1.1.3
# This is the version number of the application being deployed. This version number should be
# incremented each time you make changes to the application. Versions are not expected to

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

@ -26,7 +26,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` |
| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A |
| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A |
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A |
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
@ -212,6 +212,8 @@ service, the **Proxy Endpoint** should be set to `http://<RELEASE>-litellm:4000`
The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey`
was not provided to the helm command line, the `masterkey` is a randomly
generated string in the `sk-...` format stored in the `<RELEASE>-litellm-masterkey` Kubernetes Secret.
The key is generated once on the first install; later `helm upgrade` runs reuse the
value already in that Secret, so upgrading never rotates the master key.
```bash
kubectl -n litellm get secret <RELEASE>-litellm-masterkey -o jsonpath="{.data.masterkey}"

6
helm/litellm-helm/templates/secret-masterkey.yaml
View file

@ -1,9 +1,11 @@
{{- if not .Values.masterkeySecretName }}
{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }}
{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }}
{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }}
{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }}
apiVersion: v1
kind: Secret
metadata:
name: {{ include "litellm.fullname" . }}-masterkey
name: {{ $secretName }}
data:
masterkey: {{ $masterkey | b64enc }}
type: Opaque

42
helm/litellm-helm/tests/hpa_tests.yaml
View file

@ -1,4 +1,4 @@
suite: "hpa with behavior"
suite: "hpa"
templates:
- hpa.yaml
tests:
@ -23,14 +23,44 @@ tests:
- equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 }
- equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 }
---
suite: "hpa without behavior"
templates:
- hpa.yaml
tests:
- it: "does not render behavior when not set"
set:
autoscaling.enabled: true
asserts:
- isKind: { of: HorizontalPodAutoscaler }
- isNull: { path: spec.behavior }
- it: "scales on cpu at the documented 60 percent by default"
set:
autoscaling.enabled: true
asserts:
- isKind: { of: HorizontalPodAutoscaler }
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
- equal: { path: "spec.metrics[0].resource.target.type", value: Utilization }
- equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 60 }
- it: "does not scale on memory by default"
set:
autoscaling.enabled: true
asserts:
- lengthEqual: { path: spec.metrics, count: 1 }
- it: "honours an explicit cpu target override"
set:
autoscaling.enabled: true
autoscaling.targetCPUUtilizationPercentage: 75
asserts:
- equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 75 }
- it: "renders a memory metric only when a memory target is set"
set:
autoscaling.enabled: true
autoscaling.targetMemoryUtilizationPercentage: 80
asserts:
- lengthEqual: { path: spec.metrics, count: 2 }
- equal: { path: "spec.metrics[1].resource.name", value: memory }
- equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 }
- it: "renders no hpa when autoscaling is disabled"
asserts:
- hasDocuments: { count: 0 }

47
helm/litellm-helm/tests/masterkey-secret_tests.yaml
View file

@ -15,6 +15,53 @@ tests:
# Note: The masterkey is generated as "sk-<18-random-chars>" in plain text,
# but stored as base64 encoded in Kubernetes secret (requirement).
# "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern.
- it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade
template: secret-masterkey.yaml
set:
masterkeySecretName: ""
kubernetesProvider:
scheme:
"v1/Secret":
gvr:
version: "v1"
resource: "secrets"
namespaced: true
objects:
- kind: Secret
apiVersion: v1
metadata:
name: RELEASE-NAME-litellm-masterkey
namespace: NAMESPACE
data:
masterkey: c2stZXhpc3Rpbmcta2V5
asserts:
- equal:
path: data.masterkey
value: c2stZXhpc3Rpbmcta2V5
- it: should let an explicit masterkey value override the one already stored in the cluster
template: secret-masterkey.yaml
set:
masterkeySecretName: ""
masterkey: sk-explicit
kubernetesProvider:
scheme:
"v1/Secret":
gvr:
version: "v1"
resource: "secrets"
namespaced: true
objects:
- kind: Secret
apiVersion: v1
metadata:
name: RELEASE-NAME-litellm-masterkey
namespace: NAMESPACE
data:
masterkey: c2stZXhpc3Rpbmcta2V5
asserts:
- equal:
path: data.masterkey
value: c2stZXhwbGljaXQ=
- it: should not create a secret if masterkeySecretName is set
template: secret-masterkey.yaml
set:

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

@ -200,7 +200,16 @@ autoscaling:
enabled: false
minReplicas: 1
maxReplicas: 100
targetCPUUtilizationPercentage: 80
# 60 is the documented recommendation. See "Recommended Machine Specifications"
# in https://docs.litellm.ai/docs/proxy/prod. A new replica clears the startupProbe
# above only after up to failureThreshold x periodSeconds = 300 seconds, so a target
# high enough to trip near saturation adds capacity minutes after it was needed.
targetCPUUtilizationPercentage: 60
# Deliberately left unset rather than given a value. The prisma query engine's
# resident memory is a high-water mark that ratchets to the pod's worst-ever write
# and is never returned, so a memory target reads the largest write a pod ever did
# rather than what it is doing now, and replicas ratchet up without scaling back in.
# Memory is a floor to provision under 'resources', not a signal to scale on.
# targetMemoryUtilizationPercentage: 80
# behavior: {}

13
helm/litellm/templates/_helpers.tpl
View file

@ -428,3 +428,16 @@ envFrom:
{{- end }}
{{- end }}
{{- end -}}
{{/*
ingress-nginx's admission webhook rejects a dot in an Exact or Prefix path
(strict-validate-path-type) and serves ImplementationSpecific as a plain
prefix location, so a dotted path takes that type there.
*/}}
{{- define "litellm.ingress.pathType" -}}
{{- if and (eq .controller "nginx") (contains "." .path) -}}
ImplementationSpecific
{{- else -}}
{{- .pathType -}}
{{- end -}}
{{- end -}}

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

@ -5,6 +5,10 @@
{{- $gatewayPort := .Values.gateway.service.port -}}
{{- $backendPort := .Values.backend.service.port -}}
{{- $uiPort := .Values.ui.service.port -}}
{{- $controller := .Values.ingress.controller | default "alb" -}}
{{- if not (has $controller (list "alb" "nginx")) }}
{{- fail (printf "ingress.controller: unknown controller %q, expected one of alb, nginx" $controller) }}
{{- end }}
{{/*
Backends addressable from ingress.extraPaths, keyed by the `service` field.
*/}}
@ -27,10 +31,11 @@
/litellm-asset-prefix, so without /*.txt they fall to the backend catch-all
→ 404 → client-side navigation never settles and the login flow spins in an
infinite redirect loop (/ ⇄ /ui/login). ui/nginx.conf already serves *.txt
from the export; the rule only routes the request to it. Needs an ingress
controller whose ImplementationSpecific path is a wildcard pattern
(AWS ALB: `*` = 0+ chars); this chart targets the AWS Load Balancer
Controller.
from the export; the rule only routes the request to it. It needs an
ingress controller whose ImplementationSpecific path is a wildcard pattern
(AWS ALB: `*` = 0+ chars), so it is rendered for ingress.controller=alb
only: ingress-nginx serves ImplementationSpecific as a literal prefix
location, where /*.txt can never match.
*/}}
{{- $uiPaths := list
(dict "path" "/" "pathType" "Exact")
@ -38,8 +43,10 @@
(dict "path" "/litellm-asset-prefix" "pathType" "Prefix")
(dict "path" "/_next" "pathType" "Prefix")
(dict "path" "/ui" "pathType" "Prefix")
(dict "path" "/*.txt" "pathType" "ImplementationSpecific")
-}}
{{- if eq $controller "alb" }}
{{- $uiPaths = append $uiPaths (dict "path" "/*.txt" "pathType" "ImplementationSpecific") }}
{{- end }}
{{/*
Gateway data-plane prefixes — must mirror gateway/routes/allowlist.py.
Versioned paths are listed explicitly to avoid routing management routes
@ -83,12 +90,6 @@
adding to it.
*/}}
{{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}}
{{- range $uiPaths }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" .path .pathType) }}
{{- end }}
{{- range $gatewayPrefixes }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|Prefix" .) }}
{{- end }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
@ -115,8 +116,10 @@ spec:
paths:
# --- UI (Next.js static export) ---
{{- range $uiPaths }}
{{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" .path "pathType" .pathType) }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" .path $pathType) }}
- path: {{ .path }}
pathType: {{ .pathType }}
pathType: {{ $pathType }}
backend:
service:
name: {{ $uiName }}
@ -134,8 +137,10 @@ spec:
port:
number: {{ $gatewayPort }}
{{- range $gatewayPrefixes }}
{{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" . "pathType" "Prefix") }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" . $pathType) }}
- path: {{ . }}
pathType: Prefix
pathType: {{ $pathType }}
backend:
service:
name: {{ $gatewayName }}
@ -147,10 +152,11 @@ spec:
Rendered after every built-in path so an entry can never take
precedence over a default, and before the backend catch-all.
Position only decides the match on controllers that honour manifest
order: the AWS Load Balancer Controller this chart targets sorts
Exact paths first and Prefix paths longest-first, but keeps
order: the AWS Load Balancer Controller (ingress.controller=alb)
sorts Exact paths first and Prefix paths longest-first, but keeps
ImplementationSpecific paths in manifest order, which is what the
/*.txt rule above already depends on.
/*.txt rule above already depends on. ingress-nginx ignores order
and serves the longest matching location.
*/}}
{{- range $idx, $extra := .Values.ingress.extraPaths }}
{{- if not (kindIs "map" $extra) }}
@ -164,10 +170,11 @@ spec:
{{- if not $target }}
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown service %q, expected one of backend, gateway, ui" $idx $extra.path $service) }}
{{- end }}
{{- $pathType := $extra.pathType | default "Prefix" }}
{{- if not (has $pathType (list "Prefix" "Exact" "ImplementationSpecific")) }}
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown pathType %q, expected one of Exact, ImplementationSpecific, Prefix" $idx $extra.path $pathType) }}
{{- $requestedPathType := $extra.pathType | default "Prefix" }}
{{- if not (has $requestedPathType (list "Prefix" "Exact" "ImplementationSpecific")) }}
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown pathType %q, expected one of Exact, ImplementationSpecific, Prefix" $idx $extra.path $requestedPathType) }}
{{- end }}
{{- $pathType := include "litellm.ingress.pathType" (dict "controller" $controller "path" $extra.path "pathType" $requestedPathType) }}
{{- if eq $extra.path "/" }}
{{- fail (printf "ingress.extraPaths[%d]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture" $idx) }}
{{- end }}

205
helm/litellm/tests/ingress_controller_tests.yaml Normal file
View file

@ -0,0 +1,205 @@
suite: test ingress.controller
templates:
- ingress.yaml
values:
- ./values/required.yaml
tests:
- it: keeps the AWS Load Balancer Controller path types by default
set:
ingress.enabled: true
asserts:
- contains:
path: spec.rules[0].http.paths
content:
path: /favicon.ico
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- contains:
path: spec.rules[0].http.paths
content:
path: /eu.assemblyai
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- contains:
path: spec.rules[0].http.paths
content:
path: /*.txt
pathType: ImplementationSpecific
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- it: renders no dotted Exact or Prefix path for ingress-nginx, whose admission webhook rejects them
set:
ingress.enabled: true
ingress.controller: nginx
asserts:
- notMatchRegexRaw:
pattern: 'path: /\S*\.\S*\n\s+pathType: (Exact|Prefix)\n'
- contains:
path: spec.rules[0].http.paths
content:
path: /favicon.ico
pathType: ImplementationSpecific
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- contains:
path: spec.rules[0].http.paths
content:
path: /eu.assemblyai
pathType: ImplementationSpecific
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- it: drops the /*.txt wildcard for ingress-nginx and keeps every other route as is
set:
ingress.enabled: true
ingress.controller: nginx
asserts:
- notContains:
path: spec.rules[0].http.paths
content:
path: /*.txt
any: true
- contains:
path: spec.rules[0].http.paths
content:
path: /ui
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- contains:
path: spec.rules[0].http.paths
content:
path: /test
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- equal:
path: spec.rules[0].http.paths[-1]
value:
path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-backend
port:
number: 4001
- it: rejects an extraPaths entry that repeats a built-in path at the pathType ingress-nginx renders it with
set:
ingress.enabled: true
ingress.controller: nginx
ingress.extraPaths:
- path: /favicon.ico
service: ui
pathType: ImplementationSpecific
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /favicon.ico with pathType ImplementationSpecific is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: rejects a controller it has no path types for
set:
ingress.enabled: true
ingress.controller: traefik
asserts:
- failedTemplate:
errorMessage: 'ingress.controller: unknown controller "traefik", expected one of alb, nginx'
- it: rejects an extraPaths entry that repeats a built-in path once ingress-nginx normalizes its pathType
set:
ingress.enabled: true
ingress.controller: nginx
ingress.extraPaths:
- path: /favicon.ico
service: ui
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /favicon.ico with pathType ImplementationSpecific is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: renders a dotted extraPaths entry as ImplementationSpecific for ingress-nginx
set:
ingress.enabled: true
ingress.controller: nginx
ingress.extraPaths:
- path: /eu.assemblyai.custom
service: gateway
- path: /robots.txt
service: ui
pathType: Exact
asserts:
- notMatchRegexRaw:
pattern: 'path: "?/\S*\.\S*"?\n\s+pathType: (Exact|Prefix)\n'
- contains:
path: spec.rules[0].http.paths
content:
path: /eu.assemblyai.custom
pathType: ImplementationSpecific
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- contains:
path: spec.rules[0].http.paths
content:
path: /robots.txt
pathType: ImplementationSpecific
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- it: keeps the requested pathType of a dotted extraPaths entry for the AWS Load Balancer Controller
set:
ingress.enabled: true
ingress.extraPaths:
- path: /eu.assemblyai.custom
service: gateway
- path: /robots.txt
service: ui
pathType: Exact
asserts:
- contains:
path: spec.rules[0].http.paths
content:
path: /eu.assemblyai.custom
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- contains:
path: spec.rules[0].http.paths
content:
path: /robots.txt
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000

15
helm/litellm/values.yaml
View file

@ -10,6 +10,18 @@ imagePullSecrets: []
ingress:
enabled: false
className: ""
# Which ingress controller serves this Ingress. Controllers disagree on the
# pathTypes they accept, so this picks the pathType of the dotted paths, the
# built-in ones and any dotted extraPaths entry alike:
# alb AWS Load Balancer Controller (default): Exact and Prefix paths plus
# the /*.txt wildcard that routes the UI's RSC payloads.
# nginx ingress-nginx: its admission webhook rejects a dot in an Exact or
# Prefix path (strict-validate-path-type, on by default from v1.12.0
# until v1.12.6 / v1.13.2 allowed dots again), so /favicon.ico and
# /eu.assemblyai render as ImplementationSpecific, which nginx serves
# as a plain prefix location. /*.txt is dropped: nginx has no
# wildcard pathType, so that rule could never match there.
controller: alb
annotations: {}
host: "" # optional; if set, becomes the rule's host
tls: []
@ -26,7 +38,8 @@ ingress:
#
# path required; the HTTP path to route
# service which component serves it: gateway (default), backend, or ui
# pathType Prefix (default), Exact, or ImplementationSpecific
# pathType Prefix (default), Exact, or ImplementationSpecific; a dotted
# path renders as ImplementationSpecific when controller is nginx
#
# The target component only answers paths its own route allowlist keeps, so
# a path here still has to be one that component serves.

91
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
View file

@ -14,18 +14,38 @@ then fails on a Node binary that was never written. Deleting a cache directory
that exists without a Node binary is what turns a killed bootstrap back into a
recoverable one.
Both budgets are overridable so an operator can widen them without a release:
``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and
``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command.
``prisma migrate deploy`` is the other command whose runtime is not a
constant: it grows with the number of pending migrations, so a fresh database
that has to replay every migration this package ships overruns a per-command
budget sized for the short bookkeeping commands, on a laptop as much as on a
slow CI runner. Migrate deploy therefore runs under its own budget.
The Python ``prisma`` wrapper spawns Node, which spawns the Rust schema
engine, so killing only the wrapper on timeout leaves the engine running with
no parent: it keeps mutating the database after the proxy has given up, holds
Prisma's advisory lock so every retry and every later boot queues behind it,
and dies mid-migration once its pipes close, leaving a half-applied ledger row.
Every Prisma command therefore runs in a process group of its own, and a
timeout kills the whole group.
All three budgets are overridable so an operator can widen them without a
release: ``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install,
``LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT`` for ``prisma migrate deploy`` and
``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every other Prisma command. The
per-command budget used to bound migrate deploy as well, so a deployment that
raised it above the deploy default keeps that larger budget for deploy unless
the deploy override says otherwise.
"""
import math
import os
import shutil
import signal
import subprocess
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from typing import IO, Optional, Union
from litellm_proxy_extras._logging import logger
@ -36,10 +56,12 @@ except ImportError:
PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
BOOTSTRAP_ARG = "--version"
@ -88,6 +110,15 @@ def prisma_bootstrap_timeout() -> float:
)
def prisma_migrate_deploy_timeout() -> float:
"""Seconds one ``prisma migrate deploy`` may run for, however many migrations are pending."""
if os.getenv(PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR) is not None:
return _timeout_from_env(
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT
)
return max(DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT, prisma_command_timeout())
def nodeenv_cache_dir() -> Optional[Path]:
"""Where Prisma installs its private Node runtime, or None if unknowable."""
override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR)
@ -143,6 +174,49 @@ def heal_incomplete_nodeenv_cache() -> bool:
return True
def _kill_process_group(process: "subprocess.Popen[str]") -> None:
if os.name == "nt":
process.kill()
return
try:
os.killpg(process.pid, signal.SIGKILL)
except ProcessLookupError:
return
def run_prisma(
argv: Sequence[str],
*,
timeout: float,
env: Mapping[str, str],
stdout: Union[IO[str], int, None] = subprocess.PIPE,
stderr: Optional[int] = subprocess.PIPE,
) -> "subprocess.CompletedProcess[str]":
"""Run one Prisma CLI command in its own process group, bounded by ``timeout``.
Raises ``subprocess.TimeoutExpired`` once the budget is spent, after killing
the command together with every process it spawned, and
``subprocess.CalledProcessError`` on a non-zero exit. Output is captured as
text unless ``stdout``/``stderr`` say otherwise.
"""
with subprocess.Popen(
argv,
env=env,
stdout=stdout,
stderr=stderr,
text=True,
start_new_session=True,
) as process:
try:
out, err = process.communicate(timeout=timeout)
except BaseException:
_kill_process_group(process)
raise
if process.returncode:
raise subprocess.CalledProcessError(process.returncode, process.args, out, err)
return subprocess.CompletedProcess(process.args, process.returncode, out, err)
def ensure_prisma_toolchain(
prisma_command: str, prisma_env: dict[str, str]
) -> ToolchainBootstrap:
@ -155,14 +229,7 @@ def ensure_prisma_toolchain(
timeout = prisma_bootstrap_timeout()
logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout)
try:
subprocess.run(
[prisma_command, BOOTSTRAP_ARG],
timeout=timeout,
check=True,
capture_output=True,
text=True,
env=prisma_env,
)
run_prisma([prisma_command, BOOTSTRAP_ARG], timeout=timeout, env=prisma_env)
except subprocess.TimeoutExpired:
logger.warning(
"Preparing the Prisma CLI toolchain timed out after %ss. Raise %s "

7
litellm-proxy-extras/litellm_proxy_extras/replica_identity.py
View file

@ -16,7 +16,7 @@ import tempfile
from pathlib import Path
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout
from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout, run_prisma
REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"
@ -66,7 +66,7 @@ def apply_replica_identity_full(
with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir:
sql_path = Path(tmp_dir) / "replica_identity_full.sql"
sql_path.write_text(REPLICA_IDENTITY_FULL_SQL)
subprocess.run(
run_prisma(
[
prisma_command,
"db",
@ -77,9 +77,6 @@ def apply_replica_identity_full(
schema_path,
],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
env=prisma_env,
)
except subprocess.CalledProcessError as e:

1
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"]
}

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

@ -6,17 +6,22 @@ import shutil
import subprocess
import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
ensure_prisma_toolchain,
prisma_command_timeout,
prisma_migrate_deploy_timeout,
)
@ -40,6 +45,40 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@dataclass(frozen=True)
class _MigrateAttemptBudget:
"""Retries left, and the recoveries already run.
A recovery that lands something new costs nothing, so a database full of
objects `prisma db push` created works through them one per pass. Anything
that made no progress spends an attempt, so a stuck run still gives up.
"""
attempts_left: int
recoveries: frozenset[str] = frozenset()
@property
def exhausted(self) -> bool:
return self.attempts_left <= 0
@property
def attempt_number(self) -> int:
return MAX_MIGRATE_DEPLOY_ATTEMPTS - self.attempts_left + 1
def spend(self) -> "_MigrateAttemptBudget":
return replace(self, attempts_left=self.attempts_left - 1)
def after_recovery(self, recovery: str) -> "_MigrateAttemptBudget":
if recovery in self.recoveries:
return self.spend()
return replace(self, recoveries=self.recoveries | {recovery})
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
@ -193,7 +232,7 @@ class ProxyExtrasDBManager:
# 1. Generate migration SQL file by comparing empty state to current db state
logger.info("Generating baseline migration...")
migration_file = init_dir / "migration.sql"
subprocess.run(
prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"migrate",
@ -204,14 +243,13 @@ class ProxyExtrasDBManager:
"--script",
],
stdout=open(migration_file, "w"),
check=True,
timeout=prisma_command_timeout(),
env=prisma_env,
)
# 3. Mark the migration as applied since it represents current state
logger.info("Marking baseline migration as applied...")
subprocess.run(
prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"migrate",
@ -219,7 +257,6 @@ class ProxyExtrasDBManager:
"--applied",
"0_init",
],
check=True,
timeout=prisma_command_timeout(),
env=prisma_env,
)
@ -248,7 +285,7 @@ class ProxyExtrasDBManager:
"""Mark a specific migration as rolled back"""
# Set up environment for offline mode if configured
prisma_env = _get_prisma_env()
subprocess.run(
prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"migrate",
@ -257,20 +294,60 @@ class ProxyExtrasDBManager:
migration_name,
],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
env=prisma_env,
)
@staticmethod
def _roll_back_migration_best_effort(migration_name: str) -> None:
"""Mark a migration rolled back, tolerating a concurrent resolver
having already done it."""
try:
ProxyExtrasDBManager._roll_back_migration(migration_name)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
pass
@staticmethod
def _failed_migration_logs(migration_name: str) -> Optional[str]:
"""Return failed migration logs, or None if the ledger is unavailable."""
database_url = os.getenv("DATABASE_URL")
if not database_url:
return None
try:
import psycopg
except ImportError:
return None
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
ledger_table = psycopg.sql.SQL("{}.{}").format(
psycopg.sql.Identifier(
ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
),
psycopg.sql.Identifier("_prisma_migrations"),
)
try:
with psycopg.connect(
cleaned_url, connect_timeout=10, autocommit=True
) as conn:
row = conn.execute(
psycopg.sql.SQL(
"SELECT logs FROM {} "
"WHERE migration_name = %s AND finished_at IS NULL "
"AND rolled_back_at IS NULL"
).format(ledger_table),
(migration_name,),
).fetchone()
except (psycopg.OperationalError, psycopg.DatabaseError):
return None
return (row[0] or "") if row else ""
@staticmethod
def _resolve_specific_migration(migration_name: str):
"""Mark a specific migration as applied"""
prisma_env = _get_prisma_env()
subprocess.run(
prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "resolve", "--applied", migration_name],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
env=prisma_env,
)
@ -354,7 +431,7 @@ class ProxyExtrasDBManager:
try:
logger.info("Generating migration diff between DB and schema.prisma...")
with open(diff_sql_path, "w") as f:
subprocess.run(
prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"migrate",
@ -365,7 +442,6 @@ class ProxyExtrasDBManager:
schema_path,
"--script",
],
check=True,
timeout=prisma_command_timeout(),
stdout=f,
env=_get_prisma_env(),
@ -388,7 +464,7 @@ class ProxyExtrasDBManager:
migration_files = sorted(Path(migrations_dir).glob("*/migration.sql"))
for mig_file in migration_files:
try:
subprocess.run(
prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"db",
@ -399,9 +475,6 @@ class ProxyExtrasDBManager:
schema_path,
],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
env=_get_prisma_env(),
)
logger.info(f"Applied migration: {mig_file.parent.name}")
@ -434,7 +507,7 @@ class ProxyExtrasDBManager:
applied_ok = False
try:
logger.info("Running prisma db execute to apply the migration diff...")
result = subprocess.run(
result = prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"db",
@ -445,9 +518,6 @@ class ProxyExtrasDBManager:
schema_path,
],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
env=_get_prisma_env(),
)
logger.info(f"prisma db execute stdout: {result.stdout}")
@ -476,7 +546,7 @@ class ProxyExtrasDBManager:
for migration_name in migration_names:
try:
logger.info(f"Resolving migration: {migration_name}")
subprocess.run(
prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"migrate",
@ -485,9 +555,6 @@ class ProxyExtrasDBManager:
migration_name,
],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
env=_get_prisma_env(),
)
logger.debug(f"Resolved migration: {migration_name}")
@ -658,7 +725,8 @@ class ProxyExtrasDBManager:
v2 migration resolver (opt-in via --use_v2_migration_resolver).
Runs `prisma migrate deploy` and handles standard recovery paths
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
(P3005 baseline, P3009/P3018 idempotent errors, deadlocks against a
concurrent migrate deploy). Critically, it does
NOT call `_resolve_all_migrations` — the diff-and-force recovery that
caused schema thrashing when two LiteLLM versions contended for the
same DB during rolling deploys.
@ -666,6 +734,9 @@ class ProxyExtrasDBManager:
Ahead-of-HEAD state (DB has migrations newer than this build ships)
is logged as a warning, not a fatal error — users whose DBs got into
weird shapes from the old thrashing should still be able to start.
The retry budget only counts attempts that made no progress: see
_MigrateAttemptBudget.
"""
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
@ -676,11 +747,12 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
subprocess.run(
prisma_toolchain.run_prisma(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),
check=True,
env=_get_prisma_env(),
stdout=None,
stderr=None,
)
return True
except (
@ -698,136 +770,172 @@ class ProxyExtrasDBManager:
original_dir = os.getcwd()
os.chdir(migrations_dir)
deploy_timeout = prisma_migrate_deploy_timeout()
budget = _MigrateAttemptBudget(attempts_left=MAX_MIGRATE_DEPLOY_ATTEMPTS)
try:
for attempt in range(4):
while not budget.exhausted:
try:
result = subprocess.run(
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
timeout=deploy_timeout,
env=_get_prisma_env(),
)
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
return True
except subprocess.TimeoutExpired:
logger.info(
f"prisma migrate deploy attempt {attempt + 1} timed out, retrying"
logger.warning(
"prisma migrate deploy attempt %s timed out after %ss, retrying. "
"Raise %s if this database needs longer to apply its pending migrations.",
budget.attempt_number,
deploy_timeout,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
)
time.sleep(random.randrange(5, 15))
continue
next_budget = budget.spend()
except subprocess.CalledProcessError as e:
stderr = e.stderr or ""
next_budget = ProxyExtrasDBManager._budget_after_deploy_failure(
e, budget, schema_path
)
if "P3005" in stderr and "database schema is not empty" in stderr:
logger.info(
"Schema exists but no migrations ledger — creating baseline"
)
ProxyExtrasDBManager._create_baseline_migration(schema_path)
continue
if "P3009" in stderr:
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} failed idempotently — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
# We're already inside the outer
# `except CalledProcessError` handler —
# re-raising CalledProcessError from here
# would escape as itself, bypassing
# proxy_cli.py's `except RuntimeError`.
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
continue
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
if "P3018" in stderr:
if ProxyExtrasDBManager._is_permission_error(stderr):
raise RuntimeError(
"Database migration failed due to insufficient "
"permissions. Please grant the required privileges "
f"and retry.\n\nPrisma error:\n{stderr}"
) from e
migration_match = re.search(
r"Migration name: (\d+_\S+)", stderr
)
if (
migration_match
and ProxyExtrasDBManager._is_idempotent_error(stderr)
):
name = migration_match.group(1)
logger.info(
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
)
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
continue
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from e
if next_budget.attempts_left < budget.attempts_left:
time.sleep(random.randrange(5, 15))
budget = next_budget # rebind-ok: the loop carries the budget from one migrate deploy pass to the next
raise RuntimeError(
"Database migration failed after 4 attempts (retry loop "
"exhausted by timeouts or repeated idempotent-recovery "
"continues). Check database connectivity, load, and "
"_prisma_migrations ledger state."
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
"attempts that made no progress (timeouts, deadlock retries, or a "
"recovery that had already run once). Check database connectivity, "
"load, and _prisma_migrations ledger state, and raise "
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
)
finally:
os.chdir(original_dir)
@staticmethod
def _budget_after_deploy_failure(
error: subprocess.CalledProcessError,
budget: "_MigrateAttemptBudget",
schema_path: str,
) -> "_MigrateAttemptBudget":
"""Recover from one failed `prisma migrate deploy`, and price the pass.
Returns the budget the next pass runs under, or raises when the failure
is not one this resolver knows how to recover from.
"""
stderr = error.stderr or ""
if "P3005" in stderr and "database schema is not empty" in stderr:
logger.info("Schema exists but no migrations ledger — creating baseline")
if ProxyExtrasDBManager._create_baseline_migration(schema_path):
return budget.after_recovery("baseline")
return budget.spend()
if "P3009" in stderr:
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
name = migration_match.group(1)
logger.info(
f"Migration {name} failed idempotently — marking applied and retrying"
)
ProxyExtrasDBManager._mark_migration_applied(name)
return budget.after_recovery(f"resolved:{name}")
if migration_match:
migration_name = migration_match.group(1)
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
if ledger_logs is not None and (
ledger_logs == "" or _MIGRATION_DEADLOCK_MARKER in ledger_logs
):
logger.info(
"Migration %s failed in a concurrent migrate deploy "
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._roll_back_migration_best_effort(migration_name)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
if "P3018" in stderr:
if ProxyExtrasDBManager._is_permission_error(stderr):
raise RuntimeError(
"Database migration failed due to insufficient "
"permissions. Please grant the required privileges "
f"and retry.\n\nPrisma error:\n{stderr}"
) from error
migration_match = re.search(r"Migration name: (\d+_\S+)", stderr)
if migration_match and ProxyExtrasDBManager._is_idempotent_error(stderr):
name = migration_match.group(1)
logger.info(
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
)
ProxyExtrasDBManager._mark_migration_applied(name)
return budget.after_recovery(f"resolved:{name}")
if migration_match and _MIGRATION_DEADLOCK_MARKER in stderr:
logger.info(
"Migration %s deadlocked against a concurrent "
"migrate deploy, rolling its ledger row back "
"and retrying",
migration_match.group(1),
)
ProxyExtrasDBManager._roll_back_migration_best_effort(
migration_match.group(1)
)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
if _MIGRATION_DEADLOCK_MARKER in stderr:
logger.info(
"prisma migrate deploy attempt %s deadlocked against "
"a concurrent migrate deploy, retrying",
budget.attempt_number,
)
return budget.spend()
if "P1002" in stderr and "advisory lock" in stderr:
logger.info(
"prisma migrate deploy attempt %s timed out waiting for "
"the advisory lock a concurrent migrate deploy holds, retrying",
budget.attempt_number,
)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
@staticmethod
def _mark_migration_applied(name: str) -> None:
"""Roll a failed ledger row back if it is still there, then mark it applied."""
try:
ProxyExtrasDBManager._roll_back_migration(name)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
pass # may already be rolled-back
try:
ProxyExtrasDBManager._resolve_specific_migration(name)
except (
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
) as resolve_err:
# We're called from inside an `except CalledProcessError` handler —
# re-raising CalledProcessError from here would escape as itself,
# bypassing proxy_cli.py's `except RuntimeError`.
raise RuntimeError(
f"Failed to mark migration {name} as applied "
f"after idempotent recovery. Manual "
f"intervention may be required.\n\n"
f"Detail: {resolve_err}"
) from resolve_err
@staticmethod
def apply_replica_identity_full_if_requested() -> bool:
"""
@ -906,12 +1014,9 @@ class ProxyExtrasDBManager:
logger.info("Running prisma migrate deploy")
try:
# Set migrations directory for Prisma
result = subprocess.run(
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
timeout=prisma_migrate_deploy_timeout(),
env=_get_prisma_env(),
)
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
@ -983,7 +1088,7 @@ class ProxyExtrasDBManager:
f"Found failed migration: {failed_migration}, marking as rolled back"
)
# Mark the failed migration as rolled back
subprocess.run(
prisma_toolchain.run_prisma(
[
_get_prisma_command(),
"migrate",
@ -992,9 +1097,6 @@ class ProxyExtrasDBManager:
failed_migration,
],
timeout=prisma_command_timeout(),
check=True,
capture_output=True,
text=True,
env=_get_prisma_env(),
)
logger.info(
@ -1119,14 +1221,20 @@ class ProxyExtrasDBManager:
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
# Use prisma db push with increased timeout
subprocess.run(
prisma_toolchain.run_prisma(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),
check=True,
stdout=None,
stderr=None,
env=_get_prisma_env(),
)
return True
except subprocess.TimeoutExpired:
logger.info(f"Attempt {attempt + 1} timed out")
logger.warning(
"Attempt %s timed out. Raise %s if this database needs longer to apply its schema.",
attempt + 1,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR,
)
time.sleep(random.randrange(5, 15))
except subprocess.CalledProcessError as e:
attempts_left = 3 - attempt

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

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.92"
version = "0.4.93"
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.92"
version = "0.4.93"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

235
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
View file

@ -42,7 +42,7 @@ def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
"Error: P3018\nMigration name: 20250326162113_baseline\n"
"Database error code: 42501\npermission denied for schema public"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="permission"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
@ -60,7 +60,7 @@ def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
'Reason: syntax error at or near "BRKN" LINE 42'
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
@ -124,7 +124,7 @@ def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
def fake_resolve(*args, **kwargs):
resolve_called["n"] += 1
monkeypatch.setattr("subprocess.run", fake_run)
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", fake_run)
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
@ -139,7 +139,7 @@ def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_pat
(tmp_path / "schema.prisma").write_text("// stub")
stderr = "db push error"
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="prisma db push failed"):
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
@ -209,7 +209,7 @@ def test_v2_resolve_specific_migration_failure_raises_runtime_error(
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
"relation already exists"
)
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(
RuntimeError, match="Failed to mark migration .* as applied"
):
@ -228,7 +228,7 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
stdout = "Applied migration.\n"
stderr = ""
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", lambda *a, **kw: FakeResult())
resolve_called = {"n": 0}
monkeypatch.setattr(
@ -240,3 +240,226 @@ def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
assert ok is True
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
_DEADLOCK_P3018_STDERR = (
"Error: P3018\n"
"Migration name: 20260415120000_health_check_latest_per_model_index\n"
"Database error code: 40P01\n"
"deadlock detected"
)
def _stub_v2_env(monkeypatch, tmp_path):
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
monkeypatch.setattr(
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
)
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
(tmp_path / "schema.prisma").write_text("// stub")
monkeypatch.setattr("time.sleep", lambda _: None)
def _succeed_after(failures: int, stderr: str):
calls = {"n": 0}
class _OkResult:
stdout = "Applied migration.\n"
stderr = ""
def _run(*args, **kwargs):
if "deploy" not in args[0]:
return _OkResult()
calls["n"] += 1
if calls["n"] <= failures:
raise subprocess.CalledProcessError(
returncode=1, cmd=args[0], stderr=stderr, output=""
)
return _OkResult()
return _run
def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
"""v2: losing the migrate deploy deadlock race against a concurrent
instance rolls the ledger row back and retries instead of dying."""
_stub_v2_env(monkeypatch, tmp_path)
rolled_back = []
monkeypatch.setattr(
ProxyExtrasDBManager,
"_roll_back_migration",
lambda name: rolled_back.append(name),
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_specific_migration",
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
)
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, _DEADLOCK_P3018_STDERR))
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
assert ok is True
assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
"""v2: a deadlock on every attempt still fails after the retry budget."""
_stub_v2_env(monkeypatch, tmp_path)
monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None)
with patch(
"litellm_proxy_extras.prisma_toolchain.run_prisma",
side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR),
):
with pytest.raises(RuntimeError, match="after 4 attempts"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path):
"""v2: the surviving instance sees the victim's failed ledger row as P3009.
When that row's logs show a deadlock, roll it back and retry."""
_stub_v2_env(monkeypatch, tmp_path)
stderr = (
"Error: P3009\n"
"migrate found failed migrations in the target database\n"
"The `20260415120000_health_check_latest_per_model_index` migration "
"started at 2026-09-01 18:46:13 UTC failed"
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_failed_migration_logs",
lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock",
)
rolled_back = []
monkeypatch.setattr(
ProxyExtrasDBManager,
"_roll_back_migration",
lambda name: rolled_back.append(name),
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_specific_migration",
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
)
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
assert ok is True
assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path):
"""v2: empty failed ledger logs mean a concurrent deploy moved it on."""
_stub_v2_env(monkeypatch, tmp_path)
stderr = (
"Error: P3009\n"
"migrate found failed migrations in the target database\n"
"The `20260415120000_health_check_latest_per_model_index` migration "
"started at 2026-09-01 18:46:13 UTC failed"
)
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "")
rolled_back = []
monkeypatch.setattr(
ProxyExtrasDBManager,
"_roll_back_migration",
lambda name: rolled_back.append(name),
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_resolve_specific_migration",
lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
)
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
assert ok is True
assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]
def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
"""v2: an unreadable ledger cannot establish that P3009 was a deadlock."""
_stub_v2_env(monkeypatch, tmp_path)
stderr = (
"Error: P3009\n"
"migrate found failed migrations in the target database\n"
"The `20260415120000_health_check_latest_per_model_index` migration "
"started at 2026-09-01 18:46:13 UTC failed"
)
monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_roll_back_migration",
lambda name: pytest.fail("an unreadable ledger must not trigger a retry"),
)
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
"""v2: a failed ledger row whose logs show a real SQL error stays fatal."""
_stub_v2_env(monkeypatch, tmp_path)
stderr = (
"Error: P3009\n"
"migrate found failed migrations in the target database\n"
"The `20260101000000_genuinely_broken` migration started at "
"2026-09-01 18:46:13 UTC failed"
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"_failed_migration_logs",
lambda name: 'ERROR: syntax error at or near "BRKN"',
)
with patch("litellm_proxy_extras.prisma_toolchain.run_prisma", side_effect=_fake_migrate_deploy_failure(1, stderr)):
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path):
"""v2: a deadlock reported without a Prisma error code (the advisory-lock
waiter as victim) is retried, not fatal."""
_stub_v2_env(monkeypatch, tmp_path)
monkeypatch.setattr(
"litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, "Database error: deadlock detected")
)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
assert ok is True
_P1002_ADVISORY_LOCK_STDERR = (
"Error: P1002\n\n"
"The database server at `127.0.0.1`:`45743` was reached but timed out.\n\n"
"Context: Timed out trying to acquire a postgres advisory lock "
"(SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms."
)
def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path):
"""v2: the advisory-lock waiter that times out while a peer's retry holds
the lock retries instead of dying."""
_stub_v2_env(monkeypatch, tmp_path)
monkeypatch.setattr(
"litellm_proxy_extras.prisma_toolchain.run_prisma",
_succeed_after(2, _P1002_ADVISORY_LOCK_STDERR),
)
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
assert ok is True
def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_path):
"""v2: a plain P1002 (database unreachable) stays fatal."""
_stub_v2_env(monkeypatch, tmp_path)
stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out."
monkeypatch.setattr("litellm_proxy_extras.prisma_toolchain.run_prisma", _succeed_after(1, stderr))
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)

2
litellm-rust/ADDING_A_PROVIDER.md
View file

@ -26,4 +26,4 @@ variants of it. The test for a good abstraction is that adding the next provider
is a few declarative lines, not a new file of duplicated flow. Only diverge from
the base when behavior is genuinely different, and say so explicitly in the PR.
**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.
**Calling:** hosts invoke the core entrypoint — the Python bridge and the `ai-gateway` route service both call `litellm_core::messages::messages`. Never add a provider handler to `ai-gateway`. Register new modules in `lib.rs` / `mod.rs`, then run the commands under "Checks" in [CLAUDE.md](CLAUDE.md).

10
litellm-rust/AGENTS.md
View file

@ -1,16 +1,18 @@
# AGENTS.md
litellm-rust has exactly THREE crates. A crate is a LAYER, 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-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
| 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 (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
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
@ -28,7 +30,7 @@ core/src/messages/
Handlers never live in `ai-gateway`. `ocr`, `audio_transcription`, and `realtime` are still hosted there from before this rule; they move to `core` as they are touched.
Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
Adding a crate: default to a module. A new crate requires a real trigger: separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.

16
litellm-rust/CLAUDE.md
View file

@ -21,12 +21,14 @@ variants of it. The test for a good abstraction is that adding the next provider
is a few declarative lines, not a new file of duplicated flow. Only diverge from
the base when behavior is genuinely different, and say so explicitly in the PR.
## Crates (exactly three — see AGENTS.md)
## 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. A crate is a **layer**, not
a route — add modules, not crates.
`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
is a layer or shared foundation, not a route; add modules, not crates.
## Core Boundary
@ -173,10 +175,14 @@ for changes under `litellm-rust/`.
```bash
cd litellm-rust
cargo fmt --check
cargo clippy --workspace --all-targets -- -D warnings
cargo clippy -p litellm-core --all-targets --features bedrock-auth -- -D warnings
# the ai-gateway binary + server code is behind the `server` feature
cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings
cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings
cargo clippy -p litellm-ai-gateway --all-targets --all-features -- -D warnings
cargo test --workspace
cargo test -p litellm-core --features bedrock-auth
# the `auth`, `routes`, `state` and `realtime` tests only exist under `server`
cargo test -p litellm-ai-gateway --features server
```
When a Rust path is exposed through Python, add Python parity tests that compare

200
litellm-rust/Cargo.lock generated
View file

@ -919,6 +919,12 @@ version = "0.3.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109"
[[package]]
name = "futures-timer"
version = "3.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968"
[[package]]
name = "futures-util"
version = "0.3.33"
@ -972,6 +978,12 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "glob"
version = "0.3.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b"
[[package]]
name = "h2"
version = "0.3.27"
@ -1380,6 +1392,12 @@ dependencies = [
"wasm-bindgen",
]
[[package]]
name = "lazy_static"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
[[package]]
name = "libc"
version = "0.2.186"
@ -1394,8 +1412,8 @@ dependencies = [
"base64",
"futures-channel",
"futures-util",
"litellm-config",
"litellm-core",
"pyo3",
"reqwest",
"serde",
"serde_json",
@ -1404,6 +1422,17 @@ dependencies = [
"tokio",
"tokio-tungstenite",
"tower",
"tracing",
]
[[package]]
name = "litellm-config"
version = "0.1.0"
dependencies = [
"litellm-core",
"pyo3",
"serde_json",
"thiserror 2.0.19",
]
[[package]]
@ -1416,13 +1445,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]]
@ -1430,14 +1463,28 @@ name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"criterion",
"futures-util",
"litellm-ai-gateway",
"litellm-core",
"litellm-python-interop",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
"serde",
"serde_json",
"tokio",
"tokio-tungstenite",
"tracing",
]
[[package]]
name = "litellm-python-interop"
version = "0.1.0"
dependencies = [
"pyo3",
"pythonize",
"rstest",
"serde",
"serde_json",
]
[[package]]
@ -1476,6 +1523,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"
@ -1627,6 +1684,15 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "proc-macro-crate"
version = "3.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f"
dependencies = [
"toml_edit",
]
[[package]]
name = "proc-macro2"
version = "1.0.107"
@ -1638,9 +1704,9 @@ dependencies = [
[[package]]
name = "pyo3"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"libc",
"once_cell",
@ -1666,18 +1732,18 @@ dependencies = [
[[package]]
name = "pyo3-build-config"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
checksum = "f41027e41b4bd03f6e60f9f417fe24a6341a6bb744edd62b6f709f2a52ea30e9"
dependencies = [
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
checksum = "e591a95526fead067432c3b3a33fc74770b87b1e04e73671090d9c2055a2b327"
dependencies = [
"libc",
"pyo3-build-config",
@ -1685,9 +1751,9 @@ dependencies = [
[[package]]
name = "pyo3-macros"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
checksum = "73225868fc1cd84eef2c3c230ddb91273bf1de46aeb8a4248da76d32a0924a1c"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
@ -1697,9 +1763,9 @@ dependencies = [
[[package]]
name = "pyo3-macros-backend"
version = "0.29.0"
version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
checksum = "571575aa3749fa6216757dd47d2a3e7ef360f329a40f0666a9fbd14889024952"
dependencies = [
"heck",
"proc-macro2",
@ -1899,6 +1965,12 @@ version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
[[package]]
name = "relative-path"
version = "1.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2"
[[package]]
name = "reqwest"
version = "0.12.28"
@ -1919,6 +1991,7 @@ dependencies = [
"hyper-util",
"js-sys",
"log",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"quinn",
@ -1956,6 +2029,35 @@ dependencies = [
"windows-sys 0.52.0",
]
[[package]]
name = "rstest"
version = "0.26.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49"
dependencies = [
"futures-timer",
"futures-util",
"rstest_macros",
]
[[package]]
name = "rstest_macros"
version = "0.26.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0"
dependencies = [
"cfg-if",
"glob",
"proc-macro-crate",
"proc-macro2",
"quote",
"regex",
"relative-path",
"rustc_version",
"syn 2.0.119",
"unicode-ident",
]
[[package]]
name = "rustc-hash"
version = "2.1.3"
@ -2210,6 +2312,15 @@ dependencies = [
"digest 0.11.3",
]
[[package]]
name = "sharded-slab"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6"
dependencies = [
"lazy_static",
]
[[package]]
name = "shlex"
version = "2.0.1"
@ -2348,6 +2459,15 @@ dependencies = [
"syn 3.0.0",
]
[[package]]
name = "thread_local"
version = "1.1.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070"
dependencies = [
"cfg-if",
]
[[package]]
name = "time"
version = "0.3.53"
@ -2488,6 +2608,36 @@ dependencies = [
"tokio",
]
[[package]]
name = "toml_datetime"
version = "1.1.1+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7"
dependencies = [
"serde_core",
]
[[package]]
name = "toml_edit"
version = "0.25.13+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b"
dependencies = [
"indexmap",
"toml_datetime",
"toml_parser",
"winnow",
]
[[package]]
name = "toml_parser"
version = "1.1.3+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56"
dependencies = [
"winnow",
]
[[package]]
name = "tower"
version = "0.5.3"
@ -2566,6 +2716,17 @@ dependencies = [
"once_cell",
]
[[package]]
name = "tracing-subscriber"
version = "0.3.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319"
dependencies = [
"sharded-slab",
"thread_local",
"tracing-core",
]
[[package]]
name = "try-lock"
version = "0.2.5"
@ -2598,6 +2759,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"
@ -2903,6 +3070,15 @@ version = "0.52.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
[[package]]
name = "winnow"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81"
dependencies = [
"memchr",
]
[[package]]
name = "writeable"
version = "0.6.3"

13
litellm-rust/Cargo.toml
View file

@ -1,7 +1,9 @@
[workspace]
members = [
"crates/core",
"crates/config",
"crates/ai-gateway",
"crates/python-interop",
"crates/python-bridge",
]
resolver = "2"
@ -13,16 +15,21 @@ license = "MIT"
repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
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"
pyo3 = "0.29.0"
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"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
serde_json = { version = "1.0", features = ["float_roundtrip"] }
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"

21
litellm-rust/README.md
View file

@ -25,10 +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-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
| 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 (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
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
@ -37,8 +39,10 @@ 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-bridge/ PyO3 bridge for Python LiteLLM.
python-interop/ Domain-neutral PyO3 conversion and GIL primitives.
python-bridge/ PyO3 API adapter for Python LiteLLM.
```
The folder shape follows the Python provider tree:
@ -47,11 +51,6 @@ function per top-level route, mirroring the core entrypoints.
## Checks
Run these before pushing Rust changes. GitHub Actions runs the same checks for
changes under `litellm-rust/`.
```bash
cargo fmt --check
cargo clippy --workspace --all-targets -- -D warnings
cargo test --workspace
```
Run the commands under "Checks" in [CLAUDE.md](CLAUDE.md) before pushing Rust
changes. That list is the single source of truth and matches what GitHub Actions
runs for changes under `litellm-rust/`.

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

@ -45,15 +45,9 @@ 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
25. Run, and keep green:
```bash
cd litellm-rust
cargo fmt --check
cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings
cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings
cargo test --workspace
```
25. Run, and keep green, the commands under "Checks" in `litellm-rust/CLAUDE.md`.
That list is the single source of truth and matches what GitHub Actions runs.

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

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

@ -14,7 +14,9 @@ path = "src/main.rs"
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
@ -30,14 +32,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"

4
litellm-rust/crates/ai-gateway/Dockerfile
View file

@ -36,12 +36,12 @@ FROM chef AS builder
# whenever only gateway source changes.
COPY --from=planner /build/litellm-rust/recipe.json recipe.json
RUN cargo chef cook --locked --release \
-p litellm-ai-gateway --features python-config \
-p litellm-ai-gateway --features server,python-config \
--recipe-path recipe.json
# Now copy the real sources and build the gateway binary. Deps are already cooked
# above, so this step only recompiles the gateway crate.
COPY litellm-rust/ .
RUN cargo build --locked --release -p litellm-ai-gateway --features python-config
RUN cargo build --locked --release -p litellm-ai-gateway --bin litellm-ai-gateway --features server,python-config
# ---- Runtime ----------------------------------------------------------------
# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3

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

@ -6,24 +6,30 @@ dials OpenAI upstream, and splices the two sockets frame-by-frame.
## Crates
`litellm-rust` is exactly three crates (a crate is a **layer**, 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-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — marshals Python objects and calls core entrypoints. |
| 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 (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
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
@ -42,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
@ -81,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
@ -99,7 +106,7 @@ Worker tuning, rarely needed: `LITELLM_LOG_CHANNEL_CAPACITY` (4096),
## Build & run with Docker
The image is built `--features python-config` and installs litellm **from this
The image is built `--features server,python-config` and installs litellm **from this
repo's source** (the config reader is newer than any PyPI release), so the build
**context is the repo root**:
@ -134,10 +141,10 @@ docker run --rm -p 4001:4001 \
```bash
# config.yaml mode — needs litellm importable in the active python env
LITELLM_CONFIG_PATH=./crates/ai-gateway/config.yaml \
cargo run --release -p litellm-ai-gateway --features python-config
cargo run --release -p litellm-ai-gateway --features server,python-config
# env stand-in mode — no python, no config
cargo run --release -p litellm-ai-gateway
cargo run --release -p litellm-ai-gateway --features server
```
## Deploy on Render

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

48
litellm-rust/crates/ai-gateway/src/audio_transcription/common_utils.rs
View file

@ -1,48 +0,0 @@
use std::collections::BTreeMap;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::error::CoreError;
use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use serde_json::{Map, Value};
pub(super) fn audio_transcription_provider_config(
provider: &str,
) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
match provider {
"bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
pub(super) fn string_headers(
headers: Option<Map<String, Value>>,
) -> CoreResult<BTreeMap<String, String>> {
headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
"audio transcription extra_headers.{key} must be a string"
))
})
})
.collect()
}
pub(super) fn has_header(headers: &BTreeMap<String, String>, name: &str) -> bool {
headers.keys().any(|key| key.eq_ignore_ascii_case(name))
}
pub(super) fn truncate_error_body(body: &str) -> String {
let truncated: String = body.chars().take(256).collect();
if truncated.chars().count() == body.chars().count() {
truncated
} else {
format!("{truncated}... (truncated)")
}
}

89
litellm-rust/crates/ai-gateway/src/audio_transcription/handler.rs
View file

@ -1,89 +0,0 @@
use std::time::SystemTime;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::error::CoreError;
use litellm_core::providers::bedrock::audio_transcription::aws_auth_config;
use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
use serde_json::Value;
use super::common_utils::truncate_error_body;
use super::types::ProviderAudioTranscriptionRequest;
use crate::client::http_client;
pub(crate) async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
) -> CoreResult<Value> {
let body = serde_json::to_vec(&request.body).map_err(|error| {
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
})?;
let mut request_builder = http_client().post(&request.url).body(body.clone());
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request_builder
.send()
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text).map_err(|error| {
CoreError::InvalidResponse(format!("invalid audio response JSON: {error}"))
})?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.into_json())
}
pub(crate) async fn sign_request(
request: &ProviderAudioTranscriptionRequest,
optional_params: &serde_json::Map<String, Value>,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
let env_lookup = environment_lookup;
let auth = request
.config
.auth_strategy(&request.model, optional_params, &env_lookup)?;
let body = serde_json::to_vec(&request.body).map_err(|error| {
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
})?;
let mut headers = super::common_utils::string_headers(None)?;
headers.insert("Content-Type".to_string(), "application/json".to_string());
headers.extend(request.upstream_headers.iter().cloned());
match auth {
AudioTranscriptionAuth::Bearer => {}
AudioTranscriptionAuth::AwsSigV4 { region, .. } => {
let credentials =
resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup)
.await?;
headers.extend(sign_bedrock_post(
&request.url,
&body,
&headers,
&region,
&credentials,
SystemTime::now(),
)?);
}
}
Ok(ProviderAudioTranscriptionRequest {
upstream_headers: headers.into_iter().collect(),
..request.clone()
})
}
pub(super) fn environment_lookup(key: &str) -> Option<String> {
std::env::var(key).ok()
}

138
litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs
View file

@ -1,15 +1,14 @@
use litellm_core::audio_transcription::{
AudioTranscriptionRequest as CoreAudioTranscriptionRequest, ProviderAudioTranscriptionRequest,
prepare_audio_transcription_provider_call,
};
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::Error;
use serde_json::{Map, Value, json};
use std::future::Future;
use std::pin::Pin;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::CoreError;
use serde_json::{Map, Value, json};
use super::common_utils::{audio_transcription_provider_config, has_header, string_headers};
use super::handler::sign_request;
use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use super::types::PreparedAudioTranscriptionRequest;
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
@ -26,7 +25,7 @@ pub(crate) struct AudioTranscriptionLifecycleHooks {
request_metadata: RequestMetadata,
}
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
type AudioLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl AudioTranscriptionLifecycleHooks {
@ -45,7 +44,7 @@ impl AudioTranscriptionLifecycleHooks {
async fn run_pre_call_guardrails(
&self,
request: PreparedAudioTranscriptionRequest,
) -> CoreResult<PreparedAudioTranscriptionRequest> {
) -> Result<PreparedAudioTranscriptionRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
@ -63,17 +62,17 @@ impl AudioTranscriptionLifecycleHooks {
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(CoreError::InvalidRequest(
return Err(Error::InvalidRequest(
"audio transcription pre_call guardrail must return an object".to_string(),
));
};
let audio = data.remove("audio").ok_or_else(|| {
CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string())
Error::InvalidRequest("audio transcription guardrail removed audio".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(value)) => value,
Some(_) => {
return Err(CoreError::InvalidRequest(
return Err(Error::InvalidRequest(
"audio transcription optional_params must be an object".to_string(),
));
}
@ -89,53 +88,36 @@ impl AudioTranscriptionLifecycleHooks {
async fn prepare_provider_request(
&self,
request: PreparedAudioTranscriptionRequest,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
let config = audio_transcription_provider_config(&request.custom_llm_provider)
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
let env_lookup = super::handler::environment_lookup;
let headers = string_headers(request.extra_headers)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let body = config.transform_transcription_request(
&request.model,
request.audio,
filtered_params,
)?;
let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?;
let mut upstream_headers = headers.into_iter().collect::<Vec<_>>();
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(
&upstream_headers
.iter()
.cloned()
.collect::<std::collections::BTreeMap<_, _>>(),
"authorization",
)
&& let Some(api_key) = request.api_key.as_deref()
{
upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
}
let provider_request = ProviderAudioTranscriptionRequest {
model: request.model,
config,
url,
body: body.body,
upstream_headers,
timeout: request.timeout,
};
let provider_request = self.run_during_call_guardrails(provider_request).await?;
sign_request(&provider_request, &request.optional_params).await
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let PreparedAudioTranscriptionRequest {
model,
custom_llm_provider,
audio,
api_key,
api_base,
extra_headers,
optional_params,
timeout,
..
} = request;
let provider_request =
prepare_audio_transcription_provider_call(CoreAudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: Some(&custom_llm_provider),
extra_headers,
optional_params,
timeout,
})?;
self.run_during_call_guardrails(provider_request).await
}
async fn run_during_call_guardrails(
&self,
request: ProviderAudioTranscriptionRequest,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
) -> Result<ProviderAudioTranscriptionRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
@ -144,23 +126,23 @@ impl AudioTranscriptionLifecycleHooks {
.run_during_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": "bedrock",
"url": request.url,
"body": request.body,
"model": request.model(),
"custom_llm_provider": request.custom_llm_provider(),
"url": request.url(),
"body": request.body(),
})),
)
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(CoreError::InvalidRequest(
return Err(Error::InvalidRequest(
"audio transcription during_call guardrail must return an object".to_string(),
));
};
let body = data.remove("body").ok_or_else(|| {
CoreError::InvalidRequest("audio transcription guardrail removed body".to_string())
Error::InvalidRequest("audio transcription guardrail removed body".to_string())
})?;
Ok(ProviderAudioTranscriptionRequest { body, ..request })
Ok(request.with_body(body))
}
fn logging_payload(
@ -241,7 +223,7 @@ impl CallLifecycleHooks<PreparedAudioTranscriptionRequest, ProviderAudioTranscri
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a CoreError,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -281,22 +263,22 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
}
}
fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &CoreError) -> &'static str {
fn core_error_kind(error: &Error) -> &'static str {
match error {
CoreError::Auth(_) => "AuthError",
CoreError::InvalidProvider(_) => "InvalidProvider",
CoreError::InvalidRequest(_) => "InvalidRequest",
CoreError::InvalidType { .. } => "InvalidType",
CoreError::MissingField(_) => "MissingField",
CoreError::Http { .. } => "HttpError",
CoreError::InvalidResponse(_) => "InvalidResponse",
CoreError::Network(_) => "NetworkError",
CoreError::Connect(_) => "ConnectError",
CoreError::Routing(_) => "RoutingError",
CoreError::Unsupported(_) => "UnsupportedRequest",
Error::Auth(_) => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",
Error::Connect(_) => "ConnectError",
Error::Routing(_) => "RoutingError",
Error::Unsupported(_) => "UnsupportedRequest",
}
}

8
litellm-rust/crates/ai-gateway/src/audio_transcription/mod.rs
View file

@ -1,19 +1,17 @@
use litellm_core::CoreResult;
use litellm_core::Error;
use litellm_core::audio_transcription::execute_audio_transcription_provider_call;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
mod common_utils;
mod handler;
mod hooks;
mod prepare;
mod types;
pub use types::AudioTranscriptionRequest;
use handler::execute_audio_transcription_provider_call;
use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult<Value> {
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
let PreparedAudioTranscriptionCall { request, hooks } =
prepare_audio_transcription_call(request);
CallLifecycle::default()

11
litellm-rust/crates/ai-gateway/src/audio_transcription/types.rs
View file

@ -1,7 +1,6 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use serde_json::{Map, Value};
@ -46,13 +45,3 @@ impl CallLifecycleRequest for PreparedAudioTranscriptionRequest {
)
}
}
#[derive(Clone)]
pub(crate) struct ProviderAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) config: &'static dyn AudioTranscriptionProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

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

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

@ -15,8 +15,7 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::error::Error;
use litellm_core::realtime::transformation::RealtimeProviderConfig;
use litellm_core::realtime::types::RealtimeEvent;
use tokio::net::TcpStream;
@ -48,7 +47,7 @@ pub(crate) type UpstreamRx = SplitStream<UpstreamWs>;
/// Resolve the OpenAI API key from the explicit param or the environment.
///
/// Blank/whitespace values are treated as absent (guard at resolution time).
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|key| !key.is_empty())
@ -58,7 +57,7 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
.ok()
.filter(|key| !key.trim().is_empty())
})
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
}
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.
@ -70,24 +69,24 @@ pub(crate) async fn dial_upstream(
model: &str,
api_key: &str,
api_base: Option<&str>,
) -> CoreResult<UpstreamWs> {
) -> Result<UpstreamWs, Error> {
let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model);
let mut request = url
.as_str()
.into_client_request()
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
// GA realtime: only Authorization. The legacy OpenAI-Beta header triggers
// beta_api_shape_disabled, so we do not send it.
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {api_key}"))
.map_err(|err| CoreError::Auth(err.to_string()))?,
.map_err(|err| Error::Auth(err.to_string()))?,
);
let (upstream, _response) = connect_async(request)
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
Ok(upstream)
}
@ -96,22 +95,22 @@ pub(crate) async fn dial_upstream(
/// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an
/// error on a non-text frame, a closed socket, or undecodable JSON so the pool can
/// discard a misbehaving socket rather than warm it.
pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<RealtimeEvent> {
pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> Result<RealtimeEvent, Error> {
loop {
let message = upstream_rx
.next()
.await
.ok_or_else(|| CoreError::Network("upstream closed before first event".to_string()))?
.map_err(|err| CoreError::Network(err.to_string()))?;
.ok_or_else(|| Error::Network("upstream closed before first event".to_string()))?
.map_err(|err| Error::Network(err.to_string()))?;
match message {
Message::Text(text) => {
return serde_json::from_str(&text)
.map_err(|err| CoreError::InvalidResponse(err.to_string()));
.map_err(|err| Error::InvalidResponse(err.to_string()));
}
// Ignore protocol frames (ping/pong) while waiting for the first event.
Message::Ping(_) | Message::Pong(_) => continue,
Message::Close(_) => {
return Err(CoreError::Network(
return Err(Error::Network(
"upstream closed before first event".to_string(),
));
}
@ -139,7 +138,7 @@ pub(crate) async fn splice<In, Out>(
mut observe: impl FnMut(&RealtimeEvent) + Send,
mut client_in: In,
mut client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,
@ -154,7 +153,7 @@ where
client_out
.send(outbound)
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
}
}
@ -175,26 +174,26 @@ where
// inflate its own spend log. Logging observes upstream events only.
for outbound in config.transform_realtime_request(&event, model)?.events {
let payload = serde_json::to_string(&outbound)
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
.map_err(|err| Error::InvalidResponse(err.to_string()))?;
upstream_tx
.send(Message::Text(payload))
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
}
}
// upstream -> client
upstream_message = upstream_rx.next() => {
let Some(message) = upstream_message else { break }; // upstream closed
match message.map_err(|err| CoreError::Network(err.to_string()))? {
match message.map_err(|err| Error::Network(err.to_string()))? {
Message::Text(text) => {
let event: RealtimeEvent = serde_json::from_str(&text)
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
.map_err(|err| Error::InvalidResponse(err.to_string()))?;
observe(&event);
for outbound in config.transform_realtime_response(&event, model)?.events {
client_out
.send(outbound)
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
}
}
Message::Close(_) => break,
@ -225,7 +224,7 @@ pub async fn realtime<In, Out>(
observe: impl FnMut(&RealtimeEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,
@ -258,7 +257,7 @@ pub async fn realtime_warm<In, Out>(
observe: impl FnMut(&RealtimeEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,

4
litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs
View file

@ -28,7 +28,7 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use futures_util::StreamExt;
use litellm_core::CoreResult;
use litellm_core::Error;
use litellm_core::realtime::types::RealtimeEvent;
use crate::io::realtime::{
@ -438,7 +438,7 @@ impl RealtimePool {
///
/// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends
/// unprompted is `session.created`; we buffer exactly that and read nothing more.
async fn warm_one(key: &UpstreamKey) -> CoreResult<WarmConnection> {
async fn warm_one(key: &UpstreamKey) -> Result<WarmConnection, Error> {
let upstream: UpstreamWs =
dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?;
let (tx, mut rx) = upstream.split();

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

@ -4,10 +4,10 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::Error;
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
use litellm_core::{CoreError, CoreResult};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;
@ -37,51 +37,49 @@ impl ResponsesWebSocketConnection {
url: &str,
headers: &HashMap<String, String>,
timeout: Option<Duration>,
) -> CoreResult<Self> {
) -> Result<Self, Error> {
let mut request = url
.into_client_request()
.map_err(|error| CoreError::Network(error.to_string()))?;
.map_err(|error| Error::Network(error.to_string()))?;
for (name, value) in headers {
let header_name = name
.parse::<HeaderName>()
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
let header_value = HeaderValue::from_str(value)
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
CoreError::Network("Responses WebSocket connection timed out".to_string())
Error::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match error {
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => CoreError::Network(other.to_string()),
other => Error::Network(other.to_string()),
})?;
Ok(Self {
socket: Arc::new(Mutex::new(Some(socket))),
})
}
pub async fn send_text(&self, text: String) -> CoreResult<()> {
pub async fn send_text(&self, text: String) -> Result<(), Error> {
let mut socket = self.socket.lock().await;
let Some(socket) = socket.as_mut() else {
return Err(CoreError::Network(
"Responses WebSocket is closed".to_string(),
));
return Err(Error::Network("Responses WebSocket is closed".to_string()));
};
socket
.send(Message::Text(text))
.await
.map_err(|error| CoreError::Network(error.to_string()))
.map_err(|error| Error::Network(error.to_string()))
}
pub async fn recv_text(&self) -> CoreResult<Option<String>> {
pub async fn recv_text(&self) -> Result<Option<String>, Error> {
let mut socket_guard = self.socket.lock().await;
let Some(socket) = socket_guard.as_mut() else {
return Ok(None);
@ -90,27 +88,27 @@ impl ResponsesWebSocketConnection {
Some(Ok(Message::Text(text))) => Ok(Some(text)),
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
.map(Some)
.map_err(|error| CoreError::InvalidResponse(error.to_string())),
.map_err(|error| Error::InvalidResponse(error.to_string())),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(CoreError::Network(error.to_string())),
Some(Err(error)) => Err(Error::Network(error.to_string())),
}
}
pub async fn close(&self) -> CoreResult<()> {
pub async fn close(&self) -> Result<(), Error> {
let mut socket = self.socket.lock().await;
if let Some(socket) = socket.as_mut() {
socket
.close(None)
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
.map_err(|error| Error::Network(error.to_string()))?;
}
*socket = None;
Ok(())
}
}
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> Result<String, Error> {
api_key
.map(str::trim)
.filter(|value| !value.is_empty())
@ -120,38 +118,38 @@ pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
.ok()
.filter(|value| !value.trim().is_empty())
})
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
.ok_or_else(|| Error::Auth(MISSING_KEY_MESSAGE.to_string()))
}
async fn dial_upstream(
model: &str,
api_key: &str,
api_base: Option<&str>,
) -> CoreResult<ResponsesUpstreamWs> {
) -> Result<ResponsesUpstreamWs, Error> {
let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model);
let mut request = url
.as_str()
.into_client_request()
.map_err(|error| CoreError::Network(error.to_string()))?;
.map_err(|error| Error::Network(error.to_string()))?;
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {api_key}"))
.map_err(|error| CoreError::Auth(error.to_string()))?,
.map_err(|error| Error::Auth(error.to_string()))?,
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_async(request),
)
.await
.map_err(|_| CoreError::Network("Responses WebSocket connection timed out".to_string()))?;
.map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match error {
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => CoreError::Network(other.to_string()),
other => Error::Network(other.to_string()),
})
}
@ -166,7 +164,7 @@ impl ResponsesWebSocketStreaming {
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
@ -193,7 +191,7 @@ pub(crate) async fn splice<In, Out>(
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
mut client_in: In,
mut client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
@ -210,18 +208,18 @@ where
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
upstream_tx.send(Message::Text(payload))
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
.map_err(|error| Error::Network(error.to_string()))?;
}
}
message = upstream_rx.next() => {
let Some(message) = message else { break };
match message.map_err(|error| CoreError::Network(error.to_string()))? {
match message.map_err(|error| Error::Network(error.to_string()))? {
Message::Text(text) => {
let event = serde_json::from_str::<ResponsesWsEvent>(&text)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
observe(&event);
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_response(&event, model)?
@ -229,7 +227,7 @@ where
{
client_out.send(outbound)
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
.map_err(|error| Error::Network(error.to_string()))?;
}
}
Message::Close(_) => break,
@ -252,7 +250,7 @@ pub async fn async_responses_websocket<In, Out>(
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
@ -267,11 +265,11 @@ where
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
.map_err(|error| Error::InvalidResponse(error.to_string()))?;
upstream_tx
.send(Message::Text(payload))
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
.map_err(|error| Error::Network(error.to_string()))?;
}
}
ResponsesWebSocketStreaming::bidirectional_forward(
@ -296,7 +294,7 @@ pub async fn responses_ws<In, Out>(
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
@ -514,7 +512,7 @@ mod tests {
)
.await
.expect_err("status error");
assert!(matches!(error, CoreError::Http { status: 401, .. }));
assert!(matches!(error, Error::Http { status: 401, .. }));
server.await.expect("server task");
}
@ -543,7 +541,7 @@ mod tests {
)
.await
.expect_err("status error");
assert!(matches!(error, CoreError::Http { status: 500, .. }));
assert!(matches!(error, Error::Http { status: 500, .. }));
server.await.expect("server task");
}
}

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");

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

@ -3,8 +3,7 @@ use std::time::{Duration, Instant};
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::error::Error;
use litellm_core::ocr::transformation::OcrProviderConfig;
use reqwest::Url;
use serde_json::{Map, Value};
@ -13,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,
@ -33,12 +33,14 @@ pub(super) fn truncate_error_body(body: &str) -> String {
format!("{truncated}... (truncated)")
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) fn ocr_provider_config(
provider: &str,
model: &str,
) -> 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)
}
@ -56,7 +58,7 @@ fn is_azure_document_intelligence_model(model: &str) -> bool {
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<Vec<(String, String)>> {
) -> Result<Vec<(String, String)>, Error> {
extra_headers
.unwrap_or_default()
.into_iter()
@ -65,7 +67,7 @@ pub(super) fn string_headers(
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
Error::InvalidRequest(format!(
"OCR extra_headers.{key} must be a string, got {}",
litellm_core::error::json_type_name(&value)
))
@ -74,13 +76,7 @@ pub(super) fn string_headers(
.collect()
}
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
headers
.iter()
.any(|(key, _)| key.eq_ignore_ascii_case(name))
}
fn document_url_field(document: &Value) -> CoreResult<Option<(&str, &str)>> {
fn document_url_field(document: &Value) -> Result<Option<(&str, &str)>, Error> {
let Some(object) = document.as_object() else {
return Ok(None);
};
@ -138,13 +134,13 @@ fn is_blocked_ip(ip: IpAddr) -> bool {
}
}
fn blocked_url_error(url: &Url) -> CoreError {
CoreError::InvalidRequest(format!(
fn blocked_url_error(url: &Url) -> Error {
Error::InvalidRequest(format!(
"OCR document URL rejected by SSRF protection: {url}"
))
}
async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
async fn validate_safe_fetch_url(url: &Url) -> Result<(), Error> {
if !matches!(url.scheme(), "http" | "https") {
return Err(blocked_url_error(url));
}
@ -162,7 +158,7 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
.ok_or_else(|| blocked_url_error(url))?;
let addresses = tokio::net::lookup_host((host, port))
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
let mut saw_address = false;
for address in addresses {
saw_address = true;
@ -176,25 +172,25 @@ async fn validate_safe_fetch_url(url: &Url) -> CoreResult<()> {
Ok(())
}
fn redirect_location(response: &reqwest::Response, url: &Url) -> CoreResult<Url> {
fn redirect_location(response: &reqwest::Response, url: &Url) -> Result<Url, Error> {
let location = response
.headers()
.get(reqwest::header::LOCATION)
.and_then(|value| value.to_str().ok())
.ok_or_else(|| {
CoreError::InvalidResponse("OCR document redirect missing Location header".to_string())
Error::InvalidResponse("OCR document redirect missing Location header".to_string())
})?;
url.join(location)
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR document redirect: {err}")))
.map_err(|err| Error::InvalidResponse(format!("invalid OCR document redirect: {err}")))
}
async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)> {
async fn safe_get_document_url(url: &str) -> Result<(Url, reqwest::Response), Error> {
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
let mut current_url = Url::parse(url)
.map_err(|err| CoreError::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
.map_err(|err| Error::InvalidRequest(format!("invalid OCR document URL: {err}")))?;
for _ in 0..MAX_SAFE_FETCH_REDIRECTS {
validate_safe_fetch_url(&current_url).await?;
@ -202,28 +198,28 @@ async fn safe_get_document_url(url: &str) -> CoreResult<(Url, reqwest::Response)
.get(current_url.clone())
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
if !response.status().is_redirection() {
return Ok((current_url, response));
}
current_url = redirect_location(&response, &current_url)?;
}
Err(CoreError::InvalidRequest(
Err(Error::InvalidRequest(
"Too many redirects while fetching OCR document URL".to_string(),
))
}
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> CoreResult<()> {
fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Result<(), Error> {
if max_bytes == 0 {
return Err(CoreError::InvalidRequest(format!(
return Err(Error::InvalidRequest(format!(
"OCR document URL download is disabled (MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0). url={url}"
)));
}
if content_length > max_bytes {
let size_mb = content_length as f64 / (1024.0 * 1024.0);
let max_size_mb = max_bytes as f64 / (1024.0 * 1024.0);
return Err(CoreError::InvalidRequest(format!(
return Err(Error::InvalidRequest(format!(
"OCR document size ({size_mb:.2}MB) exceeds maximum allowed size ({max_size_mb:.2}MB). url={url}"
)));
}
@ -233,7 +229,7 @@ fn enforce_download_size(content_length: u64, max_bytes: u64, url: &Url) -> Core
async fn read_response_with_limit(
mut response: reqwest::Response,
url: &Url,
) -> CoreResult<Vec<u8>> {
) -> Result<Vec<u8>, Error> {
let max_bytes = max_document_download_bytes();
if let Some(content_length) = response.content_length() {
enforce_download_size(content_length, max_bytes, url)?;
@ -246,7 +242,7 @@ async fn read_response_with_limit(
while let Some(chunk) = response
.chunk()
.await
.map_err(|err| CoreError::Network(err.to_string()))?
.map_err(|err| Error::Network(err.to_string()))?
{
bytes_downloaded += chunk.len() as u64;
enforce_download_size(bytes_downloaded, max_bytes, url)?;
@ -255,7 +251,7 @@ async fn read_response_with_limit(
Ok(bytes)
}
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreResult<Value> {
pub(super) async fn convert_document_url_to_data_uri(document: Value) -> Result<Value, Error> {
let Some((field, url)) = document_url_field(&document)? else {
return Ok(document);
};
@ -267,7 +263,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(CoreError::Http {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&body),
});
@ -290,7 +286,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes
let mut transformed = document
.as_object()
.cloned()
.ok_or_else(|| CoreError::InvalidRequest("OCR document must be an object".to_string()))?;
.ok_or_else(|| Error::InvalidRequest("OCR document must be an object".to_string()))?;
transformed.insert(field.to_string(), Value::String(data_uri));
Ok(Value::Object(transformed))
}
@ -316,11 +312,11 @@ fn retry_after_secs(response: &reqwest::Response) -> u64 {
.unwrap_or(2)
}
fn operation_status(response_json: &Value) -> CoreResult<&str> {
fn operation_status(response_json: &Value) -> Result<&str, Error> {
let status = response_json
.get("status")
.and_then(Value::as_str)
.ok_or(CoreError::MissingField("status"))?;
.ok_or(Error::MissingField("status"))?;
match status {
"succeeded" => Ok("succeeded"),
"running" | "notStarted" => Ok("running"),
@ -330,24 +326,25 @@ fn operation_status(response_json: &Value) -> CoreResult<&str> {
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.unwrap_or("Unknown error");
Err(CoreError::InvalidResponse(format!(
Err(Error::InvalidResponse(format!(
"Azure Document Intelligence analysis failed: {message}"
)))
}
other => Err(CoreError::InvalidResponse(format!(
other => Err(Error::InvalidResponse(format!(
"Unknown operation status: {other}"
))),
}
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn poll_document_intelligence(
operation_url: &str,
original_url: &str,
headers: &[(String, String)],
timeout: Option<Duration>,
) -> CoreResult<Value> {
) -> Result<Value, Error> {
if !same_origin(operation_url, original_url) {
return Err(CoreError::InvalidResponse(
return Err(Error::InvalidResponse(
"Azure Document Intelligence: rejected cross-origin polling URL".to_string(),
));
}
@ -358,7 +355,7 @@ pub(super) async fn poll_document_intelligence(
));
loop {
if start.elapsed() > timeout {
return Err(CoreError::Network(format!(
return Err(Error::Network(format!(
"Azure Document Intelligence operation polling timed out after {} seconds",
timeout.as_secs()
)));
@ -373,21 +370,21 @@ pub(super) async fn poll_document_intelligence(
let response = request_builder
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
let retry_after = retry_after_secs(&response);
let status = response.status();
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text).map_err(|err| {
CoreError::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
Error::InvalidResponse(format!("invalid Azure DI poll response JSON: {err}"))
})?;
if operation_status(&response_json)? == "succeeded" {
return Ok(response_json);
@ -398,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()));
@ -426,7 +425,7 @@ mod tests {
assert!(matches!(
error,
CoreError::InvalidRequest(message)
Error::InvalidRequest(message)
if message.contains("SSRF protection")
));
}
@ -444,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()
)
);
}
}

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

@ -1,13 +1,19 @@
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::error::Error;
use litellm_core::http_utils::http_request;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::Value;
use super::common_utils::{poll_document_intelligence, truncate_error_body};
use super::types::ProviderOcrRequest;
use super::hooks::OcrLifecycleHooks;
use super::types::PreparedOcrRequest;
use crate::client::http_client;
pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) async fn execute_ocr_provider_call(
request: PreparedOcrRequest,
hooks: &OcrLifecycleHooks,
) -> Result<Value, Error> {
let request = hooks.prepare_provider_request(request).await?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
@ -16,10 +22,9 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co
request_builder = request_builder.timeout(duration);
}
let response = request_builder
.send()
let response = http_request(request_builder)
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
let status = response.status();
if request.config.response_handling() == OcrResponseHandling::AzureDocumentIntelligencePoll
@ -31,7 +36,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co
.and_then(|value| value.to_str().ok())
.map(str::to_string)
.ok_or_else(|| {
CoreError::InvalidResponse(
Error::InvalidResponse(
"Azure Document Intelligence returned 202 but no Operation-Location header found"
.to_string(),
)
@ -45,27 +50,35 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co
.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());
}
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let response_json: Value = serde_json::from_str(&text)
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
.map_err(|err| Error::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
Ok(request
.config
.transform_ocr_response(&request.model, response_json)?
.transform_ocr_response_with_params(
&request.model,
response_json,
&request.optional_params,
)?
.into_json())
}

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

@ -1,16 +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 litellm_core::CoreResult;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrAuthStrategy;
use serde_json::{Map, Value, json};
use super::common_utils::{
convert_document_url_to_data_uri, has_header, ocr_provider_config, 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,
};
@ -27,7 +26,7 @@ pub(crate) struct OcrLifecycleHooks {
request_metadata: RequestMetadata,
}
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
type OcrFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
type OcrLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl OcrLifecycleHooks {
@ -46,7 +45,7 @@ impl OcrLifecycleHooks {
async fn run_pre_call_guardrails(
&self,
request: PreparedOcrRequest,
) -> CoreResult<PreparedOcrRequest> {
) -> Result<PreparedOcrRequest, Error> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
@ -64,6 +63,10 @@ impl OcrLifecycleHooks {
.await
.map_err(guardrail_error_to_core_error)?;
let (document, optional_params) = parse_ocr_pre_call_guardrail_request(guardrail_request)?;
let optional_params = match &request.config {
Ok(config) => config.map_ocr_params(&optional_params),
Err(_) => optional_params,
};
Ok(PreparedOcrRequest {
document,
optional_params,
@ -71,44 +74,58 @@ impl OcrLifecycleHooks {
})
}
async fn prepare_provider_request(
pub(crate) async fn prepare_provider_request(
&self,
request: PreparedOcrRequest,
) -> CoreResult<ProviderOcrRequest> {
let config = ocr_provider_config(&request.custom_llm_provider, &request.model)
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
) -> Result<ProviderOcrRequest, Error> {
let config = request.config?;
let env_lookup = |key: &str| std::env::var(key).ok();
let headers = string_headers(request.extra_headers)?;
let auth_strategy = config.auth_strategy();
let api_key = (!has_header(&headers, auth_strategy.header_name()))
.then(|| config.resolve_api_key(request.api_key.as_deref(), &env_lookup))
.transpose()?;
let upstream_headers = config.validate_environment(
string_headers(request.extra_headers)?,
request.api_key.as_deref(),
&env_lookup,
)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_ocr_params(&request.optional_params);
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, filtered_params)?
.transform_ocr_request(&request.model, document, optional_params.clone())?
.data;
let upstream_headers = upstream_headers(&headers, auth_strategy, api_key.as_deref());
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,
})
@ -120,7 +137,7 @@ impl OcrLifecycleHooks {
custom_llm_provider: &str,
url: &str,
body: Value,
) -> CoreResult<Value> {
) -> Result<Value, Error> {
if self.guardrail_runner.is_empty() {
return Ok(body);
}
@ -169,9 +186,66 @@ impl OcrLifecycleHooks {
}
}
impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for 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, ProviderOcrRequest>;
type DuringCallFuture<'a> = OcrFuture<'a, PreparedOcrRequest>;
type SuccessFuture<'a> = OcrLogFuture<'a>;
type FailureFuture<'a> = OcrLogFuture<'a>;
@ -188,7 +262,7 @@ impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLi
_context: &'a CallLifecycleContext,
request: PreparedOcrRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { self.prepare_provider_request(request).await })
Box::pin(async move { Ok(request) })
}
fn async_log_success_event<'a>(
@ -217,7 +291,7 @@ impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLi
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a CoreError,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -249,21 +323,6 @@ impl CallLifecycleHooks<PreparedOcrRequest, ProviderOcrRequest, Value> for OcrLi
}
}
fn upstream_headers(
headers: &[(String, String)],
auth_strategy: OcrAuthStrategy,
api_key: Option<&str>,
) -> Vec<(String, String)> {
api_key
.map(|api_key| match auth_strategy {
OcrAuthStrategy::Bearer => ("Authorization".to_string(), format!("Bearer {api_key}")),
OcrAuthStrategy::Header(header_name) => (header_name.to_string(), api_key.to_string()),
})
.into_iter()
.chain(headers.iter().cloned())
.collect()
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Ocr,
@ -278,19 +337,19 @@ fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
fn parse_ocr_pre_call_guardrail_request(
request: GuardrailRequest,
) -> CoreResult<(Value, Map<String, Value>)> {
) -> Result<(Value, Map<String, Value>), Error> {
let Value::Object(mut data) = request.data else {
return Err(CoreError::InvalidRequest(
return Err(Error::InvalidRequest(
"OCR pre_call guardrail must return an object".to_string(),
));
};
let document = data.remove("document").ok_or_else(|| {
CoreError::InvalidRequest("OCR pre_call guardrail removed document".to_string())
Error::InvalidRequest("OCR pre_call guardrail removed document".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(params)) => params,
Some(_) => {
return Err(CoreError::InvalidRequest(
return Err(Error::InvalidRequest(
"OCR pre_call guardrail optional_params must be an object".to_string(),
));
}
@ -299,33 +358,32 @@ fn parse_ocr_pre_call_guardrail_request(
Ok((document, optional_params))
}
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> CoreResult<Value> {
fn parse_ocr_during_call_guardrail_request(request: GuardrailRequest) -> Result<Value, Error> {
let Value::Object(mut data) = request.data else {
return Err(CoreError::InvalidRequest(
return Err(Error::InvalidRequest(
"OCR during_call guardrail must return an object".to_string(),
));
};
data.remove("body").ok_or_else(|| {
CoreError::InvalidRequest("OCR during_call guardrail removed body".to_string())
})
data.remove("body")
.ok_or_else(|| Error::InvalidRequest("OCR during_call guardrail removed body".to_string()))
}
fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
fn guardrail_error_to_core_error(error: GuardrailError) -> Error {
Error::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &CoreError) -> &'static str {
fn core_error_kind(error: &Error) -> &'static str {
match error {
CoreError::Auth(_) => "AuthError",
CoreError::InvalidProvider(_) => "InvalidProvider",
CoreError::InvalidRequest(_) => "InvalidRequest",
CoreError::InvalidType { .. } => "InvalidType",
CoreError::MissingField(_) => "MissingField",
CoreError::Http { .. } => "HttpError",
CoreError::InvalidResponse(_) => "InvalidResponse",
CoreError::Network(_) => "NetworkError",
CoreError::Connect(_) => "ConnectError",
CoreError::Routing(_) => "RoutingError",
CoreError::Unsupported(_) => "UnsupportedRequest",
Error::Auth(_) => "AuthError",
Error::InvalidProvider(_) => "InvalidProvider",
Error::InvalidRequest(_) => "InvalidRequest",
Error::InvalidType { .. } => "InvalidType",
Error::MissingField(_) => "MissingField",
Error::Http { .. } => "HttpError",
Error::InvalidResponse(_) => "InvalidResponse",
Error::Network(_) => "NetworkError",
Error::Connect(_) => "ConnectError",
Error::Routing(_) => "RoutingError",
Error::Unsupported(_) => "UnsupportedRequest",
}
}

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

@ -1,4 +1,4 @@
use litellm_core::CoreResult;
use litellm_core::Error;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
@ -13,12 +13,162 @@ pub use types::OcrRequest;
use handler::execute_ocr_provider_call;
use prepare::{PreparedOcrCall, prepare_ocr_call};
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn ocr(request: OcrRequest<'_>) -> Result<Value, Error> {
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);
CallLifecycle::default()
.run_request(request, &hooks, execute_ocr_provider_call)
.run_request(request, &hooks, |request| {
execute_ocr_provider_call(request, &hooks)
})
.await
}
#[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""#));
}
}

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

@ -2,7 +2,9 @@ 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;
use super::types::{OcrRequest, PreparedOcrRequest};
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
@ -13,6 +15,7 @@ pub(crate) struct PreparedOcrCall {
pub(crate) hooks: OcrLifecycleHooks,
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
let call_id = request
.litellm_call_id
@ -25,9 +28,41 @@ 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()))
.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();
let mut mapped = config.map_ocr_params(
&request
.optional_params
.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,
};
PreparedOcrCall {
request: PreparedOcrRequest {
config,
model,
custom_llm_provider,
litellm_call_id: call_id,
@ -35,7 +70,7 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall {
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
optional_params,
timeout: request.timeout,
},
hooks: OcrLifecycleHooks::new(
@ -46,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);
@ -55,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`"))
);
}
}

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

@ -25,6 +25,7 @@ pub struct OcrRequest<'a> {
}
pub(crate) struct PreparedOcrRequest {
pub(crate) config: Result<&'static dyn OcrProviderConfig, litellm_core::Error>,
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) litellm_call_id: String,
@ -52,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.

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

@ -1,39 +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::CoreResult;
use litellm_core::error::CoreError;
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) -> CoreResult<Router> {
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| CoreError::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| CoreError::Routing(format!("serializing model_list failed: {err}")))?;
let deployments: Vec<Deployment> = serde_json::from_str(&model_list_json)
.map_err(|err| CoreError::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;

44
litellm-rust/crates/ai-gateway/src/realtime/streaming.rs
View file

@ -106,16 +106,16 @@ impl RealTimeStreaming {
/// `litellm_call_id`, replacing the gateway-generated fallback.
fn on_session(&mut self, event: &RealtimeEvent) {
let session = event.data.get("session").and_then(Value::as_object);
if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str) {
if !id.is_empty() {
self.id = id.to_string();
self.litellm_call_id = id.to_string();
}
if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str)
&& !id.is_empty()
{
self.id = id.to_string();
self.litellm_call_id = id.to_string();
}
if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str) {
if !model.is_empty() {
self.model = model.to_string();
}
if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str)
&& !model.is_empty()
{
self.model = model.to_string();
}
}
@ -323,6 +323,32 @@ mod tests {
assert_eq!(streaming.dropped(), 0);
}
#[test]
fn blank_session_id_and_model_keep_the_gateway_fallbacks() {
let mut streaming = RealTimeStreaming::new(
Vec::new(),
"call_fallback".to_string(),
"gpt-realtime".to_string(),
RequestMetadata::default(),
);
streaming.observe(&event(
r#"{"type":"session.created","session":{"id":"","model":""}}"#,
));
let payload = streaming.build_payload();
assert_eq!(payload.id, "call_fallback");
assert_eq!(payload.litellm_call_id, "call_fallback");
assert_eq!(payload.model, "gpt-realtime");
streaming.observe(&event(
r#"{"type":"session.updated","session":{"id":"sess_002","model":""}}"#,
));
let payload = streaming.build_payload();
assert_eq!(payload.id, "sess_002");
assert_eq!(payload.litellm_call_id, "sess_002");
assert_eq!(payload.model, "gpt-realtime");
}
#[test]
fn payload_serializes_with_camelcase_times_and_realtime_call_type() {
let mut streaming = RealTimeStreaming::new(

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

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

@ -9,7 +9,7 @@ use axum::http::StatusCode;
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue};
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use litellm_core::CoreError;
use litellm_core::Error;
use serde_json::{Map, Value};
use crate::auth::RequireMasterKey;
@ -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>,
@ -46,7 +52,7 @@ fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRout
let mut response = Response::builder()
.status(
StatusCode::from_u16(upstream.status().as_u16()).map_err(|error| {
MessagesRouteError(CoreError::InvalidResponse(format!(
MessagesRouteError(Error::InvalidResponse(format!(
"invalid upstream response status: {error}"
)))
})?,
@ -58,13 +64,13 @@ fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRout
response
.body(Body::from_stream(upstream.bytes_stream()))
.map_err(|error| {
MessagesRouteError(CoreError::InvalidResponse(format!(
MessagesRouteError(Error::InvalidResponse(format!(
"failed to build streaming response: {error}"
)))
})
}
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, CoreError> {
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, Error> {
let forwarded = headers
.iter()
.filter(|(name, _)| {
@ -74,19 +80,19 @@ fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>,
})
.map(|(name, value)| {
let value = value.to_str().map_err(|_| {
CoreError::InvalidRequest(format!("invalid value for header {}", name.as_str()))
Error::InvalidRequest(format!("invalid value for header {}", name.as_str()))
})?;
Ok((name.to_string(), Value::String(value.to_string())))
})
.collect::<Result<Map<_, _>, CoreError>>()?;
.collect::<Result<Map<_, _>, Error>>()?;
Ok((!forwarded.is_empty()).then_some(forwarded))
}
#[derive(Debug)]
struct MessagesRouteError(CoreError);
struct MessagesRouteError(Error);
impl From<CoreError> for MessagesRouteError {
fn from(error: CoreError) -> Self {
impl From<Error> for MessagesRouteError {
fn from(error: Error) -> Self {
Self(error)
}
}
@ -94,28 +100,28 @@ impl From<CoreError> for MessagesRouteError {
impl IntoResponse for MessagesRouteError {
fn into_response(self) -> Response {
let (status, message) = match self.0 {
CoreError::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
CoreError::InvalidProvider(_) | CoreError::Routing(_) => (
Error::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
Error::InvalidProvider(_) | Error::Routing(_) => (
StatusCode::NOT_FOUND,
"no messages deployment is configured for this model".to_string(),
),
CoreError::Auth(_) => (
Error::Auth(_) => (
StatusCode::BAD_GATEWAY,
"messages provider authentication failed".to_string(),
),
CoreError::Http { .. }
| CoreError::Network(_)
| CoreError::Connect(_)
| CoreError::InvalidResponse(_)
| CoreError::InvalidType { .. }
| CoreError::MissingField(_) => (
Error::Http { .. }
| Error::Network(_)
| Error::Connect(_)
| Error::InvalidResponse(_)
| Error::InvalidType { .. }
| Error::MissingField(_) => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),
// The gateway has no Python implementation to decline to, so a
// request the core cannot serve is reported to the caller. The
// reason is a fixed internal string, never provider content.
CoreError::Unsupported(reason) => (
Error::Unsupported(reason) => (
StatusCode::BAD_REQUEST,
format!("messages request is not supported: {reason}"),
),

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

@ -1,10 +1,10 @@
use std::sync::Arc;
use litellm_core::Error;
use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER;
use litellm_core::messages::types::MessagesRequest;
use litellm_core::messages::{messages, messages_stream};
use litellm_core::router::Router;
use litellm_core::{CoreError, CoreResult};
use serde_json::{Map, Value};
pub(crate) enum MessagesResponse {
@ -12,20 +12,26 @@ 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,
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<MessagesResponse> {
) -> Result<MessagesResponse, Error> {
let model = body
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|model| !model.is_empty())
.ok_or_else(|| CoreError::InvalidRequest("messages body requires a model".to_string()))?;
let deployment = router.get_available_deployment(model).ok_or_else(|| {
CoreError::Routing(format!("no deployment available for model '{model}'"))
})?;
.ok_or_else(|| Error::InvalidRequest("messages body requires a model".to_string()))?;
let deployment = router
.get_available_deployment(model)
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
let provider_model = deployment.litellm_params.model.as_str();
let upstream_model = provider_model
.split_once('/')
@ -37,7 +43,7 @@ pub async fn run(
};
let mut body = body;
body.as_object_mut()
.ok_or_else(|| CoreError::InvalidRequest("messages body must be an object".to_string()))?
.ok_or_else(|| Error::InvalidRequest("messages body must be an object".to_string()))?
.insert(
"model".to_string(),
Value::String(upstream_model.to_string()),
@ -60,6 +66,6 @@ pub async fn run(
serde_json::to_value(response)
.map(MessagesResponse::Json)
.map_err(|err| {
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
Error::InvalidResponse(format!("failed to serialize messages response: {err}"))
})
}

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

34
litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs
View file

@ -11,8 +11,7 @@ use std::time::Duration;
use crate::io::realtime_pool::{RealtimePool, upstream_key};
use futures_util::{Sink, Stream};
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::error::Error;
use litellm_core::realtime::types::RealtimeEvent;
use litellm_core::router::Router;
@ -29,15 +28,15 @@ pub async fn run<In, Out>(
observe: impl FnMut(&RealtimeEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = RealtimeEvent> + Unpin + Send,
Out: Sink<RealtimeEvent> + Unpin + Send,
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
{
let deployment = router.get_available_deployment(model).ok_or_else(|| {
CoreError::Routing(format!("no deployment available for model '{model}'"))
})?;
let deployment = router
.get_available_deployment(model)
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
let params = &deployment.litellm_params;
// Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model.
let provider_model = params
@ -51,18 +50,17 @@ where
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
) {
if let Some(handoff) = pool.take(&key) {
return crate::io::realtime::realtime_warm(
provider_model,
handoff,
idle_timeout,
observe,
client_in,
client_out,
)
.await;
}
) && let Some(handoff) = pool.take(&key)
{
return crate::io::realtime::realtime_warm(
provider_model,
handoff,
idle_timeout,
observe,
client_in,
client_out,
)
.await;
}
// Cold path: fresh dial (the original behavior).

12
litellm-rust/crates/ai-gateway/src/routes/responses/service.rs
View file

@ -2,13 +2,13 @@ use std::sync::Arc;
use std::time::Duration;
use futures_util::{Sink, Stream};
use litellm_core::Error;
use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext};
use litellm_core::responses::instrumentation::{
ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
ResponsesWsMetadata,
};
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::{CoreError, CoreResult};
use crate::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
@ -26,22 +26,22 @@ pub async fn run<In, Out>(
metadata: RequestMetadata,
client_in: In,
client_out: Out,
) -> CoreResult<()>
) -> Result<(), Error>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let deployment = router.get_available_deployment(model).ok_or_else(|| {
CoreError::Routing(format!("no deployment available for model '{model}'"))
})?;
let deployment = router
.get_available_deployment(model)
.ok_or_else(|| Error::Routing(format!("no deployment available for model '{model}'")))?;
let params = &deployment.litellm_params;
let provider_model = params
.model
.strip_prefix("openai/")
.unwrap_or(&params.model);
if params.model.contains('/') && !params.model.starts_with("openai/") {
return Err(CoreError::InvalidProvider(
return Err(Error::InvalidProvider(
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));
}

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

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

@ -1,22 +1,19 @@
use std::sync::{Arc, Mutex};
use std::time::Duration;
use litellm_core::error::CoreError;
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::{has_header, 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;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -135,6 +132,7 @@ struct RecordingOcrGuardrail {
hooks: Vec<GuardrailEventHook>,
events: Mutex<Vec<&'static str>>,
block_pre_call: bool,
block_during_call: bool,
}
impl RecordingOcrGuardrail {
@ -143,6 +141,7 @@ impl RecordingOcrGuardrail {
hooks,
events: Mutex::new(Vec::new()),
block_pre_call: false,
block_during_call: false,
}
}
@ -151,6 +150,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,
}
}
@ -192,91 +201,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]
@ -395,7 +408,7 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
.await
.expect_err("provider error propagates");
assert!(matches!(err, CoreError::Http { status: 500, .. }));
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
assert_eq!(
logger.events(),
@ -439,7 +452,7 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() {
.await
.expect_err("guardrail blocks request");
assert!(matches!(err, CoreError::InvalidRequest(_)));
assert!(matches!(err, Error::InvalidRequest(_)));
assert_eq!(guardrail.events(), vec!["async_pre_call_hook"]);
assert_eq!(
logger.events(),
@ -594,21 +607,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,
CoreError::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(_)));
}
}

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

@ -6,11 +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 }
@ -29,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

14
litellm-rust/crates/core/src/audio_transcription/client.rs Normal file
View file

@ -0,0 +1,14 @@
use std::sync::OnceLock;
use std::time::Duration;
use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS;
pub(super) fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS))
.build()
.unwrap_or_else(|_| reqwest::Client::new())
})
}

91
litellm-rust/crates/core/src/audio_transcription/handler.rs Normal file
View file

@ -0,0 +1,91 @@
use serde_json::Value;
use crate::error::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
use super::types::ProviderAudioTranscriptionRequest;
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
) -> Result<Value, Error> {
let body = serde_json::to_vec(&request.body)
.map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
let headers = signed_headers(&request, &body).await?;
let mut request_builder = http_client().post(&request.url).body(body);
for (key, value) in headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder)
.await
.map_err(|error| Error::Network(error.to_string()))?;
let status = response.status();
let text = 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(&text),
});
}
let response_json = serde_json::from_str(&text)
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.into_json())
}
#[cfg(feature = "bedrock-auth")]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::collections::BTreeMap;
use std::time::SystemTime;
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
use crate::providers::bedrock::audio_transcription::aws_auth_config;
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
return Ok(request.upstream_headers.clone());
};
let env_lookup = |key: &str| std::env::var(key).ok();
let credentials = resolve_credentials(
aws_auth_config(&request.optional_params, &env_lookup),
&env_lookup,
)
.await?;
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
let signature = sign_bedrock_post(
&request.url,
body,
&unsigned,
region,
&credentials,
SystemTime::now(),
)?;
Ok(unsigned.into_iter().chain(signature).collect())
}
#[cfg(not(feature = "bedrock-auth"))]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
_body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
match request.auth {
AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()),
}
}

19
litellm-rust/crates/core/src/audio_transcription/mod.rs
View file

@ -1,2 +1,21 @@
use crate::Error;
mod client;
mod handler;
mod prepare;
pub mod transformation;
pub mod types;
use serde_json::Value;
pub use handler::execute_audio_transcription_provider_call;
pub use prepare::prepare_audio_transcription_provider_call;
pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?)
.await
}
#[cfg(test)]
mod tests;

74
litellm-rust/crates/core/src/audio_transcription/prepare.rs Normal file
View file

@ -0,0 +1,74 @@
use crate::error::Error;
use crate::http_utils::{has_header, string_headers};
#[cfg(feature = "bedrock-auth")]
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
#[cfg(feature = "bedrock-auth")]
if provider == "bedrock" {
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
}
let _ = provider;
None
}
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub fn prepare_audio_transcription_provider_call(
request: AudioTranscriptionRequest<'_>,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.or_else(|| {
request
.custom_llm_provider
.map(|provider| CustomLlmProvider {
model: request.model,
custom_llm_provider: provider,
})
})
.ok_or_else(|| {
Error::InvalidProvider(
"unable to resolve custom_llm_provider for audio transcription request".to_string(),
)
})?;
let model = provider_info.model.to_string();
let config = provider_config(provider_info.custom_llm_provider)
.ok_or_else(|| Error::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers("audio transcription", request.extra_headers)?;
let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?;
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(&headers, "authorization")
&& let Some(api_key) = request.api_key
{
headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
}
if !has_header(&headers, "content-type") {
headers.push(("Content-Type".to_string(), "application/json".to_string()));
}
let url = config.complete_url(
request.api_base,
&model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let transformed =
config.transform_transcription_request(&model, request.audio, filtered_params)?;
Ok(ProviderAudioTranscriptionRequest {
model,
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
config,
url,
body: transformed.body,
upstream_headers: headers,
auth,
#[cfg(feature = "bedrock-auth")]
optional_params: request.optional_params,
timeout: request.timeout,
})
}

50
litellm-rust/crates/core/src/audio_transcription/tests.rs Normal file
View file

@ -0,0 +1,50 @@
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use serde_json::{Map, json};
use super::audio_transcription;
use super::types::AudioTranscriptionRequest;
#[tokio::test]
async fn bedrock_request_is_signed_and_contains_audio() {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let address = listener.local_addr().expect("address");
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("connection");
let mut request = Vec::new();
let mut buffer = [0_u8; 16_384];
let count = stream.read(&mut buffer).expect("request");
request.extend_from_slice(&buffer[..count]);
let request = String::from_utf8_lossy(&request);
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
assert!(request.contains("x-amz-date:"));
assert!(request.contains("\"bytes\":\"AQI=\""));
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
stream.write_all(response).expect("response");
});
let optional_params = Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let api_base = format!("http://{address}");
let response = audio_transcription(AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
api_key: None,
api_base: Some(&api_base),
custom_llm_provider: Some("bedrock"),
extra_headers: None,
optional_params,
timeout: None,
})
.await
.expect("transcription");
assert_eq!(response, json!({"text": "hello"}));
server.join().expect("server");
}

12
litellm-rust/crates/core/src/audio_transcription/transformation.rs
View file

@ -1,7 +1,6 @@
use crate::Error;
use serde_json::{Map, Value};
use crate::CoreResult;
use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData};
#[derive(Clone, Debug, PartialEq, Eq)]
@ -16,6 +15,7 @@ pub enum AudioTranscriptionAuth {
pub trait AudioTranscriptionProviderConfig: Sync {
fn supported_transcription_params(&self) -> &'static [&'static str];
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
fn map_transcription_params(&self, params: &Map<String, Value>) -> Map<String, Value> {
params
.iter()
@ -32,13 +32,13 @@ pub trait AudioTranscriptionProviderConfig: Sync {
model: &str,
audio: Value,
optional_params: Map<String, Value>,
) -> CoreResult<AudioTranscriptionRequestData>;
) -> Result<AudioTranscriptionRequestData, Error>;
fn transform_transcription_response(
&self,
model: &str,
response_json: Value,
) -> CoreResult<AudioTranscriptionResponseData>;
) -> Result<AudioTranscriptionResponseData, Error>;
fn complete_url(
&self,
@ -46,12 +46,12 @@ pub trait AudioTranscriptionProviderConfig: Sync {
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String>;
) -> Result<String, Error>;
fn auth_strategy(
&self,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<AudioTranscriptionAuth>;
) -> Result<AudioTranscriptionAuth, Error>;
}

53
litellm-rust/crates/core/src/audio_transcription/types.rs
View file

@ -1,5 +1,56 @@
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_json::{Map, Value};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
pub audio: Value,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}
#[derive(Clone)]
pub struct ProviderAudioTranscriptionRequest {
pub(super) model: String,
pub(super) custom_llm_provider: String,
pub(super) config: &'static dyn AudioTranscriptionProviderConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: AudioTranscriptionAuth,
#[cfg(feature = "bedrock-auth")]
pub(super) optional_params: Map<String, Value>,
pub(super) timeout: Option<Duration>,
}
impl ProviderAudioTranscriptionRequest {
pub fn model(&self) -> &str {
&self.model
}
pub fn custom_llm_provider(&self) -> &str {
&self.custom_llm_provider
}
pub fn url(&self) -> &str {
&self.url
}
pub fn body(&self) -> &Value {
&self.body
}
pub fn with_body(self, body: Value) -> Self {
Self { body, ..self }
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AudioTranscriptionRequestData {

34
litellm-rust/crates/core/src/call_lifecycle/mod.rs
View file

@ -1,7 +1,7 @@
use std::future::Future;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use crate::{CoreError, CoreResult};
use crate::Error;
pub mod types;
@ -11,14 +11,14 @@ pub use types::{
};
pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
type PreCallFuture<'a>: Future<Output = CoreResult<InitialReq>> + Send + 'a
type PreCallFuture<'a>: Future<Output = Result<InitialReq, Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
ProviderReq: 'a,
Resp: 'a;
type DuringCallFuture<'a>: Future<Output = CoreResult<ProviderReq>> + Send + 'a
type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
@ -56,7 +56,7 @@ pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a CoreError,
error: &'a Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a>;
}
@ -86,12 +86,12 @@ impl<'a> CallLifecycle<'a> {
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> CoreResult<Resp>
) -> Result<Resp, Error>
where
InitialReq: CallLifecycleRequest,
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = CoreResult<Resp>>,
ProviderFuture: Future<Output = Result<Resp, Error>>,
{
let context = request.lifecycle_context();
self.run(context, request, hooks, provider_call).await
@ -103,11 +103,11 @@ impl<'a> CallLifecycle<'a> {
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> CoreResult<Resp>
) -> Result<Resp, Error>
where
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = CoreResult<Resp>>,
ProviderFuture: Future<Output = Result<Resp, Error>>,
{
let call_start = epoch_seconds();
let mut phases = Vec::new();
@ -166,7 +166,7 @@ impl<'a> CallLifecycle<'a> {
&self,
context: &CallLifecycleContext,
hooks: &Hooks,
error: &CoreError,
error: &Error,
call_start: f64,
phases: &mut Vec<CallLifecyclePhaseTiming>,
) where
@ -251,8 +251,8 @@ mod tests {
}
impl CallLifecycleHooks<String, String, String> for RecordingHooks {
type PreCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
type PreCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
@ -294,7 +294,7 @@ mod tests {
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a CoreError,
_error: &'a Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -304,8 +304,8 @@ mod tests {
}
impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
type PreCallFuture<'a> = BoxFuture<'a, CoreResult<RecordingRequest>>;
type DuringCallFuture<'a> = BoxFuture<'a, CoreResult<String>>;
type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
@ -345,7 +345,7 @@ mod tests {
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a CoreError,
_error: &'a Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -383,13 +383,13 @@ mod tests {
"request".to_string(),
&hooks,
|_request| async move {
Err::<String, CoreError>(CoreError::Network("provider down".to_string()))
Err::<String, Error>(Error::Network("provider down".to_string()))
},
)
.await
.expect_err("call fails");
assert_eq!(error, CoreError::Network("provider down".to_string()));
assert_eq!(error, Error::Network("provider down".to_string()));
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
}

8
litellm-rust/crates/core/src/chat_completions/common_utils.rs
View file

@ -1,13 +1,13 @@
use serde_json::{Map, Value};
use crate::error::CoreResult;
use crate::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use serde_json::{Map, Value};
use super::transformation::ChatCompletionsProviderConfig;
const HEADER_CONTEXT: &str = "chat completions";
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) fn chat_completions_provider_config(
provider: &str,
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
@ -23,6 +23,6 @@ pub(super) fn chat_completions_provider_config(
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<Vec<(String, String)>> {
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers)
}

40
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -1,19 +1,23 @@
use serde_json::Value;
use crate::error::{CoreError, CoreResult};
use crate::http_utils::truncate_error_body;
use crate::error::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::client::http_client;
use super::prepare::prepare_provider_request;
use super::transformation::ChatCompletionsAuth;
use super::types::{
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
ResolvedChatCompletionsRequest,
};
#[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)]
pub(super) async fn execute_chat_completions_provider_call(
request: ProviderChatCompletionsRequest,
) -> CoreResult<ChatCompletionsResponse> {
request: ResolvedChatCompletionsRequest<'_>,
) -> Result<ChatCompletionsResponse, Error> {
let request = prepare_provider_request(request)?;
let body = serde_json::to_vec(&request.body).map_err(|err| {
CoreError::InvalidRequest(format!(
Error::InvalidRequest(format!(
"failed to serialize chat completions request: {err}"
))
})?;
@ -27,14 +31,14 @@ pub(super) async fn execute_chat_completions_provider_call(
request_builder = request_builder.timeout(duration);
}
let response = request_builder.send().await.map_err(|err| {
let response = http_request(request_builder).await.map_err(|err| {
// Failing to establish the connection means the request never went out,
// so the host can still serve it. Everything else here, a timeout
// above all, may have reached the provider and been answered.
if err.is_connect() || err.is_builder() {
CoreError::Connect(err.to_string())
Error::Connect(err.to_string())
} else {
CoreError::Network(err.to_string())
Error::Network(err.to_string())
}
})?;
@ -42,17 +46,17 @@ pub(super) async fn execute_chat_completions_provider_call(
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
.map_err(|err| Error::Network(err.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
return Err(Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let body: Value = serde_json::from_str(&text).map_err(|err| {
CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
Error::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
})?;
request
.config
@ -69,10 +73,10 @@ pub(super) async fn execute_chat_completions_provider_call(
/// second kind has already been billed, and a host that keeps a reference
/// implementation must not retry those, so collapse them to one variant that
/// can only mean the provider was already called.
pub(super) fn as_response_error(err: CoreError) -> CoreError {
pub(super) fn as_response_error(err: Error) -> Error {
match err {
already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already,
other => CoreError::InvalidResponse(other.to_string()),
already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already,
other => Error::InvalidResponse(other.to_string()),
}
}
@ -80,7 +84,7 @@ pub(super) fn as_response_error(err: CoreError) -> CoreError {
pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
body: &[u8],
) -> CoreResult<Vec<(String, String)>> {
) -> Result<Vec<(String, String)>, Error> {
use std::collections::BTreeMap;
use std::time::SystemTime;
@ -101,7 +105,7 @@ pub(super) async fn signed_headers(
.iter()
.any(|(name, _)| is_sigv4_computed_header(name))
{
return Err(CoreError::Unsupported(
return Err(Error::Unsupported(
"request forwards a header AWS SigV4 computes",
));
}
@ -137,9 +141,9 @@ pub(super) async fn signed_headers(
pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
_body: &[u8],
) -> CoreResult<Vec<(String, String)>> {
) -> Result<Vec<(String, String)>, Error> {
match &request.auth {
ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported(
ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
_ => Ok(request.upstream_headers.clone()),

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