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

Merge branch 'litellm_internal_staging' into litellm_mcp_lifecycle_e2e

Browse source
This commit is contained in:
yuneng-jiang 2026-09-08 13:55:51 -07:00 • committed by GitHub
commit 1625f8e7f5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
427 changed files with 19288 additions and 5024 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
.circleci/config.yml
View file

@ -1440,6 +1440,7 @@ jobs:
TEST_FILES=$(printf "%s\n" \
tests/local_testing/test_dual_cache.py \
tests/local_testing/test_redis_batch_optimizations.py \
tests/local_testing/test_redis_increment_with_floor.py \
tests/local_testing/test_router_utils.py)
echo "$TEST_FILES" | circleci tests run \
--verbose \
@ -2648,6 +2649,19 @@ jobs:
name: Start mock LLM server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
background: true
- run:
name: Start mock Presidio server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
background: true
- run:
name: Wait for mock Presidio server
command: |
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
sleep 1
done
echo "Mock Presidio server never answered /health on port 8091" >&2
exit 1
- run:
name: Start LiteLLM proxy
environment:
@ -2778,6 +2792,19 @@ jobs:
name: Start mock LLM server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
background: true
- run:
name: Start mock Presidio server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
background: true
- run:
name: Wait for mock Presidio server
command: |
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
sleep 1
done
echo "Mock Presidio server never answered /health on port 8091" >&2
exit 1
- run:
name: Start LiteLLM proxy under a server root path
environment:

45
.github/workflows/cost-map-guard.yml vendored Normal file
View file

@ -0,0 +1,45 @@
name: Cost map guard
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed
pull_request_target:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
cancel-in-progress: true
jobs:
cost-map-guard:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Fetch the pull request head and its merge base
id: revisions
env:
GH_TOKEN: ${{ github.token }}
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')"
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Run the guard
env:
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
HEAD_REF: ${{ github.event.pull_request.head.ref }}
run: |
uv run --frozen python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"

6
.github/workflows/publish-basedpyright-base-counts.yml vendored
View file

@ -1,6 +1,6 @@
name: Publish basedpyright base counts
# Every commit on litellm_internal_staging is some branch's future merge-base.
# Every commit on main or litellm_internal_staging can become a future merge-base.
# Publishing its per-rule basedpyright counts as an artifact lets
# scripts/type_check_gate.py download them in seconds instead of paying a
# 60-110s second basedpyright pass on every fresh worktree or moved merge-base.
@ -10,13 +10,13 @@ name: Publish basedpyright base counts
on:
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
inputs:
ref:
description: "Ref to compute and publish base counts for"
description: "Ref to compute and publish base counts for (defaults to the workflow run's commit)"
required: false
default: litellm_internal_staging
permissions:
contents: read

6
.github/workflows/sync-together-ai-models.yml vendored
View file

@ -13,10 +13,12 @@ jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
env:
BASE_BRANCH: ${{ github.event.repository.default_branch }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: litellm_internal_staging
ref: ${{ env.BASE_BRANCH }}
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
@ -63,6 +65,6 @@ jobs:
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base litellm_internal_staging
--base "$BASE_BRANCH"
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

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

@ -12,6 +12,7 @@ on:
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
concurrency:

37
.github/workflows/test-model-map.yml vendored
View file

@ -1,37 +0,0 @@
name: Validate model_prices_and_context_window.json
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
validate-model-prices-json:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Validate model_prices_and_context_window.json
run: |
jq empty model_prices_and_context_window.json
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Check model_prices_and_context_window.schema.json is in sync
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py --check

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

@ -117,6 +117,9 @@ jobs:
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- name: Run pytest tests/test_litellm_rust with the compiled extension
run: make test-rust-extension
- run: >-
uv build --wheel --out-dir panic-dist
--config-setting "maturin.build-args=--features panic-test,extension-module"

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

@ -116,6 +116,7 @@ jobs:
tests/test_litellm/rerank_api
tests/test_litellm/rust_bridge
tests/test_litellm/sandbox
tests/test_litellm/skills
tests/test_litellm/test_router
tests/test_litellm/vector_stores
tests/test_litellm/videos

16
ARCHITECTURE.md
View file

@ -149,7 +149,7 @@ graph TD
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |
| `litellm_skills` | `proxy/hooks/skills_injection.py` | Skills injection |
| `litellm_skills` | `proxy/hooks/litellm_skills/main.py` | Skills injection |
To add a new proxy hook, implement `CustomLogger` and register in `PROXY_HOOKS`.
@ -220,20 +220,20 @@ graph LR
| Job | Interval | Purpose | Key Files |
|-----|----------|---------|-----------|
| `update_spend` | 60s | Batch write spend logs to PostgreSQL | `proxy/db/db_spend_update_writer.py` |
| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/management_helpers/budget_reset_job.py` |
| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/common_utils/reset_budget_job.py` |
| `add_deployment` | 10s | Sync new model deployments from DB | `proxy/proxy_server.py` (`ProxyConfig`) |
| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/management_helpers/spend_log_cleanup.py` |
| `check_batch_cost` | 30min | Calculate costs for batch jobs | `proxy/management_helpers/check_batch_cost_job.py` |
| `check_responses_cost` | 30min | Calculate costs for responses API | `proxy/management_helpers/check_responses_cost_job.py` |
| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/management_helpers/key_rotation_manager.py` |
| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/db/db_transaction_queue/spend_log_cleanup.py` |
| `check_batch_cost` | 30min | Calculate costs for batch jobs | `enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py` |
| `check_responses_cost` | 30min | Calculate costs for responses API | `enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py` |
| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/common_utils/key_rotation_manager.py` |
| `_run_background_health_check` | continuous | Health check model deployments | `proxy/proxy_server.py` |
| `send_weekly_spend_report` | weekly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) |
| `send_monthly_spend_report` | monthly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) |
**Cost Attribution Flow:**
1. LLM response returns to `utils.py` wrapper after `litellm.acompletion()` completes
2. `update_response_metadata()` (`llm_response_utils/response_metadata.py`) is called
3. `logging_obj._response_cost_calculator()` (`litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`)
2. `update_response_metadata()` (`litellm_core_utils/llm_response_utils/response_metadata.py`) is called
3. `logging_obj._response_cost_calculator()` (`litellm_core_utils/litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`)
4. Cost is stored in `response._hidden_params["response_cost"]`
5. `proxy/common_request_processing.py` extracts cost from `hidden_params` and adds to response headers (`x-litellm-response-cost`)
6. `logging_obj.async_success_handler()` triggers callbacks including `_ProxyDBLogger.async_log_success_event()`

6
CLAUDE.md
View file

@ -29,7 +29,7 @@ Never test structure of code only function of it
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on `litellm_internal_staging` in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. If your branch already carries a budget edit, drop it before opening the PR
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
@ -70,7 +70,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch

4
CONTRIBUTING.md
View file

@ -315,10 +315,12 @@ Ensure the UI builds successfully before submitting your PR:
npm run build
```
Local lint and budget checks follow origin's current default branch. They refresh it from the remote instead of trusting cached `origin/HEAD`. For an intentional comparison against another branch or commit, use `make check BASE_REF=<ref>` or the standalone gate's `--base <ref>` option. An explicit ref can also be used offline once it has been fetched locally. Without an override, unavailable remote metadata stops the check
## Submitting Your PR
1. **Push your branch**: `git push origin your-feature-branch`
2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`.
2. **Create a PR**: Go to GitHub and open a pull request against the repository's current default branch. Run `python3 scripts/default_branch.py --branch` to check its name
3. **Fill out the PR template**: Provide clear description of changes
4. **Wait for review**: Maintainers will review and provide feedback
5. **Address feedback**: Make requested changes and push updates

73
Makefile
View file

@ -4,6 +4,7 @@
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
test-rust-extension \
info lint lint-inner lint-dev lint-checks format \
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
@ -34,7 +35,7 @@ help:
@echo " make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed"
@echo " make lint-format - Check ruff format formatting (matches CI)"
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)"
@echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
@echo " make lint-test-quality - Gate the test suite against test-quality-budget.json"
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)"
@ -54,12 +55,16 @@ help:
@echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)"
@echo " make test-integration - Run integration tests"
@echo " make test-unit-helm - Run helm unit tests"
@echo " make test-rust-extension - Build the Rust extension and run its public Python tests"
@echo ""
@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
UV := uv
UV_RUN := $(UV) run --no-sync
BASE_REF ?=
export BASE_REF
RESOLVE_BASE = python3 scripts/default_branch.py --base "$(BASE_REF)"
# Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so
# it runs before any venv exists. See scripts/gate_slot_lock.py.
@ -67,7 +72,7 @@ GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py
LINT_DEP_INSTALL ?= install-dev
LINT_E2E_DEP_INSTALL ?= lint-install
LINT_DEP_BASE ?= lint-fetch-base
LINT_DEP_BASE ?=
LINT_JOBS := $(shell sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4)
LINT_OUTPUT_SYNC := $(if $(filter output-sync,$(.FEATURES)),--output-sync=target,)
@ -130,10 +135,8 @@ format: install-dev
format-check: install-dev
cd litellm && $(UV_RUN) ruff format --check --exclude '/enterprise/' . && cd ..
# Single fetch of the PR base so the delta-based gates below share one network round
# trip instead of each re-fetching when chained from `lint`.
lint-fetch-base:
git fetch origin litellm_internal_staging
@$(RESOLVE_BASE)
# Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated
# Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The
@ -150,7 +153,9 @@ lint-install:
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
@files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \
@base_ref=$$($(RESOLVE_BASE)) && \
changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \
files=$$(printf '%s\n' "$$changed" | grep -v '^litellm/enterprise/' || true) || exit $$?; \
if [ -z "$$files" ]; then \
echo "No changed litellm Python files to format-check."; \
else \
@ -167,7 +172,9 @@ lint-ruff: $(LINT_DEP_INSTALL)
# https://github.com/astral-sh/ruff/discussions/10977
# https://github.com/astral-sh/ruff/discussions/4049
lint-format-changed: install-dev
@git diff origin/main --unified=0 --no-color -- '*.py' | \
@base_ref=$$($(RESOLVE_BASE)) && \
diff=$$(git diff "$$base_ref" --unified=0 --no-color -- '*.py') && \
printf '%s\n' "$$diff" | \
perl -ne '\
if (/^diff --git a\/(.*) b\//) { $$file = $$1; } \
if (/^@@ .* \+(\d+)(?:,(\d+))? @@/) { \
@ -182,20 +189,22 @@ lint-format-changed: install-dev
done
lint-ruff-dev: install-dev
@tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
@base_ref=$$($(RESOLVE_BASE)) || exit $$?; \
tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
cd litellm && \
($(UV_RUN) ruff check . --output-format=pylint || true) > "$$tmpfile" && \
$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch=origin/main && \
$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch="$$base_ref" && \
cd .. ; \
rm -f "$$tmpfile"
lint-ruff-FULL-dev: install-dev
@files=$$(git diff --name-only origin/main -- '*.py'); \
@base_ref=$$($(RESOLVE_BASE)) && \
files=$$(git diff --name-only "$$base_ref" -- '*.py') || exit $$?; \
if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
else echo "No changed .py files to check."; fi
lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/type_check_gate.py --base "$(BASE_REF)"
lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
$(UV_RUN) basedpyright tests/e2e
@ -203,37 +212,37 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
# Type-discipline budget (mutable collections / casts / type guards / kwargs /
# unexplained suppressions), the test-linting.yml step `make lint` used to omit.
lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/type_discipline_gate.py --base "$(BASE_REF)"
# Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes,
# litellm module-global mutation, credential-gated skips, conftest snapshot
# inventory), counted across tests/ the same delta-vs-base way.
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)"
# --update lowers each limit by what this branch fixed since its branch point, so
# it needs the base ref fetched to resolve the merge-base.
lint-basedpyright-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_check_gate.py --update
lint-basedpyright-budget-update: install-dev
$(UV_RUN) python scripts/type_check_gate.py --update --base "$(BASE_REF)"
lint-format: format-check
lint-ruff-budget: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py
$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
# Strict gate, invoked the same way CI does in test-linting.yml so a local pass
# means the CI check will pass too.
lint-gate: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
lint-ruff-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/ruff_strict_gate.py --update
lint-ruff-budget-update: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py --update --base "$(BASE_REF)"
lint-type-discipline-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_discipline_gate.py --update
lint-type-discipline-budget-update: install-dev
$(UV_RUN) python scripts/type_discipline_gate.py --update --base "$(BASE_REF)"
lint-test-quality-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/test_quality_gate.py --update
lint-test-quality-budget-update: install-dev
$(UV_RUN) python scripts/test_quality_gate.py --update --base "$(BASE_REF)"
# Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update
@ -249,14 +258,15 @@ check-import-safety: $(LINT_DEP_INSTALL)
# runs the diff-scoped ruff format check, whole-tree ruff check, the strict-rule /
# type-discipline / basedpyright budgets as a delta vs the base, then the circular-import
# and import-safety checks. Steps that compare against the base resolve it the same way CI
# does (merge-base with origin/litellm_internal_staging). Setup (env sync, Prisma client,
# does (merge-base with origin's current default branch). Setup (env sync, Prisma client,
# base fetch) runs once up front; the checks themselves are independent, so a sub-make
# fans them out with -j and the fast ones finish under basedpyright's shadow.
lint:
@$(GATE_SLOT_LOCK) $(MAKE) lint-inner
lint-inner: lint-install lint-fetch-base
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-inner: lint-install
@base_ref=$$($(RESOLVE_BASE)) && \
$(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
@ -281,6 +291,17 @@ pre-commit:
@$(MAKE) check
# Testing targets
test-rust-extension:
@temporary=$$(mktemp -d) && \
trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
set -- "$$temporary"/wheels/*.whl && \
[ "$$#" -eq 1 ] && \
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
test: install-test-deps
$(UV_RUN) pytest tests/

146
ci_cd/cost_map_guard.py Normal file
View file

@ -0,0 +1,146 @@
"""Guard the cost map on pull requests.
Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file,
and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named
litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models.
"""
from __future__ import annotations
import argparse
import json
import subprocess
import sys
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
CostMap = dict[str, object]
@dataclass(frozen=True, slots=True)
class Snapshot:
cost_map: str
backup: str
schema: str
def _parse_object(text: str, path: str) -> CostMap | str:
try:
parsed: Final = json.loads(text)
except json.JSONDecodeError as error:
return f"{path} is not valid JSON: {error}"
return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
def _rendered_schema(cost_map: CostMap) -> str:
try:
return render(build_schema(cost_map))
except SystemExit as error:
return str(error)
def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]:
schema_text: Final = _rendered_schema(head_map)
if not schema_text.startswith("{"):
return (schema_text,)
backup_failure: Final = (
()
if head.backup == head.cost_map
else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
)
schema_failure: Final = (
()
if head.schema == schema_text
else (
f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
"run `python ci_cd/generate_model_prices_schema.py` and commit the result",
)
)
return (
*backup_failure,
*schema_failure,
*(
f"{COST_MAP_PATH} does not validate against its schema: {error}"
for error in validation_errors(head_map, json.loads(schema_text))[:20]
),
)
def _entries(cost_map: CostMap) -> dict[str, dict[str, object]]:
return {key: entry for key, entry in cost_map.items() if isinstance(entry, dict)}
def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str]) -> tuple[str, ...]:
base_map: Final = _parse_object(base.cost_map, COST_MAP_PATH)
if isinstance(base_map, str):
return (f"merge base: {base_map}",)
base_entries: Final = _entries(base_map)
head_entries: Final = _entries(head_map)
removed_fields: Final = tuple(
f"{key}.{field}"
for key, entry in base_entries.items()
if key in head_entries
for field in entry
if field not in head_entries[key]
)
return (
*(
f"bot PRs may only change the cost map files, not {path}"
for path in changed_files
if path not in GUARDED_PATHS
),
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),
*(
f"bot PRs may not change {key}"
for key in sorted(SPECIAL_ROOT_KEYS)
if base_map.get(key) != head_map.get(key)
),
)
def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]:
head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH)
if isinstance(head_map, str):
return (head_map,)
return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ()))
def _git(*args: str) -> str:
result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True)
return result.stdout if result.returncode == 0 else ""
def snapshot(revision: str) -> Snapshot:
return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS))
def main(argv: Sequence[str]) -> int:
parser: Final = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--base", required=True, help="merge base of the pull request")
parser.add_argument("--head", required=True, help="head commit of the pull request")
parser.add_argument("--head-ref", required=True, help="head branch name of the pull request")
args: Final = parser.parse_args(argv)
bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX)
changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines())
failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot)
contract: Final = "bot contract enforced" if bot else "human PR, file checks only"
if failures:
print(f"cost map guard failed ({contract}):")
print("\n".join(f"- {failure}" for failure in failures))
return 1
print(f"cost map guard passed ({contract})")
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))

62
ci_cd/run_migration.py
View file

@ -6,12 +6,9 @@ import subprocess
import sys
from datetime import datetime
from pathlib import Path
import testing.postgresql
from typing import Final
DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE)
DEFAULT_BASE_BRANCH = "litellm_internal_staging"
def _find_destructive_statements(sql: str) -> list:
@ -94,31 +91,57 @@ def _print_stale_branch_refusal(base_branch: str, behind: int) -> None:
print(banner, file=out)
def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
def _default_base_branch(root_dir: Path) -> str:
try:
result: Final = subprocess.run(
[
sys.executable,
str(Path(__file__).resolve().parents[1] / "scripts" / "default_branch.py"),
"--repo-root",
str(root_dir),
"--branch",
],
check=True,
capture_output=True,
text=True,
timeout=90,
)
except (OSError, subprocess.SubprocessError) as exc:
_print_freshness_failure(
"default branch",
"Could not discover origin's default branch. Pass --base-branch <name> to choose one.",
exc.stderr if isinstance(exc, subprocess.CalledProcessError) else str(exc),
)
sys.exit(3)
return result.stdout.strip()
def _check_branch_freshness(root_dir: Path, base_branch: str | None = None) -> None:
"""Fetch origin/<base_branch> and exit 3 if HEAD is behind it."""
resolved_branch: Final = base_branch or _default_base_branch(root_dir)
cwd = str(root_dir)
try:
subprocess.run(
["git", "fetch", "origin", base_branch],
["git", "fetch", "origin", f"+refs/heads/{resolved_branch}:refs/remotes/origin/{resolved_branch}"],
check=True,
capture_output=True,
text=True,
cwd=cwd,
)
except FileNotFoundError:
_print_freshness_failure(base_branch, "git executable not found on PATH")
_print_freshness_failure(resolved_branch, "git executable not found on PATH")
sys.exit(3)
except subprocess.CalledProcessError as e:
_print_freshness_failure(
base_branch,
f"`git fetch origin {base_branch}` failed",
resolved_branch,
f"`git fetch origin {resolved_branch}` failed",
e.stderr or "",
)
sys.exit(3)
try:
result = subprocess.run(
["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"],
["git", "rev-list", "--count", f"HEAD..origin/{resolved_branch}"],
check=True,
capture_output=True,
text=True,
@ -127,23 +150,23 @@ def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
behind = int(result.stdout.strip())
except subprocess.CalledProcessError as e:
_print_freshness_failure(
base_branch,
f"`git rev-list HEAD..origin/{base_branch}` failed",
resolved_branch,
f"`git rev-list HEAD..origin/{resolved_branch}` failed",
e.stderr or "",
)
sys.exit(3)
except ValueError:
_print_freshness_failure(
base_branch,
resolved_branch,
"could not parse commit count from `git rev-list`",
)
sys.exit(3)
if behind > 0:
_print_stale_branch_refusal(base_branch, behind)
_print_stale_branch_refusal(resolved_branch, behind)
sys.exit(3)
print(f"Branch freshness OK: up to date with origin/{base_branch}.")
print(f"Branch freshness OK: up to date with origin/{resolved_branch}.")
def _print_destructive_refusal(destructive_lines: list) -> None:
@ -198,7 +221,7 @@ def _print_destructive_refusal(destructive_lines: list) -> None:
def create_migration(
migration_name: str = None,
allow_destructive: bool = False,
base_branch: str = DEFAULT_BASE_BRANCH,
base_branch: str | None = None,
skip_freshness_check: bool = False,
):
"""
@ -211,7 +234,7 @@ def create_migration(
DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this
flag, the script exits non-zero and prints guidance.
base_branch (str): Branch to check freshness against
(default: "litellm_internal_staging").
(default: origin's current default branch).
skip_freshness_check (bool): Skip the "branch is up to date" check.
Only for intentional migrations against an older base.
"""
@ -225,6 +248,8 @@ def create_migration(
else:
_check_branch_freshness(root_dir, base_branch)
import testing.postgresql
try:
migrations_dir = (
root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
@ -342,9 +367,8 @@ if __name__ == "__main__":
)
parser.add_argument(
"--base-branch",
default=DEFAULT_BASE_BRANCH,
help=(
f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). "
"Branch to check freshness against (default: origin's current default branch). "
"The script fetches origin/<base-branch> and refuses to run if HEAD "
"is behind it."
),

6
docker/component_entrypoint.sh
View file

@ -1,5 +1,11 @@
#!/bin/sh
# stale samples from a previous container incarnation would be summed into the aggregate
if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then
mkdir -p "$PROMETHEUS_MULTIPROC_DIR"
rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db
fi
case "$USE_DDTRACE" in
[Tt][Rr][Uu][Ee])
export DD_TRACE_OPENAI_ENABLED="False"

12
helm/litellm-helm/templates/deployment.yaml
View file

@ -152,6 +152,13 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
@ -189,6 +196,11 @@ spec:
- name: http
containerPort: {{ .Values.service.port }}
protocol: TCP
{{- if .Values.metricsServer.enabled }}
- name: metrics
containerPort: {{ .Values.metricsServer.port }}
protocol: TCP
{{- end }}
livenessProbe:
httpGet:
path: {{ .Values.livenessProbe.path | quote }}

17
helm/litellm-helm/templates/service-metrics.yaml Normal file
View file

@ -0,0 +1,17 @@
{{- if .Values.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.fullname" . }}-metrics
labels:
{{- include "litellm.labels" . | nindent 4 }}
spec:
type: ClusterIP
ports:
- port: {{ .Values.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.selectorLabels" . | nindent 4 }}
{{- end }}

2
helm/litellm-helm/templates/servicemonitor.yaml
View file

@ -26,7 +26,7 @@ spec:
{{- toYaml .namespaceSelector.matchNames | nindent 4 }}
{{- end }}
endpoints:
- port: http
- port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }}
path: /metrics/
interval: {{ .interval }}
scrapeTimeout: {{ .scrapeTimeout }}

106
helm/litellm-helm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,106 @@
suite: separate metrics server
templates:
- configmap-litellm.yaml
- deployment.yaml
- service.yaml
- service-metrics.yaml
- servicemonitor.yaml
tests:
- it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default
asserts:
- notContains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
any: true
template: deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- hasDocuments:
count: 0
template: service-metrics.yaml
- it: should scrape the proxy port when the metrics server is disabled
template: servicemonitor.yaml
set:
serviceMonitor.enabled: true
asserts:
- equal:
path: spec.endpoints[0].port
value: http
- it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor
set:
metricsServer.enabled: true
metricsServer.port: 4101
serviceMonitor.enabled: true
service.type: LoadBalancer
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "4101"
template: deployment.yaml
- contains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
containerPort: 4101
protocol: TCP
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-metrics
template: service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
template: service-metrics.yaml
- equal:
path: spec.endpoints[0].port
value: metrics
template: servicemonitor.yaml
- equal:
path: spec.endpoints[0].path
value: /metrics/
template: servicemonitor.yaml
- it: should reject a metrics port equal to the proxy port
template: deployment.yaml
set:
metricsServer.enabled: true
metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: metricsServer.port must differ from service.port

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

@ -180,6 +180,16 @@ proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT)
# so a scrape never runs on an inference worker. Adds a `metrics` port to the
# container and a dedicated ClusterIP `<release>-metrics` Service, and the
# ServiceMonitor scrapes it instead of the proxy port. The separate port has
# no virtual-key auth: keep it off public ingress. Needs the proxy image
# v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
{}
# Unset by default so the chart installs on small clusters such as Minikube, and so an

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

@ -441,3 +441,5 @@ ImplementationSpecific
{{- .pathType -}}
{{- end -}}
{{- end -}}
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}

53
helm/litellm/templates/gateway/deployment.yaml
View file

@ -64,14 +64,25 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
{{- if eq (int .Values.gateway.metricsServer.port) 4000 }}
{{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }}
{{- end }}
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
@ -97,16 +108,54 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: metrics
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- {{ .Values.gateway.metricsServer.port | quote }}
env:
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
ports:
- name: metrics
containerPort: {{ .Values.gateway.metricsServer.port }}
protocol: TCP
volumeMounts:
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
readinessProbe:
tcpSocket: { port: metrics }
periodSeconds: 10
livenessProbe:
tcpSocket: { port: metrics }
periodSeconds: 15
failureThreshold: 6
resources:
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
{{- end }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
emptyDir: {}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}

18
helm/litellm/templates/gateway/service-metrics.yaml Normal file
View file

@ -0,0 +1,18 @@
{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.gateway.fullname" . }}-metrics
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
spec:
type: ClusterIP
ports:
- port: {{ .Values.gateway.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.gateway.selectorLabels" . | nindent 4 }}
{{- end }}

148
helm/litellm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,148 @@
suite: test gateway metrics sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/service.yaml
- gateway/service-metrics.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, volume, env or service port when the metrics server is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
any: true
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- hasDocuments:
count: 0
template: gateway/service-metrics.yaml
- it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4101
gateway.service.type: LoadBalancer
gateway.image.tag: v1.101.0
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.101.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- "4101"
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].env
value:
- name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].ports
value:
- name: metrics
containerPort: 4101
protocol: TCP
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].volumeMounts
value:
- name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 50m
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
emptyDir: {}
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: gateway/service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-gateway-metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: gateway/service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
template: gateway/service-metrics.yaml
- it: rejects a metrics port equal to the gateway port
template: gateway/deployment.yaml
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: gateway.metricsServer.port must differ from the gateway port 4000

16
helm/litellm/values.yaml
View file

@ -268,6 +268,22 @@ gateway:
config:
create: true
proxy_config: {}
# Serve Prometheus /metrics from a `metrics` sidecar container (same image,
# `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the
# workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a
# scrape never runs on an inference worker. Adds a `metrics` port to the pod
# and a dedicated ClusterIP `<gateway>-metrics` Service; point your scrape
# config at it. The port has no virtual-key auth: keep it off public ingress.
# Needs the gateway image v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
requests:
cpu: 50m
memory: 128Mi
limits:
memory: 512Mi
image:
repository: ghcr.io/berriai/litellm-gateway
tag: "" # defaults to .Chart.AppVersion

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql Normal file
View file

@ -0,0 +1,3 @@
ALTER TABLE "LiteLLM_AutoRouterSession"
ADD COLUMN IF NOT EXISTS "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0;

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

@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
@@id([api_key, session_id, router_name])

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

@ -8,14 +8,10 @@ import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
from typing import TYPE_CHECKING, Final, Optional
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
prisma_command_timeout,
prisma_migrate_deploy_timeout,
)
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
if TYPE_CHECKING:
import psycopg
import psycopg.sql
def str_to_bool(value: Optional[str]) -> bool:
@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
_INVALID_LITELLM_INDEXES_SQL: Final = (
"SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
"FROM pg_index i "
"JOIN pg_class c ON c.oid = i.indexrelid "
"JOIN pg_class t ON t.oid = i.indrelid "
"JOIN pg_namespace n ON n.oid = t.relnamespace "
"WHERE NOT i.indisvalid "
" AND c.relkind = 'i' "
" AND n.nspname = %s "
" AND t.relname LIKE %s "
" AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) "
"ORDER BY c.relname"
)
@dataclass(frozen=True, slots=True)
class _InvalidIndex:
schema: str
name: str
table_size: str
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@ -624,7 +650,7 @@ class ProxyExtrasDBManager:
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
schema, etc.) from DATABASE_URL so psycopg can parse it."""
from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
parsed = urlparse(url)
if not parsed.query:
@ -645,7 +671,7 @@ class ProxyExtrasDBManager:
"target_session_attrs",
}
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
return urlunparse(parsed._replace(query=urlencode(kept)))
return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
@staticmethod
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
@ -719,6 +745,95 @@ class ProxyExtrasDBManager:
", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
)
@staticmethod
def _invalid_litellm_indexes(
conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
) -> tuple[_InvalidIndex, ...]:
rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
return tuple(_InvalidIndex(*row) for row in rows)
@staticmethod
def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
from psycopg import sql
target: Final = sql.Identifier(index.schema, index.name)
if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name):
return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover"
return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt"
@staticmethod
def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None:
import psycopg
statement, action = ProxyExtrasDBManager._index_repair(index)
try:
conn.execute(statement)
except psycopg.Error as e:
logger.warning(
"Could not repair invalid index %s.%s, will retry on the next startup. "
"If this keeps happening, run `%s` by hand as the index owner. Error: %s",
index.schema,
index.name,
statement.as_string(conn),
e,
)
return
logger.info("%s invalid index %s.%s", action, index.schema, index.name)
@staticmethod
def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
"""Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
INVALID (a migration deadlock between replicas is the usual cause; the
retried migration skips them because of IF NOT EXISTS). Never raises:
returns True when no invalid index remains, False when the repair was
skipped or failed and will be retried on the next startup. Looks in the
schema DATABASE_URL names, the only URL Prisma migrates through, but
connects over DIRECT_URL when set: the session settings, the advisory
lock and REINDEX CONCURRENTLY all need one server session, which a
transaction pooler does not give."""
prisma_url: Final = os.getenv("DATABASE_URL")
if not prisma_url:
return False
try:
import psycopg
from psycopg import sql
except ImportError:
logger.warning(
"psycopg is not installed; skipping the invalid index check. "
"Install the litellm[extra_proxy] extra, which includes psycopg."
)
return False
schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public"
cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url)
try:
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
conn.execute("SET statement_timeout = 0")
conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout)))
found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
if not found:
return True
logger.warning(
"Found %d invalid index(es) left by an interrupted CREATE INDEX "
"CONCURRENTLY, rebuilding: %s",
len(found),
", ".join(f"{index.name} (table size {index.table_size})" for index in found),
)
lock_row: Final = conn.execute(
"SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)
).fetchone()
if lock_row is None or not lock_row[0]:
logger.info("Another replica is already rebuilding the invalid indexes, skipping")
return False
for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
ProxyExtrasDBManager._repair_index(conn, index)
remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
except psycopg.Error as e:
logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
return False
return not remaining
@staticmethod
def _setup_database_v2(use_migrate: bool) -> bool:
"""
@ -994,6 +1109,7 @@ class ProxyExtrasDBManager:
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
if migrated:
ProxyExtrasDBManager.repair_invalid_indexes()
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return migrated

11
litellm-proxy-extras/migration_runbook.md
View file

@ -48,7 +48,7 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
## What It Does
1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check))
1. **Verifies the current branch is up to date with origin's current default branch** (see [Branch freshness](#branch-freshness-check))
2. Creates temp PostgreSQL DB
3. Applies existing migrations
4. Compares with `schema.prisma`
@ -57,11 +57,11 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
## Branch Freshness Check
Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. The default base is discovered from origin's advertised HEAD on each run, so an existing clone follows a default-branch change without trusting cached `origin/HEAD`. If discovery or fetching fails, migration generation stops. A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
Flags:
- `--base-branch <name>` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`.
- `--base-branch <name>` — check against a different base (e.g. a release branch). Defaults to origin's current default branch
- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base.
When the guard fires:
@ -69,8 +69,9 @@ When the guard fires:
1. Update your branch:
```bash
git fetch origin && git rebase origin/litellm_internal_staging
# or git merge origin/litellm_internal_staging — whichever matches your workflow
base_branch=$(python3 scripts/default_branch.py --branch) &&
git fetch origin "+refs/heads/$base_branch:refs/remotes/origin/$base_branch" &&
git rebase "origin/$base_branch"
```
2. Re-run `run_migration.py`.

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

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

4
litellm/__init__.py
View file

@ -495,6 +495,7 @@ public_model_groups: Optional[List[str]] = None
public_agent_groups: Optional[List[str]] = None
agent_search_embedding_model: Optional[str] = None
mcp_tool_search: Optional[Mapping[str, object]] = None
skill_search_embedding_model: Optional[str] = None
# Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
# New format: { "displayName": { "url": "...", "index": 0 } }
# Old format: { "displayName": "url" } (for backward compatibility)
@ -2001,6 +2002,9 @@ if TYPE_CHECKING:
from .llms.hosted_vllm.responses.transformation import (
HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig,
)
from .llms.fireworks_ai.responses.transformation import (
FireworksAIResponsesAPIConfig as FireworksAIResponsesAPIConfig,
)
from .llms.github_copilot.chat.transformation import (
GithubCopilotConfig as GithubCopilotConfig,
)

5
litellm/_lazy_imports_registry.py
View file

@ -237,6 +237,7 @@ LLM_CONFIG_NAMES: Final = (
"XAIResponsesAPIConfig",
"LiteLLMProxyResponsesAPIConfig",
"HostedVLLMResponsesAPIConfig",
"FireworksAIResponsesAPIConfig",
"VolcEngineResponsesAPIConfig",
"PerplexityResponsesConfig",
"DatabricksResponsesAPIConfig",
@ -957,6 +958,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.hosted_vllm.responses.transformation",
"HostedVLLMResponsesAPIConfig",
),
"FireworksAIResponsesAPIConfig": (
".llms.fireworks_ai.responses.transformation",
"FireworksAIResponsesAPIConfig",
),
"VolcEngineResponsesAPIConfig": (
".llms.volcengine.responses.transformation",
"VolcEngineResponsesAPIConfig",

1
litellm/batches/main.py
View file

@ -319,6 +319,7 @@ def create_batch(
timeout=timeout,
max_retries=optional_params.max_retries,
create_batch_data=_create_batch_request,
custom_endpoint=optional_params.get("custom_endpoint"),
)
else:
raise litellm.exceptions.BadRequestError(

65
litellm/caching/redis_cache.py
View file

@ -20,6 +20,8 @@ from contextvars import ContextVar
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
from pydantic import TypeAdapter
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import (
@ -80,11 +82,29 @@ class _AsyncRedisCommands(Protocol):
def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
_BREAKER_GUARD_FRAME_NAMES: Final = frozenset(
{"<lambda>", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"}
)
_INCREMENT_WITH_FLOOR_LUA: Final = (
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]) "
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end "
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end "
"return count"
)
_LUA_COUNT: Final = TypeAdapter(int)
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]:
return _OPTIONAL_COUNTS.validate_python(
tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values)
)
def _get_call_stack_info(num_frames: int = 2) -> str:
"""
@ -736,6 +756,43 @@ class RedisCache(BaseCache):
)
raise e
@_redis_circuit_breaker_guard_sync
def increment_with_floor(self, key: str, value: int, ttl: int) -> int:
"""Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call.
A counter whose key expired while a request was still in flight would otherwise be
recreated negative by that request's decrement. Clamping inside the same call is what
keeps it safe: a separate corrective write could land after another pod's increment and
erase it.
The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was
created rather than ``ttl`` after it was last touched. Refreshing it on every touch
would keep a count a dead worker never decremented alive for as long as the group
takes traffic. Returns the resulting count.
"""
namespaced_key: Final = self.check_and_fix_namespace(key=key)
count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval
_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl
)
return _LUA_COUNT.validate_python(count)
@_redis_circuit_breaker_guard_sync
def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
"""Read integer counters for ``key_list``, in order, raising when Redis cannot answer.
``batch_get_cache`` swallows every failure and returns an empty dict, which the caller
cannot tell apart from "every counter is unset". A caller that has to fall back to its
own numbers when Redis is unreachable needs the failure, not a dict of zeros.
"""
namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys))
@_redis_circuit_breaker_guard
async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
"""Async twin of ``batch_get_counts``, raising on failure the same way."""
namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys))
@_redis_circuit_breaker_guard
async def async_scan_iter(self, pattern: str, count: int = 100) -> list:
start_time: Final = time.time()
@ -1241,6 +1298,14 @@ class RedisCache(BaseCache):
result = result.decode()
return float(result)
@_redis_circuit_breaker_guard
async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int:
"""Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees."""
_redis_client: Final = self._async_commands()
namespaced_key: Final = self.check_and_fix_namespace(key=key)
count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
return _LUA_COUNT.validate_python(count)
async def flush_cache_buffer(self):
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

9
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -370,7 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
and isinstance(tool_call.get("custom"), dict)
)
for msg in messages:
leading_system_count: Final = next(
(index for index, msg in enumerate(messages) if msg.get("role") != "system"),
len(messages),
)
for index, msg in enumerate(messages):
role = msg.get("role")
content = msg.get("content", "")
tool_calls = msg.get("tool_calls")
@ -378,7 +383,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if role == "system":
# Extract system message as instructions
if isinstance(content, str):
if isinstance(content, str) and index < leading_system_count:
if instructions:
# Concatenate multiple system prompts with a space
instructions = f"{instructions} {content}"

6
litellm/constants.py
View file

@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
budget_reservation_disabled_info_emitted = False
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
@ -72,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096))
DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3))
DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1))
DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5))
DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float(
os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1")
)
DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5))
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1))
DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
@ -1457,6 +1461,8 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"})
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"

40
litellm/cost_calculator.py
View file

@ -2414,6 +2414,46 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
)
_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"})
class _ResponsesWsEventResponse(BaseModel):
usage: Mapping[str, object] | None = None
class _ResponsesWsEvent(BaseModel):
type: str = ""
response: _ResponsesWsEventResponse | None = None
class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
@staticmethod
def collect_usage_from_responses_ws_results(
results: Sequence[Mapping[str, object]],
) -> tuple[Usage, ...]:
events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results)
return tuple(
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses
event.response.usage
)
for event in events
if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES
and event.response is not None
and event.response.usage is not None
)
@staticmethod
def collect_and_combine_usage_from_responses_ws_results(
results: Sequence[Mapping[str, object]],
) -> Usage:
collected_usage_objects: Final = ResponsesWebSocketTokenUsageProcessor.collect_usage_from_responses_ws_results(
results
)
return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(
list(collected_usage_objects) # mutable-ok: combine_usage_objects requires a list parameter
)
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"

3
litellm/exceptions.py
View file

@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError):
num_retries: int | None = None,
headers: dict | None = None,
exception_status_code: int | None = None,
response: httpx.Response | None = None,
):
request: Final = httpx.Request(
method="POST",
@ -352,6 +353,8 @@ class Timeout(openai.APITimeoutError):
self.max_retries = max_retries
self.num_retries = num_retries
self.headers = headers
if response is not None:
self.response = response
# custom function to convert to str
def __str__(self):

82
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, cast
from urllib.parse import urlparse
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.custom_prompt_management import CustomPromptManagement
@ -23,6 +25,7 @@ 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.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
)
OPENAI_API_HOST: Final = "api.openai.com"
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
def _validated_object_mapping(value: object) -> dict[object, object] | None:
try:
return _OBJECT_MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return None
def _validated_object_list(value: object) -> list[object] | None:
try:
return _OBJECT_LIST_ADAPTER.validate_python(value)
except ValidationError:
return None
def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
if model_map_flag is not None:
@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_
class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _request_value(request_kwargs: object, key: str) -> object:
request_mapping: Final = _validated_object_mapping(request_kwargs)
if request_mapping is None:
return None
return request_mapping.get(key)
@staticmethod
def _request_user_agent(request_kwargs: object) -> str | None:
proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request")
proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request)
if proxy_server_request_mapping is None:
return None
headers: Final = proxy_server_request_mapping.get("headers")
headers_mapping: Final = _validated_object_mapping(headers)
if headers_mapping is None:
return None
user_agent: Final = next(
(value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"),
None,
)
return user_agent if isinstance(user_agent, str) else None
@staticmethod
def _request_system(request_kwargs: object) -> str | list[object] | None:
system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system")
if isinstance(system, str):
return system
return _validated_object_list(system)
def get_chat_completion_prompt(
self,
model: str,
@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
points: Sequence[CacheControlInjectionPoint],
messages: list[AllMessageValues],
tools: list[object] | None,
cache_control: object,
model: str,
custom_llm_provider: str | None,
api_base: object,
prompt_cache_options: object,
) -> Sequence[Mapping[str, object]] | None:
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools):
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
return None
return AnthropicCacheControlHook._stamped_with_dialect(
points, model, custom_llm_provider, api_base, prompt_cache_options
@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None,
cache_control: object = None,
) -> bool:
"""Whether configured injection points must yield to client-set cache_control.
@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
"""
if all(point.get("_litellm_judged") for point in points):
return False
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools)
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
@staticmethod
def _request_has_cache_control(
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None = None,
cache_control: object = None,
) -> bool:
"""Return True if the request already carries any client-supplied cache_control.
@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
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:
@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider: str | None,
tools: list | None = None,
enable_prompt_caching: bool | None = None,
cache_control: object = None,
request_kwargs: object = None,
) -> list[CacheControlInjectionPoint]:
"""Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
return []
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
return []
if is_claude_code_one_shot_subagent_request(
messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs)
):
return []
control: Final = AnthropicCacheControlHook._default_control()
@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
models: Iterable[str],
tools: list[AllToolParamValues] | None = None,
enable_prompt_caching: bool | None = None,
request_kwargs: object = None,
) -> list[AllMessageValues]:
"""Return the messages auto prompt caching will send, default breakpoints included.
@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
for candidate in (
AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
custom_llm_provider=None,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
system=AnthropicCacheControlHook._request_system(request_kwargs),
cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"),
request_kwargs=request_kwargs,
)
for model in models
)
@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
non_default_params["cache_control_injection_points"],
messages,
tools,
non_default_params.get("cache_control"),
model,
custom_llm_provider,
api_base,
@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider=custom_llm_provider,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
cache_control=non_default_params.get("cache_control"),
request_kwargs=non_default_params,
)
if points:
non_default_params["cache_control_injection_points"] = points
@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy
bool | None, kwargs.pop("enable_prompt_caching", None)
)
cache_control: Final = kwargs.get("cache_control")
configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools):
if configured and AnthropicCacheControlHook._should_stand_down(
configured, typed_messages, system, tools, cache_control
):
return messages, system
injection_points: list[CacheControlInjectionPoint] = configured or []
if not injection_points and model is not None:
@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
model=model,
custom_llm_provider=custom_llm_provider,
enable_prompt_caching=enable_prompt_caching,
cache_control=cache_control,
request_kwargs=kwargs,
)
if not injection_points:
return messages, system

6
litellm/integrations/callback_configs.json
View file

@ -356,6 +356,12 @@
"description": "OpenTelemetry collector endpoint URL",
"required": true
},
"otel_traces_endpoint": {
"type": "text",
"ui_name": "Traces Endpoint URL",
"description": "Complete trace export URL used verbatim when the collector does not serve /v1/traces (OTel v2 only)",
"required": false
},
"otel_headers": {
"type": "text",
"ui_name": "Headers",

20
litellm/integrations/custom_guardrail.py
View file

@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger):
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None,
supported_event_hooks: list[GuardrailEventHooks],
) -> None:
allowed_hooks: Final = frozenset(supported_event_hooks) | (
frozenset((GuardrailEventHooks.logging_only,))
if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks
else frozenset()
)
def _validate_event_hook_list_is_in_supported_event_hooks(
event_hook: list[GuardrailEventHooks] | list[str],
supported_event_hooks: list[GuardrailEventHooks],
@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger):
for hook in event_hook:
if isinstance(hook, str):
hook = GuardrailEventHooks(hook)
if hook not in supported_event_hooks:
if hook not in allowed_hooks:
raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}")
if event_hook is None:
@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger):
default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default]
_validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks)
elif isinstance(event_hook, GuardrailEventHooks):
if event_hook not in supported_event_hooks:
if event_hook not in allowed_hooks:
raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}")
@staticmethod
@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger):
def uses_apply_guardrail_interface(self) -> bool:
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
def _deployment_pre_call_target(self) -> "CustomLogger":
def _deployment_hook_target(self) -> "CustomLogger":
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
return self
try:
@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger):
# CHECK IF GUARDRAIL REJECTS THE REQUEST
if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
target: Final = self._deployment_pre_call_target()
target: Final = self._deployment_hook_target()
if target is not self:
kwargs["guardrail_to_apply"] = self
result: Final = await target.async_pre_call_hook(
@ -845,7 +851,9 @@ class CustomGuardrail(CustomLogger):
return None
# CHECK IF GUARDRAIL REJECTS THE REQUEST
result: Final = await self.async_post_call_success_hook(
target: Final = self._deployment_hook_target()
hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
@ -853,7 +861,7 @@ class CustomGuardrail(CustomLogger):
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=request_data,
data=hook_request_data,
response=response,
)

19
litellm/integrations/otel/model/config.py
View file

@ -72,6 +72,14 @@ class ExporterSpec(BaseModel):
description="console | in_memory | otlp_http | otlp_grpc | <factory kind>",
)
endpoint: str | None = None
traces_endpoint: str | None = Field(
default=None,
description=(
"Complete OTLP/HTTP trace URL, used verbatim. Set this when the "
"collector serves traces on a path other than ``/v1/traces``; "
"``endpoint`` is a base URL the signal path is appended to."
),
)
headers: str | None = None
owner: ExporterOwner | None = Field(
default=None,
@ -127,6 +135,14 @@ class OpenTelemetryV2Config(BaseSettings):
default=None,
validation_alias=AliasChoices("OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"),
)
traces_endpoint: str | None = Field(
default=None,
validation_alias=AliasChoices("OTEL_TRACES_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"),
description=(
"Complete OTLP/HTTP trace URL for the single-destination shorthand, "
"used verbatim instead of ``endpoint`` + ``/v1/traces``."
),
)
headers: str | None = Field(
default=None,
validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"),
@ -250,7 +266,7 @@ class OpenTelemetryV2Config(BaseSettings):
@model_validator(mode="after")
def _normalize(self) -> "OpenTelemetryV2Config":
# An endpoint with the default exporter kind implies OTLP/HTTP.
if self.endpoint and self.exporter == "console":
if (self.endpoint or self.traces_endpoint) and self.exporter == "console":
self.exporter = "otlp_http"
# When no explicit destinations are given, fold the single-destination
# shorthand into one spec so the provider always has a destination.
@ -259,6 +275,7 @@ class OpenTelemetryV2Config(BaseSettings):
ExporterSpec(
kind=self.exporter,
endpoint=self.endpoint,
traces_endpoint=self.traces_endpoint,
headers=self.headers,
)
]

11
litellm/integrations/otel/plumbing/providers.py
View file

@ -170,7 +170,7 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
)
return HTTPExporter(
endpoint=_otlp_traces_endpoint(spec.endpoint),
endpoint=spec.traces_endpoint or _otlp_traces_endpoint(spec.endpoint),
headers=parse_headers(spec.headers),
)
if kind in _OTLP_GRPC_KINDS:
@ -201,7 +201,14 @@ def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter:
``exporter`` / ``endpoint`` / ``headers`` fields. To configure multiple
exporters, populate ``config.exporters`` directly.
"""
return _exporter_from_spec(ExporterSpec(kind=config.exporter, endpoint=config.endpoint, headers=config.headers))
return _exporter_from_spec(
ExporterSpec(
kind=config.exporter,
endpoint=config.endpoint,
traces_endpoint=config.traces_endpoint,
headers=config.headers,
)
)
def _otlp_metrics_endpoint(endpoint: str | None) -> str | None:

43
litellm/litellm_core_utils/api_route_to_call_types.py
View file

@ -14,21 +14,48 @@ from typing import Final
from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes
def _segment_matches(route_segment: str, pattern_segment: str) -> bool:
"""
Match one concrete path segment against one pattern segment.
A bare placeholder ({param}) matches any segment; a placeholder with a
literal suffix ({model}:generateContent) requires the segment to end with
that suffix and have a non-empty value before it.
"""
if not pattern_segment.startswith("{"):
return route_segment == pattern_segment
placeholder_end: Final = pattern_segment.find("}")
if placeholder_end == -1:
return route_segment == pattern_segment
literal_suffix: Final = pattern_segment[placeholder_end + 1 :]
if not literal_suffix:
return True
return route_segment.endswith(literal_suffix) and len(route_segment) > len(literal_suffix)
def _pattern_tail_spans_segments(pattern_tail: str) -> bool:
"""
Whether the pattern's last segment is a suffixed placeholder
({model}:generateContent) that may absorb extra route segments, mirroring
FastAPI's {model_name:path} converter for slash-containing model names.
"""
return pattern_tail.startswith("{") and "}" in pattern_tail and not pattern_tail.endswith("}")
def _route_matches_pattern(route: str, pattern: str) -> bool:
"""
Return True if the concrete route matches the pattern.
Pattern segments like {param} match any single path segment.
Pattern segments like {param} match any single path segment, and a
suffixed placeholder in the last segment may span multiple segments.
"""
route_parts: Final = route.strip("/").split("/")
pattern_parts: Final = pattern.strip("/").split("/")
if len(route_parts) != len(pattern_parts):
if len(route_parts) < len(pattern_parts):
return False
for r, p in zip(route_parts, pattern_parts):
if p.startswith("{") and p.endswith("}"):
continue
if r != p:
return False
return True
if len(route_parts) > len(pattern_parts) and not _pattern_tail_spans_segments(pattern_parts[-1]):
return False
head_count: Final = len(pattern_parts) - 1
merged_parts: Final = (*route_parts[:head_count], "/".join(route_parts[head_count:]))
return all(_segment_matches(r, p) for r, p in zip(merged_parts, pattern_parts))
def get_call_types_for_route(route: str) -> Sequence[CallTypes] | None:

15
litellm/litellm_core_utils/credential_accessor.py
View file

@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem
class CredentialAccessor:
@staticmethod
def find_credential(credential_name: str) -> CredentialItem | None:
return next(
(credential for credential in litellm.credential_list if credential.credential_name == credential_name),
None,
)
@staticmethod
def get_credential_values(credential_name: str) -> dict:
"""Safe accessor for credentials."""
if not litellm.credential_list:
return {}
for credential in litellm.credential_list:
if credential.credential_name == credential_name:
return credential.credential_values.copy()
return {}
credential: Final = CredentialAccessor.find_credential(credential_name)
return {} if credential is None else credential.credential_values.copy()
@staticmethod
def upsert_credentials(credentials: list[CredentialItem]):

11
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -860,6 +860,7 @@ def _map_bedrock_exception(
message=mantle_context_window_message,
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
)
if (
"too many tokens" in error_str
@ -873,6 +874,7 @@ def _map_bedrock_exception(
message=f"BedrockException: Context Window Error - {error_str}",
model=model,
llm_provider="bedrock",
response=getattr(original_exception, "response", None),
)
elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str:
raise BadRequestError(
@ -924,12 +926,14 @@ def _map_bedrock_exception(
message=f"BedrockException: Timeout Error - {error_str}",
model=model,
llm_provider="bedrock",
response=getattr(original_exception, "response", None),
)
elif "Could not process image" in error_str:
raise litellm.InternalServerError(
message=f"BedrockException - {error_str}",
model=model,
llm_provider="bedrock",
response=getattr(original_exception, "response", None),
)
elif hasattr(original_exception, "status_code"):
if original_exception.status_code == 500:
@ -937,10 +941,7 @@ def _map_bedrock_exception(
message=f"BedrockException - {original_exception.message}",
llm_provider="bedrock",
model=model,
response=httpx.Response(
status_code=500,
request=httpx.Request(method="POST", url="https://api.openai.com/v1/"),
),
response=getattr(original_exception, "response", None),
)
elif original_exception.status_code == 401:
raise AuthenticationError(
@ -969,6 +970,7 @@ def _map_bedrock_exception(
model=model,
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
response=getattr(original_exception, "response", None),
)
elif original_exception.status_code == 422:
raise BadRequestError(
@ -1001,6 +1003,7 @@ def _map_bedrock_exception(
llm_provider=custom_llm_provider,
litellm_debug_info=extra_information,
exception_status_code=original_exception.status_code,
response=getattr(original_exception, "response", None),
)

8
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -53,12 +53,14 @@ class GetModelCostMap:
_backup_model_count: int = -1 # -1 = not yet loaded
@staticmethod
def read_local_model_cost_map_text() -> str:
return files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8")
@staticmethod
def load_local_model_cost_map() -> dict:
"""Load the local backup model cost map bundled with the package."""
content: Final = json.loads(
files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text(encoding="utf-8")
)
content: Final = json.loads(GetModelCostMap.read_local_model_cost_map_text())
return content
@classmethod

12
litellm/litellm_core_utils/litellm_logging.py
View file

@ -48,6 +48,7 @@ from litellm.constants import (
)
from litellm.cost_calculator import (
RealtimeAPITokenUsageProcessor,
ResponsesWebSocketTokenUsageProcessor,
_select_model_name_for_cost_calc,
)
from litellm.exceptions import (
@ -2028,6 +2029,17 @@ class Logging(LiteLLMLoggingBaseClass):
results=result,
)
elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
combined_ws_usage: Final = (
ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
)
logging_result = LiteLLMRealtimeStreamLoggingObject(
usage=combined_ws_usage,
results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream
)
elif (
self.call_type == CallTypes.llm_passthrough_route.value
or self.call_type == CallTypes.allm_passthrough_route.value

86
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -514,6 +514,7 @@ def _get_token_base_cost(
current_time: datetime | None = None,
*,
threshold_is_inclusive: bool = False,
missing_cache_read_uses_input: bool = False,
) -> tuple[float, float, float, float, float]:
"""
Return prompt cost, completion cost, and cache costs for a given model and usage.
@ -524,6 +525,9 @@ def _get_token_base_cost(
`threshold_is_inclusive` switches that comparison to >=, for providers such as xAI
that bill the higher tier once the prompt reaches the threshold.
`missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved
input rate instead of 0.0; an explicit 0.0 rate stays a real price either way.
Returns:
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
"""
@ -551,29 +555,16 @@ def _get_token_base_cost(
float,
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
)
cache_read_cost = cast(float, _get_cost_per_unit(model_info, cache_read_cost_key))
cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None)
## CHECK IF ABOVE THRESHOLD
# Optimization: collect threshold keys first to avoid sorting all model_info keys.
# Most models don't have threshold pricing, so we can return early.
# Exclude service_tier-specific variants (e.g. input_cost_per_token_above_200k_tokens_priority)
# so that the threshold detection loop only processes standard keys. The
# service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key.
threshold_keys: Final = [
k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES)
]
if not threshold_keys:
return _apply_off_peak_to_base_costs(
model_info,
current_time,
(
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
),
)
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
threshold: float | None = None
@ -662,10 +653,7 @@ def _get_token_base_cost(
),
)
cache_read_cost = cast(
float,
_get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost),
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
break
except (IndexError, ValueError):
@ -673,6 +661,17 @@ def _get_token_base_cost(
except Exception:
continue
if cache_read_cost is None:
cache_read_cost = (
_off_peak_rate(
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
"input_cost_per_token",
prompt_base_cost,
)
if missing_cache_read_uses_input
else 0.0
)
return _apply_off_peak_to_base_costs(
model_info,
current_time,
@ -1416,6 +1415,57 @@ def get_token_type_cost_breakdown(
)
def calculate_prompt_caching_savings(
model_info: ModelInfo,
usage: Usage,
custom_llm_provider: str | None,
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
billed_at: datetime | None = None,
) -> float:
"""Read discount minus write premium, using the biller's rate and TTL resolution.
Missing reads and unpublished (missing/zero) writes claim no saving or premium;
explicit zero reads remain free. An unpublished 1h price uses the ordinary write rate.
``billed_at`` is the request's completion time, so off-peak windows resolve as the
biller saw them rather than at the later spend write.
"""
prompt_base_cost, _, cache_creation_cost, cache_creation_cost_above_1hr, cache_read_cost = _get_token_base_cost(
model_info=model_info,
usage=usage,
service_tier=service_tier,
current_time=billed_at,
threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
missing_cache_read_uses_input=True,
)
write_rate: Final = cache_creation_cost or prompt_base_cost
write_rate_1h: Final = cache_creation_cost_above_1hr or write_rate
prompt_tokens_details: Final = parse_prompt_tokens_details(usage)
cache_read_tokens: Final = max(prompt_tokens_details["cache_hit_tokens"], 0)
cache_creation_tokens: Final = max(prompt_tokens_details["cache_creation_tokens"], 0)
details: Final = prompt_tokens_details["cache_creation_token_details"]
cache_creation_details: Final = (
CacheCreationTokenDetails(
ephemeral_5m_input_tokens=max(details.ephemeral_5m_input_tokens or 0, 0),
ephemeral_1h_input_tokens=max(details.ephemeral_1h_input_tokens or 0, 0),
)
if details is not None
else None
)
read_discount: Final = cache_read_tokens * max(prompt_base_cost - cache_read_cost, 0.0)
write_premium: Final = calculate_cache_writing_cost(
cache_creation_tokens=cache_creation_tokens,
cache_creation_token_details=cache_creation_details,
cache_creation_cost_above_1hr=write_rate_1h - prompt_base_cost,
cache_creation_cost=write_rate - prompt_base_cost,
)
uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift(
model_info, vertex_location
)
return (read_discount - write_premium) * uplift
def calculate_image_response_cost_from_usage(
model: str,
image_response: ImageResponse,

5
litellm/litellm_core_utils/llm_response_utils/get_headers.py
View file

@ -1,7 +1,8 @@
from collections.abc import Mapping
from typing import Final
def get_response_headers(_response_headers: dict | None = None) -> dict:
def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict:
"""
Sets the Appropriate OpenAI headers for the response and forward all headers as llm_provider-{header}
@ -31,7 +32,7 @@ def get_response_headers(_response_headers: dict | None = None) -> dict:
return {**llm_provider_headers, **openai_headers}
def _get_llm_provider_headers(response_headers: dict) -> dict:
def _get_llm_provider_headers(response_headers: Mapping[str, str]) -> dict:
"""
Adds a llm_provider-{header} to all headers that are not already prefixed with llm_provider

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

@ -67,6 +67,96 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
_INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
_DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
_DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
def is_claude_code_user_agent(user_agent: str) -> bool:
return user_agent.startswith("claude-cli/")
def _validated_claude_code_mapping(value: object) -> dict[object, object] | None:
try:
return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return None
def _validated_claude_code_list(value: object) -> list[object] | None:
try:
return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
except ValidationError:
return None
def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None:
stripped: Final = text.strip()
if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX):
return None
fields: Final = tuple(
field
for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";")
if (field := raw_field.strip())
)
if not fields or any("=" not in field for field in fields):
return None
parsed_fields: Final = tuple(
(parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),)
)
if any(not key or not value for key, value in parsed_fields):
return None
return parsed_fields
def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
if isinstance(system, str):
return (system,)
blocks: Final = _validated_claude_code_list(system)
if blocks is None:
return None
block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)
if any(block is None for block in block_mappings):
return None
text_values: Final = tuple(
block.get("text") for block in block_mappings if block is not None and block.get("type") == "text"
)
if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values):
return None
meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip())
return meaningful_text or None
def _is_claude_code_subagent_billing_system(system: object) -> bool:
billing_texts: Final = _claude_code_billing_texts(system)
if billing_texts is None:
return False
billing_fields: Final = tuple(
fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None
)
if len(billing_fields) != len(billing_texts):
return False
subagent_values: Final = tuple(
value for fields in billing_fields for key, value in fields if key == "cc_is_subagent"
)
return subagent_values == ("true",)
def is_claude_code_one_shot_subagent_request(
messages: list[AllMessageValues],
system: object,
tools: object,
user_agent: str | None,
) -> bool:
only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None
return (
user_agent is not None
and is_claude_code_user_agent(user_agent)
and not tools
and only_message is not None
and only_message.get("role") == "user"
and _is_claude_code_subagent_billing_system(system)
)
def _strip_bedrock_id_suffixes(model: str) -> str:

2
litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py
View file

@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp(
LiteLLM_Proxy_MCP_Handler,
)
mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools)
mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools)
if not mcp_references:
return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn(

19
litellm/llms/bedrock/chat/agentcore/transformation.py
View file

@ -667,7 +667,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=str(response.read()))
raise BedrockError(
status_code=response.status_code,
message=str(response.read()),
headers=response.headers,
response=response,
)
# LOGGING
logging_obj.post_call(
@ -690,6 +695,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
status_code=response.status_code,
message=f"AgentCore: Failed to read/parse JSON response body: {e}",
headers=response.headers,
)
parsed: Final = self._parse_json_response(response_json)
@ -880,7 +886,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=str(await response.aread()))
raise BedrockError(
status_code=response.status_code,
message=str(await response.aread()),
headers=response.headers,
response=response,
)
# LOGGING
logging_obj.post_call(
@ -903,6 +914,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
status_code=response.status_code,
message=f"AgentCore: Failed to read/parse JSON response body: {e}",
headers=response.headers,
)
parsed: Final = self._parse_json_response(response_json)
@ -1031,6 +1043,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error processing response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
def validate_environment(
@ -1046,7 +1059,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
return headers
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def should_fake_stream(
self,

23
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -22,7 +22,7 @@ from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
from ..common_utils import BedrockError, _get_all_bedrock_regions
from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
@ -66,7 +66,12 @@ def make_sync_call(
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=str(response.read()))
raise BedrockError(
status_code=response.status_code,
message=str(response.read()),
headers=response.headers,
response=response,
)
if fake_stream:
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@ -247,7 +252,12 @@ class BedrockConverseLLM(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -594,7 +604,12 @@ class BedrockConverseLLM(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

1
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -2255,6 +2255,7 @@ class AmazonConverseConfig(BaseConfig):
raise BedrockError(
message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues",
status_code=422,
headers=response.headers,
)
"""

3
litellm/llms/bedrock/chat/invoke_agent/transformation.py
View file

@ -470,6 +470,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error processing response: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
def validate_environment(
@ -485,7 +486,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
return headers
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def should_fake_stream(
self,

33
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -42,6 +42,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk
from ..common_utils import (
BedrockError,
build_bedrock_stream_error,
error_response_text,
get_bedrock_response_stream_shape,
get_bedrock_tool_name,
)
@ -184,7 +185,12 @@ async def make_call(
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=response.text)
raise BedrockError(
status_code=response.status_code,
message=error_response_text(response),
headers=response.headers,
response=response,
)
if fake_stream:
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@ -228,9 +234,16 @@ async def make_call(
)
return completion_stream, response.headers
except BedrockError:
raise
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
except Exception as e:
@ -270,7 +283,12 @@ def make_sync_call(
)
if response.status_code != 200:
raise BedrockError(status_code=response.status_code, message=response.text)
raise BedrockError(
status_code=response.status_code,
message=error_response_text(response),
headers=response.headers,
response=response,
)
if fake_stream:
model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response(
@ -314,9 +332,16 @@ def make_sync_call(
)
return completion_stream, response.headers
except BedrockError:
raise
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=error_response_text(err.response),
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
except Exception as e:

2
litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py
View file

@ -247,4 +247,4 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig):
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
"""Return the appropriate error class for Bedrock."""
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)

2
litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py
View file

@ -182,4 +182,4 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM):
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
"""Return the appropriate error class for Bedrock."""
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)

2
litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py
View file

@ -212,6 +212,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
raise BedrockError(
message=f"Error parsing response: {raw_response.text}, error: {e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
verbose_logger.debug(
@ -241,6 +242,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig):
raise BedrockError(
message=f"Error setting response content: {e}. Response: {completion_response}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
# Calculate usage from headers

10
litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
View file

@ -295,7 +295,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
try:
completion_response: Final = raw_response.json()
except Exception:
raise BedrockError(message=raw_response.text, status_code=raw_response.status_code)
raise BedrockError(
message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
verbose_logger.debug(
"bedrock invoke response % s",
json.dumps(completion_response, indent=4, default=str),
@ -363,6 +367,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error processing={raw_response.text}, Received error={e}",
status_code=422,
headers=raw_response.headers,
)
try:
@ -384,6 +389,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
raise BedrockError(
message=f"Error parsing received text={outputText}.\nError-{e}",
status_code=raw_response.status_code,
headers=raw_response.headers,
)
## CALCULATING USAGE - bedrock returns usage in the headers
@ -431,7 +437,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names)
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message)
return BedrockError(status_code=status_code, message=error_message, headers=headers)
@track_llm_api_timing()
async def get_async_custom_stream_wrapper(

11
litellm/llms/bedrock/claude_platform/common_utils.py
View file

@ -1,7 +1,10 @@
from typing import Final
import httpx
import litellm
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.secret_managers.main import get_secret_str
CLAUDE_PLATFORM_SERVICE_NAME: Final = "aws-external-anthropic"
@ -15,6 +18,14 @@ def strip_claude_platform_route(model: str) -> str:
class BedrockClaudePlatformMixin(BaseAWSLLM):
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
@staticmethod
def _get_workspace_id(optional_params: dict, litellm_params: dict) -> str | None:
workspace_id = (

47
litellm/llms/bedrock/common_utils.py
View file

@ -33,8 +33,53 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import AllMessageValues
_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs"
def error_response_text(response: httpx.Response) -> str:
try:
return response.text
except httpx.ResponseNotRead:
return response.reason_phrase
def _synthesize_error_response(
*, status_code: int, headers: dict[str, object] | httpx.Headers, request: httpx.Request | None
) -> tuple[httpx.Request, httpx.Response]:
error_request: Final = request or httpx.Request(method="POST", url=_ERROR_REQUEST_URL)
safe_headers: Final = (
headers
if isinstance(headers, httpx.Headers)
else tuple((key, value) for key, value in headers.items() if isinstance(value, (str, bytes)))
)
return error_request, httpx.Response(status_code=status_code, headers=safe_headers, request=error_request)
class BedrockError(BaseLLMException):
pass
def __init__(
self,
status_code: int,
message: str,
headers: dict[str, object] | httpx.Headers | None = None,
request: httpx.Request | None = None,
response: httpx.Response | None = None,
body: dict[str, object] | None = None,
status_code_is_synthesized: bool = False,
) -> None:
error_request, error_response = (
_synthesize_error_response(status_code=status_code, headers=headers, request=request)
if response is None and headers
else (request, response)
)
super().__init__(
status_code=status_code,
message=message,
headers=headers,
request=error_request,
response=error_response,
body=body,
status_code_is_synthesized=status_code_is_synthesized,
)
_BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = (

4
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -102,6 +102,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
raise BedrockError(
status_code=response.status_code,
message=error_text,
headers=response.headers,
response=response,
)
bedrock_response: Final = response.json()
@ -124,6 +126,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
raise BedrockError(
status_code=e.response.status_code,
message=e.response.text,
headers=e.response.headers,
response=e.response,
)
except Exception as e:
verbose_logger.error("Error in CountTokens handler: %s", e)

14
litellm/llms/bedrock/embed/embedding.py
View file

@ -132,7 +132,12 @@ class BedrockEmbedding(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -161,7 +166,12 @@ class BedrockEmbedding(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

9
litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py
View file

@ -20,6 +20,7 @@ import httpx
from litellm._logging import verbose_logger
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import FileTypes, ImageObject, ImageResponse
@ -228,6 +229,14 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
"""
return _supports_nova_canvas_image_edit_from_model_cost(model or "")
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_supported_openai_params(self, model: str) -> list:
return [
"n",

14
litellm/llms/bedrock/image_edit/handler.py
View file

@ -114,7 +114,12 @@ class BedrockImageEdit(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -156,7 +161,12 @@ class BedrockImageEdit(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

9
litellm/llms/bedrock/image_edit/stability_transformation.py
View file

@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.llms.stability import (
OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO,
@ -84,6 +85,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
return True
return False
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_supported_openai_params(self, model: str) -> list:
"""
Return list of OpenAI params supported by Bedrock Stability.

14
litellm/llms/bedrock/image_generation/image_handler.py
View file

@ -119,7 +119,12 @@ class BedrockImageGeneration(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
### FORMAT RESPONSE TO OPENAI FORMAT ###
@ -162,7 +167,12 @@ class BedrockImageGeneration(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

9
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation
AmazonInvokeConfig,
)
from litellm.llms.bedrock.common_utils import (
BedrockError,
apply_bedrock_invoke_structured_output,
ensure_bedrock_anthropic_messages_tool_names,
get_anthropic_beta_from_headers,
@ -79,6 +80,14 @@ class AmazonAnthropicClaudeMessagesConfig(
BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys())
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def __init__(self, **kwargs):
BaseAnthropicMessagesConfig.__init__(self, **kwargs)
AmazonInvokeConfig.__init__(self, **kwargs)

11
litellm/llms/bedrock/passthrough/transformation.py
View file

@ -2,13 +2,14 @@ import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Optional, cast
import httpx
from httpx import Response
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo
from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo
if TYPE_CHECKING:
from httpx import URL
@ -18,6 +19,14 @@ if TYPE_CHECKING:
class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig):
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
return "stream" in endpoint

10
litellm/llms/bedrock/realtime/transformation.py
View file

@ -9,12 +9,14 @@ import json
import uuid as uuid_lib
from typing import Final, cast
import httpx
from pydantic import BaseModel
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm
from litellm.types.llms.openai import (
OpenAIRealtimeContentPartDone,
@ -121,6 +123,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
self._cumulative_usage = BedrockUsageEvent()
self._reported_usage = BedrockUsageEvent()
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict:
"""Validate environment - no special validation needed for Bedrock."""
return headers

14
litellm/llms/bedrock/rerank/handler.py
View file

@ -46,7 +46,12 @@ class BedrockRerankHandler(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")
@ -117,7 +122,12 @@ class BedrockRerankHandler(BaseAWSLLM):
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise BedrockError(status_code=error_code, message=err.response.text)
raise BedrockError(
status_code=error_code,
message=err.response.text,
headers=err.response.headers,
response=err.response,
)
except httpx.TimeoutException:
raise BedrockError(status_code=408, message="Timeout error occurred.")

6
litellm/llms/bedrock/search/transformation.py
View file

@ -39,7 +39,6 @@ from typing import Final
import httpx
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.search.transformation import (
BaseSearchConfig,
SearchResponse,
@ -380,6 +379,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
raise BedrockError(
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
message=f"AgentCore gateway MCP error: {error}",
headers=raw_response.headers,
)
# A failed tools/call is reported in-band, as HTTP 200 with result.isError
@ -389,6 +389,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
raise BedrockError(
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}",
headers=raw_response.headers,
)
text_items: Final = tuple(
@ -440,6 +441,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
raise BedrockError(
status_code=502,
message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}",
headers=raw_response.headers,
)
def get_error_class(
@ -448,7 +450,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM):
status_code: int,
headers: dict, # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict
) -> Exception:
return BaseLLMException(
return BedrockError(
status_code=status_code,
message=error_message,
headers=headers,

9
litellm/llms/bedrock/vector_stores/transformation.py
View file

@ -8,6 +8,7 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.types.integrations.rag.bedrock_knowledgebase import (
BedrockKBContent,
BedrockKBResponse,
@ -38,6 +39,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
BaseVectorStoreConfig.__init__(self)
BaseAWSLLM.__init__(self)
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict
) -> BedrockError:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials:
return {}

9
litellm/llms/bedrock_mantle/chat/transformation.py
View file

@ -13,6 +13,8 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the
from collections.abc import AsyncIterator, Iterator
from typing import Any, Final
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
@ -24,6 +26,8 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
from ...base_llm.chat.transformation import BaseLLMException
from ...bedrock.common_utils import BedrockError
from ...openai_like.chat.transformation import OpenAILikeChatConfig
from ..common_utils import mantle_base_segment
@ -45,6 +49,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
def get_config(cls):
return super().get_config()
def get_error_class(
self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def _get_openai_compatible_provider_info(
self,
api_base: str | None,

8
litellm/llms/bedrock_mantle/responses/transformation.py
View file

@ -19,11 +19,14 @@ import json
from collections.abc import Mapping
from typing import Any, Final
import httpx
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock_mantle.common_utils import (
MANTLE_HOST_RE,
BedrockMantleAuthMixin,
@ -98,6 +101,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.BEDROCK_MANTLE
def get_error_class(
self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers
) -> BaseLLMException:
return BedrockError(status_code=status_code, message=error_message, headers=headers)
def get_complete_url(
self,
api_base: str | None,

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

@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
return self._transform_ocr_response(
provider_config=provider_config,
model=model,
@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
# Use async response transform for async operations
return await provider_config.async_transform_ocr_response(
model=model,

44
litellm/llms/fireworks_ai/common_utils.py
View file

@ -1,3 +1,5 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from httpx import Headers
@ -13,7 +15,7 @@ class FireworksAIException(BaseLLMException):
pass
def get_fireworks_session_id(litellm_params: dict) -> str | None:
def get_fireworks_session_id(litellm_params: Mapping[str, object]) -> str | None:
"""
Session id to send as `x-session-affinity`, or None when the caller gave none.
@ -23,19 +25,39 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
"""
params: Final = litellm_params
metadata: Final = params.get("metadata")
if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
if isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY):
return None
for key in ("litellm_session_id", "session_id"):
value = params.get(key)
if value:
return str(value)
if isinstance(metadata, dict):
if isinstance(metadata, Mapping):
value = metadata.get("session_id")
if value:
return str(value)
return None
def with_fireworks_session_affinity(
headers: Mapping[str, str], litellm_params: Mapping[str, object]
) -> Mapping[str, str]:
if any(key.lower() == "x-session-affinity" for key in headers):
return headers
session_id: Final = get_fireworks_session_id(litellm_params)
if not session_id:
return headers
return MappingProxyType({**headers, "x-session-affinity": session_id})
def resolve_fireworks_api_key(api_key: str | None) -> str | None:
return api_key or (
get_secret_str("FIREWORKS_API_KEY")
or get_secret_str("FIREWORKS_AI_API_KEY")
or get_secret_str("FIREWORKSAI_API_KEY")
or get_secret_str("FIREWORKS_AI_TOKEN")
)
AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-"
@ -63,13 +85,7 @@ class FireworksAIMixin:
)
def _get_api_key(self, api_key: str | None) -> str | None:
dynamic_api_key: Final = api_key or (
get_secret_str("FIREWORKS_API_KEY")
or get_secret_str("FIREWORKS_AI_API_KEY")
or get_secret_str("FIREWORKSAI_API_KEY")
or get_secret_str("FIREWORKS_AI_TOKEN")
)
return dynamic_api_key
return resolve_fireworks_api_key(api_key)
def validate_environment(
self,
@ -92,9 +108,5 @@ class FireworksAIMixin:
return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params)
def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict:
if any(key.lower() == "x-session-affinity" for key in headers):
return headers
session_id: Final = get_fireworks_session_id(litellm_params)
if not session_id:
return headers
return {**headers, "x-session-affinity": session_id}
pinned: Final = with_fireworks_session_affinity(headers, litellm_params)
return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place

102
litellm/llms/fireworks_ai/responses/transformation.py Normal file
View file

@ -0,0 +1,102 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from urllib.parse import unquote
import httpx
from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam
from litellm.llms.fireworks_ai.common_utils import (
resolve_fireworks_api_key,
resolve_fireworks_resource_name,
with_fireworks_session_affinity,
)
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import ResponseInputParam
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
FIREWORKS_AI_DEFAULT_API_BASE: Final = "https://api.fireworks.ai/inference/v1"
def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object]:
extras: Final[Mapping[str, object]] = litellm_params.model_extra or MappingProxyType({})
return MappingProxyType(
{"litellm_session_id": extras.get("litellm_session_id"), "metadata": extras.get("litellm_metadata")}
)
def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam:
if "role" not in item or item["role"] != "developer":
return item
return EasyInputMessageParam(role="system", content=item["content"], type="message")
def _developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam:
if isinstance(input, str):
return input
return [_developer_item_as_system(item) for item in input]
class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.FIREWORKS_AI
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
litellm_params: GenericLiteLLMParams | None,
) -> dict: # mutable-ok: overrides the base class signature
params: Final = litellm_params or GenericLiteLLMParams()
api_key: Final = resolve_fireworks_api_key(params.api_key)
if api_key is None:
raise ValueError("FIREWORKS_API_KEY is not set")
authorized: Final = MappingProxyType(
{"Content-Type": "application/json", **headers, "Authorization": f"Bearer {api_key}"}
)
pinned: Final = with_fireworks_session_affinity(authorized, _session_params(params))
return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place
def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/")
return f"{base}/responses"
def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam:
return _developer_items_as_system(super()._validate_input_param(input))
def transform_responses_api_request(
self,
model: str,
input: str | ResponseInputParam,
response_api_optional_request_params: dict, # mutable-ok: overrides the base class signature
litellm_params: GenericLiteLLMParams,
headers: dict, # mutable-ok: overrides the base class signature
) -> dict: # mutable-ok: overrides the base class signature
return super().transform_responses_api_request(
model=resolve_fireworks_resource_name(model),
input=input,
response_api_optional_request_params=response_api_optional_request_params,
litellm_params=litellm_params,
headers=headers,
)
def transform_delete_response_api_response(
self,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> DeleteResponseResult:
deleted_id: Final = unquote(raw_response.request.url.path.rsplit("/", 1)[-1])
return DeleteResponseResult(id=deleted_id, object="response", deleted=True)
def supports_native_websocket(self) -> bool:
return False
def supports_native_file_search(self) -> bool:
return False

20
litellm/llms/gemini/google_genai/guardrail_translation/__init__.py Normal file
View file

@ -0,0 +1,20 @@
"""Google GenAI generateContent guardrail translation handler."""
from typing import Final
from litellm.llms.gemini.google_genai.guardrail_translation.handler import (
GoogleGenAIGenerateContentHandler,
)
from litellm.types.utils import CallTypes
guardrail_translation_mappings: Final = { # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict)
CallTypes.generate_content: GoogleGenAIGenerateContentHandler,
CallTypes.agenerate_content: GoogleGenAIGenerateContentHandler,
CallTypes.generate_content_stream: GoogleGenAIGenerateContentHandler,
CallTypes.agenerate_content_stream: GoogleGenAIGenerateContentHandler,
}
__all__ = (
"GoogleGenAIGenerateContentHandler",
"guardrail_translation_mappings",
)

255
litellm/llms/gemini/google_genai/guardrail_translation/handler.py Normal file
View file

@ -0,0 +1,255 @@
"""
Google GenAI generateContent handler for Unified Guardrails.
Extracts text from generateContent requests (systemInstruction.parts[].text
and contents[].parts[].text) and responses (candidates[].content.parts[].text),
applies the guardrail, and
writes the guardrailed text back in place. Requests and responses may be
dicts (wire format) or google-genai SDK objects; streaming chunks may
additionally be raw SSE frames, which are scanned for detection (a blocking
guardrail raises) without rewriting the frames.
"""
import json
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Optional
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamTransformSink,
)
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
_EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({})
def _field(container: object, name: str) -> object | None:
if isinstance(container, dict):
return container.get(name)
return getattr(container, name, None)
def _part_text(part: object) -> str | None:
text: Final = _field(part, "text")
if isinstance(text, str) and text:
return text
return None
def _write_part_text(part: object, text: str) -> None:
if isinstance(part, dict):
part["text"] = text # rebind-ok: guardrail write-back rewrites the caller's part in place by handler contract
return
setattr(part, "text", text) # noqa: B010 # SDK parts are typed as object here; direct assignment cannot type-check
def _content_text_parts(content: object) -> tuple[object, ...]:
parts: Final = _field(content, "parts")
if not isinstance(parts, (list, tuple)):
return ()
return tuple(part for part in parts if _part_text(part) is not None)
def _system_instruction(data: Mapping[str, object]) -> object | None:
return next(
(
value
for container in (data, data.get("config"))
if container is not None
for key in ("systemInstruction", "system_instruction")
for value in (_field(container, key),)
if value is not None
),
None,
)
def _request_text_parts(data: Mapping[str, object]) -> tuple[object, ...]:
contents: Final = data.get("contents")
content_list: Final = (
(contents,) if isinstance(contents, dict) else tuple(contents) if isinstance(contents, list) else ()
)
return (
*_content_text_parts(_system_instruction(data)),
*(part for content in content_list for part in _content_text_parts(content)),
)
def _response_text_parts(response: object) -> tuple[object, ...]:
candidates: Final = _field(response, "candidates")
if not isinstance(candidates, (list, tuple)):
return ()
return tuple(part for candidate in candidates for part in _content_text_parts(_field(candidate, "content")))
def _part_texts(text_parts: Sequence[object]) -> tuple[str, ...]:
return tuple(text for part in text_parts for text in (_part_text(part),) if text is not None)
def _texts_payload(
texts: Sequence[str],
) -> list[str]: # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str]
return list(texts) # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str]
def _write_back_texts(text_parts: Sequence[object], guardrailed_texts: Sequence[str] | None) -> None:
if not guardrailed_texts or len(guardrailed_texts) != len(text_parts):
return
for part, text in zip(text_parts, guardrailed_texts):
_write_part_text(part, text)
def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None:
try:
parsed: Final = json.loads(payload)
except json.JSONDecodeError:
return None
if isinstance(parsed, dict):
return parsed
return None
def _sse_payload_texts(sse_text: str) -> tuple[str, ...]:
return tuple(
text
for line in sse_text.splitlines()
if line.startswith("data:")
for payload in (line[len("data:") :].strip(),)
if payload and payload != "[DONE]"
for parsed in (_parse_json_dict_or_none(payload),)
if parsed is not None
for text in _part_texts(_response_text_parts(parsed))
)
def _chunk_sse_text(chunk: object) -> str | None:
if isinstance(chunk, bytes):
return chunk.decode("utf-8", errors="replace")
if isinstance(chunk, str):
return chunk
return None
def _accumulated_stream_text(responses_so_far: Sequence[object]) -> str:
object_texts: Final = tuple(
text
for chunk in responses_so_far
if _chunk_sse_text(chunk) is None
for text in _part_texts(_response_text_parts(chunk))
)
sse_text: Final = "".join(sse for chunk in responses_so_far for sse in (_chunk_sse_text(chunk),) if sse is not None)
return "".join(object_texts) + "".join(_sse_payload_texts(sse_text))
class GoogleGenAIGenerateContentHandler(BaseTranslation):
"""
Guardrail translation for the google genai generateContent surface
(/models/{model}:generateContent, :streamGenerateContent, and the
litellm SDK generate_content call types).
"""
async def process_input_messages(
self,
data: dict, # mutable-ok: base handler contract passes the proxy's request dict through to apply_guardrail
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> object:
text_parts: Final = _request_text_parts(data)
if not text_parts:
verbose_proxy_logger.debug("Google GenAI guardrail: no request text found, skipping")
return data
model: Final = data.get("model")
inputs: Final = (
GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model)
if isinstance(model, str)
else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)))
)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
_write_back_texts(text_parts, guardrailed_inputs.get("texts"))
return data
async def process_output_response(
self,
response: object,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: Mapping[str, object] | None = None,
) -> object:
text_parts: Final = _response_text_parts(response)
if not text_parts:
verbose_proxy_logger.debug("Google GenAI guardrail: no response text found, skipping")
return response
guardrail_request_data: Final = self._merged_request_data(
request_data=request_data,
user_api_key_dict=user_api_key_dict,
context_key="response",
context_value=response,
)
model: Final = guardrail_request_data.get("model")
inputs: Final = (
GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)), model=model)
if isinstance(model, str)
else GenericGuardrailAPIInputs(texts=_texts_payload(_part_texts(text_parts)))
)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=guardrail_request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
_write_back_texts(text_parts, guardrailed_inputs.get("texts"))
return response
async def process_output_streaming_response(
self,
responses_so_far: Sequence[object],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: Mapping[str, object] | None = None,
stream_transform_sink: StreamTransformSink | None = None,
) -> object:
accumulated_text: Final = _accumulated_stream_text(responses_so_far)
if not accumulated_text:
return responses_so_far
guardrail_request_data: Final = self._merged_request_data(
request_data=request_data,
user_api_key_dict=user_api_key_dict,
context_key="responses_so_far",
context_value=responses_so_far,
)
_guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=_texts_payload((accumulated_text,))),
request_data=guardrail_request_data,
input_type="response",
logging_obj=litellm_logging_obj,
)
return responses_so_far
def _merged_request_data(
self,
request_data: Mapping[str, object] | None,
user_api_key_dict: Optional["UserAPIKeyAuth"],
context_key: str,
context_value: object,
) -> dict: # mutable-ok: CustomGuardrail.apply_guardrail requires a plain dict request payload
base: Final = request_data if request_data is not None else _EMPTY_REQUEST_DATA
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
context_pairs: Final = ((context_key, context_value),) if context_key not in base else ()
metadata_pairs: Final = (
(("litellm_metadata", user_metadata),) if user_metadata and "litellm_metadata" not in base else ()
)
return dict((*base.items(), *context_pairs, *metadata_pairs)) # mutable-ok: apply_guardrail takes a plain dict

5
litellm/llms/litellm_proxy/skills/constants.py
View file

@ -16,3 +16,8 @@ DEFAULT_MAX_ITERATIONS: Final[int] = 10
DEFAULT_SANDBOX_TIMEOUT: Final[int] = 120
"""Default timeout in seconds for sandbox code execution."""
MAX_SKILLS_PER_SEARCH: Final[int] = 5000
"""Upper bound on how many of the caller's accessible skills a single semantic
search embeds. Ranking runs in memory over this candidate set (no tsvector/DB-side
filtering yet), so this caps worst-case embedding cost per search request."""

19
litellm/llms/litellm_proxy/skills/handler.py
View file

@ -6,11 +6,15 @@ Used by the transformation layer and skills injection hook.
"""
import uuid
from collections.abc import Sequence
from typing import Final
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.llms.litellm_proxy.skills.constants import LITELLM_SKILL_ID_PREFIX
from litellm.llms.litellm_proxy.skills.constants import (
LITELLM_SKILL_ID_PREFIX,
MAX_SKILLS_PER_SEARCH,
)
from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth
from litellm.proxy.common_utils.resource_ownership import (
get_primary_resource_owner_scope,
@ -131,6 +135,19 @@ class LiteLLMSkillsHandler:
)
return [_prisma_skill_to_litellm(s) for s in skills]
@staticmethod
async def list_skills_for_search(
user_api_key_dict: UserAPIKeyAuth | None = None,
) -> Sequence[LiteLLM_SkillsTable]:
"""Every skill the caller can access, for ranking. Same owner-scope filter as
``list_skills``, but unpaginated (up to ``MAX_SKILLS_PER_SEARCH``) since a query
must be scored against the whole accessible set, not one page of it."""
return await LiteLLMSkillsHandler.list_skills(
limit=MAX_SKILLS_PER_SEARCH,
offset=0,
user_api_key_dict=user_api_key_dict,
)
@staticmethod
async def _load_skill(skill_id: str) -> object | None:
"""Cache-first read of the Prisma skill row. Owner-scope filtering

161
litellm/llms/litellm_proxy/skills/skill_search.py Normal file
View file

@ -0,0 +1,161 @@
"""Semantic ranking over the LiteLLM-hosted skill registry, shared by GET /v1/skills?query= and the skill_search MCP tool."""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, TypeAlias
from pydantic import BaseModel, ConfigDict
from litellm.llms.litellm_proxy.skills.constants import MAX_SKILLS_PER_SEARCH
from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler
from litellm.proxy.common_utils.semantic_text_index import (
Embedder,
EmbeddingFailed,
SemanticTextIndex,
router_embedder,
)
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
DEFAULT_SKILL_SEARCH_TOP_K: Final = 5
MAX_SKILL_SEARCH_TOP_K: Final = 100
"""Matches the ``le=100`` bound GET /v1/skills?query= enforces via FastAPI's Query
validation, so the MCP tool can't return a larger payload than the REST endpoint allows."""
MAX_SKILL_SEARCH_TEXT_CHARS: Final = 4000
"""Per-skill cap on the title + description + instructions text that gets embedded, so one
search embeds at most ``MAX_SKILLS_PER_SEARCH * MAX_SKILL_SEARCH_TEXT_CHARS`` characters no
matter how long the stored instructions are."""
@dataclass(frozen=True, slots=True)
class SkillSearchHit:
skill: LiteLLM_SkillsTable
score: float
@dataclass(frozen=True, slots=True)
class SkillSearchHits:
hits: tuple[SkillSearchHit, ...]
@dataclass(frozen=True, slots=True)
class SkillSearchNotConfigured:
reason: str
@dataclass(frozen=True, slots=True)
class SkillSearchEmbeddingFailed:
reason: str
@dataclass(frozen=True, slots=True)
class SkillSearchUnsupportedProvider:
reason: str
SkillSearchOutcome: TypeAlias = SkillSearchHits | SkillSearchNotConfigured | SkillSearchEmbeddingFailed
HostedSkillSearchOutcome: TypeAlias = SkillSearchOutcome | SkillSearchUnsupportedProvider
class SkillSearchResult(BaseModel):
model_config = ConfigDict(frozen=True)
skill_id: str
display_title: str | None
description: str | None
score: float
def skill_search_text(skill: LiteLLM_SkillsTable) -> str:
joined: Final = "\n".join(part for part in (skill.display_title, skill.description, skill.instructions) if part)
return joined[:MAX_SKILL_SEARCH_TEXT_CHARS]
def skill_search_result(hit: SkillSearchHit) -> SkillSearchResult:
return SkillSearchResult(
skill_id=hit.skill.skill_id,
display_title=hit.skill.display_title,
description=hit.skill.description,
score=hit.score,
)
class SkillSearchIndex:
"""Caches one vector per distinct skill text per embedding model, so repeat searches only embed the query."""
def __init__(self, max_entries: int = MAX_SKILLS_PER_SEARCH) -> None:
self._index: Final = SemanticTextIndex(max_entries=max_entries)
async def search(
self,
query: str,
skills: Sequence[LiteLLM_SkillsTable],
top_k: int,
embed: Embedder,
embedding_model: str,
) -> SkillSearchHits | SkillSearchEmbeddingFailed:
texts: Final = tuple(skill_search_text(skill) for skill in skills)
scores: Final = await self._index.scores(query, texts, embed, embedding_model)
if isinstance(scores, EmbeddingFailed):
return SkillSearchEmbeddingFailed(reason=scores.reason)
ranked: Final = sorted(
(SkillSearchHit(skill=skill, score=score) for skill, score in zip(skills, scores, strict=True)),
key=lambda hit: hit.score,
reverse=True,
)
return SkillSearchHits(hits=tuple(ranked[:top_k]))
global_skill_search_index: Final = SkillSearchIndex()
async def search_skills(
query: str,
skills: Sequence[LiteLLM_SkillsTable],
top_k: int,
router: Router | None,
embedding_model: str | None,
index: SkillSearchIndex,
user_api_key_dict: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
) -> SkillSearchOutcome:
if embedding_model is None:
return SkillSearchNotConfigured(
reason="skill search needs litellm_settings.skill_search_embedding_model set to an embedding model from model_list"
)
if router is None:
return SkillSearchNotConfigured(reason="skill search needs a model_list so the embedding model can be called")
embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj)
return await index.search(query, skills, top_k, embed, embedding_model)
async def search_hosted_skills(
custom_llm_provider: str | None,
query: str,
top_k: int,
router: Router | None,
embedding_model: str | None,
index: SkillSearchIndex,
user_api_key_dict: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
) -> HostedSkillSearchOutcome:
"""GET /v1/skills?query= for the skills LiteLLM hosts itself: only ``litellm_proxy`` has a registry to rank."""
if custom_llm_provider != LlmProviders.LITELLM_PROXY.value:
return SkillSearchUnsupportedProvider(reason="query is only supported for custom_llm_provider=litellm_proxy")
skills: Final = await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict=user_api_key_dict)
return await search_skills(
query=query,
skills=skills,
top_k=top_k,
router=router,
embedding_model=embedding_model,
index=index,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)

11
litellm/llms/litellm_proxy/skills/transformation.py
View file

@ -154,7 +154,7 @@ class LiteLLMSkillsTransformationHandler:
user_api_key_dict=user_api_key_dict,
)
return self._db_skill_to_response(db_skill)
return self.db_skill_to_response(db_skill)
def list_skills_handler(
self,
@ -222,7 +222,9 @@ class LiteLLMSkillsTransformationHandler:
user_api_key_dict=user_api_key_dict,
)
skills: Final = [self._db_skill_to_response(s) for s in db_skills]
skills: Final = [ # mutable-ok: ListSkillsResponse.data needs list[Skill]; never mutated after
self.db_skill_to_response(s) for s in db_skills
]
return ListSkillsResponse(
data=skills,
has_more=len(skills) >= limit,
@ -288,7 +290,7 @@ class LiteLLMSkillsTransformationHandler:
skill_id=skill_id,
user_api_key_dict=user_api_key_dict,
)
return self._db_skill_to_response(db_skill)
return self.db_skill_to_response(db_skill)
def delete_skill_handler(
self,
@ -354,7 +356,7 @@ class LiteLLMSkillsTransformationHandler:
type=result.get("type", "skill_deleted"),
)
def _db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill:
def db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill:
"""
Convert a database skill record to Anthropic-compatible Skill response.
@ -375,4 +377,5 @@ class LiteLLMSkillsTransformationHandler:
latest_version=db_skill.latest_version,
source=db_skill.source or "custom",
type="skill",
description=db_skill.description,
)

79
litellm/llms/openai/responses/transformation.py
View file

@ -1,15 +1,17 @@
from collections.abc import Mapping, Sequence
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, get_type_hints
import httpx
from openai.types.responses import ResponseReasoningItem
from pydantic import BaseModel, ValidationError
from pydantic import BaseModel, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import process_response_headers
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_safe_convert_created_field,
)
@ -42,6 +44,30 @@ _PROVIDERS_WITH_COMBINATOR_REJECTING_VALIDATOR: Final = frozenset({LlmProviders.
_PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI})
class _ReasoningSupportEntry(BaseModel):
litellm_provider: str | None = None
supports_reasoning: bool | None = None
_BUNDLED_COST_MAP: Final = TypeAdapter(dict[str, _ReasoningSupportEntry])
@lru_cache(maxsize=1)
def _bundled_openai_reasoning_models() -> frozenset[str]:
"""OpenAI models the cost map shipped with this release flags as reasoning models.
The live map can lag this release (a pinned mirror, or a proxy on newer code than the
map it fetches), and a lagging entry must never strip `reasoning` from a model this
release knows accepts it.
"""
bundled: Final = _BUNDLED_COST_MAP.validate_json(GetModelCostMap.read_local_model_cost_map_text())
return frozenset(
name
for name, entry in bundled.items()
if entry.litellm_provider == LlmProviders.OPENAI.value and entry.supports_reasoning is True
)
class _DeleteResponseBody(TypedDict):
"""Decoded body of the Responses API delete call."""
@ -95,13 +121,9 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
@staticmethod
def _supports_reasoning_effort_none(model: str) -> bool:
"""Return True if the model supports reasoning.effort='none'."""
from litellm.utils import _supports_factory
from litellm.utils import supports_none_reasoning_effort
return _supports_factory(
model=model,
custom_llm_provider=None,
key="supports_none_reasoning_effort",
)
return supports_none_reasoning_effort(model=model, custom_llm_provider=None)
@staticmethod
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
@ -117,6 +139,28 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
return OpenAIGPT5Config.effort_resolves_to_none(model, effort)
@staticmethod
def _supports_reasoning_param(model: str) -> bool:
from litellm.utils import _get_model_info_helper
try:
info: Final = _get_model_info_helper(
model=model.split("/")[-1], custom_llm_provider=LlmProviders.OPENAI.value
)
except Exception:
return True
declared: Final = info.get("supports_reasoning")
if declared is not None:
return declared
return info["key"] in _bundled_openai_reasoning_models()
@staticmethod
def _requests_reasoning_effort(reasoning: object) -> bool:
effort: Final = (
reasoning.get("effort") if isinstance(reasoning, Mapping) else getattr(reasoning, "effort", None)
)
return effort is not None
@staticmethod
def _enforce_min_max_output_tokens(max_output_tokens: "int | None") -> "int | None":
"""Raise sub-minimum max_output_tokens up to the OpenAI Responses API minimum.
@ -166,6 +210,23 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
if "max_output_tokens" in params:
params["max_output_tokens"] = self._enforce_min_max_output_tokens(params.get("max_output_tokens"))
if (
self.custom_llm_provider == LlmProviders.OPENAI
and self._requests_reasoning_effort(params.get("reasoning"))
and not self._supports_reasoning_param(model=model)
):
if drop_params or litellm.drop_params:
params.pop("reasoning", None)
else:
raise litellm.UnsupportedParamsError(
message=(
f"{model} doesn't support `reasoning.effort` "
"(its model cost map entry lacks `supports_reasoning`). "
"To drop unsupported params set `litellm.drop_params = True`"
),
status_code=400,
)
if self._is_gpt_5_model(model=model):
temperature: Final = params.get("temperature")
if temperature is not None and temperature != 1:
@ -478,7 +539,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
processed_headers: Final = process_response_headers(raw_response_headers)
try:
response = ResponsesAPIResponse.model_validate(raw_response_json)
except Exception:
except ValidationError:
verbose_logger.debug(
"Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json
)
@ -870,7 +931,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
try:
response = ResponsesAPIResponse.model_validate(raw_response_json)
except Exception:
except ValidationError:
verbose_logger.debug(
"Error constructing ResponsesAPIResponse: %s, using model_construct", raw_response_json
)

131
litellm/llms/vertex_ai/batches/handler.py
View file

@ -1,6 +1,7 @@
import json
from collections.abc import Coroutine
from collections.abc import Coroutine, Sequence
from typing import TYPE_CHECKING, Final, Protocol
from urllib.parse import urlparse
import httpx
from typing_extensions import ReadOnly, TypedDict
@ -12,11 +13,13 @@ from litellm.litellm_core_utils.url_utils import (
safe_get,
)
from litellm.llms.custom_httpx.http_handler import (
HTTPHandler,
_get_httpx_client,
get_async_httpx_client,
)
from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path
from litellm.types.llms.openai import CreateBatchRequest
from litellm.types.llms.vertex_ai import (
VERTEX_CREDENTIALS_TYPES,
@ -55,6 +58,20 @@ class _FetchedResponseView(TypedDict):
response: ReadOnly[httpx.Response]
class _VertexEndpointDeployedModel(TypedDict, total=False):
model: ReadOnly[str]
class _VertexEndpointResponse(TypedDict, total=False):
deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]]
class _VertexEndpointPayloadView(TypedDict):
"""Holds one decoded GET endpoints/<id> response so the payload reads back typed."""
payload: ReadOnly[_VertexEndpointResponse]
def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse:
return response.json()
@ -78,7 +95,17 @@ class VertexAIBatchPrediction(VertexLLM):
vertex_location: str | None,
timeout: float | httpx.Timeout,
max_retries: int | None,
custom_endpoint: bool | None = None,
) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]:
if custom_endpoint:
raise VertexAIError(
status_code=400,
message=(
"Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
"The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
"use a publisher model or fine-tuned Gemini endpoint deployment instead."
),
)
sync_handler: Final = _get_httpx_client()
access_token, project_id = self._ensure_access_token(
@ -87,6 +114,26 @@ class VertexAIBatchPrediction(VertexLLM):
custom_llm_provider="vertex_ai",
)
headers: Final = {
"Content-Type": "application/json; charset=utf-8",
"Authorization": f"Bearer {access_token}",
}
transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
request=create_batch_data,
vertex_project=vertex_project or project_id,
vertex_location=vertex_location or "us-central1",
)
)
vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
vertex_batch_request=transformed_batch_request,
headers=headers,
sync_handler=sync_handler,
api_base=api_base,
vertex_location=vertex_location or "us-central1",
)
default_api_base: Final = self.create_vertex_batch_url(
vertex_location=vertex_location or "us-central1",
vertex_project=vertex_project or project_id,
@ -111,17 +158,6 @@ class VertexAIBatchPrediction(VertexLLM):
vertex_api_version="v1",
)
headers: Final = {
"Content-Type": "application/json; charset=utf-8",
"Authorization": f"Bearer {access_token}",
}
vertex_batch_request: Final[VertexAIBatchPredictionJob] = (
VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
request=create_batch_data
)
)
if _is_async is True:
return self._async_create_batch(
vertex_batch_request=vertex_batch_request,
@ -142,6 +178,77 @@ class VertexAIBatchPrediction(VertexLLM):
)
return vertex_batch_response
@staticmethod
def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str:
"""
Builds the GET url for resolving an endpoint resource (`projects/../endpoints/<id>`).
A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the
version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a
mount prefix in front of the full default path. The `:operation` suffix convention from
`_check_custom_proxy` does not apply to a plain resource GET.
"""
default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
if not api_base:
return default_endpoint_url
api_base_path: Final = urlparse(api_base).path.rstrip("/")
if api_base_path in ("/v1", "/v1beta1"):
return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url)
return api_base.rstrip("/") + urlparse(default_endpoint_url).path
def _resolve_fine_tuned_endpoint_model(
self,
vertex_batch_request: VertexAIBatchPredictionJob,
headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers
sync_handler: HTTPHandler,
api_base: str | None,
vertex_location: str,
) -> VertexAIBatchPredictionJob:
"""
A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only
accepts Model resources, so swap the endpoint resource for its deployed tuned model
(`projects/../locations/../models/<id>`) read from GET endpoints/<id>.
"""
model: Final = vertex_batch_request.get("model", "")
if "/endpoints/" not in model:
return vertex_batch_request
endpoint_url: Final = self._build_endpoint_resolution_url(
api_base=api_base,
model=model,
vertex_location=vertex_location,
)
# ``api_base`` can come from caller-supplied request kwargs, so wrap the
# fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata
# targets before the bearer token leaves the process (mirrors retrieve_batch).
fetched: Final[_FetchedResponseView] = {
"response": safe_get(
sync_handler,
endpoint_url,
headers=headers,
)
}
response: Final = fetched["response"]
if response.status_code != 200:
raise VertexAIError(
status_code=response.status_code,
message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}",
)
payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()}
deployed_models: Final = payload_view["payload"].get("deployedModels") or ()
deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else ""
if not deployed_model:
raise VertexAIError(
status_code=400,
message=(
f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model "
"resource to run batch predictions against"
),
)
resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model}
return resolved_request
async def _async_create_batch(
self,
vertex_batch_request: VertexAIBatchPredictionJob,

70
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -22,6 +22,8 @@ class VertexAIBatchTransformation:
def transform_openai_batch_request_to_vertex_ai_batch_request(
cls,
request: CreateBatchRequest,
vertex_project: str | None = None,
vertex_location: str | None = None,
) -> VertexAIBatchPredictionJob:
"""
Transforms OpenAI Batch requests to Vertex AI Batch requests
@ -31,7 +33,11 @@ class VertexAIBatchTransformation:
if input_file_id is None:
raise ValueError("input_file_id is required, but not provided")
input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl")
model: Final[str] = cls._get_model_from_gcs_file(input_file_id)
model: Final[str] = cls._get_batch_job_model(
input_file_id=input_file_id,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
output_config: Final[OutputConfig] = OutputConfig(
predictionsFormat="jsonl",
gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)),
@ -188,6 +194,33 @@ class VertexAIBatchTransformation:
path_parts: Final = input_file_id.rsplit("/", 1)
return path_parts[0]
@classmethod
def _get_batch_job_model(
cls,
input_file_id: str,
vertex_project: str | None,
vertex_location: str | None,
) -> str:
"""
Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or
the full `projects/../locations/../endpoints/<id>` resource name for a fine-tuned endpoint.
The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource
to its deployed tuned model (`projects/../locations/../models/<id>`) before sending the job.
"""
parsed_model: Final = cls._get_model_from_gcs_file(input_file_id)
if not parsed_model.startswith("endpoints/"):
return parsed_model
if not vertex_project:
raise VertexAIError(
status_code=400,
message=(
f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require "
"`vertex_project` to build the endpoint resource name"
),
)
return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}"
@classmethod
def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str:
"""
@ -202,6 +235,9 @@ class VertexAIBatchTransformation:
gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8
returns: "publishers/google/models/gemini-1.5-flash-001"
Fine-tuned Gemini endpoints are stored as `endpoints/<numeric id>` in the uri and returned
in that form.
Raises a 400 `VertexAIError` when the uri carries no parseable model path.
"""
model: Final = cls._parse_model_from_gcs_file(gcs_file_uri)
@ -210,11 +246,13 @@ class VertexAIBatchTransformation:
status_code=400,
message=(
"Vertex AI batch creation requires the model to be part of `input_file_id`, but "
f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' path segment. "
f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' or "
"'endpoints/<numeric endpoint id>' path segment. "
"Either upload the input file through LiteLLM (POST /v1/files with "
"custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or "
"pass a uri of the form "
"gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file>"
"gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file> "
"(or gs://<bucket>/<prefix>/endpoints/<numeric endpoint id>/<file> for fine-tuned models)"
),
)
return model
@ -222,18 +260,26 @@ class VertexAIBatchTransformation:
@classmethod
def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None:
"""
Returns the `publishers/<publisher>/models/<model>` path from a gcs uri, or None if the uri
does not contain one.
Returns the `publishers/<publisher>/models/<model>` or `endpoints/<numeric id>` path from a
gcs uri, or None if the uri does not contain one.
A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is
used, so a user-configured bucket prefix that happens to contain `endpoints/<digits>` cannot
override the model path LiteLLM appended after it.
"""
_, separator, model_path = unquote(gcs_file_uri).partition("publishers/")
if not separator:
return None
unquoted_uri: Final = unquote(gcs_file_uri)
_, separator, model_path = unquoted_uri.partition("publishers/")
if separator:
parts: Final = model_path.split("/")
if len(parts) >= 3 and parts[1] == "models" and parts[2]:
return f"publishers/{'/'.join(parts[:3])}"
parts: Final = model_path.split("/")
if len(parts) < 3 or parts[1] != "models" or not parts[2]:
return None
_, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/")
endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else ""
if endpoint_id.isdigit():
return f"endpoints/{endpoint_id}"
return f"publishers/{'/'.join(parts[:3])}"
return None
@classmethod
def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:

13
litellm/llms/vertex_ai/common_utils.py
View file

@ -370,6 +370,19 @@ def get_vertex_base_model_name(model: str) -> str:
return model
def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None:
"""
Fine-tuned Gemini deployments are addressed by a numeric endpoint id,
configured as `vertex_ai/<id>` or `vertex_ai/gemini/<id>`.
Returns the endpoint id, or None when `model` is a regular publisher model.
Mirrors the online chat path in `_get_vertex_url`, which sends numeric
models to `endpoints/{id}` instead of `publishers/google/models/{model}`.
"""
candidate: Final = model.split("/")[-1] if "gemini/" in model else model
return candidate if candidate.isdigit() else None
def validate_vertex_location(vertex_location: str | None) -> str:
"""
Validate a Vertex AI location before interpolating it into a request host or

51
litellm/llms/vertex_ai/files/transformation.py
View file

@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import (
)
from litellm.llms.vertex_ai.common_utils import (
_convert_vertex_datetime_to_openai_datetime,
get_vertex_ai_fine_tuned_endpoint_id,
)
from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@ -707,20 +708,39 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
def _get_gcs_object_name_from_batch_jsonl(
self,
openai_jsonl_content: list[dict[str, Any]],
deployment_model: str | None = None,
) -> str:
"""
Gets a unique GCS object name for the VertexAI batch prediction job
named as: litellm-vertex-{model}-{uuid}
The stored model path decides which Vertex model the batch job later executes against, so
`deployment_model` (the deployment's own configured model) wins over the user-supplied
JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no
deployment config.
Fine-tuned Gemini deployments (numeric endpoint ids) are stored under
`endpoints/<id>` so the batch transformation can round-trip them into a
`projects/../locations/../endpoints/<id>` batch job model instead of a
nonexistent publisher model.
"""
_model = openai_jsonl_content[0].get("body", {}).get("model", "")
if "publishers/google/models" not in _model:
_model = f"publishers/google/models/{_model}"
safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model")
raw_model: Final = (
deployment_model.removeprefix("vertex_ai/")
if deployment_model
else openai_jsonl_content[0].get("body", {}).get("model", "")
)
endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
model_path: Final = (
f"endpoints/{endpoint_id}"
if endpoint_id is not None
else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}")
)
safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model")
object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}"
return object_name
def get_object_name(self, file_data: FileTypes, purpose: str) -> str:
def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str:
"""
Get the object name for the request.
@ -728,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
upload is never materialized just to derive the GCS object name.
"""
if purpose == "batch":
## 1. If jsonl, derive the object name from the first entry's model
## 1. If jsonl, derive the object name from the deployment model (or the first entry's)
first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None)
if first_entry is not None:
return self._get_gcs_object_name_from_batch_jsonl([first_entry])
return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model)
## 2. If not jsonl, store under a server-generated managed object name
filename, _ = extract_file_metadata(file_data)
@ -761,6 +781,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
"""
Get the complete url for the request
"""
if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"):
raise VertexAIError(
status_code=400,
message=(
"Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
"The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
"remove this deployment from the batch request (e.g. `target_model_names`) or "
"use a publisher model / fine-tuned Gemini endpoint instead."
),
)
bucket_name = self._get_configured_bucket_name(litellm_params)
bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name)
file_data: Final = data.get("file")
@ -769,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
raise ValueError("file is required")
if purpose is None:
raise ValueError("purpose is required")
object_name = self.get_object_name(file_data, purpose)
configured_model: Final = litellm_params.get("model")
object_name = self.get_object_name(
file_data,
purpose,
deployment_model=configured_model if isinstance(configured_model, str) else None,
)
if object_prefix:
object_name = f"{object_prefix}/{object_name}"
encoded_object_name: Final = encode_gcs_object_name_for_url(object_name)

8
litellm/model_prices_and_context_window_backup.json
View file

@ -30791,6 +30791,9 @@
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 3e-05,
"reasoning_effort_levels": [
"medium"
],
"source": "https://developers.openai.com/api/docs/models/chat-latest",
"supported_endpoints": [
"/v1/chat/completions",
@ -30808,6 +30811,7 @@
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -37225,6 +37229,7 @@
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -37265,6 +37270,7 @@
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -37476,6 +37482,7 @@
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -37516,6 +37523,7 @@
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,

13
litellm/models/mcp_server.py
View file

@ -6,10 +6,12 @@ Canonical definition for ``litellm_mcpservertable``. Re-exported from
"""
import enum
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import Literal
from pydantic import Field
from pydantic import Field, ValidationInfo, field_validator
from litellm.types.llms.base import LiteLLMPydanticObjectBase
from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType
@ -115,3 +117,12 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
submitted_at: datetime | None = None
reviewed_at: datetime | None = None
review_notes: str | None = None
@field_validator("static_headers", "env", mode="before")
@classmethod
def decode_stored_secret_map(cls, value: object, info: ValidationInfo) -> Mapping[str, str] | None:
from litellm.proxy.common_utils.encrypt_decrypt_utils import decode_secret_map
if value is None and info.field_name == "env":
return MappingProxyType({})
return decode_secret_map(value, key=info.field_name or "secret map")

79
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -3108,15 +3108,17 @@ class MCPRequestHandler:
@staticmethod
async def _get_allowed_mcp_servers_for_agent(
user_api_key_auth: UserAPIKeyAuth | None = None,
agent_object_permission=None,
agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
) -> list[str]:
"""
Get allowed MCP servers for an agent (from the agent's object_permission).
Returns the MCP servers from the agent's object_permission.
If agent has no object_permission, returns [] (no extra restriction). An entitlement the
agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the
resolver denies.
Returns the agent's direct servers, the servers in its access groups, and the servers reached
through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no
object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that
cannot be read, or a declared toolset that resolves to no grants, raises
``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the
agent as unrestricted.
Args:
user_api_key_auth: User auth with agent_id
@ -3126,31 +3128,30 @@ class MCPRequestHandler:
if not user_api_key_auth or not user_api_key_auth.agent_id:
return []
obj_perm = agent_object_permission
if obj_perm is None:
obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
obj_perm: Final = (
agent_object_permission
if agent_object_permission is not None
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
)
if obj_perm is None:
return []
try:
direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or []
if isinstance(direct_mcp_servers, str):
direct_mcp_servers = []
mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or []
if isinstance(mcp_access_groups, str):
mcp_access_groups = []
# Permission entries may be server_ids OR names/aliases — expand to ids.
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers))
access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups)
all_servers: Final = expanded_direct_servers + access_group_servers
return list(set(all_servers))
expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(
obj_perm.mcp_servers or []
)
access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(
obj_perm.mcp_access_groups or []
)
toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm)
return list({*expanded_direct_servers, *access_group_servers, *toolset_grants})
except Exception as e:
if isinstance(e, UnloadableEntitlementError):
raise
verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e)
return []
@ -3158,13 +3159,15 @@ class MCPRequestHandler:
async def _get_agent_tool_permissions_for_server(
server_id: str,
user_api_key_auth: UserAPIKeyAuth | None = None,
agent_object_permission=None,
agent_object_permission: LiteLLM_ObjectPermissionTable | None = None,
) -> list[str] | None:
"""
Get allowed tool names for a server from the agent's object_permission.
Returns None if agent has no tool restrictions for this server. An entitlement the agent
LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the
tool resolver turns into deny-all for the server rather than an unrestricted tool list.
Get allowed tool names for a server from the agent's object_permission: the union of its
direct tool permissions and the tools its toolsets grant on that server, mirroring the key and
team levels. Returns None if agent has no tool restrictions for this server. An entitlement the
agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises
``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the
server rather than an unrestricted tool list.
Args:
server_id: Server ID to check permissions for
@ -3175,24 +3178,30 @@ class MCPRequestHandler:
if not user_api_key_auth or not user_api_key_auth.agent_id:
return None
obj_perm = agent_object_permission
if obj_perm is None:
obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
obj_perm: Final = (
agent_object_permission
if agent_object_permission is not None
else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth)
)
if obj_perm is None:
return None
try:
mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None)
if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict):
return None
# Dict keys may be server_ids OR names/aliases; normalize before lookup.
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id)
return list(tools) if tools else None
direct_tools: Final = (
global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id)
if obj_perm.mcp_tool_permissions
else None
)
toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id)
agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools)
return list(agent_tools) if agent_tools else None
except Exception as e:
if isinstance(e, UnloadableEntitlementError):
raise
verbose_logger.warning("Failed to get agent tool permissions for server: %s", e)
return None

77
litellm/proxy/_experimental/mcp_server/db.py
View file

@ -4,7 +4,7 @@ import hashlib
import json
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -23,8 +23,11 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
SecretMapDecodeError,
_get_salt_key,
decode_secret_map,
decrypt_value_helper,
encrypt_secret_map,
encrypt_value_helper,
)
from litellm.proxy.utils import PrismaClient
@ -122,6 +125,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False):
server_id: str
OAuthGrantState = Literal["valid", "refreshable", "absent"]
class _OAuthTokenRefreshResponse(TypedDict, total=False):
access_token: str
refresh_token: str
@ -360,7 +366,7 @@ def _prepare_mcp_server_data(
# exclude_unset filter is respected. Reading back from ``data`` would
# reintroduce defaults (e.g. ``env={}``) for fields the caller never set.
if data_dict.get("static_headers") is not None:
data_dict["static_headers"] = safe_dumps(data_dict["static_headers"])
data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"])
# env_vars is read from ``data_dict`` (not ``data``) like every other JSON
# column so the exclude_unset filter is respected: a partial update that
@ -376,7 +382,7 @@ def _prepare_mcp_server_data(
data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"])
if data_dict.get("env") is not None:
data_dict["env"] = safe_dumps(data_dict["env"])
data_dict["env"] = encrypt_secret_map(data_dict["env"])
if "tool_name_to_display_name" in data_dict:
data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {})
@ -589,6 +595,19 @@ def decrypt_credentials(
return credentials
def _readable_mcp_servers(
rows: Iterable["prisma_db_models.LiteLLM_MCPServerTable"],
) -> Iterable[LiteLLM_MCPServerTable]:
for row in rows:
try:
table = LiteLLM_MCPServerTable.model_validate(row.model_dump())
except SecretMapDecodeError:
verbose_proxy_logger.warning("Skipping MCP server %s: cannot decrypt secret map", row.server_id)
continue
decrypt_global_env_var_values(table.env_vars)
yield table
async def get_all_mcp_servers(
prisma_client: PrismaClient,
approval_status: str | None = None,
@ -609,10 +628,7 @@ async def get_all_mcp_servers(
)
mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where)
tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers]
for table in tables:
decrypt_global_env_var_values(table.env_vars)
return tables
return list(_readable_mcp_servers(mcp_servers))
async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None:
@ -638,13 +654,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
"server_id": {"in": server_ids},
}
)
final_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = []
for _mcp_server in _mcp_servers:
table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump())
decrypt_global_env_var_values(table.env_vars)
final_mcp_servers.append(table)
return final_mcp_servers
return list(_readable_mcp_servers(_mcp_servers))
async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]:
@ -852,12 +862,10 @@ async def create_mcp_server(
data_dict["created_by"] = touched_by
data_dict["updated_by"] = touched_by
new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create(
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
)
new_mcp_server: Final = await MCPServerRepository(prisma_client).table.create(data=data_dict)
_decrypt_env_vars_on_returned_row(new_mcp_server)
return new_mcp_server
return LiteLLM_MCPServerTable.model_validate(new_mcp_server.model_dump())
async def create_draft_mcp_server(
@ -1066,13 +1074,13 @@ async def update_mcp_server(
data_dict["credentials"] = Json(None)
updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update(
updated_mcp_server: Final = await MCPServerRepository(prisma_client).table.update(
where={"server_id": data.server_id},
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
data=data_dict,
)
_decrypt_env_vars_on_returned_row(updated_mcp_server)
return updated_mcp_server
return LiteLLM_MCPServerTable.model_validate(updated_mcp_server.model_dump()) if updated_mcp_server else None
async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None:
@ -1144,6 +1152,13 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
if rotated_env_vars is not None:
update_data["env_vars"] = safe_dumps(rotated_env_vars)
for field in ("static_headers", "env"):
try:
if secret_map := decode_secret_map(getattr(mcp_server, field, None), key=field):
update_data[field] = encrypt_secret_map(secret_map, new_encryption_key=new_master_key)
except SecretMapDecodeError:
verbose_proxy_logger.warning("Cannot rotate MCP %s for server %s", field, mcp_server.server_id)
if not update_data:
continue
@ -1453,6 +1468,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in
return False
def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState:
"""Classify local grant readiness without attempting a refresh or checking upstream revocation."""
if not cred or not cred.get("access_token"):
return "absent"
if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
return "valid"
return "refreshable" if cred.get("refresh_token") else "absent"
async def get_user_oauth_credential(
prisma_client: PrismaClient,
user_id: str,
@ -1715,12 +1739,11 @@ async def resolve_valid_user_oauth_token(
dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh
actually happens, so the valid-token path never requires a DB handle.
"""
if not cred or not cred.get("access_token"):
grant: Final = oauth_grant_state(cred)
if cred is None or grant == "absent":
return None
if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS):
if grant == "valid":
return cred
if not cred.get("refresh_token"):
return None
if prisma_client is None:
from litellm.proxy.utils import get_prisma_client_or_throw
@ -1894,9 +1917,7 @@ async def get_mcp_submissions(
order={"submitted_at": "desc"},
take=500, # safety cap; paginate if needed in a future iteration
)
items: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows]
for item in items:
decrypt_global_env_var_values(item.env_vars)
items: Final = list(_readable_mcp_servers(rows))
pending: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review)
active: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active)

69
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import (
render_token_fault,
)
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
VendorCredentialState,
aggregate_authorize,
aggregate_token,
complete_connect_flow,
describe_connect_flow,
introspect_gateway_token,
is_gateway_dcr_client_id,
is_proxy_api_resource,
@ -798,21 +800,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC
return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302)
async def _bridge_authorize_access_denial(
litellm_user_id: str,
mcp_server: MCPServer,
redirect_uri: str,
state: str,
) -> RedirectResponse | None:
"""The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.
Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the
same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting
session can actually list and call the server's tools. Without this gate the flow completes, the
client shows connected, and every tool request fail-closes to an empty list with nothing telling
the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or
deactivated user denies like a missing grant, fail closed.
"""
async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool:
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
@ -821,13 +809,22 @@ async def _bridge_authorize_access_denial(
)
try:
admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id)
admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id)
except HTTPException as exc:
if exc.status_code >= 500:
raise
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
if mcp_server.server_id in allowed_server_ids:
return False
return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted)
async def _bridge_authorize_access_denial(
litellm_user_id: str,
mcp_server: MCPServer,
redirect_uri: str,
state: str,
) -> RedirectResponse | None:
"""The denial redirect for a signed-in user who cannot reach the target server, or None to proceed."""
if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id):
return None
return _bridge_access_denied_redirect(redirect_uri, state, mcp_server)
@ -1910,6 +1907,38 @@ async def token_endpoint(
)
async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState:
"""Whether the gateway itself can see a live vendor credential for this user and server.
The one reading of "authorized" the connect page displays and the finish step enforces, so
the button a user sees and the grant they get cannot disagree. A read fault is neither, and
fails the scoped grant closed."""
from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # circular import at module load
get_user_oauth_credential,
oauth_grant_state,
)
from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load
if prisma_client is None:
return "unavailable"
try:
credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id)
except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed
return "unavailable"
return "absent" if oauth_grant_state(credential) == "absent" else "present"
@router.get("/authorize/flow")
async def authorize_flow(request: Request, flow: str) -> Response:
return await describe_connect_flow(
request=request,
flow_handle=flow,
session_user_id=_session_cookie_user_id(request),
lookup_vendor_credential=_vendor_credential_state,
lookup_server_reachability=_user_can_reach_mcp_server,
)
@router.post("/authorize/complete")
async def authorize_complete(
request: Request,
@ -1934,6 +1963,8 @@ async def authorize_complete(
delivery=delivery,
team_id=team_id,
decision=decision,
lookup_vendor_credential=_vendor_credential_state,
lookup_server_reachability=_user_can_reach_mcp_server,
)

187
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code"
ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"]
ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]]
"""Injected live-user revalidation (the token endpoint's mirror of admission):
``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is
a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else
fails the grant closed."""
VendorCredentialState = Literal["present", "absent", "unavailable"]
"""The per-user vendor credential read has three outcomes: present, absent, or unavailable."""
_DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry"
_DB_FAULTED_DESCRIPTION: Final = (
@ -195,6 +193,16 @@ class ConsentTeam(BaseModel):
team_alias: str | None = None
class LookupVendorCredential(Protocol):
"""Injected read of a user's vendor credential for one server."""
def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ...
class LookupServerReachability(Protocol):
def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ...
class LookupConsentTeams(Protocol):
"""Injected lookup of the teams a signed-in user may bind a proxy-API credential to."""
@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr
return "unresolvable"
async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState:
return "unavailable"
async def _unreachable_server(user_id: str, server_id: str) -> bool:
return False
class GatewayDcrClient(BaseModel):
"""The registration record sealed into a gateway DCR ``client_id``.
@ -449,7 +465,10 @@ def aggregate_authorize(
A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the
flow to that one server: the scope is sealed into the flow, carried into the code, and
bound into the session token, while the connect page interlude runs exactly as before.
bound into the session token. The connect URL carries only the flow handle; the page
learns the client origin, the scoped server, and whether its vendor OAuth is done from
:func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry
steers which server the page authorizes or names on the confirmation.
Validation failures respond directly with 400 and never redirect: per RFC 6749
section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and
@ -474,10 +493,7 @@ def aggregate_authorize(
resource_server_id=scoped_server.server_id if scoped_server is not None else None,
audience=None,
)
connect_url: Final = _append_query_params(
f"{base_url}/ui/connect",
(("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))),
)
connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),))
response: Final = RedirectResponse(connect_url, status_code=303)
_set_flow_cookie(response, request, handle, flow)
return response
@ -684,6 +700,99 @@ def _origin_only(url: str) -> str:
return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else ""
def _open_flow_for(
request: Request, flow_handle: str, session_user_id: str | None, now: datetime
) -> _ConnectFlow | Response:
sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
if sealed_flow is None:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
if flow is None or now.timestamp() >= flow.exp:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
if session_user_id is None:
return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
if session_user_id != flow.user_id:
return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
return flow
async def _flow_target(
flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability
) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]:
if flow.resource_server_id is None:
return "unscoped", None
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle
MCPServerManager,
global_mcp_server_manager,
)
server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id)
if (
server is None
or not server.is_gateway_managed_oauth2
or not await lookup_server_reachability(flow.user_id, server.server_id)
):
return "stale", None
state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive"
return state, server
class ConnectFlowDescription(TypedDict):
"""What the connect page is allowed to know about one in-flight flow."""
state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]]
client_origin: ReadOnly[str]
server_id: ReadOnly[str | None]
server_name: ReadOnly[str | None]
connected: ReadOnly[bool | None]
async def _describe_opened_flow(
flow: _ConnectFlow,
lookup_vendor_credential: LookupVendorCredential,
lookup_server_reachability: LookupServerReachability,
) -> ConnectFlowDescription | Response:
state, server = await _flow_target(flow, lookup_server_reachability)
if state == "interactive" and server is not None:
credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id)
if credential == "unavailable":
return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION)
interactive_description: Final[ConnectFlowDescription] = {
"state": state,
"client_origin": _origin_only(flow.redirect_uri),
"server_id": server.server_id,
"server_name": server.server_name or server.alias or server.name,
"connected": credential == "present",
}
return interactive_description
described: Final[ConnectFlowDescription] = {
"state": state,
"client_origin": _origin_only(flow.redirect_uri),
"server_id": None if server is None else server.server_id,
"server_name": None if server is None else (server.server_name or server.alias or server.name),
"connected": state == "m2m" or None,
}
return described
async def describe_connect_flow(
request: Request,
flow_handle: str,
session_user_id: str | None,
lookup_vendor_credential: LookupVendorCredential,
lookup_server_reachability: LookupServerReachability,
) -> Response:
opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc))
if isinstance(opened, Response):
return opened
described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
return (
described
if isinstance(described, Response)
else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS)
)
async def complete_connect_flow(
request: Request,
flow_handle: str,
@ -692,56 +801,34 @@ async def complete_connect_flow(
delivery: str | None = None,
team_id: str | None = None,
decision: str | None = None,
lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential,
lookup_server_reachability: LookupServerReachability = _unreachable_server,
) -> Response:
"""The deliberate finish step of the connect flow: mint the gateway authorization
code and send the browser back to the client.
"""Mint the code only after a deliberate POST by the sealed user.
Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly
per-flow cookie plus an exact match between the signed-in user and the user sealed
into the flow: a link crafted by another party dies here with ``access_denied``
instead of minting a code for the victim's identity. The flow is single-use (an atomic
claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in.
``delivery`` chooses how the code reaches the client. Default (absent or
``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"``
renders the callback URL on a page instead, for a client whose redirect URI is a
loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box,
container): the 303 would dereference the browser machine's loopback and the code
would never arrive, so the user carries it over by pasting the URL into the client or
fetching it from the client machine's terminal. Manual delivery is honored only for
loopback redirect URIs; a routable redirect URI works from any browser by
construction, so those flows always redirect. The user who sees the page is exactly
the user the 303 would have carried the code to, and the same user already sees the
code today in the dead redirect's address bar, so the page exposes the code to no new
party. Unknown ``delivery`` values are rejected rather than defaulted: a client that
asked for manual delivery and got a dead redirect instead would silently lose its
code.
``decision`` and ``team_id`` come from the native-client consent page. ``"deny"``
burns the flow and sends the client ``error=access_denied`` so it stops waiting;
``team_id`` is sealed into the code only for proxy-API flows, where it picks which of
the user's teams the minted credential is attributed to.
A scoped flow additionally requires its sealed server to have a live vendor credential
before a code can be minted. The check happens before the single-use claim, so a
premature submit can be retried after authorization; denial deliberately bypasses it.
"""
if delivery not in (None, "redirect", "manual"):
return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'")
if decision not in (None, "approve", "deny"):
return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'")
sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle))
if sealed_flow is None:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
if flow is None:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
now: Final = datetime.now(timezone.utc)
if now.timestamp() >= flow.exp:
return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection")
if session_user_id is None:
return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
if session_user_id != flow.user_id:
return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
opened: Final = _open_flow_for(request, flow_handle, session_user_id, now)
if isinstance(opened, Response):
return opened
if decision != "deny":
described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability)
if isinstance(described, Response):
return described
if described["state"] == "stale":
return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available")
if described["connected"] is False:
return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing")
flow_refusal: Final = _claim_refusal(
await _SingleUseGuard(cache).claim(
f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
),
replayed=_oauth_error(
400, "invalid_request", "this connect flow was already completed; restart the connection"
@ -750,7 +837,7 @@ async def complete_connect_flow(
if flow_refusal is not None:
return flow_refusal
response: Final = (
_denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now)
_denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now)
)
path, secure = _cookie_path_and_secure(request)
response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax")

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

@ -6272,8 +6272,7 @@ class MCPServerManager:
]
}
)
db_mcp_servers: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in raw_rows]
verbose_logger.info("Found %s MCP servers in database", len(db_mcp_servers))
verbose_logger.info("Found %s MCP servers in database", len(raw_rows))
previous_registry: Final = self.registry
new_registry: Final[dict[str, MCPServer]] = {}
@ -6281,8 +6280,9 @@ class MCPServerManager:
# Stage one: build every server. Stage two assigns short prefixes
# against the *full* set so dedup is deterministic regardless of
# iteration order.
for server in db_mcp_servers:
for row in raw_rows:
try:
server = LiteLLM_MCPServerTable.model_validate(row.model_dump())
existing_server = previous_registry.get(server.server_id)
if (
@ -6320,8 +6320,8 @@ class MCPServerManager:
except Exception as e:
verbose_logger.exception(
"Skipping MCP server %s (%s) during DB reload: %s",
server.server_id,
getattr(server, "alias", None),
getattr(row, "server_id", None),
getattr(row, "alias", None),
e,
)

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

@ -1,6 +1,7 @@
import asyncio
import importlib
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal
@ -8,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal
import anyio
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from starlette.datastructures import Headers
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT
@ -104,6 +106,9 @@ if MCP_AVAILABLE:
from mcp.types import Tool as MCPTool
from litellm.experimental_mcp_client.client import MCPClient
from litellm.llms.litellm_proxy.skills.skill_search import (
DEFAULT_SKILL_SEARCH_TOP_K,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
_UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES,
global_mcp_server_manager,
@ -188,10 +193,12 @@ if MCP_AVAILABLE:
AGENT_SEARCH_TOOL_NAME,
DEFAULT_AGENT_SEARCH_TOP_K,
MCP_TOOL_SEARCH_TOOL_NAME,
SKILL_SEARCH_TOOL_NAME,
coerce_top_k,
handle_agent_search,
handle_mcp_tool_call,
handle_mcp_tool_search,
handle_skill_search,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj
@ -210,6 +217,14 @@ if MCP_AVAILABLE:
),
user_api_key_dict=user_api_key_dict,
)
if tool_name == SKILL_SEARCH_TOOL_NAME:
return await handle_skill_search(
query=str(tool_arguments.get("query", "")),
top_k=coerce_top_k(
tool_arguments.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K
),
user_api_key_dict=user_api_key_dict,
)
rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)
(
virtual_mcp_auth_header,
@ -1153,6 +1168,45 @@ if MCP_AVAILABLE:
scopes: Final[list[str] | None] = scopes_raw if isinstance(scopes_raw, list) else None
return client_id, client_secret, scopes
_STAGED_AUTH_VALUE_AUTH_TYPES: Final = frozenset(
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.authorization)
)
@dataclass(frozen=True, slots=True)
class _StagedServerTest:
request: NewMCPServerRequest
mcp_auth_header: str | None
oauth2_headers: dict[str, str] | None
def _stage_server_test(new_mcp_server_request: NewMCPServerRequest, headers: Headers) -> _StagedServerTest:
"""
Resolve the credentials a not-yet-saved server config carries for a preview call.
Both preview endpoints (``/test/connection`` and ``/test/tools/list``) must hand the
temporary client the same credentials, or a server that the saved connection reaches
fine fails one of them.
"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
request: Final = _inherit_credentials_from_existing_server(new_mcp_server_request)
mcp_auth_header: Final = (
request.credentials.get("auth_value")
if request.auth_type in _STAGED_AUTH_VALUE_AUTH_TYPES and isinstance(request.credentials, dict)
else None
)
# Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY):
# when the primary x-litellm-api-key header is absent, the Authorization value is the
# caller's LiteLLM key, not an upstream token, and must never be forwarded upstream.
oauth2_headers: Final = (
MCPRequestHandler._get_oauth2_headers_from_headers(headers)
if request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES
and headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY)
else None
)
return _StagedServerTest(request=request, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers)
async def _list_tools_within(client: MCPClient, deadline: float) -> list[MCPTool] | None:
with anyio.move_on_after(deadline):
return await client.list_tools(raise_on_error=True)
@ -1374,6 +1428,8 @@ if MCP_AVAILABLE:
},
)
staged: Final = _stage_server_test(new_mcp_server_request, request.headers)
async def _test_connection_operation(client):
async def _noop(session):
return "ok"
@ -1382,8 +1438,10 @@ if MCP_AVAILABLE:
return {"status": "ok"}
return await _execute_with_mcp_client(
new_mcp_server_request,
staged.request,
_test_connection_operation,
mcp_auth_header=staged.mcp_auth_header,
oauth2_headers=staged.oauth2_headers,
raw_headers=_safe_get_request_headers(request),
)
@ -1404,37 +1462,11 @@ if MCP_AVAILABLE:
},
)
new_mcp_server_request = _inherit_credentials_from_existing_server(new_mcp_server_request)
staged: Final = _stage_server_test(new_mcp_server_request, request.headers)
# For OpenAPI spec servers, generate tools from the spec directly
if new_mcp_server_request.spec_path:
return await _preview_openapi_tools(new_mcp_server_request.spec_path)
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
MCPRequestHandler,
)
headers: Final = request.headers
mcp_auth_header: str | None = None
if new_mcp_server_request.auth_type in {
MCPAuth.api_key,
MCPAuth.bearer_token,
MCPAuth.basic,
MCPAuth.authorization,
}:
credentials: Final = getattr(new_mcp_server_request, "credentials", None)
if isinstance(credentials, dict):
mcp_auth_header = credentials.get("auth_value")
# Authorization doubles as the admission fallback (LITELLM_API_KEY_HEADER_NAME_SECONDARY):
# when the primary x-litellm-api-key header is absent, the Authorization value is the
# caller's LiteLLM key, not an upstream token, and must never be forwarded upstream.
oauth2_headers: dict[str, str] | None = None
if new_mcp_server_request.auth_type in _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES and headers.get(
MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY
):
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
if staged.request.spec_path:
return await _preview_openapi_tools(staged.request.spec_path)
async def _list_tools_operation(client):
# Bound the whole pagination walk: without this the preview is limited only by the
@ -1465,9 +1497,9 @@ if MCP_AVAILABLE:
}
return await _execute_with_mcp_client(
new_mcp_server_request,
staged.request,
_list_tools_operation,
mcp_auth_header=mcp_auth_header,
oauth2_headers=oauth2_headers,
mcp_auth_header=staged.mcp_auth_header,
oauth2_headers=staged.oauth2_headers,
raw_headers=_safe_get_request_headers(request),
)

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

@ -23,6 +23,7 @@ from pydantic import AnyUrl, ConfigDict
from starlette.requests import Request as StarletteRequest
from starlette.responses import JSONResponse
from starlette.types import Message, Receive, Scope, Send
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
@ -816,6 +817,11 @@ 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
raise McpError(ErrorData(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
@ -911,15 +917,18 @@ if MCP_AVAILABLE:
Returns a CallToolResult when ``name`` is a virtual tool, else ``None`` so
the caller falls through to normal tool routing.
"""
from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
from litellm.proxy._experimental.mcp_server.tool_search import (
AGENT_SEARCH_TOOL_NAME,
DEFAULT_AGENT_SEARCH_TOP_K,
MCP_TOOL_SEARCH_TOOL_NAME,
SKILL_SEARCH_TOOL_NAME,
VIRTUAL_TOOL_NAMES,
coerce_top_k,
handle_agent_search,
handle_mcp_tool_call,
handle_mcp_tool_search,
handle_skill_search,
)
if name not in VIRTUAL_TOOL_NAMES:
@ -961,6 +970,12 @@ if MCP_AVAILABLE:
top_k=coerce_top_k(args.get("top_k", DEFAULT_AGENT_SEARCH_TOP_K), default=DEFAULT_AGENT_SEARCH_TOP_K),
user_api_key_dict=user_api_key_auth,
)
if name == SKILL_SEARCH_TOOL_NAME:
return await handle_skill_search(
query=str(args.get("query", "")),
top_k=coerce_top_k(args.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K),
user_api_key_dict=user_api_key_auth,
)
virtual_logging_obj: Final = await _build_virtual_call_logging_obj(
name=name,
arguments=args,
@ -1086,6 +1101,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_client_ip,
host_progress_callback=host_progress_callback,
**data, # for logging
)
@ -1119,7 +1135,7 @@ if MCP_AVAILABLE:
except HTTPException as e:
verbose_logger.error("HTTPException in MCP tool call: %s", e)
return CallToolResult(
content=[TextContent(text=f"Error: {e.detail}", type="text")],
content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")],
isError=True,
)
except MCPUpstreamAuthError as e:
@ -1383,7 +1399,7 @@ if MCP_AVAILABLE:
########################################################
async def _get_allowed_mcp_servers_from_mcp_server_names(
mcp_servers: list[str] | None,
mcp_servers: Sequence[str] | None,
allowed_mcp_servers: list[MCPServer],
) -> list[MCPServer]:
"""
@ -1404,13 +1420,10 @@ if MCP_AVAILABLE:
server_name_matched = False
for server in allowed_mcp_servers:
if server:
match_list = [s.lower() for s in iter_known_server_prefixes(server) if s]
if server_or_group.lower() in match_list:
filtered_server[server.server_id] = server
server_name_matched = True
break
if server and _server_answers_to(server, server_or_group):
filtered_server[server.server_id] = server
server_name_matched = True
break
if not server_name_matched:
try:
@ -1440,6 +1453,72 @@ if MCP_AVAILABLE:
return allowed_mcp_servers
def _http_detail_message(detail: object) -> str:
return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail)
def _server_answers_to(server: MCPServer, name: str) -> bool:
requested: Final = name.lower()
return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known)
class _McpDeniedDetail(TypedDict):
error: ReadOnly[str]
async def raise_denied_scoped_mcp_access(
requested_names: Sequence[str],
user_api_key_auth: UserAPIKeyAuth | None,
client_ip: str | None = None,
) -> None:
"""A scoped request (``/mcp/<name>`` path or ``x-mcp-servers`` header) resolved to zero
allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy
server with no tools. Unknown, unauthorized, and access-group names all share one generic
error so scoping cannot probe which servers exist; the agent variant fires only when the
same request resolves once the agent binding is stripped, proving the binding caused the veto."""
agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None
if user_api_key_auth is not None and agent_id:
resolved_without_agent: Final = await _get_allowed_mcp_servers(
user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})),
mcp_servers=requested_names,
client_ip=client_ip,
)
def _resolved_to_server(name: str) -> bool:
return any(_server_answers_to(server, name) for server in resolved_without_agent)
vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None)
if vetoed_server is not None:
agent_denial: Final[_McpDeniedDetail] = {
"error": (
f"MCP server '{vetoed_server}' is not available to this key: the key is bound to "
f"agent '{agent_id}', whose MCP grants do not include this server. Add the server "
f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or "
f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
)
}
raise HTTPException(status_code=403, detail=agent_denial)
vetoed_group: Final = next(
(
name
for name in requested_names
if not _resolved_to_server(name)
and any(name in (server.access_groups or ()) for server in resolved_without_agent)
),
None,
)
if vetoed_group is not None:
group_denial: Final[_McpDeniedDetail] = {
"error": (
f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to "
f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the "
f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or "
f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent."
)
}
raise HTTPException(status_code=403, detail=group_denial)
generic_denial: Final[_McpDeniedDetail] = {
"error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}"
}
raise HTTPException(status_code=403, detail=generic_denial)
def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool:
"""
Check if a tool name matches any name in the filter list.
@ -1532,7 +1611,7 @@ if MCP_AVAILABLE:
async def _get_allowed_mcp_servers(
user_api_key_auth: UserAPIKeyAuth | None,
mcp_servers: list[str] | None,
mcp_servers: Sequence[str] | None,
client_ip: str | None = None,
) -> list[MCPServer]:
"""Return allowed MCP servers for a request after applying filters.
@ -1968,6 +2047,12 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
client_ip=client_ip,
)
if mcp_servers and not allowed_mcp_servers:
await raise_denied_scoped_mcp_access(
requested_names=mcp_servers,
user_api_key_auth=user_api_key_auth,
client_ip=client_ip,
)
# Pre-fetch OAuth credentials only when at least one server uses OAuth2,
# to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers.
@ -2395,6 +2480,8 @@ if MCP_AVAILABLE:
)
verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools))
return listing
except HTTPException:
raise
except Exception as e:
verbose_logger.exception("Error getting tools from managed MCP servers: %s", e)
# Continue with an empty listing instead of failing completely
@ -3077,6 +3164,7 @@ if MCP_AVAILABLE:
mcp_server_auth_headers: dict[str, dict[str, str]] | None = None,
oauth2_headers: dict[str, str] | None = None,
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
**kwargs: Any,
) -> CallToolResult:
"""
@ -3107,6 +3195,12 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
allowed_mcp_servers=allowed_mcp_servers,
)
if mcp_servers and not allowed_mcp_servers:
await raise_denied_scoped_mcp_access(
requested_names=mcp_servers,
user_api_key_auth=user_api_key_auth,
client_ip=client_ip,
)
if not allowed_mcp_servers:
raise HTTPException(
status_code=403,

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