scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge branch 'litellm_internal_staging' into litellm_/virtual-keys-e2e-tests-f40710

Browse source
This commit is contained in:
yuneng-jiang 2026-09-07 15:43:38 -07:00 • committed by GitHub
commit 1932aef9bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
271 changed files with 16347 additions and 1902 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
.circleci/config.yml
View file

@ -2648,6 +2648,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 +2791,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:

2
.github/workflows/auto-close-duplicates.yml vendored
View file

@ -22,7 +22,7 @@ on:
permissions: {}
jobs:
test:
sweep-tests:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5

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

129
.github/workflows/report-rust-release-wheel.yml vendored
View file

@ -1,129 +0,0 @@
name: Report LiteLLM Rust release wheel
on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs
workflow_run:
workflows:
- LiteLLM Rust
types:
- completed
permissions: {}
concurrency:
group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }}
cancel-in-progress: false
jobs:
report-release-wheel:
name: report release wheel
if: >-
github.event.workflow_run.event == 'pull_request' &&
github.event.workflow_run.path == '.github/workflows/test-rust.yml' &&
github.event.workflow_run.head_repository.full_name == github.repository &&
github.event.workflow_run.pull_requests[0].number != null
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
pull-requests: write
steps:
- name: Link release wheel report on PR
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
env:
COMMENT_MARKER: "<!-- litellm-release-wheel-size -->"
with:
script: |
const marker = process.env.COMMENT_MARKER;
const workflowRun = context.payload.workflow_run;
const allowedConclusions = new Set([
"action_required",
"cancelled",
"failure",
"neutral",
"skipped",
"stale",
"startup_failure",
"success",
"timed_out",
]);
if (
!allowedConclusions.has(workflowRun.conclusion) ||
workflowRun.event !== "pull_request" ||
workflowRun.path !== ".github/workflows/test-rust.yml" ||
workflowRun.head_repository?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
workflowRun.pull_requests?.length !== 1
) {
throw new Error("unexpected source workflow");
}
const pullRequest = workflowRun.pull_requests[0];
const pullRequestNumber = pullRequest.number;
const headSha = workflowRun.head_sha;
const runId = workflowRun.id;
if (
!Number.isSafeInteger(pullRequestNumber) ||
pullRequestNumber <= 0 ||
!Number.isSafeInteger(runId) ||
runId <= 0 ||
!/^[0-9a-f]{40}$/.test(headSha) ||
pullRequest.head?.sha !== headSha
) {
throw new Error("invalid source workflow metadata");
}
const runUrl =
`${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` +
`/actions/runs/${runId}`;
const result =
workflowRun.conclusion === "success"
? "successfully"
: `with \`${workflowRun.conclusion}\``;
const body = [
marker,
"## LiteLLM Rust workflow",
"",
`Workflow completed ${result} for \`${headSha}\``,
"",
`[View workflow run](${runUrl})`,
].join("\n");
const comments = await github.paginate(github.rest.issues.listComments, {
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
per_page: 100,
});
const existing = comments.find(
(comment) =>
comment.user?.login === "github-actions[bot]" &&
comment.body?.startsWith(marker),
);
const currentPullRequest = (
await github.rest.pulls.get({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: pullRequestNumber,
})
).data;
if (
currentPullRequest.state !== "open" ||
currentPullRequest.head.repo?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
currentPullRequest.head.sha !== headSha
) {
core.info("source workflow no longer matches the current pull request head");
return;
}
if (existing) {
await github.rest.issues.updateComment({
owner: context.repo.owner,
repo: context.repo.repo,
comment_id: existing.id,
body,
});
} else {
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
body,
});
}

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

6
.github/workflows/test-code-quality.yml vendored
View file

@ -74,6 +74,12 @@ jobs:
- name: check_workflow_startup_safety
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py
- name: check_workflow_job_name_collisions
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py
- name: test_workflow_job_name_collisions
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py

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

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

@ -7,6 +7,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -22,6 +23,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -34,102 +36,89 @@ concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
env:
CARGO_TERM_COLOR: always
jobs:
rust-checks:
name: rustfmt, clippy, test
rust-lint:
runs-on: ubuntu-latest
timeout-minutes: 10
defaults:
run:
working-directory: litellm-rust
env:
CARGO_TERM_COLOR: always
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Cache Cargo registry and target
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- run: cargo fmt --check
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }}
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-
${{ runner.os }}-cargo-${{ github.job }}-
- name: Check Rust formatting
run: cargo fmt --check
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
- name: Run Clippy
run: cargo clippy --workspace --all-targets --locked -- -D warnings
- run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- name: Run Clippy with Bedrock auth
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Clippy with all gateway features
run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Rust tests
run: cargo test --workspace --locked
- name: Run core tests with Bedrock auth
run: cargo test -p litellm-core --features bedrock-auth --locked
# Not --all-features: python-config links libpython, which this job does not install.
- name: Run gateway tests with the server feature
run: cargo test -p litellm-ai-gateway --features server --locked
release-wheel:
name: release wheel
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 20
permissions:
contents: read
env:
CARGO_TERM_COLOR: always
timeout-minutes: 30
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
- uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Build release wheel
run: uv build --wheel --out-dir dist
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
- name: Build panic contract wheel
run: >-
- run: cargo test --workspace --locked
working-directory: litellm-rust
- run: cargo test -p litellm-core --features bedrock-auth --locked
working-directory: litellm-rust
- run: cargo test -p litellm-ai-gateway --features server --locked
working-directory: litellm-rust
- run: uv build --wheel --out-dir dist
- run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: >-
uv build --wheel --out-dir panic-dist
--config-setting "maturin.build-args=--features panic-test,extension-module"
- name: Smoke-test native panic unwinding
run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl
- name: Verify stripped native extension
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
- name: Test native route wheel
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl

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

4
CLAUDE.md
View file

@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
@ -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
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
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
`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

4
basedpyright-code-budget.json
View file

@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1808
"limit": 1804
},
"reportRedeclaration": {
"limit": 8
@ -135,7 +135,7 @@
"limit": 21
},
"reportUnusedFunction": {
"limit": 138
"limit": 136
},
"reportUnusedImport": {
"limit": 542

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

2
litellm-rust/Cargo.lock generated
View file

@ -1415,6 +1415,8 @@ dependencies = [
"litellm-config",
"litellm-core",
"reqwest",
"rustls 0.23.42",
"rustls-native-certs",
"serde",
"serde_json",
"sha2 0.10.9",

2
litellm-rust/Cargo.toml
View file

@ -28,6 +28,8 @@ pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
sha2 = "0.10"

4
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -20,6 +20,10 @@ litellm-config.workspace = true
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
# rustls and its root store are direct dependencies so `io::tls` can build the
# one TLS config the outbound dials use; see that module for why it has to.
rustls.workspace = true
rustls-native-certs.workspace = true
# `sync` powers the bounded mpsc channel the realtime logger drains.
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
tokio-tungstenite.workspace = true

1
litellm-rust/crates/ai-gateway/src/io/mod.rs
View file

@ -3,3 +3,4 @@ pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;
pub(crate) mod tls;

33
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream(
.map_err(|err| Error::Auth(err.to_string()))?,
);
let (upstream, _response) = connect_async(request)
let (upstream, _response) = connect_upstream(request)
.await
.map_err(|err| Error::Network(err.to_string()))?;
Ok(upstream)
@ -284,6 +286,33 @@ mod tests {
serde_json::from_str(raw).expect("valid event json")
}
/// The realtime dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result = dial_upstream(
"gpt-realtime",
"sk-test",
Some(&format!("wss://127.0.0.1:{port}")),
)
.await;
assert!(matches!(result, Err(Error::Network(_))));
}
#[test]
fn resolve_api_key_prefers_param_then_blank_falls_through() {
assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");

35
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use crate::io::tls::connect_upstream;
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection {
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match error {
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -138,13 +140,13 @@ async fn dial_upstream(
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_async(request),
connect_upstream(request),
)
.await
.map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match error {
.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -324,6 +326,29 @@ mod tests {
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
/// The Responses dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result =
dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await;
assert!(matches!(result, Err(Error::Network(_))));
}
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("local address");

80
litellm-rust/crates/ai-gateway/src/io/tls.rs Normal file
View file

@ -0,0 +1,80 @@
//! Outbound WebSocket dials over a TLS config this crate builds once and owns.
//!
//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth`
//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that
//! `tokio-tungstenite` uses when handed no connector panics rather than guess
//! between them. Naming ring on a connector of our own settles that for these
//! dials without touching the process-wide default, and building the config
//! once keeps the platform trust store, which `tokio-tungstenite` would
//! otherwise re-read on every dial, off the dial path.
use std::io;
use std::sync::{Arc, OnceLock};
use rustls::{ClientConfig, RootCertStore};
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::error::TlsError;
use tokio_tungstenite::tungstenite::handshake::client::Response;
use tokio_tungstenite::{
Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
};
static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
fn build_config() -> Result<ClientConfig, Box<Error>> {
let native = rustls_native_certs::load_native_certs();
let roots = {
let mut store = RootCertStore::empty();
let (added, _ignored) = store.add_parsable_certificates(native.certs);
if added == 0 {
return Err(Box::new(Error::Io(io::Error::other(format!(
"no usable native root certificates: {:?}",
native.errors
)))));
}
store
};
ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
.with_safe_default_protocol_versions()
.map(|builder| builder.with_root_certificates(roots).with_no_client_auth())
.map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error))))
}
fn tls_config() -> Result<Arc<ClientConfig>, Box<Error>> {
if let Some(config) = TLS_CONFIG.get() {
return Ok(Arc::clone(config));
}
let built = Arc::new(build_config()?);
Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
}
pub(crate) async fn connect_upstream<R>(
request: R,
) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
where
R: IntoClientRequest + Unpin,
{
let request = request.into_client_request().map_err(Box::new)?;
let connector = match request.uri().scheme_str() {
Some("wss") => Some(Connector::Rustls(tls_config()?)),
_ => None,
};
connect_async_tls_with_config(request, None, false, connector)
.await
.map_err(Box::new)
}
#[cfg(test)]
mod tests {
use super::build_config;
#[test]
fn builds_a_usable_config_with_both_provider_features_enabled() {
let config = build_config().expect("a client config");
assert!(!config.crypto_provider().cipher_suites.is_empty());
}
}

12
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -265,6 +265,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
@ -288,6 +294,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
})
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,

48
litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs Normal file
View file

@ -0,0 +1,48 @@
//! Guards the wiring, not just the helper: a `wss://` dial through the public
//! API has to resolve its own crypto provider, in a test binary where nothing
//! has installed a process-wide one, and has to leave it uninstalled.
use std::collections::HashMap;
use std::time::Duration;
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
use tokio::net::TcpListener;
async fn dead_tls_server() -> u16 {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
port
}
#[tokio::test]
async fn dialing_wss_returns_an_error_instead_of_panicking() {
let port = dead_tls_server().await;
let result = ResponsesWebSocketConnection::connect_url(
&format!("wss://127.0.0.1:{port}/"),
&HashMap::new(),
Some(Duration::from_secs(10)),
)
.await;
assert!(
result.is_err(),
"a plain TCP server cannot finish a TLS handshake"
);
assert!(
rustls::crypto::CryptoProvider::get_default().is_none(),
"the dial settles its provider on its own connector, not process-wide"
);
}

52
litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -11,9 +11,13 @@ use litellm_ai_gateway::integrations::custom_logger::{
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
#[cfg(feature = "trace-parity")]
use litellm_core::observability::FunctionTrace;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "trace-parity")]
use tracing::instrument::WithSubscriber;
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -320,14 +324,17 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
let response = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -339,9 +346,10 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
..Default::default()
},
litellm_call_id: Some("ocr-call-1"),
})
.await
.expect("ocr request succeeds");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let response = call.await.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
assert_eq!(
@ -359,6 +367,16 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
error_kind: None,
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["success_callback"]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
@ -388,14 +406,17 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
});
let logger = Arc::new(RecordingOcrLogger::default());
let err = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -404,9 +425,10 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-2"),
})
.await
.expect_err("provider error propagates");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let err = call.await.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
@ -421,6 +443,16 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
error_kind: Some("HttpError".to_string()),
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["failure_callback"]
);
}
#[tokio::test]

28
litellm-rust/crates/python-bridge/src/function_trace.rs
View file

@ -1,3 +1,4 @@
use std::fmt::Display;
use std::future::Future;
use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
@ -6,17 +7,32 @@ use tracing::instrument::WithSubscriber;
#[derive(Serialize)]
pub(crate) struct TracedResponse<T> {
response: T,
#[serde(skip_serializing_if = "Option::is_none")]
response: Option<T>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
trace: Vec<FunctionTraceEvent>,
}
pub(crate) async fn capture<T, E>(
future: impl Future<Output = Result<T, E>>,
) -> Result<TracedResponse<T>, E> {
) -> Result<TracedResponse<T>, E>
where
E: Display,
{
let trace = FunctionTrace::default();
let response = future.with_subscriber(trace.dispatcher()).await?;
Ok(TracedResponse {
response,
trace: trace.events(),
let result = future.with_subscriber(trace.dispatcher()).await;
let events = trace.events();
Ok(match result {
Ok(response) => TracedResponse {
response: Some(response),
error: None,
trace: events,
},
Err(error) => TracedResponse {
response: None,
error: Some(error.to_string()),
trace: events,
},
})
}

9
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -486,10 +486,11 @@ asyncio.run(exercise())
let code = CString::new(
r#"
result = routes.echo("traced")
assert result == {
"response": "traced",
"trace": [{"function": "execute_echo", "depth": 0}],
}
assert result["response"] == "traced", result
assert [event["function"] for event in result["trace"]] == ["execute_echo"], result
failure = routes.echo("error")
assert failure["error"] == "invalid request: synthetic error", failure
assert [event["function"] for event in failure["trace"]] == ["execute_echo"], failure
"#,
)
.expect("Python source should not contain null bytes");

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

21
litellm/batches/batch_utils.py
View file

@ -25,6 +25,27 @@ class BatchCostUsageResult:
failed_requests: int
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"})
def batch_cost_is_final(batch: Batch) -> bool:
"""Whether this retrieve of the batch is the one to account its cost from.
A batch still in flight has nothing to price, and a "completed" batch can report
no output_file_id for a moment before the output populates; pricing either records
$0 under the batch's single spend row and pins it there. Final means a completed
batch whose output file has arrived or whose counts prove no line succeeded, or
any other terminal status (failed, cancelled, expired).
"""
if batch.status not in _TERMINAL_BATCH_STATUSES:
return False
if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None:
return True
request_counts: Final = batch.request_counts
return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0
async def calculate_batch_cost_and_usage(
file_content_dictionary: list[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],

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

@ -5,8 +5,11 @@ Handler for transforming /chat/completions api requests to litellm.responses req
import json
import os
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
from openai.types.chat import ChatCompletion
from openai.types.responses import Response
from openai.types.responses.custom_tool_param import CustomToolParam
from openai.types.responses.response_input_param import (
FunctionCallOutput,
@ -33,7 +36,7 @@ from litellm.responses.sse_output_recovery import (
record_output_item_chunk,
record_output_text_chunk,
)
from litellm.responses.utils import normalize_responses_api_stream_options
from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options
from litellm.types.llms.openai import (
REASONING_EFFORT,
ChatCompletionAnnotation,
@ -43,6 +46,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolParamFunctionChunk,
Reasoning,
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
@ -54,7 +58,7 @@ if TYPE_CHECKING:
)
from pydantic import BaseModel
from litellm import LiteLLMLoggingObj, ModelResponse
from litellm import LiteLLMLoggingObj
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.llms.openai import (
ALL_RESPONSES_API_TOOL_PARAMS,
@ -69,6 +73,28 @@ if TYPE_CHECKING:
from litellm.types.utils import Choices
_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage"))
_RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *ResponsesAPIResponse.model_fields)) - frozenset(
ChatCompletion.model_fields
)
def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
return MappingProxyType(
{
key: value
for key, value in (response_fields.items() if response_fields else ())
if value is not None and key not in _CHAT_COMPLETION_FIELDS and key not in _RESPONSES_API_ONLY_FIELDS
}
)
def _upstream_response_id(response_id: str | None) -> str | None:
if response_id is None:
return None
return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id)
class _ReasoningSummaryText(TypedDict):
type: str
text: str
@ -904,6 +930,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage),
)
model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
for key, value in _provider_metadata(raw_response.model_extra).items():
setattr(model_response, key, value)
# Preserve hidden params from the ResponsesAPIResponse, especially the headers
# which contain important provider information like x-request-id
raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {})
@ -1359,14 +1389,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
if event_type == "response.created":
# Initial response creation event
verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk)
created_response: Final = parsed_chunk.get("response")
return ModelResponseStream(
id=_upstream_response_id(created_response.get("id")) if created_response else None,
choices=[
StreamingChoices(
index=0,
delta=Delta(content=""),
finish_reason=None,
)
]
],
)
elif event_type == "response.output_item.added":
# New output item added
@ -1534,6 +1566,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
from litellm.responses.utils import ResponseAPILoggingUtils
usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage"))
provider_metadata: Final = _provider_metadata(response_data)
return ModelResponseStream(
choices=[
StreamingChoices(
@ -1546,6 +1579,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
],
usage=usage,
provider_specific_fields=dict(provider_metadata) or None, # mutable-ok: field is typed dict
)
else:
pass

10
litellm/constants.py
View file

@ -1901,6 +1901,16 @@ HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset(
}
)
PROVIDER_REQUEST_ID_HEADERS: Final[tuple[str, ...]] = (
"x-amzn-requestid",
"x-request-id",
"request-id",
"x-ms-request-id",
"apim-request-id",
"x-goog-request-id",
"cf-ray",
)
# Browser-facing security headers that a malicious or misconfigured upstream
# provider must not be able to set on the proxy's own response.
BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(

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:

9
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -21,13 +21,10 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
}
)
# The per-deployment Rust opt-in.
RUST_KWARG_KEY: Final = "rust"
# Keys `completion()` forwards from its own kwargs into `get_litellm_params`,
# which are otherwise invisible to it because that call site passes explicit
# named arguments rather than `**kwargs`.
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({RUST_KWARG_KEY})
FORWARDED_KWARGS_KEYS: Final = AWS_CREDENTIAL_KWARGS_KEYS
# Pre-define optional kwargs keys as frozenset for O(1) lookups
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
@ -58,10 +55,6 @@ OPTIONAL_KWARGS_KEYS: Final = (
"itpm",
"otpm",
"use_xai_oauth",
# The per-deployment Rust opt-in. `all_litellm_params` keeps it out
# of the provider body; this keeps it *in* litellm_params, which is
# where the chat completions handlers read it from.
RUST_KWARG_KEY,
}
)
| AWS_CREDENTIAL_KWARGS_KEYS

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

49
litellm/litellm_core_utils/litellm_logging.py
View file

@ -36,12 +36,13 @@ from litellm._logging import (
verbose_logger,
)
from litellm._uuid import uuid
from litellm.batches.batch_utils import _handle_completed_batch
from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final
from litellm.caching.caching import DualCache, InMemoryCache
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
PROVIDER_REQUEST_ID_HEADERS,
SENTRY_DENYLIST,
SENTRY_PII_DENYLIST,
)
@ -255,6 +256,30 @@ _in_memory_loggers: Final[list[CustomLogger]] = []
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys())
def _get_provider_request_id(original_exception: Exception) -> str | None:
try:
error_response: Final = getattr(original_exception, "response", None)
header_sources: Final = (
_get_response_headers(original_exception),
getattr(error_response, "headers", None),
getattr(original_exception, "litellm_response_headers", None),
)
return next(
(
str(value)
for expected_header_name in PROVIDER_REQUEST_ID_HEADERS
for headers in header_sources
if isinstance(headers, Mapping)
for header_name, value in headers.items()
if isinstance(header_name, str) and header_name.lower() == expected_header_name and value
),
None,
)
except Exception:
return None
### GLOBAL VARIABLES ###
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
@ -2899,13 +2924,6 @@ class Logging(LiteLLMLoggingBaseClass):
): # polling job will query these frequently, don't spam db logs
return
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
# check if file id is a unified file id
is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(result.id)
batch_cost: Final = kwargs.get("batch_cost", None)
batch_usage = kwargs.get("batch_usage", None)
batch_models = kwargs.get("batch_models", None)
@ -2913,9 +2931,7 @@ class Logging(LiteLLMLoggingBaseClass):
batch_failed_requests: Final = kwargs.get("batch_failed_requests", None)
has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models))
should_compute_batch_data: Final = (
not is_base64_unified_file_id or not has_explicit_batch_data and result.status == "completed"
)
should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result)
if has_explicit_batch_data:
result._hidden_params["response_cost"] = batch_cost
result._hidden_params["batch_models"] = batch_models
@ -3918,11 +3934,12 @@ class Logging(LiteLLMLoggingBaseClass):
LiteLLMResponsesTransformationHandler,
)
served_id: Final = _provider_response_id(result)
try:
return LiteLLMResponsesTransformationHandler().transform_response(
translated: Final = LiteLLMResponsesTransformationHandler().transform_response(
model=self.model,
raw_response=result,
model_response=litellm.ModelResponse(id=_provider_response_id(result)),
model_response=litellm.ModelResponse(id=served_id),
logging_obj=self,
request_data={},
messages=[],
@ -3930,6 +3947,8 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_params={},
encoding=litellm.encoding,
)
translated.id = served_id or translated.id
return translated
except Exception as e:
verbose_logger.debug(
"Responses API -> ModelResponse translation failed for "
@ -3937,7 +3956,7 @@ class Logging(LiteLLMLoggingBaseClass):
"usage-only ModelResponse to keep the spend_logs row.",
str(e),
)
model_response: Final = litellm.ModelResponse(id=_provider_response_id(result))
model_response: Final = litellm.ModelResponse(id=served_id)
model_response.model = self.model
usage: Final = getattr(result, "usage", None)
if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(usage):
@ -5670,6 +5689,7 @@ class StandardLoggingPayloadSetup:
rate_limit_category: Final = validate_rate_limit_category(getattr(original_exception, "category", None))
rate_limit_type: Final = validate_rate_limit_type(getattr(original_exception, "rate_limit_type", None))
budget_error: Final = original_exception if isinstance(original_exception, BudgetExceededError) else None
provider_request_id: Final = _get_provider_request_id(original_exception) if original_exception else None
return StandardLoggingPayloadErrorInformation(
error_code=error_status,
@ -5677,6 +5697,7 @@ class StandardLoggingPayloadSetup:
llm_provider=_llm_provider_in_exception,
traceback=_redact_string(traceback_info),
error_message=_redact_string(error_message),
error_provider_request_id=provider_request_id,
error_rate_limit_category=rate_limit_category,
error_rate_limit_type=rate_limit_type,
error_budget_entity_type=budget_error.entity_type if budget_error else None,

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

@ -26,6 +26,7 @@ from litellm.types.utils import (
PromptTokensDetailsWrapper,
ServiceTier,
Usage,
text_tokens_without_nested_reasoning,
)
from litellm.utils import get_model_info
@ -513,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.
@ -523,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)
"""
@ -550,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
@ -661,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):
@ -672,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,
@ -860,7 +860,7 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
)
or 0
)
text_tokens: Final = (
reported_text_tokens: Final = (
cast(
int | None,
getattr(usage.completion_tokens_details, "text_tokens", None),
@ -882,6 +882,12 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
or 0
)
video_tokens: Final = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0))
text_tokens: Final = text_tokens_without_nested_reasoning(
completion_tokens=usage.completion_tokens,
text_tokens=reported_text_tokens,
reasoning_tokens=reasoning_tokens,
other_modality_tokens=audio_tokens + image_tokens + video_tokens,
)
return CompletionTokensDetailsResult(
audio_tokens=audio_tokens,
@ -1409,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,

202
litellm/litellm_core_utils/prompt_templates/image_handling.py
View file

@ -2,7 +2,11 @@
Helper functions to handle images passed in messages
"""
import asyncio
import base64
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from httpx import Response
@ -11,9 +15,11 @@ import litellm
from litellm import verbose_logger
from litellm.caching.caching import InMemoryCache
from litellm.constants import MAX_IMAGE_URL_DOWNLOAD_SIZE_MB
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get, safe_get
from litellm.types.llms.openai import AllMessageValues
MAX_IMGS_IN_MEMORY: Final = 10
MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20
in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY)
@ -72,6 +78,14 @@ def _process_image_response(response: Response, url: str) -> str:
return result
def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict)
return litellm.ImageFetchError(
"Error: Unable to fetch image from URL. The proxy could not resolve this host or its URL policy rejected it; "
f"an admin can check the proxy log and `user_url_allowed_hosts` in general_settings. url={url}"
)
async def async_convert_url_to_base64(url: str) -> str:
if url.startswith("data:") and ";base64," in url:
return url
@ -93,6 +107,8 @@ async def async_convert_url_to_base64(url: str) -> str:
return _process_image_response(response, url)
except litellm.ImageFetchError:
raise
except SSRFError as e:
raise _rejected_image_fetch(url, e) from e
except Exception:
pass
raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
@ -119,8 +135,192 @@ def convert_url_to_base64(url: str) -> str:
return _process_image_response(response, url)
except litellm.ImageFetchError:
raise
except SSRFError as e:
raise _rejected_image_fetch(url, e) from e
except Exception as e:
verbose_logger.exception(e)
raise litellm.ImageFetchError(
f"Error: Unable to fetch image from URL after 3 attempts. url={url}",
)
_REMOTE_URL_PREFIXES: Final = ("http://", "https://")
@dataclass(frozen=True, slots=True)
class _RemoteImage:
part: Mapping[str, object]
image_url: Mapping[str, object] | None
url: str
@dataclass(frozen=True, slots=True)
class _RemoteFile:
part: Mapping[str, object]
file: Mapping[str, object]
url: str
def _as_mapping(value: object) -> Mapping[str, object] | None:
return value if isinstance(value, Mapping) else None # pyright: ignore[reportUnknownVariableType] # fields are parsed one by one
def _remote_url(candidate: object) -> str | None:
return candidate if isinstance(candidate, str) and candidate.startswith(_REMOTE_URL_PREFIXES) else None
_ANTHROPIC_MEDIA_BLOCK_TYPES: Final = frozenset({"document", "image"})
@dataclass(frozen=True, slots=True)
class _RemoteSource:
part: Mapping[str, object]
source: Mapping[str, object]
url: str
@dataclass(frozen=True, slots=True)
class RemoteMedia:
url: str
fields: Mapping[str, object]
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
def inline_every_remote_url(_media: RemoteMedia) -> bool:
return True
def _parse_remote_image(fields: Mapping[str, object]) -> _RemoteImage | None:
if fields.get("type") != "image_url":
return None
image_url: Final = fields.get("image_url")
image_url_fields: Final = _as_mapping(image_url)
url: Final = _remote_url(image_url_fields.get("url") if image_url_fields is not None else image_url)
return _RemoteImage(fields, image_url_fields, url) if url is not None else None
def _parse_remote_file(fields: Mapping[str, object]) -> _RemoteFile | None:
file: Final = _as_mapping(fields.get("file")) if fields.get("type") == "file" else None
url: Final = _remote_url(file.get("file_id")) if file is not None else None
return _RemoteFile(fields, file, url) if file is not None and url is not None else None
def _parse_remote_source(fields: Mapping[str, object]) -> _RemoteSource | None:
source: Final = _as_mapping(fields.get("source")) if fields.get("type") in _ANTHROPIC_MEDIA_BLOCK_TYPES else None
url: Final = _remote_url(source.get("url")) if source is not None and source.get("type") == "url" else None
return _RemoteSource(fields, source, url) if source is not None and url is not None else None
def _parse_remote_part(part: object) -> _RemoteImage | _RemoteFile | _RemoteSource | None:
fields: Final = _as_mapping(part)
if fields is None:
return None
return _parse_remote_image(fields) or _parse_remote_file(fields) or _parse_remote_source(fields)
def _remote_media(remote: _RemoteImage | _RemoteFile | _RemoteSource) -> RemoteMedia:
match remote:
case _RemoteImage(_, image_url, url):
return RemoteMedia(url, image_url if image_url is not None else _NO_FIELDS)
case _RemoteFile(_, file, url):
return RemoteMedia(url, file)
case _RemoteSource(_, source, url):
return RemoteMedia(url, source)
_PDF_FORMAT: Final = MappingProxyType({"format": "application/pdf"})
def _inferred_format(file: Mapping[str, object], url: str) -> Mapping[str, str]:
return _PDF_FORMAT if "format" not in file and url.lower().endswith(".pdf") else MappingProxyType({})
def _inlined_image_url(image_url: Mapping[str, object] | None, data_url: str) -> Mapping[str, object] | str:
return {**image_url, "url": data_url} if image_url is not None else data_url # mutable-ok: json-serialized part
def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mapping[str, object]:
kept: Final = {k: v for k, v in file.items() if k != "file_id"} # mutable-ok: json-serialized message part
return {**kept, **_inferred_format(file, url), "file_data": data_url} # mutable-ok: json-serialized part
def _base64_source(url: str, data_url: str) -> Mapping[str, str]:
fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1)
media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type
return {"type": "base64", "media_type": media_type, "data": data} # mutable-ok: json-serialized message part
def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]:
match remote:
case _RemoteImage(part, image_url, _):
return {**part, "image_url": _inlined_image_url(image_url, data_url)} # mutable-ok: json-serialized part
case _RemoteFile(part, file, url):
return {**part, "file": _inlined_file(file, url, data_url)} # mutable-ok: json-serialized message part
case _RemoteSource(part, _, url):
return {**part, "source": _base64_source(url, data_url)} # mutable-ok: json-serialized message part
def _content_parts(message: Mapping[str, object]) -> tuple[object, ...]:
content: Final = message.get("content")
return tuple(content) if isinstance(content, list) else () # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType] # parts are parsed one by one
def _inline_part(part: object, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]) -> object:
remote: Final = _parse_remote_part(part)
if remote is None or not should_inline(_remote_media(remote)):
return part
data_url: Final = data_urls.get(remote.url)
return _inline(remote, data_url) if data_url is not None else part
def _inline_message(
message: AllMessageValues, data_urls: Mapping[str, str], should_inline: Callable[[RemoteMedia], bool]
) -> AllMessageValues:
parts: Final = _content_parts(message)
if not parts:
return message
inlined_parts: Final = [ # mutable-ok: content must stay a list for the transforms' isinstance checks
_inline_part(part, data_urls, should_inline) for part in parts
]
inlined_message: Final = {**message, "content": inlined_parts} # mutable-ok: json-serialized message
return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined
async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str:
async with in_flight:
return await async_convert_url_to_base64(url)
async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]:
in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls)
try:
return tuple(await asyncio.gather(*fetches))
except BaseException:
for fetch in fetches:
fetch.cancel()
await asyncio.gather(*fetches, return_exceptions=True)
raise
async def async_inline_remote_media(
messages: list[AllMessageValues], # mutable-ok: every transform_request takes list[AllMessageValues]
should_inline: Callable[[RemoteMedia], bool] = inline_every_remote_url,
) -> list[AllMessageValues]: # mutable-ok: every transform_request takes list[AllMessageValues]
remote_urls: Final = tuple(
dict.fromkeys(
remote.url
for message in messages
for part in _content_parts(message)
if (remote := _parse_remote_part(part)) is not None and should_inline(_remote_media(remote))
)
)
if not remote_urls:
return messages
data_urls: Final = await _fetch_data_urls(remote_urls)
inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
return [ # mutable-ok: transform_request takes a list
_inline_message(message, inlined, should_inline) for message in messages
]

3
litellm/litellm_core_utils/url_utils.py
View file

@ -19,6 +19,7 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config):
check but still resolve DNS and still rewrite HTTP to the resolved IP.
"""
import asyncio
import socket
from ipaddress import ip_address, ip_network
from typing import Any, Final, Protocol
@ -471,7 +472,7 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response
kwargs.pop("follow_redirects", None)
headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})}
for _ in range(_MAX_REDIRECTS):
validated_url, original_host = validate_url(url)
validated_url, original_host = await asyncio.to_thread(validate_url, url)
response = await fetcher.get(
validated_url,
headers={**headers_view["headers"], "Host": original_host},

49
litellm/llms/azure_ai/chat/transformation.py
View file

@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base
from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error
from litellm.llms.openai.openai import OpenAIConfig
from litellm.llms.xai.chat.transformation import XAIChatConfig
@ -42,12 +43,37 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = (
)
class AzureAIGPT5Config(OpenAIGPT5Config):
@classmethod
def _model_map_lookup_name(cls, model: str) -> str:
"""Normalise a Foundry routing name to its cost-map key, when the map has one.
A Foundry deployment and its OpenAI-hosted namesake are different products with
different capabilities, so ``azure_ai/<model>`` is the entry to read whenever the map
carries it. Most gpt-5-family names have no ``azure_ai/`` row, though, and prefixing
those anyway costs them every flag: ``get_llm_provider`` re-resolves an ``azure_ai/``
name to the azure provider when a global AZURE_AI_API_BASE points at an
openai.azure.com host, ``azure/<model>`` is not a key either, so the lookup lands
nowhere and every effort answer degrades to False. A missing key defers to the base
resolver instead.
"""
prefixed: Final = model if model.startswith("azure_ai/") else f"azure_ai/{model}"
return prefixed if prefixed in litellm.model_cost else super()._model_map_lookup_name(model)
azureAIGPT5Config: Final = AzureAIGPT5Config()
class AzureAIStudioConfig(OpenAIConfig):
def get_supported_openai_params(self, model: str) -> list:
model_supports_tool_choice = True # azure ai supports this by default
if not supports_tool_choice(model=f"azure_ai/{model}"):
model_supports_tool_choice = False
supported_params = super().get_supported_openai_params(model)
supported_params = (
azureAIGPT5Config.get_supported_openai_params(model)
if azureAIGPT5Config.is_model_gpt_5_model(model)
else super().get_supported_openai_params(model)
)
if not model_supports_tool_choice:
filtered_supported_params: Final = []
for param in supported_params:
@ -61,6 +87,27 @@ class AzureAIStudioConfig(OpenAIConfig):
return supported_params
def map_openai_params(
self,
non_default_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature
optional_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature
model: str,
drop_params: bool,
) -> dict[str, object]: # mutable-ok: OpenAIConfig.map_openai_params signature
if not azureAIGPT5Config.is_model_gpt_5_model(model):
return super().map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=drop_params,
)
return azureAIGPT5Config.map_openai_params(
non_default_params=non_default_params,
optional_params=optional_params,
model=model,
drop_params=drop_params,
)
def _supports_stop_reason(self, model: str) -> bool:
"""
Check if the model supports stop tokens.

4
litellm/llms/base_llm/chat/transformation.py
View file

@ -411,6 +411,10 @@ class BaseConfig(ABC):
def has_custom_stream_wrapper(self) -> bool:
return False
@property
def uses_async_transform_request(self) -> bool:
return False
@property
def supports_stream_param_in_request_body(self) -> bool:
"""

19
litellm/llms/base_llm/image_edit/transformation.py
View file

@ -1,5 +1,6 @@
import types
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any
import httpx
@ -102,6 +103,24 @@ class BaseImageEditConfig(ABC):
) -> tuple[dict, RequestFiles]:
pass
async def async_transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[dict, RequestFiles]:
return self.transform_image_edit_request(
model=model,
prompt=prompt,
image=image,
image_edit_optional_request_params=dict(image_edit_optional_request_params),
litellm_params=litellm_params,
headers=dict(headers),
)
def finalize_image_edit_request_data(self, data: dict, resolved_request_url: str) -> dict:
"""
Last pass on the request dict after ``transform_image_edit_request``, using the

61
litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -8,7 +8,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_anthropic_image_obj,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import (
async_convert_url_to_base64,
async_inline_remote_media,
convert_url_to_base64,
)
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@ -172,6 +172,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
return _anthropic_request
@property
def uses_async_transform_request(self) -> bool:
return True
async def async_transform_request(
self,
model: str,
@ -180,26 +184,14 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
litellm_params: dict,
headers: dict,
) -> dict:
_anthropic_request: Final = self._build_bedrock_anthropic_request_base(
return self.transform_request(
model=model,
messages=messages,
messages=await async_inline_remote_media(messages),
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
await self._async_convert_document_url_sources_to_base64(_anthropic_request)
beta_list: Final = self._compute_bedrock_invoke_beta_headers(
model=model,
messages=messages,
optional_params=optional_params,
headers=headers,
)
if beta_list:
_anthropic_request["anthropic_beta"] = beta_list
return _anthropic_request
def _build_bedrock_anthropic_request_base(
self,
model: str,
@ -321,45 +313,6 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
"data": image_chunk["data"],
}
async def _async_convert_document_url_sources_to_base64(self, anthropic_request: dict) -> None:
"""
Async version of document URL conversion for async completion paths.
"""
messages: Final = anthropic_request.get("messages")
if not isinstance(messages, list):
return
for message in messages:
if not isinstance(message, dict):
continue
content = message.get("content")
if not isinstance(content, list):
continue
for block in content:
if not isinstance(block, dict) or block.get("type") != "document":
continue
source = block.get("source")
if not isinstance(source, dict) or source.get("type") != "url":
continue
source_url = source.get("url")
if not isinstance(source_url, str):
continue
inferred_format: str | None = None
if source_url.lower().endswith(".pdf"):
inferred_format = "application/pdf"
base64_url = await async_convert_url_to_base64(url=source_url)
image_chunk = convert_to_anthropic_image_obj(
openai_image_url=base64_url,
format=inferred_format,
)
block["source"] = {
"type": "base64",
"media_type": image_chunk["media_type"],
"data": image_chunk["data"],
}
def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict:
"""
Convert tool search entries to the format supported by the Bedrock Invoke API.

15
litellm/llms/bedrock/chat/mantle/transformation.py
View file

@ -10,6 +10,7 @@ at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
from collections.abc import AsyncIterator, Iterator
from typing import TYPE_CHECKING, Any, Final
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeConfig,
)
@ -110,21 +111,13 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
litellm_params: dict,
headers: dict,
) -> dict:
model_id: Final = model.replace("mantle/", "", 1)
request: Final = self._build_bedrock_anthropic_request_base(
model=model_id,
messages=messages,
return self.transform_request(
model=model,
messages=await async_inline_remote_media(messages),
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
await self._async_convert_document_url_sources_to_base64(request)
return self._restore_mantle_body_fields(
request=request,
model_id=model_id,
optional_params=optional_params,
)
@staticmethod
def _restore_mantle_body_fields(request: dict, model_id: str, optional_params: dict) -> dict:

102
litellm/llms/black_forest_labs/image_edit/transformation.py
View file

@ -9,6 +9,7 @@ API Reference: https://docs.bfl.ai/
import base64
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
import httpx
@ -16,7 +17,7 @@ from httpx._types import RequestFiles
import litellm
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
from litellm.litellm_core_utils.url_utils import safe_get
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from litellm.secret_managers.main import get_secret_str
from litellm.types.images.main import ImageEditOptionalRequestParams
@ -37,6 +38,22 @@ else:
LiteLLMLoggingObj = Any
_BFL_REQUEST_PARAMS: Final = (
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"aspect_ratio",
"steps",
"guidance",
"grow_mask",
"top",
"bottom",
"left",
"right",
)
class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"""
Configuration for Black Forest Labs image editing.
@ -85,34 +102,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
BFL-specific params are passed through directly.
"""
optional_params: Final[dict[str, object]] = {}
# Pass through BFL-specific params
bfl_params: Final = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
# Kontext-specific
"aspect_ratio",
# Fill/Inpaint-specific
"steps",
"guidance",
"grow_mask",
# Expand-specific
"top",
"bottom",
"left",
"right",
]
# Convert TypedDict to regular dict for access
params_dict: Final = dict(image_edit_optional_params)
for param in bfl_params:
if param in params_dict:
value = params_dict[param]
if value is not None:
optional_params[param] = value
params: Final[Mapping[str, object]] = image_edit_optional_params
for param in _BFL_REQUEST_PARAMS:
if (value := params.get(param)) is not None:
optional_params[param] = value
# Set default output format
if "output_format" not in optional_params:
@ -251,23 +244,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
"input_image": b64_image,
}
# Add optional params (only BFL-recognized parameters)
bfl_request_params: Final = [
"seed",
"output_format",
"safety_tolerance",
"prompt_upsampling",
"aspect_ratio",
"steps",
"guidance",
"grow_mask",
"top",
"bottom",
"left",
"right",
]
for key, value in image_edit_optional_request_params.items():
if key in bfl_request_params and value is not None:
if key in _BFL_REQUEST_PARAMS and value is not None:
request_body[key] = value
# Handle mask if provided (for inpainting)
@ -277,7 +255,39 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
request_body["mask"] = base64.b64encode(mask_bytes).decode("utf-8")
# BFL uses JSON, not multipart - return empty files
return request_body, []
return request_body, ()
async def async_transform_image_edit_request(
self,
model: str,
prompt: str | None,
image: FileTypes | None,
image_edit_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: Mapping[str, str],
) -> tuple[dict, RequestFiles]:
downloaded_image: Final = await self._fetch_remote_image(image)
downloaded_mask: Final = await self._fetch_remote_image(image_edit_optional_request_params.get("mask"))
return self.transform_image_edit_request(
model=model,
prompt=prompt,
image=image if downloaded_image is None else downloaded_image,
image_edit_optional_request_params=(
dict(image_edit_optional_request_params)
if downloaded_mask is None
else {**image_edit_optional_request_params, "mask": downloaded_mask}
),
litellm_params=litellm_params,
headers=dict(headers),
)
async def _fetch_remote_image(self, image: object) -> bytes | None:
candidate: Final = image[0] if isinstance(image, list) and image else image
if not isinstance(candidate, str) or not candidate.startswith(("http://", "https://")):
return None
response: Final = await async_safe_get(litellm.module_level_aclient, candidate, timeout=60.0)
response.raise_for_status()
return response.content
def transform_image_edit_response(
self,

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

@ -116,6 +116,7 @@ from litellm.types.llms.anthropic_skills import (
Skill,
)
from litellm.types.llms.openai import (
AllMessageValues,
CreateBatchRequest,
CreateFileRequest,
FileContentRequest,
@ -163,13 +164,10 @@ from litellm.utils import (
def _rust_responses_websocket_enabled(
custom_llm_provider: str | None,
litellm_params: GenericLiteLLMParams,
) -> bool:
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return custom_llm_provider == "openai" and rust_enabled(request_override=request_override)
return custom_llm_provider == "openai" and rust_enabled()
from .http_handler import get_shared_realtime_ssl_context
@ -488,7 +486,7 @@ class BaseLLMHTTPHandler:
def completion(
self,
model: str,
messages: list,
messages: list[AllMessageValues],
api_base: str | None,
custom_llm_provider: str,
model_response: ModelResponse,
@ -507,7 +505,7 @@ class BaseLLMHTTPHandler:
shared_session: Optional["ClientSession"] = None,
):
json_mode: Final[bool] = optional_params.pop("json_mode", False)
extra_body: Final[dict | None] = optional_params.pop("extra_body", None)
extra_body: Final[Mapping[str, object] | None] = optional_params.pop("extra_body", None)
provider_config = provider_config or ProviderConfigManager.get_provider_chat_config(
model=model, provider=litellm.LlmProviders(custom_llm_provider)
@ -522,14 +520,17 @@ class BaseLLMHTTPHandler:
)
# get config from model, custom llm provider
headers = provider_config.validate_environment(
api_key=api_key,
headers=headers or {},
model=model,
messages=messages,
optional_params=optional_params,
api_base=api_base,
litellm_params=litellm_params,
request_headers: Final = cast( # cast-ok: validate_environment is declared as a bare dict
"dict[str, object]",
provider_config.validate_environment(
api_key=api_key,
headers=headers or {},
model=model,
messages=messages,
optional_params=optional_params,
api_base=api_base,
litellm_params=litellm_params,
),
)
api_base = provider_config.get_complete_url(
@ -541,93 +542,117 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
)
data: dict[str, object] = provider_config.transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
if extra_body is not None:
data = {**data, **extra_body}
headers, signed_json_body = provider_config.sign_request(
headers=headers,
optional_params={
**optional_params,
**_aws_signing_overrides(optional_params, litellm_params),
},
request_data=data,
api_base=api_base,
api_key=api_key,
stream=stream,
fake_stream=fake_stream,
model=model,
)
## LOGGING
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": headers,
},
)
# Check if stream was converted for WebSearch interception
# This is set by the async_pre_request_hook in WebSearchInterceptionLogger
if litellm_params.get("_websearch_interception_converted_stream", False):
logging_obj.model_call_details["websearch_interception_converted_stream"] = True
if acompletion is True:
if stream is True:
data = self._add_stream_param_to_request_body(
data=data,
provider_config=provider_config,
def sign_and_log(
transformed: dict[str, object], # mutable-ok: async_completion takes dict
) -> tuple[dict[str, object], dict[str, object], bytes | None]: # mutable-ok: async_completion takes dict
data: Final = {**transformed, **extra_body} if extra_body is not None else transformed
signed: Final = cast( # cast-ok: sign_request is declared as a bare dict
"tuple[dict[str, object], bytes | None]",
provider_config.sign_request(
headers=request_headers,
optional_params={
**optional_params,
**_aws_signing_overrides(optional_params, litellm_params),
},
request_data=data,
api_base=api_base,
api_key=api_key,
stream=stream,
fake_stream=fake_stream,
)
model=model,
),
)
logging_obj.pre_call(
input=messages,
api_key=api_key,
additional_args={
"complete_input_dict": data,
"api_base": api_base,
"headers": signed[0],
},
)
if litellm_params.get("_websearch_interception_converted_stream", False):
logging_obj.model_call_details["websearch_interception_converted_stream"] = True
return data, signed[0], signed[1]
def dispatch_async(
data: dict[str, object], # mutable-ok: async_completion takes dict
signed_headers: dict[str, object], # mutable-ok: async_completion takes dict
signed_json_body: bytes | None,
):
async_client: Final = client if isinstance(client, AsyncHTTPHandler) else None
if stream is True:
return self.acompletion_stream_function(
model=model,
messages=messages,
api_base=api_base,
headers=headers,
headers=signed_headers,
custom_llm_provider=custom_llm_provider,
provider_config=provider_config,
timeout=timeout,
logging_obj=logging_obj,
data=data,
data=self._add_stream_param_to_request_body(
data=data,
provider_config=provider_config,
fake_stream=fake_stream,
),
fake_stream=fake_stream,
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
client=async_client,
litellm_params=litellm_params,
json_mode=json_mode,
optional_params=optional_params,
signed_json_body=signed_json_body,
)
return self.async_completion(
custom_llm_provider=custom_llm_provider,
provider_config=provider_config,
api_base=api_base,
headers=signed_headers,
data=data,
timeout=timeout,
model=model,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
client=async_client,
json_mode=json_mode,
signed_json_body=signed_json_body,
shared_session=shared_session,
)
else:
return self.async_completion(
custom_llm_provider=custom_llm_provider,
provider_config=provider_config,
api_base=api_base,
headers=headers,
data=data,
timeout=timeout,
model=model,
model_response=model_response,
logging_obj=logging_obj,
api_key=api_key,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
encoding=encoding,
client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
json_mode=json_mode,
signed_json_body=signed_json_body,
shared_session=shared_session,
if acompletion is True and provider_config.uses_async_transform_request:
async def transform_then_dispatch():
transformed: Final = cast( # cast-ok: async_transform_request is declared as a bare dict
"dict[str, object]",
await provider_config.async_transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=request_headers,
),
)
return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
return transform_then_dispatch()
data, signed_headers, signed_json_body = sign_and_log(
provider_config.transform_request(
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=request_headers,
)
)
if acompletion is True:
return dispatch_async(data, signed_headers, signed_json_body)
if stream is True:
data = self._add_stream_param_to_request_body(
@ -641,7 +666,7 @@ class BaseLLMHTTPHandler:
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
api_base=api_base,
headers=headers,
headers=signed_headers,
data=data,
signed_json_body=signed_json_body,
messages=messages,
@ -651,7 +676,7 @@ class BaseLLMHTTPHandler:
completion_stream, headers = self.make_sync_call(
provider_config=provider_config,
api_base=api_base,
headers=headers,
headers=signed_headers,
data=data,
signed_json_body=signed_json_body,
original_data=data,
@ -684,7 +709,7 @@ class BaseLLMHTTPHandler:
sync_httpx_client=sync_httpx_client,
provider_config=provider_config,
api_base=api_base,
headers=headers,
headers=signed_headers,
data=data,
signed_json_body=signed_json_body,
timeout=timeout,
@ -2403,9 +2428,7 @@ class BaseLLMHTTPHandler:
return None
from litellm.rust_bridge.configuration import rust_enabled
raw_request_override: Final = litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
if not rust_enabled(request_override=request_override):
if not rust_enabled():
return None
if has_agentic_hook:
return None
@ -6514,7 +6537,7 @@ class BaseLLMHTTPHandler:
@asynccontextmanager
async def _backend_connection():
if _rust_responses_websocket_enabled(custom_llm_provider, litellm_params):
if _rust_responses_websocket_enabled(custom_llm_provider):
from litellm.rust_bridge import responses_websocket as rust_responses_websocket
rust_backend: Final = await rust_responses_websocket.connect(
@ -6759,7 +6782,7 @@ class BaseLLMHTTPHandler:
litellm_params=dict(litellm_params),
)
data, files = image_edit_provider_config.transform_image_edit_request(
data, files = await image_edit_provider_config.async_transform_image_edit_request(
model=model,
image=image,
prompt=prompt,

14
litellm/llms/gemini/chat/transformation.py
View file

@ -12,7 +12,7 @@ from litellm.types.llms.openai import AllMessageValues, ChatCompletionFileObject
from litellm.types.llms.vertex_ai import ContentType, PartType
from litellm.utils import supports_reasoning
from ...vertex_ai.gemini.transformation import _gemini_convert_messages_with_history
from ...vertex_ai.gemini.transformation import GEMINI_FILES_API_URI_PREFIX, _gemini_convert_messages_with_history
from ...vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
@ -127,7 +127,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
if element.get("type") == "image_url":
img_element = cast(ChatCompletionImageObject, element) # cast-ok: runtime type tag checked
_image_url, format, detail = _image_url_fields(img_element)
if _image_url and "https://" in _image_url:
if (
_image_url
and "https://" in _image_url
and not _image_url.startswith(GEMINI_FILES_API_URI_PREFIX)
):
image_obj = convert_to_anthropic_image_obj(_image_url, format=format)
converted_image_url = convert_generic_image_chunk_to_openai_image_obj(image_obj)
if detail is not None:
@ -147,7 +151,11 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
llm_provider="gemini",
)
file_id = _file_field.get("file_id")
if file_id and ("http://" in file_id or "https://" in file_id):
if (
file_id
and ("http://" in file_id or "https://" in file_id)
and not file_id.startswith(GEMINI_FILES_API_URI_PREFIX)
):
# Convert HTTP/HTTPS file URL to base64 data
try:
base64_data = convert_url_to_base64(file_id)

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

210
litellm/llms/mistral/audio_speech/transformation.py Normal file
View file

@ -0,0 +1,210 @@
"""
Support for Mistral Voxtral text-to-speech via ``/v1/audio/speech``.
API reference: https://docs.mistral.ai/api/#tag/audio/operation/audio_speech_v1_audio_speech_post
"""
import base64
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.text_to_speech.transformation import (
BaseTextToSpeechConfig,
TextToSpeechRequestData,
)
from litellm.secret_managers.main import get_secret_str
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import HttpxBinaryResponseContent
class MistralTextToSpeechException(BaseLLMException):
pass
class MistralTextToSpeechConfig(BaseTextToSpeechConfig):
TTS_BASE_URL: Final[str] = "https://api.mistral.ai/v1"
AUDIO_CONTENT_TYPES: Final[MappingProxyType[str, str]] = MappingProxyType(
{
"mp3": "audio/mpeg",
"wav": "audio/wav",
"pcm": "audio/pcm",
"flac": "audio/flac",
"opus": "audio/ogg",
}
)
DROPPED_RESPONSE_HEADERS: Final[frozenset[str]] = frozenset(
{"content-encoding", "transfer-encoding", "content-length", "content-type"}
)
OPENAI_VOICE_ALIASES: Final[MappingProxyType[str, str]] = MappingProxyType(
{
"alloy": "en_paul_neutral",
"echo": "gb_oliver_neutral",
"fable": "en_paul_cheerful",
"onyx": "en_paul_confident",
"nova": "gb_jane_sarcasm",
"shimmer": "gb_jane_sarcasm",
}
)
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class contract returns a plain list
return ["voice", "response_format"] # mutable-ok: base class contract returns a plain list
def _map_openai_voice(self, voice_id: str) -> str:
return self.OPENAI_VOICE_ALIASES.get(voice_id.lower(), voice_id)
def _resolve_voice_id(self, voice: object) -> str | None:
if isinstance(voice, str) and voice.strip():
return self._map_openai_voice(voice.strip())
if isinstance(voice, Mapping):
candidates: Final = (voice.get(key) for key in ("voice_id", "id", "name"))
resolved: Final = next(
(candidate.strip() for candidate in candidates if isinstance(candidate, str) and candidate.strip()),
None,
)
return self._map_openai_voice(resolved) if resolved else None
return None
def map_openai_params(
self,
model: str,
optional_params: Mapping[str, object],
voice: object = None,
drop_params: bool = False,
kwargs: Mapping[str, object] | None = None,
) -> tuple[str | None, dict]: # mutable-ok: base class contract returns a plain dict
response_format: Final = optional_params.get("response_format")
ref_audio: Final = kwargs.get("ref_audio") if kwargs else None
voice_id_kwarg: Final = kwargs.get("voice_id") if kwargs else None
mapped_voice: Final = self._resolve_voice_id(voice) or self._resolve_voice_id(voice_id_kwarg)
mapped_params: Final = { # mutable-ok: base class contract returns a plain dict
key: value
for key, value in (("response_format", response_format), ("ref_audio", ref_audio))
if isinstance(value, str)
}
return mapped_voice, mapped_params
def validate_environment(
self,
headers: Mapping[str, str],
model: str,
api_key: str | None = None,
api_base: str | None = None,
) -> dict: # mutable-ok: base class contract returns a plain dict
resolved_key: Final = api_key or get_secret_str("MISTRAL_API_KEY")
if resolved_key is None:
raise MistralTextToSpeechException(
status_code=401,
message="Mistral API key is required. Set MISTRAL_API_KEY or pass api_key.",
)
return { # mutable-ok: base class contract returns a plain dict
**headers,
"Authorization": f"Bearer {resolved_key}",
"Content-Type": "application/json",
}
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: Mapping[str, object],
) -> str:
configured_base: Final = (api_base or self.TTS_BASE_URL).rstrip("/")
versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1"
return f"{versioned_base}/audio/speech"
def transform_text_to_speech_request(
self,
model: str,
input: str,
voice: str | None,
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
headers: Mapping[str, str],
) -> TextToSpeechRequestData:
response_format: Final = optional_params.get("response_format")
ref_audio: Final = optional_params.get("ref_audio")
request_data: Final[TextToSpeechRequestData] = {
"dict_body": {
"model": model,
"input": input,
**({"voice_id": voice} if voice else {}),
**({"response_format": response_format} if isinstance(response_format, str) else {}),
**({"ref_audio": ref_audio} if isinstance(ref_audio, str) else {}),
},
"headers": {"Content-Type": "application/json"},
}
return request_data
def _requested_content_type(self, request: httpx.Request) -> str:
request_body: Final = json.loads(request.content or b"{}")
requested_format: Final = request_body.get("response_format")
if not isinstance(requested_format, str):
return "audio/mpeg"
return self.AUDIO_CONTENT_TYPES.get(requested_format, "audio/mpeg")
def transform_text_to_speech_response(
self,
model: str,
raw_response: httpx.Response,
logging_obj: "LiteLLMLoggingObj",
) -> "HttpxBinaryResponseContent":
from litellm.types.llms.openai import HttpxBinaryResponseContent
try:
response_json: Final = raw_response.json()
except (json.JSONDecodeError, ValueError):
raise MistralTextToSpeechException(
status_code=raw_response.status_code,
message=f"Non-JSON response from Mistral speech API: {raw_response.text[:500]}",
headers=raw_response.headers,
)
audio_b64: Final = response_json.get("audio_data")
if not isinstance(audio_b64, str) or not audio_b64:
raise MistralTextToSpeechException(
status_code=500,
message=f"No audio_data in Mistral speech response. Response keys: {tuple(response_json.keys())}",
headers=raw_response.headers,
)
try:
audio_bytes: Final = base64.b64decode(audio_b64, validate=True)
except ValueError:
raise MistralTextToSpeechException(
status_code=500,
message="Invalid base64 audio_data in Mistral speech response.",
headers=raw_response.headers,
)
retained_headers: Final = tuple(
(key, value)
for key, value in raw_response.headers.items()
if key.lower() not in self.DROPPED_RESPONSE_HEADERS
)
response_headers: Final = retained_headers + (
("content-length", str(len(audio_bytes))),
("content-type", self._requested_content_type(raw_response.request)),
)
binary_response: Final = httpx.Response(
status_code=200,
headers=response_headers,
content=audio_bytes,
request=raw_response.request,
)
return HttpxBinaryResponseContent(binary_response)
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict | httpx.Headers, # mutable-ok: BaseLLMException takes a plain dict or httpx.Headers
) -> BaseLLMException:
return MistralTextToSpeechException(
message=error_message,
status_code=status_code,
headers=headers,
)

52
litellm/llms/oci/chat/cohere.py
View file

@ -290,19 +290,14 @@ def handle_cohere_stream_chunk(
) -> ModelResponseStream:
"""Parse a single Cohere SSE chunk into a LiteLLM ModelResponseStream.
``prior_tool_calls_emitted`` lets the caller signal whether tool calls
were already emitted in earlier chunks of the same stream. When set, the
terminal consolidation chunk's tool calls are suppressed (they would
duplicate prior deltas); otherwise they are passed through so a stream
that delivers tool calls only on the terminal chunk doesn't silently
drop them.
``prior_text_emitted`` plays the analogous role for the ``text`` field:
when set, the terminal consolidation chunk's ``text`` is suppressed
(it would re-emit the full assembled response on top of prior deltas);
when unset (e.g. a degenerate stream that delivers the entire response
in a single SSE event carrying both ``chatHistory`` and ``finishReason``),
the text is passed through so the response content isn't silently lost.
OCI Cohere streams the answer as single-token ``text`` deltas, then restates
the whole assembled ``text`` on every chunk that carries ``toolCalls`` or
``chatHistory`` (the tool-calls event and the terminal event). Once the
caller reports that earlier chunks already emitted text
(``prior_text_emitted``), those restatements are dropped so the client does
not see the answer twice; a stream whose only text lives on such a chunk
keeps it. ``prior_tool_calls_emitted`` plays the same role for the tool
calls the terminal ``chatHistory`` chunk repeats.
"""
try:
typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk)
@ -315,33 +310,10 @@ def handle_cohere_stream_chunk(
if typed_chunk.index is None:
typed_chunk.index = 0
# OCI Cohere's terminal SSE event re-sends the full assembled response in
# `text` alongside a populated `chatHistory` and a non-null `finishReason`.
# Emitting that text would concatenate the whole response onto the
# already-streamed deltas. We require both signals to be present so that a
# future API change which adds `chatHistory` to intermediate chunks (or a
# rare early-populated case) doesn't silently drop legitimate token deltas.
is_terminal_consolidation: Final = typed_chunk.chatHistory is not None and typed_chunk.finishReason is not None
# On non-terminal text-free chunks (e.g. tool-call-only or keep-alive
# chunks) emit ``content=None`` rather than ``content=""`` so downstream
# stream-mergers that distinguish "no text in this delta" from "an
# explicitly empty text delta" behave correctly.
#
# We only suppress the terminal chunk's ``text`` when the caller has
# confirmed that text deltas were already emitted earlier — otherwise
# (e.g. a degenerate stream that delivers the whole response in a
# single SSE event), passing it through is the only chance to surface it.
text: Final[str | None] = None if (is_terminal_consolidation and prior_text_emitted) else typed_chunk.text
# Tool calls on the terminal consolidation chunk (whether from
# `typed_chunk.toolCalls` or from `chatHistory`) typically restate what
# was already streamed in intermediate chunks. Re-emitting them would
# mint fresh `uuid4` IDs and cause downstream consumers to execute each
# tool call twice. We only suppress when the caller has confirmed that
# tool calls were already emitted earlier — otherwise (e.g. a short
# response that delivers tool calls exclusively on the terminal chunk),
# passing them through is the only chance to surface them.
cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls
restates_text: Final = typed_chunk.chatHistory is not None or typed_chunk.toolCalls is not None
restates_tool_calls: Final = typed_chunk.chatHistory is not None
text: Final[str | None] = None if (restates_text and prior_text_emitted) else typed_chunk.text
cohere_tool_calls: Final = None if (restates_tool_calls and prior_tool_calls_emitted) else typed_chunk.toolCalls
tool_calls: list[dict[str, object]] | None = None
if cohere_tool_calls:

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
)

16
litellm/llms/snowflake/chat/transformation.py
View file

@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
create_anthropic_image_param,
select_anthropic_content_block_type_for_file,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
@ -421,6 +422,21 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
return self._transform_request_anthropic(model, messages, optional_params, stream, extra_body)
return self._transform_request_openai(model, messages, optional_params, stream, extra_body)
@property
def uses_async_transform_request(self) -> bool:
return True
async def async_transform_request(
self,
model: str,
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, object], # mutable-ok: BaseConfig signature
litellm_params: dict[str, object], # mutable-ok: BaseConfig signature
headers: dict[str, object], # mutable-ok: BaseConfig signature
) -> dict[str, object]: # mutable-ok: BaseConfig signature
inlined_messages: Final = await async_inline_remote_media(messages) if _is_claude_model(model) else messages
return self.transform_request(model, inlined_messages, optional_params, litellm_params, headers)
def _transform_request_openai(
self,
model: str,

32
litellm/llms/vertex_ai/gemini/transformation.py
View file

@ -7,6 +7,7 @@ Why separate file? Make it easy to see how transformation works
import json
import os
import re
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from urllib.parse import quote
@ -27,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_gemini_tool_call_result,
response_schema_prompt,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import RemoteMedia, async_inline_remote_media
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
from litellm.types.files import (
@ -68,6 +70,7 @@ _GCS_METADATA_VERTEX_BASE: Any | None = None
# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings
# from litellm's HTTP stack apply (see Greptile review on PR #27278).
_GCS_METADATA_HTTP_HANDLER: HTTPHandler | None = None
GEMINI_FILES_API_URI_PREFIX: Final = "https://generativelanguage.googleapis.com/v1beta/files/"
_GEMINI_MIME_TYPE_ALIASES: Final[dict[str, str]] = {
"image/jpg": "image/jpeg",
}
@ -556,7 +559,7 @@ def _process_gemini_media(
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
part: PartType = {"file_data": file_data}
return _apply_gemini_metadata(part, model, media_resolution_enum, video_metadata)
elif image_url.startswith("https://generativelanguage.googleapis.com/v1beta/files/"):
elif image_url.startswith(GEMINI_FILES_API_URI_PREFIX):
# Gemini Files API URIs — the file is already uploaded to Google's
# servers; pass the URI through as file_data without fetching it.
# These URLs return 403 when accessed directly, so we must not try
@ -1307,6 +1310,23 @@ def sync_transform_request_body(
)
def _explicit_mime_type(fields: Mapping[str, object]) -> str | None:
hint: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type")
return hint if isinstance(hint, str) else None
def _ai_studio_inlines(media: RemoteMedia) -> bool:
return not media.url.startswith(GEMINI_FILES_API_URI_PREFIX)
def _vertex_inlines(media: RemoteMedia) -> bool:
if media.url.startswith(GEMINI_FILES_API_URI_PREFIX):
return False
return media.url.startswith("http://") or (
_explicit_mime_type(media.fields) is None and _get_image_mime_type_from_url(media.url) is None
)
async def async_transform_request_body(
gemini_api_key: str | None,
messages: list[AllMessageValues],
@ -1348,13 +1368,17 @@ async def async_transform_request_body(
vertex_auth_header=vertex_auth_header,
)
if _openai_messages_may_need_sync_gcs_metadata_fetch(messages):
inlined_messages: Final = await async_inline_remote_media(
messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
)
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
# _transform_request_body may issue a sync httpx.get (up to 5s timeout)
# via _get_gcs_object_content_type to fetch GCS object metadata. Run the
# whole sync transformation on a worker thread so it does not block the
# async event loop.
return await asyncify(_transform_request_body)(
messages=messages,
messages=inlined_messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
@ -1363,7 +1387,7 @@ async def async_transform_request_body(
)
return _transform_request_body(
messages=messages,
messages=inlined_messages,
model=model,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,

36
litellm/main.py
View file

@ -4474,7 +4474,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
acompletion: Final = ctx.acompletion
api_base: Final = ctx.api_base
api_key: Final = ctx.api_key
client = _dispatch_client_http(ctx)
injected_client: Final = _dispatch_client_http(ctx)
custom_llm_provider: Final = ctx.custom_llm_provider
headers: Final = ctx.headers
litellm_params: Final = ctx.litellm_params
@ -4486,11 +4486,11 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR
shared_session: Final = ctx.shared_session
stream: Final = ctx.stream
timeout: Final = ctx.timeout
client: Final = (
injected_client if injected_client is not None else (HTTPHandler(timeout=timeout) if stream is False else None)
) # Keep this here, otherwise, the httpx.client closes and streaming is impossible
try:
client = (
HTTPHandler(timeout=timeout) if stream is False else None
) # Keep this here, otherwise, the httpx.client closes and streaming is impossible
response: Final = base_llm_http_handler.completion(
model=model,
messages=messages,
@ -8389,6 +8389,34 @@ def speech(
client=client,
_is_async=aspeech or False,
)
elif custom_llm_provider == "mistral":
from litellm.llms.mistral.audio_speech.transformation import (
MistralTextToSpeechConfig,
)
mistral_tts_config: Final = text_to_speech_provider_config or MistralTextToSpeechConfig()
if api_base is not None:
litellm_params_dict["api_base"] = api_base
if api_key is not None:
litellm_params_dict["api_key"] = api_key
mistral_voice: Final[str | None] = voice if isinstance(voice, str) else None
response = base_llm_http_handler.text_to_speech_handler(
model=model,
input=input,
voice=mistral_voice,
text_to_speech_provider_config=mistral_tts_config,
text_to_speech_optional_params=optional_params,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
)
elif custom_llm_provider == "aws_polly":
from litellm.llms.aws_polly.text_to_speech.transformation import (
AWSPollyTextToSpeechConfig,

65
litellm/model_prices_and_context_window_backup.json
View file

@ -3485,6 +3485,55 @@
"supports_response_schema": true,
"supports_tool_choice": true
},
"azure_ai/gpt-6-astra": {
"cache_creation_input_token_cost": 1.25e-05,
"cache_creation_input_token_cost_above_272k_tokens": 2.5e-05,
"cache_read_input_token_cost": 1e-06,
"cache_read_input_token_cost_above_272k_tokens": 2e-06,
"input_cost_per_token": 1e-05,
"input_cost_per_token_above_272k_tokens": 2e-05,
"litellm_provider": "azure_ai",
"max_input_tokens": 922000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 5e-05,
"output_cost_per_token_above_272k_tokens": 7.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"source": "https://ai.azure.com/catalog/models/gpt-6-astra",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
"supports_parallel_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_cache_breakpoint": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true,
"supports_xhigh_reasoning_effort": true
},
"azure_ai/gpt-5.5": {
"deprecation_date": "2027-10-26",
"cache_read_input_token_cost": 5e-07,
@ -7189,7 +7238,7 @@
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
@ -7455,7 +7504,7 @@
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": true,
"supports_max_reasoning_effort": false,
"supports_minimal_reasoning_effort": false,
"supports_native_streaming": true,
"supports_none_reasoning_effort": true,
@ -30742,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",
@ -30759,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,
@ -34898,9 +34951,9 @@
"supports_audio_input": true
},
"mistral/voxtral-mini-tts-latest": {
"input_cost_per_character": 1.6e-05,
"litellm_provider": "mistral",
"mode": "audio_speech",
"output_cost_per_character": 1.6e-05,
"source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03",
"supported_endpoints": [
"/v1/audio/speech"
@ -37176,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,
@ -37216,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,
@ -37427,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,
@ -37467,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,
@ -56467,9 +56524,9 @@
"supports_audio_input": true
},
"mistral/voxtral-mini-tts-2603": {
"input_cost_per_character": 1.6e-05,
"litellm_provider": "mistral",
"mode": "audio_speech",
"output_cost_per_character": 1.6e-05,
"source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03",
"supported_endpoints": [
"/v1/audio/speech"

85
litellm/ocr/main.py
View file

@ -29,6 +29,8 @@ from litellm.llms.base_llm.ocr.transformation import (
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.rust_bridge import ocr as rust_ocr_bridge
from litellm.rust_bridge.bindings import native_exception_types
from litellm.rust_bridge.configuration import rust_enabled
from litellm.types.router import GenericLiteLLMParams
from litellm.utils import ProviderConfigManager, client
@ -196,12 +198,6 @@ def _rust_ocr_supported(prepared_request: _PreparedOCRRequest) -> bool:
return prepared_request.custom_llm_provider in _RUST_OCR_PROVIDERS
def _rust_ocr_enabled(prepared_request: _PreparedOCRRequest) -> bool:
raw_request_override: Final = prepared_request.litellm_params.get("rust")
request_override: Final = raw_request_override if isinstance(raw_request_override, bool) else None
return rust_ocr_bridge.rust_ocr_enabled(request_override=request_override)
def _rust_bridge_optional_params(
prepared_request: _PreparedOCRRequest,
resolve_secret: Callable[[str], str | None],
@ -286,6 +282,33 @@ def _prepare_rust_ocr_call(
)
def _map_rust_ocr_error(
error: Exception,
prepared_request: _PreparedOCRRequest,
exception_types: tuple[type[BaseException], type[BaseException]] | None,
) -> Exception:
if exception_types is None:
return error
_, upstream_error = exception_types
if not isinstance(error, upstream_error):
return error
error_args: Final = cast( # cast-ok: BaseException.args is typed with Any in the standard library stubs
tuple[object, ...], error.args
)
status_value: Final = error_args[0] if error_args else 0
message_value: Final = error_args[1] if len(error_args) > 1 else str(error)
status: Final = status_value if isinstance(status_value, int) else 0
message: Final = message_value if isinstance(message_value, str) else str(message_value)
error_factory: Final = cast( # cast-ok: the legacy provider interface leaves callable parameters untyped
Callable[..., Exception], prepared_request.provider_config.get_error_class
)
return error_factory(
error_message=message,
status_code=status or 500,
headers={}, # mutable-ok: provider error factories require a concrete header dict
)
def _run_rust_ocr(
prepared_request: _PreparedOCRRequest,
resolve_api_key: Callable[[str], str | None],
@ -296,16 +319,19 @@ def _run_rust_ocr(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
rust_response: Final = rust_ocr_bridge.ocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
try:
rust_response: Final = rust_ocr_bridge.ocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
except Exception as error:
raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
if rust_response is None:
return None
return OCRResponse.model_validate(rust_response)
@ -321,16 +347,19 @@ async def _run_rust_aocr(
prepared_request=prepared_request,
resolve_api_key=resolve_api_key,
)
rust_response: Final = await rust_ocr_bridge.aocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
try:
rust_response: Final = await rust_ocr_bridge.aocr(
model=prepared_request.model,
document=prepared_request.document,
api_key=prepared.api_key,
api_base=prepared.api_base,
custom_llm_provider=prepared_request.custom_llm_provider,
extra_headers=prepared.headers,
optional_params=prepared.optional_params,
timeout=prepared_request.effective_timeout,
)
except Exception as error:
raise _map_rust_ocr_error(error, prepared_request, native_exception_types()) from error
if rust_response is None:
return None
return OCRResponse.model_validate(rust_response)
@ -430,7 +459,7 @@ async def aocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
if _rust_ocr_supported(prepared) and rust_enabled():
from litellm.secret_managers.main import get_secret_str
rust_response: Final = await _run_rust_aocr(
@ -702,7 +731,7 @@ def ocr(
custom_llm_provider = prepared.custom_llm_provider
completion_kwargs.update({"model": model, "custom_llm_provider": custom_llm_provider})
if _rust_ocr_supported(prepared) and _rust_ocr_enabled(prepared):
if _rust_ocr_supported(prepared) and rust_enabled():
from litellm.secret_managers.main import get_secret_str
rust_response: Final = _run_rust_ocr(

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

@ -163,7 +163,10 @@ from litellm.proxy.common_utils.user_api_key_cache import get_management_object_
from litellm.proxy.management_endpoints.sso.id_jag_assertion_capture import (
id_jag_assertion_capture_gap_at_startup,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging, get_server_root_path
from litellm.proxy.middleware.per_request_root_path_middleware import (
get_request_root_path,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.table_repositories import MCPServerRepository
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.mcp import (
@ -3858,7 +3861,7 @@ class MCPServerManager:
if err.tag == "unauthorized" and isinstance(spec.config, AuthorizationCodeConfig):
# authorization_code's missing per-user token -> the per-server browser-OAuth
# challenge, built here where the full MCPServer is in hand.
raise_user_oauth_challenge(server, root_path=get_server_root_path())
raise_user_oauth_challenge(server, root_path=get_request_root_path())
if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig):
# token_exchange (OBO): a missing/rejected subject token -> the RFC 9728 challenge
# pointing at the IdP the client must SSO with to obtain one, rather than an opaque
@ -3866,7 +3869,7 @@ class MCPServerManager:
# Access) threads its claims blob into the challenge for the client to satisfy.
raise_token_exchange_challenge(
server,
root_path=get_server_root_path(),
root_path=get_request_root_path(),
claims=err.unauthorized.claims,
)
raise_public(err)
@ -3914,7 +3917,7 @@ class MCPServerManager:
if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
return
if subject_token is None and isinstance(spec.config, TokenExchangeConfig):
raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
raise_token_exchange_challenge(resolved_server, root_path=get_request_root_path())
match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
case Ok(_):
return
@ -3922,7 +3925,7 @@ class MCPServerManager:
if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig):
raise_token_exchange_challenge(
resolved_server,
root_path=get_server_root_path(),
root_path=get_request_root_path(),
claims=err.unauthorized.claims,
)
raise_public(err)

26
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
View file

@ -12,6 +12,7 @@ every other mode so the caller defers to v1 (parity-safe); it grows one branch p
from __future__ import annotations
import base64
import os
from typing import TYPE_CHECKING, Final, Literal, NoReturn
from fastapi import HTTPException
@ -310,13 +311,30 @@ def raise_public(error: CredError) -> NoReturn:
def oauth_protected_resource_path(root_path: str, server: MCPServer) -> str:
"""The server's RFC 9728 Protected Resource Metadata path, the shared anchor of both challenges.
``root_path`` is the proxy's ``SERVER_ROOT_PATH``, resolved by the caller (the imperative shell)
so this stays a pure function of its inputs; ``"/"`` and ``""`` both mean no prefix. The path is
relative, so it resolves against the caller's own host (correct even behind a reverse proxy).
``root_path`` is the prefix the request was routed under, resolved by the caller (the imperative
shell); ``"/"`` and ``""`` both mean no prefix. The path is relative, so it resolves against the
caller's own host (correct even behind a reverse proxy).
URL structure depends on how the prefix is served:
- The scalar ``SERVER_ROOT_PATH`` deployment registers the well-known routes with the prefix
*inserted* into the path (via :func:`well_known_root_suffix` at import time), matching RFC 8414
§3 well-known path insertion. When ``root_path`` equals ``SERVER_ROOT_PATH`` the URL must use
the same insertion or a client fetching it 404s.
- The per-request ``SERVER_ROOT_PATHS`` deployment can't register routes per prefix (the prefix
set is dynamic and could contain many entries); the middleware strips the prefix from
``scope["path"]`` and the router matches the un-inserted well-known route. The URL must place
the prefix *before* ``.well-known`` so the strip leaves a matching path.
Picking the wrong form 404s the client's discovery fetch — the discovery document and the 401
challenge would then disagree on where the resource metadata lives.
"""
prefix: Final = "" if root_path == "/" else root_path
name: Final = server.alias or server.server_name or server.name or server.server_id
return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}"
scalar_env: Final = os.getenv("SERVER_ROOT_PATH", "").rstrip("/")
if not prefix or (scalar_env and prefix == scalar_env):
return f"/.well-known/oauth-protected-resource{prefix}/mcp/{name}"
return f"{prefix}/.well-known/oauth-protected-resource/mcp/{name}"
def raise_user_oauth_challenge(server: MCPServer, *, root_path: str) -> NoReturn:

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

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

@ -911,15 +911,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 +964,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,
@ -3848,12 +3857,14 @@ if MCP_AVAILABLE:
# then exchanges. A tool-call-time 401 would be wrapped into a JSON-RPC error and the
# header lost, so the discovery flow needs this pre-emptive challenge.
if server and server.auth_type == MCPAuth.oauth2_token_exchange and not oauth2_headers:
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import ( # noqa: PLC0415 # lazy: adapter pulls MCP subgraph
raise_token_exchange_challenge,
)
from litellm.proxy.utils import get_server_root_path # noqa: PLC0415
from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports proxy utils
get_request_root_path,
)
raise_token_exchange_challenge(server, root_path=get_server_root_path())
raise_token_exchange_challenge(server, root_path=get_request_root_path())
# Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run
# the exchange here at the transport edge, so a rejected subject raises the RFC 9728

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

@ -11,6 +11,7 @@ from pydantic import ValidationError
from typing_extensions import ReadOnly, Required, assert_never
import litellm
from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
from litellm.proxy.agent_endpoints.agent_search import DEFAULT_AGENT_SEARCH_TOP_K
from litellm.proxy.common_utils.semantic_text_index import (
Embedder,
@ -30,7 +31,10 @@ MCP_TOOL_SEARCH_SETTINGS_KEY: Final[str] = "mcp_tool_search"
MCP_TOOL_SEARCH_TOOL_NAME: Final[str] = "mcp_tool_search"
MCP_TOOL_CALL_TOOL_NAME: Final[str] = "mcp_tool_call"
AGENT_SEARCH_TOOL_NAME: Final[str] = "agent_search"
VIRTUAL_TOOL_NAMES: Final = frozenset((MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME))
SKILL_SEARCH_TOOL_NAME: Final[str] = "skill_search"
VIRTUAL_TOOL_NAMES: Final = frozenset(
(MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, SKILL_SEARCH_TOOL_NAME)
)
def coerce_top_k(value: Any, default: int = 5) -> int:
@ -199,8 +203,28 @@ _AGENT_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
}
_SKILL_SEARCH_DEFINITION: Final[VirtualToolDefinition] = {
"name": SKILL_SEARCH_TOOL_NAME,
"description": "Find registered skills by describing what you need in natural language. Returns the best "
"matching skills you can access, ranked by semantic similarity, each with its skill_id, display_title, "
"description, and score.",
"inputSchema": {
"type": "object",
"properties": {
"query": {"type": "string", "description": "What you need the skill to do, in natural language."},
"top_k": {
"type": "integer",
"description": "Maximum number of skills to return.",
"default": DEFAULT_SKILL_SEARCH_TOP_K,
},
},
"required": _json_array("query"),
},
}
def get_virtual_tool_definitions() -> tuple[VirtualToolDefinition, ...]:
return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION)
return (_MCP_TOOL_SEARCH_DEFINITION, _MCP_TOOL_CALL_DEFINITION, _AGENT_SEARCH_DEFINITION, _SKILL_SEARCH_DEFINITION)
def _text_tool_result(text: str, is_error: bool) -> CallToolResult:
@ -223,7 +247,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI
)
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.proxy_server import llm_router
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
await check_feature_access_for_user(user_api_key_dict, "agents")
outcome: Final = await search_agents(
@ -234,6 +258,7 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI
embedding_model=litellm.agent_search_embedding_model,
index=global_agent_search_index,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
match outcome:
case AgentSearchHits(hits):
@ -245,6 +270,39 @@ async def handle_agent_search(query: str, top_k: int, user_api_key_dict: UserAPI
assert_never(outcome)
async def handle_skill_search(query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth) -> CallToolResult:
from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler
from litellm.llms.litellm_proxy.skills.skill_search import (
MAX_SKILL_SEARCH_TOP_K,
SkillSearchEmbeddingFailed,
SkillSearchHits,
SkillSearchNotConfigured,
global_skill_search_index,
search_skills,
skill_search_result,
)
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
outcome: Final = await search_skills(
query=query,
skills=await LiteLLMSkillsHandler.list_skills_for_search(user_api_key_dict),
top_k=min(max(top_k, 1), MAX_SKILL_SEARCH_TOP_K),
router=llm_router,
embedding_model=litellm.skill_search_embedding_model,
index=global_skill_search_index,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
match outcome:
case SkillSearchHits(hits):
results: Final = tuple(skill_search_result(hit).model_dump() for hit in hits)
return _text_tool_result(json.dumps(results), is_error=False)
case SkillSearchNotConfigured(reason) | SkillSearchEmbeddingFailed(reason):
return _text_tool_result(reason, is_error=True)
case _:
assert_never(outcome)
async def handle_mcp_tool_search(
query: str,
top_k: int,
@ -257,7 +315,7 @@ async def handle_mcp_tool_search(
raw_headers: dict[str, str] | None = None,
) -> CallToolResult:
from litellm.proxy._experimental.mcp_server.server import _list_mcp_tools
from litellm.proxy.proxy_server import llm_router
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
settings: Final = mcp_tool_search_settings()
if isinstance(settings, ValidationError):
@ -271,7 +329,7 @@ async def handle_mcp_tool_search(
)
ranker: Final = (
SemanticToolRanker(
embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict),
embed=router_embedder(llm_router, settings.embedding_model, user_api_key_dict, proxy_logging_obj),
embedding_model=settings.embedding_model,
index=global_mcp_tool_search_index,
)

19
litellm/proxy/_lazy_features.py
View file

@ -295,15 +295,18 @@ class LazyFeatureMiddleware:
# Short-circuit once every feature has loaded.
if scope["type"] in ("http", "websocket") and len(self._loaded) < len(self._features):
path = scope.get("path", "")
# Strip SERVER_ROOT_PATH so prefix matching works under a server
# root path. Without this, requests like /api/v1/policies/... never
# match the registered prefixes (/policies/...) and lazy features
# stay unloaded — every endpoint under them returns 404. The
# Strip the request's root_path so prefix matching works under a
# server root path. Without this, requests like /api/v1/policies/...
# never match the registered prefixes (/policies/...) and lazy
# features stay unloaded — every endpoint under them returns 404.
# scope["root_path"] wins over the cached env scalar: FastAPI
# stamps SERVER_ROOT_PATH there, and PerRequestRootPathMiddleware
# resolves SERVER_ROOT_PATHS prefixes there per request. The
# `+ "/"` boundary prevents false-positive matches (e.g. /apiv2
# against root /api). If the path doesn't start with the prefix
# (e.g. a reverse proxy already stripped it), we leave it alone.
if self._root_path and path.startswith(self._root_path + "/"):
path = path[len(self._root_path) :]
# against root /api); a pre-stripped path is left alone.
root_path: Final = str(scope.get("root_path", "")).rstrip("/") or self._root_path
if root_path and path.startswith(root_path + "/"):
path = path[len(root_path) :] # rebind-ok: local strip after the boundary check above
for feat in self._features:
if feat.module_path in self._loaded:
continue

69
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -4687,6 +4687,17 @@
"title": "Created At",
"type": "string"
},
"description": {
"anyOf": [
{
"type": "string"
},
{
"type": "null"
}
],
"title": "Description"
},
"display_title": {
"anyOf": [
{
@ -4713,6 +4724,17 @@
],
"title": "Latest Version"
},
"search_score": {
"anyOf": [
{
"type": "number"
},
{
"type": "null"
}
],
"title": "Search Score"
},
"source": {
"title": "Source",
"type": "string"
@ -4781,7 +4803,7 @@
"paths": {
"/v1/skills": {
"get": {
"description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nReturns: ListSkillsResponse with list of skills",
"description": "List skills on Anthropic.\n\nRequires `?beta=true` query parameter.\n\nModel-based routing (for multi-account support):\n- Pass model via header: `x-litellm-model: claude-account-1`\n- Pass model via query: `?model=claude-account-1`\n- Pass model via body: `{\"model\": \"claude-account-1\"}`\n\nExample usage:\n```bash\n# Basic usage\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\"\n\n# With model-based routing\ncurl \"http://localhost:4000/v1/skills?beta=true&limit=10\" -H \"Authorization: Bearer your-key\" -H \"x-litellm-model: claude-account-1\"\n```\n\nPass `?custom_llm_provider=litellm_proxy&query=<task>` to rank the LiteLLM-hosted skills you can\naccess by semantic similarity instead of paging through the whole registry:\n```bash\ncurl \"http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5\" -H \"Authorization: Bearer your-key\"\n```\n\nReturns: ListSkillsResponse with list of skills",
"operationId": "list_skills_v1_skills_get",
"parameters": [
{
@ -4849,6 +4871,39 @@
"default": "anthropic",
"title": "Custom Llm Provider"
}
},
{
"description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.",
"in": "query",
"name": "query",
"required": false,
"schema": {
"anyOf": [
{
"minLength": 1,
"type": "string"
},
{
"type": "null"
}
],
"description": "Describe what you need in natural language to rank the skills you can access by semantic similarity over their title and description. Each result carries a search_score. Only supported for custom_llm_provider=litellm_proxy. Requires litellm_settings.skill_search_embedding_model.",
"title": "Query"
}
},
{
"description": "With query: the maximum number of ranked skills to return.",
"in": "query",
"name": "top_k",
"required": false,
"schema": {
"default": 5,
"description": "With query: the maximum number of ranked skills to return.",
"maximum": 100,
"minimum": 1,
"title": "Top K",
"type": "integer"
}
}
],
"responses": {
@ -9597,7 +9652,9 @@
{
"enum": [
"warn",
"end_session"
"end_session",
"block",
"alert"
],
"type": "string"
},
@ -9605,7 +9662,7 @@
"type": "null"
}
],
"description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.",
"description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.",
"title": "On Violation"
},
"only_scan_new_messages": {
@ -11914,7 +11971,9 @@
{
"enum": [
"warn",
"end_session"
"end_session",
"block",
"alert"
],
"type": "string"
},
@ -11922,7 +11981,7 @@
"type": "null"
}
],
"description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.",
"description": "For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection. For guardrail='mcp_security': 'block' rejects the request; 'alert' only logs a warning.",
"title": "On Violation"
},
"only_scan_new_messages": {

2
litellm/proxy/_types.py
View file

@ -848,6 +848,7 @@ class LiteLLMRoutes(enum.Enum):
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
"/vector_store/info",
# Project read routes - endpoint scopes results to caller's teams (non-admin)
"/project/list",
"/project/info",
@ -1276,6 +1277,7 @@ class UpdateKeyRequest(KeyRequestBase):
# else they will get overwritten
duration: str | None = None
spend: float | None = None
soft_budget: float | None = None
metadata: dict | None = None
temp_budget_increase: float | None = None
temp_budget_expiry: datetime | None = None

7
litellm/proxy/agent_endpoints/agent_search.py
View file

@ -18,6 +18,7 @@ from litellm.types.agents import AgentResponse
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
DEFAULT_AGENT_SEARCH_TOP_K: Final = 5
@ -132,6 +133,7 @@ async def search_agents(
embedding_model: str | None,
index: AgentSearchIndex,
user_api_key_dict: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
) -> AgentSearchOutcome:
if embedding_model is None:
return AgentSearchNotConfigured(
@ -139,6 +141,5 @@ async def search_agents(
)
if router is None:
return AgentSearchNotConfigured(reason="agent search needs a model_list so the embedding model can be called")
return await index.search(
query, agents, top_k, router_embedder(router, embedding_model, user_api_key_dict), embedding_model
)
embed: Final = router_embedder(router, embedding_model, user_api_key_dict, proxy_logging_obj)
return await index.search(query, agents, top_k, embed, embedding_model)

3
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -249,7 +249,7 @@ def _agent_search_error(status_code: int, error: str, message: str) -> HTTPExcep
async def _rank_agents_by_query(
query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth
) -> tuple[AgentResponse, ...]:
from litellm.proxy.proxy_server import llm_router
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
outcome: Final = await search_agents(
query=query,
@ -259,6 +259,7 @@ async def _rank_agents_by_query(
embedding_model=litellm.agent_search_embedding_model,
index=global_agent_search_index,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
match outcome:
case AgentSearchHits(hits):

97
litellm/proxy/anthropic_endpoints/skills_endpoints.py
View file

@ -2,11 +2,23 @@
Anthropic Skills API endpoints - /v1/skills
"""
from typing import Final
from types import MappingProxyType
from typing import Annotated, Final
import orjson
from fastapi import APIRouter, Depends, Request, Response
from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm.llms.litellm_proxy.skills.skill_search import (
DEFAULT_SKILL_SEARCH_TOP_K,
SkillSearchEmbeddingFailed,
SkillSearchHits,
SkillSearchNotConfigured,
SkillSearchUnsupportedProvider,
global_skill_search_index,
search_hosted_skills,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
@ -23,6 +35,51 @@ from litellm.types.llms.anthropic_skills import (
router: Final = APIRouter()
class _SkillSearchErrorDetail(TypedDict):
error: ReadOnly[str]
message: ReadOnly[str]
def _skill_search_error(status_code: int, error: str, message: str) -> HTTPException:
detail: Final[_SkillSearchErrorDetail] = {"error": error, "message": message}
return HTTPException(status_code=status_code, detail=detail)
async def _search_skills(
custom_llm_provider: str | None, query: str, top_k: int, user_api_key_dict: UserAPIKeyAuth
) -> ListSkillsResponse:
from litellm.llms.litellm_proxy.skills.transformation import (
LiteLLMSkillsTransformationHandler,
)
from litellm.proxy.proxy_server import llm_router, proxy_logging_obj
outcome: Final = await search_hosted_skills(
custom_llm_provider=custom_llm_provider,
query=query,
top_k=top_k,
router=llm_router,
embedding_model=litellm.skill_search_embedding_model,
index=global_skill_search_index,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
to_response: Final = LiteLLMSkillsTransformationHandler().db_skill_to_response
match outcome:
case SkillSearchHits(hits):
skills: Final = [ # mutable-ok: ListSkillsResponse.data requires list[Skill]; never mutated after
to_response(hit.skill).model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits
]
return ListSkillsResponse(data=skills, has_more=False, next_page=None)
case SkillSearchUnsupportedProvider(reason):
raise _skill_search_error(400, "skill_search_unsupported_provider", reason)
case SkillSearchNotConfigured(reason):
raise _skill_search_error(400, "skill_search_not_configured", reason)
case SkillSearchEmbeddingFailed(reason):
raise _skill_search_error(503, "skill_search_unavailable", reason)
case _:
assert_never(outcome)
@router.post(
"/v1/skills",
tags=["[beta] Anthropic Skills API"],
@ -134,32 +191,58 @@ async def list_skills(
after_id: str | None = None,
before_id: str | None = None,
custom_llm_provider: str | None = "anthropic",
query: Annotated[
str | None,
Query(
min_length=1,
description="Describe what you need in natural language to rank the skills you can access by "
"semantic similarity over their title and description. Each result carries a search_score. "
"Only supported for custom_llm_provider=litellm_proxy. Requires "
"litellm_settings.skill_search_embedding_model.",
),
] = None,
top_k: Annotated[
int,
Query(ge=1, le=100, description="With query: the maximum number of ranked skills to return."),
] = DEFAULT_SKILL_SEARCH_TOP_K,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
List skills on Anthropic.
Requires `?beta=true` query parameter.
Model-based routing (for multi-account support):
- Pass model via header: `x-litellm-model: claude-account-1`
- Pass model via query: `?model=claude-account-1`
- Pass model via body: `{"model": "claude-account-1"}`
Example usage:
```bash
# Basic usage
curl "http://localhost:4000/v1/skills?beta=true&limit=10" \
-H "Authorization: Bearer your-key"
# With model-based routing
curl "http://localhost:4000/v1/skills?beta=true&limit=10" \
-H "Authorization: Bearer your-key" \
-H "x-litellm-model: claude-account-1"
```
Pass `?custom_llm_provider=litellm_proxy&query=<task>` to rank the LiteLLM-hosted skills you can
access by semantic similarity instead of paging through the whole registry:
```bash
curl "http://localhost:4000/v1/skills?custom_llm_provider=litellm_proxy&query=summarize+a+pdf&top_k=5" \
-H "Authorization: Bearer your-key"
```
Returns: ListSkillsResponse with list of skills
"""
if query is not None:
return await _search_skills(
custom_llm_provider=custom_llm_provider, query=query, top_k=top_k, user_api_key_dict=user_api_key_dict
)
from litellm.proxy.proxy_server import (
general_settings,
llm_router,

19
litellm/proxy/auth/auth_utils.py
View file

@ -1981,9 +1981,28 @@ def get_model_from_request(
if vertex_match:
model = vertex_match.group(1)
if route.lower().startswith("/bedrock"):
bedrock_model: Final = _model_from_bedrock_route(route)
return model if bedrock_model is None else bedrock_model
return model
def _model_from_bedrock_route(route: str) -> str | None:
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_extract_model_from_bedrock_endpoint,
is_bedrock_count_tokens_endpoint,
)
bedrock_endpoint: Final = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE)
if is_bedrock_count_tokens_endpoint(bedrock_endpoint):
return None
try:
return _extract_model_from_bedrock_endpoint(bedrock_endpoint)
except ValueError:
return None
def abbreviate_api_key(api_key: str) -> str:
if len(api_key) < MINIMUM_CUSTOM_KEY_LENGTH:
return "sk-..."

126
litellm/proxy/common_request_processing.py
View file

@ -756,7 +756,7 @@ class _UpstreamClosingStreamingResponse(StreamingResponse):
content: AsyncGenerator[str, None],
*,
media_type: str | None = None,
headers: dict | None = None,
headers: Mapping[str, str] | None = None,
status_code: int = status.HTTP_200_OK,
upstream_generator: AsyncGenerator[str, None] | None = None,
) -> None:
@ -888,25 +888,39 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]:
return f"data: {json.dumps({'error': error_obj})}\n\n", "data: [DONE]\n\n"
def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
"""`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384)."""
return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS})
async def _resolve_stream_headers(
headers: Mapping[str, str], refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None
) -> Mapping[str, str]:
if refresh_headers is None:
return headers
try:
return await refresh_headers()
except Exception as e: # noqa: BLE001 # a stream whose first chunk is already paid for must not fail over its headers
verbose_proxy_logger.exception("Error refreshing streaming response headers: %s", e)
return headers
async def create_response(
generator: AsyncGenerator[str, None],
media_type: str,
headers: dict,
headers: Mapping[str, str],
default_status_code: int = status.HTTP_200_OK,
request: Request | None = None,
refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None,
) -> StreamingResponse | JSONResponse:
"""
Create streaming response, checking if the first chunk is an error.
If the first chunk is an error, return a standard JSON error response.
Otherwise, return StreamingResponse and stream all content.
``refresh_headers`` is consulted once the first chunk has been buffered, for
callers whose headers can only be known then.
"""
# Tell buffering reverse proxies (nginx, ingress-nginx, Envoy) to flush SSE
# immediately instead of releasing the whole stream in one batch (issue #28384).
streaming_headers: Final = {
**headers,
"Cache-Control": "no-cache",
"X-Accel-Buffering": "no",
}
first_chunk_value: str | None = None
final_status_code = default_status_code
@ -917,6 +931,7 @@ async def create_response(
# Now get the first chunk from the actual generator
first_chunk_value = await _buffer_first_chunk_honoring_disconnect(generator, request)
resolved_headers: Final = await _resolve_stream_headers(headers, refresh_headers)
if first_chunk_value is not None:
try:
@ -943,7 +958,7 @@ async def create_response(
return JSONResponse(
status_code=final_status_code,
content={"error": error_dict},
headers=headers,
headers=resolved_headers,
)
except Exception as e:
verbose_proxy_logger.debug("Error parsing first chunk value: %s", e)
@ -972,7 +987,7 @@ async def create_response(
return StreamingResponse(
empty_gen(),
media_type=media_type,
headers=streaming_headers,
headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)),
status_code=default_status_code,
)
except Exception as e:
@ -988,7 +1003,7 @@ async def create_response(
return StreamingResponse(
error_gen_message(),
media_type=media_type,
headers=streaming_headers,
headers=_sse_stream_headers(await _resolve_stream_headers(headers, refresh_headers)),
status_code=error_status,
)
@ -1010,7 +1025,7 @@ async def create_response(
return _UpstreamClosingStreamingResponse(
combined_generator(),
media_type=media_type,
headers=streaming_headers,
headers=_sse_stream_headers(resolved_headers),
status_code=final_status_code,
upstream_generator=generator,
)
@ -1535,7 +1550,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _merge_passthrough_streaming_headers(
response_headers: httpx.Headers | dict | None,
custom_headers: dict,
custom_headers: Mapping[str, str],
) -> dict:
"""
Merge upstream passthrough headers with proxy/custom headers.
@ -2143,14 +2158,45 @@ class ProxyBaseLLMRequestProcessing:
return fallback_model_group
@staticmethod
def _get_model_id_from_response(hidden_params: dict, data: dict) -> str:
def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
"""Extract model_id from hidden_params with fallback to litellm_metadata."""
model_id = hidden_params.get("model_id", None) or ""
if not model_id:
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
model_info: Final = litellm_metadata.get("model_info", {}) or {}
model_id = model_info.get("id", "") or ""
return model_id
litellm_metadata: Final = data.get("litellm_metadata")
model_info: Final = litellm_metadata.get("model_info") if isinstance(litellm_metadata, Mapping) else None
model_id = (model_info.get("id") or "") if isinstance(model_info, Mapping) else ""
return str(model_id) if model_id else ""
def _stream_response_headers(
self,
*,
hidden_params: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
logging_obj: LiteLLMLoggingObj,
version: str | None,
callback_headers: Mapping[str, str],
) -> Mapping[str, str]:
"""The streaming response headers describing `hidden_params`' deployment."""
return MappingProxyType(
{
**ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
model_id=self._get_model_id_from_response(hidden_params, self.data),
cache_key=hidden_params.get("cache_key") or "",
api_base=hidden_params.get("api_base") or "",
version=version,
response_cost=hidden_params.get("response_cost") or "",
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"),
request_data=self.data,
hidden_params=hidden_params,
litellm_logging_obj=logging_obj,
**(hidden_params.get("additional_headers") or MappingProxyType({})),
),
**callback_headers,
}
)
@staticmethod
def _get_deployment_model_name(
@ -2419,31 +2465,32 @@ class ProxyBaseLLMRequestProcessing:
if self._is_streaming_request(
data=self.data, is_streaming_request=is_streaming_request
) or self._is_streaming_response(response): # use generate_responses to stream responses
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=logging_obj.litellm_call_id,
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
version=version,
response_cost=response_cost,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
fastest_response_batch_completion=fastest_response_batch_completion,
request_data=self.data,
hidden_params=hidden_params,
litellm_logging_obj=logging_obj,
**additional_headers,
)
# Call response headers hook for streaming success
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
stream_callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook(
data=self.data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=dict(request.headers),
)
if callback_headers:
custom_headers.update(callback_headers)
custom_headers: Final = self._stream_response_headers(
hidden_params=hidden_params,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
callback_headers=stream_callback_headers or MappingProxyType({}),
)
async def refresh_stream_headers() -> Mapping[str, str]:
"""`custom_headers` rebuilt for whichever deployment served the stream."""
if not getattr(response, "fallback_headers_adopted", False):
return custom_headers
return self._stream_response_headers(
hidden_params=get_hidden_params_dict(response),
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
version=version,
callback_headers=stream_callback_headers or MappingProxyType({}),
)
# Preserve the original client-requested model (pre-alias mapping) for downstream
# streaming generators. Pre-call processing can rewrite `self.data["model"]` for
@ -2581,6 +2628,7 @@ class ProxyBaseLLMRequestProcessing:
media_type="text/event-stream",
headers=custom_headers,
request=request,
refresh_headers=refresh_stream_headers,
)
### CALL HOOKS ### - modify outgoing data
@ -3032,7 +3080,7 @@ class ProxyBaseLLMRequestProcessing:
response: Any,
proxy_logging_obj: "ProxyLogging",
user_api_key_dict: "UserAPIKeyAuth",
custom_headers: dict,
custom_headers: Mapping[str, str],
request_headers: dict[str, str],
) -> Response | None:
if not self._has_post_call_guardrails_for_passthrough():

79
litellm/proxy/common_utils/semantic_text_index.py
View file

@ -5,10 +5,11 @@ from __future__ import annotations
import math
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from itertools import chain
from itertools import chain, islice
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
from fastapi import HTTPException
from openai import OpenAIError
from pydantic import BaseModel, ConfigDict
@ -16,10 +17,14 @@ from litellm.exceptions import BudgetExceededError
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
Vector: TypeAlias = tuple[float, ...]
DEFAULT_MAX_CACHED_VECTORS: Final = 5000
"""Ceiling on how many (embedding model, text) vectors one index keeps; the least recently searched are evicted first."""
class Embedder(Protocol):
def __call__(self, texts: Sequence[str]) -> Awaitable[Sequence[Vector]]: ...
@ -42,6 +47,16 @@ class _EmbeddingData(BaseModel):
data: tuple[_EmbeddingItem, ...]
class _EmbeddingRequest(BaseModel):
"""The /embeddings-shaped request as the pre-call hooks (rate limits, budgets, guardrails) hand it back."""
model_config = ConfigDict(frozen=True, extra="ignore")
model: str
input: tuple[str, ...]
metadata: dict[str, object] # mutable-ok: the router mutates the metadata dict it is handed
def cosine_similarity(left: Vector, right: Vector) -> float:
dot: Final = sum(a * b for a, b in zip(left, right, strict=True))
norms: Final = math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right))
@ -57,23 +72,40 @@ def embedding_spend_metadata(user_api_key_dict: UserAPIKeyAuth) -> dict[str, obj
}
def router_embedder(router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth) -> Embedder:
def router_embedder(
router: Router, embedding_model: str, user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging
) -> Embedder:
"""Embeds through the router after the same key rate-limit, budget and guardrail pre-call hooks /embeddings runs."""
async def embed(texts: Sequence[str]) -> Sequence[Vector]:
batch: Final = list(texts) # mutable-ok: Router.aembedding accepts only str | list input
request: Final = { # mutable-ok: pre_call_hook mutates the request dict in place
"model": embedding_model,
"input": list(texts), # mutable-ok: Router.aembedding accepts only str | list input
"metadata": embedding_spend_metadata(user_api_key_dict),
}
processed: Final = _EmbeddingRequest.model_validate(
await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=request, call_type="aembedding"
)
)
response: Final = await router.aembedding(
model=embedding_model, input=batch, metadata=embedding_spend_metadata(user_api_key_dict)
model=processed.model,
input=list(processed.input), # mutable-ok: Router.aembedding accepts only str | list input
metadata=processed.metadata,
)
return tuple(item.embedding for item in _EmbeddingData.model_validate(response.model_dump()).data)
return embed
_NO_VECTORS: Final[Mapping[str, Vector]] = MappingProxyType({})
_CacheKey: TypeAlias = tuple[str, str]
async def _embed_all(embed: Embedder, texts: Sequence[str]) -> tuple[Vector, ...] | EmbeddingFailed:
try:
vectors: Final = tuple(await embed(texts))
except HTTPException:
raise
except (OpenAIError, ValueError, BudgetExceededError) as exc:
return EmbeddingFailed(reason=f"embedding the search query failed: {exc}")
if len(vectors) != len(texts):
@ -111,20 +143,34 @@ async def _embed_query_and_texts(
class SemanticTextIndex:
"""Caches one vector per distinct text per embedding model, so repeat searches only embed the query."""
"""Caches one vector per distinct text per embedding model, so repeat searches only embed the query.
def __init__(self) -> None:
self._vectors: Mapping[str, Mapping[str, Vector]] = MappingProxyType({})
Holds at most ``max_entries`` vectors across all models: once full, the texts no recent search touched go first."""
def _merged(self, embedding_model: str, embedded: _Embedded) -> Mapping[str, Vector]:
kept: Final = MappingProxyType(
def __init__(self, max_entries: int = DEFAULT_MAX_CACHED_VECTORS) -> None:
self._max_entries: Final = max_entries
self._vectors: Mapping[_CacheKey, Vector] = MappingProxyType({})
def _cached(self, embedding_model: str) -> Mapping[str, Vector]:
return MappingProxyType(
{text: vector for (model, text), vector in self._vectors.items() if model == embedding_model}
)
def _merged(self, embedding_model: str, embedded: _Embedded, texts: Sequence[str]) -> Mapping[_CacheKey, Vector]:
dimension: Final = len(embedded.query_vector)
touched: Final = MappingProxyType({(embedding_model, text): embedded.vectors[text] for text in texts})
untouched: Final = MappingProxyType(
{
text: vector
for text, vector in self._vectors.get(embedding_model, _NO_VECTORS).items()
if len(vector) == len(embedded.query_vector)
key: vector
for key, vector in chain(
self._vectors.items(),
(((embedding_model, text), vector) for text, vector in embedded.vectors.items()),
)
if key not in touched and (key[0] != embedding_model or len(vector) == dimension)
}
)
return MappingProxyType({**kept, **embedded.vectors})
ordered: Final = MappingProxyType({**untouched, **touched})
return MappingProxyType(dict(islice(ordered.items(), max(len(ordered) - self._max_entries, 0), None)))
async def scores(
self, query: str, texts: Sequence[str], embed: Embedder, embedding_model: str
@ -132,11 +178,10 @@ class SemanticTextIndex:
"""Cosine similarity of `query` to each entry of `texts`, in the same order."""
if not texts:
return ()
cached: Final = self._vectors.get(embedding_model, _NO_VECTORS)
embedded: Final = await _embed_query_and_texts(embed, query, texts, cached)
embedded: Final = await _embed_query_and_texts(embed, query, texts, self._cached(embedding_model))
if isinstance(embedded, EmbeddingFailed):
return embedded
if not _same_dimension(embedded.query_vector, embedded.vectors, texts):
return EmbeddingFailed(reason=f"embedding model {embedding_model} returned vectors of mixed dimensions")
self._vectors = MappingProxyType({**self._vectors, embedding_model: self._merged(embedding_model, embedded)})
self._vectors = self._merged(embedding_model, embedded, texts)
return tuple(cosine_similarity(embedded.query_vector, embedded.vectors[text]) for text in texts)

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

@ -37,7 +37,9 @@ CACHE_TTL_1H_SECONDS: Final = 3600
AUTOROUTER_BENCHMARKS_SQL: Final = """
WITH windowed AS (
SELECT * FROM "LiteLLM_AutoRouterSession"
WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp
WHERE last_turn_at >= $1::timestamp
AND first_turn_at < $2::timestamp
AND ($3::text IS NULL OR api_key = $3::text)
),
tier_maps AS (
SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns

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

@ -14,6 +14,7 @@ import time
import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
import litellm
@ -84,6 +85,25 @@ else:
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
def _is_batch_cost_row(payload: SpendLogsPayload) -> bool:
return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success"
_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"})
def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]:
"""Reduce a batch's cost row to what tells the retrieves apart when logging is off.
A proxy run with spend logs disabled still needs one row per batch to charge it once,
so the row is written either way, but it carries no request of its own: no metadata,
no requester IP, no key, model, or token counts (LIT-7048).
"""
if disable_spend_logs is False:
return payload
return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS})
class _SpendBatch(Protocol):
litellm_usertable: BatchTable
litellm_verificationtoken: BatchTable
@ -215,7 +235,12 @@ class DBSpendUpdateWriter:
start_time: datetime | None,
end_time: datetime | None,
response_cost: float | None,
) -> None:
) -> bool:
"""Record the request's spend, answering whether its cost still needs charging.
False only for a batch retrieve whose cost row another retrieve already wrote,
so the caller leaves the key, team, and user counters alone (LIT-7048).
"""
from litellm.proxy.proxy_server import (
disable_spend_logs,
litellm_proxy_budget_name,
@ -232,7 +257,7 @@ class DBSpendUpdateWriter:
team_id,
)
if ProxyUpdateSpend.disable_spend_updates() is True:
return
return True
if token is not None and isinstance(token, str) and token.startswith("sk-"):
hashed_token = hash_token(token=token)
else:
@ -262,11 +287,12 @@ class DBSpendUpdateWriter:
if team_id is not None and team_id != "":
payload["team_id"] = team_id
if not await self._record_spend_log(
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
):
return False
if disable_spend_logs is False:
await self._insert_spend_log_to_db(
payload=payload,
prisma_client=prisma_client,
)
await self._enqueue_tool_usage_transaction(
payload=payload,
completion_response=completion_response,
@ -306,6 +332,7 @@ class DBSpendUpdateWriter:
)
verbose_proxy_logger.debug("Runs spend update on all tables")
return True
except Exception:
spend_log_error(
"Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue "
@ -318,7 +345,102 @@ class DBSpendUpdateWriter:
org_id,
end_user_id,
)
return
return True
async def _record_spend_log(
self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool
) -> bool:
if prisma_client is not None and _is_batch_cost_row(payload):
return await self._claim_batch_cost_spend_log(
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
)
if disable_spend_logs is False:
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
return True
async def _claim_batch_cost_spend_log(
self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool
) -> bool:
"""Write the batch's cost row now, or learn that another retrieve already did.
Every retrieve of one batch shares this row, so the insert that lands first owns
the charge and every later one finds the row and charges nothing (LIT-7048). Only
a row that recorded a charge counts: a failed retrieve, a request whose client
picked the batch id as its call id, and the $0 row an older proxy left behind
while the batch was still running all leave the charge to be made.
"""
from litellm.repositories.table_repositories import SpendLogsRepository
request_id: Final = payload["request_id"]
row: Final = _batch_cost_row_to_write(payload, disable_spend_logs)
spend_logs: Final = SpendLogsRepository(prisma_client).table
try:
claimed: Final = await spend_logs.create_many(
data=[prisma_client.jsonify_object(row)], # mutable-ok: prisma create_many takes a list
skip_duplicates=True,
)
if claimed == 1:
return True
existing: Final = await spend_logs.find_unique(
where={"request_id": request_id} # mutable-ok: prisma where clause
)
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log
verbose_proxy_logger.warning(
"Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e
)
await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client)
return True
if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success":
verbose_proxy_logger.warning(
"Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own",
request_id,
getattr(existing, "call_type", None),
)
return True
if existing.spend > 0:
verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id)
return False
return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row)
async def _take_over_uncharged_batch_cost_row(
self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object]
) -> bool:
"""Take the batch's cost row over from the poll that left it charging nothing.
A pre-upgrade proxy wrote that row every time it polled the batch while it was still
running, so the charge is still to be made and the row still has to end up carrying
it. The row stops matching the moment it carries a charge, so it is one retrieve that
takes it over and charges, and every later one reads the charge and charges nothing.
"""
from litellm.repositories.table_repositories import SpendLogsRepository
request_id: Final = payload["request_id"]
if payload["spend"] <= 0:
verbose_proxy_logger.debug(
"Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id
)
return False
try:
taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many(
data=prisma_client.jsonify_object(
MappingProxyType({field: value for field, value in row.items() if field != "request_id"})
),
where={ # mutable-ok: prisma where clause
"request_id": request_id,
"call_type": CallTypes.aretrieve_batch.value,
"status": "success",
"spend": 0.0,
},
)
except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over
verbose_proxy_logger.warning(
"Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e
)
return False
if taken_over == 0:
verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id)
return False
return True
async def _enqueue_tool_usage_transaction(
self,
@ -380,6 +502,7 @@ class DBSpendUpdateWriter:
llm_router=get_llm_router,
cost_breakdown=metadata.get("cost_breakdown"),
recorded_autorouter_savings=metadata.get("autorouter_savings"),
billed_at=payload.get("endTime"),
)
transaction: Final = build_autorouter_turn_transaction(
payload=payload,
@ -2066,6 +2189,7 @@ class DBSpendUpdateWriter:
usage_object=usage_obj,
cost_breakdown=_metadata.get("cost_breakdown"),
recorded_autorouter_savings=_metadata.get("autorouter_savings"),
billed_at=payload.get("endTime"),
)
daily_transaction: Final = BaseDailySpendTransaction(

30
litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
View file

@ -16,7 +16,12 @@ keeps the batched-DELETE path, so existing deployments are untouched.
import re
from collections.abc import Callable
from datetime import date, datetime, timedelta, timezone
from typing import TYPE_CHECKING, Final, TypeAlias
from typing import (
TYPE_CHECKING,
Final,
TypeAlias,
cast, # noqa: TID251 # db.tx is reached through untyped __getattr__ delegation
)
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
@ -25,6 +30,8 @@ from litellm.constants import (
)
if TYPE_CHECKING:
from prisma.client import TransactionManager
from litellm.proxy.utils import PrismaClient
SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs"
@ -116,6 +123,21 @@ def select_partitions_to_drop(partitions: list[tuple[str, datetime | None]], cut
return [name for name, upper in partitions if upper is not None and upper <= cutoff]
_TX_COMMIT_SLACK: Final = timedelta(seconds=5)
def _bounded_tx(prisma_client: "PrismaClient", timeout_ms: int) -> "TransactionManager":
"""
Open an interactive transaction that outlives the statement bound it
carries. prisma's default 5s transaction timeout would close it mid
lock-wait, after which the engine answers the next call with a 422.
"""
return cast( # cast-ok: PrismaWrapper delegates tx via __getattr__ (untyped)
"TransactionManager",
prisma_client.db.tx(timeout=timedelta(milliseconds=timeout_ms) + _TX_COMMIT_SLACK),
)
class SpendLogsPartitionManager:
def __init__(
self,
@ -137,7 +159,7 @@ class SpendLogsPartitionManager:
if budget_ms is None:
return False
try:
async with prisma_client.db.tx() as tx:
async with _bounded_tx(prisma_client, budget_ms) as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}")
rows: Final = await tx.query_raw(
"""
@ -172,7 +194,7 @@ class SpendLogsPartitionManager:
wait for the lock and statement_timeout bounds the work itself, so a
partition this run cannot get is simply left for the next one.
"""
async with prisma_client.db.tx() as tx:
async with _bounded_tx(prisma_client, timeout_ms) as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
await tx.execute_raw(statement)
@ -209,7 +231,7 @@ class SpendLogsPartitionManager:
async def _list_partitions(
self, prisma_client: "PrismaClient", timeout_ms: int
) -> list[tuple[str, datetime | None]]:
async with prisma_client.db.tx() as tx:
async with _bounded_tx(prisma_client, timeout_ms) as tx:
await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
rows: Final = await tx.query_raw(
"""

10
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union,
from urllib.parse import urlparse
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel
from pydantic import BaseModel, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@ -1202,7 +1202,13 @@ async def patch_guardrail(
litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True)
litellm_params_dict.update(requested_litellm_params)
merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict)
litellm_params = LitellmParams(**merged_litellm_params)
try:
litellm_params = LitellmParams(**merged_litellm_params)
except ValidationError as validation_error:
raise HTTPException(
status_code=422,
detail=f"Invalid guardrail configuration, update rejected: {validation_error}",
) from validation_error
# Update guardrail_info if provided
guardrail_info: Final = (

7
litellm/proxy/guardrails/guardrail_hooks/mcp_security/__init__.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Final, Literal, Optional, cast
from typing import TYPE_CHECKING, Final, Literal, Optional
import litellm
from litellm.proxy.guardrails.guardrail_hooks.mcp_security.mcp_security_guardrail import (
@ -20,10 +20,7 @@ def initialize_guardrail(
if not guardrail_name:
raise ValueError("MCP Security: guardrail_name is required")
on_violation: Final[Literal["block", "alert"]] = cast(
Literal["block", "alert"],
getattr(litellm_params, "on_violation", "block"),
)
on_violation: Final[Literal["block", "alert"]] = "block" if litellm_params.on_violation == "block" else "alert"
mcp_security_guardrail: Final = MCPSecurityGuardrail(
guardrail_name=guardrail_name,

1
litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py
View file

@ -21,6 +21,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None),
block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None),
)
litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback)

39
litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
View file

@ -93,6 +93,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
check_tool_results: bool | None = None,
file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
file_sanitization_fail_open: bool | None = None,
block_on_file_modify: bool | None = None,
**kwargs,
):
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
@ -124,6 +125,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
self.poll_interval = 2 # Seconds between polling attempts
self.file_sanitization_timeout = file_sanitization_timeout
self.file_sanitization_fail_open = file_sanitization_fail_open is not False
self.block_on_file_modify = block_on_file_modify is not False
super().__init__(**kwargs)
@ -372,13 +374,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
result = await self.sanitize_file_content(
file_data, filename, user_api_key_alias=user_api_key_alias
)
if result.get("action") == "block":
violations = result.get("violations", [])
raise HTTPException(
status_code=400,
detail=f"Image blocked by Prompt Security. Violations: {', '.join(violations)}",
)
self._raise_if_file_blocked(result, "Image")
except HTTPException:
raise
except Exception as e:
@ -408,7 +404,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
file_data: bytes,
filename: str,
user_api_key_alias: str | None = None,
) -> dict:
) -> _SanitizeResult:
"""
Sanitize file content using Prompt Security API.
Returns: dict with keys 'action', 'content', 'metadata'
@ -528,6 +524,17 @@ class PromptSecurityGuardrail(CustomGuardrail):
raise HTTPException(status_code=408, detail="File sanitization timeout")
def _raise_if_file_blocked(self, sanitization_result: _SanitizeResult, resource_name: str) -> None:
action: Final = sanitization_result.get("action")
if action != "block" and not (action == "modify" and self.block_on_file_modify):
return
violations: Final = sanitization_result.get("violations", ())
raise HTTPException(
status_code=400,
detail=f"{resource_name} blocked by Prompt Security. Violations: {', '.join(violations)}",
)
async def _process_image_url_item(self, item: dict, user_api_key_alias: str | None) -> dict:
"""Process and sanitize image_url items."""
image_url_data: Final = item.get("image_url", {})
@ -547,13 +554,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
file_data, filename, user_api_key_alias=user_api_key_alias
)
action: Final = sanitization_result.get("action")
if action == "block":
violations: Final = sanitization_result.get("violations", [])
raise HTTPException(
status_code=400,
detail=f"File blocked by Prompt Security. Violations: {', '.join(violations)}",
)
self._raise_if_file_blocked(sanitization_result, "File")
if action == "modify":
sanitized_content: Final = sanitization_result.get("content", "")
@ -615,13 +616,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
file_data, filename, user_api_key_alias=user_api_key_alias
)
action: Final = sanitization_result.get("action")
if action == "block":
violations: Final = sanitization_result.get("violations", [])
raise HTTPException(
status_code=400,
detail=f"Document blocked by Prompt Security. Violations: {', '.join(violations)}",
)
self._raise_if_file_blocked(sanitization_result, "Document")
if action == "modify":
sanitized_content: Final = sanitization_result.get("content", "")

6
litellm/proxy/hooks/litellm_skills/__init__.py
View file

@ -21,10 +21,7 @@ from litellm.llms.litellm_proxy.skills import (
code_execution_handler,
get_litellm_code_execution_tool,
)
from litellm.proxy.hooks.litellm_skills.main import (
SkillsInjectionHook,
skills_injection_hook,
)
from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook
__all__ = [
"LITELLM_CODE_EXECUTION_TOOL",
@ -35,5 +32,4 @@ __all__ = [
"SkillsSandboxExecutor",
"code_execution_handler",
"get_litellm_code_execution_tool",
"skills_injection_hook",
]

11
litellm/proxy/hooks/litellm_skills/main.py
View file

@ -29,6 +29,7 @@ import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Protocol
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
@ -493,7 +494,6 @@ class SkillsInjectionHook(CustomLogger):
Returns the final response with generated files inline.
"""
import litellm
from litellm.llms.litellm_proxy.skills.code_execution import (
LiteLLMInternalTools,
)
@ -723,7 +723,6 @@ print('No executable skill module found')
Returns the final response with generated files inline.
"""
import litellm
from litellm.llms.litellm_proxy.skills.code_execution import (
LiteLLMInternalTools,
)
@ -913,11 +912,3 @@ print('No executable skill module found')
verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to response", len(generated_files))
return response
# Global instance for registration
skills_injection_hook: Final = SkillsInjectionHook()
import litellm
litellm.logging_callback_manager.add_litellm_callback(skills_injection_hook)

26
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import batch_cost_is_final
from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
@ -37,6 +38,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
CallTypes,
LiteLLMBatch,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@ -248,6 +250,18 @@ class _ProxyDBLogger(CustomLogger):
)
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
if (
isinstance(completion_response, LiteLLMBatch)
and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
and not batch_cost_is_final(completion_response)
):
verbose_proxy_logger.debug(
"Cost tracking deferred for batch %s still in status %s",
completion_response.id,
completion_response.status,
)
await _release_budget_reservation(budget_reservation=budget_reservation)
return
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
@ -285,7 +299,7 @@ class _ProxyDBLogger(CustomLogger):
call_type=call_type,
):
## UPDATE DATABASE
await _update_database_and_spend_counters(
charged: Final = await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=increment_spend_counters,
user_api_key=user_api_key,
@ -302,6 +316,8 @@ class _ProxyDBLogger(CustomLogger):
request_tags=tags,
model_access_groups=model_access_groups,
)
if not charged:
return
# update cache (fire-and-forget for backward compat:
# cached object fields, soft budget alerts, etc.)
@ -578,9 +594,9 @@ async def _update_database_and_spend_counters(
budget_reservation: dict | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
) -> None:
) -> bool:
try:
await proxy_logging_obj.db_spend_update_writer.update_database(
charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
token=user_api_key,
response_cost=response_cost,
user_id=user_id,
@ -605,6 +621,9 @@ async def _update_database_and_spend_counters(
"Failed to invalidate budget reservation counters after release failed"
)
raise
if not charged:
await _release_budget_reservation(budget_reservation=budget_reservation)
return False
try:
await increment_spend_counters(
@ -630,6 +649,7 @@ async def _update_database_and_spend_counters(
finally:
budget_reservation["finalized"] = True
raise
return True
async def _release_budget_reservation(budget_reservation: dict | None) -> None:

40
litellm/proxy/litellm_pre_call_utils.py
View file

@ -2882,6 +2882,28 @@ def _add_guardrails_from_policies_in_metadata(
)
def add_guardrails_from_auth_metadata(
user_api_key_dict: UserAPIKeyAuth,
data: dict, # mutable-ok: writes guardrails into the live request dict, same contract as the helpers it wraps
metadata_variable_name: str,
) -> None:
"""Resolve key, team, and project guardrails, direct and via policies, onto the request metadata."""
_add_guardrails_from_key_or_team_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=user_api_key_dict.project_metadata,
data=data,
metadata_variable_name=metadata_variable_name,
)
_add_guardrails_from_policies_in_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=user_api_key_dict.project_metadata,
data=data,
metadata_variable_name=metadata_variable_name,
)
async def move_guardrails_to_metadata(
data: dict,
_metadata_variable_name: str,
@ -2914,22 +2936,8 @@ async def move_guardrails_to_metadata(
data.pop("policies", None)
return
# Check key/team/project-level guardrails
_add_guardrails_from_key_or_team_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=project_metadata,
data=data,
metadata_variable_name=_metadata_variable_name,
)
#########################################################################################
# Add guardrails from policies attached to key/team/project metadata
#########################################################################################
_add_guardrails_from_policies_in_metadata(
key_metadata=user_api_key_dict.metadata,
team_metadata=user_api_key_dict.team_metadata,
project_metadata=project_metadata,
add_guardrails_from_auth_metadata(
user_api_key_dict=user_api_key_dict,
data=data,
metadata_variable_name=_metadata_variable_name,
)

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

@ -645,6 +645,7 @@ async def get_auto_router_benchmarks(
str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to 30 days before end_date)")
] = None,
end_date: Annotated[str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to today)")] = None,
api_key: Annotated[str | None, Query(description="Filter to one virtual key token hash")] = None,
) -> AutoRouterBenchmarksResponse:
"""
Benchmarks for the auto-router dashboard: session shape, savings against the configured
@ -681,6 +682,7 @@ async def get_auto_router_benchmarks(
AUTOROUTER_BENCHMARKS_SQL,
start_day.isoformat(),
(end_day + timedelta(days=1)).isoformat(),
api_key,
)
rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ())
groups: Final = (

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

@ -196,6 +196,38 @@ class _ModelRowWhere(TypedDict):
model_id: ReadOnly[str]
class _KeyUpdateResult(TypedDict):
token: ReadOnly[str]
data: ReadOnly[Mapping[str, object]]
class _KeyRowWhere(TypedDict):
token: ReadOnly[str]
class _BudgetRowWhere(TypedDict):
budget_id: ReadOnly[str]
class _BudgetRowSoftBudgetUpdate(TypedDict):
soft_budget: ReadOnly[float | None]
updated_by: ReadOnly[str]
class _BudgetRowSoftBudgetCreate(TypedDict):
soft_budget: ReadOnly[float]
created_by: ReadOnly[str]
updated_by: ReadOnly[str]
class _KeyUpdateTx(Protocol):
@property
def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ...
@property
def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ...
class _ConfigTableActions(Protocol):
"""Config table surface this module needs; the shared repository seam exposes no ``update``."""
@ -1812,11 +1844,7 @@ async def generate_key_fn(
status_code=400,
detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
)
if data.soft_budget is not None and (not math.isfinite(data.soft_budget) or data.soft_budget < 0):
raise HTTPException(
status_code=400,
detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"},
)
_validate_soft_budget_value(data.soft_budget)
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
_custom_key_generate_hook(proxy_server)
@ -2121,6 +2149,88 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_
return non_default_values
def _validate_soft_budget_value(soft_budget: float | None) -> None:
if soft_budget is not None and (not math.isfinite(soft_budget) or soft_budget < 0):
raise HTTPException(
status_code=400,
detail={"error": f"soft_budget must be a non-negative finite number. Received: {soft_budget}"},
)
async def _update_key_soft_budget(
db: _KeyUpdateTx,
existing_key_row: LiteLLM_VerificationToken,
soft_budget: float | None,
changed_by: str,
) -> str | None:
existing_budget_id: Final = existing_key_row.budget_id
if existing_budget_id is not None:
budget_update: Final[_BudgetRowSoftBudgetUpdate] = {"soft_budget": soft_budget, "updated_by": changed_by}
budget_where: Final[_BudgetRowWhere] = {"budget_id": existing_budget_id}
await db.litellm_budgettable.update(where=budget_where, data=budget_update)
return existing_budget_id
if soft_budget is None:
return None
budget_create: Final[_BudgetRowSoftBudgetCreate] = {
"soft_budget": soft_budget,
"created_by": changed_by,
"updated_by": changed_by,
}
created_budget: Final = await db.litellm_budgettable.create(data=budget_create)
return created_budget.budget_id
async def _apply_soft_budget_update(
data: UpdateKeyRequest,
non_default_values: Mapping[str, object],
db: _KeyUpdateTx,
existing_key_row: LiteLLM_VerificationToken,
changed_by: str,
) -> Mapping[str, object]:
remaining: Final = MappingProxyType({k: v for k, v in non_default_values.items() if k != "soft_budget"})
updated_budget_id: Final = await _update_key_soft_budget(
db=db,
existing_key_row=existing_key_row,
soft_budget=data.soft_budget,
changed_by=changed_by,
)
if updated_budget_id is not None and existing_key_row.budget_id is None:
return MappingProxyType({**remaining, "budget_id": updated_budget_id})
return remaining
async def _update_key_row_with_soft_budget(
prisma_client: PrismaClient,
key: str,
data: UpdateKeyRequest,
non_default_values: Mapping[str, object],
existing_key_row: LiteLLM_VerificationToken,
changed_by: str,
) -> _KeyUpdateResult:
hashed_token: Final = _hash_token_if_needed(key)
key_where: Final[_KeyRowWhere] = {"token": hashed_token}
tx: _KeyUpdateTx
async with prisma_client.tx() as tx:
update_values: Final = await _apply_soft_budget_update(
data=data,
non_default_values=non_default_values,
db=tx,
existing_key_row=existing_key_row,
changed_by=changed_by,
)
updated_row: Final = await tx.litellm_verificationtoken.update(
where=key_where,
data=with_settings_updated_at(
prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token}))
),
)
updated_data: Final[Mapping[str, object]] = (
updated_row.model_dump() if updated_row is not None else MappingProxyType({})
)
result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data}
return result
async def prepare_key_update_data(
data: UpdateKeyRequest | RegenerateKeyRequest,
existing_key_row: LiteLLM_VerificationToken,
@ -2659,6 +2769,7 @@ async def _validate_update_key_data(
(data.max_budget is not None and data.max_budget != existing_key_row.max_budget)
or data.spend is not None
or "budget_limits" in data.model_fields_set
or "soft_budget" in data.model_fields_set
)
_existing_metadata: Final = getattr(existing_key_row, "metadata", None)
@ -2840,6 +2951,11 @@ async def update_key_fn(
"""
Update an existing API key's parameters.
The body is a merge patch: a field left out keeps its stored value, and on the key's own columns
an explicit null clears it. The metadata-backed fields below are the exception, merging into the
stored metadata instead: passing one as null leaves it unchanged, while `metadata` itself
replaces the stored metadata wholesale.
Parameters:
- key: Optional[str] - The key to update. Either key or key_alias must be provided.
- key_alias: Optional[str] - User-friendly key alias. If key is omitted, also identifies the key to update (must match exactly one key, same as /key/delete's key_aliases)
@ -2857,7 +2973,7 @@ async def update_key_fn(
- model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
- budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
- budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.)
- soft_budget: Optional[float] - [TODO] Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached.
- soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget.
- max_parallel_requests: Optional[int] - Rate limit for parallel requests
- metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"}
- tpm_limit: Optional[int] - Tokens per minute limit
@ -2913,6 +3029,7 @@ async def update_key_fn(
"""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
@ -2928,6 +3045,8 @@ async def update_key_fn(
detail={"error": f"max_budget must be a non-negative finite number. Received: {data.max_budget}"},
)
_validate_soft_budget_value(data.soft_budget)
# get the row from db
existing_key_row: Final = await _get_and_validate_existing_key(
token=data.key,
@ -2984,10 +3103,22 @@ async def update_key_fn(
existing_key_alias=existing_key_row.key_alias,
)
_data: Final = {**non_default_values, "token": key}
if prisma_client is None:
raise Exception("Not connected to DB!")
response: Final = await prisma_client.update_data(token=key, data=_data)
changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
response: Final = (
await _update_key_row_with_soft_budget(
prisma_client=prisma_client,
key=key,
data=data,
non_default_values=non_default_values,
existing_key_row=existing_key_row,
changed_by=changed_by,
)
if "soft_budget" in data.model_fields_set
else await prisma_client.update_data(token=key, data=MappingProxyType({**non_default_values, "token": key}))
)
# Delete - key from cache, since it's been updated!
# key updated - a new model could have been added to this key. it should not block requests after this is done
@ -6427,7 +6558,7 @@ async def _list_key_helper(
{"token": "desc"}, # fallback sort
]
),
include={"object_permission": True},
include={"object_permission": True, "litellm_budget_table": True},
)
verbose_proxy_logger.debug("Fetched %s keys", len(keys))

122
litellm/proxy/middleware/per_request_root_path_middleware.py Normal file
View file

@ -0,0 +1,122 @@
"""Per-request ``root_path`` resolution from ``SERVER_ROOT_PATHS``.
``SERVER_ROOT_PATH`` is a startup scalar, so one deployment serves exactly one
client-visible URL path prefix; a request under any other prefix 404s before a
handler runs. When the ingress preserves several prefixes into one pod (e.g.
``/tenant-a/*`` and ``/tenant-b/*``), the matched prefix becomes that
request's ``scope["root_path"]`` instead: Starlette strips it during route
matching and rebuilds it into ``request.base_url``, so every emitted URL —
the MCP OAuth discovery ``resource`` (RFC 9728 §3) and the 401 challenges'
``resource_metadata`` among them — lands under the prefix the client called.
Opt-in: with ``SERVER_ROOT_PATHS`` unset the middleware is not added at all.
"""
import os
from collections.abc import Sequence
from contextvars import ContextVar
from typing import Final
from starlette.types import ASGIApp, Receive, Scope, Send
from litellm._logging import verbose_proxy_logger
SERVER_ROOT_PATHS_ENV: Final = "SERVER_ROOT_PATHS"
# The effective ``root_path`` for the currently-handled request. Populated by
# ``PerRequestRootPathMiddleware`` from the (possibly-mutated) scope so code
# that emits URLs off the request path — the 401 challenges' resource_metadata
# and ``get_custom_url``'s SSO callbacks among them — can pick up the prefix
# the client actually called without threading scope through every call site.
# ``None`` means "middleware did not run" (the ``SERVER_ROOT_PATHS`` env is
# unset, so no per-request prefix exists); readers fall back to the scalar
# ``SERVER_ROOT_PATH`` in that case, which matches the pre-middleware behavior.
_request_root_path_var: Final[ContextVar[str | None]] = ContextVar("_request_root_path_var", default=None)
def get_request_root_path() -> str:
"""Return the effective ``root_path`` for the current request.
Reads the value ``PerRequestRootPathMiddleware`` stashed for this request;
falls back through :func:`~litellm.proxy.utils.get_server_root_path` (i.e.
the ``SERVER_ROOT_PATH`` env) when the middleware did not run — the
scalar-only deployment. Delegating to the existing helper keeps every
existing ``monkeypatch.setattr("litellm.proxy.utils.get_server_root_path"``
test override working, and keeps a single source of truth for the scalar.
"""
value: Final = _request_root_path_var.get()
if value is not None:
return value
# Lazy import: utils.py imports this module (via the lazy import inside
# get_custom_url), so a top-level import would build a cycle at load time.
from litellm.proxy.utils import get_server_root_path # noqa: PLC0415 # lazy import breaks a two-way dep
return get_server_root_path()
def normalize_root_paths(raw_paths: Sequence[str]) -> tuple[str, ...]:
"""Strip whitespace and trailing slashes, dedupe, order longest-first;
warn and drop entries missing a leading ``/`` and the bare root."""
kept: Final[list[str]] = [] # mutable-ok: local accumulator; escapes only as a tuple
for entry in raw_paths:
candidate = entry.strip()
if not candidate:
continue
if not candidate.startswith("/"):
verbose_proxy_logger.warning(
"%s entry %r does not start with '/' and will be ignored.",
SERVER_ROOT_PATHS_ENV,
entry,
)
continue
candidate = candidate.rstrip("/")
if not candidate:
verbose_proxy_logger.warning(
"%s entry %r is the bare root and will be ignored; a root-mounted deployment needs no entry.",
SERVER_ROOT_PATHS_ENV,
entry,
)
continue
if candidate not in kept:
kept.append(candidate)
return tuple(sorted(kept, key=len, reverse=True))
def get_server_root_paths() -> tuple[str, ...]:
"""The normalized ``SERVER_ROOT_PATHS`` prefixes, empty when unset."""
configured: Final = os.getenv(SERVER_ROOT_PATHS_ENV, "")
if not configured.strip():
return ()
return normalize_root_paths(configured.split(","))
class PerRequestRootPathMiddleware:
"""Sets ``scope["root_path"]`` to the configured prefix matching the
request path on a whole-segment boundary. ``scope["path"]`` is left
untouched (Starlette strips ``root_path`` at route-match time). Must be
the outermost middleware so inner middlewares and the router see the
resolved value; a matched prefix overrides a scalar ``SERVER_ROOT_PATH``
for that request.
"""
def __init__(self, app: ASGIApp, root_paths: Sequence[str]) -> None:
self.app = app
self.root_paths: Final = normalize_root_paths(root_paths)
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] in ("http", "websocket"):
path: Final = scope.get("path", "")
for prefix in self.root_paths:
if path == prefix or path.startswith(prefix + "/"):
scope["root_path"] = prefix # rebind-ok: ASGI middleware contract; Router and base_url read it
break
# Stash the effective root_path (matched prefix, or the scope's
# existing value when nothing matched — i.e. FastAPI's scalar
# SERVER_ROOT_PATH) so code that emits URLs off the request path
# picks the same prefix the router will resolve the request under.
token: Final = _request_root_path_var.set(str(scope.get("root_path", "")))
try:
await self.app(scope, receive, send)
finally:
_request_root_path_var.reset(token)
return
await self.app(scope, receive, send)

8
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -15,6 +15,7 @@ from typing import (
runtime_checkable,
)
from litellm.batches.batch_utils import batch_cost_is_final
from litellm.proxy._types import ProxyException
from litellm.repositories.table_repositories import (
ManagedFileRepository,
@ -1357,12 +1358,7 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool:
enumerated the batch and none succeeded. A zero or unknown total means counts
are unreported, so stay eligible and let the next poller pass revisit it. (#37713)
"""
if response.output_file_id is not None:
return True
request_counts = response.request_counts
if request_counts is None:
return False
return request_counts.total > 0 and request_counts.completed == 0
return batch_cost_is_final(response)
async def update_batch_in_database(

7
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -712,6 +712,10 @@ BEDROCK_ENDPOINT_ACTIONS: Final = {
BEDROCK_STREAMING_ACTIONS: Final = {"invoke-with-response-stream", "converse-stream"}
def is_bedrock_count_tokens_endpoint(endpoint: str) -> bool:
return "count_tokens" in endpoint or "count-tokens" in endpoint
def _extract_model_from_bedrock_endpoint(endpoint: str) -> str:
"""
Extract model name from Bedrock endpoint path.
@ -986,8 +990,7 @@ async def bedrock_llm_proxy_route(
request_body: Final = await _read_request_body(request=request)
# Special handling for count_tokens endpoints
if "count_tokens" in endpoint or "count-tokens" in endpoint:
if is_bedrock_count_tokens_endpoint(endpoint):
return await handle_bedrock_count_tokens(
endpoint=endpoint,
request=request,

2
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -322,7 +322,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
def get_response_headers(
headers: httpx.Headers,
litellm_call_id: str | None = None,
custom_headers: dict | None = None,
custom_headers: Mapping[str, str] | None = None,
) -> dict:
# Exclude headers that uvicorn writes itself (server, date) and
# encoding/length headers that don't survive re-serialization.

20
litellm/proxy/proxy_server.py
View file

@ -598,6 +598,10 @@ from litellm.proxy.middleware.admission_control_middleware import (
from litellm.proxy.middleware.in_flight_requests_middleware import (
InFlightRequestsMiddleware,
)
from litellm.proxy.middleware.per_request_root_path_middleware import (
PerRequestRootPathMiddleware,
get_server_root_paths,
)
from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
from litellm.proxy.middleware.request_size_limit_middleware import (
RequestSizeLimitMiddleware,
@ -18384,6 +18388,22 @@ app.add_middleware(
get_settings=lambda: get_admission_control_settings(general_settings),
state=admission_control_state,
)
# Added last on purpose - last-added is outermost, and the client-visible URL
# prefix must be resolved into scope["root_path"] before any inner middleware
# or the router inspects the path. Only added when SERVER_ROOT_PATHS is
# configured, so the default deployment's middleware stack is unchanged.
_server_root_paths: Final = get_server_root_paths()
if _server_root_paths:
if server_root_path and server_root_path != "/":
verbose_proxy_logger.warning(
"Both SERVER_ROOT_PATH=%r and SERVER_ROOT_PATHS=%r are set. A request "
"matching a SERVER_ROOT_PATHS prefix overrides the scalar root_path for "
"that request; unmatched requests keep SERVER_ROOT_PATH. Configure one "
"mechanism or the other.",
server_root_path,
_server_root_paths,
)
app.add_middleware(PerRequestRootPathMiddleware, root_paths=_server_root_paths)
async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse":

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

@ -688,6 +688,13 @@
],
"default_model_placeholder": "gpt-3.5-turbo"
},
{
"provider": "CHATGPT",
"provider_display_name": "ChatGPT Subscription",
"litellm_provider": "chatgpt",
"credential_fields": [],
"default_model_placeholder": "chatgpt/gpt-5.4"
},
{
"provider": "CLARIFAI",
"provider_display_name": "Clarifai",

37
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -1389,12 +1389,15 @@ def _approximate_input_size(request_body: Mapping[str, object]) -> int:
def _count_input_tokens(request_body: dict, model: str) -> int | None:
try:
if "messages" in request_body:
return litellm.token_counter(
model=model,
messages=request_body.get("messages") or [],
tools=request_body.get("tools"),
tool_choice=request_body.get("tool_choice"),
)
try:
return litellm.token_counter(
model=model,
messages=request_body.get("messages") or (),
tools=request_body.get("tools"),
tool_choice=request_body.get("tool_choice"),
)
except ValueError:
return _count_text_tokens(model=model, text=request_body.get("messages"))
if "prompt" in request_body:
return _count_text_tokens(model=model, text=request_body.get("prompt"))
if "input" in request_body:
@ -1442,11 +1445,7 @@ def _estimate_output_tokens(
if _is_input_only_route(route=route):
return 0
requested: int | None = None
for key in ("max_completion_tokens", "max_tokens", "max_output_tokens"):
requested = _to_int(request_body.get(key))
if requested is not None:
break
requested: Final = _requested_output_tokens(request_body)
# Clamp at min(requested-or-default, model_max-or-default). Two purposes:
# (1) Without an explicit cap we still need a finite reservation so the
@ -1457,9 +1456,19 @@ def _estimate_output_tokens(
# at the cap — the model can only physically emit max_output_tokens
# anyway, so reserving more is both wasteful and a DoS surface.
model_ceiling: Final = _to_int(model_info.get("max_output_tokens")) or DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
if requested is None:
requested = DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK
return min(requested, model_ceiling)
return min(DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK if requested is None else requested, model_ceiling)
_OUTPUT_TOKEN_FIELDS: Final = ("max_completion_tokens", "max_tokens", "max_output_tokens")
def _requested_output_tokens(request_body: Mapping[str, object]) -> int | None:
inference_config: Final = request_body.get("inferenceConfig")
candidates: Final = (
*(request_body.get(field) for field in _OUTPUT_TOKEN_FIELDS),
inference_config.get("maxTokens") if isinstance(inference_config, Mapping) else None,
)
return next((tokens for tokens in map(_to_int, candidates) if tokens is not None), None)
def _count_text_tokens(model: str, text: object) -> int:

105
litellm/proxy/spend_tracking/savings.py
View file

@ -9,12 +9,17 @@ have been aggregated across models.
"""
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Final, NamedTuple
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_cost_per_unit, generic_cost_per_token
from litellm.litellm_core_utils.llm_cost_calc.utils import (
_get_cost_per_unit,
calculate_prompt_caching_savings,
generic_cost_per_token,
)
from litellm.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@ -32,42 +37,13 @@ class SavingsSpend(NamedTuple):
gateway_injected_caching: float = 0.0
def _input_cache_read_and_write_cost(info: ModelInfo | None) -> tuple[float, float, float]:
"""
Return ``(input_cost, cache_read_cost, cache_write_cost)`` per token.
``info`` is whatever pricing the caller resolved -- deployment rates when the
request came through a router deployment, public rates otherwise -- so a
negotiated price is honoured here rather than silently replaced by the list rate.
``None`` falls open to ``(0.0, 0.0, 0.0)`` so savings degrade to zero rather than
raising inside the spend writer.
Prices are read through ``_get_cost_per_unit``, the same accessor the cost
calculator uses, which coerces the string prices a ``config.yaml`` can produce
(``"3e-7"``) and resolves service-tier suffixes.
An absent cache price mirrors the input cost, which yields a zero discount on the
read leg and a zero premium on the write leg. Mirroring rather than taking
``_get_cost_per_unit``'s 0.0 default is load-bearing on the write leg: a zero write
price would make the premium ``0 - input_cost``, turning a model that simply has no
write pricing into a spurious extra saving.
The two legs then differ on an explicit ``0.0``, and the asymmetry is deliberate. A
free cache *write* does not exist -- entries carrying a literal zero (``deepseek-chat``
does) mean "no separate price", so a falsy write price also mirrors input. A free
cache *read* is real: 15 models charge for input and serve reads for nothing, which
is the largest discount available, so the read leg keeps its literal zero.
"""
if info is None:
return 0.0, 0.0, 0.0
input_cost: Final = _get_cost_per_unit(info, "input_cost_per_token") or 0.0
cache_read_cost: Final = _get_cost_per_unit(info, "cache_read_input_token_cost", default_value=None)
cache_write_cost: Final = _get_cost_per_unit(info, "cache_creation_input_token_cost", default_value=None)
return (
input_cost,
input_cost if cache_read_cost is None else cache_read_cost,
cache_write_cost if cache_write_cost else input_cost,
)
def _coerce_billed_at(value: datetime | str | None) -> datetime | None:
if isinstance(value, datetime) or value is None:
return value
try:
return datetime.fromisoformat(value.replace("Z", "+00:00"))
except ValueError:
return None
class _ModelIdentity(NamedTuple):
@ -586,6 +562,7 @@ def compute_savings_spend(
llm_router: "Callable[[], Router | None] | None" = None,
cost_breakdown: Mapping[str, object] | None = None,
recorded_autorouter_savings: object = None,
billed_at: datetime | str | None = None,
) -> SavingsSpend:
"""
Dollar savings for one request, split by optimization driver.
@ -595,24 +572,10 @@ def compute_savings_spend(
premium paid to write those entries, both derived here from ``usage_object`` so no
caller can hand in a count that disagrees with the usage record.
The net form follows from what the request would have cost with caching off. The
provider reports ``prompt_tokens`` as the inclusive total of three disjoint
partitions (uncached text, cache reads, cache writes), so an uncached counterfactual
bills every one of those tokens at the flat input rate::
would_have_cost = (text + reads + writes) * input
actually_cost = text * input + reads * read_rate + writes * write_rate
savings = reads * (input - read_rate) - writes * (write_rate - input)
So the write leg subtracts the write PREMIUM, not the whole write cost: those tokens
had to be sent either way, and the counterfactual already pays the input rate for
them. The premium stays signed, because a handful of models price writes below their
input rate and there the write is a genuine extra saving.
A request that only writes cache and gets no hits therefore reports negative savings,
which is accurate: it really did cost more than the uncached call would have. The
daily rollup increments arithmetically, so those rows offset positive ones in the
same bucket.
The uncached counterfactual pays the ordinary input rate for the same prompt size
and tier. Cache writes subtract only the premium over that rate, split by TTL.
Savings stay signed: a write-only request can lose money, and daily rollups net
those losses against read savings.
Caching is reported twice. ``prompt_caching`` is every net dollar caching saved,
whoever caused it, which is what a customer means by "what did caching save me".
@ -638,12 +601,9 @@ def compute_savings_spend(
calls this and only auto-routed ones need one, so looking it up eagerly at the call
site would fetch and discard it on the rest.
``cost_breakdown`` is what the cost calculator recorded for this request, and it
carries both what the request really cost and the tier and region it was priced on.
Only the auto-router driver reads it. Compression and prompt caching price a
hypothetical token delta off flat rate keys, so they are blind to tiered pricing in
the same way; that is pre-existing behaviour on two shipped drivers rather than
something introduced here, and moving those numbers is its own change.
``cost_breakdown`` supplies the biller's tier and region to caching and auto-router
savings. Caching also uses the logged prompt size and TTL split. Compression retains
its flat input-rate estimate; changing that counterfactual is a separate concern.
``recorded_autorouter_savings`` is the figure the logging path stamped on the spend
log's metadata, honoured over recomputation so the rollup, the turn table and the
@ -658,13 +618,24 @@ def compute_savings_spend(
pricing: Final = _effective_model_info(router_instance, model_id, model or "") or (
_model_info(identity) if identity else None
)
input_cost, cache_read_cost, cache_write_cost = _input_cache_read_and_write_cost(pricing)
input_cost: Final = (_get_cost_per_unit(pricing, "input_cost_per_token") or 0.0) if pricing else 0.0
compression: Final = max(compression_saved_tokens, 0) * input_cost
cache_read_input_tokens: Final = extract_cache_read_tokens(usage_object)
cache_creation_input_tokens: Final = extract_cache_creation_tokens(usage_object)
read_discount: Final = max(cache_read_input_tokens, 0) * max(input_cost - cache_read_cost, 0.0)
write_premium: Final = max(cache_creation_input_tokens, 0) * (cache_write_cost - input_cost)
prompt_caching: Final = read_discount - write_premium
usage: Final = _usage_from_spend_log(usage_object)
basis: Final = _pricing_basis(cost_breakdown)
billed_at_datetime: Final = _coerce_billed_at(billed_at)
prompt_caching: Final = (
calculate_prompt_caching_savings(
model_info=pricing,
usage=usage,
custom_llm_provider=identity.provider if identity else custom_llm_provider,
service_tier=basis.service_tier,
data_residency=basis.data_residency,
vertex_location=basis.vertex_location,
billed_at=billed_at_datetime,
)
if pricing is not None and usage is not None
else 0.0
)
gateway_injected_caching: Final = prompt_caching if gateway_injected_cache else 0.0
# The figure the logging path recorded wins, before the usage gate on purpose: a row

24
litellm/proxy/utils.py
View file

@ -152,7 +152,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
from litellm.proxy.hooks.sensitive_data_routing import (
_PROXY_SensitiveDataRoutingHandler,
)
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.repositories.budget_repository import BudgetRepository
@ -924,7 +924,13 @@ class ProxyLogging:
"incoming_bearer_token": kwargs.get("incoming_bearer_token"),
"metadata": {"headers": kwargs.get("headers") or {}},
}
user_api_key_auth: Final = kwargs.get("user_api_key_auth")
if isinstance(user_api_key_auth, UserAPIKeyAuth):
add_guardrails_from_auth_metadata(
user_api_key_dict=user_api_key_auth,
data=synthetic_data,
metadata_variable_name="metadata",
)
return synthetic_data
def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None:
@ -7206,7 +7212,19 @@ def get_custom_url(request_base_url: str, route: str | None = None) -> str:
else:
base_url = request_base_url
server_root_path: Final = get_server_root_path()
# get_request_root_path() returns the prefix the router is actually
# resolving this request under: the matched SERVER_ROOT_PATHS entry when
# PerRequestRootPathMiddleware ran, otherwise the SERVER_ROOT_PATH scalar.
# This keeps the emitted URL under one prefix — the one the client called —
# instead of stacking the scalar onto a request already living under a
# dynamic prefix (which would produce /tenant-a/legacy/... — a path that
# doesn't exist). join_paths()'s tail-dedup then collapses the append when
# base_url (i.e. request.base_url) already ends in the same prefix.
from litellm.proxy.middleware.per_request_root_path_middleware import ( # noqa: PLC0415 # lazy: middleware imports utils
get_request_root_path,
)
server_root_path: Final = get_request_root_path()
if route is not None:
if server_root_path != "":
# First join base_url with server_root_path, then with route

9
litellm/proxy/vector_store_endpoints/utils.py
View file

@ -161,6 +161,8 @@ async def can_user_access_vector_store(
this vector store id.
5. The caller's team_id matches the vector store's team_id.
A dashboard session credential is evaluated against the same effective
contexts as listing (its own grants plus each real team of the user).
Otherwise access is denied.
"""
if _is_proxy_admin(user_api_key_dict):
@ -169,7 +171,8 @@ async def can_user_access_vector_store(
if vector_store.get("team_id") is None:
return True
return await _is_vector_store_granted(vector_store, user_api_key_dict)
auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
return await _is_vector_store_granted_to_any(vector_store, auth_contexts)
async def _is_vector_store_granted(
@ -219,7 +222,7 @@ async def _team_auth_context(team_id: str, user_api_key_dict: UserAPIKeyAuth) ->
)
async def _vector_store_listing_auth_contexts(
async def _vector_store_auth_contexts(
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[UserAPIKeyAuth, ...]:
if not is_ui_session_credential(user_api_key_dict):
@ -250,7 +253,7 @@ async def filter_listable_vector_stores(
if _is_proxy_admin(user_api_key_dict):
return tuple(vector_stores)
auth_contexts: Final = await _vector_store_listing_auth_contexts(user_api_key_dict)
auth_contexts: Final = await _vector_store_auth_contexts(user_api_key_dict)
return tuple([vs for vs in vector_stores if await _is_vector_store_granted_to_any(vs, auth_contexts)])

3
litellm/repositories/model_repository.py
View file

@ -4,7 +4,6 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
import json
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol
from litellm.models.model import LiteLLM_ProxyModelTable
@ -109,7 +108,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]:
"""Find every model except the row currently being updated."""
records: Final = await self.table.find_many(
where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})})
where={"model_id": {"not": model_id}} # mutable-ok: Prisma requires plain dicts for query serialization
)
return tuple(self._to_model_list(records))

25
litellm/responses/utils.py
View file

@ -25,9 +25,15 @@ from litellm.types.utils import (
PromptTokensDetailsWrapper,
SpecialEnums,
Usage,
text_tokens_without_nested_reasoning,
)
def _output_token_detail(details: object, field: str) -> int | None:
value: Final = getattr(details, field, None)
return value if isinstance(value, int) else None
def _is_object_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a list is a Sequence of anything
return isinstance(value, list)
@ -1137,11 +1143,22 @@ class ResponseAPILoggingUtils:
response_api_usage, "output_tokens_details", None
)
if output_tokens_details:
reasoning_tokens: Final = _output_token_detail(output_tokens_details, "reasoning_tokens")
image_tokens: Final = _output_token_detail(output_tokens_details, "image_tokens")
audio_tokens: Final = _output_token_detail(output_tokens_details, "audio_tokens")
reported_text_tokens: Final = _output_token_detail(output_tokens_details, "text_tokens")
completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None),
image_tokens=getattr(output_tokens_details, "image_tokens", None),
text_tokens=getattr(output_tokens_details, "text_tokens", None),
audio_tokens=getattr(output_tokens_details, "audio_tokens", None),
reasoning_tokens=reasoning_tokens,
image_tokens=image_tokens,
text_tokens=None
if reported_text_tokens is None
else text_tokens_without_nested_reasoning(
completion_tokens=completion_tokens,
text_tokens=reported_text_tokens,
reasoning_tokens=reasoning_tokens or 0,
other_modality_tokens=(audio_tokens or 0) + (image_tokens or 0),
),
audio_tokens=audio_tokens,
)
extra_usage_fields: Final = {

152
litellm/router.py
View file

@ -403,6 +403,10 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]:
"""
Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still
@ -612,6 +616,38 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(
RETRY_BREADCRUMB_LIMIT: Final = 4
class FallbackAwareStreamWrapper(CustomStreamWrapper):
"""Base for the Router's chat-completion stream wrappers, which are built around the
attempt the Router picked first and have to repoint themselves when a fallback takes over."""
fallback_headers_adopted: bool = False
def adopt_fallback_response_headers(
self,
fallback_response: object,
prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]],
) -> None:
"""Repoint this wrapper at the deployment that served the stream.
Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit
counters, `model_id` and `api_base` cannot reach the proxy's response headers or
its callbacks.
"""
self._response_headers = getattr(fallback_response, "_response_headers", None)
fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params
if fallback_hidden_params:
self._hidden_params = { # mutable-ok: the rest of litellm writes into _hidden_params
**fallback_hidden_params,
# dict() because add_retry_fallback_headers mutates additional_headers in place
"additional_headers": dict(fallback_headers), # mutable-ok: see above
}
self._base_hidden_params = { # mutable-ok: CustomStreamWrapper keeps this snapshot as a dict
**self._hidden_params,
"response_cost": None,
}
self.fallback_headers_adopted = True
class Router:
model_names: set = set()
cache_responses: bool | None = False
@ -2572,6 +2608,18 @@ class Router:
return fallback_hidden_params, {}
return fallback_hidden_params, cast("dict[str, object]", fallback_headers)
@staticmethod
def _adopt_fallback_response_headers(
wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]",
fallback_response: object,
) -> tuple[dict[str, object], dict[str, object]]:
"""Repoint the wrapper at `fallback_response`, returning its prepared hidden params."""
prepared: Final = Router._prepare_fallback_hidden_params(fallback_response)
adopting_wrapper: Final = wrapper_ref()
if adopting_wrapper is not None:
adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared)
return prepared
@staticmethod
def _apply_fallback_hidden_params_to_item(
fallback_item: object,
@ -2611,7 +2659,7 @@ class Router:
held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack()
class FallbackStreamWrapper(CustomStreamWrapper):
class FallbackStreamWrapper(FallbackAwareStreamWrapper):
def __init__(self, async_generator: AsyncGenerator):
# Copy attributes from the original model_response
super().__init__(
@ -2619,6 +2667,7 @@ class Router:
model=model_response.model,
custom_llm_provider=model_response.custom_llm_provider,
logging_obj=model_response.logging_obj,
_response_headers=getattr(model_response, "_response_headers", None),
)
self._async_generator = async_generator
inner_chunks: Final[object] = getattr(model_response, "chunks", None)
@ -2695,8 +2744,17 @@ class Router:
# If fallback returns a streaming response, iterate over it
if hasattr(fallback_response, "__aiter__"):
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
wrapper_ref, fallback_response
)
fallback_headers_are_settled = False
async for fallback_item in fallback_response:
if not fallback_headers_are_settled:
fallback_headers_are_settled = True # rebind-ok: one-shot latch
# a fallback that failed over again only repoints itself once it yields
prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields
Router._adopt_fallback_response_headers(wrapper_ref, fallback_response)
)
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
if (
fallback_item
@ -2738,7 +2796,11 @@ class Router:
e,
)
return FallbackStreamWrapper(stream_with_fallbacks())
wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks())
# weak, so the generator closing over it does not keep the wrapper out of
# refcount teardown and delay the `finally` that releases the deployment slot
wrapper_ref: Final = weakref.ref(wrapped_response)
return wrapped_response
@staticmethod
def _extract_partial_responses_usage(
@ -3167,13 +3229,14 @@ class Router:
"""
from litellm.exceptions import MidStreamFallbackError
class SyncFallbackStreamWrapper(CustomStreamWrapper):
class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper):
def __init__(self, sync_generator: Generator):
super().__init__(
completion_stream=sync_generator,
model=model_response.model,
custom_llm_provider=model_response.custom_llm_provider,
logging_obj=model_response.logging_obj,
_response_headers=getattr(model_response, "_response_headers", None),
)
self._sync_generator = sync_generator
if hasattr(model_response, "_hidden_params"):
@ -3229,8 +3292,17 @@ class Router:
)
if hasattr(fallback_response, "__iter__"):
prepared_fallback_hidden_params = Router._prepare_fallback_hidden_params(fallback_response)
prepared_fallback_hidden_params = Router._adopt_fallback_response_headers(
wrapper_ref, fallback_response
)
fallback_headers_are_settled = False
for fallback_item in fallback_response:
if not fallback_headers_are_settled:
fallback_headers_are_settled = True # rebind-ok: one-shot latch
# a fallback that failed over again only repoints itself once it yields
prepared_fallback_hidden_params = ( # rebind-ok: re-read once the fallback yields
Router._adopt_fallback_response_headers(wrapper_ref, fallback_response)
)
Router._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params)
if (
fallback_item
@ -3268,7 +3340,10 @@ class Router:
close_err,
)
return SyncFallbackStreamWrapper(stream_with_fallbacks())
wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks())
# weak, for the same reason as the async twin
wrapper_ref: Final = weakref.ref(wrapped_response)
return wrapped_response
async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
"""
@ -4449,7 +4524,7 @@ class Router:
self.fail_calls[model_name] += 1
raise e
async def aspeech(self, model: str, input: str, voice: str, **kwargs):
async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
"""
Example Usage:
@ -4501,7 +4576,7 @@ class Router:
)
raise e
async def _aspeech(self, model: str, input: str, voice: str, **kwargs):
async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
model_name: Final = model
try:
verbose_router_logger.debug("Inside _aspeech()- model: %s; kwargs: %s", model, kwargs)
@ -4525,7 +4600,7 @@ class Router:
**{
**data,
"input": input,
"voice": voice,
"voice": data.get("voice") if voice is None else voice,
"client": model_client,
**kwargs,
}
@ -7458,6 +7533,21 @@ class Router:
Context_Policy_Fallbacks={content_policy_fallbacks}",
)
@staticmethod
def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]:
failed_deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None)
status_code: Final = getattr(exception, "status_code", None)
if not failed_deployment_id or not isinstance(status_code, int):
return ()
if litellm._should_retry(status_code): # pyright: ignore[reportPrivateUsage] # as in should_retry_this_error
return ()
already_skipped_ids: Final = _as_retry_skipped_deployment_ids(already_skipped)
skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id))))
verbose_router_logger.debug(
"Retry skips deployments that already answered %s to this request: %s", status_code, skipped
)
return skipped
@tracer.wrap()
async def async_function_with_retries(self, *args, **kwargs):
verbose_router_logger.debug("Inside async function with retries.")
@ -7553,6 +7643,12 @@ class Router:
## LOGGING
if num_retries > 0:
kwargs = self.log_retry(kwargs=kwargs, e=original_exception)
first_skipped_ids: Final = self._deployment_ids_to_skip_on_retry(
exception=original_exception,
already_skipped=kwargs.get("_retry_skipped_deployment_ids"),
)
if first_skipped_ids:
kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it
else:
raise
@ -7622,6 +7718,12 @@ class Router:
except Exception:
raise e
skipped_ids = self._deployment_ids_to_skip_on_retry(
exception=e,
already_skipped=kwargs.get("_retry_skipped_deployment_ids"),
)
if skipped_ids:
kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it
_timeout = self._time_to_sleep_before_retry(
e=e,
remaining_retries=remaining_retries,
@ -12454,7 +12556,7 @@ class Router:
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order: Final = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils._get_order_filtered_deployments(
healthy_deployments = litellm.utils.get_order_filtered_deployments(
cast(list[dict], healthy_deployments), target_order=_target_order
)
@ -12462,11 +12564,24 @@ class Router:
## this request via weighted-failover. Always honored, regardless of the
## router-level flag, so a stale exclusion key on kwargs cannot escape.
_excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None)
healthy_deployments = litellm.utils._get_excluded_filtered_deployments(
healthy_deployments = litellm.utils.get_excluded_filtered_deployments(
cast(list[dict], healthy_deployments),
excluded_deployment_ids=_excluded_deployment_ids,
)
## RETRY SKIP ## -> drop deployments that already refused this request with a
## non-retryable status, unless that leaves nothing, so the caller still gets
## the provider's own error instead of a no-deployments error.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (
litellm.utils.get_excluded_filtered_deployments(
healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids
)
or healthy_deployments
)
if len(healthy_deployments) == 0:
exception: Final = await async_raise_no_deployment_exception(
litellm_router_instance=self,
@ -13359,7 +13474,7 @@ class Router:
## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2)
_target_order: Final = (request_kwargs or {}).pop("_target_order", None)
healthy_deployments = litellm.utils._get_order_filtered_deployments(
healthy_deployments = litellm.utils.get_order_filtered_deployments(
healthy_deployments, target_order=_target_order
)
@ -13367,11 +13482,22 @@ class Router:
## this request via weighted-failover. See async counterpart in
## async_get_healthy_deployments for details.
_excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None)
healthy_deployments = litellm.utils._get_excluded_filtered_deployments(
healthy_deployments = litellm.utils.get_excluded_filtered_deployments(
healthy_deployments,
excluded_deployment_ids=_excluded_deployment_ids,
)
## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments.
_retry_skipped_deployment_ids: Final = _as_retry_skipped_deployment_ids(
request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None
)
healthy_deployments = (
litellm.utils.get_excluded_filtered_deployments(
healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids
)
or healthy_deployments
)
if len(healthy_deployments) == 0:
model_ids = self.get_model_ids(model_name=model)
_cooldown_time = self.cooldown_cache.get_min_cooldown(

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