chore: merge litellm_internal_staging into litellm_lit_7039_least_busy_shared_counts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-08 17:11:05 +00:00
commit eeef03f122
467 changed files with 24099 additions and 5991 deletions

View file

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

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"

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

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

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

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

View file

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

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

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,92 @@ 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
- name: Run pytest tests/test_litellm_rust with the compiled extension
run: make test-rust-extension
- run: >-
uv build --wheel --out-dir panic-dist
--config-setting "maturin.build-args=--features panic-test,extension-module"
- 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

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

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

View file

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

View file

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

View file

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

View file

@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5532
"limit": 5570
},
"reportMissingTypeArgument": {
"limit": 15273
"limit": 15281
},
"reportMissingTypeStubs": {
"limit": 40
@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1808
"limit": 1804
},
"reportRedeclaration": {
"limit": 8
@ -99,19 +99,19 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44357
"limit": 44358
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38208
"limit": 38271
},
"reportUnknownParameterType": {
"limit": 19532
"limit": 19584
},
"reportUnknownVariableType": {
"limit": 29751
"limit": 29814
},
"reportUnnecessaryCast": {
"limit": 110
@ -135,7 +135,7 @@
"limit": 21
},
"reportUnusedFunction": {
"limit": 137
"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:]))

View file

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

View file

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

View file

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

View file

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

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

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"

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

View file

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

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

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

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

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,

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

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]

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

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

View file

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

View file

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

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

View file

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

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

View file

@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
budget_reservation_disabled_info_emitted = False
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
@ -1901,6 +1902,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(

View file

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

View file

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

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:

View file

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

View file

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

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

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

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,

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,

View file

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

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
]

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

View file

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

View file

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

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.

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

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

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.

View file

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

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:

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

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,

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,

View file

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

View file

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

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)

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

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

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

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

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

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

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

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:

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
)

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,

View file

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

View file

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

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