mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
Merge pull request #26430 from BerriAI/litellm_internal_staging
Some checks failed
Unit Tests: Security / security (push) Has been cancelled
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / schema-migration (push) Has been cancelled
Some checks failed
Unit Tests: Security / security (push) Has been cancelled
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / schema-migration (push) Has been cancelled
merge main
This commit is contained in:
commit
7734770bd3
204 changed files with 20050 additions and 2278 deletions
1831
.circleci/config.yml
1831
.circleci/config.yml
File diff suppressed because it is too large
Load diff
37
.github/workflows/_test-unit-services-base.yml
vendored
37
.github/workflows/_test-unit-services-base.yml
vendored
|
|
@ -32,41 +32,39 @@ on:
|
|||
required: false
|
||||
type: boolean
|
||||
default: false
|
||||
dist:
|
||||
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
|
||||
required: false
|
||||
type: string
|
||||
default: "loadscope"
|
||||
artifact-name:
|
||||
description: "Unique name for the coverage artifact (must be unique per run)"
|
||||
required: false
|
||||
type: string
|
||||
default: "run"
|
||||
secrets:
|
||||
DATABASE_URL:
|
||||
required: false
|
||||
POSTGRES_USER:
|
||||
required: false
|
||||
POSTGRES_PASSWORD:
|
||||
required: false
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
# The postgres service container below is spawned per-job on localhost and
|
||||
# destroyed with the job. Nothing outside the runner can reach it. The
|
||||
# user/password/database here are not secrets — they're bootstrap values
|
||||
# for a throwaway container — so we hardcode them instead of attaching
|
||||
# every matrix shard to a GHA environment just to read three "secrets"
|
||||
# (which also produces a "temporarily deployed to …" notification on the
|
||||
# PR timeline per shard per push).
|
||||
jobs:
|
||||
run:
|
||||
name: Run tests
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: ${{ inputs.timeout-minutes }}
|
||||
# Environment is derived from the enable-* flags, not caller-controllable.
|
||||
# This prevents callers from passing arbitrary environment names to bypass secret scoping.
|
||||
environment: >-
|
||||
${{
|
||||
inputs.enable-postgres && 'integration-postgres' ||
|
||||
''
|
||||
}}
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres@sha256:705a5d5b5836f3fcba0d02c4d281e6a7dd9ed2dd4078640f08a1e1e9896e097d # postgres:14
|
||||
env:
|
||||
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
|
||||
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: litellm
|
||||
POSTGRES_DB: litellm_test
|
||||
ports:
|
||||
- 5432:5432
|
||||
|
|
@ -114,7 +112,7 @@ jobs:
|
|||
- name: Run Prisma migrations
|
||||
if: ${{ inputs.enable-postgres }}
|
||||
env:
|
||||
DATABASE_URL: ${{ secrets.DATABASE_URL }}
|
||||
DATABASE_URL: "postgresql://litellm:litellm@localhost:5432/litellm_test"
|
||||
run: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
|
||||
|
|
@ -124,7 +122,8 @@ jobs:
|
|||
MAX_FAILURES: ${{ inputs.max-failures }}
|
||||
WORKERS: ${{ inputs.workers }}
|
||||
RERUNS: ${{ inputs.reruns }}
|
||||
DATABASE_URL: ${{ inputs.enable-postgres && secrets.DATABASE_URL || '' }}
|
||||
DIST: ${{ inputs.dist }}
|
||||
DATABASE_URL: ${{ inputs.enable-postgres && 'postgresql://litellm:litellm@localhost:5432/litellm_test' || '' }}
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
|
|
@ -143,7 +142,7 @@ jobs:
|
|||
-n "${WORKERS}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--dist=loadscope \
|
||||
--dist="${DIST}" \
|
||||
--durations=20 \
|
||||
--cov=litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
|
|
|
|||
2
.github/workflows/check_duplicate_issues.yml
vendored
2
.github/workflows/check_duplicate_issues.yml
vendored
|
|
@ -39,7 +39,7 @@ jobs:
|
|||
if: github.event.action == 'opened'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.11"
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Auto-close if high-confidence duplicate
|
||||
if: github.event.action == 'opened'
|
||||
|
|
|
|||
65
.github/workflows/create-release-branch.yml
vendored
Normal file
65
.github/workflows/create-release-branch.yml
vendored
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
name: Create Release Branch
|
||||
|
||||
on:
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
tag:
|
||||
description: "Release tag (e.g. v1.83.0-stable) — branch will be named release/<tag>"
|
||||
required: true
|
||||
type: string
|
||||
commit_hash:
|
||||
description: "Full 40-char commit SHA the branch should point to"
|
||||
required: true
|
||||
type: string
|
||||
workflow_call:
|
||||
inputs:
|
||||
tag:
|
||||
description: "Release tag"
|
||||
required: true
|
||||
type: string
|
||||
commit_hash:
|
||||
description: "Full 40-char commit SHA the branch should point to"
|
||||
required: true
|
||||
type: string
|
||||
|
||||
permissions: {}
|
||||
|
||||
jobs:
|
||||
create-branch:
|
||||
name: Create Release Branch
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: write
|
||||
steps:
|
||||
- name: Validate inputs
|
||||
env:
|
||||
TAG: ${{ inputs.tag }}
|
||||
COMMIT_HASH: ${{ inputs.commit_hash }}
|
||||
run: |
|
||||
if ! echo "${COMMIT_HASH}" | grep -qE '^[0-9a-f]{40}$'; then
|
||||
echo "::error::commit_hash must be a full 40-character commit SHA"
|
||||
exit 1
|
||||
fi
|
||||
if ! echo "${TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.[0-9]+'; then
|
||||
echo "::error::tag must start with vX.Y.Z"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Create release branch
|
||||
env:
|
||||
TAG: ${{ inputs.tag }}
|
||||
COMMIT_HASH: ${{ inputs.commit_hash }}
|
||||
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
|
||||
with:
|
||||
script: |
|
||||
const tag = process.env.TAG;
|
||||
const commitHash = process.env.COMMIT_HASH;
|
||||
const branchName = `release/${tag}`;
|
||||
|
||||
await github.rest.git.createRef({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
ref: `refs/heads/${branchName}`,
|
||||
sha: commitHash,
|
||||
});
|
||||
core.info(`Created branch ${branchName} at ${commitHash}`);
|
||||
11
.github/workflows/create-release.yml
vendored
11
.github/workflows/create-release.yml
vendored
|
|
@ -102,6 +102,17 @@ jobs:
|
|||
body: updatedBody,
|
||||
draft: false,
|
||||
});
|
||||
|
||||
} catch (error) {
|
||||
core.setFailed(error.message);
|
||||
}
|
||||
|
||||
create-branch:
|
||||
name: Create Release Branch
|
||||
needs: release
|
||||
permissions:
|
||||
contents: write
|
||||
uses: ./.github/workflows/create-release-branch.yml
|
||||
with:
|
||||
tag: ${{ inputs.tag }}
|
||||
commit_hash: ${{ inputs.commit_hash }}
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ jobs:
|
|||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.11"
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
|
|
|
|||
2
.github/workflows/scan_duplicate_issues.yml
vendored
2
.github/workflows/scan_duplicate_issues.yml
vendored
|
|
@ -29,7 +29,7 @@ jobs:
|
|||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.13"
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Scan for duplicate issues
|
||||
env:
|
||||
|
|
|
|||
136
.github/workflows/test-code-quality.yml
vendored
Normal file
136
.github/workflows/test-code-quality.yml
vendored
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
name: Code Quality Checks
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
code-quality:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Checkout litellm-docs (for documentation_tests)
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
repository: BerriAI/litellm-docs
|
||||
path: _litellm_docs_checkout
|
||||
persist-credentials: false
|
||||
|
||||
- name: Wire up docs path expected by documentation_tests/*
|
||||
run: |
|
||||
# documentation_tests scripts read from docs/my-website/docs/...
|
||||
# In litellm-docs the same files live at docs/... (repo root).
|
||||
# Point docs/my-website -> litellm-docs checkout so the paths resolve.
|
||||
rm -rf docs/my-website
|
||||
ln -s ../_litellm_docs_checkout docs/my-website
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Install dependencies
|
||||
run: uv sync --frozen --all-groups --all-extras
|
||||
|
||||
- name: check_licenses
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_licenses.py
|
||||
|
||||
- name: check_provider_folders_documented
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py
|
||||
|
||||
- name: router_code_coverage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py
|
||||
|
||||
- name: test_chat_completion_imports
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py
|
||||
|
||||
- name: info_log_check
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/info_log_check.py
|
||||
|
||||
- name: check_guardrail_apply_decorator
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_guardrail_apply_decorator.py
|
||||
|
||||
- name: test_ban_set_verbose
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_ban_set_verbose.py
|
||||
|
||||
- name: code_qa_check_tests
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/code_qa_check_tests.py
|
||||
|
||||
- name: check_get_model_cost_key_performance
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_get_model_cost_key_performance.py
|
||||
|
||||
- name: test_proxy_types_import
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_proxy_types_import.py
|
||||
|
||||
- name: callback_manager_test
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/callback_manager_test.py
|
||||
|
||||
- name: recursive_detector
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/recursive_detector.py
|
||||
|
||||
- name: test_router_strategy_async
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_router_strategy_async.py
|
||||
|
||||
- name: litellm_logging_code_coverage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/litellm_logging_code_coverage.py
|
||||
|
||||
- name: ensure_async_clients_test
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/ensure_async_clients_test.py
|
||||
|
||||
- name: enforce_llms_folder_style
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/enforce_llms_folder_style.py
|
||||
|
||||
- name: prevent_key_leaks_in_exceptions
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/prevent_key_leaks_in_exceptions.py
|
||||
|
||||
- name: check_unsafe_enterprise_import
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
|
||||
|
||||
- name: ban_copy_deepcopy_kwargs
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
|
||||
|
||||
- name: check_fastuuid_usage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
|
||||
|
||||
- name: memory_test
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py
|
||||
|
||||
- name: documentation_test_env_keys
|
||||
run: uv run --no-sync python ./tests/documentation_tests/test_env_keys.py
|
||||
|
||||
- name: documentation_test_router_settings
|
||||
run: uv run --no-sync python ./tests/documentation_tests/test_router_settings.py
|
||||
|
||||
- name: documentation_test_api_docs
|
||||
run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py
|
||||
39
.github/workflows/test-semgrep.yml
vendored
Normal file
39
.github/workflows/test-semgrep.yml
vendored
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
name: Semgrep
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
semgrep:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Run Semgrep (custom rules)
|
||||
run: uv tool run --from 'semgrep==1.157.0' semgrep scan --config .semgrep/rules . --error
|
||||
216
.github/workflows/test-unit-proxy-db.yml
vendored
216
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -12,8 +12,74 @@ concurrency:
|
|||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
# Semantic matrix: each shard groups tests by concern (auth, server, logging, …)
|
||||
# rather than alphabetical letter ranges. Adding a new test file means adding it
|
||||
# to whichever group it belongs to, not reshuffling slices.
|
||||
#
|
||||
# Design targets:
|
||||
# * Every shard runs in <= 7 minutes of wall-clock on the default runner.
|
||||
# Most of a shard's time is pytest plugin load + xdist worker imports +
|
||||
# pytest-cov instrumentation, not the tests themselves. Keeping per-shard
|
||||
# work low and matching worker count to runner cores is what controls it.
|
||||
# * workers: 4 matches the 4-core ubuntu-latest runner. -n 8 on 4 cores
|
||||
# oversubscribes 2x and workers fight for CPU during their cold-start
|
||||
# imports (measured ~441% CPU for -n 8 locally, i.e. ~55% effective).
|
||||
# * test_key_generate_prisma.py stays serial (workers=0) — it has event-loop
|
||||
# conflicts with the logging worker when run in parallel.
|
||||
# * test_proxy_utils.py runs as a single shard with --dist=worksteal so
|
||||
# xdist balances its 188 parametrized cases across workers instead of
|
||||
# pinning the whole file to one worker (the default --dist=loadscope
|
||||
# behavior for single-file targets).
|
||||
# * test_db_schema_migration.py is isolated because one test in it
|
||||
# (test_aaaasschema_migration_check) takes ~170s — by itself it
|
||||
# determines the shard's wall-clock floor.
|
||||
jobs:
|
||||
# Fast guard — fails the workflow if a test_*.py file under
|
||||
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
|
||||
# The semantic-shard design (no catch-all "remaining" bucket) relies on
|
||||
# every test file being explicitly assigned; this guard prevents a new
|
||||
# file from silently dropping out of CI.
|
||||
assert-shard-coverage:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 2
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
- name: Assert every test_*.py is in a matrix shard
|
||||
run: |
|
||||
python3 - <<'PY'
|
||||
import pathlib, sys, yaml
|
||||
wf = yaml.safe_load(open(".github/workflows/test-unit-proxy-db.yml"))
|
||||
matrix = wf["jobs"]["proxy-db"]["strategy"]["matrix"]["include"]
|
||||
referenced = set()
|
||||
for entry in matrix:
|
||||
for token in entry["test-path"].split():
|
||||
if token.startswith("tests/proxy_unit_tests/"):
|
||||
referenced.add(pathlib.PurePosixPath(token).name)
|
||||
actual = {p.name for p in pathlib.Path("tests/proxy_unit_tests").iterdir()
|
||||
if p.name.startswith("test_") and (p.suffix == ".py" or p.is_dir())
|
||||
and p.name != "test_configs"}
|
||||
orphans = sorted(actual - referenced)
|
||||
if orphans:
|
||||
print("ERROR: the following files/dirs under tests/proxy_unit_tests/")
|
||||
print(" are not assigned to any shard in test-unit-proxy-db.yml:")
|
||||
for o in orphans:
|
||||
print(f" - {o}")
|
||||
print()
|
||||
print("Add each to whichever semantic shard it belongs to.")
|
||||
sys.exit(1)
|
||||
print(f"OK: all {len(actual)} files assigned to a shard.")
|
||||
PY
|
||||
|
||||
proxy-db:
|
||||
needs: assert-shard-coverage
|
||||
# Display only the semantic shard name in the checks UI instead of GHA's
|
||||
# default "proxy-db (key-generation, tests/proxy_unit_tests/…, 0, loadscope, 20)"
|
||||
# which includes every matrix field and gets truncated past the test-path.
|
||||
name: ${{ matrix.test-group }}
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
|
|
@ -22,19 +88,146 @@ jobs:
|
|||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
# Key generation tests must NOT run in parallel (event loop conflicts with logging worker)
|
||||
# Must run serially — event-loop conflict with the logging worker.
|
||||
- test-group: key-generation
|
||||
test-path: "tests/proxy_unit_tests/test_key_generate_prisma.py"
|
||||
workers: 0
|
||||
timeout: 30
|
||||
- test-group: auth-checks
|
||||
test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py"
|
||||
workers: 8
|
||||
dist: loadscope
|
||||
timeout: 20
|
||||
- test-group: remaining
|
||||
test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py"
|
||||
workers: 8
|
||||
timeout: 30
|
||||
|
||||
# ---- auth: split into 2 shards ----
|
||||
- test-group: auth-checks
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_auth_checks.py
|
||||
tests/proxy_unit_tests/test_user_api_key_auth.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
- test-group: jwt-and-keys
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_jwt.py
|
||||
tests/proxy_unit_tests/test_jwt_key_mapping.py
|
||||
tests/proxy_unit_tests/test_proxy_custom_auth.py
|
||||
tests/proxy_unit_tests/test_key_generate_dynamodb.py
|
||||
tests/proxy_unit_tests/test_deployed_proxy_keygen.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
|
||||
# ---- test_proxy_utils.py, single shard, worksteal distribution ----
|
||||
- test-group: proxy-utils
|
||||
test-path: "tests/proxy_unit_tests/test_proxy_utils.py"
|
||||
workers: 4
|
||||
dist: worksteal
|
||||
timeout: 15
|
||||
|
||||
# ---- proxy server: split into 2 shards ----
|
||||
- test-group: proxy-server-core
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_proxy_server.py
|
||||
tests/proxy_unit_tests/test_proxy_server_keys.py
|
||||
tests/proxy_unit_tests/test_proxy_server_caching.py
|
||||
tests/proxy_unit_tests/test_proxy_server_langfuse.py
|
||||
tests/proxy_unit_tests/test_proxy_server_spend.py
|
||||
tests/proxy_unit_tests/test_aproxy_startup.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
- test-group: proxy-runtime
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_proxy_config_unit_test.py
|
||||
tests/proxy_unit_tests/test_proxy_routes.py
|
||||
tests/proxy_unit_tests/test_proxy_gunicorn.py
|
||||
tests/proxy_unit_tests/test_server_root_path.py
|
||||
tests/proxy_unit_tests/test_proxy_pass_user_config.py
|
||||
tests/proxy_unit_tests/test_proxy_token_counter.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
|
||||
# ---- logging: split into 2 shards ----
|
||||
- test-group: custom-logging
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_custom_callback_input.py
|
||||
tests/proxy_unit_tests/test_custom_logger_s3_gcs.py
|
||||
tests/proxy_unit_tests/test_proxy_custom_logger.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
- test-group: logging-misc
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_proxy_reject_logging.py
|
||||
tests/proxy_unit_tests/test_audit_logs_proxy.py
|
||||
tests/proxy_unit_tests/test_search_api_logging.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
|
||||
# ---- db-and-spend: isolate the 170s schema-migration test ----
|
||||
# test_db_schema_migration.py has exactly one test, and that test
|
||||
# is mostly waiting on `prisma migrate deploy` / `prisma migrate
|
||||
# diff` subprocesses (~170s). It does no CPU-bound Python work
|
||||
# inside the test. Running with workers=0 (serial, no xdist)
|
||||
# skips the 4-worker cold-start cost we'd otherwise pay for a
|
||||
# single test, saving ~4 minutes of wall-clock.
|
||||
- test-group: schema-migration
|
||||
test-path: "tests/proxy_unit_tests/test_db_schema_migration.py"
|
||||
workers: 0
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
- test-group: db-and-spend
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_prisma_client_backoff_retry.py
|
||||
tests/proxy_unit_tests/test_db_schema_changes.py
|
||||
tests/proxy_unit_tests/test_e2e_pod_lock_manager.py
|
||||
tests/proxy_unit_tests/test_skills_db.py
|
||||
tests/proxy_unit_tests/test_update_daily_tag_spend.py
|
||||
tests/proxy_unit_tests/test_update_spend.py
|
||||
tests/proxy_unit_tests/test_project_endpoints_prisma.py
|
||||
tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
|
||||
# ---- guardrails + budget + hooks: split into 2 ----
|
||||
- test-group: guardrails-hooks
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_proxy_setting_guardrails.py
|
||||
tests/proxy_unit_tests/test_banned_keyword_list.py
|
||||
tests/proxy_unit_tests/test_unit_test_proxy_hooks.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
- test-group: budgets
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_default_end_user_budget_simple.py
|
||||
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
|
||||
tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
|
||||
- test-group: endpoints-and-responses
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_blog_posts_endpoint.py
|
||||
tests/proxy_unit_tests/test_models_fallback_endpoint.py
|
||||
tests/proxy_unit_tests/test_google_endpoint_routing.py
|
||||
tests/proxy_unit_tests/test_google_gemini_proxy_request.py
|
||||
tests/proxy_unit_tests/test_get_favicon.py
|
||||
tests/proxy_unit_tests/test_get_image.py
|
||||
tests/proxy_unit_tests/test_ui_path_detection.py
|
||||
tests/proxy_unit_tests/test_prompt_test_endpoint.py
|
||||
tests/proxy_unit_tests/test_check_batch_cost.py
|
||||
tests/proxy_unit_tests/test_check_responses_cost.py
|
||||
tests/proxy_unit_tests/test_response_polling_handler.py
|
||||
tests/proxy_unit_tests/test_response_polling_pre_call_checks.py
|
||||
tests/proxy_unit_tests/test_realtime_cache.py
|
||||
tests/proxy_unit_tests/test_proxy_exception_mapping.py
|
||||
tests/proxy_unit_tests/test_custom_tokenizer_bug.py
|
||||
tests/proxy_unit_tests/test_model_response_typing
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
with:
|
||||
test-path: ${{ matrix.test-path }}
|
||||
|
|
@ -42,8 +235,5 @@ jobs:
|
|||
reruns: 2
|
||||
timeout-minutes: ${{ matrix.timeout }}
|
||||
enable-postgres: true
|
||||
dist: ${{ matrix.dist }}
|
||||
artifact-name: proxy-db-${{ matrix.test-group }}
|
||||
secrets:
|
||||
DATABASE_URL: ${{ secrets.DATABASE_URL }}
|
||||
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
|
||||
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
|
||||
|
|
|
|||
|
|
@ -36,6 +36,8 @@ jobs:
|
|||
tests/test_litellm/proxy/health_endpoints
|
||||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
workers: 2
|
||||
reruns: 2
|
||||
|
|
|
|||
8
.github/workflows/test-unit-security.yml
vendored
8
.github/workflows/test-unit-security.yml
vendored
|
|
@ -1,6 +1,8 @@
|
|||
name: "Unit Tests: Security"
|
||||
|
||||
# Uses DATABASE_URL secret — only runs on trusted branches, not PRs.
|
||||
# Kept push-only (was previously required by DATABASE_URL secret scoping;
|
||||
# now the postgres credentials are ephemeral localhost values but the
|
||||
# push-trigger stays to match the proxy-db workflow cadence).
|
||||
on:
|
||||
push:
|
||||
branches: [main, "litellm_**"]
|
||||
|
|
@ -24,7 +26,3 @@ jobs:
|
|||
timeout-minutes: 20
|
||||
enable-postgres: true
|
||||
artifact-name: security
|
||||
secrets:
|
||||
DATABASE_URL: ${{ secrets.DATABASE_URL }}
|
||||
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
|
||||
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}
|
||||
|
|
|
|||
13
Dockerfile
13
Dockerfile
|
|
@ -27,10 +27,8 @@ RUN apk add --no-cache \
|
|||
npm \
|
||||
libsndfile
|
||||
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
# Copy dependency metadata first for layer caching
|
||||
|
|
@ -94,11 +92,14 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervi
|
|||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
|
||||
WORKDIR /app
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
|
||||
COPY --from=builder /root/.cache /root/.cache
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
|
|
|
|||
|
|
@ -1,22 +1,231 @@
|
|||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
import testing.postgresql
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import testing.postgresql
|
||||
|
||||
|
||||
def create_migration(migration_name: str = None):
|
||||
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:
|
||||
"""Return SQL lines containing DROP COLUMN, DROP TABLE, or DROP INDEX."""
|
||||
return [
|
||||
line.strip() for line in sql.splitlines() if DESTRUCTIVE_PATTERN.search(line)
|
||||
]
|
||||
|
||||
|
||||
def _print_freshness_failure(
|
||||
base_branch: str, reason: str, stderr_text: str = ""
|
||||
) -> None:
|
||||
"""Loudly refuse to run when the freshness check can't be completed."""
|
||||
banner = "=" * 72
|
||||
out = sys.stderr
|
||||
print(banner, file=out)
|
||||
print(f" FRESHNESS CHECK FAILED — COULD NOT VERIFY origin/{base_branch}", file=out)
|
||||
print(banner, file=out)
|
||||
print("", file=out)
|
||||
print(f"Reason: {reason}", file=out)
|
||||
if stderr_text:
|
||||
print("", file=out)
|
||||
print("git stderr:", file=out)
|
||||
for line in stderr_text.rstrip().splitlines():
|
||||
print(f" {line}", file=out)
|
||||
print("", file=out)
|
||||
print("Common causes:", file=out)
|
||||
print(" - No network access (offline)", file=out)
|
||||
print(" - 'origin' remote not configured, or base branch name is wrong", file=out)
|
||||
print(" - Not a git repository", file=out)
|
||||
print("", file=out)
|
||||
print("Options:", file=out)
|
||||
print(
|
||||
f" - Fix the above and re-run, OR pass --base-branch <name> if your", file=out
|
||||
)
|
||||
print(
|
||||
f" base branch is not '{base_branch}', OR pass --skip-freshness-check",
|
||||
file=out,
|
||||
)
|
||||
print(" to bypass (only if you fully understand the risk).", file=out)
|
||||
print(banner, file=out)
|
||||
|
||||
|
||||
def _print_stale_branch_refusal(base_branch: str, behind: int) -> None:
|
||||
"""Loudly refuse to run when HEAD is behind the base branch."""
|
||||
banner = "=" * 72
|
||||
out = sys.stderr
|
||||
print(banner, file=out)
|
||||
print(
|
||||
f" STALE BRANCH — {behind} commit(s) behind origin/{base_branch}",
|
||||
file=out,
|
||||
)
|
||||
print(banner, file=out)
|
||||
print("", file=out)
|
||||
print(
|
||||
f"Your branch is {behind} commit(s) behind origin/{base_branch}. Generating a\n"
|
||||
"migration from a stale branch is how newly-added columns get silently\n"
|
||||
"dropped — the exact incident this guard exists to prevent.",
|
||||
file=out,
|
||||
)
|
||||
print("", file=out)
|
||||
print("Before re-running:", file=out)
|
||||
print(f" git fetch origin && git rebase origin/{base_branch}", file=out)
|
||||
print(" # or merge, whichever matches your workflow", file=out)
|
||||
print("", file=out)
|
||||
print(
|
||||
"If you have a genuine reason to bypass (e.g. intentionally generating",
|
||||
file=out,
|
||||
)
|
||||
print("a migration against an older base), pass --skip-freshness-check.", file=out)
|
||||
print("", file=out)
|
||||
print(banner, file=out)
|
||||
print(
|
||||
" AI AGENTS / ASSISTANTS: DO NOT automatically run `git rebase` or pass\n"
|
||||
" --skip-freshness-check. Surface this error to the human operator and\n"
|
||||
" wait for their explicit confirmation. Auto-rebasing can drop the\n"
|
||||
" human's in-progress schema edits via a bad conflict resolution.",
|
||||
file=out,
|
||||
)
|
||||
print(banner, file=out)
|
||||
|
||||
|
||||
def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
|
||||
"""Fetch origin/<base_branch> and exit 3 if HEAD is behind it."""
|
||||
cwd = str(root_dir)
|
||||
try:
|
||||
subprocess.run(
|
||||
["git", "fetch", "origin", base_branch],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=cwd,
|
||||
)
|
||||
except FileNotFoundError:
|
||||
_print_freshness_failure(base_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",
|
||||
e.stderr or "",
|
||||
)
|
||||
sys.exit(3)
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"],
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
cwd=cwd,
|
||||
)
|
||||
behind = int(result.stdout.strip())
|
||||
except subprocess.CalledProcessError as e:
|
||||
_print_freshness_failure(
|
||||
base_branch,
|
||||
f"`git rev-list HEAD..origin/{base_branch}` failed",
|
||||
e.stderr or "",
|
||||
)
|
||||
sys.exit(3)
|
||||
except ValueError:
|
||||
_print_freshness_failure(
|
||||
base_branch,
|
||||
"could not parse commit count from `git rev-list`",
|
||||
)
|
||||
sys.exit(3)
|
||||
|
||||
if behind > 0:
|
||||
_print_stale_branch_refusal(base_branch, behind)
|
||||
sys.exit(3)
|
||||
|
||||
print(f"Branch freshness OK: up to date with origin/{base_branch}.")
|
||||
|
||||
|
||||
def _print_destructive_refusal(destructive_lines: list) -> None:
|
||||
"""Loudly refuse to write a destructive migration and explain how to proceed."""
|
||||
banner = "=" * 72
|
||||
out = sys.stderr
|
||||
print(banner, file=out)
|
||||
print(
|
||||
" DESTRUCTIVE MIGRATION DETECTED — REFUSING TO WRITE MIGRATION FILE", file=out
|
||||
)
|
||||
print(banner, file=out)
|
||||
print("", file=out)
|
||||
print(
|
||||
"The generated migration contains the following destructive statements:",
|
||||
file=out,
|
||||
)
|
||||
print("", file=out)
|
||||
for line in destructive_lines:
|
||||
print(f" {line}", file=out)
|
||||
print("", file=out)
|
||||
print(
|
||||
"This very often means your branch is OUT OF DATE, or your local\n"
|
||||
"schema.prisma is inconsistent with the migrations/ directory. A\n"
|
||||
"previous incident saw newly-added columns silently dropped this way\n"
|
||||
"and merged to main. Stop and verify before proceeding.",
|
||||
file=out,
|
||||
)
|
||||
print("", file=out)
|
||||
print("Before re-running:", file=out)
|
||||
print(
|
||||
" 1. git fetch origin && git status # confirm branch is up to date", file=out
|
||||
)
|
||||
print(
|
||||
" 2. Re-check all schema.prisma files are in sync (runbook Step 0)", file=out
|
||||
)
|
||||
print(" 3. Review EACH DROP above — is it actually intended?", file=out)
|
||||
print(" 4. If (and only if) the drops are intentional, re-run with:", file=out)
|
||||
print(" --allow-destructive", file=out)
|
||||
print("", file=out)
|
||||
print(banner, file=out)
|
||||
print(
|
||||
" AI AGENTS / ASSISTANTS: DO NOT automatically re-run this command\n"
|
||||
" with --allow-destructive. Surface this error to the human operator\n"
|
||||
" and wait for their explicit confirmation before passing the flag.\n"
|
||||
" Passing the flag without human review is the exact failure mode\n"
|
||||
" this guard exists to prevent.",
|
||||
file=out,
|
||||
)
|
||||
print(banner, file=out)
|
||||
|
||||
|
||||
def create_migration(
|
||||
migration_name: str = None,
|
||||
allow_destructive: bool = False,
|
||||
base_branch: str = DEFAULT_BASE_BRANCH,
|
||||
skip_freshness_check: bool = False,
|
||||
):
|
||||
"""
|
||||
Create a new migration SQL file in the migrations directory by comparing
|
||||
current database state with schema
|
||||
current database state with schema.
|
||||
|
||||
Args:
|
||||
migration_name (str): Name for the migration
|
||||
allow_destructive (bool): Required to write a migration that contains
|
||||
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").
|
||||
skip_freshness_check (bool): Skip the "branch is up to date" check.
|
||||
Only for intentional migrations against an older base.
|
||||
"""
|
||||
root_dir = Path(__file__).parent.parent
|
||||
|
||||
if skip_freshness_check:
|
||||
print(
|
||||
"WARNING: freshness check skipped (--skip-freshness-check). "
|
||||
"Generating a migration from a stale branch can silently drop columns."
|
||||
)
|
||||
else:
|
||||
_check_branch_freshness(root_dir, base_branch)
|
||||
|
||||
try:
|
||||
# Get paths
|
||||
root_dir = Path(__file__).parent.parent
|
||||
migrations_dir = (
|
||||
root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
|
||||
)
|
||||
|
|
@ -59,7 +268,27 @@ def create_migration(migration_name: str = None):
|
|||
check=True,
|
||||
)
|
||||
|
||||
if result.stdout.strip():
|
||||
# Prisma emits the literal "-- This is an empty migration." when
|
||||
# there's no real drift. Treat that as "no changes".
|
||||
diff_sql = result.stdout
|
||||
stripped = diff_sql.strip()
|
||||
is_empty_diff = (
|
||||
not stripped or stripped == "-- This is an empty migration."
|
||||
)
|
||||
|
||||
if not is_empty_diff:
|
||||
destructive_lines = _find_destructive_statements(diff_sql)
|
||||
if destructive_lines and not allow_destructive:
|
||||
_print_destructive_refusal(destructive_lines)
|
||||
sys.exit(2)
|
||||
if destructive_lines and allow_destructive:
|
||||
print(
|
||||
"WARNING: writing destructive migration "
|
||||
"(--allow-destructive passed). Statements:"
|
||||
)
|
||||
for line in destructive_lines:
|
||||
print(f" {line}")
|
||||
|
||||
# Generate timestamp and create migration directory
|
||||
timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
|
||||
migration_name = migration_name or "unnamed_migration"
|
||||
|
|
@ -68,7 +297,7 @@ def create_migration(migration_name: str = None):
|
|||
|
||||
# Write the SQL to migration.sql
|
||||
migration_file = migration_dir / "migration.sql"
|
||||
migration_file.write_text(result.stdout)
|
||||
migration_file.write_text(diff_sql)
|
||||
|
||||
print(f"Created migration in {migration_dir}")
|
||||
return True
|
||||
|
|
@ -90,8 +319,48 @@ def create_migration(migration_name: str = None):
|
|||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# If running directly, can optionally pass migration name as argument
|
||||
import sys
|
||||
|
||||
migration_name = sys.argv[1] if len(sys.argv) > 1 else None
|
||||
create_migration(migration_name)
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Generate a Prisma migration by diffing the temp DB "
|
||||
"(existing migrations applied) against schema.prisma."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"migration_name",
|
||||
nargs="?",
|
||||
default=None,
|
||||
help="Name for the migration (used in the generated directory name).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--allow-destructive",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Required to write a migration that contains DROP COLUMN, "
|
||||
"DROP TABLE, or DROP INDEX. Without this flag, destructive "
|
||||
"diffs are refused."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--base-branch",
|
||||
default=DEFAULT_BASE_BRANCH,
|
||||
help=(
|
||||
f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). "
|
||||
"The script fetches origin/<base-branch> and refuses to run if HEAD "
|
||||
"is behind it."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip-freshness-check",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Bypass the 'branch is up to date' check. Only for intentional "
|
||||
"migrations against an older base. Pairs poorly with automation."
|
||||
),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
create_migration(
|
||||
args.migration_name,
|
||||
allow_destructive=args.allow_destructive,
|
||||
base_branch=args.base_branch,
|
||||
skip_freshness_check=args.skip_freshness_check,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -26,10 +26,8 @@ RUN apk add --no-cache \
|
|||
npm \
|
||||
libsndfile
|
||||
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
# Copy dependency metadata first for layer caching
|
||||
|
|
@ -92,11 +90,14 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervi
|
|||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
|
||||
WORKDIR /app
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
|
||||
COPY --from=builder /root/.cache /root/.cache
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
|
|
|
|||
|
|
@ -138,7 +138,7 @@ RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
|
|||
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w "$LITELLM_PROXY_EXTRAS_PATH" || true && \
|
||||
chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets /app/.cache
|
||||
|
||||
USER nobody
|
||||
USER 65534
|
||||
|
||||
RUN prisma generate --schema=./schema.prisma
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ schemaVersion: 2.0.0
|
|||
|
||||
metadataTest:
|
||||
entrypoint: ["docker/prod_entrypoint.sh"]
|
||||
user: "nobody"
|
||||
user: "65534"
|
||||
workdir: "/app"
|
||||
|
||||
fileExistenceTests:
|
||||
|
|
|
|||
155
docs/my-website/docs/adaptive_router.md
Normal file
155
docs/my-website/docs/adaptive_router.md
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
# [BETA] Adaptive Router
|
||||
|
||||
:::info
|
||||
|
||||
Beta feature. Share feedback on [Discord](https://discord.gg/wuPM9dRgDw) or [Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA).
|
||||
|
||||
:::
|
||||
|
||||
**Requirements:** LiteLLM Proxy with a Postgres database. Quality estimates are stored in Postgres and loaded on startup — without a database the router works but forgets everything learned on restart.
|
||||
|
||||
You have a cheap model and an expensive one. You want to use the cheap one when it's good enough, and the expensive one when it actually matters — without hardcoding rules you'll spend months tuning.
|
||||
|
||||
The adaptive router does this automatically. It tracks which model performs best for each type of request (code, writing, analysis, etc.) and routes accordingly, balancing quality against cost based on weights you control.
|
||||
|
||||
## Quick start
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
model_info:
|
||||
input_cost_per_token: 0.0000025
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 3 # 1=budget, 2=mid, 3=frontier
|
||||
strengths: ["code_generation", "analytical_reasoning"]
|
||||
|
||||
- model_name: gpt-4o-mini
|
||||
litellm_params:
|
||||
model: openai/gpt-4o-mini
|
||||
model_info:
|
||||
input_cost_per_token: 0.00000015
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 2
|
||||
strengths: ["factual_lookup"]
|
||||
|
||||
- model_name: my-router
|
||||
litellm_params:
|
||||
model: auto_router/adaptive_router
|
||||
adaptive_router_config:
|
||||
available_models: ["gpt-4o", "gpt-4o-mini"]
|
||||
weights:
|
||||
quality: 0.7 # raise this if quality complaints; lower if bill too high
|
||||
cost: 0.3 # must sum to 1.0 with quality
|
||||
```
|
||||
|
||||
Route to it by setting `model` to your adaptive router's name:
|
||||
|
||||
```bash
|
||||
curl -X POST {{baseURL}}/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer $LITELLM_API_KEY" \
|
||||
-d '{
|
||||
"model": "my-router",
|
||||
"messages": [
|
||||
{"role": "user", "content": "build me a python script that parses CSV"},
|
||||
{"role": "assistant", "content": "Here is a script using csv.DictReader..."},
|
||||
{"role": "user", "content": "now add error handling for missing files"},
|
||||
{"role": "assistant", "content": "Wrap the open() call in a try/except FileNotFoundError..."},
|
||||
{"role": "user", "content": "perfect, that worked. thanks!"}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
The response includes a header telling you which model was actually picked:
|
||||
|
||||
```
|
||||
x-litellm-adaptive-router-model: gpt-4o
|
||||
```
|
||||
|
||||
The "thanks!" turn in the example above fires a satisfaction signal — that's what moves the bandit.
|
||||
|
||||
## Tuning cost vs. quality
|
||||
|
||||
The `weights` are your main lever:
|
||||
|
||||
| Goal | quality | cost |
|
||||
|---|---|---|
|
||||
| Minimize cost, quality is secondary | 0.3 | 0.7 |
|
||||
| Balanced | 0.5 | 0.5 |
|
||||
| Quality-first (default) | 0.7 | 0.3 |
|
||||
| Quality non-negotiable | 0.9 | 0.1 |
|
||||
|
||||
The router learns over time. For the first ~10 requests per model, it relies on the tiers you declared. After that, real performance data takes over.
|
||||
|
||||
## Force a minimum quality tier per request
|
||||
|
||||
If a specific request needs a frontier model regardless of cost, pass this header:
|
||||
|
||||
```
|
||||
x-litellm-min-quality-tier: 3
|
||||
```
|
||||
|
||||
You can also pass `min_quality_tier` via request metadata instead of a header.
|
||||
|
||||
## What's being learned
|
||||
|
||||
The router classifies each request into one of 7 types and tracks how each model performs on each independently. A model that's great at factual lookup but poor at code will win factual requests and lose code requests — even if it's cheaper overall.
|
||||
|
||||
| Type | Example |
|
||||
|---|---|
|
||||
| `code_generation` | "write me a Python sort function" |
|
||||
| `code_understanding` | "explain what this function does" |
|
||||
| `technical_design` | "how should I design this API?" |
|
||||
| `analytical_reasoning` | "calculate the probability that..." |
|
||||
| `writing` | "draft an email to my team about..." |
|
||||
| `factual_lookup` | "what is the capital of France?" |
|
||||
| `general` | anything else |
|
||||
|
||||
[**See classifier code**](https://github.com/BerriAI/litellm/blob/litellm_adaptive_routing/litellm/router_strategy/adaptive_router/classifier.py)
|
||||
|
||||
Learning signals are inspired by [Signals: Trajectory Sampling and Triage for Agentic Interactions](https://arxiv.org/pdf/2604.00356).
|
||||
|
||||
## Inspect the current state
|
||||
|
||||
```
|
||||
GET /adaptive_router/{router_name}/state
|
||||
```
|
||||
|
||||
Returns current quality estimates per model per request type. Useful for understanding why a model is or isn't being picked.
|
||||
|
||||
```json
|
||||
{
|
||||
"routers": [
|
||||
{
|
||||
"router_name": "smart-cheap-router",
|
||||
"available_models": ["fast", "smart"],
|
||||
"weights": { "quality": 0.7, "cost": 0.3 },
|
||||
"cells": [
|
||||
{
|
||||
"request_type": "analytical_reasoning",
|
||||
"model": "fast",
|
||||
"quality_mean": 0.5,
|
||||
"samples": 0
|
||||
},
|
||||
{
|
||||
"request_type": "analytical_reasoning",
|
||||
"model": "smart",
|
||||
"quality_mean": 0.95,
|
||||
"samples": 0
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
`quality_mean` is the key number — it's the router's current estimate of how well that model handles that request type. `samples` counts how many real observations have moved the prior (starts at 0; the cold-start prior mass is excluded).
|
||||
|
||||
## Known limitations
|
||||
|
||||
- Latency isn't scored — a slow model can still win on quality + cost
|
||||
- Signals are regex-based and English-biased — no LLM judge
|
||||
- Hard cap of 200 observations per cell; no decay yet
|
||||
- Once a model is picked for a session, other models' turns in that session don't contribute to learning
|
||||
|
|
@ -2061,7 +2061,7 @@ assert isinstance(
|
|||
|
||||
## Media Resolution Control (Images & Videos)
|
||||
|
||||
For Gemini 3+ models, LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types.
|
||||
LiteLLM supports per-part media resolution control using OpenAI's `detail` parameter for all Gemini models. This allows you to specify different resolution levels for individual images and videos in your request, whether using `image_url` or `file` content types.
|
||||
|
||||
**Supported `detail` values:**
|
||||
- `"low"` - Maps to `media_resolution: "low"` (280 tokens for images, 70 tokens per frame for videos)
|
||||
|
|
@ -2146,12 +2146,12 @@ response = completion(
|
|||
</Tabs>
|
||||
|
||||
:::info
|
||||
**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types, and is only available for Gemini 3+ models.
|
||||
**Per-Part Resolution:** Each image or video in your request can have its own `detail` setting, allowing mixed-resolution requests (e.g., a high-res chart alongside a low-res icon). This feature works with both `image_url` and `file` content types across all Gemini models.
|
||||
:::
|
||||
|
||||
## Video Metadata Control
|
||||
|
||||
For Gemini 3+ models, LiteLLM supports fine-grained video processing control through the `video_metadata` field. This allows you to specify frame extraction rates and time ranges for video analysis.
|
||||
LiteLLM supports fine-grained video processing control through the `video_metadata` field for all Gemini models (1.x, 2.x, 3+). This allows you to specify frame extraction rates and time ranges for video analysis.
|
||||
|
||||
**Supported `video_metadata` parameters:**
|
||||
|
||||
|
|
@ -2168,8 +2168,11 @@ For Gemini 3+ models, LiteLLM supports fine-grained video processing control thr
|
|||
- `fps` remains unchanged
|
||||
:::
|
||||
|
||||
:::tip
|
||||
Video clipping (`start_offset`/`end_offset`) and frame rate control (`fps`) are supported by all Gemini models, but analysis quality is significantly higher with the **Gemini 2.5 series** (e.g., `gemini-2.5-flash`, `gemini-2.5-pro`).
|
||||
:::
|
||||
|
||||
:::warning
|
||||
- **Gemini 3+ Only:** This feature is only available for Gemini 3.0 and newer models
|
||||
- **Video Files Recommended:** While `video_metadata` is designed for video files, error handling for other media types is delegated to the Vertex AI API
|
||||
- **File Formats Supported:** Works with `gs://`, `https://`, and base64-encoded video files
|
||||
:::
|
||||
|
|
|
|||
1
docs/my-website/package-lock.json
generated
1
docs/my-website/package-lock.json
generated
|
|
@ -26,6 +26,7 @@
|
|||
},
|
||||
"devDependencies": {
|
||||
"@docusaurus/module-type-aliases": "3.8.1",
|
||||
"ajv": "^8.18.0",
|
||||
"dotenv": "16.6.1"
|
||||
},
|
||||
"engines": {
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@
|
|||
},
|
||||
"devDependencies": {
|
||||
"@docusaurus/module-type-aliases": "3.8.1",
|
||||
"ajv": "^8.18.0",
|
||||
"dotenv": "16.6.1"
|
||||
},
|
||||
"browserslist": {
|
||||
|
|
|
|||
|
|
@ -1060,6 +1060,7 @@ const sidebars = {
|
|||
},
|
||||
items: [
|
||||
"routing",
|
||||
"adaptive_router",
|
||||
"scheduler",
|
||||
"proxy/auto_routing",
|
||||
"proxy/load_balancing",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,39 @@
|
|||
-- One row per (router, request_type, model). Hot path on every routing decision.
|
||||
CREATE TABLE "LiteLLM_AdaptiveRouterState" (
|
||||
router_name TEXT NOT NULL,
|
||||
request_type TEXT NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
alpha DOUBLE PRECISION NOT NULL,
|
||||
beta DOUBLE PRECISION NOT NULL,
|
||||
total_samples INTEGER NOT NULL DEFAULT 0,
|
||||
last_updated_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (router_name, request_type, model_name)
|
||||
);
|
||||
|
||||
-- One row per (session, router, model). Updated per turn via the queue.
|
||||
CREATE TABLE "LiteLLM_AdaptiveRouterSession" (
|
||||
session_id TEXT NOT NULL,
|
||||
router_name TEXT NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
classified_type TEXT NOT NULL,
|
||||
misalignment_count INTEGER NOT NULL DEFAULT 0,
|
||||
stagnation_count INTEGER NOT NULL DEFAULT 0,
|
||||
disengagement_count INTEGER NOT NULL DEFAULT 0,
|
||||
satisfaction_count INTEGER NOT NULL DEFAULT 0,
|
||||
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||
loop_count INTEGER NOT NULL DEFAULT 0,
|
||||
exhaustion_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_user_content TEXT,
|
||||
last_assistant_content TEXT,
|
||||
tool_call_history JSONB NOT NULL DEFAULT '[]',
|
||||
pending_tool_calls JSONB NOT NULL DEFAULT '{}',
|
||||
turn_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_processed_turn INTEGER NOT NULL DEFAULT -1,
|
||||
clean_credit_awarded BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
terminal_status INTEGER,
|
||||
last_activity_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (session_id, router_name, model_name)
|
||||
);
|
||||
|
||||
CREATE INDEX "idx_adaptive_router_session_activity"
|
||||
ON "LiteLLM_AdaptiveRouterSession" (last_activity_at);
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
|
|
@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
|
|||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
@@id([user_id, team_id])
|
||||
|
|
@ -1223,3 +1224,46 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
// Per-(router, request_type, model) Beta posterior for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterState {
|
||||
router_name String
|
||||
request_type String
|
||||
model_name String
|
||||
alpha Float
|
||||
beta Float
|
||||
total_samples Int @default(0)
|
||||
last_updated_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([router_name, request_type, model_name])
|
||||
}
|
||||
|
||||
// Per-(session, router, model) signal counters for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterSession {
|
||||
session_id String
|
||||
router_name String
|
||||
model_name String
|
||||
classified_type String
|
||||
|
||||
misalignment_count Int @default(0)
|
||||
stagnation_count Int @default(0)
|
||||
disengagement_count Int @default(0)
|
||||
satisfaction_count Int @default(0)
|
||||
failure_count Int @default(0)
|
||||
loop_count Int @default(0)
|
||||
exhaustion_count Int @default(0)
|
||||
|
||||
last_user_content String?
|
||||
last_assistant_content String?
|
||||
tool_call_history Json @default("[]")
|
||||
pending_tool_calls Json @default("{}")
|
||||
|
||||
turn_count Int @default(0)
|
||||
last_processed_turn Int @default(-1)
|
||||
clean_credit_awarded Boolean @default(false)
|
||||
terminal_status Int?
|
||||
last_activity_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([session_id, router_name, model_name])
|
||||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -30,6 +30,26 @@ def _get_prisma_env() -> dict:
|
|||
return prisma_env
|
||||
|
||||
|
||||
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
|
||||
|
||||
|
||||
def _migration_timestamp(name: str) -> int:
|
||||
"""Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name.
|
||||
|
||||
Returns 0 if the name doesn't match the Prisma pattern — unexpected-format
|
||||
entries sort as "oldest" and are treated as historical.
|
||||
"""
|
||||
m = _MIGRATION_TS_RE.match(name)
|
||||
return int(m.group(1)) if m else 0
|
||||
|
||||
|
||||
def _max_migration_timestamp(names) -> int:
|
||||
"""Max timestamp in a set/list of migration names (0 if empty)."""
|
||||
if not names:
|
||||
return 0
|
||||
return max(_migration_timestamp(n) for n in names)
|
||||
|
||||
|
||||
def _get_prisma_command() -> str:
|
||||
"""Get the Prisma command to use, bypassing Python wrapper in offline mode."""
|
||||
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
|
||||
|
|
@ -383,18 +403,301 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def setup_database(use_migrate: bool = False) -> bool:
|
||||
def _strip_prisma_query_params(url: str) -> str:
|
||||
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
|
||||
schema, etc.) from DATABASE_URL so psycopg can parse it."""
|
||||
from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
|
||||
|
||||
parsed = urlparse(url)
|
||||
if not parsed.query:
|
||||
return url
|
||||
libpq_params = {
|
||||
"sslmode",
|
||||
"sslcert",
|
||||
"sslkey",
|
||||
"sslrootcert",
|
||||
"sslpassword",
|
||||
"application_name",
|
||||
"connect_timeout",
|
||||
"client_encoding",
|
||||
"options",
|
||||
"service",
|
||||
"gssencmode",
|
||||
"krbsrvname",
|
||||
"target_session_attrs",
|
||||
}
|
||||
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
|
||||
return urlunparse(parsed._replace(query=urlencode(kept)))
|
||||
|
||||
@staticmethod
|
||||
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
|
||||
"""
|
||||
Log a warning if _prisma_migrations contains applied migrations with
|
||||
timestamps newer than every migration this build ships.
|
||||
|
||||
This is informational only for the v2 resolver — it tells the operator
|
||||
the DB was likely migrated by a newer deployment, which is usually a
|
||||
signal that this (older) version shouldn't run against it. We do NOT
|
||||
block startup: many users have weird _prisma_migrations state from
|
||||
prior thrashing bugs, and blocking them would be a breaking change.
|
||||
|
||||
Safe no-op if psycopg isn't installed or DB isn't reachable.
|
||||
"""
|
||||
database_url = os.getenv("DATABASE_URL")
|
||||
if not database_url:
|
||||
return
|
||||
|
||||
try:
|
||||
import psycopg
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
known = set(ProxyExtrasDBManager._get_migration_names(migrations_dir))
|
||||
|
||||
try:
|
||||
# autocommit=True keeps the SELECT outside a transaction. Without
|
||||
# it, psycopg3's `with conn` calls COMMIT on clean exit — which
|
||||
# fails after `UndefinedTable` (fresh DB) leaves the transaction
|
||||
# in an aborted state.
|
||||
with psycopg.connect(
|
||||
cleaned_url, connect_timeout=10, autocommit=True
|
||||
) as conn:
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT migration_name FROM _prisma_migrations "
|
||||
"WHERE finished_at IS NOT NULL AND rolled_back_at IS NULL"
|
||||
).fetchall()
|
||||
except psycopg.errors.UndefinedTable:
|
||||
return
|
||||
except (psycopg.OperationalError, psycopg.DatabaseError):
|
||||
# Swallow connection failures AND any other DB-layer error
|
||||
# (e.g. InsufficientPrivilege if the runtime user lacks SELECT
|
||||
# on _prisma_migrations). This is an informational check —
|
||||
# never block startup on it.
|
||||
return
|
||||
|
||||
applied = {r[0] for r in rows}
|
||||
unknown = applied - known
|
||||
if not unknown:
|
||||
return
|
||||
|
||||
head_newest_ts = _max_migration_timestamp(known)
|
||||
hostile = {
|
||||
name for name in unknown if _migration_timestamp(name) > head_newest_ts
|
||||
}
|
||||
if not hostile:
|
||||
return
|
||||
|
||||
sorted_hostile = sorted(hostile)
|
||||
logger.warning(
|
||||
"Database has %d migration(s) applied that are NEWER than any "
|
||||
"migration this LiteLLM version ships. This usually means the "
|
||||
"database was migrated by a newer LiteLLM deployment. Some API "
|
||||
"endpoints may fail because this proxy's Prisma client does not "
|
||||
"know about those schema changes. Consider upgrading this "
|
||||
"deployment. Unknown: %s",
|
||||
len(hostile),
|
||||
", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _setup_database_v2(use_migrate: bool) -> bool:
|
||||
"""
|
||||
v2 migration resolver (opt-in via --use_v2_migration_resolver).
|
||||
|
||||
Runs `prisma migrate deploy` and handles standard recovery paths
|
||||
(P3005 baseline, P3009/P3018 idempotent errors). Critically, it does
|
||||
NOT call `_resolve_all_migrations` — the diff-and-force recovery that
|
||||
caused schema thrashing when two LiteLLM versions contended for the
|
||||
same DB during rolling deploys.
|
||||
|
||||
Ahead-of-HEAD state (DB has migrations newer than this build ships)
|
||||
is logged as a warning, not a fatal error — users whose DBs got into
|
||||
weird shapes from the old thrashing should still be able to start.
|
||||
"""
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
|
||||
|
||||
if not use_migrate:
|
||||
# Preserve `prisma db push` path unchanged.
|
||||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=60,
|
||||
check=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
return True
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
subprocess.TimeoutExpired,
|
||||
) as e:
|
||||
# Re-raise as RuntimeError so proxy_cli.py's
|
||||
# `except RuntimeError` catches it and exits cleanly.
|
||||
raise RuntimeError(f"prisma db push failed.\n\nDetail: {e}") from e
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
||||
# Informational — never blocks.
|
||||
ProxyExtrasDBManager._warn_if_db_ahead_of_head(migrations_dir)
|
||||
|
||||
original_dir = os.getcwd()
|
||||
os.chdir(migrations_dir)
|
||||
try:
|
||||
for attempt in range(4):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=60,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
|
||||
return True
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.info(
|
||||
f"prisma migrate deploy attempt {attempt + 1} timed out, retrying"
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
continue
|
||||
|
||||
except subprocess.CalledProcessError as e:
|
||||
stderr = e.stderr or ""
|
||||
|
||||
if "P3005" in stderr and "database schema is not empty" in stderr:
|
||||
logger.info(
|
||||
"Schema exists but no migrations ledger — creating baseline"
|
||||
)
|
||||
ProxyExtrasDBManager._create_baseline_migration(schema_path)
|
||||
continue
|
||||
|
||||
if "P3009" in stderr:
|
||||
migration_match = re.search(r"`(\d+_\S+?)`", stderr)
|
||||
if (
|
||||
migration_match
|
||||
and ProxyExtrasDBManager._is_idempotent_error(stderr)
|
||||
):
|
||||
name = migration_match.group(1)
|
||||
logger.info(
|
||||
f"Migration {name} failed idempotently — marking applied and retrying"
|
||||
)
|
||||
try:
|
||||
ProxyExtrasDBManager._roll_back_migration(name)
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
subprocess.TimeoutExpired,
|
||||
):
|
||||
pass # may already be rolled-back
|
||||
try:
|
||||
ProxyExtrasDBManager._resolve_specific_migration(name)
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
subprocess.TimeoutExpired,
|
||||
) as resolve_err:
|
||||
# We're already inside the outer
|
||||
# `except CalledProcessError` handler —
|
||||
# re-raising CalledProcessError from here
|
||||
# would escape as itself, bypassing
|
||||
# proxy_cli.py's `except RuntimeError`.
|
||||
raise RuntimeError(
|
||||
f"Failed to mark migration {name} as applied "
|
||||
f"after idempotent recovery. Manual "
|
||||
f"intervention may be required.\n\n"
|
||||
f"Detail: {resolve_err}"
|
||||
) from resolve_err
|
||||
continue
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
if "P3018" in stderr:
|
||||
if ProxyExtrasDBManager._is_permission_error(stderr):
|
||||
raise RuntimeError(
|
||||
"Database migration failed due to insufficient "
|
||||
"permissions. Please grant the required privileges "
|
||||
f"and retry.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
migration_match = re.search(
|
||||
r"Migration name: (\d+_\S+)", stderr
|
||||
)
|
||||
if (
|
||||
migration_match
|
||||
and ProxyExtrasDBManager._is_idempotent_error(stderr)
|
||||
):
|
||||
name = migration_match.group(1)
|
||||
logger.info(
|
||||
f"Migration {name} SQL hit idempotent error — marking applied and retrying"
|
||||
)
|
||||
try:
|
||||
ProxyExtrasDBManager._roll_back_migration(name)
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
subprocess.TimeoutExpired,
|
||||
):
|
||||
pass # may already be rolled-back
|
||||
try:
|
||||
ProxyExtrasDBManager._resolve_specific_migration(name)
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
subprocess.TimeoutExpired,
|
||||
) as resolve_err:
|
||||
raise RuntimeError(
|
||||
f"Failed to mark migration {name} as applied "
|
||||
f"after idempotent recovery. Manual "
|
||||
f"intervention may be required.\n\n"
|
||||
f"Detail: {resolve_err}"
|
||||
) from resolve_err
|
||||
continue
|
||||
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
raise RuntimeError(
|
||||
"Database migration failed and cannot be auto-recovered. "
|
||||
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
|
||||
) from e
|
||||
|
||||
raise RuntimeError(
|
||||
"Database migration failed after 4 attempts (retry loop "
|
||||
"exhausted by timeouts or repeated idempotent-recovery "
|
||||
"continues). Check database connectivity, load, and "
|
||||
"_prisma_migrations ledger state."
|
||||
)
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
||||
@staticmethod
|
||||
def setup_database(
|
||||
use_migrate: bool = False, use_v2_resolver: bool = False
|
||||
) -> bool:
|
||||
"""
|
||||
Set up the database using either prisma migrate or prisma db push
|
||||
Uses migrations from litellm-proxy-extras package
|
||||
|
||||
Args:
|
||||
schema_path (str): Path to the Prisma schema file
|
||||
use_migrate (bool): Whether to use prisma migrate instead of db push
|
||||
use_migrate: Whether to use prisma migrate instead of db push
|
||||
use_v2_resolver: Opt into the v2 migration resolver (safer during
|
||||
rolling deploys; does not run the diff-and-force recovery
|
||||
that causes schema thrashing). Defaults to False for
|
||||
backwards compatibility.
|
||||
|
||||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
"""
|
||||
if use_v2_resolver:
|
||||
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
|
||||
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)
|
||||
|
||||
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
|
||||
for attempt in range(4):
|
||||
original_dir = os.getcwd()
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
|
||||
This is a runbook for creating and running database migrations for the LiteLLM proxy. For use for litellm engineers only.
|
||||
|
||||
> **AI AGENTS / ASSISTANTS:** If the script refuses with either a "STALE BRANCH" or "DESTRUCTIVE MIGRATION DETECTED" error, **do NOT** bypass it on your own (no `git rebase`, no `--skip-freshness-check`, no `--allow-destructive`). Surface the error to the human operator and wait for their explicit confirmation. See the [Branch freshness](#branch-freshness-check) and [Destructive migrations](#destructive-migrations-drop-column--drop-table) sections below.
|
||||
|
||||
## Step 0: Sync All `schema.prisma` Files
|
||||
|
||||
Before doing anything else, make sure all `schema.prisma` files in the repo are in sync. There are multiple copies that must match:
|
||||
|
|
@ -46,10 +48,50 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
|
|||
|
||||
## What It Does
|
||||
|
||||
1. Creates temp PostgreSQL DB
|
||||
2. Applies existing migrations
|
||||
3. Compares with `schema.prisma`
|
||||
4. Generates new migration if changes found
|
||||
1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check))
|
||||
2. Creates temp PostgreSQL DB
|
||||
3. Applies existing migrations
|
||||
4. Compares with `schema.prisma`
|
||||
5. Generates new migration if changes found
|
||||
6. **Refuses to write if the generated SQL contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`**, unless `--allow-destructive` is passed
|
||||
|
||||
## 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.
|
||||
|
||||
Flags:
|
||||
|
||||
- `--base-branch <name>` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`.
|
||||
- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base.
|
||||
|
||||
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
|
||||
```
|
||||
2. Re-run `run_migration.py`.
|
||||
|
||||
> **AI AGENTS / ASSISTANTS:** Do **not** auto-rebase or auto-pass `--skip-freshness-check`. A bad conflict resolution during rebase can itself drop the human's in-progress schema edits. Surface the error and wait for explicit confirmation.
|
||||
|
||||
## Destructive Migrations (DROP COLUMN / DROP TABLE / DROP INDEX)
|
||||
|
||||
If the generated diff contains `DROP COLUMN`, `DROP TABLE`, or `DROP INDEX`, `run_migration.py` exits non-zero and refuses to write the migration file. A previous incident saw newly-added columns silently dropped by a stale branch and merged to main — this guard exists to prevent a repeat.
|
||||
|
||||
When the guard fires:
|
||||
|
||||
1. Run `git fetch origin && git status` — confirm your branch is up to date with the base branch.
|
||||
2. Re-check all `schema.prisma` files are in sync (Step 0).
|
||||
3. Review EACH `DROP` statement printed in the error — is it actually intended?
|
||||
4. Only if the drops are genuinely intentional, re-run with the flag:
|
||||
|
||||
```bash
|
||||
uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_name" --allow-destructive
|
||||
```
|
||||
|
||||
> **AI AGENTS / ASSISTANTS:** Do **not** automatically re-run the command with `--allow-destructive`. If the guard fires while you are driving the runbook for a human, stop, show them the error, and wait for their explicit confirmation before passing the flag. Auto-passing `--allow-destructive` is the exact failure mode this guard exists to prevent.
|
||||
|
||||
## Common Fixes
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.67"
|
||||
version = "0.4.68"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -25,7 +25,7 @@ required-version = "==0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.67"
|
||||
version = "0.4.68"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
Normal file
242
litellm-proxy-extras/tests/test_setup_database_fail_fast.py
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
"""Regression tests for ProxyExtrasDBManager v2 migration resolver.
|
||||
|
||||
The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
|
||||
`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1
|
||||
(default) behavior is unchanged from pre-fix.
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm_proxy_extras.utils import (
|
||||
ProxyExtrasDBManager,
|
||||
_max_migration_timestamp,
|
||||
_migration_timestamp,
|
||||
)
|
||||
|
||||
|
||||
def _fake_migrate_deploy_failure(returncode: int, stderr: str):
|
||||
def _run(*args, **kwargs):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=returncode,
|
||||
cmd=args[0],
|
||||
stderr=stderr,
|
||||
output="",
|
||||
)
|
||||
|
||||
return _run
|
||||
|
||||
|
||||
def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a permission failure during migrate deploy raises RuntimeError."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3018\nMigration name: 20250326162113_baseline\n"
|
||||
"Database error code: 42501\npermission denied for schema public"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="permission"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a non-idempotent migration failure raises (no silent recovery)."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
|
||||
'Reason: syntax error at or near "BRKN" LINE 42'
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_removes_connection_limit():
|
||||
"""DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
|
||||
url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
|
||||
stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
|
||||
assert "connection_limit" not in stripped
|
||||
assert "pool_timeout" not in stripped
|
||||
assert "sslmode=require" in stripped
|
||||
|
||||
|
||||
def test_strip_prisma_query_params_passthrough_no_query():
|
||||
"""URLs without query strings are returned unchanged."""
|
||||
url = "postgresql://u:p@h:5432/db"
|
||||
assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url
|
||||
|
||||
|
||||
def test_migration_timestamp_extracts_leading_digits():
|
||||
assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
|
||||
assert _migration_timestamp("20250326162113_baseline") == 20250326162113
|
||||
|
||||
|
||||
def test_migration_timestamp_returns_zero_on_malformed():
|
||||
assert _migration_timestamp("0_init") == 0
|
||||
assert _migration_timestamp("not_a_migration") == 0
|
||||
|
||||
|
||||
def test_max_migration_timestamp():
|
||||
names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
|
||||
assert _max_migration_timestamp(names) == 20260415000000
|
||||
|
||||
|
||||
def test_max_migration_timestamp_empty_set():
|
||||
assert _max_migration_timestamp(set()) == 0
|
||||
|
||||
|
||||
def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v1 (default) continues to call _resolve_all_migrations on the happy path.
|
||||
|
||||
This is the existing buggy behavior — we're not fixing it in v1, only
|
||||
offering v2 as opt-in. This test pins the default so that a future
|
||||
inadvertent default flip is caught.
|
||||
"""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
# Stub `prisma migrate deploy` to claim success with pending migrations
|
||||
# applied, which is the code path that triggers the legacy post-migration
|
||||
# sanity check (a call to _resolve_all_migrations).
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
def fake_run(cmd, *args, **kwargs):
|
||||
return FakeResult()
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
|
||||
def fake_resolve(*args, **kwargs):
|
||||
resolve_called["n"] += 1
|
||||
|
||||
monkeypatch.setattr("subprocess.run", fake_run)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True) # v2 flag NOT set
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"
|
||||
|
||||
|
||||
def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
|
||||
"""v2: a failing `prisma db push` must raise RuntimeError, not leak
|
||||
CalledProcessError past proxy_cli.py's `except RuntimeError`."""
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
stderr = "db push error"
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(RuntimeError, match="prisma db push failed"):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
|
||||
"""_warn_if_db_ahead_of_head must never raise — it's informational.
|
||||
|
||||
Non-connection DB errors (e.g. InsufficientPrivilege from a user
|
||||
without SELECT on _prisma_migrations) must be caught, not propagated.
|
||||
"""
|
||||
import psycopg
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class _FakeConn:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def execute(self, *a, **kw):
|
||||
# Simulate an InsufficientPrivilege (subclass of DatabaseError).
|
||||
raise psycopg.errors.InsufficientPrivilege("permission denied")
|
||||
|
||||
def _fake_connect(*a, **kw):
|
||||
return _FakeConn()
|
||||
|
||||
monkeypatch.setattr("psycopg.connect", _fake_connect)
|
||||
|
||||
# Must not raise.
|
||||
ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))
|
||||
|
||||
|
||||
def test_v2_resolve_specific_migration_failure_raises_runtime_error(
|
||||
monkeypatch, tmp_path
|
||||
):
|
||||
"""If marking a migration as applied fails inside P3009 idempotent
|
||||
recovery, the subprocess error must be re-raised as RuntimeError so
|
||||
proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
|
||||
)
|
||||
|
||||
# First call: migrate deploy -> P3009 idempotent error.
|
||||
# Recovery path tries _resolve_specific_migration; that also raises.
|
||||
def _failing_resolve(*a, **kw):
|
||||
raise subprocess.CalledProcessError(
|
||||
returncode=1,
|
||||
cmd="prisma migrate resolve --applied",
|
||||
stderr="resolve failed",
|
||||
output="",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
|
||||
)
|
||||
|
||||
stderr = (
|
||||
"Error: P3009\nMigration `20260101000000_some_migration` failed\n"
|
||||
"relation already exists"
|
||||
)
|
||||
with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
|
||||
with pytest.raises(
|
||||
RuntimeError, match="Failed to mark migration .* as applied"
|
||||
):
|
||||
ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
|
||||
|
||||
def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
|
||||
"""v2 must never call _resolve_all_migrations — that's the bug it fixes."""
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
|
||||
)
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
|
||||
(tmp_path / "schema.prisma").write_text("// stub")
|
||||
|
||||
class FakeResult:
|
||||
stdout = "Applied migration.\n"
|
||||
stderr = ""
|
||||
|
||||
monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())
|
||||
|
||||
resolve_called = {"n": 0}
|
||||
monkeypatch.setattr(
|
||||
ProxyExtrasDBManager,
|
||||
"_resolve_all_migrations",
|
||||
lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
|
||||
)
|
||||
|
||||
ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
|
||||
assert ok is True
|
||||
assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"
|
||||
|
|
@ -1502,6 +1502,9 @@ if TYPE_CHECKING:
|
|||
from .llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeMessagesConfig as AmazonAnthropicClaudeMessagesConfig,
|
||||
)
|
||||
from .llms.bedrock.messages.mantle_transformation import (
|
||||
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
|
||||
)
|
||||
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
|
||||
from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig
|
||||
from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
|
|
|
|||
|
|
@ -171,6 +171,7 @@ LLM_CONFIG_NAMES = (
|
|||
"CohereChatConfig",
|
||||
"AnthropicMessagesConfig",
|
||||
"AmazonAnthropicClaudeMessagesConfig",
|
||||
"AmazonMantleMessagesConfig",
|
||||
"TogetherAIConfig",
|
||||
"NLPCloudConfig",
|
||||
"VertexGeminiConfig",
|
||||
|
|
@ -715,6 +716,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
|||
".llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation",
|
||||
"AmazonAnthropicClaudeMessagesConfig",
|
||||
),
|
||||
"AmazonMantleMessagesConfig": (
|
||||
".llms.bedrock.messages.mantle_transformation",
|
||||
"AmazonMantleMessagesConfig",
|
||||
),
|
||||
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
|
||||
"NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"),
|
||||
"VertexGeminiConfig": (
|
||||
|
|
|
|||
|
|
@ -164,6 +164,7 @@ MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset(
|
|||
LITELLM_UI_ALLOW_HEADERS = [
|
||||
"x-litellm-semantic-filter",
|
||||
"x-litellm-semantic-filter-tools",
|
||||
"x-litellm-adaptive-router-model",
|
||||
]
|
||||
|
||||
# Gemini model-specific minimal thinking budget constants
|
||||
|
|
|
|||
|
|
@ -410,6 +410,7 @@ def image_generation( # noqa: PLR0915
|
|||
litellm.LlmProviders.RUNWAYML,
|
||||
litellm.LlmProviders.VERTEX_AI,
|
||||
litellm.LlmProviders.OPENROUTER,
|
||||
litellm.LlmProviders.DASHSCOPE,
|
||||
):
|
||||
if image_generation_config is None:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from datetime import datetime
|
|||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
ClassVar,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
|
|
@ -12,6 +13,7 @@ from typing import (
|
|||
)
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.guardrails import (
|
||||
|
|
@ -81,6 +83,9 @@ class ModifyResponseException(Exception):
|
|||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: Optional[str] = None,
|
||||
|
|
@ -637,6 +642,13 @@ class CustomGuardrail(CustomLogger):
|
|||
if isinstance(item, dict):
|
||||
item.pop("secret_fields", None)
|
||||
|
||||
# Default-safe behavior: never persist raw matched spans in standard
|
||||
# guardrail logging payloads (single shared implementation; Bedrock hooks pass
|
||||
# raw provider JSON so redaction is not duplicated upstream).
|
||||
clean_guardrail_response = redact_nested_match_and_regex_keys(
|
||||
clean_guardrail_response
|
||||
)
|
||||
|
||||
slg = StandardLoggingGuardrailInformation(
|
||||
guardrail_name=self.guardrail_name,
|
||||
guardrail_provider=guardrail_provider,
|
||||
|
|
|
|||
|
|
@ -1615,6 +1615,14 @@ class OpenTelemetry(CustomLogger):
|
|||
value=response_id,
|
||||
)
|
||||
|
||||
litellm_call_id = standard_logging_payload.get("litellm_call_id")
|
||||
if litellm_call_id:
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
key="litellm.call_id",
|
||||
value=litellm_call_id,
|
||||
)
|
||||
|
||||
# The model used to generate the response.
|
||||
if response_obj and response_obj.get("model"):
|
||||
self.safe_set_attribute(
|
||||
|
|
@ -2281,6 +2289,10 @@ class OpenTelemetry(CustomLogger):
|
|||
# Remove trailing slash
|
||||
endpoint = endpoint.rstrip("/")
|
||||
|
||||
# Splunk Observability Cloud OTLP/HTTP uses /v2/trace/otlp (not /v1/traces). Do not rewrite.
|
||||
if signal_type == "traces" and "/v2/trace/otlp" in endpoint:
|
||||
return endpoint
|
||||
|
||||
# Check if endpoint already ends with the correct signal path
|
||||
target_path = f"/v1/{signal_type}"
|
||||
if endpoint.endswith(target_path):
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
# What is this?
|
||||
## Helper utilities
|
||||
import copy
|
||||
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -435,3 +436,42 @@ def filter_internal_params(
|
|||
|
||||
# Filter out internal parameters
|
||||
return {k: v for k, v in data.items() if k not in internal_params}
|
||||
|
||||
|
||||
def redact_nested_match_and_regex_keys(
|
||||
payload: Union[dict, List[Any], str, None],
|
||||
) -> Union[dict, List[Any], str, None]:
|
||||
"""
|
||||
Deep-copy `payload` and replace every `match` / `regex` string field with
|
||||
"[REDACTED]" anywhere in nested dict/list structures.
|
||||
|
||||
Used for guardrail spend/compliance logging so raw spans are not persisted.
|
||||
"""
|
||||
if payload is None or isinstance(payload, str):
|
||||
return payload
|
||||
try:
|
||||
redacted: Union[dict, List[Any], str, None] = copy.deepcopy(payload)
|
||||
except Exception:
|
||||
return payload
|
||||
|
||||
# Iterative traversal; `seen` guards against cyclic refs preserved by deepcopy.
|
||||
try:
|
||||
seen: set = set()
|
||||
stack: List[Any] = [redacted]
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
node_id = id(node)
|
||||
if node_id in seen:
|
||||
continue
|
||||
seen.add(node_id)
|
||||
if isinstance(node, dict):
|
||||
if "match" in node:
|
||||
node["match"] = "[REDACTED]"
|
||||
if "regex" in node:
|
||||
node["regex"] = "[REDACTED]"
|
||||
stack.extend(node.values())
|
||||
elif isinstance(node, list):
|
||||
stack.extend(node)
|
||||
except Exception:
|
||||
return payload
|
||||
return redacted
|
||||
|
|
|
|||
|
|
@ -5512,6 +5512,8 @@ def get_standard_logging_object_payload(
|
|||
|
||||
payload: StandardLoggingPayload = StandardLoggingPayload(
|
||||
id=str(id),
|
||||
litellm_call_id=kwargs.get("litellm_call_id")
|
||||
or litellm_params.get("litellm_call_id"),
|
||||
trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id(
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
|
|
|
|||
|
|
@ -370,11 +370,17 @@ class LoggingWorker:
|
|||
self._running_tasks.clear()
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Flush the logging queue."""
|
||||
"""Flush the logging queue.
|
||||
|
||||
Waits until every enqueued task has completed. ``queue.join()`` blocks
|
||||
on the queue's unfinished-task counter (decremented by ``task_done()``),
|
||||
so it correctly handles items that have been dequeued but whose
|
||||
callback hasn't finished yet — ``queue.empty()`` would return True in
|
||||
that window and cause us to skip the wait.
|
||||
"""
|
||||
if self._queue is None:
|
||||
return
|
||||
while not self._queue.empty():
|
||||
await self._queue.join()
|
||||
await self._queue.join()
|
||||
|
||||
async def clear_queue(self):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -452,7 +452,14 @@ def update_messages_with_model_file_ids(
|
|||
for c in content:
|
||||
if c["type"] == "file":
|
||||
file_object = cast(ChatCompletionFileObject, c)
|
||||
file_object_file_field = file_object["file"]
|
||||
file_object_file_field = file_object.get("file")
|
||||
if not isinstance(file_object_file_field, dict):
|
||||
# Content block has `type: "file"` but not the
|
||||
# OpenAI Chat Completions shape (e.g. a LangChain
|
||||
# v1 standardized file block, or a provider-native
|
||||
# shape that also uses `type: "file"`). Nothing to
|
||||
# remap here, so skip instead of crashing.
|
||||
continue
|
||||
file_id = file_object_file_field.get("file_id")
|
||||
format = file_object_file_field.get(
|
||||
"format", get_format_from_file_id(file_id)
|
||||
|
|
@ -1060,7 +1067,12 @@ def get_file_ids_from_messages(messages: List[AllMessageValues]) -> List[str]:
|
|||
for c in content:
|
||||
if c["type"] == "file":
|
||||
file_object = cast(ChatCompletionFileObject, c)
|
||||
file_object_file_field = file_object["file"]
|
||||
file_object_file_field = file_object.get("file")
|
||||
if not isinstance(file_object_file_field, dict):
|
||||
# Content block has `type: "file"` but not the
|
||||
# OpenAI Chat Completions shape. No file_id to
|
||||
# extract, so skip instead of raising KeyError.
|
||||
continue
|
||||
file_id = file_object_file_field.get("file_id")
|
||||
if file_id:
|
||||
file_ids.append(file_id)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ import litellm.types
|
|||
import litellm.types.llms
|
||||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
|
||||
from litellm.types.files import get_file_extension_from_mime_type
|
||||
from litellm.types.llms.anthropic import *
|
||||
|
|
@ -3324,7 +3325,7 @@ def _load_image_from_url(image_url):
|
|||
try:
|
||||
# Send a GET request to the image URL
|
||||
client = HTTPHandler(concurrent_limit=1)
|
||||
response = client.get(image_url)
|
||||
response = safe_get(client, image_url)
|
||||
response.raise_for_status() # Raise an exception for HTTP errors
|
||||
|
||||
# Check the response's content type to ensure it is an image
|
||||
|
|
@ -3562,7 +3563,7 @@ class BedrockImageProcessor:
|
|||
params={"concurrent_limit": 1},
|
||||
)
|
||||
# Send a GET request to the image URL
|
||||
response = await client.get(image_url, follow_redirects=True)
|
||||
response = await async_safe_get(client, image_url)
|
||||
response.raise_for_status() # Raise an exception for HTTP errors
|
||||
|
||||
return BedrockImageProcessor._post_call_image_processing(
|
||||
|
|
@ -3577,7 +3578,7 @@ class BedrockImageProcessor:
|
|||
try:
|
||||
client = HTTPHandler(concurrent_limit=1)
|
||||
# Send a GET request to the image URL
|
||||
response = client.get(image_url, follow_redirects=True)
|
||||
response = safe_get(client, image_url)
|
||||
response.raise_for_status() # Raise an exception for HTTP errors
|
||||
|
||||
return BedrockImageProcessor._post_call_image_processing(
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.constants import (
|
|||
DEFAULT_IMAGE_HEIGHT,
|
||||
DEFAULT_IMAGE_TOKEN_COUNT,
|
||||
DEFAULT_IMAGE_WIDTH,
|
||||
MAX_IMAGE_URL_DOWNLOAD_SIZE_MB,
|
||||
MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES,
|
||||
MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES,
|
||||
MAX_TILE_HEIGHT,
|
||||
|
|
@ -215,7 +216,14 @@ def get_image_dimensions(
|
|||
try:
|
||||
client = _get_httpx_client()
|
||||
response = safe_get(client, data)
|
||||
img_data = response.read()
|
||||
max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024)
|
||||
content_length = response.headers.get("Content-Length")
|
||||
if content_length is not None and int(content_length) > max_bytes:
|
||||
pass # skip download; img_data stays None
|
||||
else:
|
||||
body = response.read()
|
||||
if len(body) <= max_bytes:
|
||||
img_data = body
|
||||
except Exception:
|
||||
pass
|
||||
if img_data is None:
|
||||
|
|
|
|||
|
|
@ -106,6 +106,44 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
updated_reasoning_effort["summary"] = effective_summary
|
||||
completion_kwargs["reasoning_effort"] = updated_reasoning_effort
|
||||
|
||||
@staticmethod
|
||||
def _normalize_reasoning_effort(
|
||||
completion_kwargs: Dict[str, Any],
|
||||
) -> None:
|
||||
"""
|
||||
Normalize reasoning_effort values based on target model capabilities.
|
||||
|
||||
Handles both string ("max") and dict ({"effort": "max", "summary": ...})
|
||||
formats. Uses model registry to check supports_xhigh/supports_minimal.
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.utils import (
|
||||
normalize_reasoning_effort_value,
|
||||
)
|
||||
|
||||
reasoning_effort = completion_kwargs.get("reasoning_effort")
|
||||
if reasoning_effort is None:
|
||||
return
|
||||
|
||||
model = cast(str, completion_kwargs.get("model", ""))
|
||||
custom_llm_provider = completion_kwargs.get("custom_llm_provider")
|
||||
|
||||
if isinstance(reasoning_effort, str):
|
||||
normalized = normalize_reasoning_effort_value(
|
||||
reasoning_effort, model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
if normalized != reasoning_effort:
|
||||
completion_kwargs["reasoning_effort"] = normalized
|
||||
elif isinstance(reasoning_effort, dict) and "effort" in reasoning_effort:
|
||||
effort = reasoning_effort["effort"]
|
||||
normalized = normalize_reasoning_effort_value(
|
||||
effort, model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
if normalized != effort:
|
||||
completion_kwargs["reasoning_effort"] = {
|
||||
**reasoning_effort,
|
||||
"effort": normalized,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _prepare_completion_kwargs(
|
||||
*,
|
||||
|
|
@ -163,6 +201,12 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
if output_format:
|
||||
request_data["output_format"] = output_format
|
||||
|
||||
# Extract output_config from extra_kwargs so the translator can use it
|
||||
# (e.g. output_config.effort for adaptive thinking → reasoning_effort)
|
||||
extra_kwargs = extra_kwargs or {}
|
||||
if "output_config" in extra_kwargs:
|
||||
request_data["output_config"] = extra_kwargs["output_config"]
|
||||
|
||||
(
|
||||
openai_request,
|
||||
tool_name_mapping,
|
||||
|
|
@ -202,6 +246,14 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
):
|
||||
completion_kwargs[key] = value
|
||||
|
||||
# Normalize reasoning_effort based on model capabilities
|
||||
# (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported)
|
||||
# Must run BEFORE _route_openai_thinking, which prepends "responses/"
|
||||
# to the model name and would break get_model_info() lookups.
|
||||
LiteLLMMessagesToCompletionTransformationHandler._normalize_reasoning_effort(
|
||||
completion_kwargs
|
||||
)
|
||||
|
||||
LiteLLMMessagesToCompletionTransformationHandler._route_openai_thinking_to_responses_api_if_needed(
|
||||
completion_kwargs,
|
||||
thinking=thinking,
|
||||
|
|
|
|||
|
|
@ -317,6 +317,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
"tools",
|
||||
"thinking",
|
||||
"output_format",
|
||||
"output_config",
|
||||
]
|
||||
|
||||
def _is_web_search_tool(self, tool: Dict[str, Any]) -> bool:
|
||||
|
|
@ -694,6 +695,11 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
return "low"
|
||||
else:
|
||||
return "minimal"
|
||||
elif thinking_type == "adaptive":
|
||||
# Adaptive thinking: effort is controlled by output_config.effort,
|
||||
# not budget_tokens. Return a default; caller should override with
|
||||
# output_config.effort when available.
|
||||
return "medium"
|
||||
|
||||
return None
|
||||
|
||||
|
|
@ -776,6 +782,8 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
return ChatCompletionToolChoiceObjectParam(
|
||||
type="function", function=tc_function_param
|
||||
)
|
||||
elif tool_choice["type"] == "none":
|
||||
return "none"
|
||||
else:
|
||||
raise ValueError(
|
||||
"Incompatible tool choice param submitted - {}".format(tool_choice)
|
||||
|
|
@ -1041,6 +1049,12 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
if not reasoning_effort:
|
||||
return
|
||||
|
||||
# For adaptive thinking, override with output_config.effort if available
|
||||
if isinstance(thinking, dict) and thinking.get("type") == "adaptive":
|
||||
output_config = anthropic_message_request.get("output_config")
|
||||
if isinstance(output_config, dict) and output_config.get("effort"):
|
||||
reasoning_effort = output_config["effort"]
|
||||
|
||||
summary = thinking.get("summary") if isinstance(thinking, dict) else None
|
||||
auto_summary = is_reasoning_auto_summary_enabled()
|
||||
if summary:
|
||||
|
|
|
|||
|
|
@ -24,6 +24,8 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.utils import ProviderConfigManager, client
|
||||
|
||||
from ..utils import is_reasoning_auto_summary_enabled
|
||||
|
||||
from ..adapters.handler import LiteLLMMessagesToCompletionTransformationHandler
|
||||
from ..responses_adapters.handler import LiteLLMMessagesToResponsesAPIHandler
|
||||
from .interceptors import get_messages_interceptors
|
||||
|
|
@ -441,6 +443,17 @@ def anthropic_messages_handler(
|
|||
params=local_vars
|
||||
)
|
||||
)
|
||||
if is_reasoning_auto_summary_enabled():
|
||||
thinking_param = anthropic_messages_optional_request_params.get("thinking")
|
||||
if (
|
||||
isinstance(thinking_param, dict)
|
||||
and thinking_param.get("type") != "disabled"
|
||||
):
|
||||
anthropic_messages_optional_request_params["thinking"] = {
|
||||
**thinking_param,
|
||||
"display": "summarized",
|
||||
}
|
||||
|
||||
return base_llm_http_handler.anthropic_messages_handler(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
|
|||
|
|
@ -72,6 +72,23 @@ def _build_responses_kwargs(
|
|||
anthropic_request = AnthropicMessagesRequest(**request_data) # type: ignore[typeddict-item]
|
||||
responses_kwargs = _ADAPTER.translate_request(anthropic_request)
|
||||
|
||||
# Normalize reasoning effort based on model capabilities
|
||||
# (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported)
|
||||
reasoning = responses_kwargs.get("reasoning")
|
||||
if isinstance(reasoning, dict) and "effort" in reasoning:
|
||||
from litellm.llms.anthropic.experimental_pass_through.utils import (
|
||||
normalize_reasoning_effort_value,
|
||||
)
|
||||
|
||||
effort = reasoning["effort"]
|
||||
normalized = normalize_reasoning_effort_value(
|
||||
effort,
|
||||
model=model,
|
||||
custom_llm_provider=(extra_kwargs or {}).get("custom_llm_provider"),
|
||||
)
|
||||
if normalized != effort:
|
||||
responses_kwargs["reasoning"] = {**reasoning, "effort": normalized}
|
||||
|
||||
if stream:
|
||||
responses_kwargs["stream"] = True
|
||||
|
||||
|
|
|
|||
|
|
@ -251,25 +251,41 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
|
||||
@staticmethod
|
||||
def translate_thinking_to_reasoning(
|
||||
thinking: Dict[str, Any]
|
||||
thinking: Dict[str, Any],
|
||||
output_config: Optional[Dict[str, Any]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Convert Anthropic thinking param to Responses API reasoning param.
|
||||
|
||||
thinking.budget_tokens maps to reasoning effort:
|
||||
>= 10000 -> high, >= 5000 -> medium, >= 2000 -> low, < 2000 -> minimal
|
||||
|
||||
For adaptive thinking, uses output_config.effort if available,
|
||||
otherwise defaults to medium.
|
||||
"""
|
||||
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
|
||||
if not isinstance(thinking, dict):
|
||||
return None
|
||||
budget = thinking.get("budget_tokens", 0)
|
||||
if budget >= 10000:
|
||||
effort = "high"
|
||||
elif budget >= 5000:
|
||||
|
||||
thinking_type = thinking.get("type")
|
||||
|
||||
if thinking_type == "adaptive":
|
||||
# Use output_config.effort if available
|
||||
effort = "medium"
|
||||
elif budget >= 2000:
|
||||
effort = "low"
|
||||
if isinstance(output_config, dict) and output_config.get("effort"):
|
||||
effort = output_config["effort"]
|
||||
elif thinking_type == "enabled":
|
||||
budget = thinking.get("budget_tokens", 0)
|
||||
if budget >= 10000:
|
||||
effort = "high"
|
||||
elif budget >= 5000:
|
||||
effort = "medium"
|
||||
elif budget >= 2000:
|
||||
effort = "low"
|
||||
else:
|
||||
effort = "minimal"
|
||||
else:
|
||||
effort = "minimal"
|
||||
return None
|
||||
|
||||
auto_summary = is_reasoning_auto_summary_enabled()
|
||||
result: Dict[str, Any] = {"effort": effort}
|
||||
summary = thinking.get("summary")
|
||||
|
|
@ -346,7 +362,11 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
# thinking -> reasoning
|
||||
thinking = anthropic_request.get("thinking")
|
||||
if isinstance(thinking, dict):
|
||||
reasoning = self.translate_thinking_to_reasoning(thinking)
|
||||
output_config = anthropic_request.get("output_config")
|
||||
reasoning = self.translate_thinking_to_reasoning(
|
||||
thinking,
|
||||
output_config=cast(Optional[Dict[str, Any]], output_config),
|
||||
)
|
||||
if reasoning:
|
||||
responses_kwargs["reasoning"] = reasoning
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import os
|
||||
from typing import Optional
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import ModelInfo
|
||||
|
||||
|
||||
def is_reasoning_auto_summary_enabled() -> bool:
|
||||
|
|
@ -9,3 +11,47 @@ def is_reasoning_auto_summary_enabled() -> bool:
|
|||
litellm.reasoning_auto_summary
|
||||
or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
|
||||
)
|
||||
|
||||
|
||||
def normalize_reasoning_effort_value(
|
||||
effort: str,
|
||||
model: str,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Normalize a reasoning effort value based on model capabilities.
|
||||
|
||||
Degradation chains:
|
||||
- "max" → max / xhigh / high
|
||||
- "xhigh" → xhigh / high
|
||||
- "minimal" → minimal / low
|
||||
- other values pass through unchanged
|
||||
"""
|
||||
if effort not in ("max", "xhigh", "minimal"):
|
||||
return effort
|
||||
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
model_info: Optional[ModelInfo] = None
|
||||
try:
|
||||
model_info = get_model_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
except Exception:
|
||||
model_info = None
|
||||
|
||||
if effort == "max":
|
||||
if model_info and model_info.get("supports_max_reasoning_effort"):
|
||||
return "max"
|
||||
if model_info and model_info.get("supports_xhigh_reasoning_effort"):
|
||||
return "xhigh"
|
||||
return "high"
|
||||
elif effort == "xhigh":
|
||||
if model_info and model_info.get("supports_xhigh_reasoning_effort"):
|
||||
return "xhigh"
|
||||
return "high"
|
||||
elif effort == "minimal":
|
||||
if model_info and model_info.get("supports_minimal_reasoning_effort"):
|
||||
return "minimal"
|
||||
return "low"
|
||||
return "medium"
|
||||
|
|
|
|||
|
|
@ -40,9 +40,22 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
|
|||
Accepts both explicit gpt-5 model names and the ``gpt5_series/`` prefix
|
||||
used for manual routing.
|
||||
"""
|
||||
# gpt-5-chat* is a chat model and shouldn't go through GPT-5 reasoning restrictions.
|
||||
# The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
|
||||
# …) are regular chat models: they support temperature and tool_choice but NOT
|
||||
# reasoning_effort. They must NOT be routed through the GPT-5 reasoning path.
|
||||
#
|
||||
# Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
|
||||
# models and must stay on the GPT-5 path. The distinguishing feature is that
|
||||
# the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
|
||||
# (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
|
||||
# number (i.e. "gpt-5.<digit>-chat").
|
||||
#
|
||||
# Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
|
||||
# than a substring check) makes this boundary explicit and avoids any ambiguity
|
||||
# if future model names coincidentally contain "gpt-5-chat" as an interior run.
|
||||
_normalized = model.split("/")[-1] # strip provider prefix, e.g. "azure/"
|
||||
return (
|
||||
"gpt-5" in model and "gpt-5-chat" not in model
|
||||
"gpt-5" in model and not _normalized.startswith("gpt-5-chat")
|
||||
) or "gpt5_series" in model
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> List[str]:
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
api_key = (
|
||||
api_key
|
||||
|
|
|
|||
|
|
@ -65,6 +65,8 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Validate Azure AI Foundry environment and set up authentication
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Validate Azure AI Foundry environment and set up authentication
|
||||
|
|
|
|||
|
|
@ -67,6 +67,8 @@ class BaseImageEditConfig(ABC):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
return {}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, List, Literal, Optional, Union, cast
|
||||
|
||||
|
|
@ -294,7 +295,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
|||
raise ValueError(f"Invalid ARN format: {batch_id}")
|
||||
|
||||
region = arn_parts[3]
|
||||
# arn_parts[5] contains "model-invocation-job/{jobId}"
|
||||
if not re.match(r"^[a-z][a-z0-9-]*$", region):
|
||||
raise ValueError(f"Invalid region in ARN: {batch_id}")
|
||||
|
||||
# Build the endpoint URL for GetModelInvocationJob
|
||||
# AWS API format: GET /model-invocation-job/{jobIdentifier}
|
||||
|
|
|
|||
0
litellm/llms/bedrock/chat/mantle/__init__.py
Normal file
0
litellm/llms/bedrock/chat/mantle/__init__.py
Normal file
91
litellm/llms/bedrock/chat/mantle/transformation.py
Normal file
91
litellm/llms/bedrock/chat/mantle/transformation.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
"""
|
||||
Transformation for Bedrock Mantle (Claude Mythos Preview)
|
||||
|
||||
https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-anthropic-claude-mythos-preview.html
|
||||
|
||||
The bedrock-mantle endpoint uses the Anthropic Messages API format but is served
|
||||
at a different endpoint (bedrock-mantle.{region}.api.aws) with AWS SigV4 auth.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
MANTLE_ENDPOINT_TEMPLATE = "https://bedrock-mantle.{region}.api.aws/v1/messages"
|
||||
|
||||
|
||||
class AmazonMantleConfig(AmazonAnthropicClaudeConfig):
|
||||
"""
|
||||
Config for the bedrock-mantle endpoint (Claude Mythos Preview).
|
||||
|
||||
Uses the Anthropic Messages API format with AWS SigV4 auth, but at a
|
||||
different endpoint from bedrock-runtime. Model ID goes in the request body.
|
||||
|
||||
Usage: model="bedrock/mantle/anthropic.claude-mythos-preview"
|
||||
"""
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
region = self._get_aws_region_name(optional_params=optional_params, model=model)
|
||||
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
# Strip the "mantle/" routing prefix to get the real model ID
|
||||
model_id = model.replace("mantle/", "", 1)
|
||||
|
||||
request = self._build_bedrock_anthropic_request_base(
|
||||
model=model_id,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
# The parent strips "model" from the body (Invoke API puts it in URL).
|
||||
# The mantle endpoint (Messages API) requires "model" in the body.
|
||||
request["model"] = model_id
|
||||
return request
|
||||
|
||||
async def async_transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
model_id = model.replace("mantle/", "", 1)
|
||||
|
||||
request = self._build_bedrock_anthropic_request_base(
|
||||
model=model_id,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
await self._async_convert_document_url_sources_to_base64(request)
|
||||
request["model"] = model_id
|
||||
return request
|
||||
|
|
@ -696,6 +696,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
"agentcore",
|
||||
"async_invoke",
|
||||
"openai",
|
||||
"mantle",
|
||||
]:
|
||||
"""
|
||||
Get the bedrock route for the given model.
|
||||
|
|
@ -710,6 +711,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
"agentcore",
|
||||
"async_invoke",
|
||||
"openai",
|
||||
"mantle",
|
||||
],
|
||||
] = {
|
||||
"invoke/": "invoke",
|
||||
|
|
@ -719,6 +721,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
"agentcore/": "agentcore",
|
||||
"async_invoke/": "async_invoke",
|
||||
"openai/": "openai",
|
||||
"mantle/": "mantle",
|
||||
}
|
||||
|
||||
# Check explicit routes first
|
||||
|
|
@ -770,6 +773,13 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
"""
|
||||
return "agentcore/" in model
|
||||
|
||||
@staticmethod
|
||||
def _explicit_mantle_route(model: str) -> bool:
|
||||
"""
|
||||
Check if the model is an explicit mantle route (bedrock-mantle endpoint).
|
||||
"""
|
||||
return "mantle/" in model
|
||||
|
||||
@staticmethod
|
||||
def _explicit_converse_like_route(model: str) -> bool:
|
||||
"""
|
||||
|
|
@ -809,6 +819,16 @@ class BedrockModelInfo(BaseLLMModelInfo):
|
|||
if BedrockModelInfo._explicit_converse_route(model):
|
||||
return None
|
||||
|
||||
#########################################################
|
||||
# Mantle route uses the bedrock-mantle endpoint (not bedrock-runtime)
|
||||
#########################################################
|
||||
if BedrockModelInfo._explicit_mantle_route(model):
|
||||
from litellm.llms.bedrock.messages.mantle_transformation import (
|
||||
AmazonMantleMessagesConfig,
|
||||
)
|
||||
|
||||
return AmazonMantleMessagesConfig()
|
||||
|
||||
#########################################################
|
||||
# This goes through litellm.AmazonAnthropicClaude3MessagesConfig()
|
||||
# Since bedrock Invoke supports Native Anthropic Messages API
|
||||
|
|
@ -855,6 +875,12 @@ def get_bedrock_chat_config(model: str):
|
|||
)
|
||||
|
||||
return AmazonAgentCoreConfig()
|
||||
elif bedrock_route == "mantle":
|
||||
from litellm.llms.bedrock.chat.mantle.transformation import (
|
||||
AmazonMantleConfig,
|
||||
)
|
||||
|
||||
return AmazonMantleConfig()
|
||||
|
||||
# Handle provider-specific configs
|
||||
if bedrock_invoke_provider == "amazon":
|
||||
|
|
|
|||
|
|
@ -483,6 +483,8 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
if headers is None:
|
||||
headers = {}
|
||||
|
|
|
|||
|
|
@ -372,6 +372,8 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Validate environment for Bedrock Stability image edit.
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.llms.bedrock.common_utils import (
|
|||
remove_custom_field_from_tools,
|
||||
)
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
|
||||
from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import GenericStreamingChunk
|
||||
|
|
@ -59,6 +60,10 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
|
||||
DEFAULT_BEDROCK_ANTHROPIC_API_VERSION = "bedrock-2023-05-31"
|
||||
|
||||
BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(
|
||||
BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()
|
||||
)
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
BaseAnthropicMessagesConfig.__init__(self, **kwargs)
|
||||
AmazonInvokeConfig.__init__(self, **kwargs)
|
||||
|
|
@ -500,10 +505,6 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
anthropic_messages_request=anthropic_messages_request,
|
||||
)
|
||||
|
||||
# 5b. Strip `output_config` — Bedrock Invoke doesn't support it
|
||||
# Fixes: https://github.com/BerriAI/litellm/issues/22797
|
||||
anthropic_messages_request.pop("output_config", None)
|
||||
|
||||
# 5a. Remove `custom` field from tools (Bedrock doesn't support it)
|
||||
# Claude Code sends `custom: {defer_loading: true}` on tool definitions,
|
||||
# which causes Bedrock to reject the request with "Extra inputs are not permitted"
|
||||
|
|
@ -550,14 +551,43 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
if "tool-search-tool-2025-10-19" in beta_set:
|
||||
beta_set.add("tool-examples-2025-10-29")
|
||||
|
||||
filtered_auto_betas = filter_and_transform_beta_headers(
|
||||
beta_headers=list(beta_set - user_beta_set),
|
||||
provider="bedrock",
|
||||
filtered_betas = sorted(
|
||||
filter_and_transform_beta_headers(
|
||||
beta_headers=list(beta_set),
|
||||
provider="bedrock",
|
||||
)
|
||||
)
|
||||
filtered_betas = sorted(user_beta_set.union(set(filtered_auto_betas)))
|
||||
|
||||
dropped_user_betas = sorted(
|
||||
b
|
||||
for b in user_beta_set
|
||||
if not filter_and_transform_beta_headers([b], provider="bedrock")
|
||||
)
|
||||
if dropped_user_betas:
|
||||
verbose_logger.warning(
|
||||
"Bedrock Invoke: dropping unsupported anthropic-beta values "
|
||||
"from client headers: %s. Bedrock has no mapping entry for "
|
||||
"these; forwarding them would cause a 400.",
|
||||
dropped_user_betas,
|
||||
)
|
||||
|
||||
if filtered_betas:
|
||||
anthropic_messages_request["anthropic_beta"] = filtered_betas
|
||||
|
||||
# 7. Final safety net: filter top-level fields to the Bedrock Invoke allowlist.
|
||||
# Catches Anthropic-only extensions (context_management, output_config, speed,
|
||||
# mcp_servers, ...) and any future additions Claude Code may start sending.
|
||||
allowed = self.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS
|
||||
stripped = sorted(k for k in anthropic_messages_request if k not in allowed)
|
||||
if stripped:
|
||||
verbose_logger.debug(
|
||||
"Bedrock Invoke: stripping unsupported top-level request fields: %s",
|
||||
stripped,
|
||||
)
|
||||
anthropic_messages_request = {
|
||||
k: v for k, v in anthropic_messages_request.items() if k in allowed
|
||||
}
|
||||
|
||||
return anthropic_messages_request
|
||||
|
||||
def get_async_streaming_response_iterator(
|
||||
|
|
|
|||
69
litellm/llms/bedrock/messages/mantle_transformation.py
Normal file
69
litellm/llms/bedrock/messages/mantle_transformation.py
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
"""
|
||||
Transformation for Bedrock Mantle (Claude Mythos Preview) - /messages endpoint
|
||||
|
||||
Inherits all Messages API request/response transformations from
|
||||
AmazonAnthropicClaudeMessagesConfig. Overrides only the URL and model-prefix
|
||||
stripping that are specific to the bedrock-mantle endpoint.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeMessagesConfig,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
MANTLE_ENDPOINT_TEMPLATE = "https://bedrock-mantle.{region}.api.aws/v1/messages"
|
||||
|
||||
|
||||
class AmazonMantleMessagesConfig(AmazonAnthropicClaudeMessagesConfig):
|
||||
"""
|
||||
Config for the bedrock-mantle /messages endpoint (Claude Mythos Preview).
|
||||
|
||||
The mantle endpoint uses the Anthropic Messages API format and requires the
|
||||
model ID in the request body (unlike Bedrock Invoke which puts it in the URL).
|
||||
"""
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
region = self._get_aws_region_name(optional_params=optional_params, model=model)
|
||||
return MANTLE_ENDPOINT_TEMPLATE.format(region=region)
|
||||
|
||||
def transform_anthropic_messages_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
anthropic_messages_optional_request_params: Dict,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> Dict:
|
||||
# Strip "mantle/" routing prefix to get the real model ID
|
||||
model_id = model.replace("mantle/", "", 1)
|
||||
|
||||
request = super().transform_anthropic_messages_request(
|
||||
model=model_id,
|
||||
messages=messages,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
# Parent (AmazonAnthropicClaudeMessagesConfig) removes "model" from the
|
||||
# body (Bedrock Invoke puts model in the URL). The mantle endpoint
|
||||
# (Messages API) requires "model" in the request body.
|
||||
request["model"] = model_id
|
||||
return request
|
||||
|
|
@ -14,7 +14,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
|||
import httpx
|
||||
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.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.images.main import ImageEditOptionalRequestParams
|
||||
|
|
@ -123,6 +125,8 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Validate environment and set up headers for Black Forest Labs.
|
||||
|
|
@ -206,14 +210,14 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
|
|||
)
|
||||
elif isinstance(image, str):
|
||||
if image.startswith(("http://", "https://")):
|
||||
# Download image from URL
|
||||
response = httpx.get(image, timeout=60.0)
|
||||
response = safe_get(litellm.module_level_client, image, timeout=60.0)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
else:
|
||||
# Assume it's a file path
|
||||
with open(image, "rb") as f:
|
||||
return f.read()
|
||||
raise ValueError(
|
||||
"Unsupported image input: plain string values that are not URLs are not accepted. "
|
||||
"Provide image bytes or a file-like object."
|
||||
)
|
||||
elif hasattr(image, "read"):
|
||||
# File-like object
|
||||
pos = getattr(image, "tell", lambda: 0)()
|
||||
|
|
|
|||
|
|
@ -5515,6 +5515,8 @@ class BaseLLMHTTPHandler:
|
|||
api_key=litellm_params.api_key,
|
||||
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
|
||||
model=model,
|
||||
litellm_params=dict(litellm_params),
|
||||
api_base=litellm_params.api_base,
|
||||
)
|
||||
|
||||
if extra_headers:
|
||||
|
|
@ -5611,6 +5613,8 @@ class BaseLLMHTTPHandler:
|
|||
api_key=litellm_params.api_key,
|
||||
headers=image_edit_optional_request_params.get("extra_headers", {}) or {},
|
||||
model=model,
|
||||
litellm_params=dict(litellm_params),
|
||||
api_base=litellm_params.api_base,
|
||||
)
|
||||
|
||||
if extra_headers:
|
||||
|
|
|
|||
11
litellm/llms/dashscope/image_generation/__init__.py
Normal file
11
litellm/llms/dashscope/image_generation/__init__.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
|
||||
from .transformation import DashScopeImageGenerationConfig
|
||||
|
||||
__all__ = ["DashScopeImageGenerationConfig"]
|
||||
|
||||
|
||||
def get_dashscope_image_generation_config(model: str) -> BaseImageGenerationConfig:
|
||||
return DashScopeImageGenerationConfig()
|
||||
204
litellm/llms/dashscope/image_generation/transformation.py
Normal file
204
litellm/llms/dashscope/image_generation/transformation.py
Normal file
|
|
@ -0,0 +1,204 @@
|
|||
"""
|
||||
DashScope Image Generation Configuration
|
||||
|
||||
Handles transformation between OpenAI-compatible format and DashScope multimodal-generation API.
|
||||
|
||||
API endpoint: POST https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation
|
||||
|
||||
Request format:
|
||||
{
|
||||
"model": "qwen-image-2.0-pro",
|
||||
"input": {
|
||||
"messages": [{"role": "user", "content": [{"text": "<prompt>"}]}]
|
||||
},
|
||||
"parameters": {"size": "1024*1024", ...}
|
||||
}
|
||||
|
||||
Response format:
|
||||
{
|
||||
"output": {
|
||||
"choices": [{"message": {"content": [{"image": "<url>"}]}}]
|
||||
},
|
||||
"usage": {"input_tokens": 0, "output_tokens": 0, "width": 1024, "height": 1024, "image_count": 1}
|
||||
}
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
OpenAIImageGenerationOptionalParams,
|
||||
)
|
||||
from litellm.types.utils import ImageObject, ImageResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
DEFAULT_API_BASE = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
|
||||
|
||||
# Maps OpenAI size strings (WxH) to DashScope size strings (W*H)
|
||||
OPENAI_TO_DASHSCOPE_SIZE: dict = {
|
||||
"256x256": "256*256",
|
||||
"512x512": "512*512",
|
||||
"1024x1024": "1024*1024",
|
||||
"1792x1024": "1792*1024",
|
||||
"1024x1792": "1024*1792",
|
||||
"2048x2048": "2048*2048",
|
||||
}
|
||||
|
||||
|
||||
class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
|
||||
"""
|
||||
Configuration for DashScope image generation (qwen-image-2.0, qwen-image-2.0-pro).
|
||||
"""
|
||||
|
||||
def get_supported_openai_params(
|
||||
self, model: str
|
||||
) -> List[OpenAIImageGenerationOptionalParams]:
|
||||
return ["n", "size"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
supported_params = self.get_supported_openai_params(model)
|
||||
mapped: dict = {}
|
||||
for k, v in non_default_params.items():
|
||||
if k in optional_params:
|
||||
continue
|
||||
if k not in supported_params:
|
||||
continue
|
||||
if k == "size":
|
||||
# Convert "WxH" → "W*H"
|
||||
mapped["size"] = OPENAI_TO_DASHSCOPE_SIZE.get(v, v.replace("x", "*"))
|
||||
elif k == "n":
|
||||
mapped["image_count"] = v
|
||||
return mapped
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
return (
|
||||
api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE
|
||||
)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
final_api_key = api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
if not final_api_key:
|
||||
raise ValueError("DASHSCOPE_API_KEY is not set")
|
||||
headers["Authorization"] = f"Bearer {final_api_key}"
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
|
||||
def transform_image_generation_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
"""
|
||||
Transform OpenAI-style image generation request to DashScope multimodal-generation format.
|
||||
"""
|
||||
parameters: dict = {}
|
||||
for k, v in optional_params.items():
|
||||
parameters[k] = v
|
||||
|
||||
return {
|
||||
"model": model,
|
||||
"input": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"text": prompt}],
|
||||
}
|
||||
]
|
||||
},
|
||||
"parameters": parameters,
|
||||
}
|
||||
|
||||
def transform_image_generation_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: ImageResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
request_data: dict,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: Any,
|
||||
api_key: Optional[str] = None,
|
||||
json_mode: Optional[bool] = None,
|
||||
) -> ImageResponse:
|
||||
"""
|
||||
Transform DashScope response to litellm ImageResponse.
|
||||
|
||||
DashScope response: output.choices[0].message.content[0].image
|
||||
OpenAI response: data[0].url
|
||||
"""
|
||||
if raw_response.status_code != 200:
|
||||
raise self.get_error_class(
|
||||
error_message=raw_response.text,
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
try:
|
||||
response_data = raw_response.json()
|
||||
except Exception as e:
|
||||
raise self.get_error_class(
|
||||
error_message=f"Failed to parse DashScope image generation response: {e}",
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
# DashScope can return API-level errors in a 200 response body.
|
||||
# Example: {"code": "InvalidParameter", "message": "Size not supported"}
|
||||
if "code" in response_data and "output" not in response_data:
|
||||
raise self.get_error_class(
|
||||
error_message=str(response_data.get("message", response_data)),
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
if not model_response.data:
|
||||
model_response.data = []
|
||||
|
||||
choices = response_data.get("output", {}).get("choices", [])
|
||||
for choice in choices:
|
||||
content_list = choice.get("message", {}).get("content", [])
|
||||
for content_item in content_list:
|
||||
image_url = content_item.get("image")
|
||||
if image_url:
|
||||
model_response.data.append(ImageObject(url=image_url))
|
||||
|
||||
return model_response
|
||||
|
|
@ -54,6 +54,8 @@ class GeminiImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
final_api_key: Optional[str] = api_key or get_secret_str("GEMINI_API_KEY")
|
||||
if not final_api_key:
|
||||
|
|
|
|||
|
|
@ -8,7 +8,12 @@ class LiteLLMProxyImageEditConfig(OpenAIImageEditConfig):
|
|||
"""Configuration for image edit requests routed through LiteLLM Proxy."""
|
||||
|
||||
def validate_environment(
|
||||
self, headers: dict, model: str, api_key: Optional[str] = None
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
api_key = api_key or get_secret_str("LITELLM_PROXY_API_KEY")
|
||||
headers.update({"Authorization": f"Bearer {api_key}"})
|
||||
|
|
|
|||
|
|
@ -53,9 +53,21 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
|
||||
@classmethod
|
||||
def is_model_gpt_5_model(cls, model: str) -> bool:
|
||||
# gpt-5-chat* behaves like a regular chat model (supports temperature, etc.)
|
||||
# Don't route it through GPT-5 reasoning-specific parameter restrictions.
|
||||
return "gpt-5" in model and "gpt-5-chat" not in model
|
||||
# The gpt-5-chat* family (gpt-5-chat, gpt-5-chat-latest, gpt-5-chat-2025-08-07,
|
||||
# …) are regular chat models: they support temperature and tool_choice but NOT
|
||||
# reasoning_effort. They must NOT be routed through the GPT-5 reasoning path.
|
||||
#
|
||||
# Versioned chat models such as gpt-5.3-chat and gpt-5.1-chat ARE reasoning
|
||||
# models and must stay on the GPT-5 path. The distinguishing feature is that
|
||||
# the gpt-5-chat family has a literal "-chat" immediately after "gpt-5"
|
||||
# (i.e. "gpt-5-chat…"), while versioned chat models interpose a minor version
|
||||
# number (i.e. "gpt-5.<digit>-chat").
|
||||
#
|
||||
# Using a startswith("gpt-5-chat") prefix check on the normalized name (rather
|
||||
# than a substring check) makes this boundary explicit and avoids any ambiguity
|
||||
# if future model names coincidentally contain "gpt-5-chat" as an interior run.
|
||||
_normalized = model.split("/")[-1] # strip provider prefix, e.g. "openai/"
|
||||
return "gpt-5" in model and not _normalized.startswith("gpt-5-chat")
|
||||
|
||||
@classmethod
|
||||
def is_model_gpt_5_search_model(cls, model: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -165,6 +165,8 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
api_key = (
|
||||
api_key
|
||||
|
|
|
|||
|
|
@ -116,6 +116,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
api_key = api_key or litellm.api_key or get_secret_str("OPENROUTER_API_KEY")
|
||||
if not api_key:
|
||||
|
|
|
|||
|
|
@ -8,9 +8,8 @@ More information on our website: https://endpoints.ai.cloud.ovh.net
|
|||
from typing import Optional, Union, List
|
||||
|
||||
import httpx
|
||||
from litellm.utils import ModelResponseStream, _get_model_info_helper
|
||||
from litellm.utils import ModelResponseStream
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.ovhcloud.utils import OVHCloudException
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
|
@ -22,34 +21,6 @@ class OVHCloudChatConfig(OpenAIGPTConfig):
|
|||
def custom_llm_provider(self) -> Optional[str]:
|
||||
return "ovhcloud"
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
"""
|
||||
Details about function calling support can be found here:
|
||||
https://help.ovhcloud.com/csm/en-gb-public-cloud-ai-endpoints-function-calling?id=kb_article_view&sysparm_article=KB0071907
|
||||
"""
|
||||
supports_function_calling: Optional[bool] = None
|
||||
try:
|
||||
model_info = _get_model_info_helper(model, custom_llm_provider="ovhcloud")
|
||||
supports_function_calling = model_info.get(
|
||||
"supports_function_calling", None
|
||||
)
|
||||
if supports_function_calling is None:
|
||||
supports_function_calling = False
|
||||
except Exception as e:
|
||||
verbose_logger.debug(f"Error getting supported OpenAI params: {e}")
|
||||
supports_function_calling = False
|
||||
|
||||
optional_params = super().get_supported_openai_params(model)
|
||||
if supports_function_calling is not True:
|
||||
verbose_logger.debug(
|
||||
"You can see our models supporting function_calling in our catalog: https://endpoints.ai.cloud.ovh.net/catalog "
|
||||
)
|
||||
optional_params.remove("tools")
|
||||
optional_params.remove("tool_choice")
|
||||
optional_params.remove("function_call")
|
||||
optional_params.remove("response_format")
|
||||
return optional_params
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
|
|
|
|||
|
|
@ -81,6 +81,8 @@ class RecraftImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
final_api_key: Optional[str] = api_key or get_secret_str("RECRAFT_API_KEY")
|
||||
if not final_api_key:
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import re
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -66,6 +67,8 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
|||
aws_region_name = litellm_params.get("aws_region_name")
|
||||
if not aws_region_name:
|
||||
raise ValueError("aws_region_name is required for S3 Vectors")
|
||||
if not re.match(r"^[a-z][a-z0-9-]*$", aws_region_name):
|
||||
raise ValueError("Invalid aws_region_name format")
|
||||
return f"https://s3vectors.{aws_region_name}.api.aws"
|
||||
|
||||
def transform_search_vector_store_request(
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import re
|
||||
from typing import TYPE_CHECKING, Any, List, Optional, Tuple
|
||||
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
|
@ -61,6 +62,8 @@ class SnowflakeBaseConfig:
|
|||
account_id = get_secret_str("SNOWFLAKE_ACCOUNT_ID")
|
||||
if account_id is None:
|
||||
raise ValueError("Missing snowflake account_id")
|
||||
if not re.match(r"^[a-zA-Z0-9_-]+$", account_id):
|
||||
raise ValueError("Invalid account_id format")
|
||||
api_base = f"https://{account_id}.snowflakecomputing.com/api/v2"
|
||||
|
||||
api_base = api_base.rstrip("/")
|
||||
|
|
|
|||
|
|
@ -149,6 +149,8 @@ class StabilityImageEditConfig(BaseImageEditConfig):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Validate environment and set up headers for Stability AI.
|
||||
|
|
|
|||
|
|
@ -229,11 +229,20 @@ def get_vertex_base_url(
|
|||
) -> str:
|
||||
"""
|
||||
Get the base URL for Vertex AI API calls.
|
||||
|
||||
- ``global`` uses the global control plane host.
|
||||
- Multi-region geographies (e.g. ``us``, ``eu``) use ``aiplatform.{geo}.rep.googleapis.com``.
|
||||
- Regional locations (e.g. ``us-central1``) use ``{region}-aiplatform.googleapis.com``.
|
||||
"""
|
||||
if vertex_location == "global":
|
||||
return "https://aiplatform.googleapis.com"
|
||||
else:
|
||||
return f"https://{vertex_location}-aiplatform.googleapis.com"
|
||||
if vertex_location is None:
|
||||
raise ValueError("vertex_location is required")
|
||||
if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
|
||||
raise ValueError("Invalid vertex_location format")
|
||||
if "-" not in vertex_location:
|
||||
return f"https://aiplatform.{vertex_location}.rep.googleapis.com"
|
||||
return f"https://{vertex_location}-aiplatform.googleapis.com"
|
||||
|
||||
|
||||
def _get_embedding_url(
|
||||
|
|
|
|||
|
|
@ -132,26 +132,28 @@ def _extract_max_media_resolution_from_messages(
|
|||
return max_resolution
|
||||
|
||||
|
||||
def _apply_gemini_3_metadata(
|
||||
def _apply_gemini_metadata(
|
||||
part: PartType,
|
||||
model: Optional[str],
|
||||
media_resolution_enum: Optional[Dict[str, str]],
|
||||
video_metadata: Optional[Dict[str, Any]],
|
||||
) -> PartType:
|
||||
"""
|
||||
Apply the unique media_resolution and video_metadata parameters of Gemini 3+
|
||||
Apply media_resolution and video_metadata parameters to a Gemini part.
|
||||
|
||||
- Per-part media_resolution: Gemini 3+ only (2.x uses generation_config global).
|
||||
- video_metadata (fps, startOffset, endOffset): all Gemini models (1.x, 2.x, 3+).
|
||||
"""
|
||||
if model is None:
|
||||
return part
|
||||
|
||||
from .vertex_and_google_ai_studio_gemini import VertexGeminiConfig
|
||||
|
||||
if not VertexGeminiConfig._is_gemini_3_or_newer(model):
|
||||
return part
|
||||
|
||||
part_dict = dict(part)
|
||||
|
||||
if media_resolution_enum is not None:
|
||||
if media_resolution_enum is not None and VertexGeminiConfig._is_gemini_3_or_newer(
|
||||
model
|
||||
):
|
||||
part_dict["media_resolution"] = media_resolution_enum
|
||||
|
||||
if video_metadata is not None:
|
||||
|
|
@ -206,7 +208,7 @@ def _process_gemini_media(
|
|||
mime_type = format
|
||||
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
|
||||
part: PartType = {"file_data": file_data}
|
||||
return _apply_gemini_3_metadata(
|
||||
return _apply_gemini_metadata(
|
||||
part, model, media_resolution_enum, video_metadata
|
||||
)
|
||||
elif (
|
||||
|
|
@ -216,14 +218,14 @@ def _process_gemini_media(
|
|||
):
|
||||
file_data = FileDataType(mime_type=image_type, file_uri=image_url)
|
||||
part = {"file_data": file_data}
|
||||
return _apply_gemini_3_metadata(
|
||||
return _apply_gemini_metadata(
|
||||
part, model, media_resolution_enum, video_metadata
|
||||
)
|
||||
elif "http://" in image_url or "https://" in image_url or "base64" in image_url:
|
||||
image = convert_to_anthropic_image_obj(image_url, format=format)
|
||||
_blob: BlobType = {"data": image["data"], "mime_type": image["media_type"]}
|
||||
part = {"inline_data": cast(BlobType, _blob)}
|
||||
return _apply_gemini_3_metadata(
|
||||
return _apply_gemini_metadata(
|
||||
part, model, media_resolution_enum, video_metadata
|
||||
)
|
||||
raise Exception("Invalid image received - {}".format(image_url))
|
||||
|
|
@ -733,9 +735,9 @@ def _transform_request_body( # noqa: PLR0915
|
|||
**filtered_params
|
||||
)
|
||||
|
||||
# For Gemini 2.x models, add media_resolution to generation_config (global)
|
||||
# Gemini 3+ supports per-part media_resolution, but 2.x only supports global
|
||||
# Gemini 1.x does not support mediaResolution at all
|
||||
# For Gemini 2.x models, also add media_resolution to generation_config (global)
|
||||
# as a fallback, since some 2.x versions may not support per-part media_resolution.
|
||||
# Gemini 1.x does not support mediaResolution at all.
|
||||
if "gemini-2" in model:
|
||||
max_media_resolution = _extract_max_media_resolution_from_messages(messages)
|
||||
if max_media_resolution:
|
||||
|
|
|
|||
|
|
@ -103,10 +103,24 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
headers: dict,
|
||||
model: str,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
headers = headers or {}
|
||||
vertex_project = self._resolve_vertex_project()
|
||||
vertex_credentials = self._resolve_vertex_credentials()
|
||||
litellm_params = litellm_params or {}
|
||||
|
||||
_api_base = litellm_params.get("api_base") or api_base
|
||||
if _api_base is not None:
|
||||
return headers
|
||||
|
||||
vertex_project = (
|
||||
self.safe_get_vertex_ai_project(litellm_params)
|
||||
or self._resolve_vertex_project()
|
||||
)
|
||||
vertex_credentials = (
|
||||
self.safe_get_vertex_ai_credentials(litellm_params)
|
||||
or self._resolve_vertex_credentials()
|
||||
)
|
||||
access_token, _ = self._ensure_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
|
|
@ -123,8 +137,14 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
"""
|
||||
Get the complete URL for Vertex AI Imagen predict API
|
||||
"""
|
||||
vertex_project = self._resolve_vertex_project()
|
||||
vertex_location = self._resolve_vertex_location()
|
||||
vertex_project = (
|
||||
self.safe_get_vertex_ai_project(litellm_params)
|
||||
or self._resolve_vertex_project()
|
||||
)
|
||||
vertex_location = (
|
||||
self.safe_get_vertex_ai_location(litellm_params)
|
||||
or self._resolve_vertex_location()
|
||||
)
|
||||
|
||||
if not vertex_project or not vertex_location:
|
||||
raise ValueError(
|
||||
|
|
@ -348,13 +368,16 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
|
|||
if stream_pos is not None:
|
||||
image.seek(stream_pos)
|
||||
return data
|
||||
if isinstance(image, (str, Path)):
|
||||
path_obj = Path(image)
|
||||
if not path_obj.exists():
|
||||
raise ValueError(
|
||||
f"Mask/image path does not exist for Vertex AI Imagen image edit: {path_obj}"
|
||||
)
|
||||
return path_obj.read_bytes()
|
||||
if isinstance(image, str):
|
||||
raise ValueError(
|
||||
"Unsupported image input: plain string values are not accepted for "
|
||||
"Vertex AI Imagen image edit. Provide image bytes or a file-like object."
|
||||
)
|
||||
if isinstance(image, Path):
|
||||
raise ValueError(
|
||||
"Unsupported image input: filesystem paths are not accepted for "
|
||||
"Vertex AI Imagen image edit. Provide image bytes or a file-like object."
|
||||
)
|
||||
if hasattr(image, "read"):
|
||||
data = image.read()
|
||||
if isinstance(data, str):
|
||||
|
|
|
|||
|
|
@ -1006,7 +1006,8 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"global.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -1034,7 +1035,8 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"us.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
|
|
@ -1062,7 +1064,8 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"eu.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
|
|
@ -1090,7 +1093,8 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"au.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
|
|
@ -1118,7 +1122,8 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"anthropic.claude-opus-4-7": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -1146,7 +1151,9 @@
|
|||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"global.anthropic.claude-opus-4-7": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -1174,7 +1181,9 @@
|
|||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"us.anthropic.claude-opus-4-7": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
|
|
@ -1202,7 +1211,9 @@
|
|||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"eu.anthropic.claude-opus-4-7": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
|
|
@ -1230,7 +1241,9 @@
|
|||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"au.anthropic.claude-opus-4-7": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
|
|
@ -1258,7 +1271,9 @@
|
|||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"anthropic.claude-sonnet-4-6": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
|
|
@ -1285,7 +1300,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"global.anthropic.claude-sonnet-4-6": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
|
|
@ -1312,7 +1328,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"us.anthropic.claude-sonnet-4-6": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
|
|
@ -1339,7 +1356,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"eu.anthropic.claude-sonnet-4-6": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
|
|
@ -1366,7 +1384,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"au.anthropic.claude-sonnet-4-6": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
|
|
@ -1393,7 +1412,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_native_structured_output": true
|
||||
"supports_native_structured_output": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"anthropic.claude-sonnet-4-20250514-v1:0": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
|
|
@ -1911,7 +1931,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 159,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"azure_ai/claude-opus-4-7": {
|
||||
"input_cost_per_token": 5e-06,
|
||||
|
|
@ -1939,7 +1960,9 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 159
|
||||
"tool_use_system_prompt_tokens": 159,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"azure_ai/claude-opus-4-1": {
|
||||
"cache_creation_input_token_cost": 1.875e-05,
|
||||
|
|
@ -2003,7 +2026,8 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"azure/computer-use-preview": {
|
||||
"input_cost_per_token": 3e-06,
|
||||
|
|
@ -8909,7 +8933,8 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
|
|
@ -9103,7 +9128,8 @@
|
|||
"us": 1.1,
|
||||
"fast": 6.0
|
||||
},
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"claude-opus-4-6-20260205": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -9135,7 +9161,8 @@
|
|||
"us": 1.1,
|
||||
"fast": 6.0
|
||||
},
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"claude-opus-4-7": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -9167,7 +9194,9 @@
|
|||
"provider_specific_entry": {
|
||||
"us": 1.1,
|
||||
"fast": 6.0
|
||||
}
|
||||
},
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"claude-opus-4-7-20260416": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -9199,7 +9228,9 @@
|
|||
"provider_specific_entry": {
|
||||
"us": 1.1,
|
||||
"fast": 6.0
|
||||
}
|
||||
},
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"claude-sonnet-4-20250514": {
|
||||
"deprecation_date": "2026-05-14",
|
||||
|
|
@ -10352,6 +10383,22 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"dashscope/qwen-image-2.0": {
|
||||
"litellm_provider": "dashscope",
|
||||
"mode": "image_generation",
|
||||
"source": "https://www.alibabacloud.com/help/en/model-studio/models",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations"
|
||||
]
|
||||
},
|
||||
"dashscope/qwen-image-2.0-pro": {
|
||||
"litellm_provider": "dashscope",
|
||||
"mode": "image_generation",
|
||||
"source": "https://www.alibabacloud.com/help/en/model-studio/models",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations"
|
||||
]
|
||||
},
|
||||
"databricks/databricks-bge-large-en": {
|
||||
"input_cost_per_token": 1.0003e-07,
|
||||
"input_dbu_cost_per_token": 1.429e-06,
|
||||
|
|
@ -19226,6 +19273,42 @@
|
|||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"gpt-5.5": {
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"litellm_provider": "openai",
|
||||
"max_input_tokens": 272000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-05,
|
||||
"supported_endpoints": [
|
||||
"/v1/chat/completions",
|
||||
"/v1/batch",
|
||||
"/v1/responses"
|
||||
],
|
||||
"supported_modalities": [
|
||||
"text",
|
||||
"image"
|
||||
],
|
||||
"supported_output_modalities": [
|
||||
"text"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_native_streaming": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_service_tier": true,
|
||||
"supports_vision": true,
|
||||
"supports_none_reasoning_effort": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"gpt-5.4": {
|
||||
"cache_read_input_token_cost": 2.5e-07,
|
||||
"cache_read_input_token_cost_above_272k_tokens": 5e-07,
|
||||
|
|
@ -22872,6 +22955,22 @@
|
|||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"moonshot/kimi-k2.6": {
|
||||
"cache_read_input_token_cost": 1.6e-07,
|
||||
"input_cost_per_token": 9.5e-07,
|
||||
"litellm_provider": "moonshot",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4e-06,
|
||||
"source": "https://platform.kimi.ai/docs/pricing/chat-k26",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"moonshot/kimi-latest": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
|
|
@ -25052,7 +25151,8 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 159
|
||||
"tool_use_system_prompt_tokens": 159,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"openrouter/anthropic/claude-opus-4.5": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -25090,7 +25190,8 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"openrouter/anthropic/claude-sonnet-4.5": {
|
||||
"input_cost_per_image": 0.0048,
|
||||
|
|
@ -30118,7 +30219,8 @@
|
|||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"vercel_ai_gateway/anthropic/claude-sonnet-4": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
|
|
@ -31345,7 +31447,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"vertex_ai/claude-opus-4-6@default": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -31372,7 +31475,8 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_max_reasoning_effort": true
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"vertex_ai/claude-opus-4-7": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -31399,7 +31503,9 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"vertex_ai/claude-opus-4-7@default": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
|
|
@ -31426,7 +31532,9 @@
|
|||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
"tool_use_system_prompt_tokens": 346,
|
||||
"supports_max_reasoning_effort": true,
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"vertex_ai/claude-sonnet-4-5": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
|
|
@ -31478,7 +31586,8 @@
|
|||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
}
|
||||
},
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"vertex_ai/claude-sonnet-4-5@20250929": {
|
||||
"cache_creation_input_token_cost": 3.75e-06,
|
||||
|
|
@ -38345,7 +38454,8 @@
|
|||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
}
|
||||
},
|
||||
"supports_minimal_reasoning_effort": true
|
||||
},
|
||||
"duckduckgo/search": {
|
||||
"litellm_provider": "duckduckgo",
|
||||
|
|
|
|||
|
|
@ -323,6 +323,14 @@ async def authorize_with_server(
|
|||
)
|
||||
|
||||
parsed = urlparse(redirect_uri)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "invalid_redirect_uri",
|
||||
"message": "redirect_uri must use http or https scheme",
|
||||
},
|
||||
)
|
||||
base_url = urlunparse(parsed._replace(query=""))
|
||||
request_base_url = get_request_base_url(request)
|
||||
encoded_state = encode_state_with_base_url(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ Filters MCP tools semantically for /chat/completions and /responses endpoints.
|
|||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.utils import MCP_TOOL_PREFIX_SEPARATOR
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from semantic_router.routers import SemanticRouter
|
||||
|
|
@ -214,20 +215,89 @@ class SemanticMCPToolFilter:
|
|||
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _name_matches_canonical(client_name: str, canonical: str) -> bool:
|
||||
"""
|
||||
Return True if a client-side tool name refers to the given canonical
|
||||
MCP tool name.
|
||||
|
||||
MCP clients (e.g. opencode) commonly wrap the proxy's canonical tool
|
||||
name with an additive namespace prefix of their own
|
||||
(``<client_alias><sep><canonical>``). The prefix can use either a
|
||||
dash or an underscore as separator regardless of what
|
||||
``MCP_TOOL_PREFIX_SEPARATOR`` is set to on the proxy, because the
|
||||
client doesn't know the proxy's separator.
|
||||
|
||||
The match is anchored: ``canonical`` must form the complete suffix
|
||||
of ``client_name`` and be preceded by a separator character, so
|
||||
``rain_gear`` does not match canonical ``ear``.
|
||||
|
||||
Suffix matching is additionally gated on ``canonical`` itself
|
||||
containing ``MCP_TOOL_PREFIX_SEPARATOR``. Server-registered MCP
|
||||
tools are always emitted as
|
||||
``<server_name><MCP_TOOL_PREFIX_SEPARATOR><tool_name>`` (see
|
||||
``add_server_prefix_to_name``), so a canonical without the
|
||||
separator is not a namespaced MCP tool and falling back to
|
||||
suffix matching would spuriously collide with unrelated local
|
||||
user functions whose names end in the same characters.
|
||||
"""
|
||||
if client_name == canonical:
|
||||
return True
|
||||
if MCP_TOOL_PREFIX_SEPARATOR not in canonical:
|
||||
return False
|
||||
if len(client_name) <= len(canonical):
|
||||
return False
|
||||
if not client_name.endswith(canonical):
|
||||
return False
|
||||
separator = client_name[-len(canonical) - 1]
|
||||
return separator in ("_", "-")
|
||||
|
||||
def _get_tools_by_names(
|
||||
self, tool_names: List[str], available_tools: List[Any]
|
||||
) -> List[Any]:
|
||||
"""Get tools from available_tools by their names, preserving order."""
|
||||
# Match tools from available_tools (preserves format - dict or MCPTool)
|
||||
matched_tools = []
|
||||
for tool in available_tools:
|
||||
tool_name, _ = self._extract_tool_info(tool)
|
||||
if tool_name in tool_names:
|
||||
matched_tools.append(tool)
|
||||
"""
|
||||
Get tools from available_tools by their names, preserving the
|
||||
semantic router's ordering.
|
||||
|
||||
# Reorder to match semantic router's ordering
|
||||
tool_map = {self._extract_tool_info(t)[0]: t for t in matched_tools}
|
||||
return [tool_map[name] for name in tool_names if name in tool_map]
|
||||
Matching is tolerant of client-side namespace prefixes: if an
|
||||
incoming tool arrived as ``<client_alias>_<canonical>`` while the
|
||||
router returned ``<canonical>`` (see
|
||||
``_name_matches_canonical``), that tool is still selected. The
|
||||
returned tool object is the original from ``available_tools``, so
|
||||
the client-facing name is preserved for tool-call round-trips.
|
||||
"""
|
||||
# Build an index of incoming tools by their client-facing name.
|
||||
# Exact matches win over suffix matches when both are present, and
|
||||
# each incoming tool is returned at most once even if two canonical
|
||||
# names happen to be tail-compatible with the same incoming name.
|
||||
available_by_name: Dict[str, Any] = {}
|
||||
for tool in available_tools:
|
||||
client_name, _ = self._extract_tool_info(tool)
|
||||
if client_name and client_name not in available_by_name:
|
||||
available_by_name[client_name] = tool
|
||||
|
||||
matched: List[Any] = []
|
||||
used_ids: set = set()
|
||||
for canonical in tool_names:
|
||||
tool = available_by_name.get(canonical)
|
||||
if tool is None:
|
||||
# Prefer the shortest qualifying name. When several
|
||||
# incoming tools suffix-match the same canonical (e.g.
|
||||
# "my_search" and "my_tag_search" both end in "search"),
|
||||
# the one closest in length to the canonical is the
|
||||
# least-wrapped and most likely the intended target.
|
||||
best_name: Optional[str] = None
|
||||
for client_name in available_by_name:
|
||||
if not self._name_matches_canonical(client_name, canonical):
|
||||
continue
|
||||
if best_name is None or len(client_name) < len(best_name):
|
||||
best_name = client_name
|
||||
if best_name is not None:
|
||||
tool = available_by_name[best_name]
|
||||
if tool is not None and id(tool) not in used_ids:
|
||||
matched.append(tool)
|
||||
used_ids.add(id(tool))
|
||||
return matched
|
||||
|
||||
def extract_user_query(self, messages: List[Dict[str, Any]]) -> str:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,42 +1,83 @@
|
|||
# model_list:
|
||||
# - model_name: claude-sonnet-4-6
|
||||
# litellm_params: {model: anthropic/claude-sonnet-4-6}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 1
|
||||
# keywords: [tin]
|
||||
# - model_name: gpt-4o-mini
|
||||
# litellm_params: {model: openai/gpt-4o-mini}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 1
|
||||
# keywords: []
|
||||
# - model_name: gpt-4o
|
||||
# litellm_params: {model: openai/gpt-4o}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 2
|
||||
# keywords: [vision, function_calling]
|
||||
# - model_name: opus
|
||||
# litellm_params: {model: anthropic/claude-opus-4-7}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 3
|
||||
# keywords: ["architecture", "design"]
|
||||
# - model_name: my-quality-router
|
||||
# litellm_params:
|
||||
# model: auto_router/adaptive_router
|
||||
# adaptive_router_default_model: gpt-4o-mini
|
||||
# adaptive_router_config:
|
||||
# available_models: [gpt-4o-mini, gpt-4o, opus, claude-sonnet-4-6]
|
||||
# Example proxy config for the adaptive router (v0).
|
||||
#
|
||||
# Wires one logical router ("smart-cheap-router") that adaptively picks between
|
||||
# two real deployments ("fast" and "smart") based on per-session feedback signals.
|
||||
#
|
||||
# How to use from a client:
|
||||
# POST /v1/chat/completions { "model": "smart-cheap-router", ... }
|
||||
# Add { "metadata": { "litellm_session_id": "<your-session-id>" } } to enable
|
||||
# sticky-session routing within a conversation.
|
||||
#
|
||||
# Required env vars: OPENAI_API_KEY, DATABASE_URL.
|
||||
|
||||
model_list:
|
||||
|
||||
# OpenAI model for /v1/chat/completions test — 200x custom pricing
|
||||
- model_name: "gpt-4.1-mini"
|
||||
# ---- The adaptive router "control" deployment -------------------------
|
||||
# `model_name` is what clients call. `available_models` lists the underlying
|
||||
# deployments the router is allowed to pick from (must match other model_name
|
||||
# entries in this list).
|
||||
- model_name: smart-cheap-router
|
||||
litellm_params:
|
||||
model: openai/gpt-4.1-mini
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
model_info:
|
||||
id: gpt-4.1-mini-custom-pricing
|
||||
input_cost_per_token: 0.00004 # 100x standard ($0.40/1M = $0.0000004)
|
||||
output_cost_per_token: 0.00016 # 100x standard ($1.60/1M = $0.0000016)
|
||||
model: auto_router/adaptive_router
|
||||
adaptive_router_config:
|
||||
available_models: ["fast", "smart"]
|
||||
weights:
|
||||
quality: 0.7
|
||||
cost: 0.3
|
||||
|
||||
# OpenAI model for /v1/responses test — 100x custom pricing
|
||||
- model_name: "gpt-5"
|
||||
litellm_params:
|
||||
model: openai/gpt-5
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
model_info:
|
||||
id: gpt-5-custom-pricing
|
||||
mode: "chat"
|
||||
input_cost_per_token: 125 # 100x standard ($1.25/1M = $0.00000125)
|
||||
output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001)
|
||||
|
||||
# Anthropic model for /v1/messages test — 100x custom pricing
|
||||
- model_name: "claude-sonnet-4-6"
|
||||
# ---- Underlying deployments the router picks from ---------------------
|
||||
- model_name: fast
|
||||
litellm_params:
|
||||
model: anthropic/claude-sonnet-4-6
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
input_cost_per_token: 0.00000015
|
||||
model_info:
|
||||
id: claude-sonnet-4-custom-pricing
|
||||
input_cost_per_token: 0.0003 # 100x standard ($0.000003)
|
||||
output_cost_per_token: 0.0015 # 100x standard ($0.000015)
|
||||
- model_name: my-auto
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 2
|
||||
strengths: []
|
||||
|
||||
- model_name: smart
|
||||
litellm_params:
|
||||
model: auto_router/complexity_router
|
||||
complexity_router_config:
|
||||
tiers:
|
||||
SIMPLE: "gpt-4.1-mini"
|
||||
COMPLEX: claude-sonnet-4-6
|
||||
tier_boundaries:
|
||||
simple_medium: 0.30
|
||||
complexity_router_default_model: small-model
|
||||
model: anthropic/claude-opus-4-7
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
input_cost_per_token: 0.0000050
|
||||
model_info:
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 3
|
||||
strengths: ["code_generation", "technical_design", "analytical_reasoning"]
|
||||
|
||||
litellm_settings:
|
||||
drop_params: True
|
||||
|
||||
general_settings:
|
||||
master_key: sk-1234 # REPLACE in production
|
||||
|
|
|
|||
|
|
@ -427,7 +427,8 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/v1/skills/{skill_id}",
|
||||
]
|
||||
|
||||
mcp_routes = [
|
||||
# MCP tool-call / passthrough routes — data-plane. Gated by DISABLE_LLM_API_ENDPOINTS.
|
||||
mcp_inference_routes = [
|
||||
"/mcp",
|
||||
"/mcp/",
|
||||
"/mcp/{subpath}",
|
||||
|
|
@ -436,10 +437,18 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/mcp/tools/call",
|
||||
"/mcp-rest/tools/list",
|
||||
"/mcp-rest/tools/call",
|
||||
]
|
||||
|
||||
# MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS.
|
||||
mcp_management_routes = [
|
||||
"/v1/mcp/server",
|
||||
"/v1/mcp/server/{path:path}",
|
||||
]
|
||||
|
||||
# Backwards-compat union — virtual keys may be configured with
|
||||
# allowed_routes=["mcp_routes"], which should cover both halves.
|
||||
mcp_routes = mcp_inference_routes + mcp_management_routes
|
||||
|
||||
agent_routes = [
|
||||
"/v1/agents",
|
||||
"/v1/agents/{agent_id}",
|
||||
|
|
@ -477,7 +486,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
+ mapped_pass_through_routes
|
||||
+ passthrough_routes_wildcard
|
||||
+ apply_guardrail_routes
|
||||
+ mcp_routes
|
||||
+ mcp_inference_routes
|
||||
+ litellm_native_routes
|
||||
+ agent_routes
|
||||
)
|
||||
|
|
@ -530,40 +539,44 @@ class LiteLLMRoutes(enum.Enum):
|
|||
KeyManagementRoutes.KEY_ALIASES.value,
|
||||
]
|
||||
|
||||
management_routes = [
|
||||
# user
|
||||
"/user/new",
|
||||
"/user/update",
|
||||
"/user/bulk_update",
|
||||
"/user/delete",
|
||||
"/user/info",
|
||||
"/user/list",
|
||||
"/user/daily/activity",
|
||||
"/user/daily/activity/aggregated",
|
||||
# team
|
||||
"/team/new",
|
||||
"/team/update",
|
||||
"/team/delete",
|
||||
"/team/list",
|
||||
"/v2/team/list",
|
||||
"/team/info",
|
||||
"/team/block",
|
||||
"/team/unblock",
|
||||
"/team/available",
|
||||
"/team/permissions_list",
|
||||
"/team/permissions_update",
|
||||
"/team/daily/activity",
|
||||
# model
|
||||
"/model/new",
|
||||
"/model/update",
|
||||
"/model/delete",
|
||||
"/model/info",
|
||||
"/jwt/key/mapping/new",
|
||||
"/jwt/key/mapping/update",
|
||||
"/jwt/key/mapping/delete",
|
||||
"/jwt/key/mapping/list",
|
||||
"/jwt/key/mapping/info",
|
||||
] + key_management_routes
|
||||
management_routes = (
|
||||
[
|
||||
# user
|
||||
"/user/new",
|
||||
"/user/update",
|
||||
"/user/bulk_update",
|
||||
"/user/delete",
|
||||
"/user/info",
|
||||
"/user/list",
|
||||
"/user/daily/activity",
|
||||
"/user/daily/activity/aggregated",
|
||||
# team
|
||||
"/team/new",
|
||||
"/team/update",
|
||||
"/team/delete",
|
||||
"/team/list",
|
||||
"/v2/team/list",
|
||||
"/team/info",
|
||||
"/team/block",
|
||||
"/team/unblock",
|
||||
"/team/available",
|
||||
"/team/permissions_list",
|
||||
"/team/permissions_update",
|
||||
"/team/daily/activity",
|
||||
# model
|
||||
"/model/new",
|
||||
"/model/update",
|
||||
"/model/delete",
|
||||
"/model/info",
|
||||
"/jwt/key/mapping/new",
|
||||
"/jwt/key/mapping/update",
|
||||
"/jwt/key/mapping/delete",
|
||||
"/jwt/key/mapping/list",
|
||||
"/jwt/key/mapping/info",
|
||||
]
|
||||
+ key_management_routes
|
||||
+ mcp_management_routes
|
||||
)
|
||||
|
||||
spend_tracking_routes = [
|
||||
# spend
|
||||
|
|
@ -1997,7 +2010,12 @@ class TeamRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
||||
"""Represents user-controllable params for a LiteLLM_BudgetTable record"""
|
||||
"""Represents user-controllable params for a LiteLLM_BudgetTable record.
|
||||
|
||||
Budget-write paths use `model_fields.keys()` on this class as an allowlist
|
||||
for user input. Keep server-managed fields (e.g. `budget_reset_at`) on
|
||||
`LiteLLM_BudgetTableFull` so they aren't user-settable.
|
||||
"""
|
||||
|
||||
budget_id: Optional[str] = None
|
||||
soft_budget: Optional[float] = None
|
||||
|
|
@ -2015,7 +2033,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
|
||||
"""Represents all params for a LiteLLM_BudgetTable record"""
|
||||
"""LiteLLM_BudgetTable + server-managed fields returned on API responses."""
|
||||
|
||||
budget_reset_at: Optional[datetime] = None
|
||||
created_at: datetime
|
||||
|
|
@ -3695,7 +3713,11 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
|
|||
team_id: str
|
||||
budget_id: Optional[str] = None
|
||||
spend: Optional[float] = 0.0
|
||||
litellm_budget_table: Optional[LiteLLM_BudgetTable]
|
||||
total_spend: Optional[float] = 0.0
|
||||
# Union so Pydantic picks Full when data has server-managed fields
|
||||
# (/team/info) and Base when callers/tests construct with only
|
||||
# user-settable fields.
|
||||
litellm_budget_table: Optional[Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable]]
|
||||
|
||||
def safe_get_team_member_rpm_limit(self) -> Optional[int]:
|
||||
if self.litellm_budget_table is not None:
|
||||
|
|
@ -3898,7 +3920,7 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse):
|
|||
|
||||
|
||||
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
||||
team_member_budget_table: Optional[LiteLLM_BudgetTable] = None
|
||||
team_member_budget_table: Optional[LiteLLM_BudgetTableFull] = None
|
||||
# Resources inherited from access groups (separate from direct assignments)
|
||||
access_group_models: Optional[List[str]] = None
|
||||
access_group_mcp_server_ids: Optional[List[str]] = None
|
||||
|
|
|
|||
|
|
@ -626,11 +626,17 @@ async def common_checks( # noqa: PLR0915
|
|||
and user_object.max_budget is not None
|
||||
):
|
||||
user_budget = user_object.max_budget
|
||||
if user_budget < user_object.spend:
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
user_spend = await get_current_spend(
|
||||
counter_key=f"spend:user:{user_object.user_id}",
|
||||
fallback_spend=user_object.spend or 0.0,
|
||||
)
|
||||
if user_spend >= user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=user_object.spend,
|
||||
current_cost=user_spend,
|
||||
max_budget=user_budget,
|
||||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}",
|
||||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
|
||||
)
|
||||
|
||||
## 4.2 check team member budget, if team key
|
||||
|
|
@ -899,6 +905,63 @@ async def get_default_end_user_budget(
|
|||
return None
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_team_member_default_budget(
|
||||
budget_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
) -> Optional[LiteLLM_BudgetTable]:
|
||||
"""
|
||||
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
|
||||
|
||||
This budget is applied to team members whose TeamMembership row has no
|
||||
linked budget. Results are cached for performance.
|
||||
|
||||
Args:
|
||||
budget_id: The budget_id pulled from team.metadata["team_member_budget_id"]
|
||||
prisma_client: Database client instance
|
||||
user_api_key_cache: Cache for storing/retrieving budget data
|
||||
|
||||
Returns:
|
||||
LiteLLM_BudgetTable if found, None otherwise
|
||||
"""
|
||||
if prisma_client is None:
|
||||
return None
|
||||
|
||||
cache_key = f"team_member_default_budget:{budget_id}"
|
||||
|
||||
cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
if isinstance(cached_budget, LiteLLM_BudgetTable):
|
||||
return cached_budget
|
||||
if isinstance(cached_budget, dict):
|
||||
return LiteLLM_BudgetTable(**cached_budget)
|
||||
|
||||
try:
|
||||
budget_record = await prisma_client.db.litellm_budgettable.find_unique(
|
||||
where={"budget_id": budget_id}
|
||||
)
|
||||
|
||||
if budget_record is None:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Team-default member budget not found in database: {budget_id}"
|
||||
)
|
||||
return None
|
||||
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=budget_record.dict(),
|
||||
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
|
||||
)
|
||||
|
||||
return LiteLLM_BudgetTable(**budget_record.dict())
|
||||
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
f"Error fetching team-default member budget {budget_id}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def _apply_default_budget_to_end_user(
|
||||
end_user_obj: LiteLLM_EndUserTable,
|
||||
prisma_client: PrismaClient,
|
||||
|
|
@ -3224,13 +3287,31 @@ async def _check_team_member_budget(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Per-member override wins; otherwise fall back to the team-level
|
||||
# default configured via team.metadata["team_member_budget_id"].
|
||||
team_member_budget: Optional[float] = None
|
||||
if (
|
||||
team_membership is not None
|
||||
and team_membership.litellm_budget_table is not None
|
||||
and team_membership.litellm_budget_table.max_budget is not None
|
||||
):
|
||||
team_member_budget = team_membership.litellm_budget_table.max_budget
|
||||
team_member_spend = team_membership.spend or 0.0
|
||||
else:
|
||||
default_budget_id = (team_object.metadata or {}).get(
|
||||
"team_member_budget_id"
|
||||
)
|
||||
if isinstance(default_budget_id, str):
|
||||
default_budget = await get_team_member_default_budget(
|
||||
budget_id=default_budget_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
if default_budget is not None:
|
||||
team_member_budget = default_budget.max_budget
|
||||
|
||||
if team_member_budget is not None:
|
||||
team_member_spend = (
|
||||
team_membership.spend if team_membership is not None else 0.0
|
||||
) or 0.0
|
||||
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
|
@ -3665,12 +3746,20 @@ async def _organization_max_budget_check(
|
|||
if org_max_budget is None or org_max_budget <= 0:
|
||||
return
|
||||
|
||||
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
org_spend = await get_current_spend(
|
||||
counter_key=f"spend:org:{org_id}",
|
||||
fallback_spend=org_table.spend or 0.0,
|
||||
)
|
||||
|
||||
# Check if organization spend exceeds max budget
|
||||
if org_table.spend >= org_max_budget:
|
||||
if org_spend >= org_max_budget:
|
||||
# Trigger budget alert
|
||||
call_info = CallInfo(
|
||||
token=valid_token.token,
|
||||
spend=org_table.spend,
|
||||
spend=org_spend,
|
||||
max_budget=org_max_budget,
|
||||
user_id=valid_token.user_id,
|
||||
team_id=valid_token.team_id,
|
||||
|
|
@ -3686,9 +3775,9 @@ async def _organization_max_budget_check(
|
|||
)
|
||||
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=org_table.spend,
|
||||
current_cost=org_spend,
|
||||
max_budget=org_max_budget,
|
||||
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}",
|
||||
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -151,7 +151,15 @@ def is_request_body_safe(
|
|||
A malicious user can set the api_base to their own domain and invoke POST /chat/completions to intercept and steal the OpenAI API key.
|
||||
Relevant issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997
|
||||
"""
|
||||
banned_params = ["api_base", "base_url", "user_config"]
|
||||
banned_params = [
|
||||
"api_base",
|
||||
"base_url",
|
||||
"user_config",
|
||||
"aws_sts_endpoint",
|
||||
"aws_web_identity_token",
|
||||
"aws_role_name",
|
||||
"vertex_credentials",
|
||||
]
|
||||
|
||||
for param in banned_params:
|
||||
if (
|
||||
|
|
|
|||
|
|
@ -300,7 +300,7 @@ class RouteChecks:
|
|||
return True
|
||||
|
||||
if RouteChecks.check_route_access(
|
||||
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
|
||||
route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value
|
||||
):
|
||||
return True
|
||||
|
||||
|
|
@ -358,7 +358,9 @@ class RouteChecks:
|
|||
"""
|
||||
Check if route is a management route
|
||||
"""
|
||||
return route in LiteLLMRoutes.management_routes.value
|
||||
return RouteChecks.check_route_access(
|
||||
route=route, allowed_routes=LiteLLMRoutes.management_routes.value
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def is_info_route(route: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -433,7 +433,8 @@ def add_guardrail_to_applied_guardrails_header(
|
|||
return
|
||||
_metadata = request_data.get("metadata", None) or {}
|
||||
if "applied_guardrails" in _metadata:
|
||||
_metadata["applied_guardrails"].append(guardrail_name)
|
||||
if guardrail_name not in _metadata["applied_guardrails"]:
|
||||
_metadata["applied_guardrails"].append(guardrail_name)
|
||||
else:
|
||||
_metadata["applied_guardrails"] = [guardrail_name]
|
||||
# Ensure metadata is set back to request_data (important when metadata didn't exist)
|
||||
|
|
|
|||
|
|
@ -632,20 +632,27 @@ class ResetBudgetJob:
|
|||
|
||||
now = datetime.utcnow()
|
||||
|
||||
# Note on raw SQL: prisma-client-python does not support null-filtering
|
||||
# on `Json?` columns (no DbNull/JsonNull sentinel — see
|
||||
# RobertCraigie/prisma-client-py#714). We use `query_raw` with
|
||||
# `IS NOT NULL` so we don't materialize every key/team row on each
|
||||
# tick of the reset job. Writes still go through the ORM.
|
||||
|
||||
# --- Keys ---
|
||||
try:
|
||||
all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many(
|
||||
where={"budget_limits": {"not": None}} # type: ignore[arg-type]
|
||||
key_rows = await self.prisma_client.db.query_raw(
|
||||
'SELECT token, budget_limits FROM "LiteLLM_VerificationToken" '
|
||||
"WHERE budget_limits IS NOT NULL"
|
||||
)
|
||||
for key in all_keys:
|
||||
raw = key.budget_limits # type: ignore[attr-defined]
|
||||
for row in key_rows:
|
||||
raw = row["budget_limits"]
|
||||
if not raw:
|
||||
continue
|
||||
windows: list = raw if isinstance(raw, list) else json.loads(raw)
|
||||
changed = False
|
||||
for window in windows:
|
||||
counter_key = (
|
||||
f"spend:key:{key.token}:window:{window['budget_duration']}"
|
||||
f"spend:key:{row['token']}:window:{window['budget_duration']}"
|
||||
)
|
||||
if await ResetBudgetJob._reset_expired_window(
|
||||
window, counter_key, spend_counter_cache, now
|
||||
|
|
@ -653,7 +660,7 @@ class ResetBudgetJob:
|
|||
changed = True
|
||||
if changed:
|
||||
await self.prisma_client.db.litellm_verificationtoken.update(
|
||||
where={"token": key.token},
|
||||
where={"token": row["token"]},
|
||||
data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type]
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -663,26 +670,25 @@ class ResetBudgetJob:
|
|||
|
||||
# --- Teams ---
|
||||
try:
|
||||
all_teams = await self.prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"budget_limits": {"not": None}} # type: ignore[arg-type]
|
||||
team_rows = await self.prisma_client.db.query_raw(
|
||||
'SELECT team_id, budget_limits FROM "LiteLLM_TeamTable" '
|
||||
"WHERE budget_limits IS NOT NULL"
|
||||
)
|
||||
for team in all_teams:
|
||||
raw = team.budget_limits # type: ignore[attr-defined]
|
||||
for row in team_rows:
|
||||
raw = row["budget_limits"]
|
||||
if not raw:
|
||||
continue
|
||||
windows = raw if isinstance(raw, list) else json.loads(raw)
|
||||
changed = False
|
||||
for window in windows:
|
||||
counter_key = (
|
||||
f"spend:team:{team.team_id}:window:{window['budget_duration']}"
|
||||
)
|
||||
counter_key = f"spend:team:{row['team_id']}:window:{window['budget_duration']}"
|
||||
if await ResetBudgetJob._reset_expired_window(
|
||||
window, counter_key, spend_counter_cache, now
|
||||
):
|
||||
changed = True
|
||||
if changed:
|
||||
await self.prisma_client.db.litellm_teamtable.update(
|
||||
where={"team_id": team.team_id},
|
||||
where={"team_id": row["team_id"]},
|
||||
data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type]
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1300,7 +1300,10 @@ class DBSpendUpdateWriter:
|
|||
|
||||
batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists
|
||||
where={"team_id": team_id, "user_id": user_id},
|
||||
data={"spend": {"increment": response_cost}},
|
||||
data={
|
||||
"spend": {"increment": response_cost},
|
||||
"total_spend": {"increment": response_cost},
|
||||
},
|
||||
)
|
||||
# Transaction succeeded, break out of retry loop
|
||||
break
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._types import (
|
||||
BaseDailySpendTransaction,
|
||||
DailyAgentSpendTransaction,
|
||||
DailyEndUserSpendTransaction,
|
||||
DailyOrganizationSpendTransaction,
|
||||
|
|
@ -29,6 +30,8 @@ from litellm.proxy._types import (
|
|||
DailyTeamSpendTransaction,
|
||||
DailyUserSpendTransaction,
|
||||
DBSpendUpdateTransactions,
|
||||
Litellm_EntityType,
|
||||
SpendUpdateQueueItem,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj
|
||||
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
|
||||
|
|
@ -259,9 +262,36 @@ class RedisUpdateBuffer:
|
|||
if len(rpush_list) == 0:
|
||||
return
|
||||
|
||||
result_lengths = await self.redis_cache.async_rpush_pipeline(
|
||||
rpush_list=rpush_list,
|
||||
)
|
||||
try:
|
||||
result_lengths = await self.redis_cache.async_rpush_pipeline(
|
||||
rpush_list=rpush_list,
|
||||
)
|
||||
except Exception as e:
|
||||
# The in-memory queues were already drained above. If we let the
|
||||
# exception propagate without restoring, the aggregated spend is
|
||||
# permanently lost. Re-enqueue so the next scheduler tick retries.
|
||||
verbose_proxy_logger.error(
|
||||
"Spend tracking - failed to push aggregated spend updates to Redis. "
|
||||
"Restoring %d transaction sets to in-memory queues for retry on next tick. "
|
||||
"Error: %s",
|
||||
len(rpush_list),
|
||||
str(e),
|
||||
)
|
||||
await self._restore_spend_updates_to_in_memory_queues(
|
||||
db_spend_update_transactions=db_spend_update_transactions,
|
||||
daily_spend_update_transactions=daily_spend_update_transactions,
|
||||
daily_team_spend_update_transactions=daily_team_spend_update_transactions,
|
||||
daily_org_spend_update_transactions=daily_org_spend_update_transactions,
|
||||
daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions,
|
||||
daily_agent_spend_update_transactions=daily_agent_spend_update_transactions,
|
||||
spend_update_queue=spend_update_queue,
|
||||
daily_spend_update_queue=daily_spend_update_queue,
|
||||
daily_team_spend_update_queue=daily_team_spend_update_queue,
|
||||
daily_org_spend_update_queue=daily_org_spend_update_queue,
|
||||
daily_end_user_spend_update_queue=daily_end_user_spend_update_queue,
|
||||
daily_agent_spend_update_queue=daily_agent_spend_update_queue,
|
||||
)
|
||||
return
|
||||
|
||||
# Emit gauge events for each queue
|
||||
for i, queue_size in enumerate(result_lengths):
|
||||
|
|
@ -271,6 +301,101 @@ class RedisUpdateBuffer:
|
|||
service=service_types[i],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _restore_spend_updates_to_in_memory_queues(
|
||||
db_spend_update_transactions: Optional[DBSpendUpdateTransactions],
|
||||
daily_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]],
|
||||
daily_team_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
daily_org_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
daily_end_user_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
daily_agent_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
spend_update_queue: SpendUpdateQueue,
|
||||
daily_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_team_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_org_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_agent_spend_update_queue: DailySpendUpdateQueue,
|
||||
) -> None:
|
||||
"""
|
||||
Put drained-but-unpushed transactions back into in-memory queues.
|
||||
|
||||
Called when the Redis rpush pipeline raises. Without this, all spend
|
||||
data aggregated during the current scheduler tick is permanently lost
|
||||
because the source queues were already drained before the rpush.
|
||||
"""
|
||||
if db_spend_update_transactions is not None:
|
||||
entity_entries: List[
|
||||
Tuple[Litellm_EntityType, Optional[Dict[str, float]]]
|
||||
] = [
|
||||
(
|
||||
Litellm_EntityType.USER,
|
||||
db_spend_update_transactions.get("user_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.END_USER,
|
||||
db_spend_update_transactions.get("end_user_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.KEY,
|
||||
db_spend_update_transactions.get("key_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TEAM,
|
||||
db_spend_update_transactions.get("team_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TEAM_MEMBER,
|
||||
db_spend_update_transactions.get("team_member_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.ORGANIZATION,
|
||||
db_spend_update_transactions.get("org_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TAG,
|
||||
db_spend_update_transactions.get("tag_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.AGENT,
|
||||
db_spend_update_transactions.get("agent_list_transactions"),
|
||||
),
|
||||
]
|
||||
for entity_type, entities in entity_entries:
|
||||
if not entities:
|
||||
continue
|
||||
for entity_id, cost in entities.items():
|
||||
await spend_update_queue.add_update(
|
||||
SpendUpdateQueueItem(
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
response_cost=cost,
|
||||
)
|
||||
)
|
||||
|
||||
daily_pairs: List[
|
||||
Tuple[Optional[Dict[str, BaseDailySpendTransaction]], DailySpendUpdateQueue]
|
||||
] = [
|
||||
(daily_spend_update_transactions, daily_spend_update_queue),
|
||||
(daily_team_spend_update_transactions, daily_team_spend_update_queue),
|
||||
(daily_org_spend_update_transactions, daily_org_spend_update_queue),
|
||||
(
|
||||
daily_end_user_spend_update_transactions,
|
||||
daily_end_user_spend_update_queue,
|
||||
),
|
||||
(daily_agent_spend_update_transactions, daily_agent_spend_update_queue),
|
||||
]
|
||||
for daily_txns, daily_queue in daily_pairs:
|
||||
if daily_txns:
|
||||
await daily_queue.update_queue.put(daily_txns)
|
||||
|
||||
@staticmethod
|
||||
def _number_of_transactions_to_store_in_redis(
|
||||
db_spend_update_transactions: DBSpendUpdateTransactions,
|
||||
|
|
|
|||
|
|
@ -403,10 +403,18 @@ class PrismaManager:
|
|||
return dname
|
||||
|
||||
@staticmethod
|
||||
def setup_database(use_migrate: bool = False) -> bool:
|
||||
def setup_database(
|
||||
use_migrate: bool = False, use_v2_resolver: bool = False
|
||||
) -> bool:
|
||||
"""
|
||||
Set up the database using either prisma migrate or prisma db push
|
||||
|
||||
Args:
|
||||
use_migrate: Use `prisma migrate deploy` instead of `db push`.
|
||||
use_v2_resolver: Opt into the v2 migration resolver that avoids
|
||||
the diff-and-force recovery behavior (which caused schema
|
||||
thrashing during rolling deploys). Defaults to False.
|
||||
|
||||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
"""
|
||||
|
|
@ -427,7 +435,10 @@ class PrismaManager:
|
|||
|
||||
prisma_dir = PrismaManager._get_prisma_dir()
|
||||
|
||||
return ProxyExtrasDBManager.setup_database(use_migrate=use_migrate)
|
||||
return ProxyExtrasDBManager.setup_database(
|
||||
use_migrate=use_migrate,
|
||||
use_v2_resolver=use_v2_resolver,
|
||||
)
|
||||
else:
|
||||
# Use prisma db push with increased timeout
|
||||
subprocess.run(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,52 @@
|
|||
# Example proxy config for the adaptive router (v0).
|
||||
#
|
||||
# Wires one logical router ("smart-cheap-router") that adaptively picks between
|
||||
# two real deployments ("fast" and "smart") based on per-session feedback signals.
|
||||
#
|
||||
# How to use from a client:
|
||||
# POST /v1/chat/completions { "model": "smart-cheap-router", ... }
|
||||
# Add { "metadata": { "litellm_session_id": "<your-session-id>" } } to enable
|
||||
# sticky-session routing within a conversation.
|
||||
#
|
||||
# Required env vars: OPENAI_API_KEY, DATABASE_URL.
|
||||
|
||||
model_list:
|
||||
# ---- The adaptive router "control" deployment -------------------------
|
||||
# `model_name` is what clients call. `available_models` lists the underlying
|
||||
# deployments the router is allowed to pick from (must match other model_name
|
||||
# entries in this list).
|
||||
- model_name: smart-cheap-router
|
||||
litellm_params:
|
||||
model: auto_router/adaptive_router # required prefix -- triggers adaptive-router init
|
||||
adaptive_router_config:
|
||||
available_models: ["fast", "smart"]
|
||||
weights:
|
||||
quality: 0.7
|
||||
cost: 0.3
|
||||
|
||||
# ---- Underlying deployments the router picks from ---------------------
|
||||
- model_name: fast
|
||||
litellm_params:
|
||||
model: openai/gpt-4o-mini
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
input_cost_per_token: 0.00000015
|
||||
model_info:
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 2
|
||||
strengths: []
|
||||
|
||||
- model_name: smart
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
input_cost_per_token: 0.0000050
|
||||
model_info:
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 3
|
||||
strengths: ["code_generation", "technical_design", "analytical_reasoning"]
|
||||
|
||||
litellm_settings:
|
||||
drop_params: True
|
||||
|
||||
general_settings:
|
||||
master_key: sk-1234 # REPLACE in production
|
||||
|
|
@ -5,7 +5,6 @@
|
|||
# +-------------------------------------------------------------+
|
||||
# Thank you users! We ❤️ you! - Krrish & Ishaan
|
||||
|
||||
import copy
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
|
@ -18,6 +17,7 @@ from typing import (
|
|||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
ClassVar,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
|
|
@ -33,6 +33,7 @@ from fastapi import HTTPException
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.caching import DualCache
|
||||
from litellm.exceptions import GuardrailInterventionNormalStringError
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
|
@ -62,6 +63,7 @@ from litellm.types.utils import (
|
|||
CallTypesLiteral,
|
||||
Choices,
|
||||
GuardrailStatus,
|
||||
Message,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
|
|
@ -78,56 +80,33 @@ class GuardrailMessageFilterResult(NamedTuple):
|
|||
|
||||
|
||||
def _redact_pii_matches(response_json: dict) -> dict:
|
||||
try:
|
||||
# Create a deep copy to avoid modifying the original response
|
||||
redacted_response = copy.deepcopy(response_json)
|
||||
"""
|
||||
Redact match-like fields from a Bedrock ApplyGuardrail JSON payload.
|
||||
|
||||
# Get assessments from the response
|
||||
# NOTE: We use `.get("key") or []` instead of `.get("key", [])` because
|
||||
# the Bedrock API can return explicit `null` for list fields (e.g. "regexes": null).
|
||||
# In Python, dict.get("key", []) returns None (not []) when the key exists
|
||||
# with a None/null value. The `or []` ensures we always get an iterable,
|
||||
# preventing "TypeError: 'NoneType' object is not iterable".
|
||||
assessments = redacted_response.get("assessments") or []
|
||||
if not assessments:
|
||||
return redacted_response
|
||||
Delegates to :func:`redact_nested_match_and_regex_keys` (same rules as spend
|
||||
logging). Kept as a Bedrock-module entry point for existing unit tests.
|
||||
"""
|
||||
redacted = redact_nested_match_and_regex_keys(response_json)
|
||||
return redacted if isinstance(redacted, dict) else response_json
|
||||
|
||||
for assessment in assessments:
|
||||
# Redact PII entities in sensitive information policy
|
||||
sensitive_info_policy = assessment.get("sensitiveInformationPolicy")
|
||||
if sensitive_info_policy:
|
||||
pii_entities = sensitive_info_policy.get("piiEntities") or []
|
||||
for pii_entity in pii_entities:
|
||||
if "match" in pii_entity:
|
||||
pii_entity["match"] = "[REDACTED]"
|
||||
|
||||
# Redact regex matches
|
||||
regexes = sensitive_info_policy.get("regexes") or []
|
||||
for regex_match in regexes:
|
||||
if "match" in regex_match:
|
||||
regex_match["match"] = "[REDACTED]"
|
||||
def _redact_assessment_match_fields(assessments: List[dict]) -> List[dict]:
|
||||
"""
|
||||
Redact sensitive match-like fields from blocked assessment summaries.
|
||||
|
||||
# Redact custom word matches in word policy
|
||||
word_policy = assessment.get("wordPolicy")
|
||||
if word_policy:
|
||||
custom_words = word_policy.get("customWords") or []
|
||||
for custom_word in custom_words:
|
||||
if "match" in custom_word:
|
||||
custom_word["match"] = "[REDACTED]"
|
||||
|
||||
managed_words = word_policy.get("managedWordLists") or []
|
||||
for managed_word in managed_words:
|
||||
if "match" in managed_word:
|
||||
managed_word["match"] = "[REDACTED]"
|
||||
|
||||
return redacted_response
|
||||
except Exception as e:
|
||||
# We do not want to fail in any case so this is just a warning
|
||||
verbose_proxy_logger.warning("Guardrail log redaction failed: %s", str(e))
|
||||
return response_json
|
||||
This is used for customer-visible error payloads (HTTPException.detail) where
|
||||
we want to preserve policy/type/action metadata without echoing raw matched
|
||||
content.
|
||||
"""
|
||||
redacted = redact_nested_match_and_regex_keys(assessments)
|
||||
return redacted if isinstance(redacted, list) else assessments
|
||||
|
||||
|
||||
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
||||
# During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
|
||||
# OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
|
||||
use_native_during_call_hook: ClassVar[bool] = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrailIdentifier: Optional[str] = None,
|
||||
|
|
@ -418,6 +397,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
messages: Optional[List[AllMessageValues]] = None,
|
||||
response: Optional[Union[Any, litellm.ModelResponse]] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
logging_event_type: Optional[GuardrailEventHooks] = None,
|
||||
) -> BedrockGuardrailResponse:
|
||||
from datetime import datetime
|
||||
|
||||
|
|
@ -455,11 +435,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
prepared_request.headers,
|
||||
)
|
||||
|
||||
event_type = (
|
||||
GuardrailEventHooks.pre_call
|
||||
if source == "INPUT"
|
||||
else GuardrailEventHooks.post_call
|
||||
)
|
||||
# UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API
|
||||
# body, which must not be confused with the proxy hook (pre_call / during_call /
|
||||
# post_call). When omitted, keep legacy mapping for backward compatibility.
|
||||
if logging_event_type is not None:
|
||||
event_type = logging_event_type
|
||||
else:
|
||||
event_type = (
|
||||
GuardrailEventHooks.pre_call
|
||||
if source == "INPUT"
|
||||
else GuardrailEventHooks.post_call
|
||||
)
|
||||
|
||||
try:
|
||||
httpx_response = await self.async_handler.post(
|
||||
|
|
@ -514,9 +500,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
# Add guardrail information to request trace
|
||||
#########################################################
|
||||
_json_response = httpx_response.json()
|
||||
# Raw Bedrock JSON is passed here; match/regex redaction runs once inside
|
||||
# CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response=httpx_response.json(),
|
||||
guardrail_json_response=_json_response,
|
||||
request_data=request_data or {},
|
||||
guardrail_status=self._get_bedrock_guardrail_response_status(
|
||||
response=httpx_response
|
||||
|
|
@ -529,9 +518,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
if httpx_response.status_code == 200:
|
||||
# check if the response was flagged
|
||||
_json_response = httpx_response.json()
|
||||
redacted_response = _redact_pii_matches(_json_response)
|
||||
verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response)
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock AI response : %s",
|
||||
redact_nested_match_and_regex_keys(_json_response),
|
||||
)
|
||||
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
|
||||
if self._should_raise_guardrail_blocked_exception(
|
||||
bedrock_guardrail_response
|
||||
|
|
@ -808,7 +798,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
|
||||
assessments = self._extract_blocked_assessments(response)
|
||||
if assessments:
|
||||
detail["assessments"] = assessments
|
||||
detail["assessments"] = _redact_assessment_match_fields(assessments)
|
||||
|
||||
return HTTPException(status_code=400, detail=detail)
|
||||
|
||||
|
|
@ -830,8 +820,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
return False
|
||||
|
||||
# Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED)
|
||||
# NOTE: Use `or []` instead of default param to handle explicit null from Bedrock API.
|
||||
# See _redact_pii_matches() for detailed explanation of the null safety pattern.
|
||||
# NOTE: Use `.get("k") or []` not `.get("k", [])` — Bedrock can return explicit
|
||||
# JSON null; dict.get("k", []) then yields None, and `for x in None` raises.
|
||||
assessments = response.get("assessments") or []
|
||||
if not assessments:
|
||||
return False
|
||||
|
|
@ -951,7 +941,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
try:
|
||||
bedrock_guardrail_response = await self.make_bedrock_api_request(
|
||||
source="INPUT", messages=filtered_messages, request_data=data
|
||||
source="INPUT",
|
||||
messages=filtered_messages,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
bedrock_guardrail_response = e.message
|
||||
|
|
@ -1023,7 +1016,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
try:
|
||||
bedrock_guardrail_response = await self.make_bedrock_api_request(
|
||||
source="INPUT", messages=filtered_messages, request_data=data
|
||||
source="INPUT",
|
||||
messages=filtered_messages,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.during_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
bedrock_guardrail_response = e.message
|
||||
|
|
@ -1127,9 +1123,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
source="INPUT",
|
||||
messages=input_messages,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
output_task = self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=response, request_data=data
|
||||
source="OUTPUT",
|
||||
response=response,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
|
||||
# Execute both requests in parallel
|
||||
|
|
@ -1143,7 +1143,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# Only run OUTPUT validation (INPUT was already validated in pre_call or during_call)
|
||||
try:
|
||||
output_content_bedrock = await self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=response, request_data=data
|
||||
source="OUTPUT",
|
||||
response=response,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
output_content_bedrock = e.message
|
||||
|
|
@ -1270,9 +1273,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
source="INPUT",
|
||||
messages=input_messages,
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
) # Only input messages
|
||||
output_task = self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=assembled_model_response
|
||||
source="OUTPUT",
|
||||
response=assembled_model_response,
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
) # Only response
|
||||
|
||||
# Execute both requests in parallel
|
||||
|
|
@ -1286,7 +1293,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# Only run OUTPUT validation (INPUT was already validated in pre_call or during_call)
|
||||
try:
|
||||
output_guardrail_response = await self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=assembled_model_response
|
||||
source="OUTPUT",
|
||||
response=assembled_model_response,
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
output_guardrail_response = e.message
|
||||
|
|
@ -1563,11 +1573,50 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
|
||||
# Bedrock will throw an error if there is no text to process
|
||||
if filtered_messages:
|
||||
bedrock_response = await self.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=filtered_messages,
|
||||
request_data=request_data,
|
||||
_log_hook = (
|
||||
GuardrailEventHooks.pre_call
|
||||
if input_type == "request"
|
||||
else GuardrailEventHooks.post_call
|
||||
)
|
||||
# Map the abstract input_type to the Bedrock source parameter.
|
||||
# "request" -> INPUT (scan user-supplied content)
|
||||
# "response" -> OUTPUT (scan model-generated content)
|
||||
# Bedrock guardrail policies are often configured differently
|
||||
# for Input vs Output (e.g. PII blocking only on Output), so
|
||||
# the source MUST match where the text originated.
|
||||
bedrock_source: Literal["INPUT", "OUTPUT"] = (
|
||||
"OUTPUT" if input_type == "response" else "INPUT"
|
||||
)
|
||||
if bedrock_source == "OUTPUT":
|
||||
# Build a synthetic ModelResponse whose choices carry the
|
||||
# text(s) to scan, so _create_bedrock_output_content_request
|
||||
# can produce the correct Bedrock OUTPUT payload.
|
||||
synthetic_response = ModelResponse(
|
||||
choices=[
|
||||
Choices(
|
||||
index=_idx,
|
||||
message=Message(
|
||||
role="assistant",
|
||||
content=str(_msg.get("content") or ""),
|
||||
),
|
||||
finish_reason="stop",
|
||||
)
|
||||
for _idx, _msg in enumerate(filtered_messages)
|
||||
]
|
||||
)
|
||||
bedrock_response = await self.make_bedrock_api_request(
|
||||
source="OUTPUT",
|
||||
response=synthetic_response,
|
||||
request_data=request_data,
|
||||
logging_event_type=_log_hook,
|
||||
)
|
||||
else:
|
||||
bedrock_response = await self.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=filtered_messages,
|
||||
request_data=request_data,
|
||||
logging_event_type=_log_hook,
|
||||
)
|
||||
|
||||
# Apply any masking that was applied by the guardrail
|
||||
output_list = bedrock_response.get("output")
|
||||
|
|
|
|||
|
|
@ -21,20 +21,30 @@ class _PROXY_MaxBudgetLimiter(CustomLogger):
|
|||
):
|
||||
try:
|
||||
verbose_proxy_logger.debug("Inside Max Budget Limiter Pre-Call Hook")
|
||||
cache_key = f"{user_api_key_dict.user_id}_user_api_key_user_id"
|
||||
user_row = await cache.async_get_cache(
|
||||
cache_key, parent_otel_span=user_api_key_dict.parent_otel_span
|
||||
max_budget = user_api_key_dict.user_max_budget
|
||||
user_id = user_api_key_dict.user_id
|
||||
|
||||
if max_budget is None or user_id is None:
|
||||
return
|
||||
|
||||
# Personal budget applies only to non-team requests, matching
|
||||
# the explicit team-key exemption in common_checks section 4.1.
|
||||
if user_api_key_dict.team_id is not None:
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
curr_spend = await get_current_spend(
|
||||
counter_key=f"spend:user:{user_id}",
|
||||
fallback_spend=user_api_key_dict.user_spend or 0.0,
|
||||
)
|
||||
if user_row is None: # value not yet cached
|
||||
return
|
||||
max_budget = user_row["max_budget"]
|
||||
curr_spend = user_row["spend"]
|
||||
|
||||
if max_budget is None:
|
||||
return
|
||||
|
||||
if curr_spend is None:
|
||||
return
|
||||
verbose_proxy_logger.debug(
|
||||
"MaxBudgetLimiter: user_id=%s, spend=%.6f, max=%.6f",
|
||||
user_id,
|
||||
curr_spend,
|
||||
max_budget,
|
||||
)
|
||||
|
||||
# CHECK IF REQUEST ALLOWED
|
||||
if curr_spend >= max_budget:
|
||||
|
|
|
|||
|
|
@ -213,6 +213,7 @@ class _ProxyDBLogger(CustomLogger):
|
|||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
response_cost=response_cost,
|
||||
org_id=org_id,
|
||||
)
|
||||
|
||||
# update cache (fire-and-forget for backward compat:
|
||||
|
|
|
|||
|
|
@ -285,6 +285,13 @@ async def image_edit_api(
|
|||
if mask_files:
|
||||
data["mask"] = mask_files
|
||||
|
||||
for _field in ("image", "mask"):
|
||||
if _field in data and isinstance(data[_field], str):
|
||||
raise HTTPException(
|
||||
status_code=422,
|
||||
detail=f"'{_field}' must be provided as a multipart file upload, not a string.",
|
||||
)
|
||||
|
||||
# Ensure prompt exists in data (default to None for models that don't require it)
|
||||
if "prompt" not in data:
|
||||
data["prompt"] = None
|
||||
|
|
|
|||
|
|
@ -355,6 +355,7 @@ async def _upsert_budget_and_membership(
|
|||
tpm_limit: Optional[int] = None,
|
||||
rpm_limit: Optional[int] = None,
|
||||
allowed_models: Optional[List[str]] = None,
|
||||
team_default_budget_id: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Helper function to Create/Update or Delete the budget within the team membership
|
||||
|
|
@ -368,6 +369,11 @@ async def _upsert_budget_and_membership(
|
|||
tpm_limit: Tokens per minute limit for the team member
|
||||
rpm_limit: Requests per minute limit for the team member
|
||||
allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce.
|
||||
team_default_budget_id: The team's shared default member budget id (from
|
||||
team metadata.team_member_budget_id), if any. When the membership's
|
||||
existing_budget_id matches this, we clone-on-write so editing one
|
||||
member's budget does not mutate the shared default (and therefore
|
||||
every other member who still points at it).
|
||||
|
||||
If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership.
|
||||
If any of these values exist, a budget is updated or created and linked to the team membership.
|
||||
|
|
@ -385,7 +391,13 @@ async def _upsert_budget_and_membership(
|
|||
)
|
||||
return
|
||||
|
||||
if existing_budget_id is not None:
|
||||
is_shared_default = (
|
||||
existing_budget_id is not None
|
||||
and team_default_budget_id is not None
|
||||
and existing_budget_id == team_default_budget_id
|
||||
)
|
||||
|
||||
if existing_budget_id is not None and not is_shared_default:
|
||||
# Update the existing budget in-place to preserve fields not being changed.
|
||||
# Only write fields that the caller explicitly provided (non-None).
|
||||
update_data: Dict[str, Any] = {
|
||||
|
|
@ -405,11 +417,40 @@ async def _upsert_budget_and_membership(
|
|||
)
|
||||
return
|
||||
|
||||
# No existing budget — create a new one and link it to the membership.
|
||||
# Either there is no existing budget, OR the membership is still pointing
|
||||
# at the team's shared default member budget. In both cases we create a
|
||||
# NEW private budget for this user and (re)link the membership to it.
|
||||
create_data: Dict[str, Any] = {
|
||||
"created_by": user_api_key_dict.user_id or "",
|
||||
"updated_by": user_api_key_dict.user_id or "",
|
||||
}
|
||||
|
||||
# If we're forking off the shared default, seed the new row with the
|
||||
# default's values so fields the caller did not change carry over.
|
||||
if is_shared_default:
|
||||
default_budget_row = await tx.litellm_budgettable.find_unique(
|
||||
where={"budget_id": existing_budget_id}
|
||||
)
|
||||
if default_budget_row is not None:
|
||||
default_budget_dict = default_budget_row.model_dump()
|
||||
for field in (
|
||||
"max_budget",
|
||||
"soft_budget",
|
||||
"max_parallel_requests",
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
"model_max_budget",
|
||||
"budget_duration",
|
||||
"allowed_models",
|
||||
):
|
||||
value = default_budget_dict.get(field)
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, list) and len(value) == 0:
|
||||
continue
|
||||
create_data[field] = value
|
||||
|
||||
# Caller-provided values take precedence over the cloned defaults.
|
||||
if max_budget is not None:
|
||||
create_data["max_budget"] = max_budget
|
||||
if tpm_limit is not None:
|
||||
|
|
|
|||
|
|
@ -52,12 +52,17 @@ from litellm.proxy._experimental.mcp_server.utils import (
|
|||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
validate_and_normalize_mcp_server_payload as _base_validate_and_normalize_mcp_server_payload,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
|
||||
router = APIRouter(prefix="/v1/mcp", tags=["mcp"])
|
||||
|
||||
MCP_AVAILABLE: bool = True
|
||||
|
||||
TEMPORARY_MCP_SERVER_TTL_SECONDS = 300
|
||||
TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX = "litellm:mcp:temporary_server"
|
||||
|
||||
|
||||
def does_mcp_server_exist(
|
||||
|
|
@ -329,13 +334,115 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
return server
|
||||
|
||||
def get_cached_temporary_mcp_server(
|
||||
async def _cache_temporary_mcp_server_in_redis(
|
||||
server: MCPServer, ttl_seconds: int
|
||||
) -> None:
|
||||
"""
|
||||
Best-effort write-through to Redis so temporary MCP OAuth sessions are
|
||||
shared across proxy instances. Keep local in-memory cache as fallback.
|
||||
"""
|
||||
if litellm.cache is None or not hasattr(litellm.cache, "cache"):
|
||||
return
|
||||
cache_backend = getattr(litellm.cache, "cache", None)
|
||||
if cache_backend is None or not hasattr(cache_backend, "async_set_cache"):
|
||||
return
|
||||
|
||||
payload: Dict[str, Any] = server.model_dump(mode="json")
|
||||
payload_json = json.dumps(payload)
|
||||
try:
|
||||
encrypted_payload = encrypt_value_helper(payload_json)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Failed to encrypt temporary MCP server payload for Redis cache: {str(e)}"
|
||||
)
|
||||
return
|
||||
|
||||
if not isinstance(encrypted_payload, str):
|
||||
verbose_proxy_logger.debug(
|
||||
"Encrypted temporary MCP payload is not a string; skipping Redis cache write"
|
||||
)
|
||||
return
|
||||
|
||||
try:
|
||||
await cache_backend.async_set_cache(
|
||||
key=f"{TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX}:{server.server_id}",
|
||||
value=encrypted_payload,
|
||||
ttl=max(1, ttl_seconds),
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Failed to write temporary MCP server to Redis cache: {str(e)}"
|
||||
)
|
||||
|
||||
async def _get_temporary_mcp_server_from_redis(
|
||||
server_id: str,
|
||||
) -> Optional[MCPServer]:
|
||||
"""
|
||||
Best-effort read from Redis shared cache. Returns None on miss/errors.
|
||||
|
||||
Values must be encrypted strings (same contract as _cache_temporary_mcp_server_in_redis);
|
||||
legacy plaintext dict payloads are rejected.
|
||||
"""
|
||||
if litellm.cache is None or not hasattr(litellm.cache, "cache"):
|
||||
return None
|
||||
cache_backend = getattr(litellm.cache, "cache", None)
|
||||
if cache_backend is None or not hasattr(cache_backend, "async_get_cache"):
|
||||
return None
|
||||
|
||||
try:
|
||||
cached_server = await cache_backend.async_get_cache(
|
||||
key=f"{TEMPORARY_MCP_SERVER_REDIS_KEY_PREFIX}:{server_id}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Failed reading temporary MCP server from Redis cache: {str(e)}"
|
||||
)
|
||||
return None
|
||||
|
||||
if not isinstance(cached_server, str):
|
||||
verbose_proxy_logger.debug(
|
||||
"Temporary MCP Redis cache value must be an encrypted string; rejecting non-string payload"
|
||||
)
|
||||
return None
|
||||
|
||||
decrypted_json = decrypt_value_helper(
|
||||
value=cached_server,
|
||||
key="temporary_mcp_server",
|
||||
exception_type="debug",
|
||||
)
|
||||
if decrypted_json is None:
|
||||
return None
|
||||
try:
|
||||
loaded = json.loads(decrypted_json)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Invalid decrypted temporary MCP payload in Redis cache: {str(e)}"
|
||||
)
|
||||
return None
|
||||
if not isinstance(loaded, dict):
|
||||
return None
|
||||
payload_dict: Dict[str, Any] = loaded
|
||||
|
||||
try:
|
||||
return MCPServer(**payload_dict)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Invalid temporary MCP server payload in Redis cache: {str(e)}"
|
||||
)
|
||||
return None
|
||||
|
||||
async def get_cached_temporary_mcp_server(
|
||||
server_id: str,
|
||||
) -> Optional[MCPServer]:
|
||||
_prune_expired_temporary_mcp_servers()
|
||||
entry = _temporary_mcp_servers.get(server_id)
|
||||
if entry is None:
|
||||
return None
|
||||
redis_server = await _get_temporary_mcp_server_from_redis(server_id)
|
||||
if redis_server is None:
|
||||
return None
|
||||
# Intentionally avoid repopulating local cache from Redis to prevent
|
||||
# extending effective lifetime beyond the remaining Redis TTL.
|
||||
return redis_server
|
||||
return entry.server
|
||||
|
||||
def _redact_mcp_credentials(
|
||||
|
|
@ -1325,6 +1432,10 @@ if MCP_AVAILABLE:
|
|||
temporary_server,
|
||||
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
|
||||
)
|
||||
await _cache_temporary_mcp_server_in_redis(
|
||||
temporary_server,
|
||||
ttl_seconds=TEMPORARY_MCP_SERVER_TTL_SECONDS,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
f"Error caching temporary mcp server: {str(e)}"
|
||||
|
|
@ -1336,18 +1447,24 @@ if MCP_AVAILABLE:
|
|||
|
||||
return _redact_mcp_credentials(temp_record)
|
||||
|
||||
def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer:
|
||||
server = get_cached_temporary_mcp_server(server_id)
|
||||
async def _get_cached_temporary_mcp_server_or_404(
|
||||
server_id: str, request: Optional[Request] = None
|
||||
) -> MCPServer:
|
||||
server = await get_cached_temporary_mcp_server(server_id)
|
||||
if server is None:
|
||||
# Fall back to real DB/config server (e.g. for the user-side OAuth flow
|
||||
# which calls these endpoints with a real server_id, not a temp session id).
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request) if request else None
|
||||
server = global_mcp_server_manager.get_mcp_server_by_id(
|
||||
server_id
|
||||
) or global_mcp_server_manager.get_mcp_server_by_name(server_id)
|
||||
) or global_mcp_server_manager.get_mcp_server_by_name(
|
||||
server_id, client_ip=client_ip
|
||||
)
|
||||
if server is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
|
|
@ -1358,10 +1475,12 @@ if MCP_AVAILABLE:
|
|||
@router.get(
|
||||
"/server/oauth/{server_id}/authorize",
|
||||
include_in_schema=False,
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def mcp_authorize(
|
||||
request: Request,
|
||||
server_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
client_id: Optional[str] = None,
|
||||
redirect_uri: str = Query(...),
|
||||
state: str = "",
|
||||
|
|
@ -1370,7 +1489,9 @@ if MCP_AVAILABLE:
|
|||
response_type: Optional[str] = None,
|
||||
scope: Optional[str] = None,
|
||||
):
|
||||
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
|
||||
mcp_server = await _get_cached_temporary_mcp_server_or_404(
|
||||
server_id, request=request
|
||||
)
|
||||
# Use the server's stored client_id when the caller doesn't supply one
|
||||
resolved_client_id = mcp_server.client_id or client_id or ""
|
||||
if not resolved_client_id:
|
||||
|
|
@ -1399,10 +1520,12 @@ if MCP_AVAILABLE:
|
|||
@router.post(
|
||||
"/server/oauth/{server_id}/token",
|
||||
include_in_schema=False,
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def mcp_token(
|
||||
request: Request,
|
||||
server_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
grant_type: str = Form(...),
|
||||
code: Optional[str] = Form(None),
|
||||
redirect_uri: Optional[str] = Form(None),
|
||||
|
|
@ -1412,7 +1535,9 @@ if MCP_AVAILABLE:
|
|||
refresh_token: Optional[str] = Form(None),
|
||||
scope: Optional[str] = Form(None),
|
||||
):
|
||||
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
|
||||
mcp_server = await _get_cached_temporary_mcp_server_or_404(
|
||||
server_id, request=request
|
||||
)
|
||||
resolved_client_id = mcp_server.client_id or client_id or ""
|
||||
if not resolved_client_id:
|
||||
raise HTTPException(
|
||||
|
|
@ -1441,9 +1566,16 @@ if MCP_AVAILABLE:
|
|||
@router.post(
|
||||
"/server/oauth/{server_id}/register",
|
||||
include_in_schema=False,
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def mcp_register(request: Request, server_id: str):
|
||||
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
|
||||
async def mcp_register(
|
||||
request: Request,
|
||||
server_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
mcp_server = await _get_cached_temporary_mcp_server_or_404(
|
||||
server_id, request=request
|
||||
)
|
||||
request_data = await _read_request_body(request=request)
|
||||
data: dict = {**request_data}
|
||||
|
||||
|
|
|
|||
|
|
@ -302,14 +302,15 @@ class TeamMemberBudgetHandler:
|
|||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
"""
|
||||
Create team_memberships entries for existing members that don't have one.
|
||||
Ensure every team member has a TeamMembership row linked to the
|
||||
team_member_budget.
|
||||
|
||||
Called after team_member_budget is set/updated on a team to ensure
|
||||
members who joined before the budget was configured also get budget
|
||||
enforcement.
|
||||
|
||||
Only creates missing entries — does not touch existing memberships
|
||||
(which may carry individual per-member budgets).
|
||||
Called after team_member_budget is set/updated on a team. Creates
|
||||
rows for members who don't have one, and populates budget_id on
|
||||
existing rows where it is NULL. Rows with a non-NULL budget_id
|
||||
are left untouched, which preserves per-member overrides but also
|
||||
means rows pointing to a prior team-default budget_id are not
|
||||
migrated to the new one.
|
||||
"""
|
||||
if not members_with_roles:
|
||||
return
|
||||
|
|
@ -347,6 +348,21 @@ class TeamMemberBudgetHandler:
|
|||
_sanitize_for_log(team_member_budget_id),
|
||||
)
|
||||
|
||||
# Heal existing membership rows that predate the team_member_budget
|
||||
# configuration: populate budget_id where it is currently NULL.
|
||||
# Rows with an explicit budget_id (per-member override) are left alone.
|
||||
updated = await prisma_client.db.litellm_teammembership.update_many(
|
||||
where={"team_id": team_id, "budget_id": None},
|
||||
data={"budget_id": team_member_budget_id},
|
||||
)
|
||||
if updated:
|
||||
verbose_proxy_logger.info(
|
||||
"Populated budget_id on %d existing team_memberships for team %s with budget %s",
|
||||
updated,
|
||||
_sanitize_for_log(team_id),
|
||||
_sanitize_for_log(team_member_budget_id),
|
||||
)
|
||||
|
||||
|
||||
def _get_default_team_param(field: str) -> Any:
|
||||
"""
|
||||
|
|
@ -2608,6 +2624,15 @@ async def team_member_update(
|
|||
identified_budget_id = tm.budget_id
|
||||
break
|
||||
|
||||
# If this membership still points at the team's shared default member
|
||||
# budget, _upsert_budget_and_membership will clone-on-write so that the
|
||||
# update only touches this user (not every member sharing the default).
|
||||
team_default_budget_id: Optional[str] = None
|
||||
if team_table.metadata is not None:
|
||||
raw_default_budget_id = team_table.metadata.get("team_member_budget_id")
|
||||
if isinstance(raw_default_budget_id, str):
|
||||
team_default_budget_id = raw_default_budget_id
|
||||
|
||||
### upsert new budget
|
||||
async with prisma_client.db.tx() as tx:
|
||||
await _upsert_budget_and_membership(
|
||||
|
|
@ -2620,6 +2645,7 @@ async def team_member_update(
|
|||
tpm_limit=data.tpm_limit,
|
||||
rpm_limit=data.rpm_limit,
|
||||
allowed_models=data.allowed_models,
|
||||
team_default_budget_id=team_default_budget_id,
|
||||
)
|
||||
|
||||
### update team member role
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from fastapi import HTTPException, Request
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy._types import ( # key request types; user request types; team request types; customer request types
|
||||
BudgetNewRequest,
|
||||
DeleteCustomerRequest,
|
||||
|
|
@ -140,6 +141,69 @@ async def handle_budget_for_entity(
|
|||
return existing_budget_id
|
||||
|
||||
|
||||
# Fields on LiteLLM_BudgetTable that represent the budget's *configuration*
|
||||
# (i.e. the values an admin sets). We copy these when cloning a team's
|
||||
# default member-budget into an individual member-budget so that the new
|
||||
# row starts with the same limits as the default.
|
||||
_CLONABLE_BUDGET_FIELDS: Tuple[str, ...] = (
|
||||
"max_budget",
|
||||
"soft_budget",
|
||||
"max_parallel_requests",
|
||||
"tpm_limit",
|
||||
"rpm_limit",
|
||||
"model_max_budget",
|
||||
"budget_duration",
|
||||
"allowed_models",
|
||||
)
|
||||
|
||||
|
||||
async def _clone_team_default_budget_for_member(
|
||||
prisma_client: PrismaClient,
|
||||
default_team_budget_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Create a new budget row that copies the values from the team's default
|
||||
member budget. Returns the new budget_id, or None if the default budget
|
||||
no longer exists in the DB.
|
||||
|
||||
Used when adding a new team member without an explicit per-member budget,
|
||||
so the member starts with the team default's values but gets their own
|
||||
private budget row (which can be edited independently).
|
||||
"""
|
||||
default_budget = await prisma_client.db.litellm_budgettable.find_unique(
|
||||
where={"budget_id": default_team_budget_id}
|
||||
)
|
||||
if default_budget is None:
|
||||
return None
|
||||
|
||||
default_budget_dict = default_budget.model_dump()
|
||||
cloned_data: dict = {
|
||||
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
}
|
||||
for field in _CLONABLE_BUDGET_FIELDS:
|
||||
value = default_budget_dict.get(field)
|
||||
if value is None:
|
||||
continue
|
||||
# Skip empty list defaults (e.g. allowed_models = []) so the cloned
|
||||
# row matches the "no value set" shape rather than carrying a default.
|
||||
if isinstance(value, list) and len(value) == 0:
|
||||
continue
|
||||
cloned_data[field] = value
|
||||
|
||||
# Start the member's budget window at clone time, not the pool's reset
|
||||
# timestamp — otherwise a member joining mid-cycle inherits a stale reset.
|
||||
if cloned_data.get("budget_duration"):
|
||||
cloned_data["budget_reset_at"] = get_budget_reset_time(
|
||||
cloned_data["budget_duration"]
|
||||
)
|
||||
|
||||
new_budget = await prisma_client.db.litellm_budgettable.create(data=cloned_data)
|
||||
return new_budget.budget_id
|
||||
|
||||
|
||||
async def add_new_member(
|
||||
new_member: Member,
|
||||
max_budget_in_team: Optional[float],
|
||||
|
|
@ -221,8 +285,20 @@ async def add_new_member(
|
|||
response = await prisma_client.db.litellm_budgettable.create(data=budget_data)
|
||||
|
||||
_budget_id = response.budget_id
|
||||
elif default_team_budget_id is not None:
|
||||
# No per-member budget was provided, but the team has a default member
|
||||
# budget. Clone the default budget into a new row for this user so that
|
||||
# later edits to one member's budget do not bleed into other members.
|
||||
# If the default no longer exists in the DB, fall back to no budget.
|
||||
_budget_id = await _clone_team_default_budget_for_member(
|
||||
prisma_client=prisma_client,
|
||||
default_team_budget_id=default_team_budget_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
else:
|
||||
_budget_id = default_team_budget_id
|
||||
# No per-member budget and no team default → member gets no budget.
|
||||
_budget_id = None
|
||||
|
||||
if _budget_id and returned_user is not None and returned_user.user_id is not None:
|
||||
_returned_team_membership = (
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
|
|||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import Any, Optional, Tuple, Union, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -1496,10 +1497,18 @@ class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
|
|||
|
||||
def get_vertex_base_url(vertex_location: Optional[str]) -> str:
|
||||
"""
|
||||
Returns the base URL for Vertex AI based on the provided location.
|
||||
Base URL for Vertex AI pass-through (trailing slash for URL joining).
|
||||
|
||||
Keep location rules aligned with ``litellm.llms.vertex_ai.common_utils.get_vertex_base_url``.
|
||||
"""
|
||||
if vertex_location == "global":
|
||||
return "https://aiplatform.googleapis.com/"
|
||||
if vertex_location is None:
|
||||
raise ValueError("vertex_location is required")
|
||||
if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
|
||||
raise ValueError("Invalid vertex_location format")
|
||||
if "-" not in vertex_location:
|
||||
return f"https://aiplatform.{vertex_location}.rep.googleapis.com/"
|
||||
return f"https://{vertex_location}-aiplatform.googleapis.com/"
|
||||
|
||||
|
||||
|
|
@ -1703,7 +1712,8 @@ async def _base_vertex_proxy_route(
|
|||
Base function for Vertex AI passthrough routes.
|
||||
Handles common logic for all Vertex AI services.
|
||||
|
||||
Default base_target_url is `https://{vertex_location}-aiplatform.googleapis.com/`
|
||||
Default base_target_url is derived from ``get_vertex_base_url`` in this module
|
||||
(regional, ``global``, or multi-region ``.rep.`` hosts), with a trailing slash.
|
||||
|
||||
Args:
|
||||
endpoint: The endpoint path
|
||||
|
|
@ -2275,11 +2285,7 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
return
|
||||
|
||||
host_location = resolved_location or vertex_llm_base.get_default_vertex_location()
|
||||
host = (
|
||||
"aiplatform.googleapis.com"
|
||||
if host_location == "global"
|
||||
else f"{host_location}-aiplatform.googleapis.com"
|
||||
)
|
||||
host = get_vertex_base_url(host_location).removeprefix("https://").rstrip("/")
|
||||
service_url = (
|
||||
f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
)
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue