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

test(e2e): adapt OAuth acceptance to merged SDK2

Browse source
This commit is contained in:
Joshua Valluru 2026-09-19 11:59:29 -07:00
commit 09a2e5b6ec
171 changed files with 6515 additions and 3044 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

@ -63,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
@ -125,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'

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

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

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

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

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

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

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

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

3
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",
@ -2142,6 +2143,7 @@ dependencies = [
"reqwest 0.12.28",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
@ -2174,6 +2176,7 @@ dependencies = [
"serde_json",
"serde_path_to_error",
"serde_with",
"strum",
"thiserror 2.0.19",
"time",
"tokio",

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

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 {

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

@ -24,6 +24,8 @@ pub enum Error {
#[error(transparent)]
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,4 +1,4 @@
use litellm_http::request::{http_request, truncate_error_body};
use litellm_http::request::truncate_error_body;
use serde_json::Value;
use super::{Error, client::http_client};
@ -7,17 +7,18 @@ 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| {
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();
@ -37,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())
}

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

@ -1,9 +1,7 @@
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,
};
@ -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/error.rs
View file

@ -24,6 +24,8 @@ pub enum Error {
#[error(transparent)]
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,5 +1,5 @@
use litellm_http::request::{http_request, truncate_error_body};
use litellm_llms::base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData};
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;
@ -12,22 +12,9 @@ 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.
@ -77,57 +64,24 @@ pub(super) fn as_response_error(err: Error) -> Error {
}
}
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,
})
}

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

@ -1,6 +1,6 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_http::request::has_header;
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::Value;
@ -67,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(
@ -77,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
@ -92,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
@ -105,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() {

23
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",
..
}
@ -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()
}
);

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;

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

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;

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

@ -19,7 +19,10 @@ pub(crate) fn prepare_request(
Some("MISTRAL_AZURE_API_BASE"),
),
OcrProvider::AzureAi => (None, Some("AZURE_AI_API_BASE")),
OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => (None, None),
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(|| {

36
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,
@ -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 => {
@ -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)]

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

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

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

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

@ -160,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!({
@ -250,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 {
@ -260,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/Cargo.toml
View file

@ -14,6 +14,7 @@ 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

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

@ -8,6 +8,12 @@ pub enum Error {
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 {

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

@ -1,6 +1,7 @@
mod config;
mod error;
pub mod media;
pub mod outbound;
mod pool;
mod proxy;
pub mod request;

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

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

@ -27,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")
);
}
}

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

10
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}")]
@ -100,6 +106,8 @@ pub enum Error {
Params(#[from] litellm_core_utils::params::Error),
#[error(transparent)]
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
}
impl From<litellm_host::machine::MachineFault> for Error {
@ -153,8 +161,10 @@ impl Error {
| Self::DotModel
| Self::InvalidRequest(_)
| Self::InvalidProvider(_)
| Self::InvalidModel { .. }
| Self::Params(_)
| Self::Headers(_)
| Self::Http(_)
)
}

64
litellm-rust/crates/llms/src/base_llm/ocr/handler.rs
View file

@ -5,7 +5,7 @@ use litellm_host::event::WireRequest;
use litellm_http::{
ClientVariant, HttpClientConfig, HttpClientPool,
media::{MediaFetcher, UrlPolicy},
request::{HeaderPolicy, execute_http_request, with_headers},
outbound::{OutboundRequest, RequestSigner},
transport,
};
use serde::{Serialize, de::DeserializeOwned};
@ -117,8 +117,9 @@ pub async fn ocr<C: BaseOcrConfig>(
) -> Result<LiteLLMOcrResponse, Error> {
let http = config.prepare_request(request, client, hooks).await?;
let url = http.url().to_string();
let headers = request_headers(&http)?;
let response = execute_http_request(client.provider_http(), http)
let headers = http.headers().to_vec();
let response = http
.send(client.provider_http())
.await
.map_err(transport_error)?;
if !response.status().is_success() {
@ -153,21 +154,6 @@ pub async fn ocr<C: BaseOcrConfig>(
.await
}
fn request_headers(request: &reqwest::Request) -> Result<Vec<(String, String)>, Error> {
request
.headers()
.iter()
.map(|(name, value)| {
value
.to_str()
.map(|value| (name.to_string(), value.to_string()))
.map_err(|_| Error::RequestField {
path: "headers".into(),
})
})
.collect()
}
pub async fn read_json_response<T: DeserializeOwned>(
response: reqwest::Response,
native: bool,
@ -222,13 +208,13 @@ pub fn transport_error(error: reqwest::Error) -> Error {
pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
config: &C,
client: &OcrClient,
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: B,
signer: Option<&dyn RequestSigner>,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
let composed = litellm_core_utils::call_arguments::compose_body(
&request.optional_params,
&body,
@ -244,7 +230,17 @@ pub async fn transform_request_body<C: BaseOcrConfig, B: Serialize>(
});
}
config.validate_request_body(&changed.body)?;
build_http_request(client, request, url, &changed.headers, &changed.body)
let timeout = Some(request.connection.timeout);
Ok(match signer {
Some(signer) => OutboundRequest::signed_json(
url.into(),
changed.headers,
&changed.body,
timeout,
signer,
),
None => OutboundRequest::json(url.into(), changed.headers, &changed.body, timeout),
}?)
}
fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireRequest {
@ -255,22 +251,18 @@ fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireReq
}
}
pub fn build_http_request<B: Serialize>(
client: &OcrClient,
pub fn build_http_request(
request: &PreparedOcrRequest,
url: &str,
headers: &[(String, String)],
body: &B,
) -> Result<reqwest::Request, Error> {
let builder = client
.provider_http()
.post(url)
.json(body)
.timeout(request.connection.timeout);
with_headers(builder, headers, HeaderPolicy::All)
.build()
.map_err(transport::Error::from)
.map_err(Error::from)
url: String,
headers: Vec<(String, String)>,
body: &impl Serialize,
) -> Result<OutboundRequest, Error> {
Ok(OutboundRequest::json(
url,
headers,
body,
Some(request.connection.timeout),
)?)
}
pub async fn guardrail_document(

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

@ -6,6 +6,7 @@ use litellm_core_utils::{
serde_compat::{FiniteF64, LaxI64},
settings::ProcessEnvironment,
};
use litellm_http::outbound::{OutboundRequest, RequestSigner};
use serde::{
Deserialize, Serialize,
de::{DeserializeOwned, IntoDeserializer},
@ -394,6 +395,10 @@ const HEALTH_CHECK_PDF_DATA_URI: &str = "data:application/pdf;base64,JVBERi0xLjQ
/// (headers at minimum; Vertex also carries the project id).
pub trait OcrEnvironment: Send + Sync {
fn headers(&self) -> &[(String, String)];
fn signer(&self) -> Option<&dyn RequestSigner> {
None
}
}
impl OcrEnvironment for Vec<(String, String)> {
@ -536,7 +541,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static {
request: &PreparedOcrRequest,
client: &OcrClient,
hooks: &dyn CallHooks<Error>,
) -> impl Future<Output = Result<reqwest::Request, Error>> + Send {
) -> impl Future<Output = Result<OutboundRequest, Error>> + Send {
async move {
let params = self.map_ocr_params(&request.optional_params, &request.model)?;
let environment = self.validate_environment(request, client).await?;
@ -554,7 +559,16 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static {
},
)
.await?;
transform_request_body(self, client, request, &url, headers, body, hooks).await
transform_request_body(
self,
request,
&url,
headers,
body,
environment.signer(),
hooks,
)
.await
}
}
}

8
litellm-rust/crates/llms/src/bedrock/audio_transcription/mod.rs
View file

@ -8,8 +8,8 @@ use serde_json::{Map, Value, json};
use crate::base_llm::{
audio_transcription::transformation::{
AudioTranscriptionAuth, AudioTranscriptionRequestData, AudioTranscriptionResponseData,
BaseAudioTranscriptionConfig,
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
BaseAudioTranscriptionConfig, RequestAuth,
},
chat::transformation::Error,
};
@ -136,9 +136,9 @@ impl BaseAudioTranscriptionConfig for BedrockAudioTranscriptionConfig {
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AudioTranscriptionAuth, Error> {
) -> Result<RequestAuth, Error> {
let (_, model_region) = bedrock_model_id_and_region(model);
Ok(AudioTranscriptionAuth::AwsSigV4 {
Ok(RequestAuth::AwsSigV4 {
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
service: BEDROCK_SERVICE,
})

13
litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs
View file

@ -1,6 +1,6 @@
use litellm_auth_aws::{
bedrock_model_id_and_region,
constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE},
constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE},
resolve_bedrock_region,
};
use litellm_core_utils::{
@ -17,8 +17,8 @@ use litellm_types::{
use serde_json::{Map, Value, json};
use crate::base_llm::chat::transformation::{
BaseConfig, ChatCompletionsAuth, Error, ProviderChatRequestData, ProviderChatResponseData,
Unsupported, unsupported_message, unsupported_param,
BaseConfig, Error, ProviderChatRequestData, ProviderChatResponseData, RequestAuth, Unsupported,
unsupported_message, unsupported_param,
};
/// Converse parameter names, post `map_openai_params`, that the Rust path can
@ -186,7 +186,7 @@ impl BaseConfig for AmazonConverseConfig {
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
) -> Result<RequestAuth, Error> {
// Python reads `api_key` as the Bedrock bearer token and consults the
// env only when the caller passed none, so a caller-supplied empty key
// falls through to SigV4 without reaching for the environment. An
@ -199,11 +199,12 @@ impl BaseConfig for AmazonConverseConfig {
}
.filter(|token| !token.is_empty());
if let Some(token) = bearer {
return Ok(ChatCompletionsAuth::Bearer { token });
return Ok(RequestAuth::Bearer { token });
}
let (_, model_region) = bedrock_model_id_and_region(model);
Ok(ChatCompletionsAuth::AwsSigV4 {
Ok(RequestAuth::AwsSigV4 {
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
service: BEDROCK_SERVICE,
})
}

10
litellm-rust/crates/llms/src/bedrock/chat/tests.rs
View file

@ -281,8 +281,9 @@ fn signs_with_sigv4_in_the_resolved_region() {
&|_| None
)
.expect("auth resolves"),
ChatCompletionsAuth::AwsSigV4 {
region: "eu-central-1".to_string()
RequestAuth::AwsSigV4 {
region: "eu-central-1".to_string(),
service: "bedrock",
}
);
}
@ -306,11 +307,12 @@ fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() {
)
.expect("auth resolves")
};
let bearer = |token: &str| ChatCompletionsAuth::Bearer {
let bearer = |token: &str| RequestAuth::Bearer {
token: token.to_string(),
};
let sigv4 = ChatCompletionsAuth::AwsSigV4 {
let sigv4 = RequestAuth::AwsSigV4 {
region: "eu-central-1".to_string(),
service: "bedrock",
};
// A caller-supplied key is the bearer token, and outranks the env.

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

@ -1,4 +1,5 @@
pub mod anthropic;
pub mod aws_textract;
pub mod azure_ai;
pub mod base_llm;
pub mod bedrock;

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

@ -5,6 +5,7 @@ use litellm_core_utils::{
params::OpaqueParams,
url_utils::ApiUrl,
};
use litellm_http::outbound::OutboundRequest;
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value, json};
@ -166,7 +167,7 @@ impl BaseOcrConfig for ReductoParseV3Config {
request: &PreparedOcrRequest,
client: &OcrClient,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
prepare_upload_request(self, request, client, hooks).await
}
}
@ -251,7 +252,7 @@ impl BaseOcrConfig for ReductoParseLegacyConfig {
request: &PreparedOcrRequest,
client: &OcrClient,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
prepare_upload_request(self, request, client, hooks).await
}
}
@ -264,7 +265,7 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
request: &PreparedOcrRequest,
client: &OcrClient,
hooks: &dyn CallHooks<Error>,
) -> Result<reqwest::Request, Error> {
) -> Result<OutboundRequest, Error> {
let params = config.map_ocr_params(&request.optional_params, &request.model)?;
let headers = config.validate_environment(request, client).await?;
let url = config.get_complete_url(request, &params, &headers)?;
@ -286,7 +287,7 @@ async fn prepare_upload_request<C: BaseOcrConfig<Environment = Vec<(String, Stri
&body,
config.get_supported_ocr_params(&request.model),
)?;
build_http_request(client, request, &url, &headers, &body)
build_http_request(request, url, headers, &body)
}
fn uploaded_file_id(document: OcrDocument) -> Result<ReductoFileId, Error> {

65
litellm-rust/crates/llms/src/vertex_ai/ocr/deepseek_transformation.rs
View file

@ -19,6 +19,13 @@ const MODEL_PREFIX: &str = "deepseek-ai/";
const DEFAULT_LOCATION: &str = "us-central1";
const DEEPSEEK_OCR_PARAMS: &[&str] = &["stream", "temperature", "max_tokens", "top_p", "n", "stop"];
/// DeepSeek-OCR is a transcription model: at the endpoint's default sampling temperature it
/// hallucinates extra text, so requests are greedy unless the caller sets a temperature.
const DEFAULT_TEMPERATURE: f64 = 0.0;
/// Greedy decoding on dense screenshots falls into repetition loops that run to the token limit;
/// a mild penalty breaks them without changing clean-document output.
const DEFAULT_REPETITION_PENALTY: f64 = 1.05;
pub type DeepSeekOcrParams = OpaqueParams;
#[derive(Clone, Debug, Serialize, Deserialize)]
@ -171,11 +178,19 @@ impl BaseOcrConfig for VertexAIDeepSeekOCRConfig {
image_url: document.source().to_string(),
}],
}],
params: optional_params
.iter()
.filter(|(name, _)| DEEPSEEK_OCR_PARAMS.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
params: [
("temperature", DEFAULT_TEMPERATURE),
("repetition_penalty", DEFAULT_REPETITION_PENALTY),
]
.into_iter()
.map(|(name, value)| (name.to_string(), Value::from(value)))
.chain(
optional_params
.iter()
.filter(|(name, _)| DEEPSEEK_OCR_PARAMS.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone())),
)
.collect(),
})
}
}
@ -484,6 +499,46 @@ mod tests {
assert!(result.get("ignored").is_none());
}
#[test]
fn request_uses_greedy_defaults_unless_the_caller_overrides_them() {
let request = |params: DeepSeekOcrParams| {
serde_json::to_value(
VertexAIDeepSeekOCRConfig
.transform_ocr_request(
"deepseek-ai/deepseek-ocr-maas",
document(),
&params,
&[],
)
.unwrap(),
)
.unwrap()
};
let defaults = request(DeepSeekOcrParams::default());
assert_eq!(defaults["temperature"], 0.0);
assert_eq!(defaults["repetition_penalty"], 1.05);
assert_eq!(
request(serde_json::from_value(json!({"temperature":0.7})).unwrap())["temperature"],
0.7
);
}
#[test]
fn caller_temperature_argument_overrides_the_greedy_default_in_the_composed_body() {
let arguments = serde_json::from_value(json!({"temperature":0.7})).unwrap();
let body = VertexAIDeepSeekOCRConfig
.transform_ocr_request(
"deepseek-ai/deepseek-ocr-maas",
document(),
&DeepSeekOcrParams::default(),
&[],
)
.unwrap();
let composed =
litellm_core_utils::call_arguments::compose_body(&arguments, &body, &[]).unwrap();
assert_eq!(composed["temperature"], 0.7);
}
#[rstest]
#[case(json!({"type":"image_url","image_url":"data:image/png;base64,AA=="}))]
#[case(json!({"type":"document_url","document_url":"data:application/pdf;base64,AA=="}))]

3
litellm-rust/crates/python-bridge/src/errors.rs
View file

@ -58,6 +58,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr {
audio_transcription::Error::InvalidProvider(_)
| audio_transcription::Error::InvalidRequest(_)
| audio_transcription::Error::Headers(_)
| audio_transcription::Error::Http(_)
| audio_transcription::Error::InvalidType { .. }
| audio_transcription::Error::MissingField(_)
| audio_transcription::Error::Aws(_) => true,
@ -68,6 +69,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr {
chat_completions::Error::InvalidProvider(_)
| chat_completions::Error::InvalidRequest(_)
| chat_completions::Error::Headers(_)
| chat_completions::Error::Http(_)
| chat_completions::Error::InvalidType { .. }
| chat_completions::Error::MissingField(_)
| chat_completions::Error::Aws(_) => true,
@ -105,6 +107,7 @@ pub(crate) fn chat_completions_error_to_pyerr(error: chat_completions::Error) ->
| Error::InvalidType { .. }
| Error::MissingField(_)
| Error::Headers(_)
| Error::Http(_)
| Error::Transport(TransportError::Connect(_)) => {
RustBridgeDeclined::new_err(error.to_string())
}

6
litellm-rust/crates/python-bridge/src/routes/ocr/document.rs
View file

@ -7,7 +7,8 @@ use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
pybacked::PyBackedBytes,
types::{PyBytes, PyString},
sync::PyOnceLock,
types::{PyBytes, PyString, PyType},
};
#[derive(Debug)]
@ -84,7 +85,8 @@ impl FromPyObject<'_, '_> for FileDocumentInput {
"OCR file input does not accept bare str values. Pass bytes, a pathlib.Path, or a file-like object.",
));
}
if file.is_instance(&py.import("os")?.getattr("PathLike")?)? {
static PATH_LIKE: PyOnceLock<Py<PyType>> = PyOnceLock::new();
if file.is_instance(PATH_LIKE.import(py, "os", "PathLike")?)? {
return Ok(Self {
input: OcrDocumentInput::Path {
path: file.extract::<PathBuf>()?,

13
litellm/experimental_mcp_client/Readme.md
View file

@ -1,6 +1,17 @@
# LiteLLM MCP Client
LiteLLM MCP Client is a client that allows you to use MCP tools with LiteLLM.
LiteLLM MCP Client allows you to use MCP tools with LiteLLM
## MCP Python SDK compatibility
The `mcp` and `proxy` extras require MCP Python SDK 2.2 or newer within the 2.x release line. Installing core LiteLLM without these extras does not require MCP
Existing MCP SDK1 clients can continue connecting to the gateway over the supported legacy MCP protocols. The client and gateway can use different SDK versions in separate Python environments. Modern protocol advertisement remains disabled during the Phase 0 upgrade. An initialize body requesting `2026-07-28` falls back to the supported legacy version `2025-11-25`; an explicit `MCP-Protocol-Version: 2026-07-28` HTTP header is rejected with HTTP 400
Code sharing the gateway's Python environment must support SDK2. Its Python API has breaking changes, including renamed imports and snake_case model attributes such as `input_schema`, `is_error`, and `structured_content`. This also applies to callers consuming SDK objects returned by LiteLLM's experimental MCP client. MCP JSON fields retain their protocol spelling, such as `inputSchema` and `isError`
Upgrade SDK1-dependent libraries before installing them alongside `litellm[mcp]` or `litellm[proxy]`, or keep those clients in a separate environment and connect over the network. For example, `langchain-mcp-adapters==0.2.1` uses SDK1 Python APIs and is tested as a separate legacy client, not as a shared SDK2 dependency
The shared unit-test workflow runs the MCP integration suite once, with SDK2 in the gateway environment and an isolated SDK1 peer. Keep the SDK1 list/call compatibility test while SDK1 clients are supported; remove it when that support is explicitly retired and the client migration is documented
See the official [SDK migration guide](https://py.sdk.modelcontextprotocol.io/migration/) for Python API changes

170
litellm/experimental_mcp_client/client.py
View file

@ -9,57 +9,30 @@ import json
import os
from collections.abc import Awaitable, Callable, Generator
from contextlib import AbstractAsyncContextManager
from datetime import timedelta
from functools import partial
from importlib import metadata
from types import MappingProxyType
from typing import Any, Final, Protocol, TypeAlias, TypeVar
from typing import Any, Final, TypeAlias, TypeVar
import httpx
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServerParameters
import httpx2
from httpx2._client import UseClientDefault
from httpx2._types import AuthTypes
from mcp import ClientSession, MCPError, ReadResourceResult, Resource, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from typing_extensions import Unpack
_TransportStreams: TypeAlias = tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
Unpack[tuple[object, ...]],
ReadStream[SessionMessage | Exception],
WriteStream[SessionMessage],
]
_TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams]
class _StreamableHttpClientFactory(Protocol):
"""The ``streamable_http_client`` entry point this module calls on the installed MCP SDK."""
def __call__(self, *, url: str, http_client: httpx.AsyncClient | None) -> _TransportContext: ...
streamable_http_client: _StreamableHttpClientFactory | None = None
try:
import mcp.client.streamable_http as streamable_http_module
streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None)
except ImportError:
pass
MCP_STREAMABLE_HTTP_REQUIREMENT: Final = "mcp>=1.28.1"
def missing_streamable_http_client_error() -> ImportError:
return ImportError(
f"MCP streamable HTTP transport requires {MCP_STREAMABLE_HTTP_REQUIREMENT}, but the installed "
f"mcp {metadata.version('mcp')} does not provide streamable_http_client. "
"Fix with: pip install 'litellm[mcp]' (or upgrade mcp directly: pip install -U mcp)"
)
from mcp.types import (
METHOD_NOT_FOUND,
ClientResult,
REQUEST_TIMEOUT,
GetPromptRequestParams,
GetPromptResult,
ListPromptsResult,
@ -68,7 +41,6 @@ from mcp.types import (
Prompt,
ResourceTemplate,
ServerNotification,
ServerRequest,
TextContent,
)
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
@ -153,23 +125,21 @@ def _first_non_cancelled_cause(exc: BaseException) -> BaseException | None:
return None
_SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT)
"""The code the MCP SDK puts on its own elapsed read timeout, an HTTP status in a field that
otherwise carries JSON-RPC error codes."""
_SDK_READ_TIMEOUT_CODE: Final = REQUEST_TIMEOUT
"""The code the MCP SDK puts on its own elapsed read timeout."""
def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
"""Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``.
The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a
field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error
through that same class and field. The numeric code alone therefore cannot separate the two, and
an upstream answering with application code 408 would be reported as a gateway timeout it never
caused. The SDK raises its own from inside an ``except TimeoutError``, so the elapsed timeout is
The SDK reports its own elapsed read timeout as ``MCPError`` carrying ``REQUEST_TIMEOUT`` in a
field that also carries relayed upstream JSON-RPC errors. The numeric code alone therefore
cannot separate the two, and an upstream answering with the same application code would be
reported as a gateway timeout it never caused. The SDK raises its own from inside an ``except TimeoutError``, so the elapsed timeout is
on the context chain, while a relayed error is built from a received message and has no such
chain; that is the discriminator.
"""
if not isinstance(exc, McpError) or exc.error.code != _SDK_READ_TIMEOUT_CODE:
if not isinstance(exc, MCPError) or exc.error.code != _SDK_READ_TIMEOUT_CODE:
return None
if not isinstance(exc.__context__, TimeoutError):
return None
@ -179,9 +149,25 @@ def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
TSessionResult = TypeVar("TSessionResult")
class MCPSigV4Auth(httpx.Auth):
class _MCPHTTPClient(httpx2.AsyncClient):
async def send(
self,
request: httpx2.Request,
*,
stream: bool = False,
auth: AuthTypes | UseClientDefault | None = httpx2.USE_CLIENT_DEFAULT,
follow_redirects: bool | UseClientDefault = httpx2.USE_CLIENT_DEFAULT,
) -> httpx2.Response:
response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects)
if request.method == "POST" and response.is_error and response.status_code != 404:
await response.aclose()
response.raise_for_status()
return response
class MCPSigV4Auth(httpx2.Auth):
"""
httpx Auth class that signs each request with AWS SigV4.
httpx2 Auth class that signs each request with AWS SigV4.
This is used for MCP servers that require AWS SigV4 authentication,
such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow()
for every outgoing request, enabling per-request signature computation.
@ -270,7 +256,7 @@ class MCPSigV4Auth(httpx.Auth):
token=sts_creds["SessionToken"],
)
def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
@ -314,8 +300,8 @@ class MCPClient:
stdio_config: MCPStdioConfig | None = None,
extra_headers: dict[str, str] | None = None,
ssl_verify: VerifyTypes | None = None,
aws_auth: httpx.Auth | None = None,
resolved_auth: httpx.Auth | None = None,
aws_auth: httpx2.Auth | None = None,
resolved_auth: httpx2.Auth | None = None,
sampling_callback: Callable | None = None,
elicitation_callback: Callable | None = None,
logging_callback: Callable | None = None,
@ -333,10 +319,10 @@ class MCPClient:
self.stdio_config: MCPStdioConfig | None = stdio_config
self.extra_headers: dict[str, str] | None = extra_headers
self.ssl_verify: VerifyTypes | None = ssl_verify
self._aws_auth: httpx.Auth | None = aws_auth
# A pre-resolved httpx.Auth (e.g. from the v2 credential resolver) attached to the
self._aws_auth: httpx2.Auth | None = aws_auth
# A pre-resolved httpx2.Auth (e.g. from the v2 credential resolver) attached to the
# upstream client's auth= slot, taking precedence over the SigV4 aws_auth.
self._resolved_auth: httpx.Auth | None = resolved_auth
self._resolved_auth: httpx2.Auth | None = resolved_auth
self._last_initialize_instructions: str | None = None
self._sampling_callback: Callable | None = sampling_callback
self._elicitation_callback: Callable | None = elicitation_callback
@ -348,9 +334,11 @@ class MCPClient:
async def discovery_auth_fingerprint(self) -> str:
return self._hash_discovery_auth(await self.prepare_request_auth())
async def prepare_request_auth(self) -> httpx.Request:
async def prepare_request_auth(self) -> httpx2.Request:
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
request: Final = httpx2.Request(
"POST", self.server_url or "http://localhost/", headers=self._get_auth_headers()
)
if self._resolved_auth is None:
return request
flow: Final = self._resolved_auth.async_auth_flow(request)
@ -361,20 +349,20 @@ class MCPClient:
await flow.aclose()
@staticmethod
def _hash_discovery_auth(request: httpx.Request) -> str:
def _hash_discovery_auth(request: httpx2.Request) -> str:
material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items()))))
return hashlib.sha256(material.encode()).hexdigest()
def _create_transport_context(
self,
) -> tuple[_TransportContext, httpx.AsyncClient | None]:
) -> tuple[_TransportContext, httpx2.AsyncClient | None]:
"""
Create the appropriate transport context based on transport type.
Returns:
Tuple of (transport_context, http_client).
http_client is only set for HTTP transport and needs cleanup.
"""
http_client: httpx.AsyncClient | None = None
http_client: httpx2.AsyncClient | None = None
if self.transport_type == MCPTransport.stdio:
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
@ -397,14 +385,12 @@ class MCPClient:
None,
)
# HTTP transport (default)
if streamable_http_client is None:
raise missing_streamable_http_client_error()
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
http_client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(self.timeout),
timeout=httpx2.Timeout(self.timeout),
)
transport_ctx: Final = streamable_http_client(
url=self.server_url,
@ -473,13 +459,14 @@ class MCPClient:
transport: Final = await transport_ctx.__aenter__()
in_flight_error: BaseException | None = None
try:
read_stream, write_stream = transport[0], transport[1]
read_stream: Final = transport[0]
write_stream: Final = transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
async def receive_message(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
message: ServerNotification | Exception,
) -> None:
if not isinstance(message, (ValueError, httpx.RequestError, OSError)):
if not isinstance(message, (ValueError, httpx2.HTTPError, OSError)):
return
if not stream_error.done():
stream_error.set_result(message)
@ -499,7 +486,7 @@ class MCPClient:
session_ctx: Final = ClientSession(
read_stream,
write_stream,
read_timeout_seconds=timedelta(seconds=self.timeout),
read_timeout_seconds=self.timeout,
message_handler=receive_message,
**session_kwargs,
)
@ -512,7 +499,7 @@ class MCPClient:
if isinstance(ins, str) and ins.strip():
self._last_initialize_instructions = ins.strip()
return await operation(session)
except McpError:
except MCPError:
if stream_error.done():
raise stream_error.result()
raise
@ -544,7 +531,7 @@ class MCPClient:
quiet_on_error demotes the failure line to debug for callers that own the exception
(call_tool / list_tools under raise_on_error), so an expected pass-through re-auth does
not emit a warning per call; every other caller keeps the operator-visible warning."""
http_client: httpx.AsyncClient | None = None
http_client: httpx2.AsyncClient | None = None
try:
self._last_initialize_instructions = None
transport_ctx, http_client = self._create_transport_context()
@ -609,7 +596,7 @@ class MCPClient:
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
# signing (including the body hash), so it uses httpx.Auth flow instead
# signing (including the body hash), so it uses httpx2.Auth flow instead
# of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
# update the headers with the extra headers
if self.extra_headers:
@ -623,9 +610,11 @@ class MCPClient:
headers.update(injected or {})
return _strip_header_whitespace(headers)
def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]:
def _create_httpx_client_factory(
self, *, transport: httpx2.AsyncBaseTransport | None = None
) -> Callable[..., httpx2.AsyncClient]:
"""
Create a custom httpx client factory that uses LiteLLM's SSL configuration.
Create a custom httpx2 client factory that uses LiteLLM's SSL configuration.
This factory follows the same CA bundle path logic as http_handler.py:
1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle)
2. Check SSL_VERIFY environment variable
@ -636,10 +625,10 @@ class MCPClient:
def factory(
*,
headers: dict[str, str] | None = None,
timeout: httpx.Timeout | None = None,
auth: httpx.Auth | None = None,
) -> httpx.AsyncClient:
"""Create an httpx.AsyncClient with LiteLLM's SSL configuration."""
timeout: httpx2.Timeout | None = None,
auth: httpx2.Auth | None = None,
) -> httpx2.AsyncClient:
"""Create an httpx2.AsyncClient with LiteLLM's SSL configuration."""
# Get unified SSL configuration using the same logic as http_handler.py
ssl_config: Final = get_ssl_configuration(self.ssl_verify)
verbose_logger.debug("MCP client using SSL configuration: %s", type(ssl_config).__name__)
@ -649,7 +638,8 @@ class MCPClient:
fallback_auth: Final = self._resolved_auth if self._resolved_auth is not None else self._aws_auth
effective_auth: Final = auth if auth is not None else fallback_auth
guard: Final = credential_redirect_hook(self.server_url, self._credential_slot)
return httpx.AsyncClient(
return _MCPHTTPClient(
transport=transport,
headers=headers,
timeout=timeout,
auth=effective_auth,
@ -723,7 +713,7 @@ class MCPClient:
"""The error result ``call_tool`` returns when it swallows a failure (no re-execution)."""
return MCPCallToolResult(
content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc}")],
isError=True,
is_error=True,
)
async def call_tool(
@ -808,12 +798,12 @@ class MCPClient:
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult:
capabilities: Final = session.get_server_capabilities()
capabilities: Final = session.server_capabilities
if capabilities is not None and capabilities.prompts is None:
return ListPromptsResult(prompts=[])
try:
return await session.list_prompts()
except McpError as error:
except MCPError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
verbose_logger.debug(
@ -898,12 +888,12 @@ class MCPClient:
verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio")
async def _list_resources_operation(session: ClientSession) -> ListResourcesResult:
capabilities: Final = session.get_server_capabilities()
capabilities: Final = session.server_capabilities
if capabilities is not None and capabilities.resources is None:
return ListResourcesResult(resources=[])
try:
return await session.list_resources()
except McpError as error:
except MCPError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
verbose_logger.debug(
@ -947,30 +937,30 @@ class MCPClient:
verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio")
async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult:
capabilities: Final = session.get_server_capabilities()
capabilities: Final = session.server_capabilities
if capabilities is not None and capabilities.resources is None:
return ListResourceTemplatesResult(resourceTemplates=[])
return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload
try:
return await session.list_resource_templates()
except McpError as error:
except MCPError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
verbose_logger.debug(
"MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error
)
return ListResourceTemplatesResult(resourceTemplates=[])
return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload
try:
result: Final = await self.run_with_session(_list_resource_templates_operation)
resource_template_count: Final = len(result.resourceTemplates)
resource_template_names: Final = [resourceTemplate.name for resourceTemplate in result.resourceTemplates]
resource_template_count: Final = len(result.resource_templates)
resource_template_names: Final = [resource_template.name for resource_template in result.resource_templates]
verbose_logger.info(
"MCP client listed %s resource templates from %s: %s",
resource_template_count,
self.server_url or "stdio",
resource_template_names,
)
return result.resourceTemplates
return result.resource_templates
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_resource_templates was cancelled")
raise
@ -1000,7 +990,7 @@ class MCPClient:
async def _read_resource_operation(session: ClientSession):
verbose_logger.debug("MCP client sending read_resource request to session")
return await session.read_resource(url)
return await session.read_resource(str(url))
try:
read_resource_result: Final = await self.run_with_session(_read_resource_operation)

8
litellm/experimental_mcp_client/tools.py
View file

@ -26,7 +26,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall
########################################################
def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam:
"""Convert an MCP tool to an OpenAI tool."""
normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.inputSchema)
normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.input_schema)
return ChatCompletionToolParam(
type="function",
@ -73,7 +73,7 @@ def transform_mcp_tool_to_openai_responses_api_tool(
mcp_tool: MCPTool,
) -> FunctionToolParam:
"""Convert an MCP tool to an OpenAI Responses API tool."""
normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.inputSchema)
normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.input_schema)
return FunctionToolParam(
name=mcp_tool.name,
@ -93,7 +93,7 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
return AnthropicMessagesTool(
name=mcp_tool.name,
description=mcp_tool.description or "",
input_schema=sanitize_input_schema_for_anthropic(mcp_tool.inputSchema),
input_schema=sanitize_input_schema_for_anthropic(mcp_tool.input_schema),
type="custom",
)
@ -129,7 +129,7 @@ async def list_tools_with_pagination(
)
tools.extend(result.tools)
next_cursor = getattr(result, "nextCursor", None)
next_cursor = getattr(result, "next_cursor", None)
if not isinstance(next_cursor, str) or not next_cursor:
return tools
if next_cursor in seen_cursors:

95
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -25,7 +25,10 @@ from litellm.integrations.prompt_management_base import PromptManagementClient
from litellm.litellm_core_utils.prompt_templates.common_utils import (
with_prompt_cache_breakpoint,
)
from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
from litellm.llms.anthropic.common_utils import (
is_claude_code_one_shot_subagent_request,
supports_anthropic_cache_control,
)
from litellm.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@ -574,8 +577,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider: str | None,
api_base: object,
prompt_cache_options: object,
request_kwargs: object,
) -> Sequence[Mapping[str, object]] | None:
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control, request_kwargs):
return None
return AnthropicCacheControlHook._stamped_with_dialect(
points, model, custom_llm_provider, api_base, prompt_cache_options
@ -612,6 +616,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
system: str | list | None,
tools: list | None,
cache_control: object = None,
request_kwargs: object = None,
) -> bool:
"""Whether configured injection points must yield to client-set cache_control.
@ -624,7 +629,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
"""
if all(point.get("_litellm_judged") for point in points):
return False
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
return AnthropicCacheControlHook._request_has_cache_control(
messages, system, tools, cache_control, request_kwargs
)
@staticmethod
def _request_has_cache_control(
@ -632,31 +639,29 @@ class AnthropicCacheControlHook(CustomPromptManagement):
system: str | list | None,
tools: list | None = None,
cache_control: object = None,
request_kwargs: object = None,
) -> bool:
"""Return True if the request already carries any client-supplied cache_control.
When the client (e.g. Claude Code) already marks its own breakpoints we
stand down entirely rather than add more, per the auto-caching contract.
Tools count: they are a breakpoint the client can mark, they count toward
the provider's four-block limit, and caching only the tool definitions is
a common pattern, so injecting alongside them can exceed the cap. Tools
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if cache_control is not None:
return True
if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
return True
if tools is not None:
return any(
isinstance(tool, dict)
and (
tool.get("cache_control") is not None
or (isinstance(tool.get("function"), dict) and tool["function"].get("cache_control") is not None)
)
for tool in tools
"""Client breakpoints own caching in both the request and its extra_body envelope."""
bodies: Final = (
{"messages": messages, "system": system, "tools": tools, "cache_control": cache_control},
_validated_object_mapping(AnthropicCacheControlHook._request_value(request_kwargs, "extra_body")) or {},
)
return any(
body.get("cache_control") is not None
or AnthropicCacheControlHook.count_request_cache_breakpoints(
_validated_object_list(body.get("messages")) or (), body.get("system")
)
return False
> 0
or any(
AnthropicCacheControlHook._request_value(tool, "cache_control") is not None
or AnthropicCacheControlHook._request_value(
AnthropicCacheControlHook._request_value(tool, "function"), "cache_control"
)
is not None
for tool in (_validated_object_list(body.get("tools")) or ())
)
for body in bodies
)
@staticmethod
def get_default_injection_points(
@ -676,36 +681,19 @@ class AnthropicCacheControlHook(CustomPromptManagement):
even when the global flag is off. Caches the system prompt and the
trailing turn, so the stable prefix (system + tools + history) is
reused while the breakpoint advances with the conversation. Returns []
(stand down) when neither flag is on, the provider does not consume
cache_control breakpoints (only anthropic / bedrock do), the model
lacks prompt-caching support, or the request already carries
client-supplied cache_control.
(stand down) when neither flag is on, the model is not Claude on a
supported explicit-cache transport, the model lacks prompt-caching
support, or the request already carries client-supplied cache_control.
"""
import litellm
if litellm.enable_anthropic_prompt_caching is not True and enable_prompt_caching is not True:
return []
provider = custom_llm_provider
if provider is None:
from litellm.litellm_core_utils.get_llm_provider_logic import (
get_llm_provider,
)
try:
_, provider, _, _ = get_llm_provider(model=model)
except Exception: # noqa: BLE001 # unroutable model must never block the call, just skip auto-caching
return []
if provider not in ("anthropic", "bedrock"):
if not supports_anthropic_cache_control(model, custom_llm_provider):
return []
from litellm.utils import supports_prompt_caching
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
return []
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control, request_kwargs):
return []
if is_claude_code_one_shot_subagent_request(
@ -737,13 +725,15 @@ class AnthropicCacheControlHook(CustomPromptManagement):
prompt and trailing turn) do not depend on which deployment serves the call. Returns the
input list itself when auto-injection would not apply
"""
import litellm
points: Final = next(
(
candidate
for candidate in (
AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
model=model,
model=litellm.model_alias_map.get(model, model),
custom_llm_provider=None,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
@ -789,6 +779,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
prompt-management gate and the AnthropicCacheControlHook run
unchanged.
"""
import litellm
if non_default_params.get("cache_control_injection_points"):
judged: Final = AnthropicCacheControlHook._judged_configured_points(
non_default_params["cache_control_injection_points"],
@ -799,6 +791,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider,
api_base,
non_default_params.get("prompt_cache_options"),
non_default_params,
)
if judged is None:
non_default_params.pop("cache_control_injection_points")
@ -808,7 +801,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
points: Final = AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
model=litellm.model_alias_map.get(model, model),
custom_llm_provider=custom_llm_provider,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
@ -925,7 +918,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
if configured and AnthropicCacheControlHook._should_stand_down(
configured, typed_messages, system, tools, cache_control
configured, typed_messages, system, tools, cache_control, kwargs
):
return messages, system
injection_points: list[CacheControlInjectionPoint] = configured or []

5
litellm/integrations/arize/_utils.py
View file

@ -1139,7 +1139,10 @@ def _set_mcp_tool_output(span: "Span", coerced_response_obj: object) -> None:
safe_set_attribute(span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.TEXT.value)
return
structured: Final[object] = coerced_response_obj.get("structuredContent")
structured: Final[object] = coerced_response_obj.get(
"structured_content",
coerced_response_obj.get("structuredContent"), # pyright: ignore[reportUnknownMemberType] # tolerant dual-spelling lookup on untyped payloads
)
payload: Final[object] = content if content else structured if structured is not None else content
if payload is None:
return

6
litellm/integrations/otel/mappers/langfuse.py
View file

@ -39,7 +39,7 @@ LANGFUSE_TRACE_TAGS: Final = "langfuse.trace.tags"
class LangfuseMapper:
_LLM_CALL_ATTRS: dict[str, Callable[[LLMCallSpanData], AttrValue | None]] = {
"langfuse.observation.type": lambda d: "generation",
"langfuse.observation.type": lambda _: "generation",
"langfuse.observation.model.name": lambda d: d.request_model or None,
"langfuse.observation.metadata.provider": lambda d: d.provider or None,
"langfuse.observation.id": lambda d: d.identity.call_id or None,
@ -68,7 +68,9 @@ class LangfuseMapper:
collect(LangfuseMapper._MODEL_PARAMS, d.request_params)
),
LANGFUSE_OBSERVATION_INPUT: lambda d: serialize_messages(d.messages_in),
LANGFUSE_OBSERVATION_OUTPUT: lambda d: serialize_messages(output_messages(d)),
LANGFUSE_OBSERVATION_OUTPUT: lambda d: (
d.embedding_output.as_json() if d.embedding_output is not None else serialize_messages(output_messages(d))
),
"langfuse.observation.usage_details": lambda d: json_if(collect(LangfuseMapper._USAGE_FIELDS, d.usage)),
"langfuse.observation.cost_details": lambda d: (
json.dumps({"total": d.response_cost}) if d.response_cost is not None else None

28
litellm/integrations/otel/model/payloads.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
import json
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from enum import Enum
from types import MappingProxyType
@ -353,6 +353,24 @@ class ToolDefinition:
parameters_json: str | None = None # JSON-serialized schema (str so it's an AttrValue)
@dataclass(frozen=True, slots=True)
class EmbeddingOutput:
count: int
dimensions: int | None
@classmethod
def from_response(cls, response: Mapping[str, object]) -> EmbeddingOutput | None:
vectors: Final = tuple(row.get("embedding") for row in _dicts(response.get("data")))
if not vectors:
return None
first: Final = vectors[0]
width: Final = len(cast(Sequence[object], first)) if isinstance(first, list) else None
return cls(count=len(vectors), dimensions=width)
def as_json(self) -> str:
return json.dumps({"count": self.count, "dimensions": self.dimensions})
@dataclass(frozen=True)
class LLMCallSpanData:
operation: GenAIOperation
@ -386,6 +404,7 @@ class LLMCallSpanData:
call_type: str | None = None
request_route: str | None = None
trace: TraceControls = field(default_factory=TraceControls)
embedding_output: EmbeddingOutput | None = None
@classmethod
def from_standard_logging_payload(
@ -413,8 +432,12 @@ class LLMCallSpanData:
# no prompt/response text.
finish_reasons: Final = _finish_reasons(choices_out)
call_type: Final = as_str(payload.get("call_type"))
operation: Final = resolve_operation(call_type)
embedding_output: Final = (
EmbeddingOutput.from_response(response) if operation is GenAIOperation.EMBEDDINGS else None
)
return cls(
operation=resolve_operation(call_type),
operation=operation,
provider=resolve_provider(as_str(payload.get("custom_llm_provider"))),
request_model=context.request_model,
response_model=context.response_model,
@ -437,6 +460,7 @@ class LLMCallSpanData:
call_type=call_type or None,
request_route=request_route or context.identity.request_route,
trace=trace or TraceControls(),
embedding_output=embedding_output if capture_content else None,
)

15
litellm/llms/anthropic/common_utils.py
View file

@ -77,6 +77,21 @@ _CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None) -> bool:
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.utils import supports_prompt_caching
try:
provider: Final = custom_llm_provider if custom_llm_provider is not None else get_llm_provider(model=model)[1]
except Exception: # noqa: BLE001 # Optional caching must not block an unroutable request
return False
return (
provider in ("anthropic", "bedrock", "vertex_ai", "azure_ai")
and "claude" in model.lower()
and supports_prompt_caching(model=model, custom_llm_provider=provider)
)
def is_claude_code_user_agent(user_agent: str) -> bool:
"""Claude Code sends its API calls through the Anthropic SDK as `claude-cli/<version>` and its own
fetches, such as gateway model discovery, as `claude-code/<version>`"""

484
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

4
litellm/ocr/dispatch.py
View file

@ -53,7 +53,9 @@ _PYTHON_AOCR: Final = cast( # cast-ok: forward the original call shape through
def _context(request: LiteLLMOcrRequest) -> Context:
return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model)
prefix, separator, _ = request.model.partition("/")
provider: Final = request.custom_llm_provider or (prefix if separator else None)
return Context(Route.OCR, provider=provider, model=request.model)
_DISPATCH: Final = PublicDispatch(

12
litellm/proxy/_experimental/mcp_server/elicitation_handler.py
View file

@ -42,9 +42,9 @@ class _DownstreamElicitSession(Protocol):
async def elicit_url(self, message: str, url: str, elicitation_id: str) -> "ElicitResult": ...
async def elicit_form(self, message: str, requestedSchema: dict[str, object]) -> "ElicitResult": ...
async def elicit_form(self, message: str, requested_schema: dict[str, object]) -> "ElicitResult": ...
async def elicit(self, message: str, requestedSchema: dict[str, object]) -> "ElicitResult": ...
async def elicit(self, message: str, requested_schema: dict[str, object]) -> "ElicitResult": ...
async def handle_elicitation_request(
@ -145,22 +145,22 @@ async def _relay_elicitation_to_downstream(
result = await downstream_session.elicit_url(
message=params.message,
url=params.url,
elicitation_id=params.elicitationId,
elicitation_id=params.elicitation_id,
)
elif isinstance(params, ElicitRequestFormParams):
# Form mode: relay structured form to client
verbose_logger.info("MCP elicitation: relaying form mode to downstream")
result = await downstream_session.elicit_form(
message=params.message,
requestedSchema=params.requestedSchema,
requested_schema=params.requested_schema,
)
else:
# Fallback for generic ElicitRequestParams — pass an empty schema
# since elicit() requires requestedSchema as a positional arg.
# since elicit() requires requested_schema as a positional arg.
verbose_logger.info("MCP elicitation: relaying generic elicitation to downstream")
result = await downstream_session.elicit(
message=getattr(params, "message", ""),
requestedSchema=getattr(params, "requestedSchema", {}),
requested_schema=getattr(params, "requested_schema", {}), # mutable-ok: elicitation default schema
)
verbose_logger.info(
"MCP elicitation: downstream responded with action=%s",

13
litellm/proxy/_experimental/mcp_server/faults/list_outcomes.py
View file

@ -14,6 +14,7 @@ from collections.abc import Iterator
from typing import Final, Literal, NamedTuple, NoReturn, TypeAlias
import httpx
import httpx2
from mcp.types import Tool as MCPTool
from pydantic import BaseModel, ConfigDict
from typing_extensions import assert_never
@ -63,8 +64,8 @@ class AggregateToolListing(NamedTuple):
outcomes: dict[str, ServerOutcome]
def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response]:
"""Yield every ``httpx.Response`` in the exception tree, in the shared traversal's deliberate
def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response | httpx2.Response]:
"""Yield every upstream ``httpx``/``httpx2`` ``Response`` in the exception tree, in the shared traversal's deliberate
order (explicit causes first, ExceptionGroup members in raise order, the incidental
``__context__`` chain last), so a response raised while handling the real failure can never
shadow one on the explicit causal chain. Consumers apply their own predicate over the stream:
@ -72,11 +73,11 @@ def _iter_upstream_responses(exc: BaseException) -> Iterator[httpx.Response]:
behind an unrelated earlier one."""
for current in iter_exception_tree(exc):
response = getattr(current, "response", None)
if isinstance(response, httpx.Response):
if isinstance(response, (httpx.Response, httpx2.Response)):
yield response
def _find_upstream_response(exc: BaseException) -> httpx.Response | None:
def _find_upstream_response(exc: BaseException) -> httpx.Response | httpx2.Response | None:
return next(_iter_upstream_responses(exc), None)
@ -136,9 +137,9 @@ def classify_list_exception(exc: BaseException) -> ServerListFault:
response: Final = _find_upstream_response(exc)
if response is not None:
return ServerListFault(tag="upstream_error", status_code=response.status_code)
if isinstance(exc, (httpx.TimeoutException,)):
if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)):
return ServerListFault(tag="timeout")
if isinstance(exc, httpx.TransportError):
if isinstance(exc, (httpx.TransportError, httpx2.TransportError)):
return ServerListFault(tag="unreachable")
return ServerListFault(tag="internal")

2
litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
View file

@ -135,7 +135,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
mcp_tool: Final = MCPTool(
name=mcp_tool_name,
description=mcp_tool_description or "",
inputSchema={}, # Call payload has no schema; guardrail gets args from request_data
input_schema={}, # mutable-ok: call payload has no schema; guardrail gets args from request_data
)
openai_tool: Final = transform_mcp_tool_to_openai_tool(mcp_tool)
fn: Final = openai_tool["function"]

18
litellm/proxy/_experimental/mcp_server/mcp_context.py
View file

@ -6,7 +6,23 @@ mcp_server_manager.py and server.py.
"""
from contextvars import ContextVar
from typing import Final
from typing import TYPE_CHECKING, Final
if TYPE_CHECKING:
from mcp.server.context import ServerRequestContext
# The SDK 1.x ``mcp.server.lowlevel.server.request_ctx`` ContextVar was removed in
# SDK 2, which hands each request handler a ``ServerRequestContext`` argument
# instead. The handlers set this var so downstream helpers (session auth caching,
# debug diagnostics, progress forwarding) can reach the same request-scoped state.
active_mcp_request_ctx_var: Final[ContextVar["ServerRequestContext | None"]] = ContextVar(
"active_mcp_request_ctx", default=None
)
def get_active_mcp_request_ctx() -> "ServerRequestContext | None":
return active_mcp_request_ctx_var.get()
# Set server-side in proxy_server.py route handlers when a request arrives via
# /toolset/{name}/mcp or the toolset fallback in dynamic_mcp_route.

45
litellm/proxy/_experimental/mcp_server/mcp_debug.py
View file

@ -100,6 +100,8 @@ Usage with curl::
http://localhost:4000/mcp/atlassian_mcp
"""
from __future__ import annotations
import asyncio
import base64
import io
@ -109,17 +111,20 @@ from collections.abc import AsyncIterator, Callable, Mapping
from http.cookies import CookieError, SimpleCookie
from itertools import islice
from types import MappingProxyType
from typing import Final
from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl, quote, quote_plus, unquote_plus, urlencode
import httpx
import httpx2
from pydantic import JsonValue, TypeAdapter
from starlette.requests import HTTPConnection
from starlette.types import Message, Send
from litellm.litellm_core_utils.secret_redaction import REDACTED, redact_string
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
if TYPE_CHECKING:
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
# Header the client sends to opt into debug mode
MCP_DEBUG_REQUEST_HEADER: Final = "x-litellm-mcp-debug"
@ -132,9 +137,9 @@ MCP_AUTH_DIAGNOSTICS_SCOPE_KEY: Final = "litellm.mcp.auth_diagnostics"
def record_auth_resolution(server_id: str, source: AuthResolution) -> None:
from mcp.server.lowlevel.server import request_ctx
from litellm.proxy._experimental.mcp_server.mcp_context import get_active_mcp_request_ctx
context: Final[object] = request_ctx.get(None)
context: Final[object] = get_active_mcp_request_ctx()
request: Final[object] = getattr(context, "request", None)
if isinstance(request, HTTPConnection):
diagnostics: Final[object] = request.scope.get(MCP_AUTH_DIAGNOSTICS_SCOPE_KEY)
@ -150,6 +155,8 @@ class MCPAuthDiagnostics:
self._outcomes = tuple(item for item in self._outcomes if item[0] != server_id) + ((server_id, resolution),)
def resolution(self) -> str:
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
match self._outcomes:
case ():
return AuthResolution.unresolved.value
@ -159,6 +166,8 @@ class MCPAuthDiagnostics:
return AuthResolution.multiple.value
def headers(self) -> Mapping[str, str]:
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
if len(self._outcomes) <= 1:
return MappingProxyType({"x-mcp-debug-auth-resolution": self.resolution()})
return MappingProxyType(
@ -372,6 +381,8 @@ class MCPDebug:
server_url: str | None = None
server_auth_type: str | None = None
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import AuthResolution
auth_resolution: Final = AuthResolution.unresolved.value
for server_name in mcp_servers or []:
@ -409,7 +420,7 @@ def _safe_text(value: str, limit: int = _BODY_PREVIEW_CHARS) -> str:
return escaped if len(escaped) <= limit else f"{escaped[:limit]}...(truncated)"
def safe_upstream_url(url: httpx.URL) -> str:
def safe_upstream_url(url: httpx.URL | httpx2.URL) -> str:
return _safe_text(str(url.copy_with(username="", password="", path="/", query=None, fragment=None)))
@ -449,10 +460,10 @@ def _header_secret_values(name: str, value: str) -> tuple[str, ...]:
return (value, credential, decoded, password, unquote_plus(password))
def _body_secret_values(request: httpx.Request) -> tuple[str, ...] | None:
def _body_secret_values(request: httpx.Request | httpx2.Request) -> tuple[str, ...] | None:
try:
raw: Final = request.content
except httpx.RequestNotRead:
except (httpx.RequestNotRead, httpx2.RequestNotRead):
return None
if not raw:
return ()
@ -478,7 +489,7 @@ def _body_secret_values(request: httpx.Request) -> tuple[str, ...] | None:
)
def _request_secret_values(request: httpx.Request) -> tuple[str, ...] | None:
def _request_secret_values(request: httpx.Request | httpx2.Request) -> tuple[str, ...] | None:
body_values: Final = _body_secret_values(request)
if body_values is None:
return None
@ -537,18 +548,18 @@ def _preview(raw: bytes, content_type: str = "", secrets: tuple[str, ...] = ())
return _safe_text(redact_string(_mask_known_values(json.dumps(parsed, separators=(",", ":")), secrets)))
def _masked_headers(headers: httpx.Headers) -> str:
def _masked_headers(headers: httpx.Headers | httpx2.Headers) -> str:
return _safe_text(", ".join(f"{name}={value}" for name, value in headers.items() if name in _SAFE_HEADER_NAMES))
def _request_body_preview(request: httpx.Request, secrets: tuple[str, ...] | None) -> str:
def _request_body_preview(request: httpx.Request | httpx2.Request, secrets: tuple[str, ...] | None) -> str:
try:
return _preview(request.content, request.headers.get("content-type", ""), secrets or ())
except httpx.RequestNotRead:
except (httpx.RequestNotRead, httpx2.RequestNotRead):
return "(streamed, not captured)"
def _response_body_preview(response: httpx.Response, secrets: tuple[str, ...] | None) -> str:
def _response_body_preview(response: httpx.Response | httpx2.Response, secrets: tuple[str, ...] | None) -> str:
if secrets is None:
return "(omitted: request credentials unavailable)"
captured: Final = response.extensions.get(_CAPTURE_EXTENSION)
@ -556,7 +567,7 @@ def _response_body_preview(response: httpx.Response, secrets: tuple[str, ...] |
return captured
try:
return _preview(response.content, response.headers.get("content-type", ""), secrets)
except httpx.ResponseNotRead:
except (httpx.ResponseNotRead, httpx2.ResponseNotRead):
return "(not read)"
@ -569,7 +580,7 @@ async def _read_error_prefix(chunks: AsyncIterator[bytes], limit: int) -> bytes:
return buffer.getvalue()
async def capture_upstream_error_response(response: httpx.Response) -> None:
async def capture_upstream_error_response(response: httpx.Response | httpx2.Response) -> None:
if not response.is_error:
return
try:
@ -584,7 +595,7 @@ async def capture_upstream_error_response(response: httpx.Response) -> None:
if secrets is not None
else "(omitted: request credentials unavailable)"
)
except (asyncio.TimeoutError, httpx.HTTPError, httpx.StreamError):
except (asyncio.TimeoutError, httpx.HTTPError, httpx.StreamError, httpx2.HTTPError, httpx2.StreamError):
response._content = b"" # pyright: ignore[reportPrivateUsage] # rebind-ok: httpx auth retries must survive diagnostic read failures
response.extensions[_CAPTURE_EXTENSION] = (
"(unavailable: error body read failed)" # rebind-ok: httpx response hooks communicate through extensions
@ -593,7 +604,7 @@ async def capture_upstream_error_response(response: httpx.Response) -> None:
response.extensions[_CAPTURE_EXTENSION] = preview # rebind-ok: httpx response hooks communicate through extensions
def describe_upstream_response(response: httpx.Response) -> str:
def describe_upstream_response(response: httpx.Response | httpx2.Response) -> str:
try:
request: Final = response.request
except RuntimeError:
@ -616,6 +627,6 @@ def describe_upstream_http_failure(exc: BaseException) -> str | None:
describe_upstream_response(response)
for current in islice(iter_exception_tree(exc), 16)
for response in (getattr(current, "response", None),)
if isinstance(response, httpx.Response)
if isinstance(response, (httpx.Response, httpx2.Response))
)
return " | ".join(lines) or None

20
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -34,6 +34,7 @@ from urllib.parse import ParseResult, urlparse
import anyio
import httpx
import httpx2
from fastapi import HTTPException
from httpx import HTTPStatusError
from mcp import ReadResourceResult, Resource
@ -194,8 +195,7 @@ from litellm.types.mcp_server.mcp_server_manager import (
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from mcp.client.session import ClientSession
from mcp.shared.context import RequestContext
from mcp.client.session import ClientRequestContext
from mcp.types import CreateMessageRequestParams
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -1297,8 +1297,8 @@ def _passthrough_token_from_mcp_auth_header(
return None
async def _materialize_auth_headers(auth: httpx.Auth | None) -> dict[str, str] | None:
"""Extract the header a resolved ``httpx.Auth`` would set, as a plain dict, or None.
async def _materialize_auth_headers(auth: httpx2.Auth | None) -> dict[str, str] | None:
"""Extract the header a resolved ``httpx2.Auth`` would set, as a plain dict, or None.
OpenAPI tool closures egress through ``AsyncHTTPHandler`` methods that accept headers but no
``auth``, so a resolved credential must be materialized into a header value. Driving one step
@ -1313,7 +1313,7 @@ async def _materialize_auth_headers(auth: httpx.Auth | None) -> dict[str, str] |
header_name: Final = getattr(auth, "header_name", None)
if not isinstance(header_name, str) or not header_name:
return None
probe: Final = httpx.Request("GET", "http://localhost/")
probe: Final = httpx2.Request("GET", "http://localhost/")
flow: Final = auth.async_auth_flow(probe)
try:
first_request: Final = await flow.__anext__()
@ -1587,7 +1587,7 @@ def _create_sampling_callback(user_api_key_auth: UserAPIKeyAuth | None = None):
return None
async def _sampling_callback(
context: "RequestContext[ClientSession, object]",
context: "ClientRequestContext",
params: "CreateMessageRequestParams",
):
import litellm
@ -4012,7 +4012,7 @@ class MCPServerManager:
subject_token: str | None,
user_api_key_auth: UserAPIKeyAuth | None,
extra_headers: dict[str, str] | None,
) -> tuple[httpx.Auth | None, dict[str, str] | None]:
) -> tuple[httpx2.Auth | None, dict[str, str] | None]:
"""Resolve a v2-owned server's upstream credential into ``(resolved_auth, extra_headers)``.
On a missing/rejected per-user credential this raises the mode's discovery challenge
@ -5552,7 +5552,7 @@ class MCPServerManager:
verbose_logger.error(error_msg)
return CallToolResult(
content=[TextContent(type="text", text=error_msg)],
isError=True,
is_error=True,
)
try:
@ -5563,7 +5563,7 @@ class MCPServerManager:
# Convert the handler result (string response) to CallToolResult format
result: Final = CallToolResult(
content=[TextContent(type="text", text=str(handler_result))],
isError=False,
is_error=False,
)
return result
@ -5579,7 +5579,7 @@ class MCPServerManager:
verbose_logger.error(error_msg)
return CallToolResult(
content=[TextContent(type="text", text=error_msg)],
isError=True,
is_error=True,
)
async def pre_call_tool_check(

9
litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py
View file

@ -34,6 +34,7 @@ from dataclasses import dataclass
from typing import Annotated, Final, Literal
import httpx
import httpx2
from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError
from typing_extensions import assert_never
@ -337,7 +338,7 @@ def _identity_key(config: ClientCredentialsConfig) -> str:
return hashlib.sha256(material.encode("utf-8")).hexdigest()
class ClientCredentialsBearerAuth(httpx.Auth):
class ClientCredentialsBearerAuth(httpx2.Auth):
"""Bearer auth that retries an upstream 401 exactly once with a freshly minted token.
The initial token was already resolved (so config/IdP failures surfaced as typed errors
@ -356,7 +357,7 @@ class ClientCredentialsBearerAuth(httpx.Auth):
self._access_token = SecretStr(access_token)
self._refetch = refetch
async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
token: Final = self._access_token.get_secret_value()
name, value = self._carrier.header(token)
request.headers[name] = value
@ -371,5 +372,5 @@ class ClientCredentialsBearerAuth(httpx.Auth):
request.headers[fresh_name] = fresh_value
yield request
def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
raise RuntimeError("ClientCredentialsBearerAuth only supports async httpx clients")
def sync_auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
raise RuntimeError("ClientCredentialsBearerAuth only supports async httpx2 clients")

18
litellm/proxy/_experimental/mcp_server/outbound_credentials/httpx_auth.py
View file

@ -1,29 +1,29 @@
"""Concrete `httpx.Auth` objects the resolver returns for the self-contained modes.
"""Concrete `httpx2.Auth` objects the resolver returns for the self-contained modes.
These are the egress credential as the SDK consumes it: an `httpx.Auth` attached to the
These are the egress credential as the SDK consumes it: an `httpx2.Auth` attached to the
upstream `AsyncClient`. The OAuth-flow modes (`authorization_code`, `client_credentials`,
`token_exchange`) return SDK-provided auth objects instead and land later.
`auth_flow` mutating the outbound request is the `httpx.Auth` contract, not a house-style
violation: the request is httpx's object, and these carry no state of their own.
`auth_flow` mutating the outbound request is the `httpx2.Auth` contract, not a house-style
violation: the request is httpx2's object, and these carry no state of their own.
"""
from __future__ import annotations
from collections.abc import Generator
import httpx
import httpx2
from pydantic import SecretStr
class NoOpAuth(httpx.Auth):
class NoOpAuth(httpx2.Auth):
"""Attaches nothing — the `none` mode (and the seam-level default)."""
def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
yield request
class StaticHeaderAuth(httpx.Auth):
class StaticHeaderAuth(httpx2.Auth):
"""Sets one fixed header on every request — the `api_key` family and `passthrough`.
The header value is a live credential (a bearer token, an API key, a forwarded user
@ -36,6 +36,6 @@ class StaticHeaderAuth(httpx.Auth):
self.header_name = header_name
self._header_value = SecretStr(header_value)
def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]:
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
request.headers[self.header_name] = self._header_value.get_secret_value()
yield request

23
litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py
View file

@ -1,7 +1,7 @@
"""The one credential resolver: dispatch on the declared mode, fail closed.
`resolve_credentials` selects exactly one arm off the server's typed `config` and either
produces an `httpx.Auth` or returns a typed `CredError`. The `match` is over the `AuthConfig`
produces an `httpx2.Auth` or returns a typed `CredError`. The `match` is over the `AuthConfig`
variant, so each arm receives its own fully-typed config with no field-presence inference and
no precedence cascade. It is wildcard-free with an `assert_never` tail, so adding a mode without
an arm fails the type gate (basedpyright `reportMatchNotExhaustive`); a bypassed gate fails loudly
@ -25,6 +25,7 @@ from functools import partial
from typing import Final
import httpx
import httpx2
from typing_extensions import assert_never
from litellm._logging import verbose_proxy_logger
@ -135,7 +136,7 @@ class UpstreamCredentialProvider:
self._client_credentials_source = client_credentials_source or ClientCredentialsTokenSource()
self._sso_assertion_store: SSOAssertionStore = sso_assertion_store or default_sso_assertion_store()
async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]:
async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx2.Auth, CredError]:
match server.config:
case NoneConfig():
return self._none(server)
@ -155,7 +156,7 @@ class UpstreamCredentialProvider:
return _not_implemented(AuthSpecKind.aws_sigv4)
assert_never(server.config)
def _none(self, server: ServerSpec) -> Result[httpx.Auth, CredError]:
def _none(self, server: ServerSpec) -> Result[httpx2.Auth, CredError]:
try:
resource: Final = httpx.URL(server.resource)
except httpx.InvalidURL:
@ -169,12 +170,12 @@ class UpstreamCredentialProvider:
Reads from the same per-user store as the ``authorization_code`` arm, so the discovery
challenge and the egress agree on whether the user is authorized. Returns a typed ``bool``
(no ``httpx.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the
(no ``httpx2.Auth``), unlike ``resolve_credentials``. A non-per-user mode has no token in the
store, so it reads as False without a per-mode branch here.
"""
return await self._authz_token(subject, server) is not None
def _passthrough(self, subject: Subject) -> Result[httpx.Auth, CredError]:
def _passthrough(self, subject: Subject) -> Result[httpx2.Auth, CredError]:
"""Forward the caller's own upstream credential verbatim; the gateway mints nothing.
The inbound token is the caller's already-disambiguated ``Authorization`` (never the LiteLLM
@ -186,7 +187,7 @@ class UpstreamCredentialProvider:
return Ok(NoOpAuth())
return Ok(StaticHeaderAuth(subject.inbound_token.get_secret_value(), header_name="Authorization"))
def _api_key(self, config: ApiKeyConfig) -> Result[httpx.Auth, CredError]:
def _api_key(self, config: ApiKeyConfig) -> Result[httpx2.Auth, CredError]:
match config.key_source:
case SharedKey() as source:
header_name, header_value = config.header(source.value.get_secret_value())
@ -196,7 +197,9 @@ class UpstreamCredentialProvider:
return Error(CredError.of_not_implemented("api_key BYOK source not implemented yet"))
assert_never(config.key_source)
async def _id_jag(self, subject: Subject, server: ServerSpec, config: IdJagConfig) -> Result[httpx.Auth, CredError]:
async def _id_jag(
self, subject: Subject, server: ServerSpec, config: IdJagConfig
) -> Result[httpx2.Auth, CredError]:
match await self._id_jag_subject_token(subject):
case Error(err):
return Error(err)
@ -261,7 +264,7 @@ class UpstreamCredentialProvider:
async def _id_jag_exchange(
self, subject: Subject, token: str, server: ServerSpec, config: IdJagConfig
) -> Result[httpx.Auth, CredError]:
) -> Result[httpx2.Auth, CredError]:
slot: Final = _id_jag_slot_key(subject, server)
fingerprint: Final = _id_jag_fingerprint(token, server.server_id, config)
@ -313,7 +316,7 @@ class UpstreamCredentialProvider:
async def _client_credentials(
self, server_id: str, config: ClientCredentialsConfig
) -> Result[httpx.Auth, CredError]:
) -> Result[httpx2.Auth, CredError]:
"""The M2M arm: resolve a cached (or freshly minted) gateway token; no user context.
The token is resolved here, before any upstream request, so a misconfigured grant or an
@ -448,7 +451,7 @@ def _client_auth_fingerprint(client_auth: ClientAuth) -> str:
assert_never(client_auth)
def _not_implemented(kind: AuthSpecKind) -> Result[httpx.Auth, CredError]:
def _not_implemented(kind: AuthSpecKind) -> Result[httpx2.Auth, CredError]:
return Error(CredError.of_not_implemented(f"{kind.value}: resolver arm not implemented yet"))

6
litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py
View file

@ -30,7 +30,7 @@ from dataclasses import dataclass, field
from enum import Enum
from typing import Annotated, Final, Literal
import httpx
import httpx2
from expression import case, tag, tagged_union
from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator
from typing_extensions import assert_never
@ -66,7 +66,7 @@ class AuthResolution(str, Enum):
@dataclass(frozen=True, slots=True)
class ResolvedCredential:
auth: httpx.Auth = field(repr=False)
auth: httpx2.Auth = field(repr=False)
source: AuthResolution
@ -110,7 +110,7 @@ class Unauthorized:
@tagged_union(frozen=True)
class CredError:
"""Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx.Auth`.
"""Why a credential could not be produced. Fail-closed: an arm yields this or an `httpx2.Auth`.
Discriminated on the `Literal` `tag`; consumers `match self.tag` (see `summary`) so the
type checker can prove exhaustiveness. Construct via the `of_*` factories.

41
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -10,6 +10,7 @@ from uuid import uuid4
import anyio
import httpx
import httpx2
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from pydantic import ValidationError
from starlette.datastructures import Headers
@ -120,20 +121,29 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout
f"within {timeout_seconds:.0f}s. Check that the LiteLLM proxy can reach this URL "
"from its network (DNS, egress rules, firewalls) and that the server answers MCP requests."
)
if isinstance(exc, httpx.LocalProtocolError):
if isinstance(exc, (httpx.LocalProtocolError, httpx2.LocalProtocolError)):
return (
"Failed to connect to MCP server: a request header is malformed. "
"Check static headers for leading/trailing spaces or illegal characters."
)
if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout)):
if isinstance(exc, (httpx.ConnectError, httpx.ConnectTimeout, httpx2.ConnectError, httpx2.ConnectTimeout)):
return (
"Failed to connect to MCP server: the server is unreachable. Check the URL and that the server is running."
)
if isinstance(exc, httpx.TimeoutException):
if isinstance(exc, (httpx.TimeoutException, httpx2.TimeoutException)):
return "Failed to connect to MCP server: the connection timed out."
if isinstance(exc, httpx.HTTPStatusError):
if isinstance(exc, (httpx.HTTPStatusError, httpx2.HTTPStatusError)):
return f"Failed to connect to MCP server: it returned HTTP {exc.response.status_code}."
if isinstance(exc, (httpx.NetworkError, httpx.RemoteProtocolError, ConnectionError)):
if isinstance(
exc,
(
httpx.NetworkError,
httpx.RemoteProtocolError,
httpx2.NetworkError,
httpx2.RemoteProtocolError,
ConnectionError,
),
):
return (
"Failed to connect to MCP server: the connection was interrupted. "
"Check the server and network connection, then retry."
@ -148,7 +158,18 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout
"Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response. "
"Check the MCP endpoint URL and the server's protocol implementation."
)
if MCP_AVAILABLE and isinstance(exc, McpError):
if MCP_AVAILABLE and isinstance(exc, MCPError):
if exc.error.message.startswith("Unexpected content type:"):
return (
"Failed to connect to MCP server: the endpoint returned an unsupported content type. "
"Check that the URL is an MCP endpoint, not a web page, and matches the selected transport."
)
if exc.error.code == -32700 or exc.error.message.startswith("Failed to parse"):
return (
f"Failed to connect to MCP server: the endpoint returned invalid JSON or an invalid MCP response "
f"(JSON-RPC code {exc.error.code}). "
"Check the MCP endpoint URL and the server's protocol implementation."
)
if exc.error.code == -32000 and exc.error.message == "Connection closed":
return (
"Failed to connect to MCP server: the connection was closed before the request completed. "
@ -168,7 +189,7 @@ def _known_connection_error_message(exc: BaseException, url: str | None, timeout
if MCP_AVAILABLE:
from mcp.shared.exceptions import McpError
from mcp.shared.exceptions import MCPError
from mcp.types import Tool as MCPTool
from litellm.experimental_mcp_client.client import MCPClient, as_mcp_read_timeout
@ -518,7 +539,7 @@ if MCP_AVAILABLE:
ListMCPToolsRestAPIResponseObject(
name=tool.name,
description=tool.description,
inputSchema=tool.inputSchema,
inputSchema=tool.input_schema,
mcp_info=enriched_mcp_info,
)
for tool in tools
@ -1484,7 +1505,7 @@ if MCP_AVAILABLE:
effective_timeout: Final = (
min(request.timeout if request.timeout is not None else MCP_CLIENT_TIMEOUT, timeout_seconds)
if any(
isinstance(cause, McpError) and as_mcp_read_timeout(cause) is not None
isinstance(cause, MCPError) and as_mcp_read_timeout(cause) is not None
for cause in iter_exception_tree(e)
)
else timeout_seconds
@ -1635,7 +1656,7 @@ if MCP_AVAILABLE:
"message": f"Timed out listing tools after {listing_deadline} seconds. "
"The MCP server may be responding slowly or paginating excessively.",
}
model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result]
model_dumped_tools: Final[list[dict]] = [tool.model_dump(by_alias=True) for tool in list_tools_result]
return {
"tools": model_dumped_tools,
"error": None,

35
litellm/proxy/_experimental/mcp_server/sampling_handler.py
View file

@ -18,8 +18,7 @@ if typing.TYPE_CHECKING:
from collections.abc import Awaitable, Callable
from fastapi import Request
from mcp.client.session import ClientSession
from mcp.shared.context import RequestContext
from mcp.client.session import ClientRequestContext
from mcp.types import (
ContentBlock,
CreateMessageResult,
@ -333,14 +332,14 @@ def _convert_single_content(
return {"type": "text", "text": content.text}
elif content_type == "image":
image_data: Final[str] = getattr(content, "data", "")
image_mime_type: Final[str] = getattr(content, "mimeType", "image/png")
image_mime_type: Final[str] = getattr(content, "mime_type", "image/png")
return {
"type": "image_url",
"image_url": {"url": f"data:{image_mime_type};base64,{image_data}"},
}
elif content_type == "audio":
audio_data: Final[str] = getattr(content, "data", "")
audio_mime_type: Final[str] = getattr(content, "mimeType", "audio/wav")
audio_mime_type: Final[str] = getattr(content, "mime_type", "audio/wav")
# Map MIME type to OpenAI audio format
format_map: Final = {
"audio/wav": "wav",
@ -375,7 +374,7 @@ def _convert_single_content(
# ToolResultContent → proper OpenAI tool-role message.
# Marked so the message-level converter can emit it as a
# separate ``{"role": "tool", ...}`` message.
tool_result_use_id: Final = getattr(content, "toolUseId", "")
tool_result_use_id: Final = getattr(content, "tool_use_id", "")
nested_content: Final[Sequence[ContentBlock]] = getattr(content, "content", [])
if isinstance(nested_content, list):
text_parts = [getattr(c, "text", str(c)) for c in nested_content if getattr(c, "type", None) == "text"]
@ -538,7 +537,7 @@ def _extract_tool_results(
results: Final = []
for item in items:
if getattr(item, "type", None) == "tool_result":
tool_use_id = getattr(item, "toolUseId", "")
tool_use_id = getattr(item, "tool_use_id", "")
# Extract text from nested content
nested_content: Sequence[ContentBlock] = getattr(item, "content", [])
if isinstance(nested_content, list):
@ -573,7 +572,7 @@ def _convert_mcp_tools_to_openai(
"function": {
"name": tool.name,
"description": tool.description or "",
"parameters": tool.inputSchema
"parameters": tool.input_schema
or {
"type": "object",
"properties": {},
@ -718,7 +717,7 @@ def _convert_openai_response_to_mcp_result(
role="assistant",
content=content_parts,
model=actual_model,
stopReason=stop_reason,
stop_reason=stop_reason,
)
# Simple text response
text: Final = message.content or ""
@ -726,7 +725,7 @@ def _convert_openai_response_to_mcp_result(
role="assistant",
content=TextContent(type="text", text=text),
model=actual_model,
stopReason=stop_reason,
stop_reason=stop_reason,
)
@ -1066,21 +1065,21 @@ async def _build_completion_kwargs(
) -> dict[str, Any]:
openai_messages: Final = _convert_mcp_messages_to_openai(
messages=params.messages,
system_prompt=params.systemPrompt,
system_prompt=params.system_prompt,
)
completion_kwargs: Final[dict[str, object]] = {
"model": model,
"messages": openai_messages,
"max_tokens": params.maxTokens,
"max_tokens": params.max_tokens,
}
if params.temperature is not None:
completion_kwargs["temperature"] = params.temperature
if params.stopSequences:
completion_kwargs["stop"] = params.stopSequences
if params.stop_sequences:
completion_kwargs["stop"] = params.stop_sequences
openai_tools: Final = _convert_mcp_tools_to_openai(params.tools)
if openai_tools:
completion_kwargs["tools"] = openai_tools
openai_tool_choice: Final = _convert_mcp_tool_choice_to_openai(params.toolChoice)
openai_tool_choice: Final = _convert_mcp_tool_choice_to_openai(params.tool_choice)
if openai_tool_choice is not None:
completion_kwargs["tool_choice"] = openai_tool_choice
completion_kwargs["metadata"] = {"mcp_metadata": params.metadata} if params.metadata else {}
@ -1137,7 +1136,7 @@ async def _run_guardrails_and_call_llm(
async def handle_sampling_create_message(
context: "RequestContext[ClientSession, object]",
context: "ClientRequestContext",
params: "CreateMessageRequestParams",
default_model: str | None = None,
user_api_key_auth: "UserAPIKeyAuth | None" = None,
@ -1180,13 +1179,13 @@ async def handle_sampling_create_message(
try:
model: Final = _resolve_model_from_preferences(
model_preferences=params.modelPreferences,
model_preferences=params.model_preferences,
default_model=default_model,
)
verbose_logger.info(
"MCP sampling: resolved model=%s from preferences=%s",
model,
params.modelPreferences,
params.model_preferences,
)
access_denial: Final = await _check_model_access(model, user_api_key_auth)
@ -1228,7 +1227,7 @@ async def handle_sampling_create_message(
verbose_logger.info(
"MCP sampling: completed successfully, model=%s, stopReason=%s",
getattr(result, "model", "unknown"),
getattr(result, "stopReason", "unknown"),
getattr(result, "stop_reason", "unknown"),
)
return result
except Exception as e:

323
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -15,13 +15,13 @@ import traceback
import types
import uuid
from collections import Counter
from collections.abc import AsyncIterator, Callable, Mapping, Sequence
from collections.abc import AsyncIterator, Callable, Iterable, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol
import httpx
from fastapi import FastAPI, HTTPException
from pydantic import AnyUrl, ConfigDict, TypeAdapter, ValidationError
from pydantic import AnyUrl, ConfigDict, Field, TypeAdapter, ValidationError
from starlette.requests import Request as StarletteRequest
from starlette.responses import JSONResponse
from starlette.types import Message, Receive, Scope, Send
@ -64,6 +64,8 @@ from litellm.proxy._experimental.mcp_server.mcp_context import (
_mcp_gateway_initialize_instructions,
_mcp_gateway_server_name,
_mcp_proxy_mode, # pyright: ignore[reportPrivateUsage] # server-owned request mode
active_mcp_request_ctx_var,
get_active_mcp_request_ctx,
)
from litellm.proxy._experimental.mcp_server.mcp_debug import (
MCP_AUTH_DIAGNOSTICS_SCOPE_KEY,
@ -137,6 +139,24 @@ _MCP_ROUTING_PEEK_MAX_BYTES: Final = 4096
# ASGI scope keys carrying OTel request state into a stateful MCP message handler.
_MCP_TRANSPORT_SPAN_SCOPE_KEY: Final = "litellm_otel_transport_span"
_MCP_DESTINATIONS_SCOPE_KEY: Final = "litellm_otel_request_destinations"
_MCP_PROTOCOL_VERSION_HEADER: Final = b"mcp-protocol-version"
def unsupported_protocol_version(scope: Scope) -> str | None:
"""Return the unsupported ``MCP-Protocol-Version`` header value, if any.
SDK 2's ``StreamableHTTPSessionManager`` routes any version outside
``HANDSHAKE_PROTOCOL_VERSIONS`` to the modern single-exchange path, which
bypasses litellm's session/auth model, so the ASGI entry rejects it.
"""
headers: Final[Iterable[tuple[bytes, bytes]]] = scope.get("headers") or ()
values: Final = tuple(
raw.decode("latin-1").strip() for key, raw in headers if key.lower() == _MCP_PROTOCOL_VERSION_HEADER
)
for value in values:
if value and value not in HANDSHAKE_PROTOCOL_VERSIONS:
return value
return None
async def _invalidate_byok_cred_cache(user_id: str, server_id: str) -> None:
@ -156,14 +176,12 @@ try:
from mcp import ReadResourceResult, Resource
from mcp.server import Server
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.server.session import ServerSession as _McpServerSession
from mcp.types import (
BlobResourceContents,
GetPromptResult,
ResourceTemplate,
TextResourceContents,
Tool,
)
# Robust auth lookup keyed by session_object.
@ -176,7 +194,6 @@ except ImportError as e:
# so they will never be accessed at runtime
BlobResourceContents = None
GetPromptResult = None
ReadResourceContents = None
ReadResourceResult = None
Resource = None
ResourceTemplate = None
@ -277,8 +294,8 @@ def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None:
span's identity attribution.
"""
meta: Final = getattr(req_ctx, "meta", None)
extra: Final = getattr(meta, "model_extra", None)
if not isinstance(extra, dict):
extra: Final = meta if isinstance(meta, Mapping) else getattr(meta, "model_extra", None)
if not isinstance(extra, Mapping):
return None
carrier: Final = {key: extra[key] for key in ("traceparent", "tracestate") if isinstance(extra.get(key), str)}
return carrier or None
@ -456,6 +473,7 @@ if MCP_AVAILABLE:
AuthContextMiddleware,
auth_context_var,
)
from mcp.server.context import ServerRequestContext
from mcp.server.lowlevel.server import NotificationOptions
from mcp.server.models import InitializationOptions
@ -464,14 +482,23 @@ if MCP_AVAILABLE:
except ImportError:
StreamableHTTPSessionManager = None
from mcp.types import (
INVALID_REQUEST,
CallToolRequestParams,
CallToolResult,
GetPromptRequestParams,
Implementation,
InitializeRequest,
ListPromptsResult,
ListResourcesResult,
ListResourceTemplatesResult,
ListToolsResult,
PaginatedRequestParams,
Prompt,
ReadResourceRequestParams,
TextContent,
)
from mcp.types import Tool as MCPTool
from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS
from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import (
MCPAuthenticatedUser,
@ -520,46 +547,20 @@ if MCP_AVAILABLE:
Object returned by the /tools/list REST API route.
"""
mcp_info: MCPInfo | None = None
mcp_info: MCPInfo | None = Field(default=None, alias="mcp_info")
model_config = ConfigDict(arbitrary_types_allowed=True)
def _normalize_resource_contents(contents: list) -> list[ReadResourceContents]:
"""Normalize ResourceContents to ReadResourceContents, preserving meta (MCP 1.26.0+)."""
normalized: Final[list[ReadResourceContents]] = []
for content in contents:
meta = getattr(content, "meta", None)
if meta is None and hasattr(content, "model_dump"):
d = content.model_dump()
meta = d.get("meta")
if meta is None:
meta = d.get("_meta")
if isinstance(content, TextResourceContents):
normalized.append(
ReadResourceContents(
content=content.text,
mime_type=content.mimeType,
meta=meta,
)
)
elif isinstance(content, BlobResourceContents):
normalized.append(
ReadResourceContents(
content=content.blob,
mime_type=content.mimeType,
meta=meta,
)
)
return normalized
def _gateway_create_initialization_options(
self,
notification_options: NotificationOptions | None = None,
experimental_capabilities: dict[str, dict[str, object]] | None = None,
extensions: dict[str, dict[str, object]] | None = None,
) -> InitializationOptions:
base_options: Final = Server.create_initialization_options(
self,
notification_options=notification_options,
experimental_capabilities=experimental_capabilities or {},
extensions=extensions,
)
opts: Final = (
base_options.model_copy(
@ -817,8 +818,7 @@ if MCP_AVAILABLE:
############### MCP Server Routes #######################
########################################################
@server.list_tools()
async def handle_list_tools() -> "ListToolsResult | list[Tool]":
async def handle_list_tools(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListToolsResult:
"""
List all available tools, with each server's listing outcome attached to the result's
``_meta`` (SERVER_OUTCOMES_META_KEY) so a broken upstream is distinguishable from a healthy
@ -826,12 +826,9 @@ if MCP_AVAILABLE:
pass the result through unwrapped, which is what lets the ``_meta`` survive to the client.
Also captures the active session for propagation to callbacks.
"""
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
_session_reset_token = None
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
req_ctx: Final = ctx
_ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx)
_session_reset_token: Final = active_mcp_session_var.set(ctx.session)
_trace_token = None
_transport_token = None
_destinations_token = None
@ -864,13 +861,13 @@ if MCP_AVAILABLE:
)
if _mcp_proxy_mode.get():
return [Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()] # mutable-ok: MCP SDK list
return ListToolsResult(tools=[Tool.model_validate(d) for d in get_mcp_proxy_tool_definitions()])
if getattr(
getattr(user_api_key_auth, "object_permission", None),
"mcp_tool_search_enabled",
False,
):
return [Tool.model_validate(d) for d in get_virtual_tool_definitions()]
return ListToolsResult(tools=[Tool.model_validate(d) for d in get_virtual_tool_definitions()])
# Get mcp_servers from context variable
verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools")
@ -886,7 +883,7 @@ if MCP_AVAILABLE:
)
verbose_logger.info("MCP list_tools - Successfully returned %s tools", len(listing.tools))
if not listing.outcomes:
return listing.tools
return ListToolsResult(tools=listing.tools)
outcome_meta: Final = {
SERVER_OUTCOMES_META_KEY: {
key: outcome_wire_value(outcome) for key, outcome in listing.outcomes.items()
@ -894,36 +891,32 @@ if MCP_AVAILABLE:
}
return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta})
except HTTPException as e:
from mcp.shared.exceptions import McpError
from mcp.types import INVALID_REQUEST, ErrorData
from mcp.shared.exceptions import MCPError
from mcp.types import INVALID_REQUEST
raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e
raise MCPError(code=INVALID_REQUEST, message=_http_detail_message(e.detail)) from e
except Exception as e:
verbose_logger.exception("Error in list_tools endpoint: %s", e)
# Return empty list instead of failing completely
# This prevents the HTTP stream from failing and allows the client to get a response
return []
return ListToolsResult(tools=[]) # mutable-ok: MCP result payload
finally:
_otel_reset_mcp_request_destinations(_destinations_token)
_otel_reset_mcp_transport_span(_transport_token)
_otel_reset_mcp_trace_carrier(_trace_token)
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
active_mcp_session_var.reset(_session_reset_token)
active_mcp_request_ctx_var.reset(_ctx_reset_token)
def _capture_host_progress_callback(host_server) -> Callable | None:
def _capture_host_progress_callback(ctx: ServerRequestContext) -> Callable | None:
"""Return a progress-forwarding callback bound to the host MCP session.
Returns ``None`` when the host did not supply a progress token.
"""
try:
host_ctx: Final = host_server.request_context
except Exception as e:
verbose_logger.warning("Could not capture host progress context: %s", e)
return None
host_ctx: Final = ctx
if not (host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta):
return None
host_token: Final = getattr(host_ctx.meta, "progressToken", None)
host_token: Final = host_ctx.meta.get("progress_token")
if host_token is None or not (hasattr(host_ctx, "session") and host_ctx.session):
return None
host_session: Final = host_ctx.session
@ -944,10 +937,10 @@ if MCP_AVAILABLE:
return forward_progress
def _reject_mcp_proxy_operation() -> NoReturn:
from mcp.shared.exceptions import McpError
from mcp.types import METHOD_NOT_FOUND, ErrorData
from mcp.shared.exceptions import MCPError
from mcp.types import METHOD_NOT_FOUND
raise McpError(ErrorData(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy"))
raise MCPError(code=METHOD_NOT_FOUND, message="Operation unavailable on /mcp/proxy")
async def _build_virtual_call_logging_obj(
name: str,
@ -1022,7 +1015,7 @@ if MCP_AVAILABLE:
content=[ # mutable-ok: MCP result content
TextContent(type="text", text=f"Tool {name} is unavailable on /mcp/proxy")
],
isError=True,
is_error=True,
)
if _mcp_proxy_mode.get() and name in MCP_PROXY_TOOL_NAMES:
@ -1104,7 +1097,7 @@ if MCP_AVAILABLE:
text=f"Tool {name} requires mcp_tool_search_enabled on the key",
)
],
isError=True,
is_error=True,
)
args: Final = arguments or {}
@ -1154,29 +1147,24 @@ if MCP_AVAILABLE:
litellm_logging_obj=virtual_logging_obj,
)
@server.call_tool()
async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult:
async def mcp_server_tool_call(ctx: ServerRequestContext, params: CallToolRequestParams) -> CallToolResult:
"""
Call a specific tool with the provided arguments
Args:
name (str): Name of the tool to call
arguments (Dict[str, Any] | None): Arguments to pass to the tool
ctx: SDK request context carrying the client session and HTTP request
params (CallToolRequestParams): Tool name and arguments
Returns:
List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]: Tool execution results
Raises:
HTTPException: If tool not found or arguments missing
CallToolResult: Tool execution results
"""
from mcp.server.lowlevel.server import request_ctx
from mcp.types import CallToolResult
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.proxy.proxy_server import proxy_config
req_ctx: Final = request_ctx.get(None)
_session_reset_token = None
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
req_ctx: Final = ctx
_ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx)
_session_reset_token: Final = active_mcp_session_var.set(ctx.session)
_trace_token = None
_transport_token = None
_destinations_token = None
@ -1207,8 +1195,8 @@ if MCP_AVAILABLE:
# Inside this try so virtual-tool errors convert to isError
# CallToolResult instead of raising out of the protocol handler.
virtual_tool_result: Final = await _dispatch_virtual_mcp_tool(
name=name,
arguments=arguments,
name=params.name,
arguments=params.arguments,
user_api_key_auth=user_api_key_auth,
client_ip=_client_ip,
mcp_servers=mcp_servers,
@ -1220,9 +1208,9 @@ if MCP_AVAILABLE:
if virtual_tool_result is not None:
return virtual_tool_result
host_progress_callback: Final = _capture_host_progress_callback(server)
host_progress_callback: Final = _capture_host_progress_callback(ctx)
# Create a body date for logging
body_data: Final = {"name": name, "arguments": arguments}
body_data: Final = {"name": params.name, "arguments": params.arguments} # mutable-ok: logging payload
# Set trace/session id from raw_headers so spend logs and logging_obj stay consistent (same as A2A)
chain_id: Final = get_chain_id_from_headers(raw_headers)
if chain_id:
@ -1247,7 +1235,7 @@ if MCP_AVAILABLE:
# Authorization is unaffected: it ran before this, and the union is resolved
# from the untouched auth object passed to call_mcp_tool below.
user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call(
user_api_key_auth, tool_name=name
user_api_key_auth, tool_name=params.name
),
proxy_config=proxy_config,
)
@ -1273,7 +1261,7 @@ if MCP_AVAILABLE:
)
return CallToolResult(
content=[TextContent(text=str(e), type="text")],
isError=True,
is_error=True,
)
except BlockedPiiEntityError as e:
verbose_logger.error("BlockedPiiEntityError in MCP tool call: %s", e)
@ -1284,19 +1272,19 @@ if MCP_AVAILABLE:
type="text",
)
],
isError=True,
is_error=True,
)
except GuardrailRaisedException as e:
verbose_logger.error("GuardrailRaisedException in MCP tool call: %s", e)
return CallToolResult(
content=[TextContent(text=f"Error: Guardrail violation - {e}", type="text")],
isError=True,
is_error=True,
)
except HTTPException as e:
verbose_logger.error("HTTPException in MCP tool call: %s", e)
return CallToolResult(
content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")],
isError=True,
is_error=True,
)
except MCPUpstreamAuthError as e:
# The MCP session manager serializes handler exceptions as JSON-RPC errors, so a
@ -1312,13 +1300,13 @@ if MCP_AVAILABLE:
type="text",
)
],
isError=True,
is_error=True,
)
except Exception as e:
verbose_logger.exception("MCP mcp_server_tool_call - error: %s", e)
return CallToolResult(
content=[TextContent(text=f"Error: {e}", type="text")],
isError=True,
is_error=True,
)
return response
@ -1326,22 +1314,17 @@ if MCP_AVAILABLE:
_otel_reset_mcp_request_destinations(_destinations_token)
_otel_reset_mcp_transport_span(_transport_token)
_otel_reset_mcp_trace_carrier(_trace_token)
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
active_mcp_session_var.reset(_session_reset_token)
active_mcp_request_ctx_var.reset(_ctx_reset_token)
@server.list_prompts()
async def list_prompts() -> list[Prompt]:
async def list_prompts(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListPromptsResult:
"""
List all available prompts
"""
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
_session_reset_token = None
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
_ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx)
_session_reset_token: Final = active_mcp_session_var.set(ctx.session)
try:
# Get user authentication from context variable
@ -1371,36 +1354,24 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
verbose_logger.info("MCP list_prompts - Successfully returned %s prompts", len(prompts))
return prompts
return ListPromptsResult(prompts=prompts)
except Exception as e:
verbose_logger.exception("Error in list_prompts endpoint: %s", e)
# Return empty list instead of failing completely
# This prevents the HTTP stream from failing and allows the client to get a response
return []
return ListPromptsResult(prompts=[]) # mutable-ok: MCP result payload
finally:
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
active_mcp_session_var.reset(_session_reset_token)
active_mcp_request_ctx_var.reset(_ctx_reset_token)
@server.get_prompt()
async def get_prompt(name: str, arguments: dict[str, str] | None) -> GetPromptResult:
async def get_prompt(ctx: ServerRequestContext, params: GetPromptRequestParams) -> GetPromptResult:
"""
Get a specific prompt with the provided arguments
Args:
name (str): Name of the prompt to get
arguments (Dict[str, Any] | None): Arguments to pass to the prompt
Returns:
GetPromptResult: Getting prompt execution results
"""
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
_session_reset_token = None
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
_ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx)
_session_reset_token: Final = active_mcp_session_var.set(ctx.session)
try:
(
@ -1415,8 +1386,8 @@ if MCP_AVAILABLE:
verbose_logger.debug("MCP mcp_server_tool_call - User API Key Auth from context: %s", user_api_key_auth)
return await mcp_get_prompt(
name=name,
arguments=arguments,
name=params.name,
arguments=params.arguments,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
@ -1425,20 +1396,15 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
finally:
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
active_mcp_session_var.reset(_session_reset_token)
active_mcp_request_ctx_var.reset(_ctx_reset_token)
@server.list_resources()
async def list_resources() -> list[Resource]:
async def list_resources(ctx: ServerRequestContext, params: PaginatedRequestParams) -> ListResourcesResult:
"""List all available resources."""
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
_session_reset_token = None
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
_ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx)
_session_reset_token: Final = active_mcp_session_var.set(ctx.session)
try:
(
@ -1466,25 +1432,22 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
verbose_logger.info("MCP list_resources - Successfully returned %s resources", len(resources))
return resources
return ListResourcesResult(resources=resources)
except Exception as e:
verbose_logger.exception("Error in list_resources endpoint: %s", e)
return []
return ListResourcesResult(resources=[]) # mutable-ok: MCP result payload
finally:
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
active_mcp_session_var.reset(_session_reset_token)
active_mcp_request_ctx_var.reset(_ctx_reset_token)
@server.list_resource_templates()
async def list_resource_templates() -> list[ResourceTemplate]:
async def list_resource_templates(
ctx: ServerRequestContext, params: PaginatedRequestParams
) -> ListResourceTemplatesResult:
"""List all available resource templates."""
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
_session_reset_token = None
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
_ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx)
_session_reset_token: Final = active_mcp_session_var.set(ctx.session)
try:
(
@ -1514,24 +1477,19 @@ if MCP_AVAILABLE:
verbose_logger.info(
"MCP list_resource_templates - Successfully returned %s resource templates", len(resource_templates)
)
return resource_templates
return ListResourceTemplatesResult(resource_templates=resource_templates)
except Exception as e:
verbose_logger.exception("Error in list_resource_templates endpoint: %s", e)
return []
return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload
finally:
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
active_mcp_session_var.reset(_session_reset_token)
active_mcp_request_ctx_var.reset(_ctx_reset_token)
@server.read_resource()
async def read_resource(url: AnyUrl) -> list[ReadResourceContents]:
async def read_resource(ctx: ServerRequestContext, params: ReadResourceRequestParams) -> ReadResourceResult:
if _mcp_proxy_mode.get():
_reject_mcp_proxy_operation()
from mcp.server.lowlevel.server import request_ctx
req_ctx: Final = request_ctx.get(None)
_session_reset_token = None
if req_ctx:
_session_reset_token = active_mcp_session_var.set(req_ctx.session)
_ctx_reset_token: Final = active_mcp_request_ctx_var.set(ctx)
_session_reset_token: Final = active_mcp_session_var.set(ctx.session)
try:
(
@ -1545,7 +1503,7 @@ if MCP_AVAILABLE:
) = await get_or_extract_auth_context()
read_resource_result: Final = await mcp_read_resource(
url=url,
url=params.uri,
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
@ -1554,10 +1512,18 @@ if MCP_AVAILABLE:
raw_headers=raw_headers,
)
return _normalize_resource_contents(read_resource_result.contents)
return read_resource_result
finally:
if _session_reset_token is not None:
active_mcp_session_var.reset(_session_reset_token)
active_mcp_session_var.reset(_session_reset_token)
active_mcp_request_ctx_var.reset(_ctx_reset_token)
server.add_request_handler("tools/list", PaginatedRequestParams, handle_list_tools)
server.add_request_handler("tools/call", CallToolRequestParams, mcp_server_tool_call)
server.add_request_handler("prompts/list", PaginatedRequestParams, list_prompts)
server.add_request_handler("prompts/get", GetPromptRequestParams, get_prompt)
server.add_request_handler("resources/list", PaginatedRequestParams, list_resources)
server.add_request_handler("resources/templates/list", PaginatedRequestParams, list_resource_templates)
server.add_request_handler("resources/read", ReadResourceRequestParams, read_resource)
########################################################
############ End of MCP Server Routes ##################
@ -3296,11 +3262,11 @@ if MCP_AVAILABLE:
Guardrails run before the success/failure logging so the masked text, not
the raw one, is what gets logged.
A result with ``isError=True`` is logged as a failure (``status="failure"``
A result with ``is_error=True`` is logged as a failure (``status="failure"``
payload, so OTel marks the span ERROR) while the HTTP wire behavior stays
200 + ``isError: true`` per the MCP spec. The error check runs after
``async_post_mcp_tool_call_hook`` because guardrails may flip the result
to ``isError=True`` in that hook. Raised exceptions never reach here (the
to ``is_error=True`` in that hook. Raised exceptions never reach here (the
``@client`` wrapper and ``call_mcp_tool``'s except path log those), so
this cannot double-log a failure.
@ -3635,10 +3601,10 @@ if MCP_AVAILABLE:
"""Execute a local-registry tool and report whether it succeeded.
Returns the result rather than bare content because the verdict is part of it: the content
alone cannot say whether the handler failed, so callers used to stamp isError=False on every
alone cannot say whether the handler failed, so callers used to stamp is_error=False on every
outcome and an upstream rejection was served as tool output.
A failure is reported as ``isError=True`` here rather than raised, because the REST surface
A failure is reported as ``is_error=True`` here rather than raised, because the REST surface
turns an unrecognized exception into a 500 and an upstream 403 or 429 is not a gateway crash.
``MCPUpstreamAuthError`` is the exception: it propagates so the caller is told to
re-authenticate, which both renderers already know how to say.
@ -3660,8 +3626,14 @@ if MCP_AVAILABLE:
raise
except Exception as e:
verbose_logger.exception("Error executing local tool %s: %s", name, e)
return CallToolResult(content=[TextContent(text=f"Error: {e}", type="text")], isError=True)
return CallToolResult(content=[TextContent(text=str(result), type="text")], isError=False)
return CallToolResult(
content=[TextContent(text=f"Error: {e}", type="text")], # mutable-ok: MCP result content
is_error=True,
)
return CallToolResult(
content=[TextContent(text=str(result), type="text")], # mutable-ok: MCP result content
is_error=False,
)
def _get_mcp_servers_in_path(path: str) -> list[str] | None:
"""
@ -3843,7 +3815,7 @@ if MCP_AVAILABLE:
def _extract_initialize_client_info(body: bytes) -> Implementation | None:
try:
return InitializeRequest.model_validate_json(body).params.clientInfo
return InitializeRequest.model_validate_json(body, by_name=False).params.client_info
except ValidationError:
return None
@ -4553,6 +4525,21 @@ if MCP_AVAILABLE:
async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None:
"""Handle MCP requests through StreamableHTTP."""
try:
bad_version: Final = unsupported_protocol_version(scope)
if bad_version is not None:
supported: Final = ", ".join(sorted(HANDSHAKE_PROTOCOL_VERSIONS))
await JSONResponse(
status_code=400,
content={ # mutable-ok: JSON-RPC error payload
"jsonrpc": "2.0",
"id": None,
"error": {
"code": INVALID_REQUEST,
"message": f"Unsupported MCP-Protocol-Version {bad_version}; supported: {supported}",
},
},
)(scope, receive, send)
return
path: Final[str] = scope.get("path", "")
(
user_api_key_auth,
@ -5179,12 +5166,8 @@ if MCP_AVAILABLE:
return None, None, None, None, None, None, None
def _get_current_session():
try:
from mcp.server.lowlevel.server import request_ctx
return request_ctx.get().session
except (LookupError, ImportError):
return None
ctx: Final = get_active_mcp_request_ctx()
return ctx.session if ctx is not None else None
def _cache_auth_context_lazily():
session: Final = _get_current_session()

23
litellm/proxy/_experimental/mcp_server/tool_search.py
View file

@ -99,11 +99,20 @@ def mcp_tool_search_settings() -> MCPToolSearchSettings | ValidationError:
def _tool_result(tool: Tool) -> ToolSearchResult:
return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema}
return {
"name": tool.name,
"description": tool.description or "",
"inputSchema": tool.input_schema,
} # mutable-ok: wire schema payload
def _scored_result(tool: Tool, score: float) -> ToolSearchResult:
return {"name": tool.name, "description": tool.description or "", "inputSchema": tool.inputSchema, "score": score}
return {
"name": tool.name,
"description": tool.description or "",
"inputSchema": tool.input_schema,
"score": score,
} # mutable-ok: wire schema payload
_MCP_PROXY_IDENTITY_META_KEY: Final[str] = "litellm.ai/proxy_tool_identity"
@ -148,11 +157,11 @@ def _proxy_schema_result(tool: Tool) -> MCPProxySchemaResult:
"tool_id": mcp_proxy_tool_id(tool),
"name": tool.name,
"description": tool.description or "",
"inputSchema": tool.inputSchema,
"inputSchema": tool.input_schema,
}
if tool.outputSchema is None:
if tool.output_schema is None:
return base
return {**base, "outputSchema": tool.outputSchema} # mutable-ok: wire schema payload
return {**base, "outputSchema": tool.output_schema} # mutable-ok: wire schema payload
def _tool_text(tool: Tool) -> str:
@ -372,7 +381,7 @@ def _text_tool_result(text: str, is_error: bool) -> CallToolResult:
return CallToolResult(
content=[TextContent(type="text", text=text)], # mutable-ok: CallToolResult accepts only list content
isError=is_error,
is_error=is_error,
)
@ -565,7 +574,7 @@ async def handle_mcp_proxy_tool(
if not isinstance(tool_arguments, dict):
return _text_tool_result("arguments must be an object", is_error=True)
try:
validate(instance=tool_arguments, schema=tool.inputSchema)
validate(instance=tool_arguments, schema=tool.input_schema)
except JsonSchemaValidationError as exc:
return _text_tool_result(f"Invalid arguments: {exc.message}", is_error=True)

17
litellm/proxy/_experimental/mcp_server/utils.py
View file

@ -536,7 +536,11 @@ def extract_mcp_tool_result_error_message(result: object) -> str | None:
Accepts both ``mcp.types.CallToolResult`` objects and their dict
equivalents, duck-typed so the ``mcp`` package is not required.
"""
is_error: Final[object] = result.get("isError") if isinstance(result, Mapping) else getattr(result, "isError", None)
is_error: Final[object] = (
(result.get("isError") if result.get("isError") is not None else result.get("is_error"))
if isinstance(result, Mapping)
else getattr(result, "is_error", None)
)
if is_error is not True:
return None
content: Final[object] = result.get("content") if isinstance(result, Mapping) else getattr(result, "content", None)
@ -870,8 +874,9 @@ def json_unrewritable_labels(value: object, path_depth: int = 0) -> tuple[str, .
def mcp_tool_result_structured_content(result: object) -> object:
"""The ``structuredContent`` of an MCP tool result, or ``None`` when it has none."""
if isinstance(result, Mapping):
return result.get("structuredContent")
return getattr(result, "structuredContent", None)
structured: Final = result.get("structuredContent")
return structured if structured is not None else result.get("structured_content")
return getattr(result, "structured_content", None)
def set_mcp_tool_result_structured_content(result: object, value: object) -> bool:
@ -882,12 +887,12 @@ def set_mcp_tool_result_structured_content(result: object, value: object) -> boo
unmasked value in the spend log and the OTel span.
"""
if isinstance(result, MutableMapping):
result["structuredContent"] = value
result["structured_content" if "structured_content" in result else "structuredContent"] = value
return True
if not hasattr(result, "structuredContent"):
if not hasattr(result, "structured_content"):
return False
try:
setattr(result, "structuredContent", value) # attribute name is fixed by the MCP result shape
setattr(result, "structured_content", value) # attribute name is fixed by the MCP result shape
return True
except (AttributeError, TypeError, ValueError):
return False

11
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -3247,6 +3247,17 @@
],
"title": "Key Alias"
},
"key_exists": {
"anyOf": [
{
"type": "boolean"
},
{
"type": "null"
}
],
"title": "Key Exists"
},
"team_id": {
"anyOf": [
{

4
litellm/proxy/config_resolvers/__init__.py
View file

@ -5,6 +5,6 @@ from litellm.proxy.config_resolvers._descriptors import (
FieldSource,
resolve_fields,
)
from litellm.proxy.config_resolvers.settings_store import SettingsStore
from litellm.proxy.config_resolvers.settings_store import SettingsStore, config_ownership_message
__all__ = ("FieldDescriptor", "FieldSource", "SettingsStore", "resolve_fields")
__all__ = ("FieldDescriptor", "FieldSource", "SettingsStore", "config_ownership_message", "resolve_fields")

38
litellm/proxy/config_resolvers/settings_store.py
View file

@ -19,10 +19,21 @@ from litellm.proxy.config_resolvers.settings_rules import (
class ConfigOwnedKeyError(RuntimeError):
def __init__(self, section: Section, key: str) -> None:
super().__init__(f"{section}.{key} is set in the config file and cannot be changed at runtime")
def __init__(self, section: Section, key: str, *, shadows_db_value: bool = False) -> None:
super().__init__(config_ownership_message(section=section, key=key, shadows_db_value=shadows_db_value))
self.section: Final = section
self.key: Final = key
self.shadows_db_value: Final = shadows_db_value
def config_ownership_message(*, section: Section, key: str, shadows_db_value: bool) -> str:
stored: Final = (
" The value stored in the database for it is ignored and will never be applied." if shadows_db_value else ""
)
return (
f"{section}.{key} is set in the config file, so the config file owns it and it cannot be changed "
f"here.{stored} Edit the config file to change it, or remove it from the file to let the database own it."
)
_EMPTY_VALUES: Final[Mapping[str, JsonValue]] = MappingProxyType({})
@ -54,6 +65,13 @@ class SettingsStore(MutableMapping[str, JsonValue]):
)
)
def shadowed_db_keys(self) -> tuple[str, ...]:
"""Keys the config file owns whose stored value differs, so the stored one never reaches a reader."""
return tuple(sorted(key for key in self._yaml_values if self._db_value_is_shadowed(key)))
def shadows_db_value(self, key: str) -> bool:
return self.owned_by_config(key) and self._db_value_is_shadowed(key)
def apply_db_row(self, row: DbRow, db_row: Mapping[str, JsonValue]) -> None:
previous_row: Final = self._database_rows.get(row, _EMPTY_VALUES)
self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))})
@ -81,7 +99,7 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def __setitem__(self, key: str, value: JsonValue) -> None:
if self.owned_by_config(key) and value != self.get(key):
raise ConfigOwnedKeyError(self._section, key)
raise ConfigOwnedKeyError(self._section, key, shadows_db_value=self._db_value_is_shadowed(key))
self._runtime_values = MappingProxyType({**self._runtime_values, key: value})
self._deleted_runtime_keys = self._deleted_runtime_keys - frozenset((key,))
@ -89,7 +107,7 @@ class SettingsStore(MutableMapping[str, JsonValue]):
if key not in self:
raise KeyError(key)
if self.owned_by_config(key):
raise ConfigOwnedKeyError(self._section, key)
raise ConfigOwnedKeyError(self._section, key, shadows_db_value=self._db_value_is_shadowed(key))
self._runtime_values = MappingProxyType(
{key_: value for key_, value in self._runtime_values.items() if key_ != key}
)
@ -136,8 +154,14 @@ class SettingsStore(MutableMapping[str, JsonValue]):
)
)
def _resolution_for(self, key: str) -> Resolved:
def _db_value(self, key: str) -> SettingValue:
rule: Final = rule_for(self._section, key)
return self._database_rows.get(rule.db_row, _EMPTY_VALUES).get(key, ABSENT)
def _db_value_is_shadowed(self, key: str) -> bool:
db_value: Final = self._db_value(key)
return not isinstance(db_value, Absent) and db_value is not None and db_value != self.get(key)
def _resolution_for(self, key: str) -> Resolved:
yaml_value: Final[SettingValue] = self._yaml_values.get(key, ABSENT)
db_value: Final[SettingValue] = self._database_rows.get(rule.db_row, _EMPTY_VALUES).get(key, ABSENT)
return resolve(yaml_value, db_value)
return resolve(yaml_value, self._db_value(key))

40
litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense_mcp.py
View file

@ -34,15 +34,26 @@ def _serialize_mcp_content_item(item: object) -> dict[str, object]:
model_dump: Final = getattr(item, "model_dump", None)
if callable(model_dump):
try:
return dict(model_dump(exclude_none=True))
dumped: Final[dict[str, object]] = model_dump(exclude_none=True, by_alias=True)
return dict(dumped)
except TypeError:
return dict(model_dump())
dumped_fallback: Final[dict[str, object]] = model_dump()
return dict(dumped_fallback)
text: Final = getattr(item, "text", None)
if isinstance(text, str):
return {"type": getattr(item, "type", "text"), "text": text}
return {"type": "text", "text": str(item)}
def _source_field(source: object, key: str, snake_key: str) -> object:
if isinstance(source, dict):
for candidate in (key, snake_key):
if candidate in source:
return source[candidate] # pyright: ignore[reportUnknownVariableType] # dict-shaped sources arrive untyped
return None
return getattr(source, snake_key, None)
class _CiscoAIDefenseMcpMixin:
"""MCP-specific instance methods for ``CiscoAIDefenseGuardrail``.
@ -219,14 +230,14 @@ class _CiscoAIDefenseMcpMixin:
if isinstance(content, list):
content[:] = replacement
structured_replacement: Final = _CiscoAIDefenseMcpMixin._replacement_structured_content(replacement)
if hasattr(response_obj, "structuredContent"):
if hasattr(response_obj, "structured_content"):
try:
setattr(response_obj, "structuredContent", structured_replacement)
setattr(response_obj, "structured_content", structured_replacement)
except (AttributeError, TypeError, ValueError):
pass
if hasattr(response_obj, "isError"):
if hasattr(response_obj, "is_error"):
try:
setattr(response_obj, "isError", True)
setattr(response_obj, "is_error", True)
except (AttributeError, TypeError, ValueError):
pass
return True
@ -487,7 +498,7 @@ class _CiscoAIDefenseMcpMixin:
model_dump: Final = getattr(response, "model_dump", None)
if callable(model_dump):
try:
dumped = model_dump(exclude_none=True)
dumped = model_dump(exclude_none=True, by_alias=True)
except TypeError:
dumped = model_dump()
if isinstance(dumped, dict):
@ -507,8 +518,8 @@ class _CiscoAIDefenseMcpMixin:
source: object = None,
) -> dict[str, object]:
result: Final[dict[str, object]] = {"content": [_serialize_mcp_content_item(item) for item in content]}
for key in ("structuredContent", "isError"):
value = source.get(key) if isinstance(source, dict) else getattr(source, key, None)
for key, snake_key in (("structuredContent", "structured_content"), ("isError", "is_error")):
value = _source_field(source, key, snake_key)
if value is not None and (key != "isError" or isinstance(value, bool)):
result[key] = value
return result
@ -549,20 +560,21 @@ class _CiscoAIDefenseMcpMixin:
and all(isinstance(item, tuple) and len(item) == 2 and isinstance(item[0], str) for item in response_obj)
):
for index, item in enumerate(response_obj):
if item[0] == "structuredContent":
if item[0] in ("structuredContent", "structured_content"):
response_obj[index] = (item[0], replacement)
replaced = True
elif hasattr(response_obj, "structuredContent"):
elif hasattr(response_obj, "structured_content"):
try:
setattr(response_obj, "structuredContent", replacement)
setattr(response_obj, "structured_content", replacement)
replaced = True
except (AttributeError, TypeError, ValueError):
pass
elif isinstance(response_obj, dict):
result: Final = response_obj.get("result")
target: Final[dict[object, object]] = result if isinstance(result, dict) else response_obj
if "structuredContent" in target:
target["structuredContent"] = replacement
structured_key: Final = "structured_content" if "structured_content" in target else "structuredContent"
if structured_key in target:
target[structured_key] = replacement
replaced = True
return replaced

3
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -127,6 +127,7 @@ class _KeyMetadataDict(TypedDict, total=False):
team_id: ReadOnly[str | None]
user_id: ReadOnly[str | None]
user_email: ReadOnly[str | None]
key_exists: ReadOnly[bool]
def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata:
@ -136,6 +137,7 @@ def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str
team_id=meta.get("team_id"),
user_id=meta.get("user_id"),
user_email=meta.get("user_email"),
key_exists=meta.get("key_exists", False),
)
@ -512,6 +514,7 @@ async def get_api_key_metadata(
"key_alias": k.key_alias,
"team_id": k.team_id,
"user_id": getattr(k, "user_id", None),
"key_exists": True,
}
for k in key_records
}

4
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -1968,7 +1968,7 @@ async def generate_key_fn(
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
- permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
@ -3323,7 +3323,7 @@ async def update_key_fn(
- policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
- disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
- throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only.
- enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only.
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
- blocked: Optional[bool] - Whether the key is blocked
- aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases)

28
litellm/proxy/proxy_server.py
View file

@ -447,7 +447,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.config_resolvers import SettingsStore, resolve_fields
from litellm.proxy.config_resolvers import SettingsStore, config_ownership_message, resolve_fields
from litellm.proxy.config_resolvers.alerting import (
EMAIL_DESCRIPTORS,
MS_TEAMS_DESCRIPTORS,
@ -4907,6 +4907,7 @@ class ProxyConfig:
self.router_settings: Final[SettingsStore] = SettingsStore("router_settings")
self.litellm_settings: Final[SettingsStore] = SettingsStore("litellm_settings")
self.environment_variables: Final[SettingsStore] = SettingsStore("environment_variables")
self._warned_shadowed_keys: frozenset[tuple[Section, str]] = frozenset()
self._settings_stores: Final[Mapping[Section, SettingsStore]] = MappingProxyType(
{
"general_settings": self.settings,
@ -5128,12 +5129,20 @@ class ProxyConfig:
f"key '{rejected[0]}' is" if len(rejected) == 1 else f"keys {', '.join(repr(key) for key in rejected)} are"
)
pronoun: Final = "it" if len(rejected) == 1 else "them"
shadowed: Final = tuple(key for key in rejected if store.shadows_db_value(key))
stored: Final = (
f" The {'value' if len(shadowed) == 1 else 'values'} already stored in the database for "
f"{', '.join(shadowed)} {'is' if len(shadowed) == 1 else 'are'} ignored and will never be applied."
if shadowed
else ""
)
raise HTTPException(
status_code=400,
detail={
"error": f"{section_name} {subject} set in the config file and cannot be changed here",
"error": f"{section_name} {subject} set in the config file and cannot be changed here.{stored}",
"keys": list(rejected),
"section": section_name,
"stored_database_values_ignored": list(shadowed),
"resolution": (
f"edit {user_config_file_path} to change {pronoun}, "
f"or remove {pronoun} from the file to let the database own {pronoun}"
@ -7430,8 +7439,19 @@ class ProxyConfig:
self._prepared_db_settings_values(section, param_value),
)
self._warn_about_shadowed_db_settings()
return self._config_with_resolved_settings(config)
def _warn_about_shadowed_db_settings(self) -> None:
shadowed: Final[frozenset[tuple[Section, str]]] = frozenset(
(section, key) for section, store in self._settings_stores.items() for key in store.shadowed_db_keys()
)
for section, key in sorted(shadowed - self._warned_shadowed_keys):
verbose_proxy_logger.warning(
"%s", config_ownership_message(section=section, key=key, shadows_db_value=True)
)
self._warned_shadowed_keys = shadowed
def _prepared_db_settings_values(self, section: Section, value: object) -> Mapping[str, SettingsJsonValue]:
if section == "environment_variables":
decrypted: Final = self._decrypt_and_set_db_env_variables(
@ -17688,8 +17708,8 @@ _GENERAL_SETTINGS_UI_LITELLM_FIELDS: Final[dict[str, GeneralSettingsUILiteLLMFie
"type": "Boolean",
"tab": "prompt_caching",
"description": (
"Auto-adds cache_control to the system prompt and trailing turn for supported Anthropic "
"and Bedrock Claude models. The cache is shared across callers on the same upstream credentials."
"Auto-adds cache_control to the system prompt and trailing turn for supported Claude models on "
"Anthropic, Bedrock, Vertex AI, and Azure AI. The cache is shared across callers on the same upstream credentials."
),
},
"anthropic_prompt_caching_ttl": {

88
litellm/proxy/public_endpoints/provider_create_fields.json
View file

@ -209,6 +209,94 @@
],
"default_model_placeholder": "claude-3-opus"
},
{
"provider": "AWS_Textract",
"provider_display_name": "Amazon Textract",
"litellm_provider": "aws_textract",
"credential_fields": [
{
"key": "aws_access_key_id",
"label": "AWS Access Key ID",
"placeholder": null,
"tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).",
"required": false,
"field_type": "password",
"options": null,
"default_value": null
},
{
"key": "aws_secret_access_key",
"label": "AWS Secret Access Key",
"placeholder": null,
"tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).",
"required": false,
"field_type": "password",
"options": null,
"default_value": null
},
{
"key": "aws_session_token",
"label": "AWS Session Token",
"placeholder": null,
"tooltip": "Temporary credentials session token. You can provide the raw token or the environment variable (e.g. `os.environ/MY_SESSION_TOKEN`).",
"required": false,
"field_type": "password",
"options": null,
"default_value": null
},
{
"key": "aws_region_name",
"label": "AWS Region Name",
"placeholder": "us-east-1",
"tooltip": "You can provide the raw key or the environment variable (e.g. `os.environ/MY_SECRET_KEY`).",
"required": false,
"field_type": "text",
"options": null,
"default_value": null
},
{
"key": "aws_session_name",
"label": "AWS Session Name",
"placeholder": "my-session",
"tooltip": "Name for the AWS session. You can provide the raw value or the environment variable (e.g. `os.environ/MY_SESSION_NAME`).",
"required": false,
"field_type": "text",
"options": null,
"default_value": null
},
{
"key": "aws_profile_name",
"label": "AWS Profile Name",
"placeholder": "default",
"tooltip": "AWS profile name to use for authentication. You can provide the raw value or the environment variable (e.g. `os.environ/MY_PROFILE_NAME`).",
"required": false,
"field_type": "text",
"options": null,
"default_value": null
},
{
"key": "aws_role_name",
"label": "AWS Role Name",
"placeholder": "MyRole",
"tooltip": "AWS IAM role name to assume. You can provide the raw value or the environment variable (e.g. `os.environ/MY_ROLE_NAME`).",
"required": false,
"field_type": "text",
"options": null,
"default_value": null
},
{
"key": "aws_web_identity_token",
"label": "AWS Web Identity Token",
"placeholder": null,
"tooltip": "Web identity token for OIDC authentication. You can provide the raw token or the environment variable (e.g. `os.environ/MY_WEB_IDENTITY_TOKEN`).",
"required": false,
"field_type": "password",
"options": null,
"default_value": null
}
],
"default_model_placeholder": "detect-document-text"
},
{
"provider": "BedrockMantle",
"provider_display_name": "Amazon Bedrock Mantle",

2
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -34,6 +34,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_set_request_parsed_body,
)
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.types.llms.openai import (
REASONING_EFFORT,
ResponsesAPIOptionalRequestParams,
@ -296,6 +297,7 @@ async def responses_api(
route_type="aresponses",
llm_router=llm_router,
)
raise_if_required_body_param_missing(route_type="aresponses", data=data)
except Exception as e:
raise await processor._handle_llm_api_exception(
e=e,

1
litellm/proxy/route_llm_request.py
View file

@ -159,6 +159,7 @@ class ProxyModelNotFoundError(HTTPException):
REQUIRED_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = {
"acompletion": ("messages",),
"aembedding": ("input",),
"aresponses": ("input",),
"acreate_batch": ("input_file_id", "endpoint", "completion_window"),
}

1
litellm/proxy/spend_tracking/key_metadata_recovery.py
View file

@ -69,6 +69,7 @@ class KeyMetadataDict(TypedDict, total=False):
team_id: ReadOnly[str | None]
user_id: ReadOnly[str | None]
user_email: ReadOnly[str | None]
key_exists: ReadOnly[bool]
class _TokenDigestRow(BaseModel):

6
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -497,12 +497,10 @@ def _store_allowed_ips(general_settings: MutableMapping[str, object], allowed_ip
raise HTTPException(
status_code=400,
detail={ # mutable-ok: HTTPException serializes its detail as json
"error": f"{owned.section} key '{owned.key}' is set in the config file and cannot be changed here",
"error": str(owned),
"keys": (owned.key,),
"section": owned.section,
"resolution": (
"edit the config file to change it, or remove it from the file to let the database own it"
),
"stored_database_value_ignored": owned.shadows_db_value,
},
) from owned

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