scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

merge: main into litellm_mcp_ui_prompts_resources

Browse source 
Resolve the MCP SDK 2 rename (McpError -> MCPError) in rest_endpoints.py, read snake_case attributes on SDK Resource/ResourceTemplate in the catalog test, and regenerate the lazy OpenAPI snapshot and dashboard API types on SDK 2 schemas

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
joshua 2026-09-19 19:00:48 +00:00
commit 300f0f92bc
469 changed files with 54796 additions and 5209 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.circleci/config.yml
View file

@ -3015,7 +3015,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" \

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:

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

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

71
.github/workflows/issue_fixed_comment.yml vendored Normal file
View file

@ -0,0 +1,71 @@
name: Issue fixed comment
on:
issues:
types: [closed]
workflow_dispatch:
inputs:
issue_number:
description: "Closed issue number to comment on manually."
required: true
pull_request:
paths:
- .github/workflows/issue_fixed_comment.yml
- scripts/comment-fixed-issue.ts
- scripts/comment-fixed-issue.test.ts
- scripts/auto-close-duplicates.ts
permissions: {}
concurrency:
group: issue-fixed-comment-${{ github.event.issue.number || github.event.inputs.issue_number || github.run_id }}
cancel-in-progress: false
jobs:
comment-fixed-issue-tests:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
bun-version: "1.4.0"
- name: Test the closer lookup, the release placement and the comment
run: bun test scripts/comment-fixed-issue.test.ts
comment-fixed-issue:
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
issues: write
steps:
- name: Checkout scripts
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: scripts
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
# Exact version, never latest: the next step holds an issues: write token
bun-version: "1.4.0"
- name: Name the release that carries the fix
run: bun run scripts/comment-fixed-issue.ts | tee -a "${GITHUB_STEP_SUMMARY}"
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
ISSUE_NUMBER: ${{ github.event.issue.number || github.event.inputs.issue_number }}
DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
DRY_RUN: ${{ vars.ISSUE_FIXED_COMMENT_ENABLED != 'true' }}

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

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

130
AGENTS.md
View file

@ -1,3 +1,131 @@
Read @CLAUDE.md for coding guidelines
Do not write comments unless they are any of:
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
- used as an input for tools to read and act on. For example:
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
- a TODO or FIXME
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
- correct
- secure
- performant
- readable
- easy to maintain/change
- modern
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0`
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
Commit and push your work when you're done without asking
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
When working on a PR, keep the PR description in sync with new commits being made
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>`
- Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: <reason>`
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match
- No monster files or god objects
- No file sprawl: deliberate file and folder structure
- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
Follow conventional commits for commit names and PR titles
## Think Before Coding
**Don't assume. Don't hide confusion. Surface tradeoffs**
Before implementing:
- State your assumptions explicitly. If uncertain, ask
- If multiple interpretations exist, present them. Don't pick silently
- If a simpler approach exists, say so. Push back when warranted
- If something is unclear, stop. Name what's confusing. Ask
## Simplicity First
**Minimum code that solves the problem. Nothing speculative**
- No features beyond what was asked
- No abstractions for single-use code
- No "flexibility" or "configurability" that wasn't requested
- No error handling for impossible scenarios
- If you write 200 lines and it could be 50, rewrite it
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check

129
CLAUDE.md
View file

@ -1,129 +0,0 @@
Do not write comments unless they are any of:
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
- used as an input for tools to read and act on. For example:
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
- a TODO or FIXME
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
- correct
- secure
- performant
- readable
- easy to maintain/change
- modern
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0`
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
Commit and push your work when you're done without asking
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
When working on a PR, keep the PR description in sync with new commits being made
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>`
- Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: <reason>`
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match
- No monster files or god objects
- No file sprawl: deliberate file and folder structure
- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
Follow conventional commits for commit names and PR titles
## Think Before Coding
**Don't assume. Don't hide confusion. Surface tradeoffs**
Before implementing:
- State your assumptions explicitly. If uncertain, ask
- If multiple interpretations exist, present them. Don't pick silently
- If a simpler approach exists, say so. Push back when warranted
- If something is unclear, stop. Name what's confusing. Ask
## Simplicity First
**Minimum code that solves the problem. Nothing speculative**
- No features beyond what was asked
- No abstractions for single-use code
- No "flexibility" or "configurability" that wasn't requested
- No error handling for impossible scenarios
- If you write 200 lines and it could be 50, rewrite it
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify

17
CONTRIBUTING.md
View file

@ -148,7 +148,7 @@ make lint
Individual linting commands:
```bash
make format-check # Check Black formatting
make format-check # Check ruff format formatting
make lint-ruff # Run Ruff linting
make lint-basedpyright # Run basedpyright type checking
make check-circular-imports # Check for circular imports
@ -160,14 +160,14 @@ Apply formatting (auto-fixes issues):
make format
```
> **Black formatting is enforced in CI.** All PRs must pass the Black formatting check.
> **Formatting is enforced in CI.** All PRs must pass the `ruff format --check` step.
>
> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): `AGENTS.md` and `CLAUDE.md` instruct agents to run `poetry run black .` before committing.
> - **VS Code users**: Install the [Black Formatter extension](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) and enable format-on-save:
> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): follow `AGENTS.md` and run `make format` before committing.
> - **VS Code users**: Install the [Ruff extension](https://marketplace.visualstudio.com/items?itemName=charliermarsh.ruff) and enable format-on-save:
> ```json
> {
> "[python]": {
> "editor.defaultFormatter": "ms-python.black-formatter",
> "editor.defaultFormatter": "charliermarsh.ruff",
> "editor.formatOnSave": true
> }
> }
@ -197,8 +197,8 @@ make help # Show all available commands
make install-dev # Install development dependencies
make install-proxy-dev # Install proxy development dependencies
make install-test-deps # Install the full local test environment
make format # Apply Black code formatting
make format-check # Check Black formatting (matches CI)
make format # Apply ruff format code formatting
make format-check # Check ruff format formatting (matches CI)
make lint # Run all linting checks
make test-unit # Run unit tests
make test-integration # Run integration tests
@ -210,8 +210,7 @@ make test-unit-helm # Run Helm unit tests
LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html).
Our automated quality checks include:
- **Black** for consistent code formatting
- **Ruff** for linting and code quality
- **Ruff** for formatting, linting, and code quality
- **basedpyright** for static type checking
- **Circular import detection**
- **Import safety validation**

2
GEMINI.md
View file

@ -1 +1 @@
Read @CLAUDE.md for coding guidelines
Read @AGENTS.md for coding guidelines

5
README.md
View file

@ -633,9 +633,8 @@ For detailed contributing guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html).
Our automated checks include:
- **Black** for code formatting
- **Ruff** for linting and code quality
- **MyPy** for type checking
- **Ruff** for formatting, linting, and code quality
- **basedpyright** for type checking
- **Circular import detection**
- **Import safety checks**

54
ci_cd/generate_model_prices_schema.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Optional
from typing import Final, Optional
import jsonschema
@ -19,6 +19,10 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
BOOLEAN: JsonSchema = {"type": "boolean"}
STRING: JsonSchema = {"type": "string"}
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
WEEKDAY_PATTERN: Final = (
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
)
EXTRA_BOOLEAN_KEYS = frozenset(
{
@ -31,7 +35,51 @@ EXTRA_BOOLEAN_KEYS = frozenset(
}
)
HOURS_UTC: Final[JsonSchema] = {
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
}
OFF_PEAK_WINDOW: Final[JsonSchema] = {
"type": "object",
"properties": {
"hours_utc": HOURS_UTC,
"weekdays": {
"type": "array",
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
"items": {
"oneOf": [
{"type": "integer", "minimum": 1, "maximum": 7},
{"type": "string", "pattern": WEEKDAY_PATTERN},
]
},
"minItems": 1,
},
},
"required": ["hours_utc"],
"additionalProperties": False,
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
"properties": {
"hours_utc": HOURS_UTC,
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
"weekday_timezone": {
"type": "string",
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
},
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
@ -327,9 +375,7 @@ def render(schema: JsonSchema) -> str:
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(
schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER
)
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)

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

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

@ -11,7 +11,7 @@ Endpoints for /project operations
#### PROJECT MANAGEMENT ####
import json
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Request
@ -22,7 +22,11 @@ from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import delete_cached_project_object
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership
_set_object_metadata_field,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
@ -82,37 +86,38 @@ async def _check_user_permission_for_project(
user_api_key_dict: UserAPIKeyAuth,
team_id: str | None,
prisma_client: PrismaClient,
general_settings: Mapping[str, object],
require_admin: bool = False,
team_object: LiteLLM_TeamTable | None = None,
) -> bool:
"""
Check if user has permission to manage a project.
Returns True if user is proxy admin or team admin (when team_id provided).
Returns True if user is proxy admin, or a team admin of ``team_id`` when the
``team_admin_editable_team_fields`` setting grants team admins the ``projects`` permission.
If require_admin=True, only proxy admins are allowed.
If team_object is provided, it will be used instead of fetching from DB
(avoids duplicate DB queries when team was already fetched for validation).
"""
is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
if require_admin:
if require_admin or is_proxy_admin:
return is_proxy_admin
if is_proxy_admin:
return True
if not team_id or not user_api_key_dict.user_id:
if not team_id or not user_api_key_dict.user_id or not team_admin_may_manage_projects(general_settings):
return False
team = team_object
if team is None:
team = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
team_row: Final = (
team_object
if team_object is not None
else await _team_table(prisma_client).find_unique(where={"team_id": team_id})
)
if team_row is None:
return False
if team and team.admins:
return user_api_key_dict.user_id in team.admins
return False
team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump())
return _is_user_team_admin(user_api_key_dict, team) or user_api_key_dict.user_id in (team.admins or [])
async def _validate_team_exists(
@ -531,6 +536,7 @@ async def new_project(
user_api_key_dict=user_api_key_dict,
team_id=data.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()),
)
@ -735,6 +741,7 @@ async def update_project(
user_api_key_dict=user_api_key_dict,
team_id=existing_project.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
)
if not has_permission:
@ -751,6 +758,7 @@ async def update_project(
user_api_key_dict=user_api_key_dict,
team_id=data.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
team_object=(
LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()) if target_team_obj else None
),
@ -877,7 +885,7 @@ async def delete_project(
}'
```
"""
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import general_settings, premium_user, prisma_client, user_api_key_cache
try:
if not premium_user:
@ -899,6 +907,7 @@ async def delete_project(
user_api_key_dict=user_api_key_dict,
team_id=None,
prisma_client=prisma_client,
general_settings=general_settings,
require_admin=True,
)

27
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",
@ -2050,8 +2052,10 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-host",
"litellm-http",
"litellm-llms",
"litellm-types",
"mime_guess",
@ -2129,6 +2133,24 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-http"
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",
]
[[package]]
name = "litellm-llms"
version = "0.1.0"
@ -2146,6 +2168,7 @@ dependencies = [
"litellm-core-utils",
"litellm-framing",
"litellm-host",
"litellm-http",
"litellm-types",
"reqwest 0.12.28",
"rstest",
@ -2153,6 +2176,7 @@ dependencies = [
"serde_json",
"serde_path_to_error",
"serde_with",
"strum",
"thiserror 2.0.19",
"time",
"tokio",
@ -2167,9 +2191,12 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-gcp",
"litellm-callbacks-legacy",
"litellm-core",
"litellm-core-utils",
"litellm-host-python",
"litellm-http",
"litellm-llms",
"litellm-token-counter",
"litellm-types",

4
litellm-rust/Cargo.toml
View file

@ -17,6 +17,7 @@ litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
litellm-http = { path = "crates/http" }
litellm-llms = { path = "crates/llms" }
litellm-types = { path = "crates/types" }
litellm-core-utils = { path = "crates/core-utils" }
@ -26,6 +27,8 @@ litellm-token-counter = { path = "crates/token-counter" }
litellm-host-python = { path = "crates/host-python" }
bytes = "1"
http = "1"
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
@ -49,6 +52,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
webpki-roots = "1"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"

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

7
litellm-rust/crates/callbacks-legacy/src/adapter.rs
View file

@ -12,7 +12,7 @@ use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyList},
types::{PyDateTime, PyDict, PyList},
};
use serde_json::Value;
@ -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 {

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"
@ -35,6 +36,7 @@ url.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-auth-gcp.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"
)
}

9
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::{
@ -14,7 +13,3 @@ pub async fn perform(
) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
perform(&OcrClient::shared()?, request).await
}

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

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

@ -1,16 +1,20 @@
use std::sync::{Arc, Mutex};
use litellm_auth_gcp::VertexAuth;
use litellm_host::{
event::{CallEvent, MachineEvent, WireRequest},
host::{Host, HostOp, HostResult},
machine::{HostFailure, Machine, MachineStep},
};
use litellm_llms::{
base_llm::ocr::{
error::Error as OcrError,
transformation::{LiteLLMOcrResponse, OCR_RESPONSE_MAX_BYTES, OcrTransportConfig},
},
custom_httpx::llm_http_handler::OcrClient,
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};
@ -170,26 +174,64 @@ 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 facade_uses_the_injected_http_client() {
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 mut default_headers = reqwest::header::HeaderMap::new();
default_headers.insert(
"x-transport-owner",
reqwest::header::HeaderValue::from_static("host"),
);
let provider_http = reqwest::Client::builder()
.default_headers(default_headers)
.build()
.unwrap();
crate::ocr::client::perform(
&OcrClient::new(provider_http).unwrap(),
wire_request("mistral/model", &base, json!({})),
)
.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("x-transport-owner: host"));
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;
let settings = HttpSettings {
user_agent: Some("host-owned/1".into()),
..HttpSettings::default()
};
let client = OcrClient::new(
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
&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!({})))
.await
.unwrap();
server.await.unwrap();
assert!(seen.lock().unwrap()[0].contains("user-agent: host-owned/1"));
}
fn event_name(event: &CallEvent) -> &'static str {
@ -620,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();
@ -672,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,

1
litellm-rust/crates/http/AGENTS.md Normal file
View file

@ -0,0 +1 @@
- https://github.com/BerriAI/litellm-docs/blob/main/docs/guides/security_settings.md

26
litellm-rust/crates/http/Cargo.toml Normal file
View file

@ -0,0 +1,26 @@
[package]
name = "litellm-http"
version = "0.1.0"
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]
rstest.workspace = true
tokio.workspace = true

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

@ -0,0 +1,321 @@
use std::{
net::{IpAddr, Ipv4Addr},
path::PathBuf,
time::Duration,
};
use crate::{
error::Error,
proxy::EnvironmentProxies,
settings::{HttpSettings, SslVerify, TcpKeepalive},
tls::{CipherSelection, KeyExchangeGroup, Tls12CipherSuite, Unsupported},
};
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum Verify {
Disabled,
CaBundle(PathBuf),
BuiltInRoots,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct HttpClientConfig {
pub verify: Verify,
pub client_certificate: Option<PathBuf>,
pub key_exchange_group: Option<KeyExchangeGroup>,
pub tls12_cipher_suites: Option<Vec<Tls12CipherSuite>>,
pub force_ipv4: bool,
pub http2: bool,
pub user_agent: Option<String>,
pub proxies: EnvironmentProxies,
pub connect_timeout: Duration,
pub tcp_keepalive: Option<TcpKeepalive>,
pub pool_idle_timeout: Duration,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Resolution {
pub config: HttpClientConfig,
pub unsupported: Vec<Unsupported>,
}
impl From<&HttpSettings> for Verify {
fn from(settings: &HttpSettings) -> Self {
match &settings.ssl_verify {
Some(SslVerify::Disabled) => Self::Disabled,
Some(SslVerify::CaBundle(path)) => Self::CaBundle(path.clone()),
Some(SslVerify::Enabled) | None => settings
.ssl_cert_file
.clone()
.map_or(Self::BuiltInRoots, Self::CaBundle),
}
}
}
impl From<&HttpSettings> for Resolution {
fn from(settings: &HttpSettings) -> Self {
let curve = settings
.ssl_ecdh_curve
.as_deref()
.map(str::parse::<KeyExchangeGroup>)
.transpose();
let ciphers = settings
.ssl_security_level
.as_deref()
.map(CipherSelection::from)
.unwrap_or_default();
Self {
config: HttpClientConfig {
verify: Verify::from(settings),
client_certificate: settings.ssl_certificate.clone(),
key_exchange_group: curve.clone().ok().flatten(),
tls12_cipher_suites: ciphers.tls12_cipher_suites,
force_ipv4: settings.force_ipv4,
http2: settings.http2,
user_agent: settings.user_agent.clone(),
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,
},
unsupported: curve.err().into_iter().chain(ciphers.unsupported).collect(),
}
}
}
impl TryFrom<&HttpClientConfig> for reqwest::ClientBuilder {
type Error = Error;
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
let base = reqwest::Client::builder()
.use_preconfigured_tls(rustls::ClientConfig::try_from(config)?)
.connect_timeout(config.connect_timeout)
.pool_idle_timeout(config.pool_idle_timeout);
let with_keepalive = match config.tcp_keepalive {
None => base,
Some(keepalive) => base
.tcp_keepalive(keepalive.idle)
.tcp_keepalive_interval(keepalive.interval)
.tcp_keepalive_retries(keepalive.retries),
};
let with_address = if config.force_ipv4 {
with_keepalive.local_address(IpAddr::V4(Ipv4Addr::UNSPECIFIED))
} else {
with_keepalive
};
let with_protocol = if config.http2 {
with_address
} else {
with_address.http1_only()
};
let with_agent = match &config.user_agent {
Some(agent) => with_protocol.user_agent(agent),
None => with_protocol,
};
Ok(config
.proxies
.reqwest_proxies()
.into_iter()
.fold(with_agent.no_proxy(), reqwest::ClientBuilder::proxy))
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn settings(ssl_verify: Option<SslVerify>, ssl_cert_file: Option<&str>) -> HttpSettings {
HttpSettings {
ssl_verify,
ssl_cert_file: ssl_cert_file.map(PathBuf::from),
..HttpSettings::default()
}
}
#[rstest]
#[case::default(settings(None, None), Verify::BuiltInRoots)]
#[case::setting_disables(
settings(Some(SslVerify::Disabled), Some("/env/roots.pem")),
Verify::Disabled
)]
#[case::setting_bundle(
settings(Some(SslVerify::CaBundle("/configured.pem".into())), Some("/env/roots.pem")),
Verify::CaBundle("/configured.pem".into())
)]
#[case::enabled_uses_cert_file(
settings(Some(SslVerify::Enabled), Some("/env/roots.pem")),
Verify::CaBundle("/env/roots.pem".into())
)]
#[case::unset_uses_cert_file(settings(None, Some("/env/roots.pem")), Verify::CaBundle("/env/roots.pem".into()))]
fn verify_follows_setting_then_cert_file(
#[case] settings: HttpSettings,
#[case] expected: Verify,
) {
let config = Resolution::from(&settings).config;
assert_eq!(config.verify, expected);
}
#[rstest]
#[case::x25519("X25519", Some(KeyExchangeGroup::X25519))]
#[case::openssl_p256("prime256v1", Some(KeyExchangeGroup::Secp256r1))]
#[case::p384("secp384r1", Some(KeyExchangeGroup::Secp384r1))]
fn ecdh_curve_selects_the_single_key_exchange_group(
#[case] curve: &str,
#[case] expected: Option<KeyExchangeGroup>,
) {
let settings = HttpSettings {
ssl_ecdh_curve: Some(curve.into()),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(resolution.config.key_exchange_group, expected);
assert_eq!(resolution.unsupported, []);
}
#[test]
fn unsupported_ecdh_curve_keeps_the_defaults_and_is_reported() {
let settings = HttpSettings {
ssl_ecdh_curve: Some("secp521r1".into()),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(resolution.config.key_exchange_group, None);
assert_eq!(
resolution.unsupported,
[Unsupported::EcdhCurve("secp521r1".into())]
);
}
#[test]
fn legacy_security_level_keeps_every_suite_and_is_reported_unsupported() {
let settings = HttpSettings {
ssl_security_level: Some("DEFAULT@SECLEVEL=1".into()),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(resolution.config.tls12_cipher_suites, None);
assert_eq!(
resolution.unsupported,
[Unsupported::SecurityLevel("@SECLEVEL=1".into())]
);
}
#[test]
fn named_suites_restrict_tls12_and_unsupported_entries_are_reported() {
let settings = HttpSettings {
ssl_security_level: Some(
"ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES128-GCM-SHA256:!aNULL:AES256-SHA@SECLEVEL=2"
.into(),
),
..HttpSettings::default()
};
let resolution = Resolution::from(&settings);
assert_eq!(
resolution.config.tls12_cipher_suites,
Some(vec![
Tls12CipherSuite::EcdheEcdsaAes128Gcm,
Tls12CipherSuite::EcdheRsaAes256Gcm
])
);
assert_eq!(
resolution.unsupported,
[
Unsupported::CipherToken("!aNULL".into()),
Unsupported::CipherToken("AES256-SHA".into())
]
);
}
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 {
idle: Duration::from_secs(60),
interval: Duration::from_secs(30),
retries: 5,
};
let settings = HttpSettings {
ssl_certificate: Some("/client.pem".into()),
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),
..HttpSettings::default()
};
let config = Resolution::from(&settings).config;
assert_eq!(
config,
HttpClientConfig {
verify: Verify::BuiltInRoots,
client_certificate: Some("/client.pem".into()),
key_exchange_group: None,
tls12_cipher_suites: None,
force_ipv4: true,
http2: true,
user_agent: Some("litellm/1.0".into()),
proxies: proxies(),
connect_timeout: Duration::from_secs(7),
tcp_keepalive: Some(keepalive),
pool_idle_timeout: Duration::from_secs(45),
}
);
}
#[test]
fn missing_ca_bundle_is_a_read_error() {
let path = std::env::temp_dir().join("litellm-http-missing-bundle.pem");
let config = HttpClientConfig {
verify: Verify::CaBundle(path.clone()),
..Resolution::from(&HttpSettings::default()).config
};
assert!(matches!(
reqwest::ClientBuilder::try_from(&config),
Err(Error::Read { path: reported, .. }) if reported == path
));
}
#[test]
fn non_pem_ca_bundle_is_an_invalid_pem_error() {
let path =
std::env::temp_dir().join(format!("litellm-http-not-pem-{}.pem", std::process::id()));
std::fs::write(&path, b"not a certificate").unwrap();
let config = HttpClientConfig {
verify: Verify::CaBundle(path.clone()),
..Resolution::from(&HttpSettings::default()).config
};
let result = reqwest::ClientBuilder::try_from(&config).map(drop);
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
));
}
}

23
litellm-rust/crates/http/src/error.rs Normal file
View file

@ -0,0 +1,23 @@
use std::path::PathBuf;
#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
#[error("could not read {}: {message}", path.display())]
Read { path: PathBuf, message: String },
#[error("{} is not a PEM file: {message}", path.display())]
InvalidPem { path: PathBuf, message: String },
#[error("could not build the HTTP client: {0}")]
Client(String),
#[error("request body could not be serialized: {0}")]
RequestBody(String),
#[error("request forwards a header the signer computes: {0}")]
ComputedHeader(String),
#[error("request signing failed: {0}")]
Signature(String),
}
impl From<reqwest::Error> for Error {
fn from(error: reqwest::Error) -> Self {
Self::Client(error.without_url().to_string())
}
}

17
litellm-rust/crates/http/src/lib.rs Normal file
View file

@ -0,0 +1,17 @@
mod config;
mod error;
pub mod media;
pub mod outbound;
mod pool;
mod proxy;
pub mod request;
mod settings;
mod tls;
pub mod transport;
pub use config::{HttpClientConfig, Resolution, Verify};
pub use error::Error;
pub use pool::{ClientVariant, HttpClientPool};
pub use proxy::EnvironmentProxies;
pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive};
pub use tls::{KeyExchangeGroup, Tls12CipherSuite, Unsupported};

231
litellm-rust/crates/llms/src/custom_httpx/media.rs → litellm-rust/crates/http/src/media.rs
View file

@ -12,7 +12,7 @@ use reqwest::{
dns::{Addrs, Name, Resolve, Resolving},
};
const MEDIA_CONNECT_TIMEOUT_SECS: u64 = 10;
use crate::{ClientVariant, HttpClientConfig, HttpClientPool};
#[derive(Debug, thiserror::Error)]
pub enum Error {
@ -33,13 +33,48 @@ pub enum Error {
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] crate::custom_httpx::transport::Error),
Transport(#[from] crate::transport::Error),
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UrlPolicy {
pub validate: bool,
pub allowed_hosts: Vec<String>,
}
impl Default for UrlPolicy {
fn default() -> Self {
Self {
validate: true,
allowed_hosts: Vec::new(),
}
}
}
impl UrlPolicy {
fn allows(&self, host: &str, port: u16) -> bool {
let host = normalize_host(host);
let with_port = format!("{host}:{port}");
self.allowed_hosts
.iter()
.map(|entry| normalize_host(entry))
.any(|entry| entry == host || entry == with_port)
}
}
fn normalize_host(host: &str) -> String {
host.to_ascii_lowercase().trim_end_matches('.').to_owned()
}
type ProxyMatch = Arc<dyn Fn(&Url) -> bool + Send + Sync>;
#[derive(Clone)]
pub struct MediaFetcher {
client: reqwest::Client,
pinned: reqwest::Client,
unpinned: reqwest::Client,
uses_proxy: ProxyMatch,
address_resolver: Arc<dyn AddressResolver>,
url_policy: UrlPolicy,
allow_private_network: bool,
}
@ -63,26 +98,34 @@ pub struct DownloadedMedia {
}
impl MediaFetcher {
pub fn new() -> Result<Self, reqwest::Error> {
Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver))
pub fn new(
pool: &HttpClientPool,
config: &HttpClientConfig,
url_policy: UrlPolicy,
) -> Result<Self, crate::Error> {
let uses_proxy: ProxyMatch = Arc::new(config.proxies.matcher());
Self::with_resolution(
pool,
config,
url_policy,
Arc::new(SystemAddressResolver),
uses_proxy,
)
}
fn with_resolvers<R>(
transport_resolver: Arc<R>,
fn with_resolution(
pool: &HttpClientPool,
config: &HttpClientConfig,
url_policy: UrlPolicy,
address_resolver: Arc<dyn AddressResolver>,
) -> Result<Self, reqwest::Error>
where
R: Resolve + 'static,
{
let client = reqwest::Client::builder()
.connect_timeout(Duration::from_secs(MEDIA_CONNECT_TIMEOUT_SECS))
.redirect(reqwest::redirect::Policy::none())
.no_proxy()
.dns_resolver(transport_resolver)
.build()?;
uses_proxy: ProxyMatch,
) -> Result<Self, crate::Error> {
Ok(Self {
client,
pinned: pool.client(config, ClientVariant::Media)?,
unpinned: pool.client(config, ClientVariant::UnpinnedMedia)?,
uses_proxy,
address_resolver,
url_policy,
allow_private_network: false,
})
}
@ -90,8 +133,11 @@ impl MediaFetcher {
#[cfg(any(test, feature = "test-support"))]
pub fn for_test(client: reqwest::Client) -> Self {
Self {
client,
pinned: client.clone(),
unpinned: client,
uses_proxy: Arc::new(|_| false),
address_resolver: Arc::new(AllowPrivateResolver),
url_policy: UrlPolicy::default(),
allow_private_network: true,
}
}
@ -112,13 +158,13 @@ impl MediaFetcher {
) -> Result<DownloadedMedia, Error> {
let mut redirects_followed = 0;
loop {
self.validate_url(&url).await?;
let mut response = self
.client
.client_for(&url)
.await?
.get(url.clone())
.send()
.await
.map_err(crate::custom_httpx::transport::Error::from)?;
.map_err(crate::transport::Error::from)?;
if response.status().is_redirection() {
if redirects_followed == policy.max_redirects {
return Err(Error::TooManyRedirects);
@ -149,7 +195,7 @@ impl MediaFetcher {
while let Some(chunk) = response
.chunk()
.await
.map_err(crate::custom_httpx::transport::Error::from)?
.map_err(crate::transport::Error::from)?
{
enforce_download_size(bytes.len() as u64 + chunk.len() as u64, policy.max_bytes)?;
bytes.extend_from_slice(&chunk);
@ -161,7 +207,10 @@ impl MediaFetcher {
}
}
async fn validate_url(&self, url: &Url) -> Result<(), Error> {
async fn client_for(&self, url: &Url) -> Result<&reqwest::Client, Error> {
if !self.url_policy.validate {
return Ok(&self.unpinned);
}
if !matches!(url.scheme(), "http" | "https")
|| !url.username().is_empty()
|| url.password().is_some()
@ -170,17 +219,33 @@ impl MediaFetcher {
}
let host = url.host_str().ok_or(Error::BlockedUrl)?;
if self.allow_private_network {
return Ok(());
}
if let Ok(ip) = host.parse::<IpAddr>() {
return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl);
return Ok(&self.pinned);
}
let port = url.port_or_known_default().ok_or(Error::BlockedUrl)?;
if self.url_policy.allows(host, port) {
return Ok(&self.unpinned);
}
self.validate_host(host, port).await?;
Ok(if (self.uses_proxy)(url) {
&self.unpinned
} else {
&self.pinned
})
}
async fn validate_host(&self, host: &str, port: u16) -> Result<(), Error> {
if let Ok(ip) = host
.trim_start_matches('[')
.trim_end_matches(']')
.parse::<IpAddr>()
{
return (!is_blocked_ip(ip)).then_some(()).ok_or(Error::BlockedUrl);
}
let addresses = self
.address_resolver
.resolve(host, port)
.await
.map_err(|error| crate::custom_httpx::transport::Error::Network(error.to_string()))?;
.map_err(|error| crate::transport::Error::Network(error.to_string()))?;
validate_addresses(&addresses)
}
}
@ -236,7 +301,7 @@ fn is_blocked_ip(ip: IpAddr) -> bool {
}
#[derive(Default)]
struct PublicDnsResolver;
pub struct PublicDnsResolver;
struct SystemAddressResolver;
@ -287,6 +352,7 @@ mod tests {
};
use super::*;
use crate::{HttpSettings, Resolution};
async fn serve(response: &'static [u8]) -> (Url, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0")
@ -364,13 +430,32 @@ mod tests {
address: SocketAddr,
blocked_hosts: HashSet<&'static str>,
) -> MediaFetcher {
MediaFetcher::with_resolvers(
Arc::new(LoopbackDnsResolver(address)),
fetcher(address, blocked_hosts, UrlPolicy::default(), false)
}
fn fetcher(
pinned_address: SocketAddr,
blocked_hosts: HashSet<&'static str>,
url_policy: UrlPolicy,
uses_proxy: bool,
) -> MediaFetcher {
let direct = Resolution::from(&HttpSettings::default()).config;
MediaFetcher::with_resolution(
&HttpClientPool::new(Arc::new(LoopbackDnsResolver(pinned_address))),
&direct,
url_policy,
Arc::new(TestAddressResolver { blocked_hosts }),
Arc::new(move |_| uses_proxy),
)
.expect("test fetcher builds")
}
const UNROUTABLE: SocketAddr =
SocketAddr::new(IpAddr::V4(std::net::Ipv4Addr::new(192, 0, 2, 1)), 9);
const OK_RESPONSE: &[u8] =
b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok";
fn policy(max_bytes: u64, max_redirects: usize) -> DownloadPolicy {
DownloadPolicy {
timeout: Duration::from_secs(1),
@ -542,12 +627,90 @@ mod tests {
#[tokio::test]
async fn rejects_url_credentials_before_network_access() {
let fetcher = MediaFetcher::new().expect("media fetcher builds");
let fetcher = MediaFetcher::new(
&HttpClientPool::new(Arc::new(PublicDnsResolver)),
&Resolution::from(&HttpSettings::default()).config,
UrlPolicy::default(),
)
.expect("media fetcher builds");
let url =
Url::parse("https://user:password@8.8.8.8/document").expect("credentialed URL parses");
assert!(matches!(
fetcher.validate_url(&url).await,
fetcher.fetch(url, policy(1, 0)).await,
Err(Error::BlockedUrl)
));
}
#[tokio::test]
async fn allowlisted_private_host_is_fetched_without_the_pinned_resolver() {
let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
let port = url.port().expect("test URL has a port");
let allowed = UrlPolicy {
validate: true,
allowed_hosts: vec![format!("LOCALHOST:{port}")],
};
let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), allowed, false)
.fetch(url, policy(2, 0))
.await
.expect("allowlisted host downloads");
server.await.expect("server completes");
assert_eq!(media.bytes, b"ok");
}
#[tokio::test]
async fn allowlist_entry_for_another_port_does_not_open_the_host() {
let (url, _server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
let other_port = UrlPolicy {
validate: true,
allowed_hosts: vec!["localhost:1".into()],
};
let result = fetcher(UNROUTABLE, HashSet::from(["localhost"]), other_port, false)
.fetch(url, policy(2, 0))
.await;
assert!(matches!(result, Err(Error::BlockedUrl)));
}
#[tokio::test]
async fn validation_off_fetches_private_hosts_and_follows_redirects() {
let (url, server, _) = serve_named(
"localhost",
vec![
b"HTTP/1.1 302 Found\r\nLocation: /moved\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
OK_RESPONSE,
],
)
.await;
let off = UrlPolicy {
validate: false,
allowed_hosts: Vec::new(),
};
let media = fetcher(UNROUTABLE, HashSet::from(["localhost"]), off, false)
.fetch(url, policy(2, 1))
.await
.expect("unvalidated download succeeds");
let requests = server.await.expect("server completes");
assert_eq!(media.bytes, b"ok");
assert!(requests[1].starts_with("GET /moved "));
}
#[tokio::test]
async fn proxied_urls_skip_the_pinned_resolver_but_keep_the_address_check() {
let (url, server, _) = serve_named("localhost", vec![OK_RESPONSE]).await;
let media = fetcher(UNROUTABLE, HashSet::new(), UrlPolicy::default(), true)
.fetch(url.clone(), policy(2, 0))
.await
.expect("public host behind a proxy downloads");
server.await.expect("server completes");
assert_eq!(media.bytes, b"ok");
let blocked = fetcher(
UNROUTABLE,
HashSet::from(["localhost"]),
UrlPolicy::default(),
true,
)
.fetch(url, policy(2, 0))
.await;
assert!(matches!(blocked, Err(Error::BlockedUrl)));
}
}

210
litellm-rust/crates/http/src/outbound.rs Normal file
View file

@ -0,0 +1,210 @@
//! The request a route hands to the transport. The body is serialized once,
//! when the request is built, and a [`RequestSigner`] sees those exact bytes.
//!
//! Host hooks may rewrite the wire request (redaction, guardrails) and a
//! signature such as AWS SigV4 covers the body, so a route builds this after
//! its hooks ran and cannot change or re-serialize it afterwards.
use std::time::Duration;
use serde::Serialize;
use crate::{
Error,
request::{HeaderPolicy, has_header, with_headers},
};
#[derive(Clone, Copy, Debug)]
pub struct UnsignedRequest<'a> {
pub url: &'a str,
pub headers: &'a [(String, String)],
pub body: &'a [u8],
}
/// Returns the headers to add to the request; it never sees a mutable request.
pub trait RequestSigner: Send + Sync {
fn sign(&self, request: UnsignedRequest<'_>) -> Result<Vec<(String, String)>, Error>;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OutboundRequest {
url: String,
headers: Vec<(String, String)>,
body: Vec<u8>,
timeout: Option<Duration>,
}
impl OutboundRequest {
pub fn json(
url: String,
headers: Vec<(String, String)>,
body: &impl Serialize,
timeout: Option<Duration>,
) -> Result<Self, Error> {
Self::build(url, headers, body, timeout, None)
}
pub fn signed_json(
url: String,
headers: Vec<(String, String)>,
body: &impl Serialize,
timeout: Option<Duration>,
signer: &dyn RequestSigner,
) -> Result<Self, Error> {
Self::build(url, headers, body, timeout, Some(signer))
}
fn build(
url: String,
headers: Vec<(String, String)>,
body: &impl Serialize,
timeout: Option<Duration>,
signer: Option<&dyn RequestSigner>,
) -> Result<Self, Error> {
let body =
serde_json::to_vec(body).map_err(|error| Error::RequestBody(error.to_string()))?;
let content_type = (!has_header(&headers, "content-type"))
.then(|| ("content-type".to_string(), "application/json".to_string()));
let unsigned: Vec<(String, String)> = headers.into_iter().chain(content_type).collect();
let signature = signer
.map(|signer| {
signer.sign(UnsignedRequest {
url: &url,
headers: &unsigned,
body: &body,
})
})
.transpose()?
.unwrap_or_default();
Ok(Self {
url,
headers: unsigned.into_iter().chain(signature).collect(),
body,
timeout,
})
}
pub fn url(&self) -> &str {
&self.url
}
pub fn headers(&self) -> &[(String, String)] {
&self.headers
}
pub fn header(&self, name: &str) -> Option<&str> {
self.headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case(name))
.map(|(_, value)| value.as_str())
}
pub fn body(&self) -> &[u8] {
&self.body
}
pub fn timeout(&self) -> Option<Duration> {
self.timeout
}
pub async fn send(self, client: &reqwest::Client) -> Result<reqwest::Response, reqwest::Error> {
let builder = with_headers(
client.post(&self.url).body(self.body),
&self.headers,
HeaderPolicy::All,
);
match self.timeout {
Some(timeout) => builder.timeout(timeout),
None => builder,
}
.send()
.await
}
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use serde_json::json;
use super::*;
#[derive(Default)]
struct Recording(Mutex<Vec<u8>>);
impl RequestSigner for Recording {
fn sign(&self, request: UnsignedRequest<'_>) -> Result<Vec<(String, String)>, Error> {
*self.0.lock().unwrap() = request.body.to_vec();
Ok(vec![("authorization".into(), "signed".into())])
}
}
#[test]
fn the_signer_sees_exactly_the_bytes_that_are_sent() {
let signer = Recording::default();
let request = OutboundRequest::signed_json(
"https://provider.test/".into(),
vec![("x-caller".into(), "kept".into())],
&json!({"b": 1, "a": [true, null]}),
None,
&signer,
)
.unwrap();
assert_eq!(request.body(), signer.0.lock().unwrap().as_slice());
assert_eq!(request.header("authorization"), Some("signed"));
assert_eq!(request.header("x-caller"), Some("kept"));
}
#[test]
fn the_content_type_is_part_of_what_the_signer_sees() {
struct RequiresContentType;
impl RequestSigner for RequiresContentType {
fn sign(&self, request: UnsignedRequest<'_>) -> Result<Vec<(String, String)>, Error> {
has_header(request.headers, "content-type")
.then(Vec::new)
.ok_or_else(|| Error::Signature("content-type was not signed".into()))
}
}
let defaulted = OutboundRequest::signed_json(
"u".into(),
Vec::new(),
&json!({}),
None,
&RequiresContentType,
)
.unwrap();
assert_eq!(defaulted.header("content-type"), Some("application/json"));
let provider = OutboundRequest::signed_json(
"u".into(),
vec![("Content-Type".into(), "application/x-amz-json-1.1".into())],
&json!({}),
None,
&RequiresContentType,
)
.unwrap();
assert_eq!(
provider.header("content-type"),
Some("application/x-amz-json-1.1")
);
assert_eq!(provider.headers().len(), 1);
}
#[test]
fn a_signer_failure_produces_no_request() {
struct Refuses;
impl RequestSigner for Refuses {
fn sign(&self, _request: UnsignedRequest<'_>) -> Result<Vec<(String, String)>, Error> {
Err(Error::ComputedHeader("authorization".into()))
}
}
assert_eq!(
OutboundRequest::signed_json("u".into(), Vec::new(), &json!({}), None, &Refuses),
Err(Error::ComputedHeader("authorization".into()))
);
}
}

379
litellm-rust/crates/http/src/pool.rs Normal file
View file

@ -0,0 +1,379 @@
use std::{
collections::HashMap,
sync::{Arc, Mutex, MutexGuard, PoisonError},
time::{Duration, Instant},
};
use reqwest::dns::Resolve;
use crate::{config::HttpClientConfig, error::Error, proxy::EnvironmentProxies};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ClientVariant {
Provider,
NoRedirect,
Media,
UnpinnedMedia,
}
const CLIENT_TTL: Duration = Duration::from_secs(3600);
struct PooledClient {
client: reqwest::Client,
built_at: Instant,
}
type Clients = HashMap<(HttpClientConfig, ClientVariant), PooledClient>;
pub struct HttpClientPool {
media_resolver: Arc<dyn Resolve>,
ttl: Duration,
clients: Mutex<Clients>,
}
impl HttpClientPool {
pub fn new(media_resolver: Arc<dyn Resolve>) -> Self {
Self::with_ttl(media_resolver, CLIENT_TTL)
}
pub fn with_ttl(media_resolver: Arc<dyn Resolve>, ttl: Duration) -> Self {
Self {
media_resolver,
ttl,
clients: Mutex::default(),
}
}
pub fn client(
&self,
config: &HttpClientConfig,
variant: ClientVariant,
) -> Result<reqwest::Client, Error> {
let effective = match variant {
ClientVariant::Media => HttpClientConfig {
client_certificate: None,
proxies: EnvironmentProxies::default(),
..config.clone()
},
ClientVariant::UnpinnedMedia => HttpClientConfig {
client_certificate: None,
..config.clone()
},
ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(),
};
let key = (effective, variant);
if let Some(pooled) = self.lock().get(&key)
&& pooled.built_at.elapsed() < self.ttl
{
return Ok(pooled.client.clone());
}
let client = self
.apply(variant, reqwest::ClientBuilder::try_from(&key.0)?)
.build()?;
self.lock().insert(
key,
PooledClient {
client: client.clone(),
built_at: Instant::now(),
},
);
Ok(client)
}
fn lock(&self) -> MutexGuard<'_, Clients> {
self.clients.lock().unwrap_or_else(PoisonError::into_inner)
}
fn apply(
&self,
variant: ClientVariant,
builder: reqwest::ClientBuilder,
) -> reqwest::ClientBuilder {
match variant {
ClientVariant::Provider => builder,
ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => {
builder.redirect(reqwest::redirect::Policy::none())
}
ClientVariant::Media => builder
.redirect(reqwest::redirect::Policy::none())
.dns_resolver2(Arc::clone(&self.media_resolver)),
}
}
}
#[cfg(test)]
mod tests {
use std::{
net::SocketAddr,
sync::atomic::{AtomicUsize, Ordering},
time::Duration,
};
use reqwest::dns::{Addrs, Name, Resolving};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
};
use super::*;
use crate::{HttpSettings, Resolution, Verify};
struct FixedResolver(SocketAddr);
impl Resolve for FixedResolver {
fn resolve(&self, _: Name) -> Resolving {
let addrs: Addrs = Box::new(std::iter::once(self.0));
Box::pin(std::future::ready(Ok(addrs)))
}
}
fn pool() -> HttpClientPool {
HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())))
}
fn config(user_agent: &str) -> HttpClientConfig {
HttpClientConfig {
user_agent: Some(user_agent.into()),
..Resolution::from(&HttpSettings::default()).config
}
}
fn proxied_through(proxy: &str) -> EnvironmentProxies {
let proxy = proxy.to_owned();
EnvironmentProxies::from_environment(&move |name: &str| {
(name == "HTTP_PROXY").then(|| proxy.clone())
})
}
async fn serve(
status_line: &'static str,
) -> (SocketAddr, Arc<AtomicUsize>, Arc<Mutex<Vec<String>>>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let address = listener.local_addr().unwrap();
let connections = Arc::new(AtomicUsize::new(0));
let requests = Arc::new(Mutex::new(Vec::new()));
let (accepted, seen) = (Arc::clone(&connections), Arc::clone(&requests));
tokio::spawn(async move {
loop {
let (mut socket, _) = listener.accept().await.unwrap();
accepted.fetch_add(1, Ordering::SeqCst);
let seen = Arc::clone(&seen);
tokio::spawn(async move {
let mut buffer = vec![0u8; 4096];
while let Ok(read) = socket.read(&mut buffer).await {
if read == 0 {
return;
}
seen.lock()
.unwrap()
.push(String::from_utf8_lossy(&buffer[..read]).into_owned());
let response = format!(
"{status_line}\r\nLocation: /elsewhere\r\nContent-Length: 0\r\n\r\n"
);
if socket.write_all(response.as_bytes()).await.is_err() {
return;
}
}
});
}
});
(address, connections, requests)
}
async fn get(
pool: &HttpClientPool,
config: &HttpClientConfig,
variant: ClientVariant,
url: &str,
) -> reqwest::Response {
pool.client(config, variant)
.unwrap()
.get(url)
.timeout(Duration::from_secs(5))
.send()
.await
.unwrap()
}
#[tokio::test]
async fn clients_are_shared_per_config_and_variant() {
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
let url = format!("http://{address}");
let pool = pool();
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 1);
get(&pool, &config("a"), ClientVariant::NoRedirect, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 2);
get(&pool, &config("b"), ClientVariant::Provider, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 3);
}
#[tokio::test]
async fn provider_clients_route_through_the_resolved_proxy_not_the_process_environment() {
let (proxy, connections, requests) = serve("HTTP/1.1 204 No Content").await;
let config = HttpClientConfig {
proxies: proxied_through(&format!("http://user:secret@{proxy}")),
..config("a")
};
let response = get(
&pool(),
&config,
ClientVariant::Provider,
"http://upstream.invalid/v1/ocr",
)
.await;
assert_eq!(response.status(), 204);
assert_eq!(connections.load(Ordering::SeqCst), 1);
let request = requests.lock().unwrap().concat();
assert!(request.starts_with("GET http://upstream.invalid/v1/ocr HTTP/1.1"));
assert!(request.contains("proxy-authorization: Basic dXNlcjpzZWNyZXQ="));
}
#[tokio::test]
async fn no_proxy_hosts_bypass_the_resolved_proxy() {
let (upstream, _, _) = serve("HTTP/1.1 204 No Content").await;
let (proxy, proxy_connections, _) = serve("HTTP/1.1 502 Bad Gateway").await;
let config = HttpClientConfig {
proxies: EnvironmentProxies::from_environment(&move |name: &str| match name {
"HTTP_PROXY" => Some(format!("http://{proxy}")),
"NO_PROXY" => Some("127.0.0.1".into()),
_ => None,
}),
..config("a")
};
let response = get(
&pool(),
&config,
ClientVariant::Provider,
&format!("http://{upstream}/v1/ocr"),
)
.await;
assert_eq!(response.status(), 204);
assert_eq!(proxy_connections.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn expired_clients_are_rebuilt() {
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
let url = format!("http://{address}");
let pool = HttpClientPool::with_ttl(
Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())),
Duration::ZERO,
);
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
get(&pool, &config("a"), ClientVariant::Provider, &url).await;
assert_eq!(connections.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn media_clients_are_shared_across_proxy_settings_they_never_use() {
let (address, connections, _) = serve("HTTP/1.1 204 No Content").await;
let pool = HttpClientPool::new(Arc::new(FixedResolver(address)));
let url = format!("http://media.invalid:{}/doc", address.port());
for proxies in [
proxied_through("http://proxy.invalid:3128"),
EnvironmentProxies::default(),
] {
let config = HttpClientConfig {
proxies,
..config("a")
};
get(&pool, &config, ClientVariant::Media, &url).await;
}
assert_eq!(connections.load(Ordering::SeqCst), 1);
}
#[test]
fn media_variant_never_loads_the_client_certificate() {
let pool = pool();
let with_identity = HttpClientConfig {
client_certificate: Some(std::env::temp_dir().join("litellm-http-absent-client.pem")),
..config("a")
};
assert!(
pool.client(&with_identity, ClientVariant::Provider)
.is_err()
);
assert!(pool.client(&with_identity, ClientVariant::Media).is_ok());
assert!(
pool.client(&with_identity, ClientVariant::UnpinnedMedia)
.is_ok()
);
}
#[test]
fn build_failures_are_not_cached() {
let pool = pool();
let missing = HttpClientConfig {
verify: Verify::CaBundle(std::env::temp_dir().join("litellm-http-absent.pem")),
..config("a")
};
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
assert!(pool.client(&missing, ClientVariant::Provider).is_err());
assert!(pool.client(&config("a"), ClientVariant::Provider).is_ok());
}
#[tokio::test]
async fn provider_client_sends_the_configured_user_agent_over_http1() {
let (address, _, requests) = serve("HTTP/1.1 204 No Content").await;
let response = get(
&pool(),
&config("litellm-test/9"),
ClientVariant::Provider,
&format!("http://{address}"),
)
.await;
assert_eq!(response.status(), 204);
assert_eq!(response.version(), reqwest::Version::HTTP_11);
let request = requests.lock().unwrap()[0].clone();
assert!(request.contains("user-agent: litellm-test/9"), "{request}");
}
#[tokio::test]
async fn no_redirect_variant_returns_the_redirect_instead_of_following_it() {
let (address, _, _) = serve("HTTP/1.1 302 Found").await;
let response = get(
&pool(),
&config("a"),
ClientVariant::NoRedirect,
&format!("http://{address}"),
)
.await;
assert_eq!(response.status(), 302);
assert_eq!(response.headers()["location"], "/elsewhere");
}
#[tokio::test]
async fn unpinned_media_variant_uses_the_system_resolver_and_returns_redirects() {
let (address, _, _) = serve("HTTP/1.1 302 Found").await;
let pool = HttpClientPool::new(Arc::new(FixedResolver(([192, 0, 2, 1], 80).into())));
let response = get(
&pool,
&config("a"),
ClientVariant::UnpinnedMedia,
&format!("http://localhost:{}/doc", address.port()),
)
.await;
assert_eq!(response.status(), 302);
}
#[tokio::test]
async fn media_variant_resolves_through_the_injected_resolver() {
let (address, _, requests) = serve("HTTP/1.1 204 No Content").await;
let pool = HttpClientPool::new(Arc::new(FixedResolver(address)));
let url = format!("http://media.invalid:{}/doc", address.port());
let response = get(&pool, &config("a"), ClientVariant::Media, &url).await;
assert_eq!(response.status(), 204);
assert!(requests.lock().unwrap()[0].contains("host: media.invalid"));
assert!(
pool.client(&config("a"), ClientVariant::Provider)
.unwrap()
.get(&url)
.timeout(Duration::from_secs(5))
.send()
.await
.is_err()
);
}
}

164
litellm-rust/crates/http/src/proxy.rs Normal file
View file

@ -0,0 +1,164 @@
use hyper_util::client::proxy::matcher::Matcher;
use litellm_core_utils::settings::Lookup;
use veil::Redact;
#[derive(Clone, Redact, Default, PartialEq, Eq, Hash)]
pub struct EnvironmentProxies {
#[redact]
all: String,
#[redact]
http: String,
#[redact]
https: String,
no: String,
}
impl EnvironmentProxies {
pub fn from_environment(env: &impl Lookup) -> Self {
Self::resolve(env, cfg!(windows))
}
fn resolve(env: &impl Lookup, names_ignore_case: bool) -> Self {
let lowercase_first = |upper: Option<&str>, lower: &str| {
env.get(lower)
.or_else(|| upper.and_then(|name| env.truthy(name)))
.unwrap_or_default()
};
let is_cgi = env.get("REQUEST_METHOD").is_some();
Self {
all: lowercase_first(Some("ALL_PROXY"), "all_proxy"),
http: if is_cgi && names_ignore_case {
String::new()
} else {
lowercase_first((!is_cgi).then_some("HTTP_PROXY"), "http_proxy")
},
https: lowercase_first(Some("HTTPS_PROXY"), "https_proxy"),
no: lowercase_first(Some("NO_PROXY"), "no_proxy"),
}
}
pub(crate) fn matcher(&self) -> impl Fn(&reqwest::Url) -> bool + Send + Sync + use<> {
let matcher = Matcher::builder()
.all(self.all.clone())
.http(self.http.clone())
.https(self.https.clone())
.no(self.no.clone())
.build();
move |url| {
url.as_str()
.parse::<http::Uri>()
.is_ok_and(|uri| matcher.intercept(&uri).is_some())
}
}
pub(crate) fn reqwest_proxies(&self) -> Vec<reqwest::Proxy> {
let no_proxy = reqwest::NoProxy::from_string(&self.no);
[
reqwest::Proxy::http(self.http.as_str()),
reqwest::Proxy::https(self.https.as_str()),
reqwest::Proxy::all(self.all.as_str()),
]
.into_iter()
.filter_map(Result::ok)
.map(|proxy| proxy.no_proxy(no_proxy.clone()))
.collect()
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
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())
}
}
fn url(value: &str) -> reqwest::Url {
reqwest::Url::parse(value).unwrap()
}
#[rstest]
#[case::http_only(&[("HTTP_PROXY", "http://proxy:3128")], "http://api.test/", true)]
#[case::http_proxy_skips_https(&[("HTTP_PROXY", "http://proxy:3128")], "https://api.test/", false)]
#[case::all_covers_https(&[("ALL_PROXY", "http://proxy:3128")], "https://api.test/", true)]
#[case::lowercase(&[("https_proxy", "http://proxy:3128")], "https://api.test/", true)]
#[case::no_proxy_bypass(&[("HTTPS_PROXY", "http://proxy:3128"), ("NO_PROXY", "api.test")], "https://api.test/", false)]
fn proxies_follow_the_injected_environment(
#[case] env: &'static [(&'static str, &'static str)],
#[case] target: &str,
#[case] expected: bool,
) {
let proxies = EnvironmentProxies::from_environment(&env_of(env));
assert_eq!(proxies.matcher()(&url(target)), expected);
}
#[rstest]
#[case::lowercase_wins(&[("HTTPS_PROXY", "http://upper:3128"), ("https_proxy", "http://lower:3128")], &[("https_proxy", "http://lower:3128")])]
#[case::empty_uppercase_falls_through_to_lowercase(&[("HTTPS_PROXY", ""), ("https_proxy", "http://lower:3128")], &[("https_proxy", "http://lower:3128")])]
#[case::empty_uppercase_alone_is_unset(&[("HTTPS_PROXY", "")], &[])]
#[case::empty_lowercase_clears_the_uppercase_value(&[("HTTPS_PROXY", "http://upper:3128"), ("https_proxy", "")], &[])]
#[case::lowercase_no_proxy_wins(&[("NO_PROXY", "upper.test"), ("no_proxy", "lower.test")], &[("no_proxy", "lower.test")])]
#[case::cgi_forgets_the_client_settable_http_proxy(&[("REQUEST_METHOD", "GET"), ("HTTP_PROXY", "http://attacker:3128")], &[])]
#[case::cgi_keeps_lowercase_http_proxy(&[("REQUEST_METHOD", "GET"), ("HTTP_PROXY", "http://attacker:3128"), ("http_proxy", "http://lower:3128")], &[("http_proxy", "http://lower:3128")])]
#[case::cgi_keeps_every_other_variable(&[("REQUEST_METHOD", "GET"), ("HTTPS_PROXY", "http://proxy:3128"), ("ALL_PROXY", "http://all:3128"), ("NO_PROXY", "internal.test")], &[("HTTPS_PROXY", "http://proxy:3128"), ("ALL_PROXY", "http://all:3128"), ("NO_PROXY", "internal.test")])]
fn variables_resolve_like_urllib_getproxies_environment(
#[case] env: &'static [(&'static str, &'static str)],
#[case] equivalent: &'static [(&'static str, &'static str)],
) {
assert_eq!(
EnvironmentProxies::from_environment(&env_of(env)),
EnvironmentProxies::from_environment(&env_of(equivalent))
);
}
#[test]
fn cgi_drops_http_proxy_entirely_where_variable_names_ignore_case() {
let windows_env = |name: &str| match name.to_ascii_uppercase().as_str() {
"REQUEST_METHOD" => Some("GET".to_string()),
"HTTP_PROXY" => Some("http://attacker:3128".to_string()),
"HTTPS_PROXY" => Some("http://proxy:3128".to_string()),
_ => None,
};
let proxies = EnvironmentProxies::resolve(&windows_env, true);
assert!(!proxies.matcher()(&url("http://api.test/")));
assert!(proxies.matcher()(&url("https://api.test/")));
assert!(EnvironmentProxies::resolve(&windows_env, false).matcher()(
&url("http://api.test/")
));
}
#[test]
fn a_cgi_request_still_proxies_https_through_the_configured_proxy() {
let proxies = EnvironmentProxies::from_environment(&env_of(&[
("REQUEST_METHOD", "GET"),
("HTTPS_PROXY", "http://proxy:3128"),
]));
assert!(proxies.matcher()(&url("https://api.test/")));
}
#[test]
fn debug_output_hides_proxy_credentials_but_shows_which_variables_are_set() {
let proxies = EnvironmentProxies::from_environment(&env_of(&[
("HTTPS_PROXY", "http://operator:hunter2@proxy.corp:3128"),
("NO_PROXY", "internal.test"),
]));
let debug = format!("{proxies:?}");
assert!(!debug.contains("hunter2") && !debug.contains("operator"));
assert!(debug.contains("internal.test"));
assert_ne!(debug, format!("{:?}", EnvironmentProxies::default()));
}
#[test]
fn an_empty_environment_proxies_nothing() {
let proxies = EnvironmentProxies::from_environment(&env_of(&[]));
assert_eq!(proxies, EnvironmentProxies::default());
assert!(proxies.reqwest_proxies().is_empty());
}
}

20
litellm-rust/crates/llms/src/custom_httpx/http_handler.rs → litellm-rust/crates/http/src/request.rs
View file

@ -13,20 +13,12 @@ use serde_json::{Map, Value};
/// before truncation, so provider bodies are bounded and data-minimized.
const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256;
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub enum HeaderPolicy<'a> {
All,
Only(&'a [&'a str]),
Except(&'a [&'a str]),
}
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub fn with_headers(
builder: reqwest::RequestBuilder,
headers: &[(String, String)],
@ -107,18 +99,6 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
})
}
#[allow(
dead_code,
reason = "used by the OCR architecture in the next stacked PR"
)]
pub fn deserialize_optional_param<'de, D, T>(deserializer: D) -> Result<Option<Option<T>>, D::Error>
where
D: serde::Deserializer<'de>,
T: serde::Deserialize<'de>,
{
<Option<T> as serde::Deserialize>::deserialize(deserializer).map(Some)
}
#[cfg(test)]
mod tests {
use serde_json::json;

419
litellm-rust/crates/http/src/settings.rs Normal file
View file

@ -0,0 +1,419 @@
use std::{
path::{Path, PathBuf},
time::Duration,
};
use litellm_core_utils::settings::{Layer, Lookup, merge};
use crate::proxy::EnvironmentProxies;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum SslVerify {
Enabled,
Disabled,
CaBundle(PathBuf),
}
impl SslVerify {
pub fn parse(value: &str) -> Self {
match value.trim().to_ascii_lowercase().as_str() {
"true" => Self::Enabled,
"false" => Self::Disabled,
_ => Self::CaBundle(PathBuf::from(value)),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct TcpKeepalive {
pub idle: Duration,
pub interval: Duration,
pub retries: u32,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct HttpSettingsLayer {
pub ssl_verify: Option<SslVerify>,
pub ssl_cert_file: Option<PathBuf>,
pub ssl_certificate: Option<PathBuf>,
pub ssl_security_level: Option<String>,
pub ssl_ecdh_curve: Option<String>,
pub force_ipv4: Option<bool>,
pub http2: Option<bool>,
pub aiohttp_trust_env: Option<bool>,
pub disable_aiohttp_trust_env: Option<bool>,
pub disable_aiohttp_transport: Option<bool>,
pub user_agent: Option<String>,
pub tcp_keepalive: Option<TcpKeepalive>,
pub pool_idle_timeout: Option<Duration>,
pub proxies: Option<EnvironmentProxies>,
}
impl HttpSettingsLayer {
pub fn from_environment(env: &impl Lookup) -> Self {
let seconds = |name: &str, default: u32| {
Duration::from_secs(u64::from(env.parsed::<u32>(name).unwrap_or(default)))
};
Self {
ssl_verify: env.get("SSL_VERIFY").map(|value| SslVerify::parse(&value)),
ssl_cert_file: env.truthy("SSL_CERT_FILE").map(PathBuf::from),
ssl_certificate: env.get("SSL_CERTIFICATE").map(PathBuf::from),
ssl_security_level: env.get("SSL_SECURITY_LEVEL"),
ssl_ecdh_curve: env.get("SSL_ECDH_CURVE"),
force_ipv4: None,
http2: env.enabled("LITELLM_HTTP2"),
aiohttp_trust_env: env.enabled("AIOHTTP_TRUST_ENV"),
disable_aiohttp_trust_env: env.enabled("DISABLE_AIOHTTP_TRUST_ENV"),
disable_aiohttp_transport: env.enabled("DISABLE_AIOHTTP_TRANSPORT"),
user_agent: env.get("LITELLM_USER_AGENT"),
tcp_keepalive: env.enabled("AIOHTTP_SO_KEEPALIVE").map(|_| TcpKeepalive {
idle: seconds("AIOHTTP_TCP_KEEPIDLE", 60),
interval: seconds("AIOHTTP_TCP_KEEPINTVL", 30),
retries: env.parsed("AIOHTTP_TCP_KEEPCNT").unwrap_or(5),
}),
pool_idle_timeout: env
.parsed::<u32>("AIOHTTP_KEEPALIVE_TIMEOUT")
.map(|timeout| Duration::from_secs(u64::from(timeout))),
proxies: Some(EnvironmentProxies::from_environment(env))
.filter(|proxies| *proxies != EnvironmentProxies::default()),
}
}
}
impl Layer for HttpSettingsLayer {
fn or(self, lower: Self) -> Self {
Self {
ssl_verify: self.ssl_verify.or(lower.ssl_verify),
ssl_cert_file: self.ssl_cert_file.or(lower.ssl_cert_file),
ssl_certificate: self.ssl_certificate.or(lower.ssl_certificate),
ssl_security_level: self.ssl_security_level.or(lower.ssl_security_level),
ssl_ecdh_curve: self.ssl_ecdh_curve.or(lower.ssl_ecdh_curve),
force_ipv4: self.force_ipv4.or(lower.force_ipv4),
http2: self.http2.or(lower.http2),
aiohttp_trust_env: self.aiohttp_trust_env.or(lower.aiohttp_trust_env),
disable_aiohttp_trust_env: self
.disable_aiohttp_trust_env
.or(lower.disable_aiohttp_trust_env),
disable_aiohttp_transport: self
.disable_aiohttp_transport
.or(lower.disable_aiohttp_transport),
user_agent: self.user_agent.or(lower.user_agent),
tcp_keepalive: self.tcp_keepalive.or(lower.tcp_keepalive),
pool_idle_timeout: self.pool_idle_timeout.or(lower.pool_idle_timeout),
proxies: self.proxies.or(lower.proxies),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HttpSettings {
pub ssl_verify: Option<SslVerify>,
pub ssl_cert_file: Option<PathBuf>,
pub ssl_certificate: Option<PathBuf>,
pub ssl_security_level: Option<String>,
pub ssl_ecdh_curve: Option<String>,
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,
}
impl Default for HttpSettings {
fn default() -> Self {
Self {
ssl_verify: None,
ssl_cert_file: None,
ssl_certificate: None,
ssl_security_level: None,
ssl_ecdh_curve: None,
force_ipv4: false,
http2: false,
user_agent: None,
trust_proxy_env: true,
proxies: EnvironmentProxies::default(),
connect_timeout: Duration::from_secs(10),
tcp_keepalive: None,
pool_idle_timeout: Duration::from_secs(120),
}
}
}
impl HttpSettings {
pub fn from_layers(
highest_precedence_first: impl IntoIterator<Item = HttpSettingsLayer>,
) -> Self {
let merged = merge(highest_precedence_first);
let defaults = Self::default();
let http2 = merged.http2.unwrap_or(defaults.http2);
Self {
ssl_verify: merged.ssl_verify,
ssl_cert_file: merged.ssl_cert_file,
ssl_certificate: merged
.ssl_certificate
.filter(|path| !path.as_os_str().is_empty()),
ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()),
ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()),
force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4),
http2,
user_agent: merged.user_agent,
trust_proxy_env: !merged.disable_aiohttp_trust_env.unwrap_or(false)
|| merged.aiohttp_trust_env.unwrap_or(false)
|| merged.disable_aiohttp_transport.unwrap_or(false)
|| http2,
tcp_keepalive: merged.tcp_keepalive,
pool_idle_timeout: merged
.pool_idle_timeout
.unwrap_or(defaults.pool_idle_timeout),
proxies: merged.proxies.unwrap_or_default(),
..defaults
}
}
pub fn without_missing_files(self, exists: &dyn Fn(&Path) -> bool) -> Self {
Self {
ssl_verify: match self.ssl_verify {
Some(SslVerify::CaBundle(path)) if !exists(&path) => Some(SslVerify::Enabled),
other => other,
},
ssl_cert_file: self.ssl_cert_file.filter(|path| exists(path)),
..self
}
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
fn no_env(_: &str) -> Option<String> {
None
}
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())
}
}
#[rstest]
#[case("true", SslVerify::Enabled)]
#[case(" True ", SslVerify::Enabled)]
#[case("FALSE", SslVerify::Disabled)]
#[case("/etc/ssl/bundle.pem", SslVerify::CaBundle("/etc/ssl/bundle.pem".into()))]
fn ssl_verify_parses_bools_and_treats_anything_else_as_a_bundle_path(
#[case] value: &str,
#[case] expected: SslVerify,
) {
assert_eq!(SslVerify::parse(value), expected);
}
#[test]
fn higher_layers_override_lower_ones() {
let configured = HttpSettingsLayer {
ssl_verify: Some(SslVerify::Enabled),
ssl_certificate: Some("/configured/client.pem".into()),
ssl_security_level: Some("configured".into()),
user_agent: Some("configured/1".into()),
..HttpSettingsLayer::default()
};
let environment = HttpSettingsLayer::from_environment(&env_of(&[
("SSL_VERIFY", "false"),
("SSL_CERT_FILE", "/env/roots.pem"),
("SSL_CERTIFICATE", "/env/client.pem"),
("SSL_SECURITY_LEVEL", "DEFAULT@SECLEVEL=1"),
("SSL_ECDH_CURVE", "X25519"),
("LITELLM_USER_AGENT", "env/2"),
]));
let settings = HttpSettings::from_layers([environment, configured]);
assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled));
assert_eq!(settings.ssl_cert_file, Some("/env/roots.pem".into()));
assert_eq!(settings.ssl_certificate, Some("/env/client.pem".into()));
assert_eq!(
settings.ssl_security_level.as_deref(),
Some("DEFAULT@SECLEVEL=1")
);
assert_eq!(settings.ssl_ecdh_curve.as_deref(), Some("X25519"));
assert_eq!(settings.user_agent.as_deref(), Some("env/2"));
}
#[test]
fn an_explicit_false_in_a_higher_layer_beats_a_lower_true() {
let higher = HttpSettingsLayer {
http2: Some(false),
force_ipv4: Some(false),
..HttpSettingsLayer::default()
};
let lower = HttpSettingsLayer {
http2: Some(true),
force_ipv4: Some(true),
..HttpSettingsLayer::default()
};
let settings = HttpSettings::from_layers([higher, lower]);
assert!(!settings.http2);
assert!(!settings.force_ipv4);
}
#[test]
fn an_empty_environment_is_an_empty_layer_so_lower_layers_and_defaults_apply() {
assert_eq!(
HttpSettingsLayer::from_environment(&no_env),
HttpSettingsLayer::default()
);
let configured = HttpSettingsLayer {
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
http2: Some(true),
user_agent: Some("configured/1".into()),
..HttpSettingsLayer::default()
};
assert_eq!(
HttpSettings::from_layers([HttpSettingsLayer::default(), configured]),
HttpSettings {
ssl_verify: Some(SslVerify::CaBundle("/configured/roots.pem".into())),
http2: true,
user_agent: Some("configured/1".into()),
..HttpSettings::default()
}
);
assert_eq!(HttpSettings::from_layers([]), HttpSettings::default());
}
#[test]
fn empty_environment_values_clear_the_setting_like_python_truthiness() {
let configured = HttpSettingsLayer {
ssl_certificate: Some("/configured/client.pem".into()),
ssl_security_level: Some("configured".into()),
ssl_ecdh_curve: Some("X25519".into()),
..HttpSettingsLayer::default()
};
let environment = HttpSettingsLayer::from_environment(&env_of(&[
("SSL_CERTIFICATE", ""),
("SSL_SECURITY_LEVEL", ""),
("SSL_ECDH_CURVE", ""),
]));
let settings = HttpSettings::from_layers([environment, configured]);
assert_eq!(settings.ssl_certificate, None);
assert_eq!(settings.ssl_security_level, None);
assert_eq!(settings.ssl_ecdh_curve, None);
}
#[test]
fn socket_keepalive_follows_the_aiohttp_variables_with_python_defaults() {
let tuned = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(&[
("AIOHTTP_SO_KEEPALIVE", "True"),
("AIOHTTP_TCP_KEEPIDLE", "45"),
("AIOHTTP_KEEPALIVE_TIMEOUT", "30"),
]))]);
assert_eq!(
tuned.tcp_keepalive,
Some(TcpKeepalive {
idle: Duration::from_secs(45),
interval: Duration::from_secs(30),
retries: 5,
})
);
assert_eq!(tuned.pool_idle_timeout, Duration::from_secs(30));
}
#[test]
fn socket_keepalive_stays_off_unless_enabled() {
let settings = HttpSettings::from_layers([HttpSettingsLayer::from_environment(&env_of(
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
))]);
assert_eq!(settings.tcp_keepalive, None);
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
}
fn proxy_flags(
aiohttp_trust_env: bool,
disable_aiohttp_trust_env: bool,
disable_aiohttp_transport: bool,
http2: bool,
) -> HttpSettingsLayer {
HttpSettingsLayer {
aiohttp_trust_env: Some(aiohttp_trust_env),
disable_aiohttp_trust_env: Some(disable_aiohttp_trust_env),
disable_aiohttp_transport: Some(disable_aiohttp_transport),
http2: Some(http2),
..HttpSettingsLayer::default()
}
}
#[rstest]
#[case::aiohttp_default(proxy_flags(false, false, false, false), true)]
#[case::aiohttp_opted_out(proxy_flags(false, true, false, false), false)]
#[case::session_trust_env_beats_opt_out(proxy_flags(true, true, false, false), true)]
#[case::http2_uses_httpx(proxy_flags(false, true, false, true), true)]
#[case::aiohttp_disabled(proxy_flags(false, true, true, false), true)]
fn environment_proxies_apply_unless_the_aiohttp_transport_opts_out(
#[case] layer: HttpSettingsLayer,
#[case] expected: bool,
) {
assert_eq!(HttpSettings::from_layers([layer]).trust_proxy_env, expected);
}
#[test]
fn a_proxy_opt_out_in_one_source_still_yields_to_trust_env_from_another() {
let environment =
HttpSettingsLayer::from_environment(&env_of(&[("DISABLE_AIOHTTP_TRUST_ENV", "true")]));
let configured = HttpSettingsLayer {
aiohttp_trust_env: Some(true),
..HttpSettingsLayer::default()
};
assert!(!HttpSettings::from_layers([environment.clone()]).trust_proxy_env);
assert!(HttpSettings::from_layers([environment, configured]).trust_proxy_env);
}
#[test]
fn missing_files_fall_back_to_default_verification() {
let settings = HttpSettings {
ssl_verify: Some(SslVerify::CaBundle("/absent/roots.pem".into())),
ssl_cert_file: Some("/absent/env.pem".into()),
..HttpSettings::default()
}
.without_missing_files(&|_| false);
assert_eq!(settings.ssl_verify, Some(SslVerify::Enabled));
assert_eq!(settings.ssl_cert_file, None);
}
#[test]
fn existing_files_are_kept() {
let settings = HttpSettings {
ssl_verify: Some(SslVerify::CaBundle("/present/roots.pem".into())),
ssl_cert_file: Some("/present/env.pem".into()),
..HttpSettings::default()
};
assert_eq!(settings.clone().without_missing_files(&|_| true), settings);
}
#[rstest]
#[case("true", Some(true))]
#[case("True", Some(true))]
#[case("false", None)]
#[case("1", None)]
fn boolean_switches_only_turn_on_for_true(
#[case] value: &'static str,
#[case] expected: Option<bool>,
) {
let env = move |name: &str| match name {
"LITELLM_HTTP2"
| "AIOHTTP_TRUST_ENV"
| "DISABLE_AIOHTTP_TRANSPORT"
| "DISABLE_AIOHTTP_TRUST_ENV" => Some(value.to_string()),
_ => None,
};
let layer = HttpSettingsLayer::from_environment(&env);
assert_eq!(layer.http2, expected);
assert_eq!(layer.aiohttp_trust_env, expected);
assert_eq!(layer.disable_aiohttp_transport, expected);
assert_eq!(layer.disable_aiohttp_trust_env, expected);
}
}

411
litellm-rust/crates/http/src/tls.rs Normal file
View file

@ -0,0 +1,411 @@
use std::{fmt, path::Path, str::FromStr, sync::Arc};
use rustls::{
CipherSuite, ClientConfig, DigitallySignedStruct, RootCertStore, SignatureScheme,
client::danger::{HandshakeSignatureValid, ServerCertVerified, ServerCertVerifier},
crypto::{CryptoProvider, SupportedKxGroup, ring},
pki_types::{CertificateDer, PrivateKeyDer, ServerName, UnixTime, pem::PemObject},
};
use crate::{
config::{HttpClientConfig, Verify},
error::Error,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum KeyExchangeGroup {
X25519,
Secp256r1,
Secp384r1,
}
impl FromStr for KeyExchangeGroup {
type Err = Unsupported;
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name.trim().to_ascii_lowercase().as_str() {
"x25519" => Ok(Self::X25519),
"prime256v1" | "secp256r1" | "p-256" => Ok(Self::Secp256r1),
"secp384r1" | "p-384" => Ok(Self::Secp384r1),
_ => Err(Unsupported::EcdhCurve(name.to_owned())),
}
}
}
impl KeyExchangeGroup {
fn supported(self) -> &'static dyn SupportedKxGroup {
match self {
Self::X25519 => ring::kx_group::X25519,
Self::Secp256r1 => ring::kx_group::SECP256R1,
Self::Secp384r1 => ring::kx_group::SECP384R1,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum Tls12CipherSuite {
EcdheEcdsaAes128Gcm,
EcdheEcdsaAes256Gcm,
EcdheEcdsaChacha20,
EcdheRsaAes128Gcm,
EcdheRsaAes256Gcm,
EcdheRsaChacha20,
}
impl FromStr for Tls12CipherSuite {
type Err = Unsupported;
fn from_str(name: &str) -> Result<Self, Self::Err> {
match name {
"ECDHE-ECDSA-AES128-GCM-SHA256" => Ok(Self::EcdheEcdsaAes128Gcm),
"ECDHE-ECDSA-AES256-GCM-SHA384" => Ok(Self::EcdheEcdsaAes256Gcm),
"ECDHE-ECDSA-CHACHA20-POLY1305" => Ok(Self::EcdheEcdsaChacha20),
"ECDHE-RSA-AES128-GCM-SHA256" => Ok(Self::EcdheRsaAes128Gcm),
"ECDHE-RSA-AES256-GCM-SHA384" => Ok(Self::EcdheRsaAes256Gcm),
"ECDHE-RSA-CHACHA20-POLY1305" => Ok(Self::EcdheRsaChacha20),
_ => Err(Unsupported::CipherToken(name.to_owned())),
}
}
}
impl Tls12CipherSuite {
fn suite(self) -> CipherSuite {
match self {
Self::EcdheEcdsaAes128Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
Self::EcdheEcdsaAes256Gcm => CipherSuite::TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
Self::EcdheEcdsaChacha20 => CipherSuite::TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256,
Self::EcdheRsaAes128Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256,
Self::EcdheRsaAes256Gcm => CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384,
Self::EcdheRsaChacha20 => CipherSuite::TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, thiserror::Error)]
pub enum Unsupported {
#[error(
"ssl_ecdh_curve {0:?} is not supported: rustls with ring only offers X25519, prime256v1 and secp384r1, so the default key exchange groups are used"
)]
EcdhCurve(String),
#[error(
"ssl_security_level {0:?} is not supported: rustls has one fixed security level, comparable to OpenSSL level 2, so legacy servers that need a lower level cannot be reached"
)]
SecurityLevel(String),
#[error(
"ssl_security_level entry {0:?} is not supported: rustls only offers ECDHE AEAD cipher suites, so the entry is ignored"
)]
CipherToken(String),
}
#[derive(Default)]
pub(crate) struct CipherSelection {
pub(crate) tls12_cipher_suites: Option<Vec<Tls12CipherSuite>>,
pub(crate) unsupported: Vec<Unsupported>,
}
enum CipherToken {
Suite(Tls12CipherSuite),
EverySuite,
Ordering,
Unsupported(Unsupported),
}
impl From<&str> for CipherToken {
fn from(token: &str) -> Self {
match token {
"DEFAULT" | "ALL" | "HIGH" => Self::EverySuite,
"@STRENGTH" | "@SECLEVEL=2" => Self::Ordering,
level if level.starts_with("@SECLEVEL=") => {
Self::Unsupported(Unsupported::SecurityLevel(level.to_owned()))
}
name => name.parse().map_or_else(Self::Unsupported, Self::Suite),
}
}
}
impl From<&str> for CipherSelection {
fn from(value: &str) -> Self {
let tokens: Vec<CipherToken> = tokenize(value)
.iter()
.map(|token| CipherToken::from(token.as_str()))
.collect();
let every_suite = tokens
.iter()
.any(|token| matches!(token, CipherToken::EverySuite));
let mut suites: Vec<Tls12CipherSuite> = tokens
.iter()
.filter_map(|token| match token {
CipherToken::Suite(suite) => Some(*suite),
_ => None,
})
.collect();
suites.sort_unstable();
suites.dedup();
CipherSelection {
tls12_cipher_suites: (!every_suite && !suites.is_empty()).then_some(suites),
unsupported: tokens
.into_iter()
.filter_map(|token| match token {
CipherToken::Unsupported(unsupported) => Some(unsupported),
_ => None,
})
.collect(),
}
}
}
fn tokenize(value: &str) -> Vec<String> {
value
.split([':', ',', ' '])
.flat_map(|entry| match entry.split_once('@') {
Some((name, command)) => vec![name.to_owned(), format!("@{command}")],
None => vec![entry.to_owned()],
})
.filter(|token| !token.is_empty())
.collect()
}
impl TryFrom<&HttpClientConfig> for ClientConfig {
type Error = Error;
fn try_from(config: &HttpClientConfig) -> Result<Self, Self::Error> {
let base = ring::default_provider();
let provider = Arc::new(CryptoProvider {
kx_groups: config
.key_exchange_group
.map_or_else(|| base.kx_groups.clone(), |group| vec![group.supported()]),
cipher_suites: base
.cipher_suites
.iter()
.copied()
.filter(|suite| {
suite.tls13().is_some()
|| config.tls12_cipher_suites.as_ref().is_none_or(|allowed| {
allowed.iter().any(|a| a.suite() == suite.suite())
})
})
.collect(),
..base
});
let builder = ClientConfig::builder_with_provider(Arc::clone(&provider))
.with_safe_default_protocol_versions()
.map_err(|error| Error::Client(error.to_string()))?;
let verified = match &config.verify {
Verify::Disabled => builder
.dangerous()
.with_custom_certificate_verifier(Arc::new(NoVerification(provider))),
Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore {
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
}),
Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?),
};
let mut tls = match &config.client_certificate {
None => verified.with_no_client_auth(),
Some(path) => {
let (chain, key) = identity(path)?;
verified
.with_client_auth_cert(chain, key)
.map_err(|error| invalid_pem(path, error))?
}
};
tls.alpn_protocols = if config.http2 {
vec![b"h2".to_vec(), b"http/1.1".to_vec()]
} else {
vec![b"http/1.1".to_vec()]
};
Ok(tls)
}
}
fn bundle_roots(path: &Path) -> Result<RootCertStore, Error> {
let certificates = certificates(path)?;
if certificates.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
}
let mut store = RootCertStore::empty();
for certificate in certificates {
store
.add(certificate)
.map_err(|error| invalid_pem(path, error))?;
}
Ok(store)
}
fn identity(path: &Path) -> Result<(Vec<CertificateDer<'static>>, PrivateKeyDer<'static>), Error> {
let chain = certificates(path)?;
if chain.is_empty() {
return Err(invalid_pem(path, "no certificates found"));
}
let key =
PrivateKeyDer::from_pem_slice(&read(path)?).map_err(|error| invalid_pem(path, error))?;
Ok((chain, key))
}
fn certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, Error> {
CertificateDer::pem_slice_iter(&read(path)?)
.collect::<Result<_, _>>()
.map_err(|error| invalid_pem(path, error))
}
fn read(path: &Path) -> Result<Vec<u8>, Error> {
std::fs::read(path).map_err(|error| Error::Read {
path: path.to_path_buf(),
message: error.to_string(),
})
}
fn invalid_pem(path: &Path, message: impl fmt::Display) -> Error {
Error::InvalidPem {
path: path.to_path_buf(),
message: message.to_string(),
}
}
#[derive(Debug)]
struct NoVerification(Arc<CryptoProvider>);
impl ServerCertVerifier for NoVerification {
fn verify_server_cert(
&self,
_end_entity: &CertificateDer<'_>,
_intermediates: &[CertificateDer<'_>],
_server_name: &ServerName<'_>,
_ocsp_response: &[u8],
_now: UnixTime,
) -> Result<ServerCertVerified, rustls::Error> {
Ok(ServerCertVerified::assertion())
}
fn verify_tls12_signature(
&self,
_message: &[u8],
_cert: &CertificateDer<'_>,
_dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn verify_tls13_signature(
&self,
_message: &[u8],
_cert: &CertificateDer<'_>,
_dss: &DigitallySignedStruct,
) -> Result<HandshakeSignatureValid, rustls::Error> {
Ok(HandshakeSignatureValid::assertion())
}
fn supported_verify_schemes(&self) -> Vec<SignatureScheme> {
self.0.signature_verification_algorithms.supported_schemes()
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use rustls::NamedGroup;
use super::*;
use crate::{HttpSettings, Resolution};
fn config(settings: HttpSettings) -> HttpClientConfig {
Resolution::from(&settings).config
}
fn offered_groups(tls: &ClientConfig) -> Vec<NamedGroup> {
tls.crypto_provider()
.kx_groups
.iter()
.map(|group| group.name())
.collect()
}
fn offered_tls12_suites(tls: &ClientConfig) -> Vec<CipherSuite> {
tls.crypto_provider()
.cipher_suites
.iter()
.filter(|suite| suite.tls13().is_none())
.map(|suite| suite.suite())
.collect()
}
#[rstest]
#[case("X25519", NamedGroup::X25519)]
#[case("prime256v1", NamedGroup::secp256r1)]
#[case("secp384r1", NamedGroup::secp384r1)]
fn ecdh_curve_is_the_only_key_exchange_group_offered(
#[case] curve: &str,
#[case] expected: NamedGroup,
) {
let tls = ClientConfig::try_from(&config(HttpSettings {
ssl_ecdh_curve: Some(curve.into()),
..HttpSettings::default()
}))
.unwrap();
assert_eq!(offered_groups(&tls), [expected]);
}
#[test]
fn default_settings_offer_every_group_and_suite_of_the_provider() {
let tls = ClientConfig::try_from(&config(HttpSettings::default())).unwrap();
let provider = ring::default_provider();
assert_eq!(offered_groups(&tls).len(), provider.kx_groups.len());
assert_eq!(
tls.crypto_provider().cipher_suites.len(),
provider.cipher_suites.len()
);
}
#[test]
fn named_suites_are_the_only_tls12_suites_offered_and_tls13_stays() {
let tls = ClientConfig::try_from(&config(HttpSettings {
ssl_security_level: Some("ECDHE-RSA-AES256-GCM-SHA384".into()),
..HttpSettings::default()
}))
.unwrap();
assert_eq!(
offered_tls12_suites(&tls),
[CipherSuite::TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384]
);
assert!(
tls.crypto_provider()
.cipher_suites
.iter()
.any(|suite| suite.tls13().is_some())
);
}
#[rstest]
#[case(true, &[b"h2".as_slice(), b"http/1.1".as_slice()])]
#[case(false, &[b"http/1.1".as_slice()])]
fn alpn_offers_h2_only_when_http2_is_on(#[case] http2: bool, #[case] expected: &[&[u8]]) {
let tls = ClientConfig::try_from(&config(HttpSettings {
http2,
..HttpSettings::default()
}))
.unwrap();
assert_eq!(tls.alpn_protocols, expected);
}
#[test]
fn client_certificate_without_a_private_key_is_an_invalid_pem_error() {
let path = std::env::temp_dir().join(format!(
"litellm-http-cert-without-key-{}.pem",
std::process::id()
));
std::fs::write(
&path,
b"-----BEGIN CERTIFICATE-----\nAA==\n-----END CERTIFICATE-----\n",
)
.unwrap();
let result = ClientConfig::try_from(&HttpClientConfig {
client_certificate: Some(path.clone()),
..config(HttpSettings::default())
})
.map(drop);
std::fs::remove_file(&path).unwrap();
assert!(matches!(
result,
Err(Error::InvalidPem { path: reported, .. }) if reported == path
));
}
}

49
litellm-rust/crates/llms/src/custom_httpx/transport.rs → litellm-rust/crates/http/src/transport.rs
View file

@ -11,7 +11,7 @@ pub enum Error {
impl Error {
pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
let message = error.without_url().to_string();
let message = describe(error);
if before_dispatch {
Self::Connect(message)
} else {
@ -22,10 +22,18 @@ impl Error {
impl From<reqwest::Error> for Error {
fn from(error: reqwest::Error) -> Self {
Self::Network(error.without_url().to_string())
Self::Network(describe(error))
}
}
fn describe(error: reqwest::Error) -> String {
let error = error.without_url();
std::iter::successors(std::error::Error::source(&error), |cause| cause.source())
.fold(error.to_string(), |message, cause| {
format!("{message}: {cause}")
})
}
#[cfg(test)]
mod tests {
#[tokio::test]
@ -38,15 +46,38 @@ mod tests {
.send()
.await
.expect_err("invalid port");
let error = crate::custom_httpx::transport::Error::from_reqwest_before_dispatch(error);
assert!(matches!(
error,
crate::custom_httpx::transport::Error::Connect(_)
));
let error = crate::transport::Error::from_reqwest_before_dispatch(error);
assert!(matches!(error, crate::transport::Error::Connect(_)));
assert!(!error.to_string().contains("secret"));
assert!(!error.to_string().contains("private"));
}
fn root_cause(error: &dyn std::error::Error) -> Option<String> {
match error.source() {
Some(cause) => root_cause(cause).or_else(|| Some(cause.to_string())),
None => None,
}
}
#[tokio::test]
async fn network_error_message_names_the_underlying_cause() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
let address = listener.local_addr().expect("address");
drop(listener);
let error = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}/private?api_key=secret"))
.send()
.await
.expect_err("nothing listens on the port");
let root_cause = root_cause(&error).expect("reqwest reports a cause");
let message = crate::transport::Error::from(error).to_string();
assert!(message.contains(&root_cause), "{message}");
assert!(!message.contains("secret"));
}
#[tokio::test]
async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
use std::time::Duration;
@ -71,8 +102,8 @@ mod tests {
let error = response.expect_err("server does not respond");
assert!(error.is_timeout());
assert!(matches!(
crate::custom_httpx::transport::Error::from_reqwest_before_dispatch(error),
crate::custom_httpx::transport::Error::Network(_)
crate::transport::Error::from_reqwest_before_dispatch(error),
crate::transport::Error::Network(_)
));
}
}

2
litellm-rust/crates/llms/AGENTS.md
View file

@ -1,4 +1,4 @@
litellm-llms mirrors `litellm/llms/`: base config traits, provider transformations, and the `custom_httpx` handlers. See `../core/AGENTS.md` for how the crates layer.
litellm-llms mirrors `litellm/llms/`: base config traits, provider transformations, and the OCR request handler in `base_llm/ocr/handler.rs`. Transport code (clients, media fetching, header helpers, transport errors) lives in `litellm-http`. See `../core/AGENTS.md` for how the crates layer.
## Python/Rust transformation pairs

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

@ -6,7 +6,7 @@ license.workspace = true
repository.workspace = true
[features]
test-support = []
test-support = ["litellm-http/test-support"]
[dependencies]
litellm-types.workspace = true
@ -17,6 +17,7 @@ litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-host.workspace = true
litellm-framing.workspace = true
litellm-http.workspace = true
base64.workspace = true
bytes.workspace = true
data-url = "0.3.2"
@ -26,6 +27,7 @@ serde.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_path_to_error = "0.1"
serde_with.workspace = true
strum.workspace = true
thiserror.workspace = true
time.workspace = true
tokio = { workspace = true, features = ["sync"] }

2
litellm-rust/crates/llms/src/anthropic/chat/tests.rs
View file

@ -428,7 +428,7 @@ fn resolves_the_messages_url_and_x_api_key_auth() {
config
.auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None)
.expect("auth resolves"),
ChatCompletionsAuth::Header {
RequestAuth::Header {
name: "x-api-key",
value: "sk-x".to_string()
}

6
litellm-rust/crates/llms/src/anthropic/chat/transformation.rs
View file

@ -16,7 +16,7 @@ use crate::{
},
},
base_llm::chat::transformation::{
BaseConfig, ChatCompletionsAuth, Error, ProviderChatRequestData, ProviderChatResponseData,
BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth,
Unsupported, unsupported_message, unsupported_param,
},
};
@ -137,8 +137,8 @@ impl BaseConfig for AnthropicConfig {
_model: &str,
_optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
Ok(ChatCompletionsAuth::Header {
) -> Result<RequestAuth, Error> {
Ok(RequestAuth::Header {
name: "x-api-key",
value: resolve_anthropic_api_key(api_key, env_lookup)?,
})

1
litellm-rust/crates/llms/src/aws_textract/mod.rs Normal file
View file

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

12
litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md Normal file
View file

@ -0,0 +1,12 @@
- https://docs.aws.amazon.com/textract/latest/APIReference/Welcome.md
- https://docs.aws.amazon.com/textract/latest/APIReference/API_Operations.md
- https://docs.aws.amazon.com/textract/latest/APIReference/API_DetectDocumentText.md
- https://docs.aws.amazon.com/textract/latest/APIReference/API_AnalyzeDocument.md
- https://docs.aws.amazon.com/textract/latest/APIReference/API_StartDocumentTextDetection.md
- https://docs.aws.amazon.com/textract/latest/APIReference/API_Document.md
- https://docs.aws.amazon.com/textract/latest/APIReference/API_Block.md
- https://docs.aws.amazon.com/textract/latest/dg/what-is.md
- https://docs.aws.amazon.com/textract/latest/dg/sync.md
- https://docs.aws.amazon.com/textract/latest/dg/async.md
- https://docs.aws.amazon.com/textract/latest/dg/how-it-works-document-layout.md
- https://docs.aws.amazon.com/textract/latest/dg/limits.md

479
litellm-rust/crates/llms/src/aws_textract/ocr/analyze_transformation.rs Normal file
View file

@ -0,0 +1,479 @@
use std::collections::{BTreeMap, BTreeSet, HashMap};
use litellm_core_utils::call_arguments::{CallArguments, parse_options};
use serde::{Deserialize, Serialize};
use super::common_utils::{
Block, BlockType, FeatureType, LayoutType, TextractDocument, TextractEnvironment,
TextractOperation, TextractResponse, document_bytes, endpoint, environment, error_class,
health_check_document, inline_document, lines_by_page, ocr_response,
};
use crate::base_llm::ocr::{
error::Error,
handler::OcrClient,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat,
PreparedOcrRequest, decode_and_normalize_response,
},
};
const DEFAULT_FEATURE_TYPES: [FeatureType; 2] = [FeatureType::Layout, FeatureType::Tables];
#[derive(Default, Deserialize)]
pub struct AnalyzeDocumentOptions {
pub feature_types: Option<Vec<FeatureType>>,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct AnalyzeDocumentRequest {
#[serde(rename = "Document")]
pub document: TextractDocument,
#[serde(rename = "FeatureTypes")]
pub feature_types: Vec<FeatureType>,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct TextractAnalyzeDocumentConfig;
impl BaseOcrConfig for TextractAnalyzeDocumentConfig {
type OcrParams = AnalyzeDocumentOptions;
type ProviderRequest = AnalyzeDocumentRequest;
type Environment = TextractEnvironment;
fn get_supported_ocr_params(&self, _model: &str) -> &'static [&'static str] {
&["feature_types"]
}
fn get_health_check_document(&self) -> OcrDocument {
health_check_document()
}
fn map_ocr_params(
&self,
non_default_params: &CallArguments,
_model: &str,
) -> Result<AnalyzeDocumentOptions, Error> {
Ok(parse_options(non_default_params)?)
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<TextractEnvironment, Error> {
environment(request, TextractOperation::AnalyzeDocument).await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &AnalyzeDocumentOptions,
environment: &TextractEnvironment,
) -> Result<String, Error> {
Ok(endpoint(request, environment))
}
fn transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
optional_params: &AnalyzeDocumentOptions,
_headers: &[(String, String)],
) -> Result<AnalyzeDocumentRequest, Error> {
Ok(AnalyzeDocumentRequest {
document: document_bytes(&document)?,
feature_types: optional_params
.feature_types
.clone()
.unwrap_or_else(|| DEFAULT_FEATURE_TYPES.to_vec()),
})
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &AnalyzeDocumentOptions,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<AnalyzeDocumentRequest, Error> {
let document = inline_document(document, context).await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
fn get_error_class(
&self,
error_message: String,
status_code: u16,
headers: Vec<(String, String)>,
) -> Error {
error_class(error_message, status_code, headers)
}
}
fn normalize_response(
model: &str,
response: TextractResponse,
) -> Result<LiteLLMOcrResponse, Error> {
let blocks = &response.blocks;
let has_layout = blocks
.iter()
.any(|block| block.block_type.layout().is_some());
let page_markdown: Vec<(i64, String)> = if has_layout {
let by_id: HashMap<&str, &Block> = blocks
.iter()
.map(|block| (block.id.as_str(), block))
.collect();
let pages: BTreeSet<i64> = blocks.iter().map(Block::page).collect();
pages
.into_iter()
.map(|page| (page, layout_markdown(blocks, page, &by_id)))
.filter(|(_, markdown)| !markdown.is_empty())
.collect()
} else {
lines_by_page(blocks)
};
Ok(ocr_response(
model,
page_markdown,
response.document_metadata,
))
}
/// Layout blocks arrive in reading order. A list's items are repeated as
/// top-level `LAYOUT_TEXT` blocks. A `LAYOUT_TABLE` that links to its `TABLE`
/// renders it; one that only links to the table's lines takes the `TABLE` at
/// the same position on the page.
fn layout_markdown(blocks: &[Block], page: i64, by_id: &HashMap<&str, &Block>) -> String {
let on_page = || blocks.iter().filter(move |block| block.page() == page);
let list_items: BTreeSet<&str> = on_page()
.filter(|block| block.block_type == BlockType::LayoutList)
.flat_map(Block::children)
.collect();
let tables: Vec<&Block> = on_page()
.filter(|block| block.block_type == BlockType::Table)
.collect();
let table_ordinal: HashMap<&str, usize> = on_page()
.filter(|block| block.block_type == BlockType::LayoutTable)
.enumerate()
.map(|(ordinal, block)| (block.id.as_str(), ordinal))
.collect();
let table_of = |layout_table: &Block| {
layout_table
.children()
.filter_map(|id| by_id.get(id).copied())
.find(|child| child.block_type == BlockType::Table)
.or_else(|| {
table_ordinal
.get(layout_table.id.as_str())
.and_then(|ordinal| tables.get(*ordinal).copied())
})
};
let sections: Vec<String> = on_page()
.filter(|block| !list_items.contains(block.id.as_str()))
.filter_map(|block| Some((block, block.block_type.layout()?)))
.map(|(block, layout)| match layout {
LayoutType::Title => format!("# {}", text_of(block, by_id, " ")),
LayoutType::SectionHeader => format!("## {}", text_of(block, by_id, " ")),
LayoutType::List => block
.children()
.filter_map(|id| by_id.get(id))
.map(|item| format!("- {}", strip_bullet(&text_of(item, by_id, " "))))
.collect::<Vec<_>>()
.join("\n"),
LayoutType::Table => match table_of(block) {
Some(table) => table_markdown(table, by_id),
None => text_of(block, by_id, "\n"),
},
LayoutType::KeyValue => text_of(block, by_id, "\n"),
LayoutType::Figure => String::new(),
LayoutType::Text | LayoutType::Header | LayoutType::Footer | LayoutType::PageNumber => {
text_of(block, by_id, " ")
}
})
.filter(|section| !section.trim().is_empty())
.collect();
sections.join("\n\n")
}
fn text_of(block: &Block, by_id: &HashMap<&str, &Block>, separator: &str) -> String {
match &block.text {
Some(text) => text.clone(),
None => block
.children()
.filter_map(|id| by_id.get(id))
.map(|child| text_of(child, by_id, separator))
.filter(|text| !text.is_empty())
.collect::<Vec<_>>()
.join(separator),
}
}
fn strip_bullet(item: &str) -> &str {
item.trim_start_matches(['-', '*', '\u{2022}', '\u{00b7}'])
.trim_start()
}
fn table_markdown(table: &Block, by_id: &HashMap<&str, &Block>) -> String {
let cells: BTreeMap<(usize, usize), String> = table
.children()
.filter_map(|id| by_id.get(id))
.filter(|cell| cell.block_type == BlockType::Cell)
.filter_map(|cell| {
Some((
(cell.row_index?, cell.column_index?),
text_of(cell, by_id, " ").replace('|', "\\|"),
))
})
.collect();
let columns = cells.keys().map(|(_, column)| *column).max().unwrap_or(0);
let rows: BTreeSet<usize> = cells.keys().map(|(row, _)| *row).collect();
let render = |row: usize| {
let values: Vec<&str> = (1..=columns)
.map(|column| cells.get(&(row, column)).map_or("", String::as_str))
.collect();
format!("| {} |", values.join(" | "))
};
let divider = format!("|{}", " --- |".repeat(columns));
rows.iter()
.enumerate()
.flat_map(|(position, row)| {
std::iter::once(render(*row)).chain((position == 0).then(|| divider.clone()))
})
.collect::<Vec<_>>()
.join("\n")
}
#[cfg(test)]
mod tests {
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::*;
const MODEL: &str = "analyze-document";
#[fixture]
fn document() -> OcrDocument {
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,aGk=".into(),
extra_fields: Default::default(),
}
}
fn child(ids: &[&str]) -> Value {
json!([{"Type": "CHILD", "Ids": ids}])
}
fn line(id: &str, text: &str) -> Value {
json!({"Id": id, "BlockType": "LINE", "Text": text})
}
fn word(id: &str, text: &str) -> Value {
json!({"Id": id, "BlockType": "WORD", "Text": text})
}
fn layout(id: &str, block_type: &str, children: &[&str]) -> Value {
json!({"Id": id, "BlockType": block_type, "Relationships": child(children)})
}
fn table(id: &str, cells: &[&str]) -> Value {
json!({"Id": id, "BlockType": "TABLE", "Relationships": child(cells)})
}
fn cell(id: &str, row: usize, column: usize, words: &[&str]) -> Value {
json!({"Id": id, "BlockType": "CELL", "RowIndex": row, "ColumnIndex": column,
"Relationships": child(words)})
}
fn on_page(page: i64, mut block: Value) -> Value {
block["Page"] = json!(page);
block
}
#[rstest]
#[case::headings_paragraphs_and_a_list_without_repeating_its_items(
json!([
line("l1", "Quarterly Report"),
line("l2", "This report lists"),
line("l3", "the invoices."),
line("l4", "Line items"),
line("l5", "- Pay within 30 days"),
line("l6", "\u{2022} Quote the number"),
layout("t", "LAYOUT_TITLE", &["l1"]),
layout("p", "LAYOUT_TEXT", &["l2", "l3"]),
layout("h", "LAYOUT_SECTION_HEADER", &["l4"]),
layout("ul", "LAYOUT_LIST", &["i1", "i2"]),
layout("i1", "LAYOUT_TEXT", &["l5"]),
layout("i2", "LAYOUT_TEXT", &["l6"])
]),
vec![(
0,
"# Quarterly Report\n\nThis report lists the invoices.\n\n## Line items\n\n- Pay within 30 days\n- Quote the number"
)]
)]
#[case::header_footer_and_page_number_stay_in_reading_order(
json!([
line("l1", "ACME Corp"), line("l2", "Body"), line("l3", "Confidential"), line("l4", "3"),
layout("hd", "LAYOUT_HEADER", &["l1"]),
layout("p", "LAYOUT_TEXT", &["l2"]),
layout("ft", "LAYOUT_FOOTER", &["l3"]),
layout("pn", "LAYOUT_PAGE_NUMBER", &["l4"])
]),
vec![(0, "ACME Corp\n\nBody\n\nConfidential\n\n3")]
)]
#[case::a_table_is_rendered_from_its_cells_in_row_and_column_order(
json!([
line("l1", "Invoice"), line("l2", "Total"), line("l3", "12345"), line("l4", "a|b"),
word("w1", "Invoice"), word("w2", "Total"), word("w3", "12345"), word("w4", "a|b"),
{"Id": "tb", "BlockType": "TABLE", "Relationships": [
{"Type": "CHILD", "Ids": ["c4", "c1", "c3", "c2"]},
{"Type": "TABLE_TITLE", "Ids": ["title"]}
]},
cell("c1", 1, 1, &["w1"]), cell("c2", 1, 2, &["w2"]),
cell("c3", 2, 1, &["w3"]), cell("c4", 2, 2, &["w4"]),
layout("lt", "LAYOUT_TABLE", &["l1", "l2", "l3", "l4"])
]),
vec![(0, "| Invoice | Total |\n| --- | --- |\n| 12345 | a\\|b |")]
)]
#[case::a_layout_table_that_links_its_table_renders_that_one(
json!([
word("w1", "first"), word("w2", "second"),
table("tb1", &["c1"]), cell("c1", 1, 1, &["w1"]),
table("tb2", &["c2"]), cell("c2", 1, 1, &["w2"]),
layout("lt", "LAYOUT_TABLE", &["tb2"])
]),
vec![(0, "| second |\n| --- |")]
)]
#[case::a_missing_cell_leaves_an_empty_column(
json!([
word("w1", "a"), word("w2", "b"), word("w3", "c"),
table("tb", &["c1", "c2", "c3"]),
cell("c1", 1, 1, &["w1"]), cell("c2", 1, 2, &["w2"]), cell("c3", 2, 2, &["w3"]),
layout("lt", "LAYOUT_TABLE", &[])
]),
vec![(0, "| a | b |\n| --- | --- |\n| | c |")]
)]
#[case::a_layout_table_without_table_blocks_keeps_its_lines(
json!([
line("l1", "Invoice Total"),
line("l2", "12345 67.89"),
layout("lt", "LAYOUT_TABLE", &["l1", "l2"])
]),
vec![(0, "Invoice Total\n12345 67.89")]
)]
#[case::key_values_keep_one_line_each(
json!([
line("l1", "Name: Ana"),
line("l2", "Date: 2024-01-01"),
layout("kv", "LAYOUT_KEY_VALUE", &["l1", "l2"])
]),
vec![(0, "Name: Ana\nDate: 2024-01-01")]
)]
#[case::a_figure_has_no_markdown(
json!([
line("l1", "Caption"),
layout("f", "LAYOUT_FIGURE", &[]),
layout("p", "LAYOUT_TEXT", &["l1"])
]),
vec![(0, "Caption")]
)]
#[case::a_block_type_added_later_is_ignored(
json!([
line("l1", "Body"),
layout("new", "LAYOUT_SIDEBAR", &["l1"]),
layout("p", "LAYOUT_TEXT", &["l1"])
]),
vec![(0, "Body")]
)]
#[case::without_layout_blocks_lines_are_used(
json!([line("l1", "first"), word("w1", "first"), line("l2", "second")]),
vec![(0, "first\nsecond")]
)]
#[case::each_page_gets_its_own_markdown_and_its_own_tables(
json!([
on_page(1, line("a", "one")),
on_page(2, line("b", "two")),
on_page(2, word("w", "cell")),
on_page(1, layout("t1", "LAYOUT_TEXT", &["a"])),
on_page(2, table("tb", &["c"])),
on_page(2, cell("c", 1, 1, &["w"])),
on_page(2, layout("lt", "LAYOUT_TABLE", &["b"]))
]),
vec![(0, "one"), (1, "| cell |\n| --- |")]
)]
fn blocks_become_markdown_pages(#[case] blocks: Value, #[case] expected: Vec<(i64, &str)>) {
let response = TextractAnalyzeDocumentConfig
.transform_ocr_response(
MODEL,
&serde_json::to_vec(&json!({"DocumentMetadata": {"Pages": 1}, "Blocks": blocks}))
.unwrap(),
OcrResponseFormat::Litellm,
)
.unwrap();
let pages: Vec<(i64, &str)> = response
.pages
.iter()
.map(|page| (page.index, page.markdown.as_str()))
.collect();
assert_eq!(pages, expected);
}
#[rstest]
#[case::hyphen("- item", "item")]
#[case::asterisk("* item", "item")]
#[case::bullet("\u{2022} item", "item")]
#[case::middle_dot("\u{00b7}item", "item")]
#[case::no_bullet("item - with a dash", "item - with a dash")]
fn list_items_lose_their_own_bullet(#[case] item: &str, #[case] expected: &str) {
assert_eq!(strip_bullet(item), expected);
}
#[rstest]
#[case::defaults_to_layout_and_tables(json!({}), json!(["LAYOUT", "TABLES"]))]
#[case::overridden(json!({"feature_types": ["FORMS", "SIGNATURES"]}), json!(["FORMS", "SIGNATURES"]))]
#[case::explicit_null_uses_the_default(json!({"feature_types": null}), json!(["LAYOUT", "TABLES"]))]
fn feature_types_reach_the_request(
document: OcrDocument,
#[case] arguments: Value,
#[case] expected: Value,
) {
let arguments: CallArguments = serde_json::from_value(arguments).unwrap();
let params = TextractAnalyzeDocumentConfig
.map_ocr_params(&arguments, MODEL)
.unwrap();
let request = TextractAnalyzeDocumentConfig
.transform_ocr_request(MODEL, document, &params, &[])
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap(),
json!({"Document": {"Bytes": "aGk="}, "FeatureTypes": expected})
);
}
#[rstest]
#[case::undocumented_feature(json!({"feature_types": ["HANDWRITING"]}))]
#[case::lowercase_feature(json!({"feature_types": ["layout"]}))]
#[case::not_a_list(json!({"feature_types": "LAYOUT"}))]
fn feature_types_outside_the_documented_values_are_refused(#[case] arguments: Value) {
let arguments: CallArguments = serde_json::from_value(arguments).unwrap();
assert!(
TextractAnalyzeDocumentConfig
.map_ocr_params(&arguments, MODEL)
.is_err()
);
}
}

678
litellm-rust/crates/llms/src/aws_textract/ocr/common_utils.rs Normal file
View file

@ -0,0 +1,678 @@
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_auth_aws::{SigV4Signer, resolve_aws_region};
use litellm_http::outbound::RequestSigner;
use serde::{Deserialize, Serialize};
use strum::{EnumString, IntoStaticStr, VariantNames};
use crate::base_llm::ocr::{
document::{InlineDocument, inline_remote_document},
error::Error,
transformation::{
LiteLLMOcrResponse, OcrDocument, OcrEnvironment, OcrPage, OcrRequestContext, OcrUsageInfo,
PreparedOcrRequest,
},
};
const TEXTRACT_SERVICE: &str = "textract";
const AWS_JSON_CONTENT_TYPE: &str = "application/x-amz-json-1.1";
const TARGET_HEADER: &str = "X-Amz-Target";
const CONTENT_TYPE_HEADER: &str = "Content-Type";
const UNSUPPORTED_DOCUMENT: &str = "UnsupportedDocumentException";
const SYNC_DOCUMENT_MAX_BYTES: usize = 10 * 1024 * 1024;
const HEALTH_CHECK_IMAGE_DATA_URI: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4//8/AAX+Av4N70a4AAAAAElFTkSuQmCC";
/// Textract has operations rather than models; the model slot of
/// `aws_textract/<model>` names the one to call.
#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, VariantNames, PartialEq, Eq)]
#[strum(serialize_all = "kebab-case", ascii_case_insensitive)]
pub enum TextractOperation {
DetectDocumentText,
AnalyzeDocument,
}
impl TextractOperation {
pub const PROVIDER: &'static str = "aws_textract";
pub fn from_model(model: &str) -> Result<Self, Error> {
model.parse().map_err(|_| Error::InvalidModel {
provider: Self::PROVIDER,
model: model.to_string(),
supported: Self::VARIANTS,
})
}
fn target(self) -> &'static str {
match self {
Self::DetectDocumentText => "Textract.DetectDocumentText",
Self::AnalyzeDocument => "Textract.AnalyzeDocument",
}
}
}
#[derive(Debug, Deserialize, Serialize)]
pub struct TextractDocument {
#[serde(rename = "Bytes")]
pub bytes: String,
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum FeatureType {
Tables,
Forms,
Queries,
Signatures,
Layout,
}
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub(super) enum BlockType {
KeyValueSet,
Page,
Line,
Word,
Table,
Cell,
SelectionElement,
MergedCell,
Title,
Query,
QueryResult,
Signature,
TableTitle,
TableFooter,
LayoutText,
LayoutTitle,
LayoutHeader,
LayoutFooter,
LayoutSectionHeader,
LayoutPageNumber,
LayoutList,
LayoutFigure,
LayoutTable,
LayoutKeyValue,
#[serde(other)]
Unknown,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum LayoutType {
Text,
Title,
Header,
Footer,
SectionHeader,
PageNumber,
List,
Figure,
Table,
KeyValue,
}
impl BlockType {
pub fn layout(self) -> Option<LayoutType> {
match self {
Self::LayoutText => Some(LayoutType::Text),
Self::LayoutTitle => Some(LayoutType::Title),
Self::LayoutHeader => Some(LayoutType::Header),
Self::LayoutFooter => Some(LayoutType::Footer),
Self::LayoutSectionHeader => Some(LayoutType::SectionHeader),
Self::LayoutPageNumber => Some(LayoutType::PageNumber),
Self::LayoutList => Some(LayoutType::List),
Self::LayoutFigure => Some(LayoutType::Figure),
Self::LayoutTable => Some(LayoutType::Table),
Self::LayoutKeyValue => Some(LayoutType::KeyValue),
Self::KeyValueSet
| Self::Page
| Self::Line
| Self::Word
| Self::Table
| Self::Cell
| Self::SelectionElement
| Self::MergedCell
| Self::Title
| Self::Query
| Self::QueryResult
| Self::Signature
| Self::TableTitle
| Self::TableFooter
| Self::Unknown => None,
}
}
}
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub(super) enum RelationshipType {
Value,
Child,
ComplexFeatures,
MergedCell,
Title,
Answer,
Table,
TableTitle,
TableFooter,
#[serde(other)]
Unknown,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub(super) struct Block {
#[serde(default)]
pub id: String,
pub block_type: BlockType,
pub text: Option<String>,
pub page: Option<i64>,
pub row_index: Option<usize>,
pub column_index: Option<usize>,
#[serde(default)]
pub relationships: Vec<Relationship>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub(super) struct Relationship {
pub r#type: RelationshipType,
#[serde(default)]
pub ids: Vec<String>,
}
impl Block {
pub fn page(&self) -> i64 {
self.page.unwrap_or(1)
}
pub fn children(&self) -> impl Iterator<Item = &str> {
self.relationships
.iter()
.filter(|relationship| relationship.r#type == RelationshipType::Child)
.flat_map(|relationship| relationship.ids.iter().map(String::as_str))
}
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub(super) struct DocumentMetadata {
pub pages: Option<i64>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct TextractResponse {
#[serde(default)]
pub(super) blocks: Vec<Block>,
pub(super) document_metadata: Option<DocumentMetadata>,
}
pub struct TextractEnvironment {
headers: Vec<(String, String)>,
region: String,
signer: SigV4Signer,
}
impl OcrEnvironment for TextractEnvironment {
fn headers(&self) -> &[(String, String)] {
&self.headers
}
fn signer(&self) -> Option<&dyn RequestSigner> {
Some(&self.signer)
}
}
pub(super) fn health_check_document() -> OcrDocument {
OcrDocument::ImageUrl {
image_url: HEALTH_CHECK_IMAGE_DATA_URI.into(),
extra_fields: Default::default(),
}
}
pub(super) async fn environment(
request: &PreparedOcrRequest,
operation: TextractOperation,
) -> Result<TextractEnvironment, Error> {
let env_lookup = |name: &str| request.connection.secret(name);
let region =
resolve_aws_region(None, &request.optional_params, &env_lookup).ok_or_else(|| {
Error::InvalidRequest(
"Missing AWS region - pass aws_region_name or set AWS_REGION_NAME or AWS_REGION"
.into(),
)
})?;
let signer = SigV4Signer::resolve(
region.clone(),
TEXTRACT_SERVICE,
&request.optional_params,
&env_lookup,
)
.await
.map_err(litellm_auth::Error::from)?;
Ok(TextractEnvironment {
headers: operation_headers(&request.connection.extra_headers, operation),
region,
signer,
})
}
/// A caller's copy of an operation header would reach the wire next to ours
/// while the signature covers only one value, which Textract rejects.
fn operation_headers(
extra_headers: &[(String, String)],
operation: TextractOperation,
) -> Vec<(String, String)> {
let operation = [
(TARGET_HEADER, operation.target()),
(CONTENT_TYPE_HEADER, AWS_JSON_CONTENT_TYPE),
];
extra_headers
.iter()
.filter(|(name, _)| {
!operation
.iter()
.any(|(operation_name, _)| name.eq_ignore_ascii_case(operation_name))
})
.cloned()
.chain(
operation
.iter()
.map(|(name, value)| (name.to_string(), value.to_string())),
)
.collect()
}
pub(super) fn endpoint(request: &PreparedOcrRequest, environment: &TextractEnvironment) -> String {
request
.connection
.api_base
.clone()
.unwrap_or_else(|| format!("https://textract.{}.amazonaws.com/", environment.region))
}
pub(super) fn document_bytes(document: &OcrDocument) -> Result<TextractDocument, Error> {
let inline = InlineDocument::parse(document.source())?.ok_or(Error::InvalidDataUri)?;
Ok(TextractDocument {
bytes: STANDARD.encode(inline.decode(SYNC_DOCUMENT_MAX_BYTES)?),
})
}
pub(super) async fn inline_document(
document: OcrDocument,
context: OcrRequestContext<'_>,
) -> Result<OcrDocument, Error> {
inline_remote_document(
context.client.document_fetcher(),
document,
context.connection,
)
.await
}
#[derive(Deserialize)]
struct AwsError {
#[serde(rename = "__type", default)]
kind: String,
#[serde(rename = "Message", alias = "message", default)]
message: String,
}
/// Textract answers both an unsupported format and a multi-page PDF or TIFF
/// with a bare "unsupported document format", which reads like a corrupt file.
/// Say what the synchronous API accepts.
pub(super) fn error_class(body: String, status: u16, headers: Vec<(String, String)>) -> Error {
let unsupported = serde_json::from_str::<AwsError>(&body)
.ok()
.filter(|error| error.kind.ends_with(UNSUPPORTED_DOCUMENT));
Error::Provider {
status,
body: match unsupported {
Some(error) => format!(
"{UNSUPPORTED_DOCUMENT}: {}. aws_textract uses Textract's synchronous API, which reads a JPEG, PNG, or a single-page PDF or TIFF; other formats and multi-page documents are not supported",
error.message
),
None => body,
},
headers,
}
}
pub(super) fn lines_by_page(blocks: &[Block]) -> Vec<(i64, String)> {
let pages: std::collections::BTreeSet<i64> = blocks.iter().map(Block::page).collect();
pages
.into_iter()
.map(|page| {
let lines: Vec<&str> = blocks
.iter()
.filter(|block| block.block_type == BlockType::Line && block.page() == page)
.filter_map(|block| block.text.as_deref())
.collect();
(page, lines.join("\n"))
})
.filter(|(_, markdown)| !markdown.is_empty())
.collect()
}
pub(super) fn ocr_response(
model: &str,
page_markdown: Vec<(i64, String)>,
document_metadata: Option<DocumentMetadata>,
) -> LiteLLMOcrResponse {
let pages: Vec<OcrPage> = page_markdown
.into_iter()
.map(|(page, markdown)| OcrPage {
index: page - 1,
markdown,
..Default::default()
})
.collect();
let pages_processed = document_metadata
.and_then(|metadata| metadata.pages)
.or_else(|| i64::try_from(pages.len()).ok());
LiteLLMOcrResponse {
usage_info: Some(OcrUsageInfo {
pages_processed,
..Default::default()
}),
..LiteLLMOcrResponse::new(model, pages)
}
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::{Value, json};
use super::*;
const HINT: &str = "other formats and multi-page documents are not supported";
fn blocks(value: Value) -> Vec<Block> {
serde_json::from_value(value).unwrap()
}
#[rstest]
#[case::detect("detect-document-text", TextractOperation::DetectDocumentText)]
#[case::analyze("analyze-document", TextractOperation::AnalyzeDocument)]
#[case::any_case("Analyze-Document", TextractOperation::AnalyzeDocument)]
fn a_model_names_its_operation(#[case] model: &str, #[case] expected: TextractOperation) {
assert_eq!(TextractOperation::from_model(model).unwrap(), expected);
}
#[rstest]
#[case::misspelled("analyse-document")]
#[case::operation_name_from_the_api("AnalyzeDocument")]
#[case::operation_litellm_does_not_call("analyze-expense")]
#[case::empty("")]
fn a_model_outside_the_operations_is_refused_with_the_supported_names(#[case] model: &str) {
let error = TextractOperation::from_model(model).unwrap_err();
assert_eq!(
error.to_string(),
format!(
"invalid model: aws_textract has no model {model:?} - use one of: detect-document-text, analyze-document"
)
);
assert_eq!(error.http_status_code(), Some(400));
}
#[rstest]
#[case::line("LINE", BlockType::Line)]
#[case::key_value_set("KEY_VALUE_SET", BlockType::KeyValueSet)]
#[case::layout_section_header("LAYOUT_SECTION_HEADER", BlockType::LayoutSectionHeader)]
#[case::layout_key_value("LAYOUT_KEY_VALUE", BlockType::LayoutKeyValue)]
#[case::added_by_textract_later("LAYOUT_SIDEBAR", BlockType::Unknown)]
fn block_type_reads_the_documented_names(#[case] wire: &str, #[case] expected: BlockType) {
let block: Block = serde_json::from_value(json!({"BlockType": wire})).unwrap();
assert_eq!(block.block_type, expected);
}
#[rstest]
#[case::layout_title(BlockType::LayoutTitle, Some(LayoutType::Title))]
#[case::layout_table(BlockType::LayoutTable, Some(LayoutType::Table))]
#[case::table_is_not_layout(BlockType::Table, None)]
#[case::title_is_not_layout(BlockType::Title, None)]
#[case::unknown_is_not_layout(BlockType::Unknown, None)]
fn only_layout_block_types_have_a_layout_type(
#[case] block_type: BlockType,
#[case] expected: Option<LayoutType>,
) {
assert_eq!(block_type.layout(), expected);
}
#[rstest]
#[case::child_only(json!([{"Type": "CHILD", "Ids": ["a", "b"]}]), vec!["a", "b"])]
#[case::other_relationships_are_skipped(
json!([
{"Type": "TABLE_TITLE", "Ids": ["t"]},
{"Type": "CHILD", "Ids": ["a"]},
{"Type": "MERGED_CELL", "Ids": ["m"]},
{"Type": "ADDED_LATER", "Ids": ["x"]},
{"Type": "CHILD", "Ids": ["b"]}
]),
vec!["a", "b"]
)]
#[case::no_relationships(json!([]), vec![])]
fn children_are_the_ids_of_child_relationships(
#[case] relationships: Value,
#[case] expected: Vec<&str>,
) {
let block: Block =
serde_json::from_value(json!({"BlockType": "LINE", "Relationships": relationships}))
.unwrap();
assert_eq!(block.children().collect::<Vec<_>>(), expected);
}
#[rstest]
#[case::tables("TABLES", Some(FeatureType::Tables))]
#[case::forms("FORMS", Some(FeatureType::Forms))]
#[case::queries("QUERIES", Some(FeatureType::Queries))]
#[case::signatures("SIGNATURES", Some(FeatureType::Signatures))]
#[case::layout("LAYOUT", Some(FeatureType::Layout))]
#[case::lowercase_is_not_a_feature("layout", None)]
#[case::undocumented("HANDWRITING", None)]
fn feature_type_accepts_only_the_documented_values(
#[case] wire: &str,
#[case] expected: Option<FeatureType>,
) {
assert_eq!(
serde_json::from_value::<FeatureType>(json!(wire)).ok(),
expected
);
if let Some(feature) = expected {
assert_eq!(serde_json::to_value(feature).unwrap(), json!(wire));
}
}
#[rstest]
#[case::image_url(
OcrDocument::ImageUrl {
image_url: "data:image/png;base64,aGVsbG8=".into(),
extra_fields: Default::default(),
},
"aGVsbG8="
)]
#[case::document_url(
OcrDocument::DocumentUrl {
document_url: "data:application/pdf;base64,YWJj".into(),
extra_fields: Default::default(),
},
"YWJj"
)]
#[case::percent_encoded_data_uri_is_re_encoded_as_base64(
OcrDocument::DocumentUrl {
document_url: "data:,abc".into(),
extra_fields: Default::default(),
},
"YWJj"
)]
fn document_bytes_are_the_base64_payload_without_the_data_uri_envelope(
#[case] document: OcrDocument,
#[case] expected: &str,
) {
assert_eq!(document_bytes(&document).unwrap().bytes, expected);
}
#[rstest]
#[case::remote_url("https://example.com/a.pdf".to_string(), Error::InvalidDataUri)]
#[case::invalid_base64("data:image/png;base64,@@@".to_string(), Error::InvalidDataUri)]
#[case::over_the_sync_limit(
format!("data:,{}", "a".repeat(SYNC_DOCUMENT_MAX_BYTES + 1)),
Error::InlineDocumentTooLarge
)]
fn document_bytes_refuse_what_the_sync_api_cannot_take(
#[case] document_url: String,
#[case] expected: Error,
) {
let error = document_bytes(&OcrDocument::DocumentUrl {
document_url,
extra_fields: Default::default(),
})
.unwrap_err();
assert_eq!(
std::mem::discriminant(&error),
std::mem::discriminant(&expected)
);
}
#[rstest]
#[case::bare_type(
r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#,
Some("Request has unsupported document format")
)]
#[case::namespaced_type(
r#"{"__type":"com.amazonaws.textract#UnsupportedDocumentException","Message":"bad"}"#,
Some("bad")
)]
#[case::lowercase_message(
r#"{"__type":"UnsupportedDocumentException","message":"bad"}"#,
Some("bad")
)]
#[case::other_exception(r#"{"__type":"AccessDeniedException","Message":"no"}"#, None)]
#[case::json_without_a_type(r#"{"Message":"no"}"#, None)]
#[case::not_json("<html>bad gateway</html>", None)]
fn only_an_unsupported_document_gains_the_sync_api_hint(
#[case] body: &str,
#[case] hinted_message: Option<&str>,
) {
let response_headers = vec![("x-amzn-requestid".to_string(), "abc".to_string())];
let Error::Provider {
status,
body: reported,
headers,
} = error_class(body.into(), 400, response_headers.clone())
else {
panic!("expected a provider error");
};
assert_eq!(status, 400);
assert_eq!(headers, response_headers);
match hinted_message {
Some(message) => {
assert!(reported.contains(message), "{reported}");
assert!(reported.contains(HINT), "{reported}");
}
None => assert_eq!(reported, body),
}
}
#[rstest]
#[case::no_caller_headers(vec![], vec![])]
#[case::unrelated_headers_are_kept(vec![("x-trace", "1")], vec![("x-trace", "1")])]
#[case::a_caller_content_type_is_replaced(
vec![("content-type", "application/json"), ("x-trace", "1")],
vec![("x-trace", "1")]
)]
#[case::a_caller_target_is_replaced(
vec![("X-AMZ-TARGET", "Textract.AnalyzeDocument")],
vec![]
)]
fn operation_headers_are_sent_once(
#[case] extra_headers: Vec<(&str, &str)>,
#[case] kept: Vec<(&str, &str)>,
) {
let owned = |headers: Vec<(&str, &str)>| -> Vec<(String, String)> {
headers
.into_iter()
.map(|(name, value)| (name.to_string(), value.to_string()))
.collect()
};
let headers =
operation_headers(&owned(extra_headers), TextractOperation::DetectDocumentText);
let mut expected = owned(kept);
expected.extend(owned(vec![
("X-Amz-Target", "Textract.DetectDocumentText"),
("Content-Type", "application/x-amz-json-1.1"),
]));
assert_eq!(headers, expected);
}
#[rstest]
#[case::words_are_not_repeated(
json!([
{"BlockType": "PAGE"},
{"BlockType": "LINE", "Text": "Invoice 12345"},
{"BlockType": "WORD", "Text": "Invoice"},
{"BlockType": "WORD", "Text": "12345"},
{"BlockType": "LINE", "Text": "total 67.89"}
]),
vec![(1, "Invoice 12345\ntotal 67.89")]
)]
#[case::pages_are_sorted_and_keep_line_order(
json!([
{"BlockType": "LINE", "Text": "second", "Page": 2},
{"BlockType": "LINE", "Text": "first", "Page": 1},
{"BlockType": "LINE", "Text": "also second", "Page": 2}
]),
vec![(1, "first"), (2, "second\nalso second")]
)]
#[case::a_page_without_lines_is_dropped(
json!([
{"BlockType": "PAGE", "Page": 1},
{"BlockType": "LINE", "Text": "only", "Page": 2}
]),
vec![(2, "only")]
)]
#[case::no_blocks(json!([]), vec![])]
fn lines_are_grouped_by_page(#[case] input: Value, #[case] expected: Vec<(i64, &str)>) {
let pages = lines_by_page(&blocks(input));
let pages: Vec<(i64, &str)> = pages
.iter()
.map(|(page, markdown)| (*page, markdown.as_str()))
.collect();
assert_eq!(pages, expected);
}
#[rstest]
#[case::metadata_wins(Some(3), Some(3))]
#[case::metadata_without_pages_falls_back_to_the_page_count(None, Some(2))]
fn pages_are_zero_indexed_and_usage_reports_pages_processed(
#[case] metadata_pages: Option<i64>,
#[case] expected: Option<i64>,
) {
let response = ocr_response(
"detect-document-text",
vec![(1, "first".into()), (3, "third".into())],
Some(DocumentMetadata {
pages: metadata_pages,
}),
);
let pages: Vec<(i64, &str)> = response
.pages
.iter()
.map(|page| (page.index, page.markdown.as_str()))
.collect();
assert_eq!(pages, vec![(0, "first"), (2, "third")]);
assert_eq!(response.usage_info.unwrap().pages_processed, expected);
}
}

3
litellm-rust/crates/llms/src/aws_textract/ocr/mod.rs Normal file
View file

@ -0,0 +1,3 @@
pub mod analyze_transformation;
pub mod common_utils;
pub mod transformation;

240
litellm-rust/crates/llms/src/aws_textract/ocr/transformation.rs Normal file
View file

@ -0,0 +1,240 @@
use litellm_core_utils::call_arguments::CallArguments;
use serde::{Deserialize, Serialize};
use super::common_utils::{
TextractDocument, TextractEnvironment, TextractOperation, TextractResponse, document_bytes,
endpoint, environment, error_class, health_check_document, inline_document, lines_by_page,
ocr_response,
};
use crate::base_llm::ocr::{
error::Error,
handler::OcrClient,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat,
PreparedOcrRequest, decode_and_normalize_response,
},
};
#[derive(Debug, Deserialize, Serialize)]
pub struct DetectDocumentTextRequest {
#[serde(rename = "Document")]
pub document: TextractDocument,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct TextractDetectTextConfig;
impl BaseOcrConfig for TextractDetectTextConfig {
type OcrParams = ();
type ProviderRequest = DetectDocumentTextRequest;
type Environment = TextractEnvironment;
fn get_health_check_document(&self) -> OcrDocument {
health_check_document()
}
fn map_ocr_params(
&self,
_non_default_params: &CallArguments,
_model: &str,
) -> Result<(), Error> {
Ok(())
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<TextractEnvironment, Error> {
environment(request, TextractOperation::DetectDocumentText).await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_optional_params: &(),
environment: &TextractEnvironment,
) -> Result<String, Error> {
Ok(endpoint(request, environment))
}
fn transform_ocr_request(
&self,
_model: &str,
document: OcrDocument,
_optional_params: &(),
_headers: &[(String, String)],
) -> Result<DetectDocumentTextRequest, Error> {
Ok(DetectDocumentTextRequest {
document: document_bytes(&document)?,
})
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &(),
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<DetectDocumentTextRequest, Error> {
let document = inline_document(document, context).await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, Error> {
decode_and_normalize_response(model, raw_response, request_format, normalize_response)
}
fn get_error_class(
&self,
error_message: String,
status_code: u16,
headers: Vec<(String, String)>,
) -> Error {
error_class(error_message, status_code, headers)
}
}
fn normalize_response(
model: &str,
response: TextractResponse,
) -> Result<LiteLLMOcrResponse, Error> {
Ok(ocr_response(
model,
lines_by_page(&response.blocks),
response.document_metadata,
))
}
#[cfg(test)]
mod tests {
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use super::*;
const MODEL: &str = "detect-document-text";
#[fixture]
fn document(#[default("data:image/png;base64,aGVsbG8=")] source: &str) -> OcrDocument {
OcrDocument::DocumentUrl {
document_url: source.into(),
extra_fields: Default::default(),
}
}
#[rstest]
#[case::one_page_without_page_numbers(
json!({
"DetectDocumentTextModelVersion": "1.0",
"DocumentMetadata": {"Pages": 1},
"Blocks": [
{"BlockType": "PAGE"},
{"BlockType": "LINE", "Text": "Invoice 12345"},
{"BlockType": "WORD", "Text": "Invoice"},
{"BlockType": "WORD", "Text": "12345"},
{"BlockType": "LINE", "Text": "total 67.89"}
]
}),
vec![(0, "Invoice 12345\ntotal 67.89")],
Some(1)
)]
#[case::pages_out_of_order(
json!({
"DocumentMetadata": {"Pages": 2},
"Blocks": [
{"BlockType": "LINE", "Text": "second", "Page": 2},
{"BlockType": "LINE", "Text": "first", "Page": 1},
{"BlockType": "LINE", "Text": "also second", "Page": 2}
]
}),
vec![(0, "first"), (1, "second\nalso second")],
Some(2)
)]
#[case::missing_metadata_counts_the_pages_with_text(
json!({"Blocks": [{"BlockType": "LINE", "Text": "only"}]}),
vec![(0, "only")],
Some(1)
)]
#[case::blank_document(json!({"DocumentMetadata": {"Pages": 1}}), vec![], Some(1))]
fn response_lines_become_one_markdown_page_per_document_page(
#[case] raw_response: Value,
#[case] expected_pages: Vec<(i64, &str)>,
#[case] expected_pages_processed: Option<i64>,
) {
let response = TextractDetectTextConfig
.transform_ocr_response(
MODEL,
&serde_json::to_vec(&raw_response).unwrap(),
OcrResponseFormat::Litellm,
)
.unwrap();
let pages: Vec<(i64, &str)> = response
.pages
.iter()
.map(|page| (page.index, page.markdown.as_str()))
.collect();
assert_eq!(pages, expected_pages);
assert_eq!(response.model, MODEL);
assert_eq!(
response.usage_info.unwrap().pages_processed,
expected_pages_processed
);
}
#[rstest]
fn the_request_is_only_the_document_bytes(document: OcrDocument) {
let request = TextractDetectTextConfig
.transform_ocr_request(MODEL, document, &(), &[])
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap(),
json!({"Document": {"Bytes": "aGVsbG8="}})
);
}
#[rstest]
fn a_remote_url_is_refused_by_the_sync_transform(
#[with("https://example.com/a.pdf")] document: OcrDocument,
) {
let error = TextractDetectTextConfig
.transform_ocr_request(MODEL, document, &(), &[])
.unwrap_err();
assert!(matches!(error, Error::InvalidDataUri));
}
#[rstest]
fn the_health_check_document_is_an_inline_image_the_request_accepts() {
let document = TextractDetectTextConfig.get_health_check_document();
assert!(
TextractDetectTextConfig
.transform_ocr_request(MODEL, document, &(), &[])
.is_ok()
);
}
#[rstest]
fn provider_errors_go_through_the_shared_textract_error_class() {
let error = TextractDetectTextConfig.get_error_class(
r#"{"__type":"UnsupportedDocumentException","Message":"Request has unsupported document format"}"#.into(),
400,
Vec::new(),
);
assert!(
error
.to_string()
.contains("multi-page documents are not supported")
);
}
}

6
litellm-rust/crates/llms/src/azure_ai/ocr/cohere_parse_transformation.rs
View file

@ -5,6 +5,7 @@ use crate::{
base_llm::ocr::{
document::{inline_remote_document, validate_inline_document},
error::Error,
handler::OcrClient,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrDocument, OcrRequestContext, OcrResponseFormat,
PreparedOcrRequest,
@ -13,7 +14,6 @@ use crate::{
cohere::ocr::transformation::{
CohereOptions, CohereParseConfig, CohereRequest, validate_document,
},
custom_httpx::llm_http_handler::OcrClient,
};
#[derive(Default)]
@ -53,7 +53,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
) -> Result<String, Error> {
let base = super::transformation::AzureAiOcrConfig::resolve_api_base(
request.connection.api_base.as_deref(),
&crate::base_llm::ocr::transformation::credential_env,
&|name: &str| request.connection.secret(name),
)?;
self.get_complete_url(&base)
}
@ -108,7 +108,7 @@ impl BaseOcrConfig for AzureAICohereParseConfig {
}
fn validate_request_body(&self, body: &Value) -> Result<(), Error> {
let document = crate::custom_httpx::llm_http_handler::body_document(body)?;
let document = crate::base_llm::ocr::handler::body_document(body)?;
validate_document(&document)?;
validate_inline_document(&document)
}

16
litellm-rust/crates/llms/src/azure_ai/ocr/common_utils.rs
View file

@ -3,7 +3,21 @@ use std::sync::OnceLock;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
use crate::base_llm::ocr::{error::Error, transformation::OcrConnection};
use crate::base_llm::ocr::{
error::Error,
transformation::{OcrConnection, PreparedOcrRequest},
};
pub(crate) fn azure_auth_inputs(request: &PreparedOcrRequest) -> Result<AzureAuthInputs, Error> {
Ok(AzureAuthInputs {
azure_ad_token_provider: request.azure_ad_token_provider.clone(),
..AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?
}
.or_configured_token_refresh(request.connection.settings.enable_azure_ad_token_refresh))
}
pub(super) async fn resolve_entra(
config: &AzureAuthInputs,

125
litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs
View file

@ -14,24 +14,20 @@ use serde_json::{Map, Value};
use serde_with::serde_as;
use tokio::time::Instant;
use crate::{
base_llm::ocr::{
document::InlineDocument,
error::Error,
transformation::{
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES,
OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage,
OcrPageDimensions, OcrResponseContext, OcrResponseFormat, OcrUsageInfo,
PreparedOcrRequest, ResolvedOcrCredentials, credential_env,
decode_and_normalize_response, decode_response,
},
use crate::base_llm::ocr::{
document::InlineDocument,
error::Error,
handler::{CallHooks, OcrClient, read_json_response},
settings::OcrSettings,
transformation::{
BaseOcrConfig, DecodedOcrResponse, LiteLLMOcrResponse, OCR_INLINE_MAX_BYTES,
OCR_POLL_RETRY_SECS, OcrConnection, OcrCredentialInputs, OcrDocument, OcrPage,
OcrPageDimensions, OcrResponseContext, OcrResponseFormat, OcrUsageInfo, PreparedOcrRequest,
ResolvedOcrCredentials, decode_and_normalize_response, decode_response,
},
custom_httpx::llm_http_handler::{CallHooks, OcrClient, read_json_response},
};
const AZURE_DI_API_VERSION: &str = "2024-11-30";
const AZURE_DI_SUBSCRIPTION_HEADER: &str = "Ocp-Apim-Subscription-Key";
const AZURE_DI_DEFAULT_DPI: i64 = 96;
const AZURE_DI_DEFAULT_WIDTH: f64 = 8.5;
const AZURE_DI_DEFAULT_HEIGHT: f64 = 11.0;
@ -178,15 +174,11 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, Error> {
let config = AzureAuthInputs {
azure_ad_token_provider: request.azure_ad_token_provider.clone(),
..AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?
};
self.resolve_headers(&request.connection, &config, &credential_env)
.await
let config = crate::azure_ai::ocr::common_utils::azure_auth_inputs(request)?;
self.resolve_headers(&request.connection, &config, &|name: &str| {
request.connection.secret(name)
})
.await
}
fn get_complete_url(
@ -196,9 +188,17 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
_environment: &Self::Environment,
) -> Result<String, Error> {
let endpoint = nonblank(request.connection.api_base.clone())
.or_else(|| nonblank(credential_env(AZURE_DI_ENDPOINT_ENV)))
.or_else(|| nonblank(request.connection.secret(AZURE_DI_ENDPOINT_ENV)))
.ok_or_else(|| Error::Auth(litellm_auth::Error::ProviderAuthentication("Missing Azure Document Intelligence API Base - Set AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT or pass api_base".into())))?;
self.build_ocr_url(&endpoint, &request.model, optional_params)
self.build_ocr_url(
&endpoint,
&request.model,
optional_params,
&request
.connection
.settings
.document_intelligence_api_version,
)
}
fn transform_ocr_request(
@ -217,12 +217,13 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
raw_response: &[u8],
request_format: OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, Error> {
decode_and_normalize_response(
model,
raw_response,
request_format,
transform_completed_response,
)
decode_and_normalize_response(model, raw_response, request_format, |model, response| {
transform_completed_response(
model,
response,
OcrSettings::default().document_intelligence_dpi,
)
})
}
async fn async_transform_ocr_response(
@ -243,7 +244,11 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig {
.await?;
Ok(LiteLLMOcrResponse {
provider_native_response: decoded.native,
..transform_completed_response(model, decoded.data)?
..transform_completed_response(
model,
decoded.data,
context.connection.settings.document_intelligence_dpi,
)?
})
}
}
@ -356,6 +361,7 @@ fn build_request(document: OcrDocument) -> Result<DocumentIntelligenceRequest, E
fn transform_completed_response(
model: &str,
response: AzureDocumentIntelligenceOperation,
dpi: i64,
) -> Result<LiteLLMOcrResponse, Error> {
if response.status != Some(OperationStatus::Succeeded) {
return Err(Error::OperationStatus(
@ -369,7 +375,7 @@ fn transform_completed_response(
let pages = result
.pages
.into_iter()
.map(transform_azure_page)
.map(|page| transform_azure_page(page, dpi))
.collect::<Result<Vec<_>, _>>()?;
let pages_processed = i64::try_from(pages.len()).map_err(|_| Error::NumericRange("pages"))?;
Ok(LiteLLMOcrResponse {
@ -384,7 +390,7 @@ fn transform_completed_response(
})
}
fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage, Error> {
fn transform_azure_page(page: AzureDocumentIntelligencePage, dpi: i64) -> Result<OcrPage, Error> {
let index = page
.page_number
.unwrap_or(1)
@ -394,6 +400,7 @@ fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage,
page.width.unwrap_or(AZURE_DI_DEFAULT_WIDTH),
page.height.unwrap_or(AZURE_DI_DEFAULT_HEIGHT),
page.unit.as_deref().unwrap_or("inch"),
dpi,
)?;
let markdown = page
.lines
@ -409,16 +416,17 @@ fn transform_azure_page(page: AzureDocumentIntelligencePage) -> Result<OcrPage,
})
}
fn convert_dimensions(width: f64, height: f64, unit: &str) -> Result<OcrPageDimensions, Error> {
let scale = if unit == "inch" {
AZURE_DI_DEFAULT_DPI as f64
} else {
1.0
};
fn convert_dimensions(
width: f64,
height: f64,
unit: &str,
dpi: i64,
) -> Result<OcrPageDimensions, Error> {
let scale = if unit == "inch" { dpi as f64 } else { 1.0 };
Ok(OcrPageDimensions {
width: Some(pixel_dimension(width, scale, "page.width")?),
height: Some(pixel_dimension(height, scale, "page.height")?),
dpi: Some(AZURE_DI_DEFAULT_DPI),
dpi: Some(dpi),
})
}
@ -440,7 +448,7 @@ async fn read_operation_response(
hooks: &dyn CallHooks<Error>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, Error> {
if response.status() != reqwest::StatusCode::ACCEPTED {
let bytes = crate::custom_httpx::llm_http_handler::read_response_bytes(
let bytes = crate::base_llm::ocr::handler::read_response_bytes(
response,
connection.max_response_bytes,
)
@ -462,11 +470,9 @@ async fn read_operation_response(
{
return Err(Error::PollOrigin);
}
let bytes = crate::custom_httpx::llm_http_handler::read_response_bytes(
response,
connection.max_response_bytes,
)
.await?;
let bytes =
crate::base_llm::ocr::handler::read_response_bytes(response, connection.max_response_bytes)
.await?;
hooks.response_received(&bytes).await?;
poll_operation(http_client, operation, headers, connection, native, hooks).await
}
@ -480,7 +486,7 @@ async fn poll_operation(
hooks: &dyn CallHooks<Error>,
) -> Result<DecodedOcrResponse<AzureDocumentIntelligenceOperation>, Error> {
let deadline = Instant::now()
.checked_add(connection.poll_timeout)
.checked_add(connection.settings.poll_timeout)
.ok_or(Error::PollTimeout)?;
loop {
@ -491,21 +497,19 @@ async fn poll_operation(
let builder = http_client
.get(url.clone())
.timeout(remaining.min(connection.timeout));
let builder = crate::custom_httpx::http_handler::with_headers(
let builder = litellm_http::request::with_headers(
builder,
headers,
crate::custom_httpx::http_handler::HeaderPolicy::Only(&[
litellm_http::request::HeaderPolicy::Only(&[
AZURE_DI_SUBSCRIPTION_HEADER,
"authorization",
]),
);
let response = tokio::time::timeout_at(
deadline,
crate::custom_httpx::http_handler::http_request(builder),
)
.await
.map_err(|_| Error::PollTimeout)?
.map_err(crate::custom_httpx::transport::Error::from)?;
let response =
tokio::time::timeout_at(deadline, litellm_http::request::http_request(builder))
.await
.map_err(|_| Error::PollTimeout)?
.map_err(litellm_http::transport::Error::from)?;
let retry = response
.headers()
.get(reqwest::header::RETRY_AFTER)
@ -551,13 +555,14 @@ impl AzureDocumentIntelligenceOcrConfig {
endpoint: &str,
model: &str,
params: &DocumentIntelligenceParams,
api_version: &str,
) -> Result<String, Error> {
let model = format!("{}:analyze", model_id(model)?);
ApiUrl::parse(endpoint)
.and_then(|url| url.complete_path(&["documentintelligence", "documentModels", &model]))
.map(|url| {
url.append_query_pairs(
[("api-version", AZURE_DI_API_VERSION)]
[("api-version", api_version)]
.into_iter()
.chain(params.pages.iter().map(|pages| ("pages", pages.as_str())))
.chain(
@ -580,8 +585,8 @@ impl AzureDocumentIntelligenceOcrConfig {
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, Error> {
if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization")
|| crate::custom_httpx::http_handler::has_header(
if litellm_http::request::has_header(&connection.extra_headers, "authorization")
|| litellm_http::request::has_header(
&connection.extra_headers,
AZURE_DI_SUBSCRIPTION_HEADER,
)

27
litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs
View file

@ -7,12 +7,12 @@ use crate::{
base_llm::ocr::{
document::{inline_remote_document, validate_inline_document},
error::Error,
handler::OcrClient,
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrConnection, OcrDocument, OcrRequestContext,
OcrResponseFormat, PreparedOcrRequest, credential_env,
OcrResponseFormat, PreparedOcrRequest,
},
},
custom_httpx::llm_http_handler::OcrClient,
mistral::ocr::transformation::{MistralOcrConfig, MistralOcrRequest},
};
@ -50,15 +50,11 @@ impl BaseOcrConfig for AzureAiOcrConfig {
request: &PreparedOcrRequest,
_client: &OcrClient,
) -> Result<Self::Environment, Error> {
let config = AzureAuthInputs {
azure_ad_token_provider: request.azure_ad_token_provider.clone(),
..AzureAuthInputs::from_sourced_optional_params(
&request.optional_params,
&request.input_sources,
)?
};
self.resolve_headers(&request.connection, &config, &credential_env)
.await
let config = crate::azure_ai::ocr::common_utils::azure_auth_inputs(request)?;
self.resolve_headers(&request.connection, &config, &|name: &str| {
request.connection.secret(name)
})
.await
}
fn get_complete_url(
@ -67,7 +63,9 @@ impl BaseOcrConfig for AzureAiOcrConfig {
_optional_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, Error> {
self.build_ocr_url(request.connection.api_base.as_deref(), &credential_env)
self.build_ocr_url(request.connection.api_base.as_deref(), &|name: &str| {
request.connection.secret(name)
})
}
fn transform_ocr_request(
@ -107,7 +105,7 @@ impl BaseOcrConfig for AzureAiOcrConfig {
}
fn validate_request_body(&self, body: &Value) -> Result<(), Error> {
validate_inline_document(&crate::custom_httpx::llm_http_handler::body_document(body)?)
validate_inline_document(&crate::base_llm::ocr::handler::body_document(body)?)
}
}
@ -134,8 +132,7 @@ impl AzureAiOcrConfig {
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Vec<(String, String)>, Error> {
Self::resolve_api_base(connection.api_base.as_deref(), env_lookup)?;
if crate::custom_httpx::http_handler::has_header(&connection.extra_headers, "authorization")
{
if litellm_http::request::has_header(&connection.extra_headers, "authorization") {
if config.azure_ad_token_provider.is_some() {
super::common_utils::resolve_entra(config, env_lookup).await?;
}

11
litellm-rust/crates/llms/src/base_llm/audio_transcription/transformation.rs
View file

@ -21,14 +21,7 @@ impl AudioTranscriptionResponseData {
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AudioTranscriptionAuth {
Bearer,
AwsSigV4 {
region: String,
service: &'static str,
},
}
pub use litellm_auth::RequestAuth;
pub trait BaseAudioTranscriptionConfig: Sync {
fn get_supported_openai_params(&self) -> &'static [&'static str];
@ -70,5 +63,5 @@ pub trait BaseAudioTranscriptionConfig: Sync {
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AudioTranscriptionAuth, Error>;
) -> Result<RequestAuth, Error>;
}

11
litellm-rust/crates/llms/src/base_llm/chat/transformation.rs
View file

@ -41,14 +41,7 @@ pub const STREAM_PARAM: &str = "stream";
/// presence does not make a request untranslatable.
const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"];
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ChatCompletionsAuth {
Header { name: &'static str, value: String },
Bearer { token: String },
AwsSigV4 { region: String },
}
pub use litellm_auth::RequestAuth;
/// Why a request cannot be served by the Rust path.
///
@ -91,7 +84,7 @@ pub trait BaseConfig: Sync {
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error>;
) -> Result<RequestAuth, Error>;
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[("content-type", "application/json")]

26
litellm-rust/crates/llms/src/base_llm/ocr/document.rs
View file

@ -1,18 +1,14 @@
use base64::{Engine, engine::general_purpose::STANDARD};
use data_url::{DataUrl, DataUrlError, forgiving_base64::DecodeError, mime::Mime};
use litellm_http::{
media::{DownloadPolicy, Error as MediaError, MediaFetcher},
transport::Error as TransportError,
};
use reqwest::Url;
use crate::{
base_llm::ocr::{
error::Error,
transformation::{
OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS, OcrConnection, OcrDocument,
},
},
custom_httpx::{
media::{DownloadPolicy, Error as MediaError, MediaFetcher},
transport::Error as TransportError,
},
use crate::base_llm::ocr::{
error::Error,
transformation::{OCR_INLINE_MAX_BYTES, OCR_MAX_FETCH_REDIRECTS, OcrConnection, OcrDocument},
};
pub struct InlineDocument<'a>(DataUrl<'a>);
@ -72,7 +68,7 @@ pub async fn inline_remote_document(
url,
DownloadPolicy {
timeout: connection.timeout,
max_bytes: connection.max_download_bytes,
max_bytes: connection.settings.max_download_bytes,
max_redirects: OCR_MAX_FETCH_REDIRECTS,
},
)
@ -196,10 +192,8 @@ mod tests {
.redirect(reqwest::redirect::Policy::none())
.build()
.unwrap();
let client = crate::custom_httpx::llm_http_handler::OcrClient::for_test(
provider_http,
document_http,
);
let client =
crate::base_llm::ocr::handler::OcrClient::for_test(provider_http, document_http);
let converted = inline_remote_document(
client.document_fetcher(),
OcrDocument::ImageUrl {

18
litellm-rust/crates/llms/src/base_llm/ocr/error.rs
View file

@ -76,6 +76,12 @@ pub enum Error {
Unsupported(&'static str),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid model: {provider} has no model {model:?} - use one of: {}", supported.join(", "))]
InvalidModel {
provider: &'static str,
model: String,
supported: &'static [&'static str],
},
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
@ -95,11 +101,13 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Params(#[from] litellm_core_utils::params::Error),
#[error(transparent)]
Headers(#[from] crate::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
}
impl From<litellm_host::machine::MachineFault> for Error {
@ -125,9 +133,7 @@ impl Error {
pub fn http_status_code(&self) -> Option<u16> {
match self {
Self::Provider { status, .. }
| Self::Transport(crate::custom_httpx::transport::Error::Http { status, .. }) => {
Some(*status)
}
| Self::Transport(litellm_http::transport::Error::Http { status, .. }) => Some(*status),
error if error.is_request() => Some(400),
_ => None,
}
@ -155,8 +161,10 @@ impl Error {
| Self::DotModel
| Self::InvalidRequest(_)
| Self::InvalidProvider(_)
| Self::InvalidModel { .. }
| Self::Params(_)
| Self::Headers(_)
| Self::Http(_)
)
}

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