scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

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

Browse source
This commit is contained in:
Yuneng Jiang 2026-09-19 23:59:26 -07:00
commit e34fd201b3
No known key found for this signature in database
1129 changed files with 72125 additions and 13464 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
.circleci/config.yml
View file

@ -1673,6 +1673,7 @@ jobs:
command: |
docker run -d \
-p 4001:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL="postgresql://postgres:postgres@host.docker.internal:5432/litellm_test" \
-e LITELLM_MASTER_KEY="sk-1234" \
--name schema-seed \
@ -1693,6 +1694,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL="postgresql://postgres:postgres@host.docker.internal:5432/litellm_test" \
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
-e DISABLE_SCHEMA_UPDATE="True" \
@ -1767,7 +1769,9 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e LITELLM_MASTER_KEY="sk-1234" \
-e USE_PRISMA_MIGRATE=True \
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
-e AZURE_API_KEY=$AZURE_API_KEY \
@ -1862,7 +1866,9 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e LITELLM_MASTER_KEY="sk-1234" \
-e AZURE_API_KEY=$AZURE_API_KEY \
-e AZURE_API_BASE=$AZURE_API_BASE \
-e AZURE_API_VERSION="2024-05-01-preview" \
@ -1950,6 +1956,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e REDIS_HOST=$REDIS_HOST \
-e REDIS_PASSWORD=$REDIS_PASSWORD \
@ -2010,6 +2017,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e REDIS_HOST=$REDIS_HOST \
-e REDIS_PASSWORD=$REDIS_PASSWORD \
@ -2087,6 +2095,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e REDIS_HOST=host.docker.internal \
-e REDIS_PORT=6379 \
@ -2169,6 +2178,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e REDIS_HOST=$REDIS_HOST \
-e REDIS_PASSWORD=$REDIS_PASSWORD \
@ -2191,6 +2201,7 @@ jobs:
command: |
docker run -d \
-p 4001:4001 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e REDIS_HOST=$REDIS_HOST \
-e REDIS_PASSWORD=$REDIS_PASSWORD \
@ -2268,6 +2279,7 @@ jobs:
docker run -d \
--restart on-failure \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e STORE_MODEL_IN_DB="True" \
-e LITELLM_MASTER_KEY="sk-1234" \
@ -2342,6 +2354,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e REDIS_HOST=$REDIS_HOST \
-e REDIS_PASSWORD=$REDIS_PASSWORD \
@ -2424,6 +2437,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e LITELLM_MASTER_KEY="sk-1234" \
-e OPENAI_API_KEY=$OPENAI_API_KEY \
@ -2515,6 +2529,7 @@ jobs:
command: |
docker run -d \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e LITELLM_MASTER_KEY="sk-1234" \
-e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
@ -2696,6 +2711,7 @@ jobs:
name: Start LiteLLM proxy
environment:
LITELLM_MASTER_KEY: "sk-1234"
LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY: "true"
MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
DISABLE_SCHEMA_UPDATE: "true"
SERVER_ROOT_PATH: ""
@ -2839,6 +2855,7 @@ jobs:
name: Start LiteLLM proxy under a server root path
environment:
LITELLM_MASTER_KEY: "sk-1234"
LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY: "true"
MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
DISABLE_SCHEMA_UPDATE: "true"
# Output flows to this step's own log, so a boot crash is visible here
@ -3017,6 +3034,7 @@ jobs:
command: |
docker run --name my-app \
-p 4000:4000 \
-e LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY=true \
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
-e DATABASE_URL="postgresql://wrong:wrong@wrong:5432/wrong" \
myapp:latest \
@ -3157,7 +3175,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, sdk, browser]
suite: [management, accounting, database, providers, extensions, sdk, cost, browser]
filters:
branches:
only:

11
.circleci/scripts/classify_changes.sh
View file

@ -1,16 +1,22 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only|mcp-dependencies>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
has_cost_map=false
has_mcp_dependencies=false
outside_cost_map_set=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
*.md | *.mdx) : ;;
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py)
has_mcp_dependencies=true ;;
esac
case "$file" in
tests/e2e/*/*.py) : ;;
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
@ -31,6 +37,9 @@ while IFS= read -r file || [ -n "$file" ]; do
done
case "$category" in
mcp-dependencies)
[ "$has_mcp_dependencies" = true ] && echo run || echo skip
;;
cost-map-only)
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
;;

19
.circleci/scripts/run_integration.sh
View file

@ -9,6 +9,7 @@ fi
suite="${1:?integration suite required}"
results="test-results/integration-${suite}"
mkdir -p "$results"
shard_timeout=11m
integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
upstream_pid=""
proxy_pid=""
@ -108,13 +109,26 @@ awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/e
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
.venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 &
upstream_pid=$!
if [ "$suite" = cost ]; then
export INTEGRATION_WORKERS=8
fi
start_proxy() {
local port="$1"
local log_name="$2"
local -a cost_map_env
if [ "$suite" = cost ]; then
cost_map_env=(
"LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map"
"MODEL_COST_MAP_MIN_MODEL_COUNT=1"
"MODEL_COST_MAP_MAX_SHRINK_RATIO=0"
)
else
cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")
fi
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
@ -158,11 +172,12 @@ if [ "$suite" = browser ]; then
exit 0
fi
timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
timeout --signal=TERM --kill-after=20s "$shard_timeout" env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \

3
.env.example
View file

@ -26,6 +26,7 @@ NOVITA_API_KEY = ""
INFINITY_API_KEY = ""
# Development Configs
LITELLM_MASTER_KEY = "sk-1234"
# Generate one with: echo "LITELLM_MASTER_KEY=sk-$(openssl rand -hex 32)"
LITELLM_MASTER_KEY = ""
DATABASE_URL = "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
STORE_MODEL_IN_DB = "True"

2
.github/actions/detect-changes/action.yml vendored
View file

@ -14,7 +14,7 @@ description: >-
inputs:
category:
description: "Which classification to apply: backend, client or ui"
description: "Which classification to apply: backend, client, ui, provider-harness, cost-map-only or mcp-dependencies"
required: false
default: backend
github-token:

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

@ -1,3 +1,5 @@
import os
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
@ -15,6 +17,7 @@ def main() -> int:
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
@ -35,9 +38,22 @@ def main() -> int:
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
for case in cases:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")):
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error", "skipped")):
continue
_ = sys.stdout.write(f" failed: {case.get('classname', '')}::{case.get('name', '')}\n")
outcome = "skipped" if case.find("skipped") is not None else "failed"
_ = sys.stdout.write(f" {outcome}: {case.get('classname', '')}::{case.get('name', '')}\n")
for prop in case.findall("./properties/property"):
name = prop.get("name", "")
value = prop.get("value", "")
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
r"[A-Za-z0-9_.:<>-]{1,240}", value
):
_ = sys.stdout.write(f" {name}: {value}\n")
if expected_count is not None and (
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
):
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
return 1
if (
selected
and not missing

1
.github/e2e-stack/select_tests.py vendored
View file

@ -5,6 +5,7 @@ from typing import Final
SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
UNSUPPORTED: Final = re.compile(
r"^tests/e2e/(ui|claude_code|load|migrations)/"
r"|^tests/e2e/mcp/test_mcp_oauth_happy_path_e2e\.py$"
r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"

393
.github/scripts/auto_merge_price_sync.py vendored
View file

@ -1,393 +0,0 @@
"""Auto-merge the provider-info-sync bot's cost-map pull requests.
Evaluates every gate (author allowlist, cost-map-only diff, required and
non-required checks, human reviews) and merges with a merge commit when
all of them hold. Every hold reason is logged; the process exits 0 on hold
and 1 only on API or programming errors.
``DRY_RUN=1`` prints the verdict without calling the merge endpoint.
"""
from __future__ import annotations
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final
REPO_ROOT: Final = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
CLASSIFY_SCRIPT: Final = os.path.join(REPO_ROOT, ".circleci", "scripts", "classify_changes.sh")
API_ROOT: Final = "https://api.github.com"
CHANGED_FILE_CEILING: Final = 3000
OK_CHECK_CONCLUSIONS: Final = frozenset({"success", "skipped", "neutral"})
@dataclass(frozen=True, slots=True)
class PullRequest:
number: int
title: str
author_login: str
state: str
draft: bool
mergeable: bool | None
mergeable_state: str
head_sha: str
@dataclass(frozen=True, slots=True)
class CheckRun:
name: str
status: str
conclusion: str | None
@dataclass(frozen=True, slots=True)
class CommitStatus:
context: str
state: str
@dataclass(frozen=True, slots=True)
class Review:
author_login: str
state: str
body: str
commit_id: str
submitted_at: datetime
@dataclass(frozen=True, slots=True)
class Verdict:
merge: bool
reasons: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class EvaluationInputs:
pr: PullRequest
changed_files: tuple[str, ...]
required_contexts: frozenset[str]
check_runs: tuple[CheckRun, ...]
statuses: tuple[CommitStatus, ...]
reviews: tuple[Review, ...]
self_check_name: str
author_allowlist: frozenset[str]
def _is_bot_login(login: str) -> bool:
return login.lower().endswith("[bot]")
def _classify(changed_files: Sequence[str]) -> str:
result: Final = subprocess.run(
["bash", CLASSIFY_SCRIPT, "cost-map-only"],
input="\n".join(changed_files),
capture_output=True,
text=True,
check=False,
)
if result.returncode != 0:
return "error"
return result.stdout.strip()
def evaluate(
inputs: EvaluationInputs,
*,
classify: Callable[[Sequence[str]], str] = _classify,
) -> Verdict:
pr: Final = inputs.pr
reasons: list[str] = []
if pr.author_login.lower() not in {login.lower() for login in inputs.author_allowlist}:
reasons.append(f"author {pr.author_login!r} not in allowlist")
if pr.state != "open":
reasons.append("pr not open")
if pr.draft:
reasons.append("pr is a draft")
if pr.mergeable is None:
reasons.append("mergeability unknown")
elif not pr.mergeable:
reasons.append("pr not mergeable")
if pr.mergeable_state == "dirty":
reasons.append("pr has merge conflicts")
if len(inputs.changed_files) > CHANGED_FILE_CEILING:
reasons.append(f"changed file count {len(inputs.changed_files)} over {CHANGED_FILE_CEILING} ceiling")
else:
decision: Final = classify(inputs.changed_files)
if decision != "run":
reasons.append("changed files outside the cost-map-only set")
green_runs: Final = frozenset(run.name for run in inputs.check_runs if run.conclusion in OK_CHECK_CONCLUSIONS)
green_statuses: Final = frozenset(status.context for status in inputs.statuses if status.state == "success")
for context in sorted(inputs.required_contexts):
if context not in green_runs and context not in green_statuses:
reasons.append(f"required check {context!r} not green")
for run in inputs.check_runs:
if run.name == inputs.self_check_name:
continue
if run.status != "completed" or run.conclusion not in OK_CHECK_CONCLUSIONS:
reasons.append(f"check run {run.name!r} is {run.status}/{run.conclusion}")
for status in inputs.statuses:
if status.state != "success":
reasons.append(f"commit status {status.context!r} is {status.state}")
latest_state_by_reviewer: Final[dict[str, str]] = {}
for review in sorted(inputs.reviews, key=lambda review: review.submitted_at):
if _is_bot_login(review.author_login):
continue
latest_state_by_reviewer[review.author_login] = review.state
for reviewer, state in latest_state_by_reviewer.items():
if state == "CHANGES_REQUESTED":
reasons.append(f"changes requested by {reviewer}")
return Verdict(merge=not reasons, reasons=tuple(reasons))
def _request(token: str, method: str, path: str, body: Mapping[str, object] | None = None) -> object:
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
data: Final = None if body is None else json.dumps(body).encode("utf-8")
request: Final = urllib.request.Request(
url,
data=data,
method=method,
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
with urllib.request.urlopen(request) as response:
return json.loads(response.read().decode("utf-8"))
def _request_allow_fail(
token: str, method: str, path: str, body: Mapping[str, object] | None = None
) -> tuple[int, object | None]:
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
data: Final = None if body is None else json.dumps(body).encode("utf-8")
request: Final = urllib.request.Request(
url,
data=data,
method=method,
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
try:
with urllib.request.urlopen(request) as response:
return response.status, json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as exc:
return exc.code, None
def _items(payload: object, key: str | None = None) -> tuple[object, ...]:
source: Final = payload.get(key) if key and isinstance(payload, Mapping) else payload
if not isinstance(source, list):
return ()
return tuple(source)
def _paginate(token: str, path: str, key: str | None = None) -> list[object]:
separator: Final = "&" if "?" in path else "?"
results: list[object] = []
for page in range(1, 10_000):
batch: Final = _items(_request(token, "GET", f"{path}{separator}per_page=100&page={page}"), key)
results.extend(batch)
if len(batch) < 100:
return results
return results
def _text(value: object) -> str:
return value if isinstance(value, str) else ""
def _int(value: object) -> int:
return value if isinstance(value, int) else 0
def _bool(value: object) -> bool:
return value is True
def _nested(value: object, *keys: str) -> object:
current: object = value
for key in keys:
if not isinstance(current, Mapping):
return None
current = current.get(key)
return current
def _parse_time(value: object) -> datetime:
text: Final = _text(value)
if not text:
return datetime.min.replace(tzinfo=timezone.utc)
return datetime.fromisoformat(text.replace("Z", "+00:00"))
def _load_pr(token: str, repo: str, number: int) -> PullRequest:
data: Final = _request(token, "GET", f"/repos/{repo}/pulls/{number}")
if not isinstance(data, Mapping):
raise RuntimeError(f"unexpected pull payload for #{number}")
return PullRequest(
number=number,
title=_text(data.get("title")),
author_login=_text(_nested(data, "user", "login")),
state=_text(data.get("state")),
draft=_bool(data.get("draft")),
mergeable=data.get("mergeable") if isinstance(data.get("mergeable"), bool) else None,
mergeable_state=_text(data.get("mergeable_state")),
head_sha=_text(_nested(data, "head", "sha")),
)
def _list_candidate_prs(token: str, repo: str, base: str, allowlist: frozenset[str]) -> list[int]:
candidates: Final = _paginate(token, f"/repos/{repo}/pulls?state=open&base={base}")
return [
_int(item.get("number"))
for item in candidates
if isinstance(item, Mapping) and _text(_nested(item, "user", "login")).lower() in allowlist
]
def _changed_files(token: str, repo: str, number: int) -> tuple[str, ...]:
files: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/files")
return tuple(_text(item.get("filename")) for item in files if isinstance(item, Mapping))
def _required_contexts(token: str, repo: str, base: str) -> frozenset[str]:
payload: Final = _request(token, "GET", f"/repos/{repo}/rules/branches/{base}")
contexts: set[str] = set()
for rule in _items(payload):
if not isinstance(rule, Mapping) or rule.get("type") != "required_status_checks":
continue
checks: Final = _nested(rule, "parameters", "required_status_checks")
for check in _items(checks):
if isinstance(check, Mapping):
context: Final = _text(check.get("context"))
if context:
contexts.add(context)
return frozenset(contexts)
def _check_runs(token: str, repo: str, sha: str) -> tuple[CheckRun, ...]:
runs: Final = _paginate(token, f"/repos/{repo}/commits/{sha}/check-runs", key="check_runs")
return tuple(
CheckRun(
name=_text(item.get("name")),
status=_text(item.get("status")),
conclusion=item.get("conclusion") if isinstance(item.get("conclusion"), str) else None,
)
for item in runs
if isinstance(item, Mapping)
)
def _statuses(token: str, repo: str, sha: str) -> tuple[CommitStatus, ...]:
payload: Final = _request(token, "GET", f"/repos/{repo}/commits/{sha}/status")
return tuple(
CommitStatus(context=_text(item.get("context")), state=_text(item.get("state")))
for item in _items(payload, "statuses")
if isinstance(item, Mapping)
)
def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]:
reviews: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/reviews")
return tuple(
Review(
author_login=_text(_nested(item, "user", "login")),
state=_text(item.get("state")),
body=_text(item.get("body")),
commit_id=_text(item.get("commit_id")),
submitted_at=_parse_time(item.get("submitted_at")),
)
for item in reviews
if isinstance(item, Mapping)
)
def _mergeable_or_refetch(token: str, repo: str, pr: PullRequest) -> PullRequest:
if pr.mergeable is not None:
return pr
time.sleep(5)
return _load_pr(token, repo, pr.number)
def _gather_inputs(
token: str,
repo: str,
number: int,
base: str,
self_check_name: str,
allowlist: frozenset[str],
) -> EvaluationInputs:
pr: Final = _mergeable_or_refetch(token, repo, _load_pr(token, repo, number))
return EvaluationInputs(
pr=pr,
changed_files=_changed_files(token, repo, number),
required_contexts=_required_contexts(token, repo, base),
check_runs=_check_runs(token, repo, pr.head_sha),
statuses=_statuses(token, repo, pr.head_sha),
reviews=_reviews(token, repo, number),
self_check_name=self_check_name,
author_allowlist=allowlist,
)
def merge_request_body(pr: PullRequest) -> dict[str, str]:
return {"merge_method": "merge", "commit_title": f"{pr.title} (#{pr.number})", "sha": pr.head_sha}
def _merge(token: str, repo: str, pr: PullRequest) -> None:
status, _ = _request_allow_fail(token, "PUT", f"/repos/{repo}/pulls/{pr.number}/merge", merge_request_body(pr))
if status in (200, 405, 409):
print(f"auto-merge-price-sync: PR #{pr.number} merge call returned {status}")
return
raise RuntimeError(f"merge call for PR #{pr.number} returned {status}")
def main() -> int:
token: Final = os.environ.get("GH_TOKEN", "")
repo: Final = os.environ.get("REPO", "")
base: Final = os.environ.get("BASE_BRANCH", "main")
dry_run: Final = os.environ.get("DRY_RUN", "") != ""
self_check_name: Final = os.environ.get("SELF_CHECK_NAME", "auto-merge-price-sync")
allowlist: Final = frozenset(login.lower() for login in os.environ.get("PR_AUTHOR_ALLOWLIST", "").split() if login)
if not token:
print("auto-merge-price-sync: app credentials not configured")
return 0
if not repo:
print("auto-merge-price-sync: REPO not set", file=sys.stderr)
return 1
pr_number_env: Final = os.environ.get("PR_NUMBER", "")
candidates: Final = [int(pr_number_env)] if pr_number_env else _list_candidate_prs(token, repo, base, allowlist)
for number in candidates:
inputs: Final = _gather_inputs(token, repo, number, base, self_check_name, allowlist)
verdict: Final = evaluate(inputs)
for reason in verdict.reasons:
print(f"auto-merge-price-sync: PR #{number} hold: {reason}")
if not verdict.merge:
continue
print(f"auto-merge-price-sync: PR #{number} all gates green")
if dry_run:
print(f"auto-merge-price-sync: DRY_RUN merge suppressed for PR #{number}")
continue
_merge(token, repo, inputs.pr)
return 0
if __name__ == "__main__":
sys.exit(main())

2
.github/template.yaml vendored
View file

@ -21,7 +21,7 @@ Parameters:
WorkerConfigParameter:
Type: String
Description: Sample environment variable
Default: '{"model": null, "alias": null, "api_base": null, "api_version": "2023-07-01-preview", "debug": false, "temperature": null, "max_tokens": null, "request_timeout": 600, "max_budget": null, "telemetry": true, "drop_params": false, "add_function_to_prompt": false, "headers": null, "save": false, "config": null, "use_queue": false}'
Default: '{"model": null, "alias": null, "api_base": null, "api_version": "2023-07-01-preview", "debug": false, "temperature": null, "max_tokens": null, "request_timeout": 600, "max_budget": null, "drop_params": false, "add_function_to_prompt": false, "headers": null, "save": false, "config": null, "use_queue": false}'
Resources:
MyUrlFunctionPermissions:

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

@ -37,6 +37,18 @@ on:
required: false
type: number
default: 60
test-timeout-seconds:
description: >-
Per-test ceiling enforced by pytest-timeout, covering fixture setup and
teardown as well as the test body. A test that hangs fails with a
traceback of where it was stuck instead of idling the shard until
`timeout-minutes` cancels it. Timed-out tests are excluded from reruns
because pytest-timeout arms its timer once per test and
pytest-rerunfailures reruns inside that same window, so a rerun of a
timed-out test would run with no timer at all.
required: false
type: number
default: 120
max-failures:
description: "Stop after this many failures"
required: false
@ -51,6 +63,11 @@ on:
description: "Unique name for the coverage artifact (must be unique per run)"
required: true
type: string
legacy-mcp-peer:
description: "Install the isolated SDK1 peer for MCP compatibility tests"
required: false
type: boolean
default: false
permissions:
contents: read
@ -113,10 +130,17 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
env:
LEGACY_MCP_PEER: ${{ inputs.legacy-mcp-peer }}
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
.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"]'
if [ "$LEGACY_MCP_PEER" = "true" ]; then
uv venv --python "${UV_PYTHON}" .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "MCP_TEST_PEER_PYTHON=$GITHUB_WORKSPACE/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
fi
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -137,6 +161,7 @@ jobs:
MAX_FAILURES: ${{ inputs.max-failures }}
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
TEST_TIMEOUT_SECONDS: ${{ inputs.test-timeout-seconds }}
DIST: ${{ inputs.dist }}
COVERAGE_CORE: sysmon
run: |
@ -146,6 +171,8 @@ jobs:
--maxfail="${MAX_FAILURES}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--timeout="${TEST_TIMEOUT_SECONDS}" \
--rerun-except "from pytest-timeout" \
--durations=20 \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml:coverage.xml \
@ -157,6 +184,8 @@ jobs:
-n "${WORKERS}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--timeout="${TEST_TIMEOUT_SECONDS}" \
--rerun-except "from pytest-timeout" \
--dist="${DIST}" \
--durations=20 \
--cov=./litellm --cov=./enterprise/litellm_enterprise \

61
.github/workflows/auto-merge-price-sync.yml vendored
View file

@ -1,61 +0,0 @@
name: auto-merge-price-sync
on:
issue_comment:
types: [created, edited]
check_suite:
types: [completed]
status: {}
schedule:
- cron: "*/30 * * * *"
workflow_dispatch:
inputs:
pr-number:
description: "Evaluate only this PR number (empty = scan all open sync-bot PRs)"
required: false
default: ""
permissions:
contents: read
pull-requests: read
checks: read
statuses: read
concurrency:
group: auto-merge-price-sync
cancel-in-progress: false
jobs:
auto-merge-price-sync:
runs-on: ubuntu-latest
timeout-minutes: 15
env:
PROVIDER_INFO_SYNC_APP_ID: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
PROVIDER_INFO_SYNC_APP_PRIVATE_KEY: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Mint app token
id: app-token
if: ${{ env.PROVIDER_INFO_SYNC_APP_ID != '' && env.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY != '' }}
uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0
with:
app-id: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
private-key: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
- name: Auto-merge eligible sync PRs
env:
GH_TOKEN: ${{ steps.app-token.outputs.token }}
REPO: ${{ github.repository }}
PR_NUMBER: ${{ (github.event.issue.pull_request && github.event.issue.number) || github.event.inputs.pr-number || '' }}
BASE_BRANCH: main
PR_AUTHOR_ALLOWLIST: "berriai-litellm-provider-info-sync[bot]"
SELF_CHECK_NAME: auto-merge-price-sync
run: python3 .github/scripts/auto_merge_price_sync.py

4
.github/workflows/codspeed.yml vendored
View file

@ -69,7 +69,7 @@ jobs:
uv run --frozen --no-default-groups
--with pytest==8.3.5
--with pytest-codspeed==4.3.0
--with "mcp>=1.26.0,<2.0"
--with "mcp>=2.2.0,<3.0"
--with "a2a-sdk>=1.1.0,<2.0"
pytest
-p pytest_codspeed.plugin
@ -86,7 +86,7 @@ jobs:
uv run --frozen --no-default-groups
--with pytest==8.3.5
--with pytest-codspeed==4.3.0
--with "mcp>=1.26.0,<2.0"
--with "mcp>=2.2.0,<3.0"
--with "a2a-sdk>=1.1.0,<2.0"
pytest
-p pytest_codspeed.plugin

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

@ -0,0 +1,98 @@
name: LiteLLM MCP Dependency Resolution
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
pull-requests: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
resolve:
runs-on: ubuntu-latest
timeout-minutes: 15
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Detect relevant changes
id: changes
uses: ./.github/actions/detect-changes
with:
category: mcp-dependencies
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: ${{ matrix.python-version }}
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-cargo-build
- name: Verify lockfile
if: steps.changes.outputs.decision != 'skip'
run: |
uv lock --check
- name: Check locked runtime installations
if: steps.changes.outputs.decision != 'skip'
run: |
for extra in core mcp proxy; do
args=()
if [ "$extra" != core ]; then args=(--extra "$extra"); fi
UV_PROJECT_ENVIRONMENT=".venv-$extra" .github/scripts/uv_sync_with_retries.sh --frozen --no-dev --no-editable --python ${{ matrix.python-version }} "${args[@]}"
uv pip check --python ".venv-$extra"
if [ "$extra" = core ]; then
checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
else
checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra")
fi
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-$extra/bin/python" "${checker[@]}")
done
- name: Build the public wheel
if: steps.changes.outputs.decision != 'skip'
run: uv build --all-packages --wheel --out-dir dist/mcp-check
- name: Check lowest direct runtime installations
if: steps.changes.outputs.decision != 'skip'
run: |
wheel=$(realpath dist/mcp-check/litellm-[0-9]*.whl)
for extra in core mcp proxy; do
args=()
if [ "$extra" != core ]; then args=(--extra "$extra"); fi
uv pip compile pyproject.toml --no-sources --find-links dist/mcp-check "${args[@]}" --python-version ${{ matrix.python-version }} --resolution lowest-direct -o "lowest-$extra.txt"
uv venv --python ${{ matrix.python-version }} ".venv-lowest-$extra"
uv pip sync --find-links dist/mcp-check --python ".venv-lowest-$extra" "lowest-$extra.txt"
uv pip install --python ".venv-lowest-$extra" --no-deps "$wheel"
uv pip check --python ".venv-lowest-$extra"
if [ "$extra" = core ]; then
checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
else
checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra")
fi
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-lowest-$extra/bin/python" "${checker[@]}")
done

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

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

63
.github/workflows/test-mcp.yml vendored
View file

@ -1,63 +0,0 @@
name: LiteLLM MCP Tests (folder - tests/mcp_tests)
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
pull-requests: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
test:
runs-on: ubuntu-latest
timeout-minutes: 25
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Detect relevant changes
id: changes
uses: ./.github/actions/detect-changes
- name: Thank You Message
run: |
echo "### 🙏 Thank you for contributing to LiteLLM!" >> $GITHUB_STEP_SUMMARY
echo "Your PR is being tested now. We appreciate your help in making LiteLLM better!" >> $GITHUB_STEP_SUMMARY
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
uv lock --check
.github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router
- name: Run MCP tests
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml --durations=5

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

@ -49,6 +49,14 @@ jobs:
fail-fast: false
matrix:
include:
- shard: mcp-integration
artifact-name: mcp-integration
test-path: "tests/mcp_tests"
workers: 2
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
- shard: core-utils
artifact-name: core-utils
test-path: "tests/test_litellm/litellm_core_utils"
@ -254,3 +262,4 @@ jobs:
timeout-minutes: ${{ matrix.timeout-minutes }}
job-timeout-minutes: ${{ matrix.job-timeout-minutes }}
artifact-name: ${{ matrix.artifact-name }}
legacy-mcp-peer: ${{ matrix.shard == 'mcp-integration' }}

5
CONTRIBUTING.md
View file

@ -268,10 +268,13 @@ If you want to build the Docker image yourself:
# Build using the non-root Dockerfile
docker build -f docker/Dockerfile.non_root -t litellm_dev .
# Generate a master key. Requests send it as the bearer token
export LITELLM_MASTER_KEY="sk-$(openssl rand -hex 32)"
# Run with your config
docker run \
-v $(pwd)/proxy_config.yaml:/app/config.yaml \
-e LITELLM_MASTER_KEY="sk-1234" \
-e LITELLM_MASTER_KEY \
-p 4000:4000 \
litellm_dev \
--config /app/config.yaml --detailed_debug

6
README.md
View file

@ -168,7 +168,7 @@ from a2a.utils.constants import TransportProtocol
from uuid import uuid4
base_url = "http://localhost:4000/a2a/my-agent" # LiteLLM proxy + agent name
headers = {"Authorization": "Bearer sk-1234"} # LiteLLM Virtual Key
headers = {"Authorization": "Bearer <your-master-key>"} # LiteLLM master key or a virtual key
async with httpx.AsyncClient(headers=headers, timeout=60.0) as http_client:
resolver = A2ACardResolver(httpx_client=http_client, base_url=base_url)
@ -233,7 +233,7 @@ async with stdio_client(server_params) as (read, write):
```bash
curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
-H 'Authorization: Bearer sk-1234' \
-H 'Authorization: Bearer <your-master-key>' \
-H 'Content-Type: application/json' \
-d '{
"model": "gpt-4o",
@ -255,7 +255,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
"LiteLLM": {
"url": "http://localhost:4000/mcp/",
"headers": {
"x-litellm-api-key": "Bearer sk-1234"
"x-litellm-api-key": "Bearer <your-master-key>"
}
}
}

1
cookbook/livekit_agent_sdk/config.example.yaml
View file

@ -15,7 +15,6 @@ model_list:
litellm_settings:
drop_params: True
telemetry: False
general_settings:
master_key: sk-1234 # Change this to a secure key

1
cookbook/misc/config.yaml
View file

@ -55,7 +55,6 @@ litellm_settings:
# budget_duration: 30d
num_retries: 5
request_timeout: 600
telemetry: False
context_window_fallbacks: [{"gpt-3.5-turbo": ["gpt-3.5-turbo-large"]}]
general_settings:

3
docker/.env.example
View file

@ -3,7 +3,8 @@
# YOU MUST CHANGE THESE BEFORE GOING INTO PRODUCTION
############
LITELLM_MASTER_KEY="sk-1234"
# Generate one with: echo "LITELLM_MASTER_KEY=sk-$(openssl rand -hex 32)"
LITELLM_MASTER_KEY=""
############
# Database - You can change these to any PostgreSQL database that has logical replication enabled.

37
enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
View file

@ -8,7 +8,7 @@
## This provides an LLM Guard Integration for content moderation on the proxy
import asyncio
from typing import Optional
from typing import Final, Optional
import aiohttp
from fastapi import HTTPException
@ -137,15 +137,20 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
return
self.print_verbose("Makes LLM Guard Check")
if call_type not in [
accepted_call_types: Final = (
"completion",
"acompletion",
"text_completion",
"atext_completion",
"embeddings",
"embedding",
"aembedding",
"image_generation",
"moderation",
"audio_transcription",
]:
"aimage_generation",
)
if call_type not in accepted_call_types:
self.print_verbose(
f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']"
f"Call Type - {call_type}, not in accepted list - {accepted_call_types}"
)
return data
@ -163,16 +168,14 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
*(self._moderate_message(message) for message in messages)
)
)
return data
input_ = data.get("input")
if input_ is not None:
data["input"] = await self._moderate_input(input_)
return data
data["input"] = await self._moderate_text_or_list(input_)
prompt = data.get("prompt")
if isinstance(prompt, str):
data["prompt"] = await self.moderation_check(text=prompt)
if prompt is not None:
data["prompt"] = await self._moderate_text_or_list(prompt)
return data
async def _moderate_message(self, message: dict) -> dict:
@ -195,17 +198,17 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
return {**part, "text": await self.moderation_check(text=part["text"])}
return part
async def _moderate_input(self, input_: object) -> object:
if isinstance(input_, str):
return await self.moderation_check(text=input_)
if isinstance(input_, list):
async def _moderate_text_or_list(self, value: object) -> object:
if isinstance(value, str):
return await self.moderation_check(text=value)
if isinstance(value, list):
return [
await self.moderation_check(text=item)
if isinstance(item, str)
else item
for item in input_
for item in value
]
return input_
return value
async def async_post_call_streaming_hook(
self, user_api_key_dict: UserAPIKeyAuth, response: str

9
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py
View file

@ -60,6 +60,11 @@ async def _get_email_settings(prisma_client) -> Dict[str, bool]:
async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
"""Helper function to save email settings to general_settings in db"""
from litellm.proxy.proxy_server import proxy_config
proxy_config.reject_config_owned_writes(
section_name="general_settings", changed_keys={"email_settings": settings}
)
try:
verbose_proxy_logger.debug(
f"Saving email settings to general_settings: {settings}"
@ -168,6 +173,8 @@ async def update_event_settings(
await _save_email_settings(prisma_client, settings_dict)
return {"message": "Email event settings updated successfully"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error updating email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@ -197,6 +204,8 @@ async def reset_event_settings(
await _save_email_settings(prisma_client, default_settings)
return {"message": "Email event settings reset to defaults"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error resetting email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))

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

@ -20,6 +20,7 @@ from typing import (
)
from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
@ -34,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from openai.types.file_deleted import FileDeleted
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
@ -59,6 +61,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_content_type_from_file_object,
get_model_id_from_unified_batch_id,
get_original_file_id,
is_litellm_executed_batch,
map_raw_file_ids_to_unified,
normalize_mime_type_for_provider,
resolve_managed_output_file_model_name,
@ -75,6 +78,7 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess
CreateFileRequest,
FileListPage,
FileObject,
HttpxBinaryResponseContent,
OpenAIFileObject,
ResponsesAPIResponse,
)
@ -86,10 +90,6 @@ from litellm.types.utils import (
SpecialEnums,
)
if TYPE_CHECKING:
from litellm.types.llms.openai import HttpxBinaryResponseContent
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from prisma.models import (
@ -204,6 +204,19 @@ def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableAct
return prisma_client.db.litellm_managedobjecttable
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
)
return MappingProxyType(
{
key: value
for key in ("storage_backend", "storage_url")
if isinstance(value := hidden_params.get(key), str)
}
)
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
@ -226,6 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
) -> None:
verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache")
storage_metadata: Final = _storage_metadata_of(file_object)
if file_object is not None:
litellm_managed_file_object = LiteLLM_ManagedFileTable(
unified_file_id=file_id,
@ -235,6 +249,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
created_by=resolve_resource_owner_id(user_api_key_dict),
team_id=user_api_key_dict.team_id,
updated_by=user_api_key_dict.user_id,
storage_backend=storage_metadata.get("storage_backend"),
storage_url=storage_metadata.get("storage_url"),
)
await self.internal_usage_cache.async_set_cache(
key=file_id,
@ -262,14 +278,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_object_json = file_object.model_dump_json()
db_data["file_object"] = file_object_json
update_data["file_object"] = file_object_json
# Extract storage metadata from hidden params if present
hidden_params = getattr(file_object, "_hidden_params", {}) or {}
if "storage_backend" in hidden_params:
db_data["storage_backend"] = hidden_params["storage_backend"]
update_data["storage_backend"] = hidden_params["storage_backend"]
if "storage_url" in hidden_params:
db_data["storage_url"] = hidden_params["storage_url"]
update_data["storage_url"] = hidden_params["storage_url"]
db_data.update(storage_metadata)
update_data.update(storage_metadata)
verbose_logger.debug(
f"Storage metadata: storage_backend={db_data.get('storage_backend')}, "
@ -314,6 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
request_tags: Sequence[str] | None = None,
persist_attribution: bool = False,
create_if_missing: bool = True,
batch_processed: bool = False,
) -> None:
"""Persist a managed object row, caching it and upserting it in the DB.
@ -328,6 +339,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
row absent from the table is left absent rather than created with the
observer as its creator, because created_by and team_id are written from
whoever calls the create branch.
batch_processed is set by callers that have already billed the batch
themselves, so CheckBatchCost skips the row instead of billing it twice.
It is written only in the upsert create branch.
"""
verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
litellm_managed_object = LiteLLM_ManagedObjectTable(
@ -379,6 +394,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
**attribution_columns,
"batch_processed": batch_processed,
},
"update": update_columns,
},
@ -1343,6 +1359,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> LLMResponseTypes:
if isinstance(response, LiteLLMBatch):
decoded_batch_id: Final = _is_base64_encoded_unified_file_id(response.id)
if decoded_batch_id and is_litellm_executed_batch(decoded_batch_id):
return response
## Check if unified_file_id is in the response
unified_file_id = response._hidden_params.get("unified_file_id") # managed file id
unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id
@ -1794,24 +1813,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Check if file deletion should be blocked due to batch references
await self._check_file_deletion_allowed(file_id)
# file_id = convert_b64_uid_to_unified_uid(file_id)
model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
if specific_model_file_id_mapping:
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
for model_id, model_file_id in specific_model_file_id_mapping.items():
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
await self._delete_storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
else:
await self._delete_provider_files(file_id, litellm_parent_otel_span, llm_router, data)
await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
@ -1820,16 +1826,53 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prom_logger.record_managed_file_deleted(result="success")
return FileDeleted(id=file_id, object="file", deleted=True)
async def _delete_storage_backend_content(self, storage_backend_name: str, storage_url: str) -> None:
try:
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
except ValueError as e:
raise HTTPException(status_code=400, detail=f"Cannot delete the stored file content: {e}") from e
await storage_backend.delete_file(storage_url)
async def _delete_provider_files(
self,
file_id: str,
litellm_parent_otel_span: Span | None,
llm_router: Router,
data: Mapping[str, object],
) -> None:
model_file_id_mapping: Final = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id)
if not specific_model_file_id_mapping:
return
filtered_data: Final = {
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
}
for model_id, model_file_id in specific_model_file_id_mapping.items():
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**filtered_data,
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
async def afile_content(
self,
file_id: str,
litellm_parent_otel_span: Optional[Span],
llm_router: Router,
**data: Dict,
) -> "HttpxBinaryResponseContent":
) -> HttpxBinaryResponseContent:
"""
Get the content of a file from first model that has it
"""
managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
return await self._storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
model_file_id_mapping = data.pop("model_file_id_mapping", None)
model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping(
[file_id], litellm_parent_otel_span
@ -1859,6 +1902,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
else:
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
async def _storage_backend_content(self, storage_backend_name: str, storage_url: str) -> HttpxBinaryResponseContent:
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
content: Final = await storage_backend.download_file(storage_url)
return HttpxBinaryResponseContent(response=httpx.Response(status_code=httpx.codes.OK, content=content))
async def _convert_storage_files_to_base64(
self,
messages: List[AllMessageValues],
@ -1889,16 +1937,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# File is stored in a storage backend, download and convert to base64
try:
from litellm.llms.base_llm.files.storage_backend_factory import (
get_storage_backend,
)
storage_backend_name = db_file.storage_backend
storage_url = db_file.storage_url
# Get storage backend (uses same env vars as callback)
try:
storage_backend = get_storage_backend(storage_backend_name)
storage_backend = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
except ValueError as e:
verbose_logger.warning(
f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}"

4
enterprise/pyproject.toml
View file

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

5
litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_autorouter_savings_estimate_coverage/migration.sql Normal file
View file

@ -0,0 +1,5 @@
ALTER TABLE "LiteLLM_AutoRouterSession"
ADD COLUMN IF NOT EXISTS "savings_estimated_turns" INTEGER NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}';

36
litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql Normal file
View file

@ -0,0 +1,36 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineComparison" (
"scope" TEXT PRIMARY KEY,
"api_key" TEXT NOT NULL,
"session_id" TEXT NOT NULL,
"router_name" TEXT NOT NULL,
"initial_equivalent" BOOLEAN NOT NULL,
"revision" BIGINT NOT NULL DEFAULT 0,
"published_revision" BIGINT NOT NULL DEFAULT 0,
"history" TEXT,
"attempted_at" TIMESTAMP(3),
"retired" BOOLEAN NOT NULL DEFAULT FALSE,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_scope"
ON "LiteLLM_AutoRouterBaselineComparison" ("api_key", "session_id", "router_name");
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_updated"
ON "LiteLLM_AutoRouterBaselineComparison" ("updated_at");
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_dirty"
ON "LiteLLM_AutoRouterBaselineComparison" ("attempted_at", "updated_at", "scope")
WHERE NOT "retired" AND "revision" <> "published_revision";
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineObservation" (
"request_id" TEXT PRIMARY KEY,
"scope" TEXT NOT NULL,
"started_at" DOUBLE PRECISION NOT NULL,
"revision" BIGINT NOT NULL,
"data" TEXT NOT NULL,
"publication" TEXT,
"conflicted" BOOLEAN NOT NULL DEFAULT FALSE
);
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_order"
ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "started_at", "request_id");
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_revision"
ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "revision", "started_at");

8
litellm-proxy-extras/litellm_proxy_extras/migrations/20260918000000_add_managed_file_content_table/migration.sql Normal file
View file

@ -0,0 +1,8 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedFileContentTable" (
"id" TEXT NOT NULL,
"content" BYTEA NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_ManagedFileContentTable_pkey" PRIMARY KEY ("id")
);

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

@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([team_id, created_at(sort: Desc)])
}
model LiteLLM_ManagedFileContentTable {
id String @id @default(uuid())
content Bytes
created_at DateTime @default(now())
}
model LiteLLM_ManagedVectorStoreTable {
id String @id @default(uuid())
unified_resource_id String @unique // The base64 encoded unified vector store ID
@ -1545,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession {
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
}
model LiteLLM_AutoRouterBaselineComparison {
scope String @id
api_key String
session_id String
router_name String
initial_equivalent Boolean
revision BigInt @default(0)
published_revision BigInt @default(0)
history String?
attempted_at DateTime?
retired Boolean @default(false)
updated_at DateTime @default(now())
@@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope")
@@index([updated_at], map: "idx_autorouter_baseline_updated")
}
model LiteLLM_AutoRouterBaselineObservation {
request_id String @id
scope String
started_at Float
revision BigInt
data String
publication String?
conflicted Boolean @default(false)
@@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order")
@@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision")
}
model LiteLLM_AutoRouterSession {
api_key String
session_id String
@ -1571,6 +1607,10 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
savings_estimated_baseline_models Json @default("{}")
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")

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

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

13
litellm-rust/Cargo.lock generated
View file

@ -1960,6 +1960,7 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-types",
"litellm-auth",
"litellm-http",
"moka",
"reqwest 0.12.28",
"serde_json",
@ -1976,6 +1977,7 @@ dependencies = [
"azure_identity",
"litellm-auth",
"moka",
"rstest",
"serde_json",
"sha2 0.10.9",
"strum",
@ -2028,7 +2030,7 @@ dependencies = [
]
[[package]]
name = "litellm-callbacks-legacy"
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
dependencies = [
"litellm-auth",
@ -2137,11 +2139,15 @@ version = "0.1.0"
dependencies = [
"http 1.4.2",
"hyper-util",
"litellm-core-utils",
"reqwest 0.12.28",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"veil",
"webpki-roots",
]
@ -2170,6 +2176,7 @@ dependencies = [
"serde_json",
"serde_path_to_error",
"serde_with",
"strum",
"thiserror 2.0.19",
"time",
"tokio",
@ -2185,8 +2192,9 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-gcp",
"litellm-callbacks-legacy",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
"litellm-host-python",
"litellm-http",
"litellm-llms",
@ -3038,7 +3046,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"h2 0.4.15",

4
litellm-rust/Cargo.toml
View file

@ -11,7 +11,7 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-core = { path = "crates/core" }
litellm-host = { path = "crates/host" }
litellm-callbacks-legacy = { path = "crates/callbacks-legacy" }
litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
litellm-framing = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
@ -34,7 +34,7 @@ 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", "multipart", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }

10
litellm-rust/clippy.toml Normal file
View file

@ -0,0 +1,10 @@
# The Tokio runtime is reached only through `host-python/src/execution.rs`, whose fork gate
# must see every entry. Going around it makes a fork-after-use hang instead of raising.
disallowed-methods = [
{ path = "pyo3_async_runtimes::tokio::get_runtime", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py_with_locals", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" },
]

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

@ -7,6 +7,7 @@ repository.workspace = true
[dependencies]
litellm-auth.workspace = true
litellm-http.workspace = true
moka = { workspace = true, features = ["sync"] }
serde_json.workspace = true

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

@ -20,8 +20,7 @@ use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
SIGV4_COMPUTED_HEADER_NAMES,
DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES,
};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
@ -451,11 +450,12 @@ pub fn is_sigv4_computed_header(name: &str) -> bool {
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
}
pub fn sign_bedrock_post(
pub fn sign_post(
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
service: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, Error> {
@ -463,7 +463,7 @@ pub fn sign_bedrock_post(
let params = v4::SigningParams::builder()
.identity(&identity)
.region(region)
.name(BEDROCK_SERVICE)
.name(service)
.time(signing_time)
.settings(SigningSettings::default())
.build()
@ -534,22 +534,28 @@ fn is_bedrock_region(value: &str) -> bool {
.all(|char| char.is_ascii_alphanumeric() || char == '-')
}
/// The region a caller configured: `aws_region_name`, then the model's own
/// region, then the environment. Each service decides what a missing one means.
pub fn resolve_aws_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
optional_params
.get("aws_region_name")
.and_then(Value::as_str)
.or(model_region)
.map(str::to_string)
.or_else(|| env_lookup(AWS_REGION_NAME))
.or_else(|| env_lookup(AWS_REGION))
}
pub fn resolve_bedrock_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
if let Some(region) = optional_params
.get("aws_region_name")
.and_then(Value::as_str)
{
return region.to_string();
}
if let Some(region) = model_region {
return region.to_string();
}
env_lookup(AWS_REGION_NAME)
.or_else(|| env_lookup(AWS_REGION))
resolve_aws_region(model_region, optional_params, env_lookup)
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
}
@ -609,11 +615,36 @@ pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option
#[cfg(test)]
mod tests {
use super::*;
use crate::constants::BEDROCK_SERVICE;
fn no_env(_: &str) -> Option<String> {
None
}
#[test]
fn a_region_comes_from_the_call_then_the_model_then_the_environment() {
let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]);
let region_name = |key: &str| (key == AWS_REGION_NAME).then(|| "ap-south-1".to_string());
let region = |key: &str| (key == AWS_REGION).then(|| "sa-east-1".to_string());
let resolved = [
resolve_aws_region(Some("us-east-2"), &params, &region_name),
resolve_aws_region(Some("us-east-2"), &Map::new(), &region_name),
resolve_aws_region(None, &Map::new(), &region_name),
resolve_aws_region(None, &Map::new(), &region),
resolve_aws_region(None, &Map::new(), &no_env),
];
assert_eq!(
resolved.map(|region| region.unwrap_or_else(|| "none".into())),
["eu-west-1", "us-east-2", "ap-south-1", "sa-east-1", "none"]
);
assert_eq!(
resolve_bedrock_region(None, &Map::new(), &no_env),
DEFAULT_BEDROCK_REGION
);
}
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
@ -811,11 +842,12 @@ mod tests {
None,
"test",
);
let signed = sign_bedrock_post(
let signed = sign_post(
&url,
&body,
&signable,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
SystemTime::UNIX_EPOCH,
)
@ -843,11 +875,12 @@ mod tests {
None,
"test",
);
let signed = sign_bedrock_post(
let signed = sign_post(
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
@ -878,11 +911,12 @@ mod tests {
None,
"test",
);
let signed = sign_bedrock_post(
let signed = sign_post(
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
@ -915,11 +949,12 @@ mod tests {
let url = format!(
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
);
let signed_headers = sign_bedrock_post(
let signed_headers = sign_post(
&url,
&body,
&headers,
region,
BEDROCK_SERVICE,
&credentials,
SystemTime::now(),
)?;

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

@ -1,6 +1,9 @@
mod aws;
pub mod constants;
mod error;
mod signer;
pub use aws::*;
pub use aws_credential_types::Credentials;
pub use error::Error;
pub use signer::SigV4Signer;

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

@ -0,0 +1,178 @@
use std::{collections::BTreeMap, time::SystemTime};
use aws_credential_types::Credentials;
use litellm_http::outbound::{RequestSigner, UnsignedRequest};
use serde_json::{Map, Value};
use crate::{
Error, aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_post,
};
#[derive(Clone, Debug)]
pub struct SigV4Signer {
region: String,
service: &'static str,
credentials: Credentials,
clock: fn() -> SystemTime,
}
impl SigV4Signer {
pub fn new(region: String, service: &'static str, credentials: Credentials) -> Self {
Self {
region,
service,
credentials,
clock: SystemTime::now,
}
}
pub fn with_clock(self, clock: fn() -> SystemTime) -> Self {
Self { clock, ..self }
}
pub async fn resolve(
region: String,
service: &'static str,
optional_params: &Map<String, Value>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Self, Error> {
let credentials = match host_supplied_credentials(optional_params) {
Some(credentials) => credentials,
None => {
resolve_credentials(aws_auth_config(optional_params, env_lookup), env_lookup)
.await?
}
};
Ok(Self::new(region, service, credentials))
}
}
impl RequestSigner for SigV4Signer {
fn sign(
&self,
request: UnsignedRequest<'_>,
) -> Result<Vec<(String, String)>, litellm_http::Error> {
if let Some((name, _)) = request
.headers
.iter()
.find(|(name, _)| is_sigv4_computed_header(name))
{
return Err(litellm_http::Error::ComputedHeader(name.clone()));
}
let headers: BTreeMap<String, String> = request.headers.iter().cloned().collect();
sign_post(
request.url,
request.body,
&aws_signature_headers(&headers),
&self.region,
self.service,
&self.credentials,
(self.clock)(),
)
.map(|signature| signature.into_iter().collect())
.map_err(|error| litellm_http::Error::Signature(error.to_string()))
}
}
#[cfg(test)]
mod tests {
use std::time::{Duration, UNIX_EPOCH};
use litellm_http::outbound::OutboundRequest;
use serde_json::json;
use super::*;
fn fixed_clock() -> SystemTime {
UNIX_EPOCH + Duration::from_secs(1_700_000_000)
}
fn signer(service: &'static str) -> SigV4Signer {
SigV4Signer::new(
"us-east-1".into(),
service,
Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
)
.with_clock(fixed_clock)
}
fn authorization(body: &Value, service: &'static str) -> String {
OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
body,
None,
&signer(service),
)
.unwrap()
.header("Authorization")
.unwrap()
.to_string()
}
#[test]
fn the_signature_verifies_against_the_bytes_that_are_sent() {
let sent = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
&json!({"Document": {"Bytes": "aGk="}}),
None,
&signer("textract"),
)
.unwrap();
let unsigned: BTreeMap<String, String> = sent
.headers()
.iter()
.filter(|(name, _)| !is_sigv4_computed_header(name))
.cloned()
.collect();
let recomputed = sign_post(
sent.url(),
sent.body(),
&aws_signature_headers(&unsigned),
"us-east-1",
"textract",
&Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
fixed_clock(),
)
.unwrap();
assert_eq!(
sent.header("Authorization"),
Some(recomputed["Authorization"].as_str())
);
}
#[test]
fn the_signature_depends_on_the_body_and_the_service() {
let original = authorization(&json!({"text": "card 4111"}), "textract");
assert_ne!(
original,
authorization(&json!({"text": "card [REDACTED]"}), "textract")
);
assert_ne!(
original,
authorization(&json!({"text": "card 4111"}), "bedrock")
);
assert!(original.contains("/us-east-1/textract/aws4_request"));
}
#[test]
fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() {
let error = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("authorization".into(), "Bearer caller".into())],
&json!({}),
None,
&signer("textract"),
)
.unwrap_err();
assert_eq!(
error,
litellm_http::Error::ComputedHeader("authorization".into())
);
}
}

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

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

48
litellm-rust/crates/auth-azure/src/types.rs
View file

@ -1,11 +1,10 @@
use serde_json::{Map, Value};
use std::collections::BTreeMap;
use strum::EnumString;
use litellm_auth::Error;
use litellm_auth::{
CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle,
CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle,
};
use serde_json::{Map, Value};
use strum::EnumString;
pub const DEFAULT_AZURE_SCOPE: &str = "https://cognitiveservices.azure.com/.default";
@ -52,6 +51,16 @@ pub struct AzureAuthInputs {
}
impl AzureAuthInputs {
pub fn or_configured_token_refresh(self, enabled: bool) -> Self {
if *self.enable_azure_ad_token_refresh.value() || !enabled {
return self;
}
Self {
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..self
}
}
#[cfg(test)]
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, Error> {
Self::from_sourced_optional_params(params, &BTreeMap::new())
@ -115,12 +124,12 @@ fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSourc
#[cfg(test)]
mod tests {
use serde_json::json;
use std::collections::BTreeMap;
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
use litellm_auth::{InputSource, Sourced};
use serde_json::json;
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
#[test]
fn selector_parsing_is_exact() {
@ -189,4 +198,29 @@ mod tests {
assert!(!debug.contains("token-value"));
assert!(!debug.contains("secret-value"));
}
#[rstest::rstest]
#[case::global_turns_refresh_on(json!({}), true, true, InputSource::Deployment)]
#[case::global_overrides_a_call_false_like_python(json!({"enable_azure_ad_token_refresh": false}), true, true, InputSource::Deployment)]
#[case::call_true_survives_a_global_false(json!({"enable_azure_ad_token_refresh": true}), false, true, InputSource::Request)]
#[case::both_off(json!({}), false, false, InputSource::Request)]
fn token_refresh_follows_the_configured_global(
#[case] params: serde_json::Value,
#[case] global: bool,
#[case] enabled: bool,
#[case] source: InputSource,
) {
let sources = BTreeMap::from([(
"enable_azure_ad_token_refresh".to_string(),
InputSource::Request,
)]);
let inputs =
AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources)
.unwrap()
.or_configured_token_refresh(global);
assert_eq!(
inputs.enable_azure_ad_token_refresh,
Sourced::new(enabled, source)
);
}
}

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

@ -1,17 +1,13 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc};
use gcp_auth::{CustomServiceAccount, TokenProvider};
use litellm_auth::{
CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential,
};
use moka::future::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use litellm_auth::http::apply_credential;
use litellm_auth::{CredentialPlacement, Error, InputSource, SecretValue, Sourced};
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS";
@ -45,6 +41,16 @@ impl VertexConfig {
})
}
pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self {
let configured =
|value: Option<&str>| value.filter(|value| !value.is_empty()).map(str::to_string);
Self {
project_id: self.project_id.or_else(|| configured(project_id)),
location: self.location.or_else(|| configured(location)),
..self
}
}
pub fn project_id(&self) -> Option<&str> {
self.project_id.as_deref()
}
@ -571,4 +577,29 @@ mod tests {
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}
#[test]
fn configured_defaults_sit_between_call_params_and_the_environment() {
let env = |name: &str| Some(format!("env-{name}"));
let from_config =
VertexConfig::default().or_configured(Some("global-project"), Some("global-location"));
assert_eq!(
get_vertex_ai_project(&from_config, &env).as_deref(),
Some("global-project")
);
assert_eq!(
get_vertex_ai_location(&from_config, &env).as_deref(),
Some("global-location")
);
let from_call =
config(json!({"vertex_project":"call-project","vertex_location":"call-location"}))
.or_configured(Some("global-project"), Some("global-location"));
assert_eq!(from_call.project_id(), Some("call-project"));
assert_eq!(from_call.location(), Some("call-location"));
let empty_global = VertexConfig::default().or_configured(Some(""), None);
assert_eq!(
get_vertex_ai_project(&empty_global, &env).as_deref(),
Some("env-VERTEXAI_PROJECT")
);
}
}

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

@ -40,13 +40,22 @@ pub fn apply_credential(
)
}
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
/// How the upstream call is authenticated. API-key strategies become headers
/// in `prepare`; SigV4 covers the serialized body, so it is applied where the
/// outbound request is built.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RequestAuth {
Header { name: &'static str, value: String },
Bearer { token: String },
AwsSigV4 { region: String },
Header {
name: &'static str,
value: String,
},
Bearer {
token: String,
},
AwsSigV4 {
region: String,
service: &'static str,
},
}
#[cfg(test)]

0
litellm-rust/crates/callbacks-legacy/AGENTS.md → litellm-rust/crates/callbacks-legacy-python/AGENTS.md
View file

2
litellm-rust/crates/callbacks-legacy/Cargo.toml → litellm-rust/crates/callbacks-legacy-python/Cargo.toml
View file

@ -1,5 +1,5 @@
[package]
name = "litellm-callbacks-legacy"
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
edition.workspace = true
license.workspace = true

0
litellm-rust/crates/callbacks-legacy/python_contract.json → litellm-rust/crates/callbacks-legacy-python/python_contract.json
View file

13
litellm-rust/crates/callbacks-legacy/src/adapter.rs → litellm-rust/crates/callbacks-legacy-python/src/adapter.rs
View file

@ -12,16 +12,16 @@ use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyList},
types::{PyDateTime, PyDict, PyList},
};
use serde_json::Value;
use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call,
legacy_python::Streaming,
prepare, setup,
finalize, is_internal_call, prepare,
python::Streaming,
setup,
};
/// What the legacy contract needs to know about the route it is logging.
@ -73,10 +73,7 @@ pub struct LegacyLogging {
}
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
py.import("datetime")?
.getattr("datetime")?
.call_method1("fromtimestamp", (epoch_seconds,))
.map(Bound::unbind)
PyDateTime::from_timestamp(py, epoch_seconds, None).map(|value| value.into_any().unbind())
}
fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {

0
litellm-rust/crates/callbacks-legacy/src/call.rs → litellm-rust/crates/callbacks-legacy-python/src/call.rs
View file

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

@ -6,8 +6,8 @@ use litellm_host::event::{RequestContext, WireRequest};
use litellm_host_python::to_py;
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
use crate::legacy_python::{Logging, Wrapper};
use crate::logger::PythonLogger;
use crate::python::{Logging, Wrapper};
pub trait LegacyCallbacks {
/// `Logging.update_from_kwargs`: what the logger is told about the request it is

0
litellm-rust/crates/callbacks-legacy/src/deferred.rs → litellm-rust/crates/callbacks-legacy-python/src/deferred.rs
View file

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

@ -13,9 +13,9 @@ mod adapter;
mod call;
mod callbacks;
mod deferred;
mod legacy_python;
mod logger;
mod preparation;
mod python;
#[cfg(test)]
#[path = "../tests/support.rs"]
mod test_support;

8
litellm-rust/crates/callbacks-legacy/src/logger.rs → litellm-rust/crates/callbacks-legacy-python/src/logger.rs
View file

@ -5,7 +5,7 @@ use pyo3::{
types::{PyDict, PyTuple},
};
use crate::legacy_python::{self, Wrapper};
use crate::python::{self, Wrapper};
/// The `Logging` instance one call fans out through.
pub struct PythonLogger {
@ -90,7 +90,7 @@ impl DeploymentHooks {
kwargs: &Py<PyDict>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
legacy_python::DeploymentHooks::BeforeDeploymentCall
python::DeploymentHooks::BeforeDeploymentCall
.call(py, (kwargs, call_type))
.map(Bound::unbind)
}
@ -101,7 +101,7 @@ impl DeploymentHooks {
response: &Option<Py<PyAny>>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
legacy_python::DeploymentHooks::AfterDeploymentSuccess
python::DeploymentHooks::AfterDeploymentSuccess
.call(py, (kwargs, response, call_type))
.map(Bound::unbind)
}
@ -112,7 +112,7 @@ impl DeploymentHooks {
error: &Py<PyBaseException>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
legacy_python::DeploymentHooks::AfterDeploymentFailure
python::DeploymentHooks::AfterDeploymentFailure
.call(py, (kwargs, error, call_type))
.map(Bound::unbind)
}

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

@ -3,7 +3,7 @@ use pyo3::{
types::{PyDict, PyList},
};
use crate::legacy_python::Wrapper;
use crate::python::Wrapper;
struct CredentialEntry<'py>(Bound<'py, PyAny>);

4
litellm-rust/crates/callbacks-legacy/src/legacy_python.rs → litellm-rust/crates/callbacks-legacy-python/src/python.rs
View file

@ -1,7 +1,7 @@
use pyo3::prelude::*;
use strum::{IntoStaticStr, VariantArray};
const MODULE: &str = "litellm.rust_bridge.legacy_callbacks";
const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python";
/// Every litellm Python internal the native call still borrows, grouped by the subsystem it
/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today
@ -9,7 +9,7 @@ const MODULE: &str = "litellm.rust_bridge.legacy_callbacks";
/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a
/// user's own callback is not borrowing and does not belong here.
///
/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and
/// `litellm/rust_bridge/callbacks_legacy_python.py` is the only Python module behind it, and
/// `python_contract.json` pins each function's parameters on both sides.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum LegacyPython {

0
litellm-rust/crates/callbacks-legacy/tests/deferred.rs → litellm-rust/crates/callbacks-legacy-python/tests/deferred.rs
View file

0
litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs → litellm-rust/crates/callbacks-legacy-python/tests/deployment_hooks.rs
View file

0
litellm-rust/crates/callbacks-legacy/tests/payload.rs → litellm-rust/crates/callbacks-legacy-python/tests/payload.rs
View file

10
litellm-rust/crates/callbacks-legacy/tests/support.rs → litellm-rust/crates/callbacks-legacy-python/tests/support.rs
View file

@ -5,12 +5,12 @@ use pyo3::types::{PyDict, PyTuple};
use crate::{LegacyLogging, LegacySurface, PublicCall};
/// The parameters of every `legacy_callbacks` function, as the real module declares them.
/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
/// `tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
/// signatures, and [`namespace`] binds every fake call against it.
pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json");
/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests
/// Stand-ins for `callbacks_legacy_python`, the only Python module the crate calls. Tests
/// share one interpreter and run concurrently, so each fake is installed idempotently and
/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`).
/// Every fake is bound against the contract first, so a call the real module would reject
@ -23,10 +23,10 @@ import sys
import traceback
import types
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'):
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'):
sys.modules.setdefault(name, types.ModuleType(name))
legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python']
CONTRACT = json.loads(python_contract)

0
litellm-rust/crates/callbacks-legacy/tests/terminal.rs → litellm-rust/crates/callbacks-legacy-python/tests/terminal.rs
View file

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

@ -6,4 +6,5 @@ pub mod params;
pub mod prompt_templates;
pub mod secret_redaction;
pub mod serde_compat;
pub mod settings;
pub mod url_utils;

144
litellm-rust/crates/core-utils/src/settings.rs Normal file
View file

@ -0,0 +1,144 @@
use std::str::FromStr;
pub trait Lookup {
fn get(&self, name: &str) -> Option<String>;
fn truthy(&self, name: &str) -> Option<String> {
self.get(name).filter(|value| !value.is_empty())
}
fn enabled(&self, name: &str) -> Option<bool> {
self.get(name)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
.then_some(true)
}
fn parsed<T: FromStr>(&self, name: &str) -> Option<T>
where
Self: Sized,
{
self.get(name).and_then(|value| value.trim().parse().ok())
}
}
impl<F: Fn(&str) -> Option<String>> Lookup for F {
fn get(&self, name: &str) -> Option<String> {
self(name)
}
}
pub struct ProcessEnvironment;
impl Lookup for ProcessEnvironment {
fn get(&self, name: &str) -> Option<String> {
std::env::var(name).ok()
}
}
pub trait Layer: Default {
fn or(self, lower: Self) -> Self;
}
pub fn merge<L: Layer>(highest_precedence_first: impl IntoIterator<Item = L>) -> L {
highest_precedence_first
.into_iter()
.reduce(L::or)
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
fn env_of(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option<String> {
move |name| {
values
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string())
}
}
#[test]
fn get_keeps_a_present_empty_value_like_os_getenv_with_a_fallback() {
let env = env_of(&[("EMPTY", "")]);
assert_eq!(env.get("EMPTY"), Some(String::new()));
assert_eq!(env.get("ABSENT"), None);
}
#[test]
fn truthy_drops_an_empty_value_like_a_python_or_chain() {
let env = env_of(&[("EMPTY", ""), ("SET", "value")]);
assert_eq!(env.truthy("EMPTY"), None);
assert_eq!(env.truthy("SET").as_deref(), Some("value"));
}
#[test]
fn enabled_only_switches_on_for_true_and_never_forces_off() {
let env = env_of(&[
("LOWER", "true"),
("PADDED", " True "),
("OFF", "false"),
("ONE", "1"),
]);
assert_eq!(env.enabled("LOWER"), Some(true));
assert_eq!(env.enabled("PADDED"), Some(true));
assert_eq!(env.enabled("OFF"), None);
assert_eq!(env.enabled("ONE"), None);
assert_eq!(env.enabled("ABSENT"), None);
}
#[test]
fn parsed_trims_and_skips_values_that_do_not_parse() {
let env = env_of(&[("PADDED", " 45 "), ("WORD", "soon"), ("FRACTION", "0.5")]);
assert_eq!(env.parsed::<u32>("PADDED"), Some(45));
assert_eq!(env.parsed::<u32>("WORD"), None);
assert_eq!(env.parsed::<f64>("FRACTION"), Some(0.5));
assert_eq!(env.parsed::<u32>("ABSENT"), None);
}
#[derive(Debug, Default, PartialEq)]
struct Pair {
first: Option<u8>,
second: Option<u8>,
}
impl Layer for Pair {
fn or(self, lower: Self) -> Self {
Self {
first: self.first.or(lower.first),
second: self.second.or(lower.second),
}
}
}
#[test]
fn merge_takes_each_field_from_the_highest_layer_that_sets_it() {
let merged = merge([
Pair {
first: Some(1),
second: None,
},
Pair {
first: Some(2),
second: Some(2),
},
Pair {
first: Some(3),
second: Some(3),
},
]);
assert_eq!(
merged,
Pair {
first: Some(1),
second: Some(2),
}
);
}
#[test]
fn merging_no_layers_yields_the_empty_layer() {
assert_eq!(merge(Vec::<Pair>::new()), Pair::default());
}
}

7
litellm-rust/crates/core/AGENTS.md
View file

@ -5,10 +5,11 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve
Each crate mirrors one top-level Python package, so a Rust path reads as its Python path with the crate name in place of the package directory. Dependencies only point down:
- `litellm-types` mirrors `litellm/types/`: pure serde data, no I/O
- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments), no network I/O
- `litellm-llms` mirrors `litellm/llms/`: `base_llm/<api>/transformation.rs`, `<provider>/<api>/transformation.rs`, and `custom_httpx/` (HTTP helpers and the OCR request handler)
- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments, settings lookup and layer merge), no network I/O
- `litellm-http` is Rust-only and route-neutral: settings resolution, the pooled `reqwest` clients, TLS, proxies, the SSRF-safe media fetcher, request and header helpers, and transport errors. Python's `litellm/llms/custom_httpx/` is split by responsibility instead of mirrored: its transport half lives here, its OCR handler in `litellm-llms`
- `litellm-llms` mirrors `litellm/llms/`: `base_llm/<api>/transformation.rs`, `<provider>/<api>/transformation.rs`, and `base_llm/ocr/handler.rs` (the OCR request handler)
- `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::custom_httpx::llm_http_handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business.

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

@ -15,6 +15,7 @@ futures-util.workspace = true
base64.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
@ -36,7 +37,6 @@ veil.workspace = true
[dev-dependencies]
litellm-auth-gcp.workspace = true
litellm-http.workspace = true
litellm-llms = { workspace = true, features = ["test-support"] }
rstest.workspace = true
rstest_reuse.workspace = true

6
litellm-rust/crates/core/src/audio_transcription/error.rs
View file

@ -20,9 +20,11 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,4 +1,4 @@
use litellm_llms::custom_httpx::http_handler::{http_request, truncate_error_body};
use litellm_http::request::truncate_error_body;
use serde_json::Value;
use super::{Error, client::http_client};
@ -7,34 +7,29 @@ use crate::audio_transcription::types::ProviderAudioTranscriptionRequest;
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::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
let response = crate::outbound::outbound_request::<Error>(
&request.auth,
request.url.clone(),
request.upstream_headers.clone(),
&request.body,
request.timeout,
&request.optional_params,
)
.await?
.send(http_client())
.await
.map_err(|error| {
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
let status = response.status();
let text = response.text().await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
if !status.is_success() {
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
return Err(Error::Transport(litellm_http::transport::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}")))?;
@ -43,33 +38,3 @@ pub async fn execute_audio_transcription_provider_call(
.transform_audio_transcription_response(&request.model, response_json)?
.into_json())
}
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::{collections::BTreeMap, time::SystemTime};
use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post};
use litellm_llms::base_llm::audio_transcription::transformation::AudioTranscriptionAuth;
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())
}

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

@ -1,10 +1,8 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_http::request::{has_header, string_headers};
use litellm_llms::{
base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
},
base_llm::audio_transcription::transformation::{BaseAudioTranscriptionConfig, RequestAuth},
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
custom_httpx::http_handler::{has_header, string_headers},
};
use super::Error;
@ -43,11 +41,14 @@ pub fn prepare_audio_transcription_provider_call(
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}")));
match &auth {
RequestAuth::Bearer { token } if !has_header(&headers, "authorization") => {
headers.push(("Authorization".to_string(), format!("Bearer {token}")));
}
RequestAuth::Header { name, value } if !has_header(&headers, name) => {
headers.push(((*name).to_string(), value.clone()));
}
RequestAuth::Bearer { .. } | RequestAuth::Header { .. } | RequestAuth::AwsSigV4 { .. } => {}
}
if !has_header(&headers, "content-type") {
headers.push(("Content-Type".to_string(), "application/json".to_string()));

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

@ -1,7 +1,7 @@
use std::time::Duration;
use litellm_llms::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
BaseAudioTranscriptionConfig, RequestAuth,
};
use serde_json::{Map, Value};
@ -24,7 +24,7 @@ pub struct ProviderAudioTranscriptionRequest {
pub url: String,
pub body: Value,
pub upstream_headers: Vec<(String, String)>,
pub auth: AudioTranscriptionAuth,
pub auth: RequestAuth,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}

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

@ -1,8 +1,8 @@
use litellm_http::request::string_headers as shared_string_headers;
use litellm_llms::{
anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
base_llm::chat::transformation::BaseConfig,
bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
custom_httpx::http_handler::string_headers as shared_string_headers,
};
use serde_json::{Map, Value};

6
litellm-rust/crates/core/src/chat_completions/error.rs
View file

@ -20,9 +20,11 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,7 +1,5 @@
use litellm_llms::{
base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData},
custom_httpx::http_handler::{http_request, truncate_error_body},
};
use litellm_http::{outbound::OutboundRequest, request::truncate_error_body};
use litellm_llms::base_llm::chat::transformation::ProviderChatResponseData;
use litellm_types::utils::ChatCompletionsResponse;
use serde_json::Value;
@ -14,50 +12,29 @@ pub(super) async fn execute_chat_completions_provider_call(
request: ResolvedChatCompletionsRequest<'_>,
) -> Result<ChatCompletionsResponse, Error> {
let request = prepare_provider_request(request)?;
let body = serde_json::to_vec(&request.body).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize chat completions request: {err}"
))
})?;
let headers = signed_headers(&request, &body).await?;
let outbound = outbound_request(&request).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(|err| {
let response = outbound.send(http_client()).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() {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect(
err.to_string(),
))
Error::Transport(litellm_http::transport::Error::Connect(err.to_string()))
} else {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
}
})?;
let status = response.status();
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
})?;
if !status.is_success() {
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
return Err(Error::Transport(litellm_http::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
}
let body: Value = serde_json::from_str(&text).map_err(|err| {
@ -82,64 +59,29 @@ pub(super) async fn execute_chat_completions_provider_call(
pub(super) fn as_response_error(err: Error) -> Error {
match err {
already @ (Error::InvalidResponse(_)
| Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
..
})) => already,
| Error::Transport(litellm_http::transport::Error::Http { .. })) => already,
other => Error::InvalidResponse(other.to_string()),
}
}
pub(super) async fn signed_headers(
pub(super) async fn outbound_request(
request: &ProviderChatCompletionsRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::{collections::BTreeMap, time::SystemTime};
use litellm_auth_aws::{
aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
};
let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else {
return Ok(request.upstream_headers.clone());
};
// Reattaching a header the signer also emits would put both copies on the
// wire, and Bedrock rejects that pair. Python instead drops the caller's
// copy and prefers a forwarded Authorization over the signature, so leave
// the request to Python rather than serving it a different way here.
if request
.upstream_headers
.iter()
.any(|(name, _)| is_sigv4_computed_header(name))
{
return Err(Error::Unsupported(
"request forwards a header AWS SigV4 computes",
));
}
let env_lookup = |key: &str| std::env::var(key).ok();
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
// A host with its own resolution chain hands the result down; only fall
// back to deriving credentials here when it supplied none.
let credentials = match host_supplied_credentials(&request.optional_params) {
Some(credentials) => credentials,
None => {
resolve_credentials(
aws_auth_config(&request.optional_params, &env_lookup),
&env_lookup,
)
.await?
) -> Result<OutboundRequest, Error> {
crate::outbound::outbound_request(
&request.auth,
request.url.clone(),
request.upstream_headers.clone(),
&request.body,
request.timeout,
&request.optional_params,
)
.await
.map_err(|error| match error {
// Python drops the caller's copy and prefers a forwarded Authorization
// over the signature, so leave the request to it.
Error::Http(litellm_http::Error::ComputedHeader(_)) => {
Error::Unsupported("request forwards a header AWS SigV4 computes")
}
};
let signature = sign_bedrock_post(
&request.url,
body,
&aws_signature_headers(&unsigned),
region,
&credentials,
SystemTime::now(),
)?;
// Every original header goes back on the wire alongside the computed ones,
// as Python reattaches them. The guard above already rejected the names
// that would collide, so no name appears twice.
Ok(unsigned.into_iter().chain(signature).collect())
other => other,
})
}

14
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -1,8 +1,6 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_llms::{
base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth},
custom_httpx::http_handler::has_header,
};
use litellm_http::request::has_header;
use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth};
use litellm_types::llms::openai::ChatMessage;
use serde_json::Value;
@ -69,7 +67,7 @@ fn validate_environment(
request: &ResolvedChatCompletionsRequest<'_>,
model: &str,
config: &dyn BaseConfig,
) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> {
) -> Result<(Vec<(String, String)>, RequestAuth), Error> {
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers(request.extra_headers.clone())?;
let auth = config.auth(
@ -79,7 +77,7 @@ fn validate_environment(
&env_lookup,
)?;
match &auth {
ChatCompletionsAuth::Header { name, value } => {
RequestAuth::Header { name, value } => {
// The deployment's credential replaces whatever the caller forwarded
// under the same name, mirroring Python's
// `{**headers, **anthropic_headers}`: letting a request header win
@ -94,7 +92,7 @@ fn validate_environment(
headers.push(((*name).to_string(), value.clone()));
}
}
ChatCompletionsAuth::Bearer { token } => {
RequestAuth::Bearer { token } => {
// Bedrock's `get_request_headers` assigns `headers["Authorization"]`
// unconditionally once a bearer token resolves, so the deployment's
// identity outranks whatever the caller forwarded. Keeping the
@ -107,7 +105,7 @@ fn validate_environment(
headers.push(("authorization".to_string(), format!("Bearer {token}")));
}
// SigV4 signs the serialized body, so the handler adds its headers.
ChatCompletionsAuth::AwsSigV4 { .. } => {}
RequestAuth::AwsSigV4 { .. } => {}
}
for (name, value) in config.default_headers() {

45
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -1,4 +1,4 @@
use litellm_llms::base_llm::chat::transformation::ChatCompletionsAuth;
use litellm_llms::base_llm::chat::transformation::RequestAuth;
use serde_json::{Map, Value, json};
use super::{
@ -90,7 +90,7 @@ fn adds_the_auth_and_default_headers() {
);
assert!(matches!(
prepared.auth,
ChatCompletionsAuth::Header {
RequestAuth::Header {
name: "x-api-key",
..
}
@ -265,7 +265,7 @@ fn rejects_non_string_extra_headers() {
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
decline(call),
Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError {
Error::Headers(litellm_http::request::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
@ -289,8 +289,9 @@ fn prepares_a_bedrock_call_without_resolving_credentials() {
);
assert_eq!(
prepared.auth,
ChatCompletionsAuth::AwsSigV4 {
region: "us-east-1".to_string()
RequestAuth::AwsSigV4 {
region: "us-east-1".to_string(),
service: "bedrock",
}
);
// SigV4 signs the serialized body, so prepare must not have added an
@ -326,15 +327,14 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
json!("abc-123"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
let signed = super::handler::outbound_request(&prepared)
.await
.expect("signs");
let authorization = signed
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.clone())
.expect("carries an authorization header");
.header("authorization")
.expect("carries an authorization header")
.to_string();
assert!(
authorization.starts_with("AWS4-HMAC-SHA256"),
"expected a SigV4 signature, got {authorization}"
@ -346,6 +346,7 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
// It still goes on the wire, it is just not part of the signature.
assert!(
signed
.headers()
.iter()
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
"forwarded header was dropped instead of reattached"
@ -376,7 +377,7 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() {
call.api_key = None;
call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
let error = super::handler::outbound_request(&prepared)
.await
.expect_err("{forwarded} should decline instead of being signed");
assert!(
@ -466,7 +467,7 @@ fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
.expect("prepares");
assert_eq!(
prepared.auth,
ChatCompletionsAuth::Bearer {
RequestAuth::Bearer {
token: "sk-test".to_string()
}
);
@ -771,10 +772,7 @@ mod round_trip {
assert!(
matches!(
err,
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status: 429,
..
})
Error::Transport(litellm_http::transport::Error::Http { status: 429, .. })
),
"expected a 429, got {err:?}"
);
@ -801,7 +799,7 @@ mod round_trip {
assert!(
matches!(
err,
Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_))
Error::Transport(litellm_http::transport::Error::Connect(_))
),
"expected a pre-send connect failure, got {err:?}"
);
@ -825,16 +823,11 @@ mod round_trip {
}
// An upstream status is already unambiguous, so it survives intact.
assert!(matches!(
as_response_error(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: 500,
body: "boom".to_string()
}
)),
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
as_response_error(Error::Transport(litellm_http::transport::Error::Http {
status: 500,
..
})
body: "boom".to_string()
})),
Error::Transport(litellm_http::transport::Error::Http { status: 500, .. })
));
}
}

4
litellm-rust/crates/core/src/chat_completions/types.rs
View file

@ -1,6 +1,6 @@
use std::time::Duration;
use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth};
use litellm_types::llms::openai::ChatMessage;
use serde_json::{Map, Value};
@ -38,7 +38,7 @@ pub struct ProviderChatCompletionsRequest {
pub url: String,
pub body: Value,
pub upstream_headers: Vec<(String, String)>,
pub auth: ChatCompletionsAuth,
pub auth: RequestAuth,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}

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

@ -4,6 +4,7 @@ pub mod constants;
pub mod error;
pub mod messages;
pub mod ocr;
mod outbound;
pub mod responses;
pub use error::Error;

6
litellm-rust/crates/core/src/messages/common_utils.rs
View file

@ -1,11 +1,9 @@
pub(super) use litellm_llms::custom_httpx::http_handler::{
has_bearer_auth, has_header, truncate_error_body,
};
use litellm_http::request::string_headers as shared_string_headers;
pub(super) use litellm_http::request::{has_bearer_auth, has_header, truncate_error_body};
use litellm_llms::{
anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig,
custom_httpx::http_handler::string_headers as shared_string_headers,
};
use serde_json::{Map, Value};

4
litellm-rust/crates/core/src/messages/error.rs
View file

@ -15,9 +15,9 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
}
impl From<LlmError> for Error {

6
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -1,9 +1,7 @@
use std::time::Duration;
use litellm_llms::{
base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig,
custom_httpx::{http_handler::http_request, transport::Error as TransportError},
};
use litellm_http::{request::http_request, transport::Error as TransportError};
use litellm_llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig;
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use serde_json::Value;

4
litellm-rust/crates/core/src/messages/tests.rs
View file

@ -82,7 +82,7 @@ fn string_headers_rejects_non_string_values() {
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError {
Error::Headers(litellm_http::request::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
@ -432,7 +432,7 @@ async fn messages_maps_provider_error_status_to_http_error() {
assert!(matches!(
err,
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { status: 401, .. })
Error::Transport(litellm_http::transport::Error::Http { status: 401, .. })
));
}

16
litellm-rust/crates/core/src/ocr/arguments.rs
View file

@ -15,6 +15,18 @@ const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[
"azure_federated_token_file",
"enable_azure_ad_token_refresh",
];
const AWS_AUTH_OPTION_FIELDS: &[&str] = &[
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"aws_region_name",
"aws_session_name",
"aws_profile_name",
"aws_role_name",
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_external_id",
];
const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[
"vertex_credentials",
"vertex_ai_credentials",
@ -35,6 +47,7 @@ pub fn consumed_optional_param_names(
let (model, config) = resolve_provider_config(model, custom_llm_provider)?;
let provider_fields = config.get_supported_ocr_params(&model);
let auth_fields: &[&str] = match config {
OcrConfigKind::AwsTextract | OcrConfigKind::AwsTextractAnalyze => AWS_AUTH_OPTION_FIELDS,
OcrConfigKind::AzureAi
| OcrConfigKind::AzureDocumentIntelligence
| OcrConfigKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS,
@ -57,6 +70,9 @@ pub(crate) fn is_secret_param(name: &str) -> bool {
| "azure_federated_token_file"
| "vertex_credentials"
| "vertex_ai_credentials"
| "aws_secret_access_key"
| "aws_session_token"
| "aws_web_identity_token"
)
}

5
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -1,6 +1,5 @@
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::OcrClient,
use litellm_llms::base_llm::ocr::{
error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse,
};
use crate::ocr::{

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

@ -1,12 +1,10 @@
use futures_util::future::BoxFuture;
use litellm_auth::SecretValue;
use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest};
use litellm_llms::{
base_llm::ocr::{
error::Error,
transformation::{LiteLLMOcrResponse, PreparedOcrRequest},
},
custom_httpx::llm_http_handler::{CallHooks, OcrClient},
use litellm_llms::base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::{LiteLLMOcrResponse, PreparedOcrRequest},
};
use serde_json::Value;
@ -24,7 +22,7 @@ pub(crate) async fn perform_ocr_request(
) -> Result<LiteLLMOcrResponse, Error> {
request.response_format()?;
let config = request.config;
let request = prepare_request(request, caller_document);
let request = prepare_request(request, caller_document, client);
let hooks = OcrCallHooks::new(host.clone(), &request, config);
config.ocr(client, &request, &hooks).await
}

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

@ -8,6 +8,10 @@ pub mod route;
pub mod types;
pub mod wire;
#[cfg(test)]
#[path = "../../tests/aws_textract_ocr.rs"]
mod aws_textract_tests;
#[cfg(test)]
#[path = "../../tests/azure_ai_ocr.rs"]
mod azure_ai_tests;

44
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -1,6 +1,7 @@
use litellm_auth::{InputSource, SecretValue, Sourced};
use litellm_llms::base_llm::ocr::transformation::{
OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env,
use litellm_llms::base_llm::ocr::{
handler::OcrClient,
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest},
};
use super::provider_config::OcrProvider;
@ -9,26 +10,34 @@ use crate::ocr::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
pub(crate) fn prepare_request(
request: ResolvedOcrRequest,
caller_document: bool,
client: &OcrClient,
) -> PreparedOcrRequest {
let credentials = request.credentials.clone();
let api_base_env = match request.config.provider() {
OcrProvider::Mistral => Some("MISTRAL_API_BASE"),
OcrProvider::AzureAi => Some("AZURE_AI_API_BASE"),
OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => None,
let (preferred_api_key_env, api_base_env) = match request.config.provider() {
OcrProvider::Mistral => (
Some("MISTRAL_AZURE_API_KEY"),
Some("MISTRAL_AZURE_API_BASE"),
),
OcrProvider::AzureAi => (None, Some("AZURE_AI_API_BASE")),
OcrProvider::AwsTextract
| OcrProvider::Cohere
| OcrProvider::Reducto
| OcrProvider::VertexAi => (None, None),
};
let secret = |name: &str| client.secrets().truthy(name);
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
credentials.api_key.clone().or_else(|| {
request
.config
.get_api_key_env_var()
.and_then(credential_env)
preferred_api_key_env
.into_iter()
.chain(request.config.get_api_key_env_var())
.find_map(secret)
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
})
});
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
credentials.api_base.clone().or_else(|| {
api_base_env
.and_then(credential_env)
.and_then(secret)
.map(|value| Sourced::new(value, InputSource::Environment))
})
});
@ -51,7 +60,12 @@ pub(crate) fn prepare_request(
PreparedOcrRequest {
model,
document,
connection: OcrConnection::new(resolved, transport),
connection: OcrConnection::new(
resolved,
transport,
client.settings().clone(),
client.secrets().clone(),
),
caller_document,
optional_params,
input_sources,
@ -61,7 +75,11 @@ pub(crate) fn prepare_request(
#[cfg(test)]
pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest {
prepare_request(request, true)
prepare_request(
request,
true,
&OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()),
)
}
#[cfg(test)]

40
litellm-rust/crates/core/src/ocr/provider_config.rs
View file

@ -1,5 +1,9 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_llms::{
aws_textract::ocr::{
analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation,
transformation::TextractDetectTextConfig,
},
azure_ai::ocr::{
cohere_parse_transformation::AzureAICohereParseConfig,
document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig,
@ -7,13 +11,13 @@ use litellm_llms::{
},
base_llm::ocr::{
error::Error,
handler::{self, CallHooks, OcrClient},
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument,
PreparedOcrRequest, ResolvedOcrCredentials,
},
},
cohere::ocr::transformation::CohereParseConfig,
custom_httpx::llm_http_handler::{self, CallHooks, OcrClient},
mistral::ocr::transformation::MistralOcrConfig,
reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config},
vertex_ai::ocr::{
@ -25,6 +29,14 @@ use strum::{EnumString, IntoStaticStr};
macro_rules! with_config {
($kind:expr, $config:ident => $body:expr) => {
match $kind {
OcrConfigKind::AwsTextract => {
let $config = TextractDetectTextConfig;
$body
}
OcrConfigKind::AwsTextractAnalyze => {
let $config = TextractAnalyzeDocumentConfig;
$body
}
OcrConfigKind::Cohere => {
let $config = CohereParseConfig;
$body
@ -67,6 +79,8 @@ macro_rules! with_config {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum OcrConfigKind {
AwsTextract,
AwsTextractAnalyze,
Cohere,
Mistral,
AzureAi,
@ -81,6 +95,7 @@ pub(crate) enum OcrConfigKind {
impl OcrConfigKind {
pub(crate) const fn provider(self) -> OcrProvider {
match self {
Self::AwsTextract | Self::AwsTextractAnalyze => OcrProvider::AwsTextract,
Self::Cohere => OcrProvider::Cohere,
Self::Mistral => OcrProvider::Mistral,
Self::AzureAi | Self::AzureCohere | Self::AzureDocumentIntelligence => {
@ -116,7 +131,7 @@ impl OcrConfigKind {
request: &PreparedOcrRequest,
hooks: &dyn CallHooks<Error>,
) -> Result<LiteLLMOcrResponse, Error> {
with_config!(self, config => llm_http_handler::ocr(&config, client, request, hooks).await)
with_config!(self, config => handler::ocr(&config, client, request, hooks).await)
}
}
@ -141,6 +156,7 @@ pub fn get_health_check_document(
#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, PartialEq, Eq)]
#[strum(serialize_all = "snake_case")]
pub(crate) enum OcrProvider {
AwsTextract,
Cohere,
Mistral,
AzureAi,
@ -162,6 +178,10 @@ pub(crate) fn resolve_provider_config(
.parse::<OcrProvider>()
.map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?;
let config = match ocr_provider {
OcrProvider::AwsTextract => match TextractOperation::from_model(provider.model)? {
TextractOperation::DetectDocumentText => OcrConfigKind::AwsTextract,
TextractOperation::AnalyzeDocument => OcrConfigKind::AwsTextractAnalyze,
},
OcrProvider::Cohere => OcrConfigKind::Cohere,
OcrProvider::Mistral => OcrConfigKind::Mistral,
OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => {
@ -419,6 +439,22 @@ mod tests {
}
#[rstest]
#[case::misspelled_operation("aws_textract/analyse-document")]
#[case::operation_name_from_the_api("aws_textract/AnalyzeDocument")]
fn textract_models_outside_its_two_operations_are_refused(#[case] model: &str) {
assert!(matches!(
resolve_provider_config(model, None),
Err(Error::InvalidModel {
provider: "aws_textract",
..
})
));
}
#[rstest]
#[case("aws_textract/detect-document-text", OcrConfigKind::AwsTextract)]
#[case("aws_textract/analyze-document", OcrConfigKind::AwsTextractAnalyze)]
#[case("aws_textract/Analyze-Document", OcrConfigKind::AwsTextractAnalyze)]
#[case("reducto/parse-legacy", OcrConfigKind::ReductoLegacy)]
#[case("reducto/future-parse-model", OcrConfigKind::ReductoV3)]
#[case("azure_ai/Cohere-parse-v5", OcrConfigKind::AzureCohere)]

5
litellm-rust/crates/core/src/ocr/route.rs
View file

@ -6,9 +6,8 @@ use litellm_host::{
machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute},
route::Route,
};
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::OcrClient,
use litellm_llms::base_llm::ocr::{
error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse,
};
use super::handler::perform_ocr_request;

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

@ -277,7 +277,7 @@ mod tests {
vec![("x-a".to_string(), "1".to_string())]
);
assert_eq!(request.transport.extra_headers_source, InputSource::Request);
assert_eq!(request.transport.timeout, Duration::from_secs(7));
assert_eq!(request.transport.timeout, Some(Duration::from_secs(7)));
assert_eq!(request.input_sources.len(), 2);
let defaulted = LiteLLMOcrRequest::from_inputs(

30
litellm-rust/crates/core/src/outbound.rs Normal file
View file

@ -0,0 +1,30 @@
use std::time::Duration;
use litellm_auth::RequestAuth;
use litellm_auth_aws::SigV4Signer;
use litellm_http::outbound::OutboundRequest;
use serde_json::{Map, Value};
/// Header credentials are already in `headers`; SigV4 is applied here, over the
/// bytes that are sent.
pub(crate) async fn outbound_request<E>(
auth: &RequestAuth,
url: String,
headers: Vec<(String, String)>,
body: &Value,
timeout: Option<Duration>,
optional_params: &Map<String, Value>,
) -> Result<OutboundRequest, E>
where
E: From<litellm_http::Error> + From<litellm_auth_aws::Error>,
{
let RequestAuth::AwsSigV4 { region, service } = auth else {
return Ok(OutboundRequest::json(url, headers, body, timeout)?);
};
let env_lookup = |key: &str| std::env::var(key).ok();
let signer =
SigV4Signer::resolve(region.clone(), service, optional_params, &env_lookup).await?;
Ok(OutboundRequest::signed_json(
url, headers, body, timeout, &signer,
)?)
}

4
litellm-rust/crates/core/src/responses/error.rs
View file

@ -11,7 +11,7 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
}

34
litellm-rust/crates/core/src/responses/websocket.rs
View file

@ -95,9 +95,7 @@ impl ResponsesWebSocketConnection {
timeout: Option<Duration>,
) -> Result<Self, Error> {
let mut request = url.into_client_request().map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
for (name, value) in headers {
let header_name = name
@ -110,7 +108,7 @@ impl ResponsesWebSocketConnection {
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
Error::Transport(litellm_http::transport::Error::Network(
"Responses WebSocket connection timed out".into(),
))
})?,
@ -118,14 +116,12 @@ impl ResponsesWebSocketConnection {
};
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
Error::Transport(litellm_http::transport::Error::Http {
status: response.status().as_u16(),
body: String::new(),
})
}
other => Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
other.to_string(),
)),
other => Error::Transport(litellm_http::transport::Error::Network(other.to_string())),
})?;
Ok(Self {
socket: Arc::new(Mutex::new(Some(socket))),
@ -135,16 +131,12 @@ impl ResponsesWebSocketConnection {
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(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Network(
"Responses WebSocket is closed".into(),
),
));
return Err(Error::Transport(litellm_http::transport::Error::Network(
"Responses WebSocket is closed".into(),
)));
};
socket.send(Message::Text(text)).await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})
}
@ -160,9 +152,9 @@ impl ResponsesWebSocketConnection {
.map_err(|error| Error::InvalidResponse(error.to_string())),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Network(error.to_string()),
)),
Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network(
error.to_string(),
))),
}
}
@ -170,9 +162,7 @@ impl ResponsesWebSocketConnection {
let mut socket = self.socket.lock().await;
if let Some(socket) = socket.as_mut() {
socket.close(None).await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
}
*socket = None;

193
litellm-rust/crates/core/tests/aws_textract_ocr.rs Normal file
View file

@ -0,0 +1,193 @@
use std::{collections::BTreeMap, time::SystemTime};
use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post};
use litellm_llms::base_llm::ocr::error::Error;
use serde_json::{Value, json};
use time::{PrimitiveDateTime, format_description};
use crate::ocr::{
route::LocalOcrHost,
test_support::{
MockResponse, header, mock_server, perform_ocr_with, request_body,
wire_request_with_document,
},
types::LiteLLMOcrRequest,
};
const ACCESS_KEY_ID: &str = "AKIDEXAMPLE";
const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY";
fn textract_request(base: &str) -> LiteLLMOcrRequest {
textract_request_for("aws_textract/detect-document-text", base)
}
fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest {
wire_request_with_document(
model,
&format!("{base}/"),
json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}),
json!({
"aws_access_key_id": ACCESS_KEY_ID,
"aws_secret_access_key": SECRET_ACCESS_KEY,
"aws_region_name": "eu-west-1"
}),
)
}
fn textract_response() -> MockResponse {
MockResponse::json(json!({
"DocumentMetadata": {"Pages": 1},
"Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}]
}))
}
/// Recomputes SigV4 over the bytes the server received, at the time the client claimed.
fn expected_authorization(url: &str, raw_request: &str) -> String {
let format =
format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z")
.unwrap();
let signed_at: SystemTime =
PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format)
.unwrap()
.assume_utc()
.into();
let headers: BTreeMap<String, String> = ["content-type", "x-amz-target"]
.into_iter()
.map(|name| {
(
name.to_string(),
header(raw_request, name).unwrap().to_string(),
)
})
.collect();
let body = raw_request.split_once("\r\n\r\n").unwrap().1;
sign_post(
url,
body.as_bytes(),
&aws_signature_headers(&headers),
"eu-west-1",
"textract",
&Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"),
signed_at,
)
.unwrap()["Authorization"]
.clone()
}
#[tokio::test]
async fn the_request_is_signed_for_textract_and_lines_become_the_page() {
let (base, seen, server) = mock_server(vec![textract_response()]).await;
let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base)))
.await
.unwrap();
server.await.unwrap();
let raw = seen.lock().unwrap()[0].clone();
assert_eq!(
header(&raw, "x-amz-target"),
Some("Textract.DetectDocumentText")
);
assert_eq!(
header(&raw, "content-type"),
Some("application/x-amz-json-1.1")
);
assert_eq!(
request_body(&raw),
json!({"Document": {"Bytes": "b3JpZ2luYWw="}})
);
assert_eq!(
header(&raw, "authorization"),
Some(expected_authorization(&format!("{base}/"), &raw).as_str())
);
assert_eq!(response.pages[0].markdown, "Invoice 12345");
assert_eq!(response.usage_info.unwrap().pages_processed, Some(1));
}
#[tokio::test]
async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() {
let (base, seen, server) = mock_server(vec![textract_response()]).await;
let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| {
assert!(
!wire
.headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("authorization")),
"the hook ran after signing"
);
wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ=");
Ok(wire)
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
let raw = seen.lock().unwrap()[0].clone();
assert_eq!(
request_body(&raw),
json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}})
);
assert_eq!(
header(&raw, "authorization"),
Some(expected_authorization(&format!("{base}/"), &raw).as_str())
);
}
#[tokio::test]
async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() {
let (base, _, server) = mock_server(vec![MockResponse {
status: 400,
headers: vec![],
body: json!({
"__type": "UnsupportedDocumentException",
"Message": "Request has unsupported document format"
}),
}])
.await;
let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base)))
.await
.unwrap_err();
server.await.unwrap();
let Error::Provider { status, body, .. } = error else {
panic!("expected a provider error, got {error:?}");
};
assert_eq!(status, 400);
assert!(
body.contains("multi-page documents are not supported"),
"{body}"
);
}
#[tokio::test]
async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"DocumentMetadata": {"Pages": 1},
"Blocks": [
{"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"},
{"Id": "t", "BlockType": "LAYOUT_TITLE",
"Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]}
]
}))])
.await;
let request = textract_request_for("aws_textract/analyze-document", &base);
let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap();
server.await.unwrap();
let raw = seen.lock().unwrap()[0].clone();
assert_eq!(
header(&raw, "x-amz-target"),
Some("Textract.AnalyzeDocument")
);
assert_eq!(
request_body(&raw)["FeatureTypes"],
json!(["LAYOUT", "TABLES"])
);
assert_eq!(
header(&raw, "authorization"),
Some(expected_authorization(&format!("{base}/"), &raw).as_str())
);
assert_eq!(response.pages[0].markdown, "# Quarterly Report");
}

60
litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
View file

@ -1,10 +1,12 @@
use litellm_host::event::{CallEvent, MachineEvent};
use litellm_llms::base_llm::ocr::error::Error;
use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings};
use rstest::rstest;
use serde_json::{Value, json};
use super::{
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
test_support::{
MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request,
},
wire::{OcrWireRequest, decode_request},
};
use crate::ocr::route::LocalOcrHost;
@ -200,6 +202,42 @@ async fn immediate_response_normalizes_pages_and_preserves_native() {
);
}
#[tokio::test]
async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded",
"analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]}
}))])
.await;
let client = ocr_client().with_settings(OcrSettings {
document_intelligence_api_version: "2099-01-01".into(),
document_intelligence_dpi: 72,
..OcrSettings::default()
});
let result = crate::ocr::client::perform(
&client,
wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})),
)
.await
.unwrap();
server.await.unwrap();
let target = seen.lock().unwrap()[0]
.split_whitespace()
.nth(1)
.unwrap()
.to_string();
assert_eq!(
query_value(&format!("{base}{target}"), "api-version").as_deref(),
Some("2099-01-01")
);
assert_eq!(
serde_json::to_value(&result.pages[0].dimensions).unwrap(),
json!({"width":612,"height":792,"dpi":72})
);
}
#[tokio::test]
async fn accepted_response_polls_to_success_with_only_credentials() {
let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}});
@ -425,13 +463,19 @@ async fn polling_deadline_bounds_retry_delay() {
},
])
.await;
let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
request.transport.poll_timeout = std::time::Duration::from_millis(100);
let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
let client = ocr_client().with_settings(OcrSettings {
poll_timeout: std::time::Duration::from_millis(100),
..OcrSettings::default()
});
let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request))
.await
.unwrap()
.unwrap_err();
let error = tokio::time::timeout(
std::time::Duration::from_secs(1),
crate::ocr::client::perform(&client, request),
)
.await
.unwrap()
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("timed out"));
}

4
litellm-rust/crates/core/tests/cohere_ocr.rs
View file

@ -38,7 +38,7 @@ mod transformation {
)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
@ -75,7 +75,7 @@ mod transformation {
)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}

65
litellm-rust/crates/core/tests/ocr.rs
View file

@ -6,16 +6,15 @@ use litellm_host::{
host::{Host, HostOp, HostResult},
machine::{HostFailure, Machine, MachineStep},
};
use litellm_http::{HttpClientPool, HttpSettings, Resolution};
use litellm_llms::{
base_llm::ocr::{
error::Error as OcrError,
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
},
custom_httpx::{
llm_http_handler::OcrClient,
media::{PublicDnsResolver, UrlPolicy},
},
use litellm_http::{
HttpClientPool, HttpSettings, Resolution,
media::{PublicDnsResolver, UrlPolicy},
};
use litellm_llms::base_llm::ocr::{
error::Error as OcrError,
handler::OcrClient,
settings::OcrSettings,
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
};
use rstest::rstest;
use serde_json::{Value, json};
@ -175,6 +174,43 @@ async fn facade_retains_native_response_when_requested() {
);
}
#[rstest]
#[case::plain_key(&[("MISTRAL_API_KEY", "plain")], "plain")]
#[case::azure_key_wins(&[("MISTRAL_AZURE_API_KEY", "azure"), ("MISTRAL_API_KEY", "plain")], "azure")]
#[case::empty_azure_key_falls_through(&[("MISTRAL_AZURE_API_KEY", ""), ("MISTRAL_API_KEY", "plain")], "plain")]
#[tokio::test]
async fn mistral_env_fallbacks_follow_python_through_the_injected_secret_source(
#[case] secrets: &'static [(&'static str, &'static str)],
#[case] expected_key: &str,
) {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let secret_base = base.clone();
let client = ocr_client().with_secrets(Arc::new(move |name: &str| match name {
"MISTRAL_AZURE_API_BASE" => Some(secret_base.clone()),
"MISTRAL_API_BASE" => Some("http://127.0.0.1:9/never-read".into()),
_ => secrets
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string()),
}));
let request = decode_request(OcrWireRequest {
model: "mistral/model".into(),
document: json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
api_key: None,
api_base: None,
custom_llm_provider: None,
extra_headers: None,
optional_params: Default::default(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap();
crate::ocr::client::perform(&client, request).await.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].contains(&format!("authorization: Bearer {expected_key}")));
}
#[tokio::test]
async fn ocr_client_uses_the_injected_http_pool_configuration() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
@ -187,6 +223,8 @@ async fn ocr_client_uses_the_injected_http_pool_configuration() {
&Resolution::from(&settings).config,
UrlPolicy::default(),
VertexAuth::default(),
OcrSettings::default(),
Arc::new(litellm_core_utils::settings::ProcessEnvironment),
)
.unwrap();
crate::ocr::client::perform(&client, wire_request("mistral/model", &base, json!({})))
@ -624,7 +662,7 @@ async fn read_bounded_response(response: Vec<u8>, limit: usize) -> Result<bytes:
.unwrap();
let result = tokio::time::timeout(
std::time::Duration::from_secs(2),
litellm_llms::custom_httpx::llm_http_handler::read_response_bytes(response, limit),
litellm_llms::base_llm::ocr::handler::read_response_bytes(response, limit),
)
.await;
server.abort();
@ -676,10 +714,7 @@ async fn oversized_error_retains_http_status_and_bounded_diagnostics_without_dra
.await
.unwrap_err();
match error {
OcrError::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status,
body,
}) => {
OcrError::Transport(litellm_http::transport::Error::Http { status, body }) => {
assert_eq!(status, 429);
assert_eq!(body, prefix);
}

7
litellm-rust/crates/core/tests/ocr/support.rs
View file

@ -2,9 +2,10 @@ use std::sync::{Arc, Mutex};
use futures_util::future::BoxFuture;
use litellm_host::event::WireRequest;
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::{CallHooks, OcrClient},
use litellm_llms::base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::LiteLLMOcrResponse,
};
use serde_json::{Value, json};
use tokio::{

51
litellm-rust/crates/core/tests/vertex_ai_ocr.rs
View file

@ -1,8 +1,8 @@
use litellm_auth::InputSource;
use litellm_llms::base_llm::ocr::transformation::OcrResponseFormat;
use litellm_llms::base_llm::ocr::{settings::OcrSettings, transformation::OcrResponseFormat};
use serde_json::{Value, json};
use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request};
use super::test_support::{MockResponse, mock_server, ocr_client, perform_ocr, wire_request};
fn request_body(request: &str) -> Value {
serde_json::from_str(request.split_once("\r\n\r\n").unwrap().1).unwrap()
@ -48,6 +48,27 @@ async fn facade_executes_vertex_mistral_with_resolved_project_and_location() {
);
}
#[tokio::test]
async fn configured_project_and_location_apply_when_the_call_sets_neither() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
let client = ocr_client().with_settings(OcrSettings {
vertex_project: Some("configured-project".into()),
vertex_location: Some("europe-west4".into()),
..OcrSettings::default()
});
crate::ocr::client::perform(
&client,
wire_request("vertex_ai/mistral-ocr-maas", &base, json!({})),
)
.await
.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].starts_with(
"POST /v1/projects/configured-project/locations/europe-west4/publishers/mistralai/models/mistral-ocr-maas:rawPredict "
));
}
#[tokio::test]
async fn supplied_authorization_is_forwarded_without_a_static_token() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await;
@ -139,17 +160,16 @@ async fn adapters_build_complete_requests_and_share_mistral_normalization() {
.prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks)
.await
.unwrap();
assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr");
assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr");
assert_eq!(
vertex_http.url().as_str(),
vertex_http.url(),
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
for http in [&direct_http, &vertex_http] {
assert_eq!(http.method(), reqwest::Method::POST);
assert_eq!(http.headers()["authorization"], "Bearer test-key");
assert_eq!(http.headers()["content-type"], "application/json");
assert_eq!(http.timeout(), Some(&Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(http.header("authorization").unwrap(), "Bearer test-key");
assert_eq!(http.header("content-type").unwrap(), "application/json");
assert_eq!(http.timeout(), Some(Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
@ -229,9 +249,9 @@ mod transformation {
.prepare_request(&vertex, &client, &crate::ocr::test_support::NoHooks)
.await
.unwrap();
assert_eq!(direct_http.url().as_str(), "https://mistral.test/v1/ocr");
assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr");
assert_eq!(
vertex_http.url().as_str(),
vertex_http.url(),
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
let http = if use_vertex {
@ -239,11 +259,10 @@ mod transformation {
} else {
&direct_http
};
assert_eq!(http.method(), reqwest::Method::POST);
assert_eq!(http.headers()["authorization"], "Bearer test-key");
assert_eq!(http.headers()["content-type"], "application/json");
assert_eq!(http.timeout(), Some(&Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body().unwrap().as_bytes().unwrap()).unwrap();
assert_eq!(http.header("authorization").unwrap(), "Bearer test-key");
assert_eq!(http.header("content-type").unwrap(), "application/json");
assert_eq!(http.timeout(), Some(Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({

16
litellm-rust/crates/host-python/src/callable.rs
View file

@ -73,13 +73,7 @@ abort = KeyboardInterrupt('cancelled')
let wrapped = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err();
assert!(wrapped.is_instance_of::<PyRuntimeError>(py));
assert!(wrapped.cause(py).unwrap().value(py).is(&original));
assert!(
wrapped
.value(py)
.getattr("__context__")
.unwrap()
.is(&original)
);
assert!(wrapped.context(py).unwrap().value(py).is(&original));
assert_eq!(
wrapped.value(py).str().unwrap().to_str().unwrap(),
"Failed to reach the caller: unavailable"
@ -115,13 +109,7 @@ original = Unformattable('cannot render')
let original = raised(&locals, "original");
let error = wrap_failure(py, TEMPLATE, failure(&original)).unwrap_err();
assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
assert!(
error
.value(py)
.getattr("__context__")
.unwrap()
.is(&original)
);
assert!(error.context(py).unwrap().value(py).is(&original));
});
}

28
litellm-rust/crates/host-python/src/driver.rs
View file

@ -445,14 +445,8 @@ where
Ok(failure) => return failure.into(),
Err(classifier_error) => classifier_error,
};
let attached = classifier_error.value(py).setattr(
"__context__",
PyRuntimeError::new_err(native).into_value(py),
);
match attached {
Ok(()) => classifier_error,
Err(error) => error,
}
classifier_error.set_context(py, Some(PyRuntimeError::new_err(native)));
classifier_error
}
fn succeeded(&mut self, py: Python<'_>, response: Py<PyAny>) -> PyResult<ExecutionStep> {
@ -1071,9 +1065,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
let error = result.unwrap_err();
assert!(error.is_instance_of::<pyo3::exceptions::PyTypeError>(py));
assert_eq!(error.value(py).to_string(), "classifier failed");
let context = error.value(py).getattr("__context__").unwrap();
assert!(context.is_instance_of::<PyRuntimeError>());
assert_eq!(context.str().unwrap().to_string(), "provider exploded");
let context = error.context(py).unwrap();
assert!(context.is_instance_of::<PyRuntimeError>(py));
assert_eq!(context.value(py).to_string(), "provider exploded");
assert_eq!(
log,
[
@ -1186,20 +1180,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri
type Failure = Classified;
fn invoke(
&mut self,
py: Python<'_>,
_: Python<'_>,
_: &Bound<'_, PyDict>,
_: &'static str,
) -> Result<String, InvokeError<Error>> {
self.0.push("route");
Err(PyErr::from_value(
py.import("asyncio")
.unwrap()
.getattr("CancelledError")
.unwrap()
.call0()
.unwrap(),
)
.into())
Err(pyo3::exceptions::asyncio::CancelledError::new_err(()).into())
}
fn chunk(
&mut self,

102
litellm-rust/crates/host-python/src/execution.rs
View file

@ -4,6 +4,7 @@ use std::pin::Pin;
use std::task::{Context, Poll, Waker};
use std::time::Duration;
use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
use crate::{Pythonized, panic_to_pyerr, release_gil};
use futures_util::FutureExt;
use pyo3::exceptions::PyRuntimeError;
@ -12,6 +13,67 @@ use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pyo3::create_exception!(
_native,
ForkedAfterNativeRuntimeStarted,
PyRuntimeError,
"This process was forked after the native runtime started. Runtime threads do not survive fork(), so native routes cannot run here."
);
pyo3::create_exception!(
_native,
ProcessReservedForForking,
PyRuntimeError,
"This process was reserved for forking workers, so native routes cannot run here."
);
static FORK_GATE: ForkGate = ForkGate::new();
/// Whether this process has started the Tokio runtime.
pub fn runtime_started() -> bool {
FORK_GATE.started(std::process::id())
}
/// Declares that this process exists to fork workers, so it must never start the runtime.
/// Fails if it already has. Workers are unaffected: the reservation is keyed by pid.
pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> {
FORK_GATE.reserve(std::process::id())
}
/// The only door to the Tokio runtime: every route reaches it through this module, which is
/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it.
fn enter_runtime() -> PyResult<()> {
FORK_GATE
.enter(std::process::id())
.map_err(|refused| match refused {
Refused::ReservedForForking => ProcessReservedForForking::new_err(
"this process is reserved for forking workers and cannot run native routes; \
move the call into a worker, after the fork",
),
Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err(
"this process was forked after the native runtime started, and runtime threads \
do not survive fork(); start workers with spawn or forkserver, or fork before \
the first native call",
),
})
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn runtime() -> PyResult<&'static Runtime> {
enter_runtime()?;
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where
F: Future<Output = PyResult<T>> + Send + 'static,
T: for<'py> IntoPyObject<'py> + Send + 'static,
{
enter_runtime()?;
pyo3_async_runtimes::tokio::future_into_py(py, future)
}
pub fn run_sync<T, E, F>(
py: Python<'_>,
future: F,
@ -22,12 +84,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
run_sync_on(py, runtime()?, future, map_error)
}
pub fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
@ -35,7 +92,7 @@ where
T: Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future)
run_sync_value_on(py, runtime()?, future)
}
fn run_sync_value_on<T, F>(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult<T>
@ -83,7 +140,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
future_into_py(py, async move {
let result = catch_future_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
@ -95,7 +152,7 @@ where
T: for<'py> IntoPyObject<'py> + Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? })
future_into_py(py, async move { catch_future_panic(future).await? })
}
pub fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
@ -103,8 +160,9 @@ where
T: Send,
F: Future<Output = PyResult<T>> + Send,
{
let runtime = runtime()?;
let result = release_gil(py, || {
let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter();
let _runtime = runtime.enter();
std::panic::catch_unwind(AssertUnwindSafe(|| {
future.poll(&mut Context::from_waker(Waker::noop()))
}))
@ -286,27 +344,25 @@ mod tests {
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
fn runtime_worker_count() -> PyResult<usize> {
Ok(runtime()?.metrics().num_workers())
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> PyResult<bool> {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
return Ok(false);
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
runtime()?.spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
Ok(heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok())
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
@ -317,6 +373,16 @@ mod tests {
.expect("result should convert")
}
#[rstest]
fn reaching_the_runtime_marks_the_process_as_started(
#[from(initialized_python)] python: &InitializedPython,
) {
python.attach(|py| {
run_sync_value(py, async { Ok(()) }).unwrap();
assert!(runtime_started());
});
}
#[rstest]
fn inline_poll_releases_gil_and_enters_runtime(
#[from(initialized_python)] python: &InitializedPython,

139
litellm-rust/crates/host-python/src/fork_gate.rs Normal file
View file

@ -0,0 +1,139 @@
use std::sync::atomic::{AtomicU32, Ordering};
const UNSET: u32 = 0;
/// Decides which process may use the Tokio runtime. Its worker threads do not survive
/// `fork()`: a child forked after they started hangs on its first native call. The gate turns
/// both halves of that hazard into errors, keyed by pid so a fork needs no hook to be seen:
/// a process reserved for forking can never start the runtime, and a child of a process that
/// did start it is refused instead of hanging.
pub(crate) struct ForkGate {
runtime_pid: AtomicU32,
fork_only_pid: AtomicU32,
}
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum Refused {
ReservedForForking,
ForkedAfterStart,
}
#[derive(Debug, PartialEq, Eq)]
pub struct RuntimeAlreadyStarted;
impl ForkGate {
pub(crate) const fn new() -> Self {
Self {
runtime_pid: AtomicU32::new(UNSET),
fork_only_pid: AtomicU32::new(UNSET),
}
}
/// Claims the runtime for `pid`. Claim first, then look for a reservation: `reserve` does
/// the mirror image, so when the two race at least one of them sees the other.
pub(crate) fn enter(&self, pid: u32) -> Result<(), Refused> {
match self
.runtime_pid
.compare_exchange(UNSET, pid, Ordering::SeqCst, Ordering::SeqCst)
{
Err(owner) if owner != pid => return Err(Refused::ForkedAfterStart),
_ => {}
}
if self.fork_only_pid.load(Ordering::SeqCst) == pid {
// Nothing was started, so the workers forked from here must still find it unclaimed.
let _ =
self.runtime_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(Refused::ReservedForForking);
}
Ok(())
}
/// Reserves `pid` for forking. Reserve first, then look for a started runtime: `enter` does
/// the mirror image, so when the two race at least one of them sees the other. A refused
/// reservation leaves the gate exactly as it was, so a process already running the runtime
/// keeps refusing the children it forks.
pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> {
self.fork_only_pid.store(pid, Ordering::SeqCst);
if self.runtime_pid.load(Ordering::SeqCst) == pid {
let _ =
self.fork_only_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(RuntimeAlreadyStarted);
}
Ok(())
}
pub(crate) fn started(&self, pid: u32) -> bool {
self.runtime_pid.load(Ordering::SeqCst) == pid
}
}
#[cfg(test)]
mod tests {
use super::*;
const MASTER: u32 = 100;
const WORKER: u32 = 101;
#[test]
fn unreserved_process_starts_the_runtime_and_stays_started() {
let gate = ForkGate::new();
assert!(!gate.started(MASTER));
assert_eq!(gate.enter(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
}
#[test]
fn reserved_process_can_never_start_the_runtime() {
let gate = ForkGate::new();
assert_eq!(gate.reserve(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert!(!gate.started(MASTER));
}
#[test]
fn workers_forked_from_a_reserved_process_start_their_own_runtime() {
let gate = ForkGate::new();
gate.reserve(MASTER).unwrap();
gate.enter(MASTER).unwrap_err();
assert_eq!(gate.enter(WORKER), Ok(()));
assert!(gate.started(WORKER));
}
#[test]
fn reserving_after_the_runtime_started_is_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
}
#[test]
fn a_refused_reservation_leaves_the_runtime_claimed_and_its_children_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
}
#[test]
fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
assert!(!gate.started(WORKER));
assert_eq!(gate.enter(MASTER), Ok(()));
}
}

8
litellm-rust/crates/host-python/src/lib.rs
View file

@ -8,6 +8,7 @@ mod argument;
mod callable;
mod driver;
mod execution;
mod fork_gate;
mod gil;
mod handle;
mod marshal;
@ -18,7 +19,12 @@ pub use adapter::{
pub use argument::lookup;
pub use callable::wrap_failure;
pub use driver::run_call;
pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value};
pub use execution::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value,
reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value,
runtime_started,
};
pub use fork_gate::RuntimeAlreadyStarted;
pub use gil::{release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep};
pub use marshal::{Pythonized, from_py, from_py_argument, panic_to_pyerr, to_py};

8
litellm-rust/crates/http/Cargo.toml
View file

@ -5,12 +5,20 @@ edition.workspace = true
license.workspace = true
repository.workspace = true
[features]
test-support = []
[dependencies]
http.workspace = true
litellm-core-utils.workspace = true
hyper-util.workspace = true
reqwest.workspace = true
rustls.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true
veil.workspace = true
webpki-roots.workspace = true
[dev-dependencies]

41
litellm-rust/crates/http/src/config.rs
View file

@ -6,6 +6,7 @@ use std::{
use crate::{
error::Error,
proxy::EnvironmentProxies,
settings::{HttpSettings, SslVerify, TcpKeepalive},
tls::{CipherSelection, KeyExchangeGroup, Tls12CipherSuite, Unsupported},
};
@ -26,7 +27,7 @@ pub struct HttpClientConfig {
pub force_ipv4: bool,
pub http2: bool,
pub user_agent: Option<String>,
pub trust_proxy_env: bool,
pub proxies: EnvironmentProxies,
pub connect_timeout: Duration,
pub tcp_keepalive: Option<TcpKeepalive>,
pub pool_idle_timeout: Duration,
@ -72,7 +73,11 @@ impl From<&HttpSettings> for Resolution {
force_ipv4: settings.force_ipv4,
http2: settings.http2,
user_agent: settings.user_agent.clone(),
trust_proxy_env: settings.trust_proxy_env,
proxies: if settings.trust_proxy_env {
settings.proxies.clone()
} else {
EnvironmentProxies::default()
},
connect_timeout: settings.connect_timeout,
tcp_keepalive: settings.tcp_keepalive,
pool_idle_timeout: settings.pool_idle_timeout,
@ -111,11 +116,11 @@ impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
Some(agent) => with_protocol.user_agent(agent),
None => with_protocol,
};
Ok(if config.trust_proxy_env {
with_agent
} else {
with_agent.no_proxy()
})
Ok(config
.proxies
.reqwest_proxies()
.into_iter()
.fold(with_agent.no_proxy(), reqwest::ClientBuilder::proxy))
}
}
@ -227,6 +232,25 @@ mod tests {
);
}
fn proxies() -> EnvironmentProxies {
EnvironmentProxies::from_environment(&|name: &str| {
(name == "HTTPS_PROXY").then(|| "http://proxy.corp:3128".to_string())
})
}
#[test]
fn proxies_are_dropped_when_the_transport_does_not_trust_the_environment() {
let settings = HttpSettings {
trust_proxy_env: false,
proxies: proxies(),
..HttpSettings::default()
};
assert_eq!(
Resolution::from(&settings).config.proxies,
EnvironmentProxies::default()
);
}
#[test]
fn connection_settings_carry_over_unchanged() {
let keepalive = TcpKeepalive {
@ -240,6 +264,7 @@ mod tests {
http2: true,
user_agent: Some("litellm/1.0".into()),
trust_proxy_env: true,
proxies: proxies(),
connect_timeout: Duration::from_secs(7),
tcp_keepalive: Some(keepalive),
pool_idle_timeout: Duration::from_secs(45),
@ -256,7 +281,7 @@ mod tests {
force_ipv4: true,
http2: true,
user_agent: Some("litellm/1.0".into()),
trust_proxy_env: true,
proxies: proxies(),
connect_timeout: Duration::from_secs(7),
tcp_keepalive: Some(keepalive),
pool_idle_timeout: Duration::from_secs(45),

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