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

chore(auto-router): merge main into JEV launch branch
Some checks failed
LiteLLM Rust / rust-lint (push) Waiting to run
Details
LiteLLM Rust / rust-test (push) Waiting to run
Details
LiteLLM Rust / rust-wheel (push) Waiting to run
Details
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Details
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Details

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Moe Khalil 2026-09-19 23:28:15 +00:00
commit e77c154c36
313 changed files with 17819 additions and 6127 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:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

4
enterprise/pyproject.toml
View file

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

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

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

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

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

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

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

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

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

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

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

1
litellm-rust/Cargo.lock generated
View file

@ -3046,7 +3046,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"h2 0.4.15",

2
litellm-rust/Cargo.toml
View file

@ -34,7 +34,7 @@ pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }

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

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

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

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

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

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

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

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

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

18
litellm-rust/crates/python-bridge/src/diagnostics.rs
View file

@ -1,5 +1,5 @@
use litellm_host_python::release_count;
use pyo3::{prelude::*, types::PyDict};
use litellm_host_python::{release_count, runtime_started};
use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict};
#[pyfunction]
pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
@ -8,6 +8,20 @@ pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
Ok(stats.into_any().unbind())
}
/// True once this process has started the native runtime, which does not survive `fork()`.
#[pyfunction]
pub(crate) fn process_state_started() -> bool {
runtime_started()
}
/// Declares that this process only forks workers: from now on every native route raises here,
/// so the runtime can never start. Raises if it already has. Forked workers are unaffected.
#[pyfunction]
pub(crate) fn reserve_process_for_forking() -> PyResult<()> {
litellm_host_python::reserve_process_for_forking()
.map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process"))
}
#[cfg(feature = "panic-test")]
#[pyfunction]
pub(crate) fn _panic_for_test() {

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

@ -13,7 +13,7 @@ mod _native {
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
#[pymodule_export]
use crate::diagnostics::gil_stats;
use crate::diagnostics::{gil_stats, process_state_started, reserve_process_for_forking};
#[pymodule_export]
use crate::errors::{RustBridgeDeclined, RustUpstreamError};
#[pymodule_export]
@ -30,6 +30,8 @@ mod _native {
use crate::routes::responses::ResponsesWebSocketConnection;
#[pymodule_export]
use crate::token_counter::TokenCounter;
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
}
use pyo3::prelude::*;
@ -50,6 +52,8 @@ mod tests {
let mut expected = vec![
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",
"ProcessReservedForForking",
"ocr",
"aocr",
"transcription",
@ -62,6 +66,8 @@ mod tests {
"ResponsesWebSocketConnection",
"TokenCounter",
"gil_stats",
"process_state_started",
"reserve_process_for_forking",
];
expected.sort_unstable();

12
litellm-rust/crates/python-bridge/src/routes/responses.rs
View file

@ -25,7 +25,7 @@ impl ResponsesWebSocketConnection {
) -> PyResult<Bound<'py, PyAny>> {
let headers = marshal_headers(headers)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
.await
.map_err(responses_error_to_pyerr)?;
@ -35,7 +35,7 @@ impl ResponsesWebSocketConnection {
fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner
.send_text(text)
.await
@ -45,14 +45,14 @@ impl ResponsesWebSocketConnection {
fn recv_text<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.recv_text().await.map_err(responses_error_to_pyerr)
})
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.close().await.map_err(responses_error_to_pyerr)
})
}
@ -68,6 +68,10 @@ mod tests {
use tokio_tungstenite::{accept_async, tungstenite::Message};
#[test]
#[expect(
clippy::disallowed_methods,
reason = "the test server shares the routes' runtime"
)]
fn responses_websocket_connection_round_trips_through_python() {
Python::initialize();
let runtime = pyo3_async_runtimes::tokio::get_runtime();

1
litellm/constants.py
View file

@ -1694,6 +1694,7 @@ LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0)
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, int(os.getenv("LITELLM_EXECUTED_BATCH_CONCURRENCY", "4")))
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"

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

22
litellm/integrations/custom_guardrail.py
View file

@ -16,7 +16,6 @@ from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
redact_nested_match_and_regex_keys,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import REQUEST_SCAN_CONTEXT_KEY
from litellm.secret_managers.main import str_to_bool
from litellm.types.guardrails import (
DynamicGuardrailParams,
@ -945,29 +944,10 @@ class CustomGuardrail(CustomLogger):
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
if response is None:
return
output_request: Final = (
scratch_request
if type(output_translation) is type(translation)
else self._chat_shaped_request(scratch_request, translation)
)
await output_translation.process_output_response(
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=output_request
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
)
def _chat_shaped_request(
self,
scratch_request: Mapping[str, object],
translation: "BaseTranslation",
) -> dict[str, object]: # mutable-ok: BaseTranslation.process_output_response contract
"""The logged request in OpenAI chat shape, for an output scan whose translation differs from the input's."""
context: Final = translation.request_scan_context(scratch_request, self)
return {
**scratch_request,
"messages": list(context.structured_messages),
"tools": list(context.tools),
REQUEST_SCAN_CONTEXT_KEY: context,
}
def supports_scan_only_tool_results(self) -> bool:
"""Whether this guardrail can scan tool-result content.

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

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

@ -3,13 +3,15 @@
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
from typing import TYPE_CHECKING, ClassVar, Final, cast
from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast
from urllib.parse import urlsplit
from typing_extensions import ReadOnly, TypedDict
from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity
from litellm.integrations.otel.model.semconv import (
GenAIOperation,
@ -25,6 +27,7 @@ from litellm.integrations.otel.model.utils import (
as_float,
as_int,
as_str,
as_str_mapping,
as_str_tuple,
)
@ -353,6 +356,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 +407,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(
@ -405,7 +427,7 @@ class LLMCallSpanData:
# plain ``.get`` — no repeated ``isinstance`` guards.
raw_response: Final = payload.get("response")
response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {})
choices_out: Final = _dicts(response.get("choices"))
choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response)
# ``finish_reasons`` is metadata, not content, so derive it from
# ``choices_out`` before gating. The raw message/choice bodies are only
# retained when content capture is enabled (see ``capture_span_content``);
@ -413,8 +435,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 +463,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,
)
@ -679,6 +706,84 @@ def _finish_reasons(choices: tuple[Mapping[str, object], ...]) -> tuple[str, ...
return tuple(r for c in choices if (r := as_str(c.get("finish_reason"))))
class _ToolFunction(TypedDict):
name: ReadOnly[str]
arguments: ReadOnly[str]
class _ToolCall(TypedDict):
id: ReadOnly[str]
type: ReadOnly[Literal["function"]]
function: ReadOnly[_ToolFunction]
class _AssistantMessage(TypedDict):
role: ReadOnly[str]
content: ReadOnly[str | None]
refusal: ReadOnly[str | None]
tool_calls: ReadOnly[tuple[_ToolCall, ...] | None]
class _Choice(TypedDict):
message: ReadOnly[_AssistantMessage]
finish_reason: ReadOnly[str | None]
_RESPONSES_TOOL_CALL_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]:
"""A Responses API ``output`` folded into one chat-shaped assistant choice."""
items: Final = _dicts(response.get("output"))
messages: Final = tuple(item for item in items if item.get("type") == "message")
parts: Final = tuple(part for item in messages for part in _dicts(item.get("content")))
tool_calls: Final = tuple(
_responses_tool_call(item) for item in items if item.get("type") in _RESPONSES_TOOL_CALL_TYPES
)
if not messages and not tool_calls:
return ()
message: Final[_AssistantMessage] = {
"role": next((role for item in messages if (role := as_str(item.get("role")))), "assistant"),
"content": _responses_parts_text(parts, "output_text", "text"),
"refusal": _responses_parts_text(parts, "refusal", "refusal"),
"tool_calls": tool_calls or None,
}
choice: Final[_Choice] = {"message": message, "finish_reason": _responses_finish_reason(response, bool(tool_calls))}
return (choice,)
def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None:
texts: Final = tuple(
text for part in parts if part.get("type") == part_type if (text := as_str(part.get(field))) is not None
)
return "".join(texts) if texts else None
def _responses_tool_call(item: Mapping[str, object]) -> _ToolCall:
custom: Final = item.get("type") == "custom_tool_call"
function: Final[_ToolFunction] = {
"name": as_str(item.get("name")) or "",
"arguments": as_str(item.get("input" if custom else "arguments")) or "",
}
tool_call: Final[_ToolCall] = {
"id": as_str(item.get("call_id")) or as_str(item.get("id")) or "",
"type": "function",
"function": function,
}
return tool_call
def _responses_finish_reason(response: Mapping[str, object], has_tool_calls: bool) -> str | None:
status: Final = as_str(response.get("status"))
if status == "completed":
return "tool_calls" if has_tool_calls else "stop"
if status != "incomplete":
return None
details: Final = as_str_mapping(response.get("incomplete_details"))
reason: Final = details.get("reason") if details is not None else None
return "content_filter" if reason == "content_filter" else "length"
def _parse_error(payload: StandardLoggingPayload) -> SpanError | None:
"""A ``SpanError`` for a failed request, or ``None`` on success."""
if payload.get("status") != "failure":

85
litellm/litellm_core_utils/litellm_logging.py
View file

@ -212,6 +212,7 @@ if TYPE_CHECKING:
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates
from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector
from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation
try:
from litellm_enterprise.enterprise_callbacks.callback_controls import (
EnterpriseCallbackControls,
@ -501,6 +502,8 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_request_debug: bool = False
streamed_anthropic_message_id: str | None = None
classifier_input: Mapping[str, JsonValue] | None = None
baseline_cache_context: "BaselineCacheContext | None" = None
baseline_observation: "CapturedBaselineObservation | None" = None
def __init__(
self,
@ -508,7 +511,7 @@ class Logging(LiteLLMLoggingBaseClass):
messages,
stream,
call_type,
start_time,
start_time: datetime.datetime,
litellm_call_id: str,
function_id: str,
litellm_trace_id: str | None = None,
@ -2181,6 +2184,7 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result,
start_time,
end_time,
build_logging_payload: bool = True,
):
"""Resolve hidden params, compute response cost, and emit the standard logging payload."""
hidden_params: Final = getattr(logging_result, "_hidden_params", {})
@ -2205,6 +2209,9 @@ class Logging(LiteLLMLoggingBaseClass):
else:
self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result)
if not build_logging_payload:
return
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
logging_result, start_time, end_time
)
@ -2215,6 +2222,19 @@ class Logging(LiteLLMLoggingBaseClass):
if standard_logging_payload is not None:
emit_standard_logging_payload(standard_logging_payload)
async def _prepare_baseline_cache_estimate(self, response_obj: object) -> None:
if self.baseline_cache_context is None:
return
from litellm.proxy.hooks.autorouter_baseline_cache import finalize_baseline_cache
await finalize_baseline_cache(self, response_obj)
async def invalidate_baseline_cache_estimate(self, reason: str, *, completed: bool = False) -> None:
"""Invalidate uncertain attempts; retire the reservation at logical completion."""
from litellm.proxy.hooks.autorouter_baseline_cache import invalidate_baseline_cache
await invalidate_baseline_cache(self, reason, completed=completed)
def _build_standard_logging_payload(
self, init_response_obj: object, start_time: Any, end_time: Any
) -> StandardLoggingPayload | None:
@ -2266,6 +2286,7 @@ class Logging(LiteLLMLoggingBaseClass):
end_time=None,
cache_hit=None,
standard_logging_object: StandardLoggingPayload | None = None,
build_logging_payload: bool = True,
):
try:
if start_time is None:
@ -2303,6 +2324,7 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result=logging_result,
start_time=start_time,
end_time=end_time,
build_logging_payload=build_logging_payload,
)
elif standard_logging_object is not None:
self.model_call_details["standard_logging_object"] = standard_logging_object
@ -3051,8 +3073,17 @@ class Logging(LiteLLMLoggingBaseClass):
result=result,
cache_hit=cache_hit,
standard_logging_object=kwargs.get("standard_logging_object", None),
build_logging_payload=self.baseline_cache_context is None,
)
if self.stream is not True and self.baseline_cache_context is not None:
await self._prepare_baseline_cache_estimate(result)
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
result, start_time, end_time
)
if (prepared_payload := self.model_call_details.get("standard_logging_object")) is not None:
emit_standard_logging_payload(prepared_payload)
## BUILD COMPLETE STREAMED RESPONSE
if "async_complete_streaming_response" in self.model_call_details:
return # break out of this.
@ -3097,6 +3128,8 @@ class Logging(LiteLLMLoggingBaseClass):
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
await self._prepare_baseline_cache_estimate(complete_streaming_response)
## STANDARDIZED LOGGING PAYLOAD
try:
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
@ -3125,6 +3158,7 @@ class Logging(LiteLLMLoggingBaseClass):
# Only build standard_logging_object if not already built by
# _success_handler_helper_fn
if self.model_call_details.get("standard_logging_object") is None:
await self._prepare_baseline_cache_estimate(result)
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
result, start_time, end_time
@ -3631,6 +3665,8 @@ class Logging(LiteLLMLoggingBaseClass):
"""
Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions.
"""
if self.baseline_cache_context is not None:
await self.invalidate_baseline_cache_estimate("failed_request")
await self.special_failure_handlers(exception=exception)
if not self.should_run_logging(event_type="async_failure"): # prevent double logging
return
@ -5528,6 +5564,10 @@ class StandardLoggingPayloadSetup:
for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
clean_metadata[key] = metadata[key]
recorded_guardrails: Final = metadata.get("applied_guardrails")
if applied_guardrails and isinstance(recorded_guardrails, list):
clean_metadata["applied_guardrails"] = list(dict.fromkeys([*applied_guardrails, *recorded_guardrails]))
user_api_key: Final = metadata.get("user_api_key")
if user_api_key and isinstance(user_api_key, str) and is_valid_sha256_hash(user_api_key):
clean_metadata["user_api_key_hash"] = user_api_key
@ -6149,6 +6189,8 @@ def _autorouter_savings_for_payload(
model_id: str | None,
usage_object: Mapping[str, object] | None,
cost_breakdown: Mapping[str, object] | None,
baseline_usage: Usage | None = None,
baseline_provenance: Literal["observed_initial", "modeled"] | None = None,
) -> float | None:
"""The auto-router savings figure for the payload, or ``None`` when there is none.
@ -6167,6 +6209,8 @@ def _autorouter_savings_for_payload(
model_id=model_id,
usage_object=usage_object,
cost_breakdown=cost_breakdown,
baseline_usage=baseline_usage,
baseline_provenance=baseline_provenance,
)
except Exception as e: # noqa: BLE001 # a savings figure must never fail request logging
verbose_logger.debug("autorouter savings skipped on logging payload: %s", e)
@ -6343,13 +6387,18 @@ def get_standard_logging_object_payload(
model_name = response_model_name
request_cost_breakdown: Final = cost_breakdown_with_guardrail(logging_obj.cost_breakdown, guardrail_cost)
autorouter_savings: Final = _autorouter_savings_for_payload(
request_metadata=metadata,
model=model_name,
custom_llm_provider=custom_llm_provider,
model_id=_model_id,
usage_object=usage_dict,
cost_breakdown=request_cost_breakdown,
captured_baseline: Final = logging_obj.baseline_observation
autorouter_savings: Final = (
None
if status != "success" or cache_hit or logging_obj.baseline_cache_context is not None
else _autorouter_savings_for_payload(
request_metadata=metadata,
model=model_name,
custom_llm_provider=custom_llm_provider,
model_id=_model_id,
usage_object=usage_dict,
cost_breakdown=request_cost_breakdown,
)
)
payload: Final[StandardLoggingPayload] = StandardLoggingPayload(
@ -6396,6 +6445,26 @@ def get_standard_logging_object_payload(
response_cost=response_cost,
cost_breakdown=request_cost_breakdown,
autorouter_savings=autorouter_savings,
autorouter_savings_estimate=(
{
"version": 3,
"status": "unknown",
"reason": "pending_projection",
} # mutable-ok: spend-log JSON serialization requires plain mappings
if captured_baseline is not None
else (
{ # mutable-ok: spend-log JSON serialization requires plain mappings
"version": 1,
"status": "estimated" if autorouter_savings is not None else "unknown",
"reason": "uncached_usage" if autorouter_savings is not None else "baseline_unavailable",
}
if metadata.get("routing_decision")
else None
)
),
autorouter_baseline_observation=(
captured_baseline.model_dump_json() if captured_baseline is not None else None
),
total_tokens=usage_dict.get("total_tokens", 0),
prompt_tokens=usage_dict.get("prompt_tokens", 0),
completion_tokens=usage_dict.get("completion_tokens", 0),

8
litellm/litellm_core_utils/redact_messages.py
View file

@ -128,6 +128,8 @@ def _redact_responses_api_output(output_items):
for content_part in output_item.content:
if getattr(content_part, "text", None) is not None:
content_part.text = REDACTED_BY_LITELLM
if getattr(content_part, "refusal", None) is not None:
content_part.refusal = REDACTED_BY_LITELLM
# Redact reasoning items in output array
if hasattr(output_item, "type") and output_item.type == "reasoning":
@ -138,6 +140,8 @@ def _redact_responses_api_output(output_items):
if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"):
output_item.arguments = REDACTED_BY_LITELLM
if hasattr(output_item, "type") and output_item.type == "custom_tool_call" and hasattr(output_item, "input"):
output_item.input = REDACTED_BY_LITELLM
def _redact_responses_api_output_dict(output_items, redacted_str: str):
@ -153,6 +157,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
for content_item in output_item["content"]:
if isinstance(content_item, dict) and content_item.get("text") is not None:
content_item["text"] = redacted_str
if isinstance(content_item, dict) and content_item.get("refusal") is not None:
content_item["refusal"] = redacted_str
if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list):
for summary_item in output_item["summary"]:
@ -161,6 +167,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
if output_item.get("type") == "function_call" and "arguments" in output_item:
output_item["arguments"] = redacted_str
if output_item.get("type") == "custom_tool_call" and "input" in output_item:
output_item["input"] = redacted_str
def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]:

35
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -31,7 +31,6 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
)
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
RequestScanContext,
StreamingScanKey,
StreamTransformSink,
)
@ -529,26 +528,6 @@ class AnthropicMessagesHandler(BaseTranslation):
)
return result if result else None
def request_scan_context(
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
) -> RequestScanContext:
if data.get("messages") is None:
return RequestScanContext()
translated: Final = self._translate_to_openai(
{key: value for key, value in data.items() if key != "system"} # mutable-ok: API message payload
)
hoisted_system_message: Final = (
None
if effective_skip_system_message_for_guardrail(guardrail_to_apply)
else self._hoisted_top_level_system_message(data)
)
return RequestScanContext.scoped(
(*(() if hoisted_system_message is None else (hoisted_system_message,)), *translated["messages"]),
tuple(tool for tool in translated.get("tools") or () if not is_provider_native_tool_dict(tool)),
guardrail_to_apply,
skip_system=False,
)
async def process_input_messages(
self,
data: dict,
@ -718,7 +697,9 @@ class AnthropicMessagesHandler(BaseTranslation):
return data
def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None:
def _hoisted_top_level_system_message(
self, data: dict
) -> AllMessageValues | None: # mutable-ok: API message payload
"""Return the system message produced by translating the top-level prompt."""
system: Final = data.get("system")
if not system:
@ -1220,7 +1201,7 @@ class AnthropicMessagesHandler(BaseTranslation):
)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -1292,7 +1273,7 @@ class AnthropicMessagesHandler(BaseTranslation):
key="response",
)
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(guardrail_inputs, prepared_request_data, guardrail_to_apply),
inputs=guardrail_inputs,
request_data=prepared_request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -1342,11 +1323,7 @@ class AnthropicMessagesHandler(BaseTranslation):
key="responses",
)
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(
GenericGuardrailAPIInputs(texts=[string_so_far]), # mutable-ok: guardrail inputs want a list
prepared_request_data,
guardrail_to_apply,
),
inputs={"texts": [string_so_far]},
request_data=prepared_request_data,
input_type="response",
logging_obj=litellm_logging_obj,

11
litellm/llms/anthropic/chat/handler.py
View file

@ -4,7 +4,8 @@ Calling + translation logic for anthropic's `/v1/messages` endpoint
import copy
import json
from collections.abc import Callable
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast
import httpx
@ -31,7 +32,6 @@ from litellm.types.llms.anthropic import (
ContentBlockStop,
MessageBlockDelta,
MessageStartBlock,
UsageDelta,
)
from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
@ -557,6 +557,7 @@ class ModelResponseIterator:
self.tool_index = -1
self.json_mode = json_mode
self.speed = speed
self._cumulative_usage: Mapping[str, object] = MappingProxyType({})
# rewritten-name -> caller's original. Built per-request from the
# forward map in AnthropicConfig._build_request_tool_name_maps; only
# contains entries we actually rewrote, so a tool legitimately named
@ -631,10 +632,12 @@ class ModelResponseIterator:
return True
return False
def _handle_usage(self, anthropic_usage_chunk: dict | UsageDelta) -> Usage:
def _handle_usage(self, anthropic_usage_chunk: Mapping[str, object]) -> Usage:
# message_delta usage is cumulative but may omit fields reported at message_start.
self._cumulative_usage = MappingProxyType({**self._cumulative_usage, **anthropic_usage_chunk})
reasoning_content: Final = "".join(self.reasoning_content_chunks) if self.reasoning_content_chunks else None
usage: Final = AnthropicConfig().calculate_usage(
usage_object=cast(dict, anthropic_usage_chunk),
usage_object=self._cumulative_usage,
reasoning_content=reasoning_content,
speed=self.speed,
)

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

20
litellm/llms/anthropic/count_tokens/handler.py
View file

@ -4,9 +4,11 @@ Anthropic CountTokens API handler.
Uses httpx for HTTP requests instead of the Anthropic SDK.
"""
from typing import Any, Final
from collections.abc import Mapping
from typing import Final
import httpx
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import verbose_logger
@ -16,6 +18,8 @@ from litellm.llms.anthropic.count_tokens.transformation import (
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
"""
@ -27,13 +31,14 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
async def handle_count_tokens_request(
self,
model: str,
messages: list[dict[str, Any]],
messages: list[dict[str, JsonValue]],
api_key: str,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, Any]] | None = None,
system: Any | None = None,
) -> dict[str, Any]:
tools: list[dict[str, JsonValue]] | None = None,
system: JsonValue = None,
optional_params: Mapping[str, JsonValue] | None = None,
) -> dict[str, JsonValue]:
"""
Handle a CountTokens request using httpx.
@ -52,7 +57,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
"""
try:
# Validate the request
self.validate_request(model, messages)
self.validate_request(model, messages, system=system, tools=tools)
verbose_logger.debug("Processing Anthropic CountTokens request for model: %s", model)
@ -62,6 +67,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
messages=messages,
tools=tools,
system=system,
optional_params=optional_params,
)
verbose_logger.debug("Transformed request: %s", request_body)
@ -97,7 +103,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
message=error_text,
)
anthropic_response: Final = response.json()
anthropic_response: Final = _COUNT_RESPONSE.validate_json(response.content)
verbose_logger.debug("Anthropic response: %s", anthropic_response)

56
litellm/llms/anthropic/count_tokens/transformation.py
View file

@ -4,10 +4,17 @@ Anthropic CountTokens API transformation logic.
This module handles the transformation of requests to Anthropic's CountTokens API format.
"""
from typing import Any, Final
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from pydantic import JsonValue, TypeAdapter
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
class AnthropicCountTokensConfig:
"""
@ -31,27 +38,31 @@ class AnthropicCountTokensConfig:
def transform_request_to_count_tokens(
self,
model: str,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
system: Any | None = None,
) -> dict[str, Any]:
messages: list[dict[str, JsonValue]],
tools: list[dict[str, JsonValue]] | None = None,
system: JsonValue = None,
optional_params: Mapping[str, JsonValue] | None = None,
) -> dict[str, JsonValue]: # mutable-ok: provider transport requires JSON dictionaries
"""
Transform request to Anthropic CountTokens format.
Includes optional system and tools fields for accurate token counting.
"""
request: Final[dict[str, Any]] = {
"model": model,
"messages": messages,
}
if system is not None:
request["system"] = system
if tools is not None:
request["tools"] = tools
return request
options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({})
return _COUNT_REQUEST.validate_python(
MappingProxyType(
{
"model": model,
"messages": messages,
**MappingProxyType(
{key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
),
**MappingProxyType(
{key: value for key, value in options.items() if key in COUNT_TOKEN_OPTION_NAMES}
),
}
)
)
def get_required_headers(self, api_key: str) -> dict[str, str]:
"""
@ -76,7 +87,14 @@ class AnthropicCountTokensConfig:
headers, _ = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
return headers
def validate_request(self, model: str, messages: list[dict[str, Any]]) -> None:
def validate_request(
self,
model: str,
messages: Sequence[Mapping[str, JsonValue]],
*,
system: JsonValue = None,
tools: list[dict[str, JsonValue]] | None = None,
) -> None:
"""
Validate the incoming count tokens request.
@ -90,7 +108,7 @@ class AnthropicCountTokensConfig:
if not model:
raise ValueError("model parameter is required")
if not messages:
if not messages and not system and not tools:
raise ValueError("messages parameter is required")
if not isinstance(messages, list):

2
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
View file

@ -1227,7 +1227,7 @@ class LiteLLMAnthropicMessagesAdapter:
self._add_system_message_to_messages(new_messages, anthropic_message_request)
new_kwargs: Final[ChatCompletionRequest] = {
"model": anthropic_message_request.get("model", ""),
"model": anthropic_message_request["model"],
"messages": new_messages,
}
## CONVERT METADATA (user_id + litellm metadata)

462
litellm/llms/anthropic/prompt_cache_prediction.py
View file

@ -1,10 +1,11 @@
from __future__ import annotations
import asyncio
import hashlib
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from dataclasses import dataclass, field
from itertools import accumulate, groupby
from types import MappingProxyType
from typing import Annotated, Final, Literal, Protocol, TypeAlias
@ -14,9 +15,14 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAda
import litellm
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
DEFAULT_ANTHROPIC_API_VERSION,
AnthropicMessagesConfig,
)
from litellm.types.router import LiteLLM_Params
from litellm.types.utils import ModelResponse
from litellm.utils import supports_thinking_cache_preservation
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_HEADERS: Final = TypeAdapter(dict[str, str])
@ -100,10 +106,7 @@ _Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminato
class _Message(_StrictModel):
role: Literal["user", "assistant"]
content: str | Annotated[tuple[_Block, ...], Field(strict=False)]
def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]:
return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content)
content: Annotated[str, Field(min_length=1, pattern=r"\S")] | Annotated[tuple[_Block, ...], Field(strict=False)]
class _Tool(_StrictModel):
@ -113,10 +116,7 @@ class _Tool(_StrictModel):
type: Literal["custom"] | None = None
class _Request(_StrictModel):
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
class _RequestOptions(_StrictModel):
model: str | None = None
max_tokens: int | None = None
stream: bool | None = None
@ -127,6 +127,289 @@ class _Request(_StrictModel):
metadata: Mapping[str, JsonValue] | None = None
class _Request(_RequestOptions):
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
class _Thinking(_StrictModel):
type: Literal["thinking"]
thinking: str
signature: str = Field(min_length=1)
_PlanBlock: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult | _Thinking, Field(discriminator="type")]
class _PlanMessage(_StrictModel):
role: Literal["user", "assistant", "system"]
content: str | Annotated[tuple[_PlanBlock, ...], Field(strict=False)]
class _PlanTool(_Tool):
cache_control: _CacheControl | None = None
class _PlanRequest(_RequestOptions):
messages: tuple[_PlanMessage, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_Text, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_PlanTool, ...], Field(strict=False)] | None = None
cache_control: _CacheControl | None = None
thinking: Mapping[str, JsonValue] | None = None
tool_choice: Mapping[str, JsonValue] | None = None
output_config: Mapping[str, JsonValue] | None = None
speed: Literal["fast", "standard"] | None = None
service_tier: Literal["auto", "standard_only"] | None = None
@dataclass(frozen=True, slots=True)
class CacheBoundary:
fingerprint: str
prefix_body: Mapping[str, JsonValue] = field(repr=False)
ttl_seconds: int
lookback_fingerprints: tuple[str, ...]
content_fingerprint: str = ""
lookback_content_fingerprints: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class PromptCachePlan:
full_body: Mapping[str, JsonValue] = field(repr=False)
breakpoints: tuple[CacheBoundary, ...]
@dataclass(frozen=True, slots=True)
class UnsupportedCachePlan:
reason: Literal[
"unsupported_prompt_shape",
"conflicting_cache_ttl",
"too_many_cache_breakpoints",
"invalid_cache_ttl_order",
"unsupported_thinking_cache_semantics",
"token_count_unavailable",
"inconsistent_prefix_token_count",
]
@dataclass(frozen=True, slots=True)
class CountedBreakpoint:
fingerprint: str
ttl_seconds: int
prefix_tokens: int
lookback_fingerprints: tuple[str, ...]
content_fingerprint: str = ""
lookback_content_fingerprints: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class CountedPromptCachePlan:
total_tokens: int
breakpoints: tuple[CountedBreakpoint, ...]
@dataclass(frozen=True, slots=True)
class _Position:
section: Literal["tools", "system", "messages"]
message_index: int
role: str
block: Mapping[str, JsonValue]
marker: _CacheControl | None
def _content_blocks(content: JsonValue) -> tuple[Mapping[str, JsonValue], ...]:
if isinstance(content, str):
return (MappingProxyType({"type": "text", "text": content}),)
return tuple(_JSON_OBJECT.validate_python(block) for block in content) if isinstance(content, list) else ()
def _position(
section: Literal["tools", "system", "messages"],
message_index: int,
role: str,
block: Mapping[str, JsonValue],
) -> _Position:
control: Final = block.get("cache_control")
return _Position(
section,
message_index,
role,
MappingProxyType({key: value for key, value in block.items() if key != "cache_control"}),
_CacheControl.model_validate(control) if control is not None else None,
)
def _positions(body: Mapping[str, JsonValue]) -> tuple[_Position, ...]:
tools: Final = body.get("tools")
messages: Final = body.get("messages")
return (
*tuple(
_position("tools", -1, "", _JSON_OBJECT.validate_python(tool))
for tool in (tools if isinstance(tools, list) else ())
),
*tuple(_position("system", -1, "", block) for block in _content_blocks(body.get("system"))),
*tuple(
_position("messages", message_index, str(message.get("role")), block)
for message_index, raw_message in enumerate(messages if isinstance(messages, list) else ())
for message in (_JSON_OBJECT.validate_python(raw_message),)
for block in _content_blocks(message.get("content"))
),
)
def _prefix_body(
body: Mapping[str, JsonValue],
positions: tuple[_Position, ...],
last_index: int,
) -> Mapping[str, JsonValue]:
prefix: Final = positions[: last_index + 1]
sections: Final = MappingProxyType(
{
section: _count_objects(tuple(position.block for position in prefix if position.section == section))
for section in ("tools", "system")
if any(position.section == section for position in prefix)
}
)
messages: Final = tuple(
MappingProxyType(
_JSON_OBJECT.validate_python(
MappingProxyType(
{"role": group[0].role, "content": _count_objects(tuple(position.block for position in group))}
)
)
)
for _, values in groupby(
(position for position in prefix if position.section == "messages"),
key=lambda position: position.message_index,
)
for group in (tuple(values),)
)
return MappingProxyType(
_JSON_OBJECT.validate_python(
MappingProxyType(
{
**MappingProxyType({key: body[key] for key in COUNT_TOKEN_OPTION_NAMES if key in body}),
**sections,
"messages": _count_objects(messages),
}
)
)
)
def _position_group(position: _Position, index: int) -> tuple[str, int, str | int]:
block_type: Final = position.block.get("type")
return (
position.section,
position.message_index,
block_type if isinstance(block_type, str) and block_type in ("tool_use", "tool_result") else index,
)
def _chain_digest(previous: str, current: str) -> str:
return _digest((previous, current))
def _cacheable_position(position: _Position) -> bool:
block_type: Final = position.block.get("type")
if block_type == "thinking":
return False
text: Final = position.block.get("text")
return block_type != "text" or (isinstance(text, str) and bool(text.strip()))
def _entry_fingerprint(fingerprint: str, ttl_seconds: int) -> str:
return _digest(("native-cache-prefix-v2", fingerprint, ttl_seconds))
def parse_cache_plan(body: Mapping[str, JsonValue]) -> PromptCachePlan | UnsupportedCachePlan:
try:
request: Final = _PlanRequest.model_validate(body)
positions: Final = _positions(body)
except ValidationError:
return UnsupportedCachePlan("unsupported_prompt_shape")
explicit: Final = tuple(
(index, position.marker) for index, position in enumerate(positions) if position.marker is not None
)
automatic_index: Final = next(
(index for index in reversed(range(len(positions))) if _cacheable_position(positions[index])), None
)
automatic_existing: Final = next((marker for index, marker in explicit if index == automatic_index), None)
if (
request.cache_control is not None
and automatic_existing is not None
and automatic_existing != request.cache_control
):
return UnsupportedCachePlan("conflicting_cache_ttl")
automatic: Final = (
((automatic_index, request.cache_control),)
if (request.cache_control is not None and automatic_index is not None and automatic_existing is None)
else ()
)
markers: Final = tuple(sorted((*explicit, *automatic), key=lambda value: value[0]))
if len(markers) > 4:
return UnsupportedCachePlan("too_many_cache_breakpoints")
ttls: Final = tuple(3600 if marker.ttl == "1h" else 300 for _, marker in markers)
if any(first < second for first, second in zip(ttls, ttls[1:])):
return UnsupportedCachePlan("invalid_cache_ttl_order")
settings: Final = MappingProxyType(
{
key: body[key]
for key in ("thinking", "output_config", "speed")
if key in body and not (key == "speed" and body[key] == "standard")
}
)
hashes: Final = tuple(
accumulate(
(
_digest(
(
position.section,
position.message_index,
position.role,
position.block,
body.get("tool_choice") if position.section == "messages" else None,
)
)
for position in positions
),
_chain_digest,
initial=_digest(settings),
)
)[1:]
groups: Final = tuple(
tuple(index for index, _ in values)
for _, values in groupby(
enumerate(positions),
key=lambda item: _position_group(item[1], item[0]),
)
)
return PromptCachePlan(
full_body=MappingProxyType(dict(body)),
breakpoints=tuple(
CacheBoundary(
fingerprint=_entry_fingerprint(hashes[index], ttl),
prefix_body=_prefix_body(body, positions, index),
ttl_seconds=ttl,
lookback_fingerprints=tuple(
_entry_fingerprint(hashes[earlier], ttl)
for group in reversed(tuple(group for group in groups if group[0] <= index)[-20:])
for earlier in reversed(group)
if earlier <= index
),
content_fingerprint=hashes[index],
lookback_content_fingerprints=tuple(
hashes[earlier]
for group in reversed(tuple(group for group in groups if group[0] <= index)[-20:])
for earlier in reversed(group)
if earlier <= index
),
)
for (index, _), ttl in zip(markers, ttls)
),
)
@dataclass(frozen=True, slots=True)
class PromptPrefix:
prefix_body: Mapping[str, JsonValue]
@ -137,68 +420,28 @@ class PromptPrefix:
def _digest(value: object) -> str:
return hashlib.sha256(
json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()
json.dumps(value, default=_json_object, separators=(",", ":"), ensure_ascii=False).encode()
).hexdigest()
def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str:
return _digest((previous, boundary))
def _json_object(value: object) -> dict[str, JsonValue]: # mutable-ok: JSON serialization requires a dictionary
return _JSON_OBJECT.validate_python(value)
def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None:
try:
request: Final = _Request.model_validate(body)
blocks: Final = tuple(message.blocks() for message in request.messages)
_Request.model_validate(body)
except ValidationError:
return None
markers: Final = tuple(
(message_index, block_index, block.cache_control)
for message_index, message_blocks in enumerate(blocks)
for block_index, block in enumerate(message_blocks)
if block.cache_control is not None
)
if len(markers) != 1:
plan: Final = parse_cache_plan(body)
if isinstance(plan, UnsupportedCachePlan) or len(plan.breakpoints) != 1:
return None
message_end, block_end, marker = markers[0]
normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True))
context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized})
boundaries: Final = tuple(
(
message_index,
request.messages[message_index].role,
_JSON_OBJECT.validate_python(
block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True)
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
)
hashes: Final = tuple(
accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl)))
)[1:]
prefix_messages: Final = tuple(
_Message(
role=request.messages[message_index].role,
content=tuple(
block
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
)
prefix: Final = plan.breakpoints[0]
return PromptPrefix(
prefix_body=MappingProxyType(
_JSON_OBJECT.validate_python(
_Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump(
mode="json", exclude_none=True
)
)
),
fingerprint=hashes[-1],
fingerprints=tuple(reversed(hashes[-20:])),
ttl_seconds=3600 if marker.ttl == "1h" else 300,
prefix_body=prefix.prefix_body,
fingerprint=prefix.fingerprint,
fingerprints=prefix.lookback_fingerprints,
ttl_seconds=prefix.ttl_seconds,
)
@ -246,6 +489,9 @@ class _CountBody(BaseModel):
messages: Sequence[Mapping[str, JsonValue]]
tools: Sequence[Mapping[str, JsonValue]] | None = None
system: str | Sequence[Mapping[str, JsonValue]] | None = None
thinking: Mapping[str, JsonValue] | None = None
tool_choice: Mapping[str, JsonValue] | None = None
output_config: Mapping[str, JsonValue] | None = None
class _CountResult(BaseModel):
@ -262,16 +508,36 @@ def _count_objects(
return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary
async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
native: Final = _CountBody.model_validate(body)
def _messages_url(model: str, api_key: str, api_base: str | None) -> str:
return AnthropicMessagesConfig().get_complete_url( # pyright: ignore[reportUnknownMemberType] # canonical native URL owner takes legacy JSON arguments
api_base=api_base,
api_key=api_key,
model=model,
optional_params=_JSON_OBJECT.validate_python(MappingProxyType({})),
litellm_params=_JSON_OBJECT.validate_python(MappingProxyType({})),
)
async def count_prompt_tokens(
model: str,
api_key: str,
body: Mapping[str, JsonValue],
api_base: str | None = None,
) -> int | None:
try:
native: Final = _CountBody.model_validate(body)
count_url: Final = _messages_url(model, api_key, api_base) + "/count_tokens"
result: Final = _CountResult.model_validate(
await _counter.handle_count_tokens_request(
model=model,
messages=_count_objects(native.messages),
tools=_count_objects(native.tools) if native.tools is not None else None,
system=native.system,
system=_JSON_OBJECT.validate_python(MappingProxyType({"system": native.system}))["system"],
api_key=api_key,
api_base=count_url,
optional_params=_JSON_OBJECT.validate_python(
MappingProxyType({key: body[key] for key in COUNT_TOKEN_OPTION_NAMES if key in body})
),
timeout=15.0,
)
)
@ -280,10 +546,55 @@ async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonV
return result.input_tokens
async def count_cache_plan(
model: str,
api_key: str,
plan: PromptCachePlan,
token_counter: TokenCounter = count_prompt_tokens,
) -> CountedPromptCachePlan | UnsupportedCachePlan:
if any(position.block.get("type") == "thinking" for position in _positions(plan.full_body)):
if not supports_thinking_cache_preservation(model, "anthropic"):
return UnsupportedCachePlan("unsupported_thinking_cache_semantics")
total: Final = await token_counter(model, api_key, plan.full_body)
if total is None:
return UnsupportedCachePlan("token_count_unavailable")
counts: Final = tuple(
await asyncio.gather(*(token_counter(model, api_key, marker.prefix_body) for marker in plan.breakpoints))
)
if any(value is None for value in counts):
return UnsupportedCachePlan("token_count_unavailable")
known: Final = tuple(value for value in counts if value is not None)
if any(value < 0 for value in (total, *known)) or any(
first > second for first, second in zip(known, (*known[1:], total))
):
return UnsupportedCachePlan("inconsistent_prefix_token_count")
return CountedPromptCachePlan(
total,
tuple(
CountedBreakpoint(
marker.fingerprint,
marker.ttl_seconds,
count,
marker.lookback_fingerprints,
marker.content_fingerprint,
marker.lookback_content_fingerprints,
)
for marker, count in zip(plan.breakpoints, known)
),
)
@dataclass(frozen=True, slots=True)
class NativePredictionTarget:
model: str
api_key: str
api_key: str = field(repr=False)
api_base: str | None = None
def supported_baseline_recipient(target: NativePredictionTarget, wire: httpx.Request) -> bool:
return wire.headers.get("x-api-key") == target.api_key and wire.url == httpx.URL(
_messages_url(target.model, target.api_key, target.api_base)
)
@dataclass(frozen=True, slots=True)
@ -297,11 +608,26 @@ class UnsupportedPredictionTarget:
def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
return _resolve_prediction_target(params, allow_configured_endpoint=False)
def resolve_baseline_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
return _resolve_prediction_target(params, allow_configured_endpoint=True)
def _resolve_prediction_target(
params: LiteLLM_Params,
*,
allow_configured_endpoint: bool,
) -> NativePredictionTarget | UnsupportedPredictionTarget:
configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True))
if configured_options - _DEPLOYMENT_OPTIONS:
return UnsupportedPredictionTarget("unsupported_deployment_configuration")
api_base: Final = AnthropicModelInfo.get_api_base(params.api_base)
if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"):
if not allow_configured_endpoint and api_base not in (
"https://api.anthropic.com",
"https://api.anthropic.com/v1/messages",
):
return UnsupportedPredictionTarget("unsupported_provider_endpoint")
try:
model, provider, _, _ = litellm.get_llm_provider(
@ -314,7 +640,7 @@ def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget
api_key: Final = AnthropicModelInfo.get_api_key(params.api_key)
if api_key is None or not _supported_provider_key(api_key):
return UnsupportedPredictionTarget("unsupported_provider_credentials")
return NativePredictionTarget(model=model, api_key=api_key)
return NativePredictionTarget(model=model, api_key=api_key, api_base=api_base)
def _supported_provider_key(api_key: str) -> bool:

9
litellm/llms/azure/chat/gpt_transformation.py
View file

@ -280,10 +280,17 @@ class AzureOpenAIConfig(BaseConfig):
ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages)
azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages))
request_params: Final = MappingProxyType(
{
key: value
for key, value in optional_params.items()
if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions")
}
)
return {
"model": model,
"messages": azure_messages,
**optional_params,
**request_params,
**sanitized_tools_update(optional_params),
}

40
litellm/llms/base_llm/files/litellm_db_storage_backend.py Normal file
View file

@ -0,0 +1,40 @@
from typing import TYPE_CHECKING, Final
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.repositories.managed_file_content_repository import ManagedFileContentRepository
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
LITELLM_DB_STORAGE_BACKEND_NAME: Final = "litellm_db"
LITELLM_DB_STORAGE_URL_PREFIX: Final = f"{LITELLM_DB_STORAGE_BACKEND_NAME}://"
def storage_url_to_row_id(storage_url: str) -> str:
if not storage_url.startswith(LITELLM_DB_STORAGE_URL_PREFIX):
raise ValueError(f"Not a {LITELLM_DB_STORAGE_BACKEND_NAME} storage url: {storage_url}")
return storage_url.removeprefix(LITELLM_DB_STORAGE_URL_PREFIX)
class LiteLLMDbStorageBackend(BaseFileStorageBackend):
def __init__(self, prisma_client: "PrismaClient") -> None:
self._contents = ManagedFileContentRepository(prisma_client)
async def upload_file(
self,
file_content: bytes,
filename: str,
content_type: str,
path_prefix: str | None = None,
file_naming_strategy: str = "uuid",
) -> str:
return f"{LITELLM_DB_STORAGE_URL_PREFIX}{await self._contents.store(file_content)}"
async def download_file(self, storage_url: str) -> bytes:
content: Final = await self._contents.load(storage_url_to_row_id(storage_url))
if content is None:
raise ValueError(f"No stored file content for {storage_url}")
return content
async def delete_file(self, storage_url: str) -> None:
await self._contents.delete(storage_url_to_row_id(storage_url))

26
litellm/llms/base_llm/files/storage_backend_factory.py
View file

@ -6,32 +6,46 @@ based on the backend type. Backends use the same configuration as their correspo
callbacks (e.g., azure_storage uses the same env vars as AzureBlobStorageLogger).
"""
from typing import TYPE_CHECKING
from litellm._logging import verbose_logger
from .azure_blob_storage_backend import AzureBlobStorageBackend
from .litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME, LiteLLMDbStorageBackend
from .storage_backend import BaseFileStorageBackend
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
def get_storage_backend(backend_type: str) -> BaseFileStorageBackend:
def get_storage_backend(backend_type: str, prisma_client: "PrismaClient | None" = None) -> BaseFileStorageBackend:
"""
Factory function to create a storage backend instance.
Backends are configured using the same environment variables as their
corresponding callbacks. For example, "azure_storage" uses the same
env vars as AzureBlobStorageLogger.
env vars as AzureBlobStorageLogger. "litellm_db" stores file bytes in the
proxy's own database and needs the connected Prisma client.
Args:
backend_type: Backend type identifier (e.g., "azure_storage")
backend_type: Backend type identifier (e.g., "azure_storage", "litellm_db")
prisma_client: The proxy's database client, required by "litellm_db"
Returns:
BaseFileStorageBackend: Instance of the appropriate storage backend
Raises:
ValueError: If backend_type is not supported
ValueError: If backend_type is not supported, or "litellm_db" is asked for without a database
"""
verbose_logger.debug("Creating storage backend: type=%s", backend_type)
if backend_type == "azure_storage":
return AzureBlobStorageBackend()
else:
raise ValueError(f"Unsupported storage backend type: {backend_type}. Supported types: azure_storage")
if backend_type == LITELLM_DB_STORAGE_BACKEND_NAME:
if prisma_client is None:
raise ValueError(f"Storage backend {LITELLM_DB_STORAGE_BACKEND_NAME} requires a database-connected proxy")
return LiteLLMDbStorageBackend(prisma_client)
raise ValueError(
f"Unsupported storage backend type: {backend_type}. "
f"Supported types: azure_storage, {LITELLM_DB_STORAGE_BACKEND_NAME}"
)

93
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -1,17 +1,8 @@
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
request_tools,
response_assistant_turn,
scoped_structured_message_indices,
)
if TYPE_CHECKING:
from fastapi import HTTPException
@ -21,43 +12,7 @@ if TYPE_CHECKING:
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.utils import GenericGuardrailAPIInputs
@dataclass(frozen=True, slots=True)
class RequestScanContext:
"""The scoped request turns and tool definitions a guardrail's request scan sees, in OpenAI chat shape."""
structured_messages: tuple["AllMessageValues", ...] = ()
tools: tuple["ChatCompletionToolParam", ...] = ()
conversation_supplied: bool = False
@staticmethod
def scoped(
structured_messages: Sequence["AllMessageValues"],
tools: Sequence["ChatCompletionToolParam"],
guardrail_to_apply: "CustomGuardrail",
*,
skip_system: bool | None = None,
) -> "RequestScanContext":
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
scoped_indices: Final = scoped_structured_message_indices(
structured_messages,
scan_only_tool_results=scan_only_tool_results,
skip_system=(
effective_skip_system_message_for_guardrail(guardrail_to_apply) if skip_system is None else skip_system
),
skip_tool=effective_skip_tool_message_for_guardrail(guardrail_to_apply),
)
return RequestScanContext(
structured_messages=tuple(structured_messages[index] for index in scoped_indices),
tools=() if scan_only_tool_results else tuple(tools),
conversation_supplied=bool(structured_messages),
)
REQUEST_SCAN_CONTEXT_KEY: Final = "litellm_request_scan_context"
from litellm.types.llms.openai import AllMessageValues
@dataclass(slots=True)
@ -302,50 +257,6 @@ class BaseTranslation(ABC):
"""
return None
def request_scan_context(
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
) -> RequestScanContext:
"""Override wherever ``process_input_messages`` scopes or translates the request differently."""
structured_messages: Final = self.get_structured_messages(
dict(data) # mutable-ok: get_structured_messages takes the request as a dict
)
return RequestScanContext.scoped(
structured_messages or (), request_tools(data.get("tools")), guardrail_to_apply
)
def with_response_context(
self,
inputs: "GenericGuardrailAPIInputs",
request_data: Mapping[str, object] | None,
guardrail_to_apply: "CustomGuardrail",
) -> "GenericGuardrailAPIInputs":
"""``inputs`` plus the scoped request conversation, closed by the scanned reply, and the request tools."""
if request_data is None:
return inputs
precomputed: Final = request_data.get(REQUEST_SCAN_CONTEXT_KEY)
context: Final = (
precomputed
if isinstance(precomputed, RequestScanContext)
else self.request_scan_context(request_data, guardrail_to_apply)
)
if not context.conversation_supplied:
return inputs
assistant_turn: Final = response_assistant_turn(inputs.get("texts") or (), inputs.get("tool_calls") or ())
contextual_inputs: Final[GenericGuardrailAPIInputs] = {
**inputs,
"structured_messages": [ # mutable-ok: GenericGuardrailAPIInputs fields are lists
*context.structured_messages,
*(() if assistant_turn is None else (assistant_turn,)),
],
}
if not context.tools:
return contextual_inputs
with_tools: Final[GenericGuardrailAPIInputs] = {
**contextual_inputs,
"tools": list(context.tools), # mutable-ok: GenericGuardrailAPIInputs fields are lists
}
return with_tools
def extract_request_tool_names(self, data: dict) -> list[str]:
"""
Extract tool names from the request body for allowlist/policy checks.

64
litellm/llms/base_llm/guardrail_translation/utils.py
View file

@ -2,24 +2,12 @@ from __future__ import annotations
import json
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
from pydantic import BaseModel
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionTextObject,
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolParam,
ResponseAPIUsage,
)
if TYPE_CHECKING:
from litellm.types.utils import ChatCompletionMessageToolCall
from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
def _anthropic_stream_chunk_events(item: object) -> list[dict]:
@ -290,57 +278,9 @@ def scoped_structured_message_indices(
)
def _assistant_tool_call(
tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall,
) -> ChatCompletionAssistantToolCall:
function: Final = stream_item_field(tool_call, "function")
tool_call_id: Final = stream_item_field(tool_call, "id")
name: Final = stream_item_field(function, "name")
arguments: Final = stream_item_field(function, "arguments")
return ChatCompletionAssistantToolCall(
id=tool_call_id if isinstance(tool_call_id, str) else None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=name if isinstance(name, str) else None,
arguments=arguments if isinstance(arguments, str) else "",
),
)
def response_assistant_turn(
texts: Sequence[str],
tool_calls: Sequence[ChatCompletionToolCallChunk] | Sequence[ChatCompletionMessageToolCall],
) -> ChatCompletionAssistantMessage | None:
"""The scanned reply as the assistant turn closing the request conversation."""
assistant_tool_calls: Final = tuple(_assistant_tool_call(tool_call) for tool_call in tool_calls)
if not texts and not assistant_tool_calls:
return None
content: Final = (
texts[0]
if len(texts) == 1
else tuple(ChatCompletionTextObject(type="text", text=text) for text in texts) or None
)
if not assistant_tool_calls:
return ChatCompletionAssistantMessage(role="assistant", content=content)
return ChatCompletionAssistantMessage(
role="assistant",
content=content,
tool_calls=list(assistant_tool_calls), # mutable-ok: the assistant message type takes a list
)
ToolT = TypeVar("ToolT")
def request_tools(raw_tools: object) -> tuple[ChatCompletionToolParam, ...]:
"""The request's ``tools`` list, as the chat completion request model already validated it upstream."""
if not isinstance(raw_tools, list):
return ()
return tuple(
cast(Sequence[ChatCompletionToolParam], raw_tools) # cast-ok: the request model validated tools upstream
)
def openai_tool_name(tool: object) -> str | None:
if not isinstance(tool, dict):
return None

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -2163,6 +2163,8 @@ class BaseLLMHTTPHandler:
e=e, litellm_params=litellm_params_dict
)
if should_retry and not hit_max_attempt:
if logging_obj.baseline_cache_context is not None:
await logging_obj.invalidate_baseline_cache_estimate("retried_request")
verbose_logger.debug(
"Anthropic /v1/messages: invalid thinking signature; "
"stripping thinking blocks and retrying (attempt %s/%s).",

6
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -453,7 +453,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["model"] = response.model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -616,7 +616,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if responses_so_far and hasattr(responses_so_far[0], "model") and responses_so_far[0].model:
inputs["model"] = responses_so_far[0].model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -797,7 +797,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if responses_so_far and getattr(responses_so_far[0], "model", None):
inputs["model"] = responses_so_far[0].model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,

31
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -48,7 +48,6 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
)
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
RequestScanContext,
StreamingScanKey,
StreamTransformSink,
)
@ -453,28 +452,6 @@ class OpenAIResponsesHandler(BaseTranslation):
)
return cast(list[AllMessageValues], messages) if messages else None
def request_scan_context(
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
) -> RequestScanContext:
raw_tools: Final = data.get("tools")
structured_messages: Final = tuple(
self.get_structured_messages(
dict(data) # mutable-ok: get_structured_messages takes the request as a dict
)
or ()
)
return RequestScanContext(
structured_messages=structured_messages,
tools=tuple(
cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list
for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(
tuple(raw_tools) if isinstance(raw_tools, list) else ()
)
for tool in form.chat_tools
),
conversation_supplied=bool(structured_messages),
)
async def process_input_messages(
self,
data: dict,
@ -778,7 +755,7 @@ class OpenAIResponsesHandler(BaseTranslation):
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -892,7 +869,7 @@ class OpenAIResponsesHandler(BaseTranslation):
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
inputs=inputs,
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -951,7 +928,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if hasattr(model_response_stream, "model") and model_response_stream.model:
inputs["model"] = model_response_stream.model
await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
inputs=inputs,
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,
@ -973,7 +950,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if response_model:
fallback_inputs["model"] = response_model
fallback_outputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=self.with_response_context(fallback_inputs, request_data, guardrail_to_apply),
inputs=fallback_inputs,
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,

482
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

15
litellm/models/autorouter_session.py
View file

@ -8,6 +8,8 @@ maintains per (api_key, session_id, router_name).
from collections.abc import Mapping
from datetime import datetime
from pydantic import Field
from litellm.types.llms.base import LiteLLMPydanticObjectBase
@ -22,18 +24,25 @@ class LiteLLM_AutoRouterSession(LiteLLMPydanticObjectBase):
turns: int
spend: float
saved_spend: float
savings_estimated_turns: int = 0
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
savings_estimated_baseline_models: Mapping[str, int] = Field(default_factory=dict)
classifier_cost: float
tier_turns: Mapping[str, int]
baseline_models: Mapping[str, int]
@property
def baseline_model(self) -> str | None:
"""The baseline most of this session's turns were priced against, or None when no turn recorded one.
"""The baseline most covered turns were priced against, or None when none were estimated.
A router reconfigured mid-session leaves turns priced against two baselines; the row keeps both
counts, and the label is the one that priced the most money-carrying turns rather than whatever the
router is configured with now.
"""
if not self.baseline_models:
if not self.savings_estimated_baseline_models:
return None
return max(self.baseline_models, key=lambda model: (self.baseline_models[model], model))
return max(
self.savings_estimated_baseline_models,
key=lambda model: (self.savings_estimated_baseline_models[model], model),
)

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": [
{

6
litellm/proxy/_types.py
View file

@ -13,6 +13,7 @@ from pydantic import (
ConfigDict,
Field,
Json,
JsonValue,
PositiveInt,
field_validator,
model_validator,
@ -123,6 +124,7 @@ class SupportedDBObjectType(str, enum.Enum):
MODEL_COST_MAP = "model_cost_map"
TOOLS = "tools"
CONFIG_OVERRIDES = "config_overrides"
WEBSEARCH_INTERCEPTION_SETTINGS = "websearch_interception_settings"
def __str__(self):
return str(self.value)
@ -3940,6 +3942,7 @@ class SpendLogsRouterMetadata(TypedDict):
class SpendLogsMetadata(TypedDict):
autorouter_baseline_observation: ReadOnly[str | None]
"""
Specific metadata k,v pairs logged to spendlogs for easier cost tracking
"""
@ -3980,7 +3983,8 @@ class SpendLogsMetadata(TypedDict):
original_model_group: ReadOnly[str | None] # Model group requested before any fallbacks
cost_breakdown: CostBreakdown | None # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.)
compression_savings: CompressionSavingsMetadata | None
autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed
autorouter_savings: ReadOnly[float | None]
autorouter_savings_estimate: ReadOnly[Mapping[str, JsonValue] | None]
litellm_gateway_injected_cache: ReadOnly[str | None]
router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model
azure_spillover: ReadOnly[AzureSpillover | None] # None = Azure did not report spillover

60
litellm/proxy/auth/auth_checks.py
View file

@ -4012,6 +4012,64 @@ async def get_org_object(
return _org_obj
def _last_known_org_cache_key(org_id: str) -> str:
return f"org_id:{org_id}:with_budget:last_known"
async def _keep_last_known_org(
org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache
) -> None:
cache_key: Final = _last_known_org_cache_key(org_id)
held_locally: Final = await user_api_key_cache.async_get_cache(
key=cache_key, local_only=True, model_type=LiteLLM_OrganizationTable
)
if held_locally is not None:
return
await user_api_key_cache.async_set_cache(
key=cache_key,
value=org,
model_type=LiteLLM_OrganizationTable,
ttl=get_management_object_ttl(user_api_key_cache),
)
async def get_org_object_for_request(
org_id: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> LiteLLM_OrganizationTable | None:
try:
org: Final = await get_org_object(
org_id=org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
include_budget_table=True,
)
except OrganizationNotFoundError:
return None
except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits
if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e):
verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True)
return None
last_known_org: Final = await user_api_key_cache.async_get_cache(
key=_last_known_org_cache_key(org_id),
model_type=LiteLLM_OrganizationTable,
)
if last_known_org is not None:
return last_known_org
if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable():
return None
raise
if org is None:
return None
await _keep_last_known_org(org, org_id, user_api_key_cache)
return org
async def _get_resources_from_access_groups(
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
@ -5680,7 +5738,7 @@ async def _project_max_budget_check(
if project_object.litellm_budget_table is not None:
max_budget = project_object.litellm_budget_table.max_budget
if max_budget is None or max_budget <= 0 or not math.isfinite(max_budget):
if max_budget is None or not math.isfinite(max_budget):
return
from litellm.proxy.proxy_server import get_current_spend

52
litellm/proxy/auth/user_api_key_auth.py
View file

@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_checks import (
get_jwt_key_mapping_object,
get_key_end_user_budget_id,
get_object_permission,
get_org_object_for_request,
get_project_object,
get_team_membership,
get_team_object,
@ -2611,6 +2612,47 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc
return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
async def _inherit_org_identity(
user_api_key_auth_obj: UserAPIKeyAuth,
team_object: LiteLLM_TeamTableCachedObj | None,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> None:
if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
user_api_key_auth_obj.org_id = team_object.organization_id
already_populated: Final = any(
value is not None
for value in (
user_api_key_auth_obj.organization_alias,
user_api_key_auth_obj.organization_max_budget,
user_api_key_auth_obj.organization_tpm_limit,
user_api_key_auth_obj.organization_rpm_limit,
user_api_key_auth_obj.organization_metadata,
)
)
if user_api_key_auth_obj.org_id is None or already_populated or prisma_client is None:
return
org_object: Final = await get_org_object_for_request(
org_id=user_api_key_auth_obj.org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if org_object is None:
return
user_api_key_auth_obj.organization_alias = org_object.organization_alias
user_api_key_auth_obj.organization_metadata = org_object.metadata
budget: Final = org_object.litellm_budget_table
if budget is None:
return
user_api_key_auth_obj.organization_max_budget = budget.max_budget
user_api_key_auth_obj.organization_tpm_limit = budget.tpm_limit
user_api_key_auth_obj.organization_rpm_limit = budget.rpm_limit
def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, object]) -> bool:
return master_key is None and not any(
general_settings.get(flag, False)
@ -2849,8 +2891,14 @@ async def _run_centralized_common_checks(
user_object=user_object,
)
if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
user_api_key_auth_obj.org_id = team_object.organization_id
await _inherit_org_identity(
user_api_key_auth_obj=user_api_key_auth_obj,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
# common_checks identifies admin via user_object, not the token
# (non_proxy_admin_allowed_routes_check). JWT admin shortcut and

177
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -7,10 +7,12 @@
import asyncio
import os
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -18,6 +20,15 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.batches_endpoints.litellm_executed_batches import (
LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE,
LiteLLMExecutedBatchRunner,
ManagedBatchStore,
batch_error,
executed_batch_runner_lost,
litellm_executed_provider_for,
resolve_litellm_executed_provider,
)
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
@ -47,16 +58,87 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_model_id_from_unified_batch_id,
get_models_from_unified_file_id,
get_original_file_id,
is_litellm_executed_batch,
prepare_data_with_credentials,
update_batch_in_database,
validate_managed_id_requirement,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_metadata
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy, is_known_model
from litellm.repositories.managed_batch_repository import ManagedBatchRepository
from litellm.repositories.table_repositories import ManagedFileRepository
from litellm.router import Router
from litellm.types.llms.openai import LiteLLMBatchCreateRequest
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from prisma.models import LiteLLM_ManagedObjectTable
router: Final = APIRouter()
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None:
metadata: Final = data.get("litellm_metadata")
if metadata is None:
return None
return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata))
def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyLogging) -> LiteLLMExecutedBatchRunner:
from litellm.proxy.proxy_server import general_settings, prisma_client
managed_files: Final = proxy_logging_obj.get_proxy_hook("managed_files")
if prisma_client is None or not isinstance(managed_files, ManagedBatchStore):
raise batch_error(
400,
"LiteLLM-executed batches need a database: set DATABASE_URL so LiteLLM can keep the batch and its files",
)
return LiteLLMExecutedBatchRunner(
llm_router=llm_router,
prisma_client=prisma_client,
managed_files=managed_files,
batches=ManagedBatchRepository(prisma_client),
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
)
async def _batch_from_database(
batch_id: str,
unified_batch_id: str | Literal[False],
executed_batch: bool,
managed_files_obj: object,
prisma_client: PrismaClient | None,
llm_router: Router | None,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple["LiteLLM_ManagedObjectTable | None", LiteLLMBatch | None]:
row, batch = await get_batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
)
updated_at: Final[object] = getattr(row, "updated_at", None)
if not executed_batch or batch is None or llm_router is None or not isinstance(updated_at, datetime):
return row, batch
if not executed_batch_runner_lost(batch.status, updated_at):
return row, batch
runner: Final = _litellm_executed_batch_runner(llm_router, proxy_logging_obj)
return row, await runner.fail_abandoned(batch, user_api_key_dict)
async def _raise_when_input_file_must_be_managed(model: str, credentials: Mapping[str, object]) -> None:
if await litellm_executed_provider_for(credentials) is None:
return
raise batch_error(
400,
f"Batches for {model} run inside LiteLLM, so the input file must be a LiteLLM managed file: "
f"{LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE}",
)
def _raise_not_found_when_openai_fallback_unservable(
@ -101,6 +183,24 @@ async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str |
return db_file.storage_url or None
async def _create_provider_batch_for_managed_file(
llm_router: Router,
create_batch_data: LiteLLMBatchCreateRequest,
input_file_id: str,
unified_file_id: str,
) -> LiteLLMBatch:
resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id)
request: Final[LiteLLMBatchCreateRequest] = {
**create_batch_data,
"input_file_id": resolved_storage_url or input_file_id,
"disable_fallbacks": True,
}
response: Final = await llm_router.acreate_batch(**request)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
return response
@router.post(
"/{provider}/v1/batches",
dependencies=[Depends(user_api_key_auth)],
@ -296,24 +396,35 @@ async def create_batch(
await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
_create_batch_data["model"] = model
resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id)
if resolved_storage_url is not None:
_create_batch_data["input_file_id"] = resolved_storage_url
if llm_router is None:
raise HTTPException(
status_code=500,
detail={"error": "LLM Router not initialized. Ensure models added to proxy."},
)
_create_batch_data.update(disable_fallbacks=True) # pyright: ignore[reportCallIssue] # router flag
response = await llm_router.acreate_batch(**_create_batch_data)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
executed_provider: Final = await resolve_litellm_executed_provider(
llm_router, model, user_api_key_dict.team_id
)
response = (
await _litellm_executed_batch_runner(llm_router, proxy_logging_obj).create(
create_request=_create_batch_data,
unified_input_file_id=input_file_id,
model=model,
provider=executed_provider,
user_api_key_dict=user_api_key_dict,
request_tags=_request_tags(_create_batch_data),
)
if executed_provider is not None
else await _create_provider_batch_for_managed_file(
llm_router, _create_batch_data, input_file_id, unified_file_id
)
)
else:
# Check if model specified via header/query/body param
model_param: Final = (
data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model")
_create_batch_data.get("model")
or request.query_params.get("model")
or request.headers.get("x-litellm-model")
)
# SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback
@ -325,6 +436,7 @@ async def create_batch(
user_api_key_dict=user_api_key_dict,
operation_context="batch creation",
)
await _raise_when_input_file_must_be_managed(model_param, credentials)
prepare_data_with_credentials(
data=_create_batch_data,
@ -486,23 +598,26 @@ async def retrieve_batch(
managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files")
from litellm.proxy.proxy_server import prisma_client
db_batch_object, response = await get_batch_from_database(
executed_batch: Final = isinstance(unified_batch_id, str) and is_litellm_executed_batch(unified_batch_id)
db_batch_object, response = await _batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
executed_batch=executed_batch,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
)
if executed_batch and response is None:
raise batch_error(404, f"No batch found with id '{batch_id}'.")
# If batch is in a terminal state, return immediately.
# Include "complete" (DB-normalized form of "completed").
if response is not None and response.status in [
"completed",
"complete",
"failed",
"cancelled",
"expired",
]:
if response is not None and (
response.status in ("completed", "complete", "failed", "cancelled", "expired") or executed_batch
):
# Call hooks and return
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
@ -978,6 +1093,17 @@ async def cancel_batch(
proxy_config=proxy_config,
)
unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) if unified_batch_id else None
if unified_model_id is not None:
resolved_unified_model: Final = (
llm_router.resolve_model_name_from_model_id(unified_model_id) if llm_router is not None else None
)
await authorize_model_for_key(
model_id=resolved_unified_model or unified_model_id,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
)
# SCENARIO 1: Batch ID is encoded with model info
if model_from_id is not None:
credentials: Final = await get_authorized_credentials_for_model(
@ -1009,6 +1135,12 @@ async def cancel_batch(
)
# SCENARIO 2: target_model_names based routing
elif unified_batch_id and is_litellm_executed_batch(unified_batch_id):
if llm_router is None:
raise batch_error(500, "LLM Router not initialized. Ensure models added to proxy.")
response = await _litellm_executed_batch_runner( # rebind-ok: each cancel path sets the route's response
llm_router, proxy_logging_obj
).cancel(batch_id, user_api_key_dict)
elif unified_batch_id:
if llm_router is None:
raise HTTPException(
@ -1022,11 +1154,6 @@ async def cancel_batch(
status_code=400,
detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."},
)
await authorize_model_for_key(
model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
)
data["model"] = model_id_from_batch
data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id)
response = await llm_router.acancel_batch(**data)

716
litellm/proxy/batches_endpoints/litellm_executed_batches.py Normal file
View file

@ -0,0 +1,716 @@
import asyncio
import json
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import pairwise
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, runtime_checkable
import httpx
from openai.types.batch import Errors
from openai.types.batch_error import BatchError
from openai.types.batch_request_counts import BatchRequestCounts
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid as uuid_module
from litellm.constants import LITELLM_EXECUTED_BATCH_CONCURRENCY
from litellm.integrations.prometheus import PrometheusLogger
from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import is_request_body_safe
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.openai_files_endpoints.common_utils import LITELLM_EXECUTED_BATCH_ID_PREFIX
from litellm.proxy.openai_files_endpoints.storage_backend_service import StorageBackendFileService
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.managed_batch_repository import ManagedBatchRepository
from litellm.types.llms.openai import LiteLLMBatchCreateRequest, OpenAIFileObject, OpenAIFilesPurpose
from litellm.types.utils import LITELLM_EXECUTED_BATCH_PROVIDERS, ExtractedFileData, LiteLLMBatch, LlmProviders
if TYPE_CHECKING:
from prisma import types as prisma_types
from litellm.router import Router
BatchEndpoint: TypeAlias = Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"]
BatchStatus: TypeAlias = Literal[
"in_progress", "finalizing", "completed", "failed", "cancelling", "cancelled", "expired"
]
TERMINAL_BATCH_STATUSES: Final[frozenset[str]] = frozenset({"completed", "failed", "cancelled", "expired"})
_STOP_STATUSES: Final[frozenset[str]] = TERMINAL_BATCH_STATUSES | frozenset({"cancelling"})
_BATCH_ENDPOINT_ADAPTER: Final[TypeAdapter[BatchEndpoint]] = TypeAdapter(BatchEndpoint)
_CANCEL_POLL_SECONDS: Final = 1.0
_HEARTBEAT_SECONDS: Final = 30.0
_STALE_AFTER_SECONDS: Final = 180.0
_FILES_API_PROBE_TIMEOUT_SECONDS: Final = 5.0
_COMPLETION_WINDOW_SECONDS: Final = 24 * 60 * 60
_RUNNER_LOST_MESSAGE: Final = "the proxy replica running this batch stopped before it finished; resubmit the batch"
_EXPIRED_MESSAGE: Final = "This request could not be executed before the completion window expired."
_ROUTER_METHODS: Final[Mapping[BatchEndpoint, str]] = MappingProxyType(
{
"/v1/chat/completions": "acompletion",
"/v1/completions": "atext_completion",
"/v1/embeddings": "aembedding",
"/v1/responses": "aresponses",
}
)
_CANCELLING_TRANSITIONS: Final[Mapping[BatchStatus, BatchStatus]] = MappingProxyType(
{"completed": "cancelled", "expired": "cancelled", "in_progress": "cancelling", "finalizing": "cancelling"}
)
LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE: Final = (
"upload it through POST /v1/files with purpose=batch and either the x-litellm-model header or the "
"target_model_names form field naming the model, so LiteLLM keeps the file and runs the batch itself"
)
_RUNNING_BATCHES: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong references keep running batch tasks alive
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
class _ErrorDetail(TypedDict):
message: ReadOnly[str]
type: ReadOnly[str]
param: ReadOnly[None]
code: ReadOnly[None]
class _ErrorBody(TypedDict):
error: ReadOnly[_ErrorDetail]
class _ResultResponse(TypedDict):
status_code: ReadOnly[int]
request_id: ReadOnly[str]
body: ReadOnly[Mapping[str, object]]
class _LineError(TypedDict):
code: ReadOnly[str]
message: ReadOnly[str]
class _ResultLine(TypedDict):
id: ReadOnly[str]
custom_id: ReadOnly[str]
response: ReadOnly[_ResultResponse | None]
error: ReadOnly[_LineError | None]
class BatchInputLine(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
custom_id: str
method: Literal["POST"]
url: str
body: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class InvalidBatchInput:
line_number: int | None
reason: str
def describe(self) -> str:
return f"line {self.line_number}: {self.reason}" if self.line_number is not None else self.reason
@dataclass(frozen=True, slots=True)
class RowOutcome:
custom_id: str
status_code: int
body: Mapping[str, object]
succeeded: bool
@dataclass(frozen=True, slots=True)
class ExpiredRow:
custom_id: str
@dataclass(frozen=True, slots=True)
class _BatchRun:
unified_batch_id: str
llm_batch_id: str
model: str
endpoint: BatchEndpoint
lines: tuple[BatchInputLine, ...]
user_api_key_dict: UserAPIKeyAuth
request_tags: tuple[str, ...]
deadline: float
@runtime_checkable
class ManagedBatchStore(Protocol):
def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: ...
async def get_unified_file_id(
self, file_id: str, litellm_parent_otel_span: object | None = None
) -> LiteLLM_ManagedFileTable | None: ...
async def store_unified_object_id(
self,
unified_object_id: str,
file_object: LiteLLMBatch,
litellm_parent_otel_span: object | None,
model_object_id: str,
file_purpose: Literal["batch", "fine-tune", "response"],
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None = None,
persist_attribution: bool = False,
batch_processed: bool = False,
) -> None: ...
class _StorageBackendFactory(Protocol):
def __call__(self, backend_type: str, prisma_client: PrismaClient | None = None) -> BaseFileStorageBackend: ...
class _ResultFileUploader(Protocol):
def __call__(
self,
file_data: Mapping[str, object],
target_storage: str,
target_model_names: Sequence[str],
purpose: OpenAIFilesPurpose,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None = None,
) -> Awaitable[OpenAIFileObject]: ...
@runtime_checkable
class _RouterCall(Protocol):
def __call__(self, **params: object) -> Awaitable[object]: ... # kwargs-ok: the request body is passed as keywords
def litellm_executed_provider_of(credentials: Mapping[str, object]) -> str | None:
explicit_provider: Final = credentials.get("custom_llm_provider")
provider: Final = (
explicit_provider if isinstance(explicit_provider, str) else _provider_of(credentials.get("model"))
)
return provider if provider in LITELLM_EXECUTED_BATCH_PROVIDERS else None
class _HttpGetter(Protocol):
async def get(
self, url: str, *, headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None
) -> httpx.Response: ...
class FilesApiProbe(Protocol):
async def __call__(self, api_base: str, api_key: str | None) -> bool: ...
class BodyRejection(Protocol):
def __call__(self, body: Mapping[str, object], /) -> str | None: ...
async def upstream_lacks_files_api(api_base: str, api_key: str | None, http_client: _HttpGetter | None = None) -> bool:
client: Final = http_client or get_async_httpx_client(llm_provider=LlmProviders.HOSTED_VLLM)
try:
response: Final = await client.get(
f"{api_base.rstrip('/')}/files",
headers=(
{"Authorization": f"Bearer {api_key}"} # mutable-ok: AsyncHTTPHandler.get wants a plain dict
if api_key
else None
),
timeout=_FILES_API_PROBE_TIMEOUT_SECONDS,
)
except httpx.HTTPError:
return False
return response.status_code == httpx.codes.NOT_FOUND
def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, str | None] | None:
model: Final = credentials.get("model")
api_base: Final = credentials.get("api_base")
api_key: Final = credentials.get("api_key")
if not isinstance(model, str):
return None
try:
_, _, resolved_api_key, resolved_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=provider,
api_base=api_base if isinstance(api_base, str) else None,
api_key=api_key if isinstance(api_key, str) else None,
)
except Exception: # noqa: BLE001 # get_llm_provider raises on a model it cannot map, which means nothing to probe
return None
return None if resolved_api_base is None else (resolved_api_base, resolved_api_key)
async def litellm_executed_provider_for(
credentials: Mapping[str, object], lacks_files_api: FilesApiProbe = upstream_lacks_files_api
) -> str | None:
provider: Final = litellm_executed_provider_of(credentials)
if provider is None:
return None
upstream: Final = _upstream_of(credentials, provider)
if upstream is None:
return None
return provider if await lacks_files_api(*upstream) else None
async def resolve_litellm_executed_provider(
llm_router: "Router",
model: str,
team_id: str | None,
lacks_files_api: FilesApiProbe = upstream_lacks_files_api,
) -> str | None:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model, team_id=team_id)
return None if credentials is None else await litellm_executed_provider_for(credentials, lacks_files_api)
def _provider_of(model: object) -> str | None:
if not isinstance(model, str):
return None
try:
return litellm.get_llm_provider(model=model)[1]
except Exception: # noqa: BLE001 # get_llm_provider raises on an unknown model, which means no provider
return None
def _validation_reason(error: ValidationError) -> str:
return "; ".join(
f"{'.'.join(str(part) for part in item['loc'])}: {item['msg']}" if item["loc"] else item["msg"]
for item in error.errors()
)
def _accept_every_body(_body: Mapping[str, object]) -> str | None:
return None
def _parse_line(
line_number: int, raw: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection
) -> BatchInputLine | InvalidBatchInput:
try:
line: Final = BatchInputLine.model_validate_json(raw)
except ValidationError as e:
return InvalidBatchInput(line_number, _validation_reason(e))
if line.url != endpoint:
return InvalidBatchInput(line_number, f"url {line.url!r} does not match the batch endpoint {endpoint!r}")
if line.body.get("stream"):
return InvalidBatchInput(line_number, "streaming requests are not supported in a batch")
rejection: Final = reject_body(line.body)
if rejection is not None:
return InvalidBatchInput(line_number, rejection)
return line
def parse_batch_input(
content: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection = _accept_every_body
) -> tuple[BatchInputLine, ...] | InvalidBatchInput:
raw_lines: Final = tuple((number, raw) for number, raw in enumerate(content.splitlines(), start=1) if raw.strip())
if not raw_lines:
return InvalidBatchInput(None, "the input file has no requests")
parsed: Final = tuple(_parse_line(number, raw, endpoint, reject_body) for number, raw in raw_lines)
first_invalid: Final = next((item for item in parsed if isinstance(item, InvalidBatchInput)), None)
if first_invalid is not None:
return first_invalid
lines: Final = tuple(item for item in parsed if isinstance(item, BatchInputLine))
custom_ids: Final = sorted(line.custom_id for line in lines)
duplicate: Final = next((first for first, second in pairwise(custom_ids) if first == second), None)
if duplicate is not None:
return InvalidBatchInput(None, f"custom_id {duplicate!r} is used more than once")
return lines
def batch_error(status_code: int, message: str) -> ProxyException:
error_type: Final = "invalid_request_error" if status_code < 500 else ProxyErrorTypes.internal_server_error.value
return ProxyException(message=message, type=error_type, param=None, code=status_code)
def _validate_endpoint(endpoint: object) -> BatchEndpoint:
try:
return _BATCH_ENDPOINT_ADAPTER.validate_python(endpoint)
except ValidationError:
raise batch_error(400, f"endpoint {endpoint!r} is not supported for a LiteLLM-executed batch")
def _status_code_of(error: Exception) -> int:
status_code: Final[object] = getattr(error, "status_code", None)
return status_code if isinstance(status_code, int) else 500
def _error_body(error: Exception) -> _ErrorBody:
body: Final[_ErrorBody] = {
"error": {"message": str(error), "type": type(error).__name__, "param": None, "code": None}
}
return body
def _line_response(outcome: RowOutcome | ExpiredRow) -> _ResultResponse | None:
if isinstance(outcome, ExpiredRow):
return None
response: Final[_ResultResponse] = {
"status_code": outcome.status_code,
"request_id": f"req_{uuid_module.uuid4().hex[:24]}",
"body": outcome.body,
}
return response
def _line_error(outcome: RowOutcome | ExpiredRow) -> _LineError | None:
if isinstance(outcome, RowOutcome):
return None
error: Final[_LineError] = {"code": "batch_expired", "message": _EXPIRED_MESSAGE}
return error
def _result_line(outcome: RowOutcome | ExpiredRow) -> _ResultLine:
line: Final[_ResultLine] = {
"id": f"batch_req_{uuid_module.uuid4().hex[:24]}",
"custom_id": outcome.custom_id,
"response": _line_response(outcome),
"error": _line_error(outcome),
}
return line
def _dump(response: object) -> Mapping[str, object]:
if isinstance(response, BaseModel):
return response.model_dump(mode="json")
raise TypeError(f"Batch rows must return a single response object, got {type(response).__name__}")
def _resolve_transition(current_status: str, requested: BatchStatus) -> BatchStatus:
if current_status != "cancelling":
return requested
return _CANCELLING_TRANSITIONS.get(requested, requested)
def executed_batch_runner_lost(status: str, updated_at: datetime) -> bool:
if status in TERMINAL_BATCH_STATUSES:
return False
return (datetime.now(timezone.utc) - updated_at).total_seconds() > _STALE_AFTER_SECONDS
class _StopWatch:
def __init__(self, load_status: Callable[[], Awaitable[str | None]], interval_seconds: float) -> None:
self._load_status = load_status
self._interval_seconds = interval_seconds
self._checked_at = float("-inf")
self._stopped = False
async def stopped(self) -> bool:
if self._stopped:
return True
now: Final = time.monotonic()
if now - self._checked_at < self._interval_seconds:
return False
self._checked_at = now
self._stopped = await self._load_status() in _STOP_STATUSES
return self._stopped
class LiteLLMExecutedBatchRunner:
def __init__(
self,
llm_router: "Router",
prisma_client: PrismaClient,
managed_files: ManagedBatchStore,
batches: ManagedBatchRepository,
proxy_logging_obj: ProxyLogging,
general_settings: Mapping[str, object],
concurrency: int = LITELLM_EXECUTED_BATCH_CONCURRENCY,
heartbeat_seconds: float = _HEARTBEAT_SECONDS,
completion_window_seconds: float = _COMPLETION_WINDOW_SECONDS,
storage_backend_factory: _StorageBackendFactory = get_storage_backend,
upload_result_file: _ResultFileUploader = StorageBackendFileService.upload_file_to_storage_backend,
) -> None:
self.llm_router = llm_router
self.prisma_client = prisma_client
self.managed_files = managed_files
self.batches = batches
self.proxy_logging_obj = proxy_logging_obj
self.general_settings = general_settings
self.concurrency = concurrency
self.heartbeat_seconds = heartbeat_seconds
self.completion_window_seconds = completion_window_seconds
self.storage_backend_factory = storage_backend_factory
self.upload_result_file = upload_result_file
async def create(
self,
create_request: LiteLLMBatchCreateRequest,
unified_input_file_id: str,
model: str,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None,
) -> LiteLLMBatch:
endpoint: Final = _validate_endpoint(create_request.get("endpoint"))
content: Final = await self._download_input(unified_input_file_id, user_api_key_dict)
parsed: Final = parse_batch_input(content, endpoint, self._body_rejection(model))
if isinstance(parsed, InvalidBatchInput):
raise batch_error(400, f"Invalid batch input file: {parsed.describe()}")
llm_batch_id: Final = f"{LITELLM_EXECUTED_BATCH_ID_PREFIX}{uuid_module.uuid4().hex}"
model_id: Final = next(iter(self.llm_router.get_model_ids(model_name=model)), model)
unified_batch_id: Final = self.managed_files.get_unified_batch_id(batch_id=llm_batch_id, model_id=model_id)
now: Final = time.time()
created_at: Final = int(now)
batch: Final = LiteLLMBatch(
id=unified_batch_id,
object="batch",
endpoint=endpoint,
input_file_id=unified_input_file_id,
completion_window="24h",
status="validating",
created_at=created_at,
expires_at=created_at + int(self.completion_window_seconds),
metadata=create_request.get("metadata"),
model=model,
request_counts=BatchRequestCounts(completed=0, failed=0, total=len(parsed)),
)
await self.managed_files.store_unified_object_id(
unified_object_id=unified_batch_id,
file_object=batch,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
model_object_id=llm_batch_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
request_tags=request_tags,
persist_attribution=True,
batch_processed=True,
)
_record_batch_created(model, provider, user_api_key_dict)
run: Final = _BatchRun(
unified_batch_id=unified_batch_id,
llm_batch_id=llm_batch_id,
model=model,
endpoint=endpoint,
lines=parsed,
user_api_key_dict=user_api_key_dict,
request_tags=tuple(request_tags or ()),
deadline=now + self.completion_window_seconds,
)
task: Final = asyncio.create_task(self._run(run))
_RUNNING_BATCHES.add(task)
task.add_done_callback(_RUNNING_BATCHES.discard)
return batch
async def cancel(self, unified_batch_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLMBatch:
current: Final = await self.batches.load_batch(unified_batch_id)
if current is None:
raise batch_error(404, f"Batch {unified_batch_id} not found")
if current.status in TERMINAL_BATCH_STATUSES:
raise batch_error(400, f"Cannot cancel a batch with status '{current.status}'")
if current.status == "cancelling":
return current
cancelling: Final = current.model_copy(
update=MappingProxyType({"status": "cancelling", "cancelling_at": int(time.time())})
)
unchanged: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {"status": current.status}
if await self.batches.compare_and_set(cancelling, unchanged, user_api_key_dict.user_id):
return cancelling
return await self.cancel(unified_batch_id, user_api_key_dict)
async def fail_abandoned(self, batch: LiteLLMBatch, user_api_key_dict: UserAPIKeyAuth) -> LiteLLMBatch:
error: Final = BatchError(message=_RUNNER_LOST_MESSAGE, code="runner_lost")
errors: Final = Errors(data=[error], object="list") # mutable-ok: Errors.data is typed as a list
failed: Final = batch.model_copy(
update=MappingProxyType({"status": "failed", "failed_at": int(time.time()), "errors": errors})
)
untouched: Final[prisma_types.DateTimeFilter] = {
"lt": datetime.now(timezone.utc) - timedelta(seconds=_STALE_AFTER_SECONDS)
}
still_abandoned: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {
"status": batch.status,
"updated_at": untouched,
}
if await self.batches.compare_and_set(failed, still_abandoned, user_api_key_dict.user_id):
return failed
return await self.batches.load_batch(batch.id) or batch
def _body_rejection(self, model: str) -> BodyRejection:
def reject(body: Mapping[str, object]) -> str | None:
try:
is_request_body_safe(
request_body=dict(body), # mutable-ok: is_request_body_safe takes a dict
general_settings=dict(self.general_settings), # mutable-ok: is_request_body_safe takes a dict
llm_router=self.llm_router,
model=model,
)
except ValueError as e:
return str(e)
return None
return reject
async def _download_input(self, unified_input_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bytes:
stored: Final = await self.managed_files.get_unified_file_id(
unified_input_file_id, litellm_parent_otel_span=user_api_key_dict.parent_otel_span
)
if stored is None or not stored.storage_backend or not stored.storage_url:
raise batch_error(
400,
f"LiteLLM does not hold the content of input file {unified_input_file_id}: "
f"{LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE}",
)
try:
backend: Final = self.storage_backend_factory(stored.storage_backend, prisma_client=self.prisma_client)
return await backend.download_file(stored.storage_url)
except ValueError as e:
raise batch_error(400, str(e))
async def _run(self, run: _BatchRun) -> None:
heartbeat: Final = asyncio.create_task(self._heartbeat(run))
try:
await self._execute(run)
except Exception as e: # noqa: BLE001 # whatever fails, the batch must end up marked failed
verbose_proxy_logger.exception("LiteLLM-executed batch %s failed: %s", run.unified_batch_id, e)
error: Final = BatchError(message=str(e), code="internal_error")
errors: Final = Errors(data=[error], object="list") # mutable-ok: Errors.data is typed as a list
try:
await self._advance(run, "failed", MappingProxyType({"errors": errors}))
except Exception as advance_error: # noqa: BLE001 # a failed status write is logged, never raised
verbose_proxy_logger.exception(
"LiteLLM-executed batch %s could not be marked failed: %s", run.unified_batch_id, advance_error
)
finally:
heartbeat.cancel()
async def _heartbeat(self, run: _BatchRun) -> None:
while True:
await asyncio.sleep(self.heartbeat_seconds)
try:
await self._touch(run)
except Exception as e: # noqa: BLE001 # a missed beat is logged and the next one retries
verbose_proxy_logger.warning("LiteLLM-executed batch %s heartbeat failed: %s", run.unified_batch_id, e)
async def _touch(self, run: _BatchRun) -> None:
await self.batches.touch(run.unified_batch_id, run.user_api_key_dict.user_id)
async def _execute(self, run: _BatchRun) -> None:
await self._advance(run, "in_progress")
watch: Final = _StopWatch(lambda: self.batches.load_status(run.unified_batch_id), _CANCEL_POLL_SECONDS)
semaphore: Final = asyncio.Semaphore(self.concurrency)
results: Final = await asyncio.gather(*(self._run_row(run, line, watch, semaphore) for line in run.lines))
outcomes: Final = tuple(outcome for outcome in results if outcome is not None)
if await self._advance(run, "finalizing") is None:
return
succeeded: Final = tuple(
outcome for outcome in outcomes if isinstance(outcome, RowOutcome) and outcome.succeeded
)
failed: Final = tuple(
outcome for outcome in outcomes if isinstance(outcome, ExpiredRow) or not outcome.succeeded
)
output_file_id: Final = await self._upload_results(run, "output", succeeded)
error_file_id: Final = await self._upload_results(run, "error", failed)
request_counts: Final = BatchRequestCounts(completed=len(succeeded), failed=len(failed), total=len(run.lines))
final_status: Final[BatchStatus] = (
"expired" if any(isinstance(outcome, ExpiredRow) for outcome in outcomes) else "completed"
)
await self._advance(
run,
final_status,
MappingProxyType(
{"output_file_id": output_file_id, "error_file_id": error_file_id, "request_counts": request_counts}
),
)
async def _run_row(
self, run: _BatchRun, line: BatchInputLine, watch: _StopWatch, semaphore: asyncio.Semaphore
) -> RowOutcome | ExpiredRow | None:
async with semaphore:
if await watch.stopped():
return None
remaining: Final = run.deadline - time.time()
if remaining <= 0:
return ExpiredRow(custom_id=line.custom_id)
try:
return await asyncio.wait_for(self._row_outcome(run, line), timeout=remaining)
except asyncio.TimeoutError:
return ExpiredRow(custom_id=line.custom_id)
async def _row_outcome(self, run: _BatchRun, line: BatchInputLine) -> RowOutcome:
try:
body: Final = await self._dispatch(run, line)
except Exception as e: # noqa: BLE001 # a provider error becomes the row's error line, never a crashed batch
return RowOutcome(
custom_id=line.custom_id, status_code=_status_code_of(e), body=_error_body(e), succeeded=False
)
return RowOutcome(custom_id=line.custom_id, status_code=200, body=body, succeeded=True)
async def _dispatch(self, run: _BatchRun, line: BatchInputLine) -> Mapping[str, object]:
params: Final = MappingProxyType(
{**line.body, "model": run.model, "metadata": self._row_metadata(run), "disable_fallbacks": True}
)
return _dump(await self._router_call(run.endpoint)(**params))
def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall:
method: Final[object] = getattr(self.llm_router, _ROUTER_METHODS[endpoint], None)
if not isinstance(method, _RouterCall):
raise TypeError(f"the router has no callable for {endpoint}")
return method
def _row_metadata(self, run: _BatchRun) -> dict[str, object]: # mutable-ok: router updates metadata in place
return { # mutable-ok: the router updates request metadata in place
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(run.user_api_key_dict),
"user_api_key": run.user_api_key_dict.api_key,
"user_api_end_user_max_budget": run.user_api_key_dict.end_user_max_budget,
"tags": list(run.request_tags), # mutable-ok: litellm types request tags as a list
"batch_id": run.unified_batch_id,
}
async def _upload_results(
self, run: _BatchRun, kind: Literal["output", "error"], outcomes: Sequence[RowOutcome | ExpiredRow]
) -> str | None:
if not outcomes:
return None
content: Final = "".join(f"{json.dumps(_result_line(outcome))}\n" for outcome in outcomes).encode()
file_data: Final[ExtractedFileData] = {
"filename": f"{run.llm_batch_id}_{kind}.jsonl",
"content": content,
"content_type": "application/jsonl",
"headers": _NO_HEADERS,
}
file_object: Final = await self.upload_result_file(
file_data=file_data,
target_storage=LITELLM_DB_STORAGE_BACKEND_NAME,
target_model_names=(run.model,),
purpose="batch_output",
proxy_logging_obj=self.proxy_logging_obj,
user_api_key_dict=run.user_api_key_dict,
prisma_client=self.prisma_client,
)
return file_object.id
async def _advance(
self, run: _BatchRun, requested: BatchStatus, fields: Mapping[str, object] = _NO_FIELDS
) -> BatchStatus | None:
current: Final = await self.batches.load_batch(run.unified_batch_id)
if current is None:
raise RuntimeError(f"Batch {run.unified_batch_id} is no longer stored")
if current.status in TERMINAL_BATCH_STATUSES:
return None
status: Final = _resolve_transition(current.status, requested)
updated: Final = current.model_copy(
update=MappingProxyType({**fields, "status": status, f"{status}_at": int(time.time())})
)
unchanged: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {"status": current.status}
if await self.batches.compare_and_set(updated, unchanged, run.user_api_key_dict.user_id):
return status
return await self._advance(run, requested, fields)
def _record_batch_created(model: str, provider: str, user_api_key_dict: UserAPIKeyAuth) -> None:
prometheus_logger: Final = PrometheusLogger.get_instance()
if prometheus_logger is None:
return
prometheus_logger.record_managed_batch_created(
model=model,
api_provider=provider,
user=user_api_key_dict.user_id or "",
user_email=user_api_key_dict.user_email or "",
api_key_alias=user_api_key_dict.key_alias or "",
)

67
litellm/proxy/client/cli/commands/statusline_script.py
View file

@ -32,6 +32,7 @@ import unicodedata
import urllib.error
import urllib.request
from collections.abc import Callable, Mapping
from math import isfinite
from pathlib import Path
from types import MappingProxyType
from typing import IO, Final, NamedTuple, Protocol
@ -43,6 +44,7 @@ FETCH_TIMEOUT_SECONDS: Final = 3
BAR_WIDTH: Final = 24
BAR_FULL: Final = "\u2588"
BAR_EMPTY: Final = "\u2591"
SEPARATOR: Final = " \u00b7 "
TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024
CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",)
CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY")
@ -63,8 +65,11 @@ class Session(NamedTuple):
router_name: str
last_model: str
spend: float
baseline_spend: float
baseline_spend: float | None
baseline_model: str | None
turns: int | None = None
savings_estimated_turns: int | None = None
savings_estimated_actual_spend: float | None = None
class Credentials(NamedTuple):
@ -205,17 +210,38 @@ def _session_from_payload(payload: Mapping[str, object]) -> Session | None:
router_name: Final = printable(payload.get("router_name"))
last_model: Final = printable(payload.get("last_model"))
spend: Final = payload.get("spend")
baseline_spend: Final = payload.get("baseline_spend")
baseline_spend: Final = payload.get("savings_estimated_baseline_spend", payload.get("baseline_spend"))
turns: Final = payload.get("turns")
estimated_turns: Final = payload.get("savings_estimated_turns")
estimated_actual: Final = payload.get("savings_estimated_actual_spend")
if not router_name or not last_model:
return None
if not isinstance(spend, (int, float)) or not isinstance(baseline_spend, (int, float)):
if not isinstance(spend, (int, float)) or isinstance(spend, bool) or not isfinite(spend):
return None
if baseline_spend is not None and (
not isinstance(baseline_spend, (int, float)) or isinstance(baseline_spend, bool) or not isfinite(baseline_spend)
):
return None
return Session(
router_name=router_name,
last_model=last_model,
spend=float(spend),
baseline_spend=float(baseline_spend),
baseline_spend=float(baseline_spend) if baseline_spend is not None else None,
baseline_model=printable(payload.get("baseline_model")) or None,
turns=turns if isinstance(turns, int) and not isinstance(turns, bool) and turns >= 0 else None,
savings_estimated_turns=(
estimated_turns
if isinstance(estimated_turns, int) and not isinstance(estimated_turns, bool) and estimated_turns >= 0
else (0 if estimated_turns is not None else None)
),
savings_estimated_actual_spend=(
float(estimated_actual)
if isinstance(estimated_actual, (int, float))
and not isinstance(estimated_actual, bool)
and isfinite(estimated_actual)
and estimated_actual >= 0
else None
),
)
@ -314,15 +340,36 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo
return f"{code}{text}{RESET}" if use_color else text
routed: Final = paint(BOLD, f"Routed to: {model}")
if session is None or session.baseline_model is None or session.baseline_spend <= 0:
if session is None:
return routed
if session.savings_estimated_turns == 0 or session.baseline_spend is None:
return f"{routed}{SEPARATOR}Savings unavailable"
if session.baseline_model is None or session.baseline_spend <= 0:
return routed
if session.savings_estimated_turns is not None and (
session.savings_estimated_actual_spend is None
or session.turns is None
or session.savings_estimated_turns > session.turns
):
return f"{routed}{SEPARATOR}Savings unavailable"
compared_spend: Final = (
session.savings_estimated_actual_spend
if session.savings_estimated_turns is not None and session.savings_estimated_actual_spend is not None
else session.spend
)
coverage: Final = (
f"{SEPARATOR}{session.savings_estimated_turns} of {session.turns} turns estimated"
if session.savings_estimated_turns is not None
else ""
)
reference: Final = baseline_label(session.baseline_model, config_dir)
pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100
delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}")
peak: Final = max(session.spend, session.baseline_spend)
pct: Final = round((session.baseline_spend - compared_spend) / session.baseline_spend * 100)
sign: Final = "-" if pct > 0 else "+" if pct < 0 else ""
delta: Final = paint(LITELLM_COLOR, f"{sign}{abs(pct)}% vs {reference}")
peak: Final = max(compared_spend, session.baseline_spend)
label_width: Final = max(_display_width(session.router_name), _display_width(reference))
rows: Final = (
(session.router_name, session.spend, LITELLM_COLOR),
(session.router_name, compared_spend, LITELLM_COLOR),
(reference, session.baseline_spend, BASELINE_COLOR),
)
lines: Final = (
@ -331,7 +378,7 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo
f"{paint(DIM, f'${amount:.2f}')}"
for label, amount, color in rows
)
return "\n".join((f"{routed} {delta}", *lines))
return "\n".join((f"{routed} {delta}{coverage}", *lines))
def color_enabled(env: Mapping[str, str]) -> bool:

8
litellm/proxy/common_request_processing.py
View file

@ -3749,6 +3749,14 @@ class ProxyBaseLLMRequestProcessing:
"async_streaming_data_generator: error closing response stream: %s",
e,
)
logging_obj: Final = request_data.get("litellm_logging_obj")
if (
not stream_completed
and isinstance(logging_obj, LiteLLMLoggingObj)
and logging_obj.baseline_cache_context is not None
and logging_obj.model_call_details.get("prompt_cache_response_complete") is not True
):
await logging_obj.invalidate_baseline_cache_estimate("incomplete_response", completed=True)
@staticmethod
async def async_streaming_data_generator(

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

57
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({})
@ -49,15 +60,25 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def rejected_writes(self, incoming: Mapping[str, JsonValue]) -> tuple[str, ...]:
return tuple(
sorted(
key for key, value in incoming.items() if self.owned_by_config(key) and value != self._yaml_values[key]
)
sorted(key for key, value in incoming.items() if self.owned_by_config(key) and value != self.get(key))
)
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)
changed: Final = frozenset(
key
for key in (*previous_row, *db_row)
if previous_row.get(key, ABSENT) != db_row.get(key, ABSENT) # pyright: ignore[reportUnknownArgumentType] # JsonValue vs Absent compare
)
self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))})
self._clear_runtime_keys(frozenset((*previous_row, *db_row)))
self._clear_runtime_keys(changed)
def resolved(self) -> Mapping[str, JsonValue]:
return MappingProxyType(dict(self))
@ -81,7 +102,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 +110,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}
)
@ -112,6 +133,9 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def __len__(self) -> int:
return sum(1 for _ in self)
def __bool__(self) -> bool:
return any(True for _ in self)
def _clear_runtime(self) -> None:
self._runtime_values = _EMPTY_VALUES
self._deleted_runtime_keys = frozenset()
@ -136,8 +160,19 @@ 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)
and db_value != self.config_value(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))

43
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Final, NamedTuple
from litellm._logging import verbose_proxy_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.db.create_views import SupportsExecuteRaw
if TYPE_CHECKING:
from litellm.proxy._types import SpendLogsPayload
@ -75,6 +76,9 @@ SELECT
COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
COALESCE(SUM(spend), 0)::float8 AS spend,
COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
COALESCE(SUM(savings_estimated_turns), 0)::int AS savings_estimated_turns,
COALESCE(SUM(savings_estimated_actual_spend), 0)::float8 AS savings_estimated_actual_spend,
COALESCE(SUM(savings_estimated_saved_spend), 0)::float8 AS savings_estimated_saved_spend,
COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost,
COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns,
COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
@ -104,6 +108,9 @@ class AutoRouterTurnTransaction:
cache_touched: bool
tier: str | None = None
baseline_model: str | None = None
savings_estimated_turns: int = 0
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
class TurnCacheFacts(NamedTuple):
@ -215,13 +222,18 @@ def build_autorouter_turn_transaction(
turn_at: Final = _turn_time_utc(str(payload.get("startTime") or ""))
if turn_at is None:
return None
from litellm.proxy.spend_tracking.savings import classifier_cost_from_decision
from litellm.proxy.spend_tracking.savings import (
classifier_cost_from_decision,
recorded_estimated_autorouter_savings,
)
usage_object_raw: Final = metadata.get("usage_object")
cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None)
tier_raw: Final = routing_decision.get("tier")
baseline_raw: Final = routing_decision.get("savings_baseline_model")
classifier_cost: Final = classifier_cost_from_decision(routing_decision)
actual_spend: Final = float(payload.get("spend") or 0.0) + (classifier_cost or 0.0)
estimated_savings: Final = recorded_estimated_autorouter_savings(metadata)
return AutoRouterTurnTransaction(
api_key=api_key,
session_id=bounded_session_id(session_id),
@ -232,13 +244,16 @@ def build_autorouter_turn_transaction(
model=model,
turn_at=turn_at,
total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0),
spend=actual_spend,
saved_spend=saved_spend,
classifier_cost=classifier_cost or 0.0,
covered=cache.covered,
cache_hit=cache.read_tokens > 0,
cache_ttl_seconds=cache.write_ttl_seconds,
cache_touched=cache.touched,
savings_estimated_turns=int(estimated_savings is not None),
savings_estimated_actual_spend=actual_spend if estimated_savings is not None else 0.0,
savings_estimated_saved_spend=estimated_savings if estimated_savings is not None else 0.0,
)
@ -263,6 +278,10 @@ _BASELINE: Final = f"{_p('baseline_model')}::text"
_BASELINE_DELTA: Final = (
f"(CASE WHEN {_BASELINE} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_BASELINE}, 1) END)"
)
_ESTIMATED_BASELINE: Final = f"{_p('savings_estimated_turns')}::int = 1 AND {_BASELINE} IS NOT NULL"
_ESTIMATED_BASELINE_DELTA: Final = (
f"(CASE WHEN {_ESTIMATED_BASELINE} THEN jsonb_build_object({_BASELINE}, 1) ELSE '{{}}'::jsonb END)"
)
_IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at"
_SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}"
@ -281,7 +300,8 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t (
same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns,
baseline_models
baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend,
savings_estimated_baseline_models
)
VALUES (
{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
@ -292,13 +312,18 @@ VALUES (
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
{_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
{_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA}
{_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA},
{_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8,
{_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA}
)
ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
turns = t.turns + 1,
total_tokens = t.total_tokens + EXCLUDED.total_tokens,
spend = t.spend + EXCLUDED.spend,
saved_spend = t.saved_spend + EXCLUDED.saved_spend,
savings_estimated_turns = t.savings_estimated_turns + EXCLUDED.savings_estimated_turns,
savings_estimated_actual_spend = t.savings_estimated_actual_spend + EXCLUDED.savings_estimated_actual_spend,
savings_estimated_saved_spend = t.savings_estimated_saved_spend + EXCLUDED.savings_estimated_saved_spend,
classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost,
classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1,
covered_turns = t.covered_turns + EXCLUDED.covered_turns,
@ -331,6 +356,10 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
baseline_models = (CASE WHEN {_BASELINE} IS NOT NULL
THEN t.baseline_models || jsonb_build_object({_BASELINE}, COALESCE((t.baseline_models ->> {_BASELINE})::int, 0) + 1)
ELSE t.baseline_models END),
savings_estimated_baseline_models = (CASE WHEN {_ESTIMATED_BASELINE}
THEN t.savings_estimated_baseline_models || jsonb_build_object(
{_BASELINE}, COALESCE((t.savings_estimated_baseline_models ->> {_BASELINE})::int, 0) + 1)
ELSE t.savings_estimated_baseline_models END),
first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at),
last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at)
"""
@ -348,6 +377,10 @@ def _upsert_params(transaction: AutoRouterTurnTransaction) -> tuple[str | float
return tuple(_as_sql_param(getattr(transaction, name)) for name in _UPSERT_PARAM_FIELDS)
async def write_autorouter_turn(db: SupportsExecuteRaw, transaction: AutoRouterTurnTransaction) -> None:
await db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction))
async def _upsert_turn_with_retry(
prisma_client: PrismaClient,
transaction: AutoRouterTurnTransaction,
@ -355,7 +388,7 @@ async def _upsert_turn_with_retry(
) -> None:
for attempt in range(n_retry_times + 1):
try:
await prisma_client.db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction))
await write_autorouter_turn(prisma_client.db, transaction)
except DB_RETRY_SAFE_ERROR_TYPES:
if attempt >= n_retry_times:
raise

640
litellm/proxy/db/baseline_accounting.py Normal file
View file

@ -0,0 +1,640 @@
from __future__ import annotations
import asyncio
import json
from collections.abc import AsyncIterator, Callable, Sequence
from datetime import datetime, timedelta
from functools import reduce
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, cast
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator
from typing_extensions import Self
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.autorouter_session_rollup import (
AutoRouterTurnTransaction,
write_autorouter_turn,
)
from litellm.proxy.db.create_views import SupportsRawQueries
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
DailySpendEntity,
SpendRow,
build_bulk_upsert,
merge_by_conflict_key,
)
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
from litellm.proxy.spend_tracking.baseline_accounting import (
BaselineEstimate,
BaselineHistory,
BaselineObservation,
advance_baseline_history,
)
from litellm.proxy.spend_tracking.savings import BaselineCosts, BaselineCostSnapshot, price_baseline_comparison
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
class DailyBaselineTarget(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
entity: DailySpendEntity
entity_id: str | None
class DailyBaselineAttribution(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
date: str
api_key: str
model: str | None = None
custom_llm_provider: str | None = None
model_group: str | None = None
endpoint: str | None = None
mcp_namespaced_tool_name: str | None = None
targets: tuple[DailyBaselineTarget, ...] = ()
def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow:
table: Final = DAILY_SPEND_TABLES[target.entity]
return MappingProxyType(
{
"date": self.date,
"api_key": self.api_key,
"model": self.model,
"custom_llm_provider": self.custom_llm_provider,
"model_group": self.model_group,
"endpoint": self.endpoint,
"mcp_namespaced_tool_name": self.mcp_namespaced_tool_name,
table.entity_id_column: target.entity_id,
"request_id": request_id,
"autorouter_savings_spend": savings_delta,
}
)
class BaselineAccountingRecord(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
scope: str = Field(pattern=r"^autorouter-baseline:v3:[a-f0-9]{64}$")
api_key: str = Field(min_length=1)
session_id: str = Field(min_length=1, max_length=256)
router_name: str = Field(min_length=1)
baseline_model: str = Field(min_length=1)
observation: BaselineObservation
pricing: BaselineCostSnapshot
turn: AutoRouterTurnTransaction | None
daily: DailyBaselineAttribution | None
@model_validator(mode="after")
def consistent_turn(self) -> Self:
turn: Final = self.turn
if turn is not None and (
(turn.api_key, turn.session_id, turn.router_name, turn.baseline_model)
!= (self.api_key, self.session_id, self.router_name, self.baseline_model)
or turn.spend != self.pricing.actual_spend + self.pricing.classifier_cost
or any(
(
turn.saved_spend,
turn.savings_estimated_turns,
turn.savings_estimated_actual_spend,
turn.savings_estimated_saved_spend,
)
)
):
raise ValueError("Baseline observation must own an unestimated turn with matching scope and actual cost")
return self
class BaselinePublication(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
version: Literal[3] = 3
comparison_id: str
comparison_started_at: float
status: Literal["estimated", "unknown"]
reason: str
provenance: Literal["observed_identical", "modeled"] | None = None
actual_spend: float | None = None
baseline_spend: float | None = None
input_tokens: int | None = None
cache_read_input_tokens: int | None = None
cache_creation_5m_input_tokens: int | None = None
cache_creation_1h_input_tokens: int | None = None
@property
def costs(self) -> BaselineCosts | None:
if self.status != "estimated" or self.actual_spend is None or self.baseline_spend is None:
return None
return BaselineCosts(self.actual_spend, self.baseline_spend)
def baseline_publication(
record: BaselineAccountingRecord, estimate: BaselineEstimate, first_at: float
) -> BaselinePublication:
costs: Final = price_baseline_comparison(record.pricing, estimate.usage, estimate.provenance)
details: Final = estimate.usage.prompt_tokens_details if estimate.usage is not None else None
writes: Final = details.cache_creation_token_details if details is not None else None
return BaselinePublication(
comparison_id=record.scope,
comparison_started_at=first_at,
status="estimated" if costs is not None else "unknown",
reason=estimate.reason if costs is not None or estimate.usage is None else "pricing_unavailable",
provenance=estimate.provenance if costs is not None else None,
actual_spend=costs.actual if costs is not None else None,
baseline_spend=costs.baseline if costs is not None else None,
input_tokens=details.text_tokens if details is not None else None,
cache_read_input_tokens=details.cached_tokens if details is not None else None,
cache_creation_5m_input_tokens=writes.ephemeral_5m_input_tokens if writes is not None else None,
cache_creation_1h_input_tokens=writes.ephemeral_1h_input_tokens if writes is not None else None,
)
class _Comparison(BaseModel):
revision: int
published_revision: int
initial_equivalent: bool
retired: bool
history: str | None
class _StoredRecord(BaseModel):
data: str
publication: str | None
conflicted: bool
started_at: float
class _Change(BaseModel):
request_id: str
publication: BaselinePublication
api_key: str
session_id: str
router_name: str
baseline_model: str
covered_delta: int
actual_delta: float
savings_delta: float
daily: DailyBaselineAttribution | None
class _TransactionManager(Protocol):
async def __aenter__(self) -> SupportsRawQueries: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
class _TransactionalDatabase(Protocol):
def tx(self, *, timeout: timedelta) -> _TransactionManager: ...
_COMPARISONS: Final = TypeAdapter(tuple[_Comparison, ...])
_RECORDS: Final = TypeAdapter(tuple[_StoredRecord, ...])
_HISTORY: Final = TypeAdapter(BaselineHistory)
_PAGE_TIMESTAMPS: Final = 128
_TRANSACTION_TIMEOUT: Final = timedelta(seconds=10)
_CREATE_COMPARISON: Final = """
INSERT INTO "LiteLLM_AutoRouterBaselineComparison"
(scope, api_key, session_id, router_name, initial_equivalent)
VALUES ($1, $2, $3, $4, NOT EXISTS (
SELECT 1 FROM "LiteLLM_AutoRouterSession"
WHERE api_key = $2 AND session_id = $3 AND router_name = $4
)) ON CONFLICT (scope) DO NOTHING
"""
_LOCK_COMPARISON: Final = """
SELECT revision, published_revision, initial_equivalent, retired, history
FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope = $1 FOR UPDATE
"""
_INSERT_RECORD: Final = """
INSERT INTO "LiteLLM_AutoRouterBaselineObservation"
(request_id, scope, started_at, revision, data)
VALUES ($1, $2, $3::float8, $4::bigint, $5)
ON CONFLICT (request_id) DO NOTHING
"""
_MARK_CONFLICT: Final = """
UPDATE "LiteLLM_AutoRouterBaselineObservation"
SET conflicted = TRUE, revision = $4::bigint
WHERE request_id = $1 AND scope = $2 AND data <> $3 AND NOT conflicted
"""
_READ_PAGE: Final = """
WITH times AS (
SELECT DISTINCT started_at FROM "LiteLLM_AutoRouterBaselineObservation"
WHERE scope = $1 AND revision > $2::bigint
AND ($3::float8 IS NULL OR started_at > $3::float8)
AND ($5::float8 IS NULL OR (
started_at >= $5::float8 AND publication::jsonb->>'status' = 'estimated'
))
ORDER BY started_at LIMIT $4::int
)
SELECT data, publication, conflicted, started_at
FROM "LiteLLM_AutoRouterBaselineObservation"
WHERE scope = $1 AND revision > $2::bigint
AND started_at IN (SELECT started_at FROM times)
AND ($5::float8 IS NULL OR publication::jsonb->>'status' = 'estimated')
ORDER BY started_at, request_id
"""
_UPDATE_LOGS: Final = """
WITH changes AS (
SELECT request_id, publication::jsonb AS publication
FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb)
)
UPDATE "LiteLLM_SpendLogs" AS logs
SET metadata = (COALESCE(logs.metadata::jsonb, '{}'::jsonb) - 'autorouter_baseline_observation') || jsonb_build_object(
'autorouter_savings_estimate', changes.publication,
'autorouter_savings', CASE WHEN changes.publication->>'status' = 'estimated' THEN
(changes.publication->>'baseline_spend')::float8 - (changes.publication->>'actual_spend')::float8
ELSE NULL END
)
FROM changes WHERE logs.request_id = changes.request_id
"""
_UPDATE_PUBLICATIONS: Final = """
UPDATE "LiteLLM_AutoRouterBaselineObservation" AS observations
SET publication = x.publication::text
FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb)
WHERE observations.request_id = x.request_id
"""
_UPDATE_SESSIONS: Final = """
WITH changes AS (
SELECT * FROM jsonb_to_recordset($1::jsonb) AS x(
api_key text, session_id text, router_name text, baseline_model text,
covered_delta int, actual_delta float8, savings_delta float8
)
), totals AS (
SELECT api_key, session_id, router_name, SUM(covered_delta)::int AS covered_delta,
SUM(actual_delta) AS actual_delta, SUM(savings_delta) AS savings_delta
FROM changes GROUP BY api_key, session_id, router_name
), models AS (
SELECT api_key, session_id, router_name, jsonb_object_agg(baseline_model, delta) AS deltas
FROM (
SELECT api_key, session_id, router_name, baseline_model, SUM(covered_delta)::int AS delta
FROM changes GROUP BY api_key, session_id, router_name, baseline_model
) grouped GROUP BY api_key, session_id, router_name
)
UPDATE "LiteLLM_AutoRouterSession" AS session
SET saved_spend = session.saved_spend + totals.savings_delta,
savings_estimated_turns = session.savings_estimated_turns + totals.covered_delta,
savings_estimated_actual_spend = session.savings_estimated_actual_spend + totals.actual_delta,
savings_estimated_saved_spend = session.savings_estimated_saved_spend + totals.savings_delta,
savings_estimated_baseline_models = (
SELECT COALESCE(jsonb_object_agg(key, value), '{}'::jsonb) FROM (
SELECT key, SUM(value::int)::int AS value FROM (
SELECT * FROM jsonb_each_text(session.savings_estimated_baseline_models)
UNION ALL SELECT * FROM jsonb_each_text(models.deltas)
) combined GROUP BY key HAVING SUM(value::int) > 0
) counts
)
FROM totals JOIN models USING (api_key, session_id, router_name)
WHERE session.api_key = totals.api_key AND session.session_id = totals.session_id
AND session.router_name = totals.router_name
"""
def _primary_transaction(client: PrismaClient) -> _TransactionManager:
primary: Final = cast(_TransactionalDatabase, writer_wrapper(client.db))
return primary.tx(timeout=_TRANSACTION_TIMEOUT)
def _serialized(model: BaseModel) -> str:
return json.dumps(model.model_dump(mode="json"), sort_keys=True, separators=(",", ":"))
def _change(record: BaselineAccountingRecord, old: BaselinePublication | None, new: BaselinePublication) -> _Change:
previous: Final = old.costs if old is not None else None
current: Final = new.costs
return _Change(
request_id=record.observation.request_id,
publication=new,
api_key=record.api_key,
session_id=record.session_id,
router_name=record.router_name,
baseline_model=record.baseline_model,
covered_delta=int(current is not None) - int(previous is not None),
actual_delta=(current.actual if current is not None else 0.0)
- (previous.actual if previous is not None else 0.0),
savings_delta=(current.savings if current is not None else 0.0)
- (previous.savings if previous is not None else 0.0),
daily=record.daily,
)
def _project_group(
previous: tuple[BaselineHistory, tuple[_Change, ...]], stored: Sequence[_StoredRecord]
) -> tuple[BaselineHistory, tuple[_Change, ...]]:
history, prior_changes = previous
records: Final = tuple(BaselineAccountingRecord.model_validate_json(item.data) for item in stored)
observations: Final = tuple(
record.observation.model_copy(
update=MappingProxyType(
{"outcome": "uncertain", "baseline_equivalent": False, "reason": "conflicting_observation"}
)
)
if row.conflicted
else record.observation
for record, row in zip(records, stored)
)
advanced, estimates = advance_baseline_history(history, observations)
publications: Final = tuple(
baseline_publication(
record, estimate, advanced.first_at if advanced.first_at is not None else observations[0].started_at
)
for record, estimate in zip(records, estimates)
)
changes: Final = tuple(
_change(record, old, publication)
for record, row, publication in zip(records, stored, publications)
for old in (BaselinePublication.model_validate_json(row.publication) if row.publication else None,)
if publication != old
)
return advanced, (*prior_changes, *changes)
async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None:
if not changes:
return
serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":"))
await db.execute_raw(_UPDATE_LOGS, serialized)
await db.execute_raw(_UPDATE_SESSIONS, serialized)
for entity, table in DAILY_SPEND_TABLES.items():
if adjustments := tuple(
change.daily.adjustment(target, change.savings_delta, change.request_id)
for change in changes
if change.daily is not None and change.savings_delta != 0
for target in change.daily.targets
if target.entity == entity
):
statement, values = build_bulk_upsert(table, merge_by_conflict_key(table, adjustments))
await db.execute_raw(statement, *values)
await db.execute_raw(_UPDATE_PUBLICATIONS, serialized)
class BaselineAccountingStore:
def __init__(self, transaction: Callable[[], _TransactionManager]) -> None:
self.transaction: Final = transaction
@classmethod
def for_client(cls, client: PrismaClient) -> BaselineAccountingStore:
def transaction() -> _TransactionManager:
return _primary_transaction(client)
return cls(transaction)
async def append(
self, record: BaselineAccountingRecord
) -> Literal["recorded", "retired", "conflict", "unavailable"]:
try:
async with self.transaction() as db:
await db.execute_raw("SET LOCAL statement_timeout = 5000")
await db.execute_raw("SET LOCAL lock_timeout = 1000")
await db.execute_raw(
_CREATE_COMPARISON, record.scope, record.api_key, record.session_id, record.router_name
)
rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, record.scope)))
if not rows:
return "unavailable"
revision: Final = rows[0].revision + 1
data: Final = _serialized(record)
inserted: Final = await db.execute_raw(
_INSERT_RECORD,
record.observation.request_id,
record.scope,
record.observation.started_at,
revision,
data,
)
if inserted and record.turn is not None:
await write_autorouter_turn(db, record.turn)
conflicted: Final = (
0
if inserted
else await db.execute_raw(
_MARK_CONFLICT, record.observation.request_id, record.scope, data, revision
)
)
canonical: Final = (
_RECORDS.validate_python(
tuple(
await db.query_raw(
'SELECT data, publication, conflicted, started_at FROM "LiteLLM_AutoRouterBaselineObservation" '
"WHERE request_id=$1 AND scope=$2",
record.observation.request_id,
record.scope,
)
)
)
if not inserted
else ()
)
if not inserted and not canonical:
return "conflict"
if rows[0].retired:
await _publish(
db,
(
_change(
BaselineAccountingRecord.model_validate_json(canonical[0].data)
if canonical
else record,
BaselinePublication.model_validate_json(canonical[0].publication)
if canonical and canonical[0].publication is not None
else None,
BaselinePublication(
comparison_id=record.scope,
comparison_started_at=canonical[0].started_at
if canonical
else record.observation.started_at,
status="unknown",
reason="comparison_retired",
),
),
),
)
return "retired"
if inserted or conflicted:
await self._withdraw(
db, record.scope, canonical[0].started_at if canonical else record.observation.started_at
)
await db.execute_raw(
'UPDATE "LiteLLM_AutoRouterBaselineComparison" SET revision = $2::bigint, '
"updated_at = CURRENT_TIMESTAMP, attempted_at = NULL WHERE scope = $1",
record.scope,
revision,
)
return "recorded"
except Exception: # noqa: BLE001 # accounting failure must not change inference or actual billing
verbose_proxy_logger.warning("Auto-router baseline observation could not be persisted")
return "unavailable"
async def _pages(
self, db: SupportsRawQueries, scope: str, after_revision: int, withdraw_from: float | None = None
) -> AsyncIterator[tuple[_StoredRecord, ...]]:
cursor: float | None = None
while page := _RECORDS.validate_python(
tuple(await db.query_raw(_READ_PAGE, scope, after_revision, cursor, _PAGE_TIMESTAMPS, withdraw_from))
):
yield page
cursor = page[-1].started_at # rebind-ok: keyset pagination advances after each complete timestamp group
async def _withdraw(self, db: SupportsRawQueries, scope: str, started_at: float) -> None:
async for page in self._pages(db, scope, 0, withdraw_from=started_at):
await _publish(
db,
tuple(
_change(
BaselineAccountingRecord.model_validate_json(row.data),
previous,
BaselinePublication(
comparison_id=scope,
comparison_started_at=min(previous.comparison_started_at, started_at),
status="unknown",
reason="pending_projection",
),
)
for row in page
if row.publication is not None
for previous in (BaselinePublication.model_validate_json(row.publication),)
),
)
async def retire_before(self, cutoff: datetime, batch_size: int, timeout_ms: int) -> None:
async with self.transaction() as db:
await db.execute_raw(f"SET LOCAL statement_timeout = {max(1, timeout_ms)}")
await db.execute_raw(f"SET LOCAL lock_timeout = {max(1, timeout_ms)}")
await db.execute_raw(
'WITH expired AS (SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison" '
"WHERE NOT retired AND updated_at < $1::timestamptz ORDER BY updated_at "
"LIMIT $2::int FOR UPDATE SKIP LOCKED) "
'UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison '
"SET retired=TRUE, history=NULL FROM expired WHERE comparison.scope=expired.scope",
cutoff,
batch_size,
)
await db.execute_raw(
'DELETE FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id IN ('
'SELECT event.request_id FROM "LiteLLM_AutoRouterBaselineObservation" AS event '
'JOIN "LiteLLM_AutoRouterBaselineComparison" AS comparison USING (scope) '
"WHERE comparison.retired AND comparison.updated_at < $1::timestamptz "
"LIMIT $2::int)",
cutoff,
batch_size,
)
async def project(self, scope: str) -> Literal["published", "unchanged", "unavailable"]:
try:
async with self.transaction() as db:
await db.execute_raw("SET LOCAL statement_timeout = 5000")
await db.execute_raw("SET LOCAL lock_timeout = 1000")
rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, scope)))
if not rows or rows[0].retired or rows[0].revision == rows[0].published_revision:
return "unchanged"
missing_log: Final = await db.query_raw(
'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" AS observation '
'WHERE scope=$1 AND publication IS NULL AND NOT EXISTS (SELECT 1 FROM "LiteLLM_SpendLogs" AS log '
"WHERE log.request_id=observation.request_id) LIMIT 1",
scope,
)
if missing_log:
return "unavailable"
state: Final = rows[0]
checkpoint: Final = (
_HISTORY.validate_json(state.history)
if state.history is not None
else BaselineHistory(equivalent=state.initial_equivalent)
)
changed: Final = await db.query_raw(
'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" '
"WHERE scope = $1 AND revision > $2::bigint AND started_at <= $3::float8 LIMIT 1",
scope,
state.published_revision,
checkpoint.last_at,
)
history = BaselineHistory(equivalent=state.initial_equivalent) if changed else checkpoint
async for page in self._pages(db, scope, 0 if changed else state.published_revision):
history, updates = reduce(
_project_group,
(tuple(group) for _, group in groupby(page, key=lambda item: item.started_at)),
(history, ()),
)
await _publish(db, updates)
await db.execute_raw(
'UPDATE "LiteLLM_AutoRouterBaselineComparison" '
"SET published_revision = revision, history = $2 WHERE scope = $1",
scope,
_HISTORY.dump_json(history).decode(),
)
return "published"
except Exception: # noqa: BLE001 # rollback leaves the durable revision dirty for a later flush
verbose_proxy_logger.warning("Auto-router baseline projection remains pending")
return "unavailable"
class _Scope(BaseModel):
scope: str
_SCOPES: Final = TypeAdapter(tuple[_Scope, ...])
_CLAIM_DIRTY: Final = """
WITH candidates AS (
SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison"
WHERE NOT retired AND revision <> published_revision
AND (attempted_at IS NULL OR attempted_at < CURRENT_TIMESTAMP - INTERVAL '30 seconds')
ORDER BY attempted_at NULLS FIRST, updated_at, scope LIMIT 32 FOR UPDATE SKIP LOCKED
)
UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison
SET attempted_at = CURRENT_TIMESTAMP FROM candidates
WHERE comparison.scope = candidates.scope RETURNING comparison.scope
"""
async def _flush_records(
store: BaselineAccountingStore, records: Sequence[BaselineAccountingRecord]
) -> tuple[BaselineAccountingRecord, ...]:
slots: Final = asyncio.Semaphore(4)
async def append(record: BaselineAccountingRecord) -> bool:
async with slots:
return await store.append(record) == "unavailable"
failed: Final = await asyncio.gather(*(append(record) for record in records))
return tuple(record for record, retry in zip(records, failed) if retry)
async def flush_baseline_accounting(client: PrismaClient) -> None:
from litellm.proxy.utils import request_spend_log_flush
store: Final = BaselineAccountingStore.for_client(client)
async with client.baseline_accounting_lock:
batch: Final = tuple(client.baseline_accounting_transactions[:32])
client.baseline_accounting_transactions = client.baseline_accounting_transactions[
32:
] # rebind-ok: drain under lock
more_queued: Final = bool(client.baseline_accounting_transactions)
try:
remaining: Final = await asyncio.wait_for(_flush_records(store, batch), timeout=5)
except (Exception, asyncio.CancelledError) as error: # noqa: BLE001 # unknown acknowledgements can be replayed safely
async with client.baseline_accounting_lock:
client.baseline_accounting_transactions.extend(batch)
if isinstance(error, asyncio.CancelledError):
raise
return
async with client.baseline_accounting_lock:
client.baseline_accounting_transactions.extend(remaining)
if more_queued and len(remaining) < len(batch):
request_spend_log_flush(client)
try:
async with store.transaction() as db:
await db.execute_raw("SET LOCAL statement_timeout = 1000")
scopes: Final = _SCOPES.validate_python(tuple(await db.query_raw(_CLAIM_DIRTY)))
slots: Final = asyncio.Semaphore(4)
async def project(item: _Scope) -> str:
async with slots:
return await store.project(item.scope)
outcomes: Final = await asyncio.wait_for(asyncio.gather(*(project(item) for item in scopes)), timeout=5)
if len(scopes) == 32 and "published" in outcomes:
request_spend_log_flush(client)
except Exception: # noqa: BLE001 # durable dirty comparisons remain eligible after the retry interval
verbose_proxy_logger.warning("Auto-router baseline projection will retry on a later spend flush")

32
litellm/proxy/db/daily_spend_bulk_upsert.py
View file

@ -14,6 +14,8 @@ from itertools import groupby
from types import MappingProxyType
from typing import Final, Literal
from pydantic import TypeAdapter
DailySpendEntity = Literal["user", "team", "org", "tag", "end_user", "agent"]
SqlValue = str | int | float | None
@ -43,6 +45,36 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP
}
)
_ENTITY_INPUT_KEYS: Final[Mapping[DailySpendEntity, str]] = MappingProxyType(
{
"user": "user",
"team": "team_id",
"org": "organization_id",
"end_user": "end_user",
"agent": "agent_id",
"tag": "request_tags",
}
)
_TAGS: Final = TypeAdapter(tuple[str, ...])
def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEntity) -> tuple[str | None, ...]:
key: Final = _ENTITY_INPUT_KEYS[entity]
if key not in payload:
return ()
value: Final = payload[key]
if entity == "tag":
if value is None:
return ()
tags: Final = _TAGS.validate_json(value) if isinstance(value, str) else _TAGS.validate_python(value)
return tuple(dict.fromkeys(tags))
if value is None:
return (None,) if entity == "user" else ()
if not isinstance(value, str) or (entity == "end_user" and not value):
return ()
return (value,)
# The unique constraint's columns after the entity id, in constraint order. A NULL can
# never match itself in a unique index, so every one of these is normalized to '': the
# conflict target has to be NULL-free or the row is re-inserted on every single flush.

146
litellm/proxy/db/db_spend_update_writer.py
View file

@ -18,6 +18,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, overload
from urllib.parse import quote, unquote
from pydantic import TypeAdapter
from typing_extensions import LiteralString, ReadOnly, TypedDict
import litellm
@ -51,6 +52,7 @@ from litellm.proxy.common_utils.user_api_key_cache import project_cache_key
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
build_bulk_upsert,
daily_spend_entity_ids,
merge_by_conflict_key,
)
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
@ -82,6 +84,8 @@ from litellm.repositories.prisma_protocols import BatchTable
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution
from litellm.proxy.utils import PrismaClient, ProxyLogging
else:
PrismaClient = Any
@ -89,6 +93,7 @@ else:
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
_SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _org_member_transaction_key(org_id: str, user_id: str) -> str:
@ -579,25 +584,31 @@ class DBSpendUpdateWriter:
metadata_raw: Final = payload.get("metadata")
if not metadata_raw:
return
metadata: Final = json.loads(metadata_raw)
if not isinstance(metadata, dict) or not metadata.get("routing_decision"):
metadata: Final = _SPEND_METADATA_ADAPTER.validate_json(metadata_raw)
routing_decision: Final = metadata.get("routing_decision")
if not isinstance(routing_decision, Mapping) or not routing_decision:
return
from litellm.proxy.db.autorouter_session_rollup import (
build_autorouter_turn_transaction,
)
usage_object_raw: Final = metadata.get("usage_object")
cost_breakdown: Final = metadata.get("cost_breakdown")
savings_estimate: Final = metadata.get("autorouter_savings_estimate")
savings_spend: Final = compute_savings_spend(
model=payload.get("model"),
custom_llm_provider=payload.get("custom_llm_provider"),
compression_saved_tokens=0,
gateway_injected_cache=marks_gateway_injection(metadata, payload.get("model_id")),
routing_decision=metadata.get("routing_decision"),
routing_decision=routing_decision,
usage_object=usage_object_raw if isinstance(usage_object_raw, dict) else None,
model_id=payload.get("model_id"),
llm_router=get_llm_router,
cost_breakdown=metadata.get("cost_breakdown"),
cost_breakdown=cost_breakdown if isinstance(cost_breakdown, Mapping) else None,
recorded_autorouter_savings=metadata.get("autorouter_savings"),
recorded_autorouter_savings_estimate=(
savings_estimate if isinstance(savings_estimate, Mapping) else None
),
billed_at=payload.get("endTime"),
)
transaction: Final = build_autorouter_turn_transaction(
@ -605,6 +616,11 @@ class DBSpendUpdateWriter:
metadata=metadata,
saved_spend=savings_spend.autorouter,
)
try:
if await self._enqueue_baseline_accounting(payload, metadata, transaction, prisma_client):
return
except Exception: # noqa: BLE001 # optional baseline capture must preserve the original actual-spend rollup
verbose_proxy_logger.warning("Auto-router baseline observation was unavailable; actual turn retained")
if transaction is None:
return
async with prisma_client._autorouter_turn_transactions_lock:
@ -612,6 +628,95 @@ class DBSpendUpdateWriter:
except Exception as e: # noqa: BLE001 # a metrics enqueue must never fail the spend write
verbose_proxy_logger.debug("_enqueue_autorouter_turn_transaction error (non-blocking): %s", e)
async def _enqueue_baseline_accounting(
self,
payload: SpendLogsPayload,
metadata: Mapping[str, object],
turn: "AutoRouterTurnTransaction | None",
prisma_client: "PrismaClient",
) -> bool:
from litellm.proxy.db.baseline_accounting import (
BaselineAccountingRecord,
)
from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation
from litellm.proxy.spend_tracking.savings import baseline_cost_snapshot
serialized: Final = metadata.get("autorouter_baseline_observation")
if not isinstance(serialized, str):
return False
captured: Final = CapturedBaselineObservation.model_validate_json(serialized)
if captured.api_key != payload["api_key"] or captured.session_id != payload["session_id"]:
return False
decision: Final = _SPEND_METADATA_ADAPTER.validate_python(
metadata.get("routing_decision") or MappingProxyType({})
)
breakdown: Final = _SPEND_METADATA_ADAPTER.validate_python(
metadata.get("cost_breakdown") or MappingProxyType({})
)
daily: Final = await self._baseline_daily_attribution(payload, prisma_client)
record: Final = BaselineAccountingRecord(
scope=captured.scope,
api_key=captured.api_key,
session_id=captured.session_id,
router_name=captured.router_name,
baseline_model=captured.baseline_model,
observation=captured.observation.model_copy(update=MappingProxyType({"request_id": payload["request_id"]})),
pricing=baseline_cost_snapshot(captured.model, captured.prices, payload["spend"], breakdown, decision),
turn=turn,
daily=daily,
)
async with prisma_client.baseline_accounting_lock:
if len(prisma_client.baseline_accounting_transactions) >= 10000:
verbose_proxy_logger.warning("Auto-router baseline observation queue is full")
return False
prisma_client.baseline_accounting_transactions.append(record)
from litellm.proxy.utils import request_spend_log_flush
request_spend_log_flush(prisma_client)
return True
async def _baseline_daily_attribution(
self,
payload: SpendLogsPayload,
prisma_client: "PrismaClient",
) -> "DailyBaselineAttribution | None":
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution, DailyBaselineTarget
normalized: Final = cast(SpendLogsPayload, MappingProxyType({**payload, "end_user_id": payload["end_user"]}))
bases: Final = tuple(
zip(
DAILY_SPEND_TABLES,
await asyncio.gather(
*(
self._common_add_spend_log_transaction_to_daily_transaction( # pyright: ignore[reportUnknownMemberType] # legacy payload union; this caller supplies a validated spend payload
normalized,
prisma_client,
"request_tags" if entity == "tag" else entity,
)
for entity in DAILY_SPEND_TABLES
)
),
)
)
base: Final = next((base for _, base in bases if base is not None), None)
if base is None:
return None
return DailyBaselineAttribution(
date=base["date"],
api_key=base["api_key"],
model=base.get("model"),
custom_llm_provider=base.get("custom_llm_provider"),
model_group=base.get("model_group"),
endpoint=base.get("endpoint"),
mcp_namespaced_tool_name=base.get("mcp_namespaced_tool_name"),
targets=tuple(
DailyBaselineTarget(entity=entity, entity_id=identity)
for entity, values in bases
if values is not None
for identity in daily_spend_entity_ids(payload, entity)
),
)
def _enqueue_tool_registry_upsert(
self,
kwargs: dict | None,
@ -2322,21 +2427,13 @@ class DBSpendUpdateWriter:
prisma_client: PrismaClient,
type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user",
) -> BaseDailySpendTransaction | None:
common_expected_keys: Final = ["startTime", "api_key"]
if type == "user":
expected_keys = ["user", *common_expected_keys]
elif type == "team":
expected_keys = ["team_id", *common_expected_keys]
elif type == "org":
expected_keys = ["organization_id", *common_expected_keys]
elif type == "request_tags":
expected_keys = ["request_tags", *common_expected_keys]
elif type == "end_user":
expected_keys = ["end_user_id", *common_expected_keys]
elif type == "agent":
expected_keys = ["agent_id", *common_expected_keys]
else:
raise ValueError(f"Invalid type: {type}")
entity: Final = "tag" if type == "request_tags" else type
identity_payload: Final = (
MappingProxyType({**payload, "end_user": payload.get("end_user_id")}) if type == "end_user" else payload
)
if not daily_spend_entity_ids(identity_payload, entity):
return None
expected_keys: Final = ("startTime", "api_key")
if not all(key in payload for key in expected_keys):
verbose_proxy_logger.debug(
"Missing expected keys: %s, in payload, skipping from daily_user_spend_transactions", expected_keys
@ -2399,6 +2496,7 @@ class DBSpendUpdateWriter:
usage_object=usage_obj,
cost_breakdown=_metadata.get("cost_breakdown"),
recorded_autorouter_savings=_metadata.get("autorouter_savings"),
recorded_autorouter_savings_estimate=_metadata.get("autorouter_savings_estimate"),
billed_at=payload.get("endTime"),
)
timed_duration_ms: Final = _timed_request_duration_ms(payload, request_status, is_internal_call)
@ -2597,14 +2695,10 @@ class DBSpendUpdateWriter:
verbose_proxy_logger.debug("request_tags is None for request. Skipping incrementing tag spend.")
return
request_tags: Sequence[str] = []
if isinstance(payload["request_tags"], str):
request_tags = json.loads(payload["request_tags"])
elif isinstance(payload["request_tags"], list):
request_tags = payload["request_tags"]
else:
raise ValueError(f"Invalid request_tags: {payload['request_tags']}")
request_tags: Final = daily_spend_entity_ids(payload, "tag")
for tag in request_tags:
if tag is None:
continue
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyTagSpendTransaction(

11
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -549,6 +549,17 @@ class SpendLogCleanup:
Prune auto-router session rollup rows, which carry their own retention horizon.
"""
session_cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds))
from litellm.proxy.db.baseline_accounting import BaselineAccountingStore
if remaining_ms := self._remaining_timeout_ms(deadline)():
try:
await BaselineAccountingStore.for_client(prisma_client).retire_before(
session_cutoff,
self.batch_size,
remaining_ms,
)
except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job
verbose_proxy_logger.warning("Auto-router baseline retention remains pending")
sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline)
verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted)
return (sessions_result,)

3
litellm/proxy/guardrails/guardrail_hooks/akto/akto.py
View file

@ -232,8 +232,7 @@ class AktoGuardrail(CustomGuardrail):
"""
request_path: Final = self.extract_request_path(request_data)
request_headers: Final = self.build_request_headers(request_data)
request_inputs: Final = GenericGuardrailAPIInputs(model=inputs.get("model")) if include_response else inputs
request_body: Final = self.build_request_body(request_inputs, request_data)
request_body: Final = self.build_request_body(inputs, request_data)
tag: Final = self.build_tag_metadata(request_data)
response_payload = json.dumps({}) # Empty body wrapper when no response yet

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

5
litellm/proxy/guardrails/guardrail_hooks/crowdstrike_aidr/crowdstrike_aidr.py
View file

@ -425,7 +425,10 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
def _build_guard_input_for_response(self, inputs: GenericGuardrailAPIInputs) -> _GuardInput:
output_texts: Final[list[str]] = inputs.get("texts", [])
return _GuardInput(messages=[_Message(role="assistant", content=text) for text in output_texts], tools=[])
return _GuardInput(
messages=[_Message(role="assistant", content=text) for text in output_texts],
tools=inputs.get("tools", []),
)
def _extract_transformed_texts(self, guard_output: _GuardInput, num_assistant_messages: int) -> list[str]:
tail: Final = guard_output.messages[-num_assistant_messages:] if num_assistant_messages > 0 else []

2
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -286,7 +286,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
project_id: Final = headers.get("hl-project-id")
if input_type == "request" and (scan_params := inputs.get("structured_messages")):
if scan_params := inputs.get("structured_messages"):
last_msg: Final = scan_params[-1]
result: _HiddenlayerResponse = await self._call_hiddenlayer(
project_id,

2
litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py
View file

@ -197,7 +197,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
text_to_moderate: str | None = None
# Prefer structured_messages if available (has role context)
if input_type == "request" and (structured_messages := inputs.get("structured_messages")):
if structured_messages := inputs.get("structured_messages"):
text_to_moderate = self.get_user_prompt(structured_messages)
# Fall back to texts

2
litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py
View file

@ -129,7 +129,7 @@ class PromptGuardGuardrail(CustomGuardrail):
) -> GenericGuardrailAPIInputs:
texts: Final = inputs.get("texts", [])
images: Final = inputs.get("images", [])
structured_messages: Final = inputs.get("structured_messages") if input_type == "request" else None
structured_messages: Final = inputs.get("structured_messages", [])
model: Final = inputs.get("model")
if structured_messages:

2
litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py
View file

@ -452,7 +452,7 @@ class QualifireGuardrail(CustomGuardrail):
dynamic_params: Final = self.get_guardrail_dynamic_request_body_params(request_data=request_data)
# Extract messages from structured_messages or request_data
messages: list[AllMessageValues] | None = inputs.get("structured_messages") if input_type == "request" else None
messages: list[AllMessageValues] | None = inputs.get("structured_messages")
if not messages:
messages = request_data.get("messages")

5
litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py
View file

@ -380,12 +380,11 @@ class StraikerGuardrail(CustomGuardrail):
call_id: Final = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None
event_id: Final = f"{call_id or 'litellm'}:{input_type}"
is_request: Final = input_type == "request"
content: Final = StraikerWebhookContent(
texts=list(inputs.get("texts") or []),
images=list(inputs.get("images") or []),
structured_messages=_opaque_dict_list(inputs.get("structured_messages")) if is_request else None,
tools=_opaque_dict_list(inputs.get("tools")) if is_request else None,
structured_messages=_opaque_dict_list(inputs.get("structured_messages")),
tools=_opaque_dict_list(inputs.get("tools")),
tool_calls=_opaque_dict_list(inputs.get("tool_calls")),
)

2
litellm/proxy/hooks/__init__.py
View file

@ -2,6 +2,7 @@ import os
from typing import Final, Literal
from . import *
from .autorouter_baseline_cache import AutoRouterBaselineCache
from .cache_control_check import _PROXY_CacheControlCheck
from .litellm_skills import SkillsInjectionHook
from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
@ -25,6 +26,7 @@ PROXY_HOOKS: Final = {
"max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler,
"sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler,
"prompt_cache_prediction": PromptCacheObserver,
"autorouter_baseline_cache": AutoRouterBaselineCache,
}
## FEATURE FLAG HOOKS ##

344
litellm/proxy/hooks/autorouter_baseline_cache.py Normal file
View file

@ -0,0 +1,344 @@
from __future__ import annotations
import asyncio
import hashlib
import json
import time
from collections.abc import Callable, Mapping
from dataclasses import dataclass, replace
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
get_litellm_metadata_from_kwargs, # pyright: ignore[reportUnknownVariableType] # legacy metadata boundary validated below
)
from litellm.llms.anthropic.prompt_cache_prediction import (
CountedPromptCachePlan,
NativePredictionTarget,
TokenCounter,
UnsupportedCachePlan,
UnsupportedPredictionTarget,
count_cache_plan,
count_prompt_tokens,
parse_cache_plan,
resolve_baseline_prediction_target,
supported_baseline_recipient,
supported_prediction_headers,
)
from litellm.proxy.spend_tracking.baseline_accounting import BaselineObservation
from litellm.proxy.spend_tracking.savings import (
_effective_model_info, # pyright: ignore[reportPrivateUsage] # existing deployment-price owner
_proxy_llm_router, # pyright: ignore[reportPrivateUsage] # existing optional proxy-router owner
)
from litellm.types.router import BaselineRouteStamp
from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.utils import get_prompt_cache_min_tokens
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
_METADATA: Final = TypeAdapter(Mapping[str, object])
_PRICES: Final[TypeAdapter[ModelInfo | None]] = TypeAdapter(ModelInfo | None)
_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue])
_COUNT_TIMEOUT: Final = 3.0
_MAX_COUNTS: Final = 4096
class CapturedBaselineObservation(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
scope: str
api_key: str
session_id: str
router_name: str
baseline_model: str
model: str
prices: ModelInfo | None
observation: BaselineObservation
@dataclass(frozen=True, slots=True)
class BaselineCacheContext:
collector: AutoRouterBaselineCache
capture: CapturedBaselineObservation
target: NativePredictionTarget | UnsupportedPredictionTarget
baseline_deployment_id: str
invalidated: str | None = None
class _Metadata(BaseModel):
model_config = ConfigDict(strict=True, arbitrary_types_allowed=True)
route: BaselineRouteStamp = Field(alias="_autorouter_baseline_route")
user_api_key_hash: str = Field(min_length=1)
session_id: str | None = None
class _WireEvent(BaseModel):
model_config = ConfigDict(strict=True, arbitrary_types_allowed=True)
httpx_response: httpx.Response
api_call_start_time: datetime
completion_start_time: datetime
custom_llm_provider: str
stream: bool = False
prompt_cache_response_complete: bool = False
class _ResponseUsage(BaseModel):
model_config = ConfigDict(strict=True, from_attributes=True)
usage: Usage | None = None
def _digest(value: object) -> str:
return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
class AutoRouterBaselineCache(CustomLogger):
def __init__(
self,
prisma_client: PrismaClient | None,
router: Callable[[], Router | None] = _proxy_llm_router,
token_counter: TokenCounter | None = None,
clock: Callable[[], float] = time.time,
) -> None:
super().__init__() # pyright: ignore[reportUnknownMemberType] # legacy callback constructor
self.router: Final = router
self.token_counter: Final = token_counter
self.clock: Final = clock
self.count_slots: Final = asyncio.Semaphore(8)
self.counts: Mapping[str, tuple[int, float]] = MappingProxyType({})
async def async_pre_call_deployment_hook(self, kwargs: Mapping[str, object], call_type: CallTypes | None) -> None:
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj: Final = kwargs.get("litellm_logging_obj")
if not isinstance(logging_obj, Logging) or call_type != CallTypes.anthropic_messages:
return
try:
metadata: Final = _METADATA.validate_python(
get_litellm_metadata_from_kwargs(
{"litellm_params": kwargs} # mutable-ok: legacy metadata owner requires a dictionary
)
)
if metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
return
if logging_obj.baseline_cache_context is not None:
await invalidate_baseline_cache(logging_obj, "retried_request")
return
request: Final = _Metadata.model_validate(metadata)
session: Final = kwargs.get("litellm_session_id") or request.session_id or logging_obj.litellm_session_id
if not isinstance(session, str) or not session or len(session) > 256:
return
router: Final = self.router()
deployment: Final = router.get_deployment(request.route.baseline_deployment_id) if router else None
if deployment is None:
return
target: Final = resolve_baseline_prediction_target(deployment.litellm_params)
prices: Final = _PRICES.validate_python(
_effective_model_info(router, request.route.baseline_deployment_id, request.route.baseline_model)
)
scope: Final = "autorouter-baseline:v3:" + _digest(
(
request.user_api_key_hash,
session,
request.route.router_name,
request.route.baseline_deployment_id,
deployment.litellm_params.model_dump(mode="json"),
prices,
)
)
started: Final = logging_obj.start_time.timestamp()
capture: Final = CapturedBaselineObservation(
scope=scope,
api_key=request.user_api_key_hash,
session_id=session,
router_name=request.route.router_name,
baseline_model=request.route.baseline_model,
model=target.model if isinstance(target, NativePredictionTarget) else request.route.baseline_model,
prices=prices,
observation=BaselineObservation(
request_id=logging_obj.litellm_call_id,
started_at=started,
available_at=started,
outcome="uncertain",
baseline_equivalent=False,
reason="incomplete_response",
),
)
logging_obj.baseline_cache_context = BaselineCacheContext(
self, capture, target, request.route.baseline_deployment_id
)
except Exception: # noqa: BLE001 # optional observation cannot fail inference
verbose_proxy_logger.warning("Auto-router baseline observation could not be initialized")
async def _count(self, target: NativePredictionTarget, body: Mapping[str, JsonValue]) -> int | None:
key: Final = _digest((target.model, target.api_key, target.api_base, _JSON_BODY.validate_python(body)))
now: Final = self.clock()
cached: Final = self.counts.get(key)
if cached is not None and cached[1] > now:
return cached[0]
async with self.count_slots:
tokens: Final = (
await self.token_counter(target.model, target.api_key, body)
if self.token_counter is not None
else await count_prompt_tokens(target.model, target.api_key, body, api_base=target.api_base)
)
if tokens is None or tokens < 0:
return None
retained: Final = tuple((k, v) for k, v in self.counts.items() if v[1] > now and k != key)[-(_MAX_COUNTS - 1) :]
self.counts = MappingProxyType(dict((*retained, (key, (tokens, now + 3600)))))
return tokens
async def plan(
self, target: NativePredictionTarget, wire: httpx.Request, body: Mapping[str, JsonValue], usage: Usage | None
) -> tuple[CountedPromptCachePlan | None, str | None]:
if not supported_prediction_headers(wire.headers):
return None, "unsupported_request_headers"
plan: Final = parse_cache_plan(body)
if isinstance(plan, UnsupportedCachePlan):
return None, plan.reason
details: Final = usage.prompt_tokens_details if usage is not None else None
if (
not plan.breakpoints
and details is not None
and ((details.cached_tokens or 0) + (details.cache_creation_tokens or 0))
):
return None, "implicit_cache_without_breakpoints"
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
return await self._count(target, body)
try:
counted: Final = await asyncio.wait_for(
count_cache_plan(target.model, target.api_key, plan, token_counter=count), timeout=_COUNT_TIMEOUT
)
return (None, counted.reason) if isinstance(counted, UnsupportedCachePlan) else (counted, None)
except TimeoutError:
return None, "token_count_timeout"
except Exception: # noqa: BLE001 # token counting cannot fail a completed request
return None, "token_count_unavailable"
async def invalidate_baseline_cache(logging_obj: Logging, reason: str, *, completed: bool = False) -> None:
context: Final = logging_obj.baseline_cache_context
if context is not None:
logging_obj.baseline_cache_context = replace(
context, invalidated=reason
) # rebind-ok: request-owned retry marker
logging_obj.baseline_observation = context.capture.model_copy(
update=MappingProxyType(
{ # rebind-ok: capture uncertainty for failure logging
"observation": context.capture.observation.model_copy(
update=MappingProxyType(
{
"available_at": max(context.capture.observation.started_at, context.collector.clock()),
"reason": reason,
}
)
),
}
)
)
async def finalize_baseline_cache(logging_obj: Logging, response_obj: object) -> None:
context: Final = logging_obj.baseline_cache_context
if context is None:
return
try:
capture: Final = await _capture(context, logging_obj, response_obj)
if logging_obj.baseline_cache_context is context:
logging_obj.baseline_observation = capture # rebind-ok: attach only to the captured request owner
except Exception: # noqa: BLE001 # observation failures must preserve inference and billing
await invalidate_baseline_cache(logging_obj, "observation_unavailable")
async def _capture(
context: BaselineCacheContext, logging_obj: Logging, response_obj: object
) -> CapturedBaselineObservation:
original: Final = context.capture.observation
details: Final = _METADATA.validate_python(logging_obj.model_call_details)
if details.get("cache_hit") is True:
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": original.model_copy(
update=MappingProxyType({"outcome": "response_cache", "reason": "response_cache_hit"})
)
}
)
)
event: Final = _WireEvent.model_validate(details)
wire: Final = event.httpx_response.request
usage: Final = _ResponseUsage.model_validate(response_obj).usage
complete: Final = (
event.custom_llm_provider == "anthropic"
and event.httpx_response.status_code == 200
and (not event.stream or event.prompt_cache_response_complete)
)
started: Final = original.started_at
available: Final = event.completion_start_time.timestamp()
if context.invalidated or not complete or not started <= available <= context.collector.clock():
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": original.model_copy(
update=MappingProxyType(
{
"available_at": max(started, context.collector.clock()),
"reason": context.invalidated or "incomplete_response",
}
)
)
}
)
)
target: Final = context.target
if isinstance(target, UnsupportedPredictionTarget) or not supported_baseline_recipient(target, wire):
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": original.model_copy(
update=MappingProxyType(
{
"available_at": available,
"reason": target.reason
if isinstance(target, UnsupportedPredictionTarget)
else "unsupported_baseline_recipient",
}
)
)
}
)
)
body: Final = _JSON_BODY.validate_json(wire.content)
same: Final = (
logging_obj.get_router_model_id() == context.baseline_deployment_id and body.get("model") == target.model
)
plan, reason = await context.collector.plan(target, wire, body, usage)
minimum: Final = get_prompt_cache_min_tokens(target.model)
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": BaselineObservation(
request_id=original.request_id,
started_at=started,
available_at=available,
outcome="complete",
baseline_equivalent=same,
usage=usage,
plan=plan,
minimum_cache_tokens=minimum,
reason=reason,
)
}
)
)

42
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -558,6 +558,9 @@ class _SessionAggRow(BaseModel):
total_tokens: int
spend: float
saved_spend: float
savings_estimated_turns: int = 0
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
classifier_cost: float
classifier_cost_recorded_turns: int
session_seconds: float
@ -584,9 +587,19 @@ def _cache_bucket(turns: int, hits: int) -> AutoRouterCacheBucket:
return AutoRouterCacheBucket(turns=turns, hits=hits, hit_rate_pct=_pct(hits, turns))
def _savings_cohort(
turns: int, estimated_turns: int, actual_spend: float, saved_spend: float
) -> tuple[float | None, float | None]:
if turns > 0 and estimated_turns == 0:
return None, None
return saved_spend, actual_spend + saved_spend
def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
return_misses: Final = row.return_turns - row.return_hits
baseline_spend: Final = row.spend + row.saved_spend
saved_spend, baseline_spend = _savings_cohort(
row.turns, row.savings_estimated_turns, row.savings_estimated_actual_spend, row.savings_estimated_saved_spend
)
sessions: Final = row.sessions
return AutoRouterBenchmarkTotals(
sessions=sessions,
@ -595,11 +608,15 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
avg_session_seconds=row.session_seconds / sessions if sessions else 0.0,
avg_tokens_per_session=row.total_tokens / sessions if sessions else 0.0,
spend=row.spend,
saved_spend=row.saved_spend,
savings_estimated_turns=row.savings_estimated_turns,
savings_estimated_actual_spend=row.savings_estimated_actual_spend,
saved_spend=saved_spend,
classifier_cost=row.classifier_cost if row.classifier_cost_recorded_turns == row.turns else None,
baseline_spend=baseline_spend,
saved_pct=_pct(row.saved_spend, baseline_spend),
saved_per_session=row.saved_spend / sessions if sessions else 0.0,
saved_pct=_pct(saved_spend, baseline_spend) if saved_spend is not None and baseline_spend is not None else None,
saved_per_session=(row.savings_estimated_saved_spend / sessions if sessions else 0.0)
if row.savings_estimated_turns == row.turns
else None,
cache=AutoRouterCacheStats(
coverage_pct=_pct(row.covered_turns, row.turns),
hit_rate_pct=_pct(row.cache_hits, row.covered_turns),
@ -629,6 +646,8 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup:
avg_tokens_per_session=totals.avg_tokens_per_session,
spend=totals.spend,
saved_spend=totals.saved_spend,
savings_estimated_turns=totals.savings_estimated_turns,
savings_estimated_actual_spend=totals.savings_estimated_actual_spend,
classifier_cost=totals.classifier_cost,
baseline_spend=totals.baseline_spend,
saved_pct=totals.saved_pct,
@ -660,6 +679,9 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
total_tokens=sum(row.total_tokens for row in rows),
spend=sum(row.spend for row in rows),
saved_spend=sum(row.saved_spend for row in rows),
savings_estimated_turns=sum(row.savings_estimated_turns for row in rows),
savings_estimated_actual_spend=sum(row.savings_estimated_actual_spend for row in rows),
savings_estimated_saved_spend=sum(row.savings_estimated_saved_spend for row in rows),
classifier_cost=sum(row.classifier_cost for row in rows),
classifier_cost_recorded_turns=sum(row.classifier_cost_recorded_turns for row in rows),
session_seconds=sum(row.session_seconds for row in rows),
@ -809,6 +831,9 @@ async def get_auto_router_session(
raise HTTPException(
status_code=404, detail=f"No auto-routed turns recorded for session {session_id!r} under this key"
)
saved_spend, baseline_spend = _savings_cohort(
row.turns, row.savings_estimated_turns, row.savings_estimated_actual_spend, row.savings_estimated_saved_spend
)
return AutoRouterSessionResponse(
session_id=session_id,
router_name=row.router_name,
@ -816,10 +841,13 @@ async def get_auto_router_session(
turns=row.turns,
last_model=row.last_model,
spend=row.spend,
saved_spend=row.saved_spend,
baseline_spend=row.spend + row.saved_spend,
savings_estimated_turns=row.savings_estimated_turns,
savings_estimated_actual_spend=row.savings_estimated_actual_spend,
saved_spend=saved_spend,
baseline_spend=baseline_spend if row.savings_estimated_turns == row.turns else None,
savings_estimated_baseline_spend=baseline_spend,
baseline_model=row.baseline_model,
baseline_models=row.baseline_models,
baseline_models=row.savings_estimated_baseline_models,
)

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