mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_gpt54_mini_nano_versioned_models
This commit is contained in:
commit
55ea431c05
281 changed files with 28194 additions and 3153 deletions
1844
.circleci/config.yml
1844
.circleci/config.yml
File diff suppressed because it is too large
Load diff
BIN
.github/screenshots/after_org_assigned.png
vendored
Normal file
BIN
.github/screenshots/after_org_assigned.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 96 KiB |
BIN
.github/screenshots/after_org_detail.png
vendored
Normal file
BIN
.github/screenshots/after_org_detail.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 103 KiB |
BIN
.github/screenshots/before_403_error.png
vendored
Normal file
BIN
.github/screenshots/before_403_error.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 123 KiB |
BIN
.github/screenshots/before_no_org.png
vendored
Normal file
BIN
.github/screenshots/before_no_org.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 95 KiB |
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
|
||||
|
|
|
|||
|
|
@ -15,29 +15,21 @@ COPY --from=uvbin /uv /usr/local/bin/uv
|
|||
COPY --from=uvbin /uvx /usr/local/bin/uvx
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache \
|
||||
python3 \
|
||||
python3-dev \
|
||||
clang \
|
||||
llvm \
|
||||
lld \
|
||||
gcc \
|
||||
linux-headers \
|
||||
build-base \
|
||||
bash \
|
||||
coreutils \
|
||||
curl \
|
||||
openssl \
|
||||
openssl-dev \
|
||||
nodejs \
|
||||
npm \
|
||||
libsndfile && break || sleep 5; \
|
||||
apk add --no-cache \
|
||||
python3 \
|
||||
python3-dev \
|
||||
gcc \
|
||||
bash \
|
||||
coreutils \
|
||||
curl \
|
||||
openssl \
|
||||
libsndfile \
|
||||
nodejs && break || sleep 5; \
|
||||
done
|
||||
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
NVM_DIR=/root/.nvm \
|
||||
PATH="/root/.nvm/versions/node/v20.20.2/bin:/app/.venv/bin:${PATH}" \
|
||||
PATH="/app/.venv/bin:${PATH}" \
|
||||
LITELLM_NON_ROOT=true \
|
||||
PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
PRISMA_CLI_BINARY_TARGETS="debian-openssl-3.0.x" \
|
||||
|
|
@ -49,7 +41,8 @@ COPY enterprise/pyproject.toml enterprise/
|
|||
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
|
||||
|
||||
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
|
||||
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
||||
uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--extra extra_proxy \
|
||||
|
|
@ -62,38 +55,12 @@ COPY . .
|
|||
# Set non-root flag for build time consistency
|
||||
ENV LITELLM_NON_ROOT=true
|
||||
|
||||
# Build Admin UI once and stage the static output for the runtime image.
|
||||
# NOTE: .npmrc files (which may set ignore-scripts=true and min-release-age=3d)
|
||||
# are temporarily renamed during npm install/ci so they don't block lifecycle
|
||||
# scripts needed by the build. This is safe because npm ci installs from
|
||||
# package-lock.json with pinned versions + integrity hashes.
|
||||
# Stage the pre-built Admin UI from the checked-in Next.js static export.
|
||||
# _experimental/out/ is regenerated as part of the release runbook.
|
||||
# Restructure extensionless routes (foo.html -> foo/index.html) to match the layout
|
||||
# proxy_server.py expects, and drop a readiness marker.
|
||||
RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
|
||||
([ -f /app/.npmrc ] && mv /app/.npmrc /app/.npmrc.bak || true) && \
|
||||
NVM_VERSION="v0.40.4" && \
|
||||
NVM_CHECKSUM="4b7412c49960c7d31e8df72da90c1fb5b8cccb419ac99537b737028d497aba4f" && \
|
||||
NODE_VERSION="v20.20.2" && \
|
||||
NVM_SCRIPT="/tmp/install-nvm.sh" && \
|
||||
curl -fsSL "https://raw.githubusercontent.com/nvm-sh/nvm/${NVM_VERSION}/install.sh" -o "$NVM_SCRIPT" && \
|
||||
echo "${NVM_CHECKSUM} ${NVM_SCRIPT}" | sha256sum -c - && \
|
||||
bash "$NVM_SCRIPT" && \
|
||||
export NVM_DIR="$HOME/.nvm" && \
|
||||
. "$NVM_DIR/nvm.sh" && \
|
||||
nvm install "${NODE_VERSION}" && \
|
||||
nvm use "${NODE_VERSION}" && \
|
||||
npm install -g npm@11.12.1 && \
|
||||
npm install -g node-gyp@12.2.0 && \
|
||||
ln -sf "$(npm root -g)/node-gyp" "$(npm root -g)/npm/node_modules/node-gyp" && \
|
||||
npm cache clean --force && \
|
||||
cd /app/ui/litellm-dashboard && \
|
||||
if [ -f "/app/enterprise/enterprise_ui/enterprise_colors.json" ]; then \
|
||||
cp /app/enterprise/enterprise_ui/enterprise_colors.json ./ui_colors.json; \
|
||||
fi && \
|
||||
([ -f .npmrc ] && mv .npmrc .npmrc.bak || true) && \
|
||||
npm ci --no-audit --no-fund && \
|
||||
([ -f .npmrc.bak ] && mv .npmrc.bak .npmrc || true) && \
|
||||
([ -f /app/.npmrc.bak ] && mv /app/.npmrc.bak /app/.npmrc || true) && \
|
||||
npm run build && \
|
||||
cp -r /app/ui/litellm-dashboard/out/* /var/lib/litellm/ui/ && \
|
||||
cp -r /app/litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \
|
||||
cp /app/litellm/proxy/logo.jpg /var/lib/litellm/assets/logo.jpg && \
|
||||
( cd /var/lib/litellm/ui && \
|
||||
for html_file in *.html; do \
|
||||
|
|
@ -103,10 +70,10 @@ RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
|
|||
mv "$html_file" "$folder_name/index.html"; \
|
||||
fi; \
|
||||
done && \
|
||||
touch .litellm_ui_ready ) && \
|
||||
cd /app/ui/litellm-dashboard && rm -rf ./out
|
||||
touch .litellm_ui_ready )
|
||||
|
||||
RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
|
||||
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
||||
if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
|
||||
uv sync --frozen --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
|
|
@ -123,10 +90,7 @@ RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
|
|||
--python python3; \
|
||||
fi
|
||||
|
||||
RUN mkdir -p /app/.cache/npm && \
|
||||
prisma generate --schema=./schema.prisma && \
|
||||
prisma --version && \
|
||||
prisma migrate diff --from-empty --to-schema-datamodel ./schema.prisma --script > /dev/null 2>&1 || true
|
||||
RUN prisma generate --schema=./schema.prisma
|
||||
|
||||
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
|
||||
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
|
@ -137,33 +101,11 @@ WORKDIR /app
|
|||
USER root
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk upgrade --no-cache && break || sleep 5; \
|
||||
apk upgrade --no-cache && break || sleep 5; \
|
||||
done && \
|
||||
for i in 1 2 3; do \
|
||||
apk add --no-cache python3 bash openssl tzdata nodejs npm supervisor libsndfile && break || sleep 5; \
|
||||
done && \
|
||||
apk upgrade --no-cache nodejs && \
|
||||
npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \
|
||||
GLOBAL="$(npm root -g)" && \
|
||||
find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
|
||||
done && \
|
||||
find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \
|
||||
sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \
|
||||
npm cache clean --force && \
|
||||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
apk add --no-cache python3 bash openssl tzdata supervisor libsndfile nodejs && break || sleep 5; \
|
||||
done
|
||||
|
||||
COPY --from=builder /app /app
|
||||
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
|
||||
|
|
@ -179,15 +121,10 @@ ENV PATH="/app/.venv/bin:${PATH}" \
|
|||
PRISMA_SKIP_POSTINSTALL_GENERATE=1 \
|
||||
PRISMA_HIDE_UPDATE_MESSAGE=1 \
|
||||
PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \
|
||||
NPM_CONFIG_CACHE=/app/.cache/npm \
|
||||
NPM_CONFIG_PREFER_OFFLINE=true \
|
||||
PRISMA_OFFLINE_MODE=true
|
||||
|
||||
RUN sed -i 's/\r$//' docker/entrypoint.sh && \
|
||||
sed -i 's/\r$//' docker/prod_entrypoint.sh && \
|
||||
chmod +x docker/entrypoint.sh docker/prod_entrypoint.sh && \
|
||||
mkdir -p /nonexistent /.npm /var/lib/litellm/assets /var/lib/litellm/ui /tmp/.npm && \
|
||||
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent /.npm /tmp/.npm && \
|
||||
RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
|
||||
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent && \
|
||||
PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
|
||||
chown -R nobody:nogroup "$PRISMA_PATH" && \
|
||||
LITELLM_PKG_MIGRATIONS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))' 2>/dev/null || echo '')/migrations" && \
|
||||
|
|
@ -201,7 +138,7 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && \
|
|||
[ -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:
|
||||
|
|
|
|||
172
docs/my-website/blog/gemini_embedding_2_ga/index.md
Normal file
172
docs/my-website/blog/gemini_embedding_2_ga/index.md
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
---
|
||||
slug: gemini_embedding_2_ga
|
||||
title: "Gemini Embedding 2 (GA): Multimodal Embeddings on LiteLLM"
|
||||
date: 2026-04-24T10:00:00
|
||||
authors:
|
||||
- sameer
|
||||
description: "Use generally available gemini-embedding-2 for multimodal embeddings on LiteLLM via Gemini API and Vertex AI—the same flows as preview, stable model id."
|
||||
tags: [gemini, embeddings, multimodal, vertex ai]
|
||||
hide_table_of_contents: false
|
||||
---
|
||||
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Gemini Embedding 2 (GA): Multimodal Embeddings
|
||||
|
||||
Litellm now fully supports Gemini Embedding 2 GA.
|
||||
|
||||
:::info
|
||||
For end-to-end behavior, input shapes, and MIME types, see the [Gemini Embedding 2 Preview walkthrough](/blog/gemini_embedding_2_multimodal). This post focuses on **GA naming**, **cost map** coverage.
|
||||
:::
|
||||
|
||||
{/* truncate */}
|
||||
|
||||
## Supported Input Types
|
||||
|
||||
| Modality | Supported Formats |
|
||||
|----------|-------------------|
|
||||
| **Text** | Plain text |
|
||||
| **Image** | PNG, JPEG |
|
||||
| **Audio** | MP3, WAV |
|
||||
| **Video** | MP4, MOV |
|
||||
| **Documents** | PDF |
|
||||
|
||||
## Input Formats
|
||||
|
||||
LiteLLM accepts three input formats for multimodal content:
|
||||
|
||||
1. **Data URIs** – Base64-encoded inline: `data:image/png;base64,<encoded_data>`
|
||||
2. **GCS URLs** – Cloud Storage paths (Vertex AI): `gs://bucket/path/to/file.png`
|
||||
3. **Gemini File References** – Pre-uploaded files (Gemini API): `files/abc123`
|
||||
|
||||
## Quick Start
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="gemini" label="Gemini API">
|
||||
|
||||
```python
|
||||
from litellm import embedding
|
||||
import os
|
||||
|
||||
os.environ["GEMINI_API_KEY"] = "your-api-key"
|
||||
|
||||
# Text + Image (base64)
|
||||
response = embedding(
|
||||
model="gemini/gemini-embedding-2",
|
||||
input=[
|
||||
"The food was delicious and the waiter...",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
],
|
||||
)
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="vertex" label="Vertex AI">
|
||||
|
||||
```python
|
||||
import litellm
|
||||
from litellm import embedding
|
||||
|
||||
litellm.vertex_project = "your-project-id"
|
||||
litellm.vertex_location = "us-central1"
|
||||
|
||||
# Text + Image (GCS URL)
|
||||
response = embedding(
|
||||
model="vertex_ai/gemini-embedding-2",
|
||||
input=[
|
||||
"Describe this image",
|
||||
"gs://my-bucket/images/photo.png"
|
||||
],
|
||||
)
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="proxy" label="LiteLLM Proxy">
|
||||
|
||||
**1. Config (config.yaml)**
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gemini-embedding-2
|
||||
litellm_params:
|
||||
model: gemini/gemini-embedding-2
|
||||
api_key: os.environ/GEMINI_API_KEY
|
||||
- model_name: vertex-gemini-embedding-2
|
||||
litellm_params:
|
||||
model: vertex_ai/gemini-embedding-2
|
||||
vertex_project: os.environ/VERTEXAI_PROJECT
|
||||
vertex_location: global
|
||||
|
||||
general_settings:
|
||||
master_key: sk-1234
|
||||
```
|
||||
|
||||
**2. Start proxy**
|
||||
|
||||
```bash
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
**3. Call embeddings** (OpenAI-compatible **`POST /v1/embeddings`** on the proxy)
|
||||
|
||||
```bash
|
||||
curl -sS -X POST http://localhost:4000/v1/embeddings \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "gemini-embedding-2",
|
||||
"input": [
|
||||
"The food was delicious and the waiter...",
|
||||
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wSzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAAAASUVORK5CYII"
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Input Format Examples
|
||||
|
||||
| Format | Example | Provider |
|
||||
|--------|---------|----------|
|
||||
| **Data URI** | `data:image/png;base64,...` | Gemini, Vertex AI |
|
||||
| **GCS URL** | `gs://bucket/path/image.png` | Vertex AI |
|
||||
| **File reference** | `files/abc123` | Gemini API only |
|
||||
|
||||
### Supported MIME Types for Data URIs
|
||||
|
||||
- **Images:** `image/png`, `image/jpeg`
|
||||
- **Audio:** `audio/mpeg`, `audio/wav`
|
||||
- **Video:** `video/mp4`, `video/quicktime`
|
||||
- **Documents:** `application/pdf`
|
||||
|
||||
### GCS URL MIME Inference
|
||||
|
||||
For Vertex AI, MIME types are inferred from file extensions:
|
||||
|
||||
- `.png` → `image/png`
|
||||
- `.jpg` / `.jpeg` → `image/jpeg`
|
||||
- `.mp3` → `audio/mpeg`
|
||||
- `.wav` → `audio/wav`
|
||||
- `.mp4` → `video/mp4`
|
||||
- `.mov` → `video/quicktime`
|
||||
- `.pdf` → `application/pdf`
|
||||
|
||||
## Optional Parameters
|
||||
|
||||
| Parameter | Description | Maps to |
|
||||
|-----------|-------------|---------|
|
||||
| `dimensions` | Output embedding size | `outputDimensionality` |
|
||||
|
||||
```python
|
||||
response = embedding(
|
||||
model="gemini/gemini-embedding-2",
|
||||
input=["text to embed"],
|
||||
dimensions=768, # Optional: control output vector size
|
||||
)
|
||||
```
|
||||
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
|
||||
|
|
@ -10,6 +10,7 @@ Supported Providers:
|
|||
- Vertex AI (`vertex_ai/`, `vertex_ai_beta/`)
|
||||
- Bedrock (`bedrock/`, `bedrock/invoke/`, `bedrock/converse`) ([All models bedrock supports prompt caching on](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html))
|
||||
- Deepseek API (`deepseek/`)
|
||||
- xAI (`xai/`)
|
||||
|
||||
For the supported providers, LiteLLM follows the OpenAI prompt caching usage object format:
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ The function keeps high-relevance and recent context, replaces low-relevance con
|
|||
|
||||
```python
|
||||
import litellm
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a coding assistant."},
|
||||
|
|
@ -19,6 +20,7 @@ messages = [
|
|||
compressed = litellm.compress(
|
||||
messages=messages,
|
||||
model="gpt-4o",
|
||||
call_type=CallTypes.completion,
|
||||
compression_trigger=1000,
|
||||
compression_target=500,
|
||||
)
|
||||
|
|
@ -45,6 +47,7 @@ response = litellm.completion(
|
|||
|
||||
- `messages` (`List[dict]`, required): input conversation messages
|
||||
- `model` (`str`, required): model name used for token counting
|
||||
- `call_type` (`CallTypes`, default `CallTypes.completion`): the LiteLLM call type whose message schema these messages follow. Supported values: `CallTypes.completion` / `CallTypes.acompletion` (OpenAI chat-completions shape) and `CallTypes.anthropic_messages` (Anthropic Messages shape)
|
||||
- `compression_trigger` (`int`, default `200000`): compress only if input token count exceeds this
|
||||
- `compression_target` (`Optional[int]`, default `70% of compression_trigger`): desired post-compression token budget
|
||||
- `embedding_model` (`Optional[str]`): if set, combines BM25 + embedding relevance scoring
|
||||
|
|
@ -70,6 +73,28 @@ args = json.loads(tool_call.function.arguments)
|
|||
full_content = compressed["cache"][args["key"]]
|
||||
```
|
||||
|
||||
## Server-side Callback Loop (`/v1/messages`)
|
||||
|
||||
You can enable callback-based compression interception to make retrieval loops
|
||||
transparent for Anthropic Messages calls:
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
callbacks: ["compression_interception"]
|
||||
compression_interception_params:
|
||||
enabled: true
|
||||
compression_trigger: 10000
|
||||
compression_target: 7000
|
||||
```
|
||||
|
||||
With this enabled, LiteLLM runs the following server-side flow:
|
||||
|
||||
1. Compresses inbound messages before the first provider call.
|
||||
2. Injects the `litellm_content_retrieve` tool.
|
||||
3. Detects retrieval `tool_use` blocks in the model response.
|
||||
4. Resolves retrieval keys from the compression cache.
|
||||
5. Reruns the model via agentic loop and returns the final answer.
|
||||
|
||||
## Performance
|
||||
|
||||
Benchmarked on [SWE-bench Lite](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_bm25_27K) (real GitHub issues with ~27k tokens of BM25-retrieved repo context per problem).
|
||||
|
|
|
|||
|
|
@ -60,3 +60,44 @@ curl http://localhost:4000/chat/completions \
|
|||
## Supported features
|
||||
|
||||
Scaleway provider supports all features in [Generative APIs reference documentation ↗](https://www.scaleway.com/en/developers/api/generative-apis/), such as streaming, structured outputs and tool calling.
|
||||
|
||||
## Audio transcription
|
||||
|
||||
Scaleway's `/audio/transcriptions` endpoint is OpenAI-compatible and works with Whisper models.
|
||||
|
||||
### Python SDK
|
||||
|
||||
```python
|
||||
import os
|
||||
from litellm import transcription
|
||||
|
||||
os.environ["SCW_SECRET_KEY"] = "your-scaleway-secret-key"
|
||||
|
||||
with open("speech.mp3", "rb") as audio_file:
|
||||
response = transcription(
|
||||
model="scaleway/whisper-large-v3",
|
||||
file=audio_file,
|
||||
)
|
||||
print(response.text)
|
||||
```
|
||||
|
||||
### Proxy config
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: scaleway-whisper
|
||||
litellm_params:
|
||||
model: scaleway/whisper-large-v3
|
||||
api_key: "os.environ/SCW_SECRET_KEY"
|
||||
```
|
||||
|
||||
### Proxy request
|
||||
|
||||
```bash
|
||||
curl http://localhost:4000/v1/audio/transcriptions \
|
||||
-H "Authorization: Bearer YOUR_LITELLM_MASTER_KEY" \
|
||||
-F model="scaleway-whisper" \
|
||||
-F file="@speech.mp3"
|
||||
```
|
||||
|
||||
Supported optional params: `language`, `prompt`, `response_format`, `temperature`, `timestamp_granularities`.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
:::
|
||||
|
|
|
|||
95
docs/my-website/docs/proxy/agentic_loop_hook.md
Normal file
95
docs/my-website/docs/proxy/agentic_loop_hook.md
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
# Agentic Loop Hook
|
||||
|
||||
Build a `CustomLogger` callback that intercepts a model response, fulfills tool calls server-side, and reruns the model — transparently to the caller.
|
||||
|
||||
:::info Supported call types
|
||||
- `async` only (sync calls do not trigger the hook)
|
||||
- Non-streaming only (streaming responses cannot be inspected for tool calls)
|
||||
- Works on both `/v1/messages` and `/v1/chat/completions`
|
||||
:::
|
||||
|
||||
## Implement the callback
|
||||
|
||||
Override two methods on `CustomLogger`:
|
||||
|
||||
```python
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
|
||||
|
||||
MY_TOOL = "my_tool"
|
||||
|
||||
class MyToolCallback(CustomLogger):
|
||||
|
||||
async def async_should_run_agentic_loop(
|
||||
self, response, model, messages, tools, stream, custom_llm_provider, kwargs
|
||||
):
|
||||
# Return (True, context_dict) if there are tool calls to handle
|
||||
content = getattr(response, "content", None) or []
|
||||
calls = [b for b in content if isinstance(b, dict)
|
||||
and b.get("type") == "tool_use" and b.get("name") == MY_TOOL]
|
||||
if not calls:
|
||||
return False, {}
|
||||
return True, {"tool_calls": calls}
|
||||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self, tools, model, messages, response,
|
||||
anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params,
|
||||
logging_obj, stream, kwargs,
|
||||
):
|
||||
calls = tools["tool_calls"]
|
||||
results = [f"result for {c['input']}" for c in calls] # your logic here
|
||||
|
||||
follow_up = messages + [
|
||||
{"role": "assistant", "content": [
|
||||
{"type": "tool_use", "id": c["id"], "name": c["name"], "input": c["input"]}
|
||||
for c in calls
|
||||
]},
|
||||
{"role": "user", "content": [
|
||||
{"type": "tool_result", "tool_use_id": c["id"], "content": results[i]}
|
||||
for i, c in enumerate(calls)
|
||||
]},
|
||||
]
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=AgenticLoopRequestPatch(messages=follow_up),
|
||||
)
|
||||
```
|
||||
|
||||
For `/v1/chat/completions`, override `async_build_chat_completion_agentic_loop_plan` instead — same idea, `optional_params` replaces `anthropic_messages_optional_request_params`.
|
||||
|
||||
## Register it
|
||||
|
||||
```python
|
||||
import litellm
|
||||
litellm.callbacks = [MyToolCallback()]
|
||||
```
|
||||
|
||||
Or in `config.yaml`:
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
callbacks: ["my_module.MyToolCallback"]
|
||||
```
|
||||
|
||||
## `AgenticLoopPlan` fields
|
||||
|
||||
| Field | Effect |
|
||||
|---|---|
|
||||
| `run_agentic_loop=True` + `request_patch` | Reruns the model with the patched request |
|
||||
| `response_override` | Returns this value directly to the caller (no rerun) |
|
||||
| `terminate=True` | Stops the loop, returns the current response |
|
||||
| `run_agentic_loop=False` (default) | Skips; next callback is checked |
|
||||
|
||||
`AgenticLoopRequestPatch` accepts: `model`, `messages`, `tools`, `max_tokens`, `optional_params`, `kwargs`.
|
||||
|
||||
## Loop safety
|
||||
|
||||
- Default max reruns: `3` — override per-request with `kwargs["max_agentic_loops"]`
|
||||
- Identical tool-call fingerprints abort the loop automatically
|
||||
- Current depth is in `kwargs["_agentic_loop_depth"]`
|
||||
|
||||
## Examples in this repo
|
||||
|
||||
- `litellm/integrations/compression_interception/handler.py`
|
||||
- `litellm/integrations/websearch_interception/handler.py`
|
||||
|
|
@ -1505,6 +1505,84 @@ curl http://localhost:4000/v1/responses \
|
|||
|
||||
|
||||
|
||||
### Opt-in bridge for `openai/` models with custom `api_base`
|
||||
|
||||
If you're using an **OpenAI-compatible third-party provider** (e.g. llama.cpp, vLLM, LM Studio) via `openai/` prefix with a custom `api_base`, LiteLLM will normally forward `/responses` requests directly to that endpoint. If the provider only supports `/chat/completions`, the request will fail.
|
||||
|
||||
Use either of these to force the `/responses` → `/chat/completions` bridge:
|
||||
|
||||
1. **`use_chat_completions_api: true`** — makes it explicit that LiteLLM will call the provider’s chat-completions API.
|
||||
2. **`openai/chat_completions/<model_name>`** — same pattern as `responses/` on chat completions: the model id encodes the routing choice.
|
||||
|
||||
#### Python SDK Usage
|
||||
|
||||
```python showLineNumbers title="Force bridge for custom openai/ endpoint (flag)"
|
||||
import litellm
|
||||
|
||||
response = litellm.responses(
|
||||
model="openai/my-custom-model",
|
||||
input="Hello!",
|
||||
api_base="http://localhost:8080",
|
||||
api_key="fake-key",
|
||||
use_chat_completions_api=True,
|
||||
)
|
||||
|
||||
print(response)
|
||||
```
|
||||
|
||||
Or encode it in the model id:
|
||||
|
||||
```python showLineNumbers title="Force bridge via openai/chat_completions/ model prefix"
|
||||
import litellm
|
||||
|
||||
response = litellm.responses(
|
||||
model="openai/chat_completions/my-custom-model",
|
||||
input="Hello!",
|
||||
api_base="http://localhost:8080",
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
print(response)
|
||||
```
|
||||
|
||||
#### LiteLLM Proxy Usage
|
||||
|
||||
**Setup Config:**
|
||||
|
||||
```yaml showLineNumbers title="config.yaml — bridge for custom openai/ endpoint"
|
||||
model_list:
|
||||
- model_name: my-local-model
|
||||
litellm_params:
|
||||
model: openai/my-custom-model
|
||||
api_base: http://localhost:8080/v1
|
||||
api_key: fake-key
|
||||
use_chat_completions_api: true
|
||||
```
|
||||
|
||||
Alternatively set `model: openai/chat_completions/my-custom-model` instead of the flag.
|
||||
|
||||
**Start Proxy:**
|
||||
|
||||
```bash showLineNumbers title="Start LiteLLM Proxy"
|
||||
litellm --config /path/to/config.yaml
|
||||
|
||||
# RUNNING on http://0.0.0.0:4000
|
||||
```
|
||||
|
||||
**Make Request:**
|
||||
|
||||
```bash showLineNumbers title="Request via bridge"
|
||||
curl http://localhost:4000/v1/responses \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "my-local-model",
|
||||
"input": "Hello!"
|
||||
}'
|
||||
```
|
||||
|
||||
This is particularly useful when connecting clients that hardcode the `/responses` endpoint (e.g. OpenAI Codex CLI with `wire_api = "responses"`) to local or third-party OpenAI-compatible providers that only expose `/chat/completions`.
|
||||
|
||||
## Server-side compaction
|
||||
|
||||
For long-running conversations, you can enable **server-side compaction** so that when the rendered context size crosses a threshold, the server automatically runs compaction in-stream and emits a compaction item—no separate `POST /v1/responses/compact` call is required.
|
||||
|
|
|
|||
|
|
@ -8,6 +8,22 @@ Reduce costs by up to 90% by using LiteLLM to auto-inject prompt caching checkpo
|
|||
|
||||
<Image img={require('../../img/auto_prompt_caching.png')} style={{ width: '800px', height: 'auto' }} />
|
||||
|
||||
Supported Providers (`cache_control` marker):
|
||||
- Anthropic API (`anthropic/`)
|
||||
- AWS Bedrock - Claude (`bedrock/`)
|
||||
- Vertex AI - Claude and Gemini (`vertex_ai/`)
|
||||
- Google AI Studio - Gemini (`gemini/`)
|
||||
- Azure AI - Claude (`azure_ai/`)
|
||||
- OpenRouter - Claude, Gemini, MiniMax, GLM, z-ai routes (`openrouter/`)
|
||||
- Databricks - Claude (`databricks/`)
|
||||
- DashScope / Qwen (`dashscope/`)
|
||||
- MiniMax (`minimax/`)
|
||||
- Z.ai / GLM (`zai/`)
|
||||
|
||||
Provider Managed (automatic, no marker needed):
|
||||
- OpenAI (`openai/`)
|
||||
- DeepSeek (`deepseek/`)
|
||||
- xAI (`xai/`)
|
||||
|
||||
## How it works
|
||||
|
||||
|
|
|
|||
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": {
|
||||
|
|
|
|||
|
|
@ -536,6 +536,7 @@ const sidebars = {
|
|||
description: "Modify requests, responses, and more",
|
||||
items: [
|
||||
"proxy/call_hooks",
|
||||
"proxy/agentic_loop_hook",
|
||||
"proxy/rules",
|
||||
]
|
||||
},
|
||||
|
|
@ -1059,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"
|
||||
|
|
@ -148,6 +148,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"vantage",
|
||||
"posthog",
|
||||
"levo",
|
||||
"compression_interception",
|
||||
]
|
||||
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
|
||||
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
|
||||
|
|
@ -1501,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": (
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
"""
|
||||
Main compress() function — orchestrates BM25/embedding scoring, message stubbing,
|
||||
and retrieval tool injection.
|
||||
Main compress() function — normalizes input messages, orchestrates BM25/embedding
|
||||
scoring, message stubbing, and retrieval tool injection.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Set, Union, cast
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.compression.message_stubbing import (
|
||||
|
|
@ -15,27 +15,196 @@ from litellm.compression.retrieval_tool import build_retrieval_tool
|
|||
from litellm.compression.scoring.bm25 import bm25_score_messages
|
||||
from litellm.litellm_core_utils.token_counter import token_counter
|
||||
from litellm.types.compression import CompressedResult
|
||||
from litellm.types.utils import AllMessageValues, Message
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
# CallTypes that produce Anthropic-shaped messages (structured content blocks).
|
||||
# Everything else is treated as OpenAI chat-completions shape.
|
||||
_ANTHROPIC_CALL_TYPES = frozenset({CallTypes.anthropic_messages.value})
|
||||
# CallTypes that are valid targets for compression. Compression operates on
|
||||
# message-shaped inputs, so we only accept call types whose payload is a list
|
||||
# of role/content messages.
|
||||
_SUPPORTED_CALL_TYPES = frozenset(
|
||||
{
|
||||
CallTypes.completion.value,
|
||||
CallTypes.acompletion.value,
|
||||
CallTypes.anthropic_messages.value,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _normalize_call_type(call_type: Union[CallTypes, str]) -> str:
|
||||
"""Return the string value for a ``CallTypes`` enum or a raw string."""
|
||||
if isinstance(call_type, CallTypes):
|
||||
return call_type.value
|
||||
return call_type
|
||||
|
||||
|
||||
def _is_anthropic_call_type(call_type: str) -> bool:
|
||||
return call_type in _ANTHROPIC_CALL_TYPES
|
||||
|
||||
|
||||
def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]:
|
||||
"""
|
||||
Build retrieval tool definitions in the target request schema.
|
||||
|
||||
- Chat-completions call types: keep OpenAI function-tool schema.
|
||||
- Anthropic messages call type: remap to Anthropic's custom tool schema.
|
||||
"""
|
||||
if not keys:
|
||||
return []
|
||||
|
||||
openai_tools = [build_retrieval_tool(keys)]
|
||||
if not _is_anthropic_call_type(call_type):
|
||||
return openai_tools
|
||||
|
||||
# Lazy import to avoid introducing provider transformation imports during
|
||||
# module import for non-Anthropic call paths.
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools)
|
||||
return cast(List[dict], anthropic_tools)
|
||||
|
||||
|
||||
def _content_to_text(content: Any) -> str:
|
||||
"""
|
||||
Convert OpenAI/Anthropic message content blocks to plain text.
|
||||
|
||||
Text extraction policy:
|
||||
- Include text-bearing fields only (`text` blocks + string values).
|
||||
- For `tool_result`, expand into nested `content` items.
|
||||
- Ignore non-textual blocks (images/documents/tool metadata/thinking metadata).
|
||||
|
||||
Implemented iteratively (stack-based) to avoid unbounded recursion.
|
||||
"""
|
||||
parts: List[str] = []
|
||||
stack: List[Any] = [content]
|
||||
while stack:
|
||||
item = stack.pop()
|
||||
if isinstance(item, str):
|
||||
parts.append(item)
|
||||
elif isinstance(item, list):
|
||||
# Push list items in reverse order so they are processed left-to-right.
|
||||
for element in reversed(item):
|
||||
stack.append(element)
|
||||
elif isinstance(item, dict):
|
||||
item_type = item.get("type")
|
||||
if item_type == "text":
|
||||
parts.append(str(item.get("text", "")))
|
||||
elif item_type == "tool_result":
|
||||
stack.append(item.get("content", ""))
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def _normalize_messages_for_compression(
|
||||
messages: List[dict],
|
||||
call_type: str,
|
||||
) -> Tuple[List[dict], List[dict]]:
|
||||
"""
|
||||
Normalize each original message to a text-surrogate content for scoring.
|
||||
|
||||
Returns:
|
||||
(normalized_messages, original_messages_copy)
|
||||
"""
|
||||
if call_type not in _SUPPORTED_CALL_TYPES:
|
||||
raise ValueError(
|
||||
f"Unsupported call_type={call_type!r} for compression. "
|
||||
f"Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}."
|
||||
)
|
||||
|
||||
original_messages: List[Dict[str, Any]] = [dict(m) for m in messages]
|
||||
|
||||
normalized_messages: List[dict] = []
|
||||
for msg in original_messages:
|
||||
normalized_messages.append(
|
||||
{
|
||||
**msg,
|
||||
"content": _content_to_text(msg.get("content", "")),
|
||||
}
|
||||
)
|
||||
return normalized_messages, original_messages
|
||||
|
||||
|
||||
def _extract_last_user_message(messages: List[dict]) -> str:
|
||||
"""Return the text content of the last user message."""
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") == "user":
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
parts.append(part.get("text", ""))
|
||||
elif isinstance(part, str):
|
||||
parts.append(part)
|
||||
return " ".join(parts)
|
||||
return _content_to_text(msg.get("content", ""))
|
||||
return ""
|
||||
|
||||
|
||||
def _extract_tool_use_ids(content: Any) -> List[str]:
|
||||
if not isinstance(content, list):
|
||||
return []
|
||||
tool_use_ids: List[str] = []
|
||||
for part in content:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
if part.get("type") != "tool_use":
|
||||
continue
|
||||
tool_use_id = part.get("id")
|
||||
if isinstance(tool_use_id, str) and tool_use_id:
|
||||
tool_use_ids.append(tool_use_id)
|
||||
return tool_use_ids
|
||||
|
||||
|
||||
def _extract_tool_result_ids(content: Any) -> Set[str]:
|
||||
if not isinstance(content, list):
|
||||
return set()
|
||||
tool_result_ids: Set[str] = set()
|
||||
for part in content:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
if part.get("type") != "tool_result":
|
||||
continue
|
||||
tool_use_id = part.get("tool_use_id")
|
||||
if isinstance(tool_use_id, str) and tool_use_id:
|
||||
tool_result_ids.add(tool_use_id)
|
||||
return tool_result_ids
|
||||
|
||||
|
||||
def _extract_anthropic_tool_exchange_spans(
|
||||
messages: List[dict],
|
||||
) -> Tuple[List[Set[int]], Optional[str]]:
|
||||
"""
|
||||
Return atomic 2-message spans for Anthropic tool exchanges.
|
||||
|
||||
Each assistant message containing `tool_use` must be immediately followed by a
|
||||
user message containing matching `tool_result` blocks for all tool_use ids.
|
||||
"""
|
||||
spans: List[Set[int]] = []
|
||||
i = 0
|
||||
while i < len(messages):
|
||||
current = messages[i]
|
||||
if current.get("role") != "assistant":
|
||||
i += 1
|
||||
continue
|
||||
|
||||
tool_use_ids = _extract_tool_use_ids(current.get("content"))
|
||||
if not tool_use_ids:
|
||||
i += 1
|
||||
continue
|
||||
|
||||
if i + 1 >= len(messages):
|
||||
return [], "invalid_anthropic_tool_sequence"
|
||||
|
||||
next_msg = messages[i + 1]
|
||||
if next_msg.get("role") != "user":
|
||||
return [], "invalid_anthropic_tool_sequence"
|
||||
|
||||
tool_result_ids = _extract_tool_result_ids(next_msg.get("content"))
|
||||
if not tool_result_ids:
|
||||
return [], "invalid_anthropic_tool_sequence"
|
||||
|
||||
for tool_use_id in tool_use_ids:
|
||||
if tool_use_id not in tool_result_ids:
|
||||
return [], "invalid_anthropic_tool_sequence"
|
||||
|
||||
spans.append({i, i + 1})
|
||||
i += 2
|
||||
|
||||
return spans, None
|
||||
|
||||
|
||||
def _get_protected_indices(messages: List[dict]) -> List[int]:
|
||||
"""
|
||||
Return indices of messages that must never be compressed:
|
||||
|
|
@ -87,9 +256,98 @@ def _combine_scores(
|
|||
return [bm25_weight * b + emb_weight * e for b, e in zip(norm_bm25, norm_emb)]
|
||||
|
||||
|
||||
def _select_kept_indices_for_budget(
|
||||
normalized_messages: List[dict],
|
||||
original_messages: List[dict],
|
||||
combined_scores: List[float],
|
||||
compression_target: int,
|
||||
model: str,
|
||||
initial_kept_indices: Set[int],
|
||||
tool_exchange_spans: List[Set[int]],
|
||||
) -> Tuple[Set[int], Dict[int, dict]]:
|
||||
kept_indices = set(initial_kept_indices)
|
||||
current_tokens = 0
|
||||
for i in kept_indices:
|
||||
current_tokens += token_counter(
|
||||
model=model,
|
||||
text=cast(str, normalized_messages[i].get("content", "") or ""),
|
||||
)
|
||||
|
||||
# Fill token budget from highest-scoring units.
|
||||
# A unit is either:
|
||||
# 1) a single message index, or
|
||||
# 2) an Anthropic tool-exchange span that must be kept/dropped atomically.
|
||||
truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict
|
||||
span_id_by_index: Dict[int, int] = {}
|
||||
for span_id, span in enumerate(tool_exchange_spans):
|
||||
for idx in span:
|
||||
span_id_by_index[idx] = span_id
|
||||
|
||||
# Build single-message candidate units (non-span messages).
|
||||
candidate_units: List[Tuple[float, Tuple[int, ...], bool]] = []
|
||||
for idx in range(len(normalized_messages)):
|
||||
if idx in span_id_by_index or idx in kept_indices:
|
||||
continue
|
||||
candidate_units.append((combined_scores[idx], (idx,), True))
|
||||
|
||||
# Build span candidate units (atomic keep/drop for tool exchanges).
|
||||
for span in tool_exchange_spans:
|
||||
span_indices = tuple(sorted(span))
|
||||
if any(idx in kept_indices for idx in span_indices):
|
||||
continue
|
||||
span_score = max(combined_scores[idx] for idx in span_indices)
|
||||
candidate_units.append((span_score, span_indices, False))
|
||||
|
||||
# Sort by descending relevance score.
|
||||
candidate_units.sort(key=lambda item: item[0], reverse=True)
|
||||
|
||||
for _score, indices, can_truncate in candidate_units:
|
||||
if any(idx in kept_indices for idx in indices):
|
||||
continue
|
||||
msg_tokens = 0
|
||||
for idx in indices:
|
||||
msg_tokens += token_counter(
|
||||
model=model,
|
||||
text=cast(str, normalized_messages[idx].get("content", "") or ""),
|
||||
)
|
||||
remaining = compression_target - current_tokens
|
||||
|
||||
if remaining <= 0:
|
||||
break # budget exhausted
|
||||
|
||||
if current_tokens + msg_tokens <= compression_target:
|
||||
# Fits entirely
|
||||
kept_indices.update(indices)
|
||||
current_tokens += msg_tokens
|
||||
elif can_truncate and len(indices) == 1 and remaining >= 100:
|
||||
# Too large to fit whole single message, but we have budget — truncate it.
|
||||
idx = indices[0]
|
||||
truncated = truncate_message(original_messages[idx], remaining)
|
||||
truncated_tokens = token_counter(
|
||||
model=model,
|
||||
text=truncated.get("content", "") or "",
|
||||
)
|
||||
truncated_overrides[idx] = truncated
|
||||
kept_indices.add(idx)
|
||||
current_tokens += truncated_tokens
|
||||
|
||||
return kept_indices, truncated_overrides
|
||||
|
||||
|
||||
def _get_dropped_tool_span_indices(
|
||||
kept_indices: Set[int], tool_exchange_spans: List[Set[int]]
|
||||
) -> Set[int]:
|
||||
dropped_tool_span_indices: Set[int] = set()
|
||||
for span in tool_exchange_spans:
|
||||
if not any(idx in kept_indices for idx in span):
|
||||
dropped_tool_span_indices.update(span)
|
||||
return dropped_tool_span_indices
|
||||
|
||||
|
||||
def compress(
|
||||
messages: List[dict],
|
||||
model: str,
|
||||
call_type: Union[CallTypes, str] = CallTypes.completion,
|
||||
compression_trigger: int = 200_000,
|
||||
compression_target: Optional[int] = None,
|
||||
embedding_model: Optional[str] = None,
|
||||
|
|
@ -108,6 +366,12 @@ def compress(
|
|||
Parameters:
|
||||
messages: The conversation messages to (potentially) compress.
|
||||
model: The LLM model name — used for token counting.
|
||||
call_type: The LiteLLM call type whose message schema these messages
|
||||
follow. Supported values:
|
||||
- ``CallTypes.completion`` / ``CallTypes.acompletion`` — OpenAI
|
||||
chat-completions shape (default)
|
||||
- ``CallTypes.anthropic_messages`` — Anthropic Messages shape
|
||||
(structured content blocks + atomic tool exchanges)
|
||||
compression_trigger: Only compress if input exceeds this token count.
|
||||
compression_target: Target token count after compression.
|
||||
Defaults to ``compression_trigger // 2``.
|
||||
|
|
@ -122,29 +386,37 @@ def compress(
|
|||
A ``CompressedResult`` dict containing compressed messages, token
|
||||
counts, a cache of original content, and the retrieval tool definition.
|
||||
"""
|
||||
call_type_str = _normalize_call_type(call_type)
|
||||
normalized_messages, original_messages = _normalize_messages_for_compression(
|
||||
messages=messages,
|
||||
call_type=call_type_str,
|
||||
)
|
||||
|
||||
if compression_target is None:
|
||||
compression_target = compression_trigger * 7 // 10
|
||||
|
||||
original_tokens = token_counter(
|
||||
model=model, messages=cast(List[Union[AllMessageValues, Message]], messages)
|
||||
model=model,
|
||||
messages=cast(List[Any], original_messages),
|
||||
)
|
||||
|
||||
# Pass through if below trigger
|
||||
if original_tokens <= compression_trigger:
|
||||
return CompressedResult(
|
||||
messages=messages,
|
||||
messages=original_messages,
|
||||
original_tokens=original_tokens,
|
||||
compressed_tokens=original_tokens,
|
||||
compression_ratio=0.0,
|
||||
cache={},
|
||||
tools=[],
|
||||
compression_skipped_reason="below_trigger",
|
||||
)
|
||||
|
||||
# Extract query for relevance scoring
|
||||
query = _extract_last_user_message(messages)
|
||||
query = _extract_last_user_message(normalized_messages)
|
||||
|
||||
# Score each message
|
||||
bm25_scores = bm25_score_messages(query, messages)
|
||||
bm25_scores = bm25_score_messages(query, normalized_messages)
|
||||
|
||||
if embedding_model:
|
||||
from litellm.compression.scoring.embedding_scorer import (
|
||||
|
|
@ -153,7 +425,7 @@ def compress(
|
|||
|
||||
emb_scores = embedding_score_messages(
|
||||
query,
|
||||
messages,
|
||||
normalized_messages,
|
||||
model=embedding_model,
|
||||
cache=compression_cache,
|
||||
embedding_model_params=embedding_model_params,
|
||||
|
|
@ -162,85 +434,69 @@ def compress(
|
|||
else:
|
||||
combined_scores = bm25_scores
|
||||
|
||||
# Sort message indices by score descending
|
||||
ranked_indices = sorted(
|
||||
range(len(messages)),
|
||||
key=lambda i: combined_scores[i],
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
# Protected messages are never compressed
|
||||
protected_indices = _get_protected_indices(messages)
|
||||
protected_indices = _get_protected_indices(normalized_messages)
|
||||
kept_indices: Set[int] = set(protected_indices)
|
||||
|
||||
# Count tokens for protected messages
|
||||
current_tokens = 0
|
||||
for i in kept_indices:
|
||||
current_tokens += token_counter(
|
||||
model=model, text=messages[i].get("content", "") or ""
|
||||
tool_exchange_spans: List[Set[int]] = []
|
||||
if _is_anthropic_call_type(call_type_str):
|
||||
tool_exchange_spans, tool_sequence_error = (
|
||||
_extract_anthropic_tool_exchange_spans(original_messages)
|
||||
)
|
||||
|
||||
# Fill token budget from highest-scoring messages.
|
||||
# For each candidate (ranked by relevance):
|
||||
# - If it fits entirely → keep it as-is.
|
||||
# - If it doesn't fit but there's meaningful remaining budget → truncate it
|
||||
# to fill as much of the budget as possible.
|
||||
# - Otherwise → stub it (pointer only, content goes to cache).
|
||||
# Multiple messages may be truncated so we preserve partial content from
|
||||
# several high-scoring messages rather than fully stubbing all but one.
|
||||
truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict
|
||||
|
||||
for idx in ranked_indices:
|
||||
if idx in kept_indices:
|
||||
continue
|
||||
msg_content = messages[idx].get("content", "") or ""
|
||||
msg_tokens = token_counter(model=model, text=msg_content)
|
||||
remaining = compression_target - current_tokens
|
||||
|
||||
if remaining <= 0:
|
||||
break # budget exhausted
|
||||
|
||||
if current_tokens + msg_tokens <= compression_target:
|
||||
# Fits entirely
|
||||
kept_indices.add(idx)
|
||||
current_tokens += msg_tokens
|
||||
elif remaining >= 100:
|
||||
# Too large to fit whole, but we have budget — truncate it.
|
||||
truncated = truncate_message(messages[idx], remaining)
|
||||
truncated_tokens = token_counter(
|
||||
model=model,
|
||||
text=truncated.get("content", "") or "",
|
||||
if tool_sequence_error is not None:
|
||||
return CompressedResult(
|
||||
messages=original_messages,
|
||||
original_tokens=original_tokens,
|
||||
compressed_tokens=original_tokens,
|
||||
compression_ratio=0.0,
|
||||
cache={},
|
||||
tools=[],
|
||||
compression_skipped_reason=tool_sequence_error,
|
||||
)
|
||||
truncated_overrides[idx] = truncated
|
||||
kept_indices.add(idx)
|
||||
current_tokens += truncated_tokens
|
||||
|
||||
for span in tool_exchange_spans:
|
||||
# If any message in the span is protected, keep the whole span.
|
||||
if any(idx in kept_indices for idx in span):
|
||||
kept_indices.update(span)
|
||||
|
||||
kept_indices, truncated_overrides = _select_kept_indices_for_budget(
|
||||
normalized_messages=normalized_messages,
|
||||
original_messages=original_messages,
|
||||
combined_scores=combined_scores,
|
||||
compression_target=compression_target,
|
||||
model=model,
|
||||
initial_kept_indices=kept_indices,
|
||||
tool_exchange_spans=tool_exchange_spans,
|
||||
)
|
||||
|
||||
# Build compressed messages and cache
|
||||
compressed_messages: List[dict] = []
|
||||
cache: Dict[str, str] = {}
|
||||
used_keys: Set[str] = set()
|
||||
dropped_tool_span_indices = _get_dropped_tool_span_indices(
|
||||
kept_indices=kept_indices, tool_exchange_spans=tool_exchange_spans
|
||||
)
|
||||
|
||||
for i, msg in enumerate(messages):
|
||||
for i, msg in enumerate(original_messages):
|
||||
if i in dropped_tool_span_indices:
|
||||
continue
|
||||
if i in kept_indices:
|
||||
# Use the truncated version if we made one, otherwise the original
|
||||
compressed_messages.append(truncated_overrides.get(i, msg))
|
||||
else:
|
||||
key = extract_key(msg, fallback_index=i, used_keys=used_keys)
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, list):
|
||||
content = " ".join(
|
||||
p.get("text", "") if isinstance(p, dict) else str(p)
|
||||
for p in content
|
||||
)
|
||||
key = extract_key(
|
||||
normalized_messages[i], fallback_index=i, used_keys=used_keys
|
||||
)
|
||||
content = _content_to_text(msg.get("content", ""))
|
||||
cache[key] = content
|
||||
compressed_messages.append(stub_message(msg, key))
|
||||
|
||||
# Build retrieval tool
|
||||
tools = [build_retrieval_tool(list(cache.keys()))] if cache else []
|
||||
# Build retrieval tool in the target request schema
|
||||
tools = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str)
|
||||
|
||||
compressed_tokens = token_counter(
|
||||
model=model,
|
||||
messages=cast(List[Union[AllMessageValues, Message]], compressed_messages),
|
||||
messages=cast(List[Any], compressed_messages),
|
||||
)
|
||||
|
||||
return CompressedResult(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
14
litellm/integrations/compression_interception/__init__.py
Normal file
14
litellm/integrations/compression_interception/__init__.py
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
"""
|
||||
Compression Interception Module
|
||||
|
||||
Provides server-side prompt compression + retrieval tool fulfillment for
|
||||
Anthropic Messages agentic loops.
|
||||
"""
|
||||
|
||||
from litellm.integrations.compression_interception.handler import (
|
||||
CompressionInterceptionLogger,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CompressionInterceptionLogger",
|
||||
]
|
||||
399
litellm/integrations/compression_interception/handler.py
Normal file
399
litellm/integrations/compression_interception/handler.py
Normal file
|
|
@ -0,0 +1,399 @@
|
|||
"""
|
||||
Compression Interception Handler
|
||||
|
||||
CustomLogger that compresses inbound Anthropic Messages requests and fulfills
|
||||
litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
|
||||
"""
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Tuple, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.compression import compress
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.integrations.compression_interception import (
|
||||
CompressionInterceptionConfig,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
)
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
LITELLM_CONTENT_RETRIEVE_TOOL_NAME = "litellm_content_retrieve"
|
||||
_CACHE_TTL_SECONDS = 15 * 60
|
||||
|
||||
|
||||
class CompressionInterceptionLogger(CustomLogger):
|
||||
"""
|
||||
CustomLogger that implements transparent prompt compression + retrieval loops.
|
||||
|
||||
Flow:
|
||||
1. Compress inbound /v1/messages requests in pre-call hook.
|
||||
2. Inject litellm_content_retrieve tool and persist compressed cache by call_id.
|
||||
3. Detect retrieval tool_use blocks in first model response.
|
||||
4. Build typed rerun plan with tool_result blocks from the compressed cache.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool = True,
|
||||
compression_trigger: int = 200_000,
|
||||
compression_target: Optional[int] = None,
|
||||
embedding_model: Optional[str] = None,
|
||||
embedding_model_params: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.enabled = enabled
|
||||
self.compression_trigger = compression_trigger
|
||||
self.compression_target = compression_target
|
||||
self.embedding_model = embedding_model
|
||||
self.embedding_model_params = embedding_model_params
|
||||
self._compression_cache_by_call_id: Dict[str, Tuple[Dict[str, str], float]] = {}
|
||||
|
||||
@classmethod
|
||||
def from_config_yaml(
|
||||
cls, config: CompressionInterceptionConfig
|
||||
) -> "CompressionInterceptionLogger":
|
||||
return cls(
|
||||
enabled=bool(config.get("enabled", True)),
|
||||
compression_trigger=int(config.get("compression_trigger", 200_000)),
|
||||
compression_target=config.get("compression_target"),
|
||||
embedding_model=config.get("embedding_model"),
|
||||
embedding_model_params=config.get("embedding_model_params"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def initialize_from_proxy_config(
|
||||
litellm_settings: Dict[str, Any],
|
||||
callback_specific_params: Dict[str, Any],
|
||||
) -> "CompressionInterceptionLogger":
|
||||
compression_params: CompressionInterceptionConfig = {}
|
||||
if "compression_interception_params" in litellm_settings:
|
||||
compression_params = litellm_settings["compression_interception_params"]
|
||||
elif "compression_interception" in callback_specific_params:
|
||||
compression_params = callback_specific_params["compression_interception"]
|
||||
return CompressionInterceptionLogger.from_config_yaml(compression_params)
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
|
||||
) -> Optional[dict]:
|
||||
if not self.enabled:
|
||||
return None
|
||||
if call_type is not None and call_type != CallTypes.anthropic_messages:
|
||||
return None
|
||||
if int(kwargs.get("_agentic_loop_depth", 0) or 0) > 0:
|
||||
return None
|
||||
|
||||
messages = kwargs.get("messages")
|
||||
model = kwargs.get("model")
|
||||
if not isinstance(messages, list) or not isinstance(model, str):
|
||||
return None
|
||||
|
||||
if self._has_retrieval_tool(kwargs.get("tools")):
|
||||
return None
|
||||
|
||||
self._prune_expired_cache()
|
||||
|
||||
compressed = compress( # type: ignore
|
||||
messages=messages,
|
||||
model=model,
|
||||
call_type=CallTypes.anthropic_messages,
|
||||
compression_trigger=self.compression_trigger,
|
||||
compression_target=self.compression_target,
|
||||
embedding_model=self.embedding_model,
|
||||
embedding_model_params=self.embedding_model_params,
|
||||
)
|
||||
|
||||
cache = cast(Dict[str, str], compressed.get("cache", {}))
|
||||
skip_reason = cast(Optional[str], compressed.get("compression_skipped_reason"))
|
||||
compressed_tools = cast(List[Dict[str, Any]], compressed.get("tools", []))
|
||||
|
||||
# Only mutate kwargs when compression actually produced a result.
|
||||
# If compression was a no-op (below trigger, invalid tool sequence, etc.),
|
||||
# leave ``messages`` and ``tools`` untouched — injecting an empty
|
||||
# ``tools: []`` onto a request that originally had no tools breaks
|
||||
# Anthropic Messages requests.
|
||||
if cache:
|
||||
kwargs["messages"] = compressed["messages"]
|
||||
if compressed_tools:
|
||||
kwargs["tools"] = self._merge_tools(
|
||||
existing_tools=cast(
|
||||
Optional[List[Dict[str, Any]]], kwargs.get("tools")
|
||||
),
|
||||
compressed_tools=compressed_tools,
|
||||
)
|
||||
call_id = cast(Optional[str], kwargs.get("litellm_call_id"))
|
||||
if not call_id:
|
||||
call_id = str(uuid.uuid4())
|
||||
kwargs["litellm_call_id"] = call_id
|
||||
self._compression_cache_by_call_id[call_id] = (cache, time.time())
|
||||
verbose_logger.debug(
|
||||
"CompressionInterception: compressed request [call_id=%s original=%d compressed=%d cached_keys=%d]",
|
||||
call_id,
|
||||
compressed.get("original_tokens"),
|
||||
compressed.get("compressed_tokens"),
|
||||
len(cache),
|
||||
)
|
||||
elif skip_reason is not None:
|
||||
verbose_logger.debug(
|
||||
"CompressionInterception: compression skipped [reason=%s original=%d compressed=%d]",
|
||||
skip_reason,
|
||||
compressed.get("original_tokens"),
|
||||
compressed.get("compressed_tokens"),
|
||||
)
|
||||
|
||||
return kwargs
|
||||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
tools: Optional[List[Dict]],
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
kwargs: Dict,
|
||||
) -> Tuple[bool, Dict]:
|
||||
if not self.enabled:
|
||||
return False, {}
|
||||
if not self._has_retrieval_tool(tools):
|
||||
return False, {}
|
||||
|
||||
tool_calls, thinking_blocks = self._extract_retrieval_tool_calls(
|
||||
response=response
|
||||
)
|
||||
if not tool_calls:
|
||||
return False, {}
|
||||
|
||||
return True, {
|
||||
"tool_calls": tool_calls,
|
||||
"thinking_blocks": thinking_blocks,
|
||||
"tool_type": "compression_retrieval",
|
||||
}
|
||||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self,
|
||||
tools: Dict,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: Dict,
|
||||
logging_obj: Any,
|
||||
stream: bool,
|
||||
kwargs: Dict,
|
||||
) -> AgenticLoopPlan:
|
||||
self._prune_expired_cache()
|
||||
tool_calls = cast(List[Dict[str, Any]], tools.get("tool_calls", []))
|
||||
thinking_blocks = cast(List[Dict[str, Any]], tools.get("thinking_blocks", []))
|
||||
|
||||
call_id = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs)
|
||||
cache = self._get_cache(call_id=call_id)
|
||||
retrieval_results = [
|
||||
self._resolve_retrieval_content(tc, cache) for tc in tool_calls
|
||||
]
|
||||
|
||||
assistant_message = {
|
||||
"role": "assistant",
|
||||
"content": thinking_blocks
|
||||
+ [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": tc.get("id"),
|
||||
"name": tc.get("name", LITELLM_CONTENT_RETRIEVE_TOOL_NAME),
|
||||
"input": tc.get("input", {}),
|
||||
}
|
||||
for tc in tool_calls
|
||||
],
|
||||
}
|
||||
user_message = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": tool_calls[i].get("id"),
|
||||
"content": retrieval_results[i],
|
||||
}
|
||||
for i in range(len(tool_calls))
|
||||
],
|
||||
}
|
||||
follow_up_messages = messages + [assistant_message, user_message]
|
||||
|
||||
max_tokens = cast(
|
||||
Optional[int],
|
||||
anthropic_messages_optional_request_params.get("max_tokens")
|
||||
or kwargs.get("max_tokens"),
|
||||
)
|
||||
optional_params_without_max_tokens = {
|
||||
k: v
|
||||
for k, v in anthropic_messages_optional_request_params.items()
|
||||
if k != "max_tokens"
|
||||
}
|
||||
|
||||
full_model_name = model
|
||||
if logging_obj is not None:
|
||||
agentic_params = logging_obj.model_call_details.get(
|
||||
"agentic_loop_params", {}
|
||||
)
|
||||
full_model_name = cast(str, agentic_params.get("model", model))
|
||||
|
||||
request_patch = AgenticLoopRequestPatch(
|
||||
model=full_model_name,
|
||||
messages=follow_up_messages,
|
||||
max_tokens=max_tokens,
|
||||
optional_params=optional_params_without_max_tokens,
|
||||
kwargs=self._prepare_followup_kwargs(kwargs=kwargs),
|
||||
)
|
||||
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={"tool_type": "compression_retrieval", "call_id": call_id or ""},
|
||||
)
|
||||
|
||||
def _prune_expired_cache(self) -> None:
|
||||
now = time.time()
|
||||
self._compression_cache_by_call_id = {
|
||||
call_id: (cache, created_at)
|
||||
for call_id, (
|
||||
cache,
|
||||
created_at,
|
||||
) in self._compression_cache_by_call_id.items()
|
||||
if now - created_at <= _CACHE_TTL_SECONDS
|
||||
}
|
||||
|
||||
def _get_cache(self, call_id: Optional[str]) -> Dict[str, str]:
|
||||
if not call_id:
|
||||
return {}
|
||||
cache_entry = self._compression_cache_by_call_id.get(call_id)
|
||||
if cache_entry is None:
|
||||
return {}
|
||||
return cache_entry[0]
|
||||
|
||||
def _resolve_call_id(
|
||||
self, logging_obj: Any, kwargs: Dict[str, Any]
|
||||
) -> Optional[str]:
|
||||
if logging_obj is not None:
|
||||
logging_call_id = getattr(logging_obj, "litellm_call_id", None)
|
||||
if isinstance(logging_call_id, str) and logging_call_id:
|
||||
return logging_call_id
|
||||
kwargs_call_id = kwargs.get("litellm_call_id")
|
||||
return cast(
|
||||
Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None
|
||||
)
|
||||
|
||||
def _resolve_retrieval_content(
|
||||
self, tool_call: Dict[str, Any], cache: Dict[str, str]
|
||||
) -> str:
|
||||
raw_input = tool_call.get("input", {})
|
||||
key = ""
|
||||
if isinstance(raw_input, dict):
|
||||
key = str(raw_input.get("key", "") or "")
|
||||
if not key:
|
||||
return "No retrieval key provided."
|
||||
if key in cache:
|
||||
return cache[key]
|
||||
return f"[compressed content key '{key}' not found]"
|
||||
|
||||
def _extract_retrieval_tool_calls(
|
||||
self, response: Any
|
||||
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
|
||||
if isinstance(response, dict):
|
||||
content = response.get("content", [])
|
||||
else:
|
||||
content = getattr(response, "content", []) or []
|
||||
|
||||
if not isinstance(content, list):
|
||||
return [], []
|
||||
|
||||
tool_calls: List[Dict[str, Any]] = []
|
||||
thinking_blocks: List[Dict[str, Any]] = []
|
||||
|
||||
for block in content:
|
||||
if isinstance(block, dict):
|
||||
block_type = block.get("type")
|
||||
block_name = block.get("name")
|
||||
if block_type in ("thinking", "redacted_thinking"):
|
||||
thinking_blocks.append(block)
|
||||
if (
|
||||
block_type == "tool_use"
|
||||
and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
|
||||
):
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": block.get("id"),
|
||||
"type": "tool_use",
|
||||
"name": block_name,
|
||||
"input": block.get("input", {}),
|
||||
}
|
||||
)
|
||||
else:
|
||||
block_type = getattr(block, "type", None)
|
||||
block_name = getattr(block, "name", None)
|
||||
if block_type == "thinking":
|
||||
thinking_blocks.append(
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": getattr(block, "thinking", ""),
|
||||
"signature": getattr(block, "signature", ""),
|
||||
}
|
||||
)
|
||||
elif block_type == "redacted_thinking":
|
||||
thinking_blocks.append(
|
||||
{
|
||||
"type": "redacted_thinking",
|
||||
"data": getattr(block, "data", ""),
|
||||
}
|
||||
)
|
||||
if (
|
||||
block_type == "tool_use"
|
||||
and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
|
||||
):
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": getattr(block, "id", None),
|
||||
"type": "tool_use",
|
||||
"name": block_name,
|
||||
"input": getattr(block, "input", {}) or {},
|
||||
}
|
||||
)
|
||||
|
||||
return tool_calls, thinking_blocks
|
||||
|
||||
def _prepare_followup_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
internal_keys = {"litellm_logging_obj"}
|
||||
return {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if not k.startswith("_compression_interception") and k not in internal_keys
|
||||
}
|
||||
|
||||
def _has_retrieval_tool(self, tools: Any) -> bool:
|
||||
if not isinstance(tools, list):
|
||||
return False
|
||||
for tool in tools:
|
||||
if not isinstance(tool, dict):
|
||||
continue
|
||||
function = tool.get("function")
|
||||
if tool.get("type") == "function" and isinstance(function, dict):
|
||||
if function.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME:
|
||||
return True
|
||||
if (
|
||||
tool.get("type") == "custom"
|
||||
and tool.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _merge_tools(
|
||||
self,
|
||||
existing_tools: Optional[List[Dict[str, Any]]],
|
||||
compressed_tools: List[Dict[str, Any]],
|
||||
) -> List[Dict[str, Any]]:
|
||||
merged = list(existing_tools or [])
|
||||
if self._has_retrieval_tool(merged):
|
||||
return merged
|
||||
merged.extend(compressed_tools)
|
||||
return merged
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
|
|||
from litellm.types.integrations.argilla import ArgillaItem
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.integrations.custom_logger import AgenticLoopPlan
|
||||
from litellm.types.utils import (
|
||||
AdapterCompletionStreamWrapper,
|
||||
CallTypes,
|
||||
|
|
@ -239,7 +240,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
self,
|
||||
model: str,
|
||||
request_kwargs: Dict,
|
||||
messages: Optional[List[Dict[str, str]]] = None,
|
||||
messages: Optional[List[Dict[str, Any]]] = None,
|
||||
input: Optional[Union[str, List]] = None,
|
||||
specific_deployment: Optional[bool] = False,
|
||||
) -> Optional[PreRoutingHookResponse]:
|
||||
|
|
@ -676,6 +677,26 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
"""
|
||||
pass
|
||||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self,
|
||||
tools: Dict,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: Dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
stream: bool,
|
||||
kwargs: Dict,
|
||||
) -> AgenticLoopPlan:
|
||||
"""
|
||||
Build a typed rerun plan for Anthropic Messages agentic loops.
|
||||
|
||||
Override this method to separate callback decision/tool execution from
|
||||
follow-up request execution (handled by BaseLLMHTTPHandler).
|
||||
"""
|
||||
return AgenticLoopPlan(run_agentic_loop=False)
|
||||
|
||||
async def async_should_run_chat_completion_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
|
|
@ -707,6 +728,22 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
"""
|
||||
pass
|
||||
|
||||
async def async_build_chat_completion_agentic_loop_plan(
|
||||
self,
|
||||
tools: Dict,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
response: Any,
|
||||
optional_params: Dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
stream: bool,
|
||||
kwargs: Dict,
|
||||
) -> AgenticLoopPlan:
|
||||
"""
|
||||
Build a typed rerun plan for chat-completions agentic loops.
|
||||
"""
|
||||
return AgenticLoopPlan(run_agentic_loop=False)
|
||||
|
||||
# Useful helpers for custom logger classes
|
||||
|
||||
def truncate_standard_logging_payload_content(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ if TYPE_CHECKING:
|
|||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
|
||||
class PrometheusLogger(CustomLogger):
|
||||
# Class variables or attributes
|
||||
|
||||
|
|
@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger):
|
|||
amount: float = 1.0,
|
||||
) -> None:
|
||||
_labels = prometheus_label_factory(
|
||||
supported_enum_labels=self.get_labels_for_metric(
|
||||
metric_name=metric_name
|
||||
),
|
||||
supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name),
|
||||
enum_values=enum_values,
|
||||
label_context=label_context,
|
||||
)
|
||||
|
|
@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
user_api_key = hash_token(user_api_key)
|
||||
|
||||
label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request.
|
||||
label_context = PrometheusLabelFactoryContext(
|
||||
enum_values
|
||||
) # amortized per request.
|
||||
|
||||
# increment total LLM requests and spend metric
|
||||
self._increment_top_level_request_and_spend_metrics(
|
||||
|
|
@ -3490,7 +3491,9 @@ def _prometheus_labels_from_context(
|
|||
}
|
||||
|
||||
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
|
||||
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user()
|
||||
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = (
|
||||
ctx.get_resolved_end_user()
|
||||
)
|
||||
|
||||
for sk, val in ctx._custom_by_sanitized_key.items():
|
||||
if sk in supported_enum_labels:
|
||||
|
|
|
|||
|
|
@ -51,8 +51,7 @@ class PrometheusLabelFactoryContext:
|
|||
self.enum_values = enum_values
|
||||
enum_dict = enum_values.model_dump()
|
||||
self._sanitized_enum: Dict[str, Optional[str]] = {
|
||||
k: _sanitize_prometheus_label_value(v)
|
||||
for k, v in enum_dict.items()
|
||||
k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items()
|
||||
}
|
||||
self._custom_by_sanitized_key: Dict[str, Optional[str]] = {}
|
||||
if enum_values.custom_metadata_labels is not None:
|
||||
|
|
|
|||
|
|
@ -28,6 +28,10 @@ from litellm.integrations.websearch_interception.transformation import (
|
|||
from litellm.types.integrations.websearch_interception import (
|
||||
WebSearchInterceptionConfig,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
|
@ -573,6 +577,35 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self,
|
||||
tools: Dict,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: Dict,
|
||||
logging_obj: Any,
|
||||
stream: bool,
|
||||
kwargs: Dict,
|
||||
) -> AgenticLoopPlan:
|
||||
tool_calls = tools["tool_calls"]
|
||||
thinking_blocks = tools.get("thinking_blocks", [])
|
||||
request_patch = await self._build_anthropic_request_patch(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
thinking_blocks=thinking_blocks,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={"tool_type": "websearch", "response_format": "anthropic"},
|
||||
)
|
||||
|
||||
async def async_run_chat_completion_agentic_loop(
|
||||
self,
|
||||
tools: Dict,
|
||||
|
|
@ -608,6 +641,33 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
response_format=response_format,
|
||||
)
|
||||
|
||||
async def async_build_chat_completion_agentic_loop_plan(
|
||||
self,
|
||||
tools: Dict,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
response: Any,
|
||||
optional_params: Dict,
|
||||
logging_obj: Any,
|
||||
stream: bool,
|
||||
kwargs: Dict,
|
||||
) -> AgenticLoopPlan:
|
||||
tool_calls = tools["tool_calls"]
|
||||
response_format = tools.get("response_format", "openai")
|
||||
request_patch = await self._build_chat_completion_request_patch(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs,
|
||||
response_format=response_format,
|
||||
)
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={"tool_type": "websearch", "response_format": response_format},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_max_tokens(
|
||||
optional_params: Dict,
|
||||
|
|
@ -672,7 +732,48 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
stream: bool,
|
||||
kwargs: Dict,
|
||||
) -> Any:
|
||||
"""Execute litellm.search() and make follow-up request"""
|
||||
"""Legacy path: execute search + build patch + run follow-up call."""
|
||||
request_patch = await self._build_anthropic_request_patch(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
thinking_blocks=thinking_blocks,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
if request_patch.messages is None:
|
||||
raise ValueError("WebSearchInterception: missing follow-up messages")
|
||||
|
||||
optional_params = dict(anthropic_messages_optional_request_params)
|
||||
optional_params.update(request_patch.optional_params)
|
||||
max_tokens = request_patch.max_tokens
|
||||
if max_tokens is None:
|
||||
max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None))
|
||||
else:
|
||||
optional_params.pop("max_tokens", None)
|
||||
if max_tokens is None:
|
||||
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
|
||||
|
||||
return await anthropic_messages.acreate(
|
||||
max_tokens=max_tokens,
|
||||
messages=request_patch.messages,
|
||||
model=request_patch.model or model,
|
||||
**optional_params,
|
||||
**request_patch.kwargs,
|
||||
)
|
||||
|
||||
async def _build_anthropic_request_patch(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
tool_calls: List[Dict],
|
||||
thinking_blocks: List[Dict],
|
||||
anthropic_messages_optional_request_params: Dict,
|
||||
logging_obj: Any,
|
||||
kwargs: Dict,
|
||||
) -> AgenticLoopRequestPatch:
|
||||
"""Execute litellm.search() and build follow-up request patch."""
|
||||
|
||||
# Extract search queries from tool_use blocks
|
||||
search_tasks = []
|
||||
|
|
@ -721,20 +822,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
thinking_blocks=thinking_blocks,
|
||||
)
|
||||
|
||||
# Make follow-up request with search results
|
||||
# Type cast: user_message is a Dict for Anthropic format (default response_format)
|
||||
follow_up_messages = messages + [assistant_message, cast(Dict, user_message)]
|
||||
|
||||
verbose_logger.debug(
|
||||
"WebSearchInterception: Making follow-up request with search results"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Last message (tool_result): {user_message}"
|
||||
)
|
||||
|
||||
# Correlation context for structured logging
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get(
|
||||
"litellm_call_id", "unknown"
|
||||
|
|
@ -742,61 +831,41 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
full_model_name = model # safe default before try block
|
||||
|
||||
# Use anthropic_messages.acreate for follow-up request
|
||||
try:
|
||||
max_tokens = self._resolve_max_tokens(
|
||||
anthropic_messages_optional_request_params, kwargs
|
||||
)
|
||||
max_tokens = self._resolve_max_tokens(
|
||||
anthropic_messages_optional_request_params, kwargs
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request"
|
||||
)
|
||||
|
||||
# Create a copy of optional params without max_tokens (since we pass it explicitly)
|
||||
optional_params_without_max_tokens = {
|
||||
k: v
|
||||
for k, v in anthropic_messages_optional_request_params.items()
|
||||
if k != "max_tokens"
|
||||
}
|
||||
optional_params_without_max_tokens = {
|
||||
k: v
|
||||
for k, v in anthropic_messages_optional_request_params.items()
|
||||
if k != "max_tokens"
|
||||
}
|
||||
kwargs_for_followup = self._prepare_followup_kwargs(kwargs)
|
||||
|
||||
kwargs_for_followup = self._prepare_followup_kwargs(kwargs)
|
||||
|
||||
# Get model from logging_obj.model_call_details["agentic_loop_params"]
|
||||
# This preserves the full model name with provider prefix (e.g., "bedrock/invoke/...")
|
||||
if logging_obj is not None:
|
||||
agentic_params = logging_obj.model_call_details.get(
|
||||
"agentic_loop_params", {}
|
||||
)
|
||||
full_model_name = agentic_params.get("model", model)
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Using model name: {full_model_name}"
|
||||
if logging_obj is not None:
|
||||
agentic_params = logging_obj.model_call_details.get(
|
||||
"agentic_loop_params", {}
|
||||
)
|
||||
|
||||
final_response = await anthropic_messages.acreate(
|
||||
max_tokens=max_tokens,
|
||||
messages=follow_up_messages,
|
||||
model=full_model_name,
|
||||
**optional_params_without_max_tokens,
|
||||
**kwargs_for_followup,
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Final response: {final_response}"
|
||||
)
|
||||
return final_response
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"WebSearchInterception: Follow-up request failed "
|
||||
"[call_id=%s model=%s messages=%d searches=%d]: %s",
|
||||
_call_id,
|
||||
full_model_name,
|
||||
len(follow_up_messages),
|
||||
len(final_search_results),
|
||||
str(e),
|
||||
)
|
||||
raise
|
||||
full_model_name = agentic_params.get("model", model)
|
||||
verbose_logger.debug(
|
||||
"WebSearchInterception: Built anthropic request patch "
|
||||
"[call_id=%s model=%s messages=%d searches=%d]",
|
||||
_call_id,
|
||||
full_model_name,
|
||||
len(follow_up_messages),
|
||||
len(final_search_results),
|
||||
)
|
||||
return AgenticLoopRequestPatch(
|
||||
model=full_model_name,
|
||||
messages=follow_up_messages,
|
||||
max_tokens=max_tokens,
|
||||
optional_params=optional_params_without_max_tokens,
|
||||
kwargs=kwargs_for_followup,
|
||||
)
|
||||
|
||||
async def _execute_search(self, query: str) -> str:
|
||||
"""Execute a single web search using router's search tools"""
|
||||
|
|
@ -883,7 +952,36 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs: Dict,
|
||||
response_format: str = "openai",
|
||||
) -> Any:
|
||||
"""Execute litellm.search() and make follow-up chat completion request"""
|
||||
"""Legacy path: execute search + build patch + run follow-up call."""
|
||||
request_patch = await self._build_chat_completion_request_patch(
|
||||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs,
|
||||
response_format=response_format,
|
||||
)
|
||||
if request_patch.messages is None:
|
||||
raise ValueError("WebSearchInterception: missing follow-up messages")
|
||||
params = dict(optional_params)
|
||||
params.update(request_patch.optional_params)
|
||||
return await litellm.acompletion(
|
||||
model=request_patch.model or model,
|
||||
messages=request_patch.messages,
|
||||
**params,
|
||||
**request_patch.kwargs,
|
||||
)
|
||||
|
||||
async def _build_chat_completion_request_patch( # noqa: PLR0915
|
||||
self,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
tool_calls: List[Dict],
|
||||
optional_params: Dict,
|
||||
kwargs: Dict,
|
||||
response_format: str = "openai",
|
||||
) -> AgenticLoopRequestPatch:
|
||||
"""Execute litellm.search() and build chat-completion rerun patch."""
|
||||
|
||||
# Extract search queries from tool_calls
|
||||
search_tasks = []
|
||||
|
|
@ -963,74 +1061,56 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}"
|
||||
)
|
||||
|
||||
# Use litellm.acompletion for follow-up request
|
||||
try:
|
||||
# Remove internal parameters that shouldn't be passed to follow-up request
|
||||
internal_params = {
|
||||
"_websearch_interception",
|
||||
"acompletion",
|
||||
"litellm_logging_obj",
|
||||
"custom_llm_provider",
|
||||
# Remove internal parameters that shouldn't be passed to follow-up request
|
||||
internal_params = {
|
||||
"_websearch_interception",
|
||||
"acompletion",
|
||||
"litellm_logging_obj",
|
||||
"custom_llm_provider",
|
||||
"model_alias_map",
|
||||
"stream_response",
|
||||
"custom_prompt_dict",
|
||||
}
|
||||
kwargs_for_followup = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if not k.startswith("_websearch_interception") and k not in internal_params
|
||||
}
|
||||
|
||||
full_model_name = model
|
||||
if "custom_llm_provider" in kwargs:
|
||||
custom_llm_provider = kwargs["custom_llm_provider"]
|
||||
if not model.startswith(custom_llm_provider) and "/" not in model:
|
||||
full_model_name = f"{custom_llm_provider}/{model}"
|
||||
|
||||
verbose_logger.debug(
|
||||
"WebSearchInterception: Built chat completion request patch model=%s messages=%d",
|
||||
full_model_name,
|
||||
len(follow_up_messages),
|
||||
)
|
||||
|
||||
tools_param = optional_params.get("tools")
|
||||
optional_params_clean = {
|
||||
k: v
|
||||
for k, v in optional_params.items()
|
||||
if k
|
||||
not in {
|
||||
"tools",
|
||||
"extra_body",
|
||||
"model_alias_map",
|
||||
"stream_response",
|
||||
"custom_prompt_dict",
|
||||
}
|
||||
kwargs_for_followup = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if not k.startswith("_websearch_interception")
|
||||
and k not in internal_params
|
||||
}
|
||||
}
|
||||
if tools_param is not None:
|
||||
optional_params_clean["tools"] = tools_param
|
||||
|
||||
# Get full model name from kwargs
|
||||
full_model_name = model
|
||||
if "custom_llm_provider" in kwargs:
|
||||
custom_llm_provider = kwargs["custom_llm_provider"]
|
||||
# Reconstruct full model name with provider prefix if needed
|
||||
if not model.startswith(custom_llm_provider):
|
||||
# Check if model already has a provider prefix
|
||||
if "/" not in model:
|
||||
full_model_name = f"{custom_llm_provider}/{model}"
|
||||
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Using model name: {full_model_name}"
|
||||
)
|
||||
|
||||
# Prepare tools for follow-up request (same as original)
|
||||
tools_param = optional_params.get("tools")
|
||||
|
||||
# Remove tools and extra_body from optional_params to avoid issues
|
||||
# extra_body often contains internal LiteLLM params that shouldn't be forwarded
|
||||
optional_params_clean = {
|
||||
k: v
|
||||
for k, v in optional_params.items()
|
||||
if k
|
||||
not in {
|
||||
"tools",
|
||||
"extra_body",
|
||||
"model_alias_map",
|
||||
"stream_response",
|
||||
"custom_prompt_dict",
|
||||
}
|
||||
}
|
||||
|
||||
final_response = await litellm.acompletion(
|
||||
model=full_model_name,
|
||||
messages=follow_up_messages,
|
||||
tools=tools_param,
|
||||
**optional_params_clean,
|
||||
**kwargs_for_followup,
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}"
|
||||
)
|
||||
return final_response
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
f"WebSearchInterception: Follow-up request failed: {str(e)}"
|
||||
)
|
||||
raise
|
||||
return AgenticLoopRequestPatch(
|
||||
model=full_model_name,
|
||||
messages=follow_up_messages,
|
||||
optional_params=optional_params_clean,
|
||||
kwargs=kwargs_for_followup,
|
||||
)
|
||||
|
||||
async def _create_empty_search_result(self) -> str:
|
||||
"""Create an empty search result for tool calls without queries"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -296,6 +296,15 @@ def get_supported_openai_params( # noqa: PLR0915
|
|||
return OVHCloudAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
elif custom_llm_provider == "scaleway":
|
||||
if request_type == "transcription":
|
||||
from litellm.llms.scaleway.audio_transcription.transformation import (
|
||||
ScalewayAudioTranscriptionConfig,
|
||||
)
|
||||
|
||||
return ScalewayAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
elif custom_llm_provider == "elevenlabs":
|
||||
if request_type == "transcription":
|
||||
from litellm.llms.elevenlabs.audio_transcription.transformation import (
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -684,7 +684,7 @@ def generic_cost_per_token( # noqa: PLR0915
|
|||
- cache_creation
|
||||
- image_tokens
|
||||
)
|
||||
# Clamp to zero: inconsistent streaming usage
|
||||
# Clamp to zero: inconsistent streaming usage
|
||||
if text_tokens < 0:
|
||||
text_tokens = 0
|
||||
prompt_tokens_details["text_tokens"] = text_tokens
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.types.llms.anthropic import (
|
|||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolParam,
|
||||
)
|
||||
|
|
@ -67,6 +68,32 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
super().__init__()
|
||||
self.adapter = LiteLLMAnthropicMessagesAdapter()
|
||||
|
||||
def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
|
||||
"""Translate Anthropic request to OpenAI chat completion format."""
|
||||
(
|
||||
chat_completion_compatible_request,
|
||||
_tool_name_mapping,
|
||||
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
|
||||
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
|
||||
)
|
||||
return chat_completion_compatible_request
|
||||
|
||||
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
|
||||
"""
|
||||
Convert Anthropic messages request data to OpenAI-spec structured messages.
|
||||
|
||||
Uses the Anthropic-to-OpenAI adapter to translate message format.
|
||||
"""
|
||||
messages = data.get("messages")
|
||||
if messages is None:
|
||||
return None
|
||||
chat_completion_compatible_request = self._translate_to_openai(data)
|
||||
result = cast(
|
||||
List[AllMessageValues],
|
||||
chat_completion_compatible_request.get("messages", []),
|
||||
)
|
||||
return result if result else None
|
||||
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -82,13 +109,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply)
|
||||
|
||||
(
|
||||
chat_completion_compatible_request,
|
||||
_tool_name_mapping,
|
||||
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
|
||||
# Use a shallow copy to avoid mutating request data (pop on litellm_metadata).
|
||||
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
|
||||
)
|
||||
chat_completion_compatible_request = self._translate_to_openai(data)
|
||||
|
||||
structured_messages = cast(
|
||||
List[AllMessageValues],
|
||||
|
|
@ -103,8 +124,6 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
chat_completion_compatible_request.get("tools", [])
|
||||
)
|
||||
task_mappings: List[Tuple[int, Optional[int]]] = []
|
||||
# Track (message_index, content_index) for each text
|
||||
# content_index is None for string content, int for list content
|
||||
|
||||
# Step 1: Extract all text content and images
|
||||
for msg_idx, message in enumerate(messages):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,322 @@
|
|||
"""
|
||||
Agentic Streaming Iterator for Anthropic Messages
|
||||
|
||||
Wraps the raw SSE byte stream from the Anthropic pass-through endpoint,
|
||||
yields every chunk to the caller (preserving real streaming), collects
|
||||
all bytes, and on stream exhaustion rebuilds the full Anthropic response
|
||||
to run through agentic completion hooks. If an agentic hook fires, the
|
||||
follow-up response is chained as Phase 2 of the same iterator.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SSE parsing helpers (module-level to keep the class lean)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _parse_sse_events(raw: bytes) -> List[tuple]:
|
||||
"""Return a list of (event_type, parsed_data_dict) from raw SSE bytes."""
|
||||
text = raw.decode("utf-8", errors="replace")
|
||||
lines = text.split("\n")
|
||||
events: List[tuple] = []
|
||||
current_event_type: Optional[str] = None
|
||||
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if stripped.startswith("event:"):
|
||||
current_event_type = stripped[len("event:") :].strip()
|
||||
continue
|
||||
if not stripped.startswith("data:"):
|
||||
continue
|
||||
data_str = stripped[len("data:") :].strip()
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
event_type = current_event_type or data.get("type", "")
|
||||
current_event_type = None
|
||||
events.append((event_type, data))
|
||||
return events
|
||||
|
||||
|
||||
def _handle_message_start(data: Dict, response: Dict) -> None:
|
||||
msg = data.get("message", {})
|
||||
response["id"] = msg.get("id", response["id"])
|
||||
response["model"] = msg.get("model", response["model"])
|
||||
response["role"] = msg.get("role", response["role"])
|
||||
usage = msg.get("usage", {})
|
||||
if usage:
|
||||
response["usage"]["input_tokens"] = usage.get("input_tokens", 0)
|
||||
for key in ("cache_creation_input_tokens", "cache_read_input_tokens"):
|
||||
if key in usage:
|
||||
response["usage"][key] = usage[key]
|
||||
|
||||
|
||||
def _handle_content_block_start(data: Dict, content_blocks: Dict[int, Dict]) -> None:
|
||||
idx = data.get("index", len(content_blocks))
|
||||
block = data.get("content_block", {})
|
||||
block_type = block.get("type", "text")
|
||||
|
||||
_BLOCK_TEMPLATES: Dict[str, Dict] = {
|
||||
"text": {"type": "text", "text": ""},
|
||||
"thinking": {"type": "thinking", "thinking": "", "signature": ""},
|
||||
"redacted_thinking": {
|
||||
"type": "redacted_thinking",
|
||||
"data": block.get("data", ""),
|
||||
},
|
||||
}
|
||||
if block_type == "tool_use":
|
||||
content_blocks[idx] = {
|
||||
"type": "tool_use",
|
||||
"id": block.get("id", ""),
|
||||
"name": block.get("name", ""),
|
||||
"input": {},
|
||||
"_partial_json": "",
|
||||
}
|
||||
elif block_type in _BLOCK_TEMPLATES:
|
||||
content_blocks[idx] = dict(_BLOCK_TEMPLATES[block_type])
|
||||
else:
|
||||
content_blocks[idx] = dict(block)
|
||||
|
||||
|
||||
def _handle_content_block_delta(data: Dict, content_blocks: Dict[int, Dict]) -> None:
|
||||
idx = data.get("index", 0)
|
||||
delta = data.get("delta", {})
|
||||
delta_type = delta.get("type", "")
|
||||
block = content_blocks.get(idx)
|
||||
if block is None:
|
||||
return
|
||||
|
||||
if delta_type == "text_delta":
|
||||
block["text"] = block.get("text", "") + delta.get("text", "")
|
||||
elif delta_type == "input_json_delta":
|
||||
block["_partial_json"] = block.get("_partial_json", "") + delta.get(
|
||||
"partial_json", ""
|
||||
)
|
||||
elif delta_type == "thinking_delta":
|
||||
block["thinking"] = block.get("thinking", "") + delta.get("thinking", "")
|
||||
elif delta_type == "signature_delta":
|
||||
block["signature"] = delta.get("signature", block.get("signature", ""))
|
||||
|
||||
|
||||
def _handle_content_block_stop(data: Dict, content_blocks: Dict[int, Dict]) -> None:
|
||||
idx = data.get("index", 0)
|
||||
block = content_blocks.get(idx)
|
||||
if block and block.get("type") == "tool_use":
|
||||
partial = block.pop("_partial_json", "")
|
||||
if partial:
|
||||
try:
|
||||
block["input"] = json.loads(partial)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
block["input"] = {"_raw": partial}
|
||||
|
||||
|
||||
def _handle_message_delta(data: Dict, response: Dict) -> None:
|
||||
delta = data.get("delta", {})
|
||||
if "stop_reason" in delta:
|
||||
response["stop_reason"] = delta["stop_reason"]
|
||||
if "stop_sequence" in delta:
|
||||
response["stop_sequence"] = delta["stop_sequence"]
|
||||
usage = data.get("usage", {})
|
||||
if usage.get("output_tokens") is not None:
|
||||
response["usage"]["output_tokens"] = usage["output_tokens"]
|
||||
for key in (
|
||||
"input_tokens",
|
||||
"cache_creation_input_tokens",
|
||||
"cache_read_input_tokens",
|
||||
):
|
||||
if key in usage:
|
||||
response["usage"][key] = usage[key]
|
||||
|
||||
|
||||
class AgenticAnthropicStreamingIterator:
|
||||
"""
|
||||
Two-phase async iterator that enables agentic hooks on streaming
|
||||
Anthropic Messages pass-through responses.
|
||||
|
||||
Phase 1: Yield raw SSE bytes from the upstream response while
|
||||
accumulating them. When the inner iterator is exhausted,
|
||||
rebuild the full Anthropic response dict and call agentic hooks.
|
||||
|
||||
Phase 2: If an agentic hook fires and returns a follow-up response
|
||||
(streaming or non-streaming), yield those bytes to the caller.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
completion_stream: AsyncIterator,
|
||||
http_handler: Any,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: Dict,
|
||||
logging_obj: Any,
|
||||
custom_llm_provider: str,
|
||||
kwargs: Dict,
|
||||
):
|
||||
self._inner = completion_stream.__aiter__()
|
||||
self._http_handler = http_handler
|
||||
self._model = model
|
||||
self._messages = messages
|
||||
self._anthropic_messages_provider_config = anthropic_messages_provider_config
|
||||
self._anthropic_messages_optional_request_params = (
|
||||
anthropic_messages_optional_request_params
|
||||
)
|
||||
self._logging_obj = logging_obj
|
||||
self._custom_llm_provider = custom_llm_provider
|
||||
self._kwargs = kwargs
|
||||
|
||||
self._collected_bytes: List[bytes] = []
|
||||
self._stream_exhausted = False
|
||||
self._hook_processing_done = False
|
||||
self._follow_up_iterator: Optional[AsyncIterator] = None
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
# Phase 1: yield from upstream, collect bytes
|
||||
if not self._stream_exhausted:
|
||||
try:
|
||||
chunk = await self._inner.__anext__()
|
||||
self._collected_bytes.append(chunk)
|
||||
return chunk
|
||||
except StopAsyncIteration:
|
||||
self._stream_exhausted = True
|
||||
await self._process_agentic_hooks()
|
||||
# Fall through to Phase 2
|
||||
|
||||
# Phase 2: yield from follow-up stream if one was created
|
||||
if self._follow_up_iterator is not None:
|
||||
chunk = await self._follow_up_iterator.__anext__()
|
||||
return chunk
|
||||
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def _process_agentic_hooks(self) -> None:
|
||||
"""Rebuild the Anthropic response from collected SSE bytes and call hooks."""
|
||||
if self._hook_processing_done:
|
||||
return
|
||||
self._hook_processing_done = True
|
||||
|
||||
if not self._collected_bytes:
|
||||
return
|
||||
|
||||
try:
|
||||
rebuilt = self._rebuild_anthropic_response_from_sse(self._collected_bytes)
|
||||
if rebuilt is None:
|
||||
verbose_logger.debug(
|
||||
"AgenticStreamingIterator: Could not rebuild response from SSE bytes"
|
||||
)
|
||||
return
|
||||
|
||||
[
|
||||
(
|
||||
f"{b.get('type')}({b.get('name', '')})"
|
||||
if b.get("type") == "tool_use"
|
||||
else b.get("type")
|
||||
)
|
||||
for b in rebuilt.get("content", [])
|
||||
]
|
||||
|
||||
result = await self._http_handler._call_agentic_completion_hooks(
|
||||
response=rebuilt,
|
||||
model=self._model,
|
||||
messages=self._messages,
|
||||
anthropic_messages_provider_config=self._anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=self._anthropic_messages_optional_request_params,
|
||||
logging_obj=self._logging_obj,
|
||||
stream=True,
|
||||
custom_llm_provider=self._custom_llm_provider,
|
||||
kwargs=self._kwargs,
|
||||
)
|
||||
|
||||
if result is None:
|
||||
return
|
||||
|
||||
if hasattr(result, "__aiter__"):
|
||||
self._follow_up_iterator = result.__aiter__()
|
||||
elif isinstance(result, dict):
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
||||
fake = FakeAnthropicMessagesStreamIterator(
|
||||
response=cast(AnthropicMessagesResponse, result)
|
||||
)
|
||||
self._follow_up_iterator = fake.__aiter__()
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"AgenticStreamingIterator: Unexpected result type from hooks: %s",
|
||||
type(result).__name__,
|
||||
)
|
||||
except Exception as e:
|
||||
_call_id = getattr(self._logging_obj, "litellm_call_id", "unknown")
|
||||
verbose_logger.exception(
|
||||
"AgenticStreamingIterator: Error in agentic hook processing "
|
||||
"[call_id=%s model=%s]: %s",
|
||||
_call_id,
|
||||
self._model,
|
||||
str(e),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _rebuild_anthropic_response_from_sse(
|
||||
raw_bytes: List[bytes],
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Parse collected SSE bytes into an Anthropic Messages response dict.
|
||||
|
||||
Processes SSE events in order:
|
||||
- message_start -> envelope (id, model, role, usage)
|
||||
- content_block_start -> new content block
|
||||
- content_block_delta -> accumulate text/json/thinking deltas
|
||||
- content_block_stop -> finalize block
|
||||
- message_delta -> stop_reason, output usage
|
||||
- message_stop -> end
|
||||
"""
|
||||
events = _parse_sse_events(b"".join(raw_bytes))
|
||||
|
||||
response: Dict[str, Any] = {
|
||||
"id": "",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 0, "output_tokens": 0},
|
||||
}
|
||||
content_blocks: Dict[int, Dict[str, Any]] = {}
|
||||
saw_message_start = False
|
||||
|
||||
for event_type, data in events:
|
||||
if event_type == "message_start":
|
||||
saw_message_start = True
|
||||
_handle_message_start(data, response)
|
||||
elif event_type == "content_block_start":
|
||||
_handle_content_block_start(data, content_blocks)
|
||||
elif event_type == "content_block_delta":
|
||||
_handle_content_block_delta(data, content_blocks)
|
||||
elif event_type == "content_block_stop":
|
||||
_handle_content_block_stop(data, content_blocks)
|
||||
elif event_type == "message_delta":
|
||||
_handle_message_delta(data, response)
|
||||
|
||||
if not saw_message_start:
|
||||
return None
|
||||
|
||||
for idx in sorted(content_blocks.keys()):
|
||||
block = content_blocks[idx]
|
||||
block.pop("_partial_json", None)
|
||||
response["content"].append(block)
|
||||
|
||||
return response
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ if TYPE_CHECKING:
|
|||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
|
||||
class BaseTranslation(ABC):
|
||||
|
|
@ -101,6 +102,16 @@ class BaseTranslation(ABC):
|
|||
"""
|
||||
return responses_so_far
|
||||
|
||||
def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]:
|
||||
"""
|
||||
Convert request data to OpenAI-spec structured messages.
|
||||
|
||||
Override in subclasses for format-specific conversion.
|
||||
|
||||
Returns None if no convertible content is found.
|
||||
"""
|
||||
return None
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> List[str]:
|
||||
"""
|
||||
Extract tool names from the request body for allowlist/policy checks.
|
||||
|
|
|
|||
|
|
@ -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)()
|
||||
|
|
|
|||
|
|
@ -78,6 +78,10 @@ from litellm.types.containers.main import (
|
|||
DeleteContainerResult,
|
||||
)
|
||||
from litellm.types.files import TwoStepFileUploadConfig
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
|
@ -2047,7 +2051,23 @@ class BaseLLMHTTPHandler:
|
|||
request_body=request_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
initial_response = completion_stream
|
||||
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
|
||||
AgenticAnthropicStreamingIterator,
|
||||
)
|
||||
|
||||
initial_response = AgenticAnthropicStreamingIterator(
|
||||
completion_stream=completion_stream,
|
||||
http_handler=self,
|
||||
model=model,
|
||||
messages=messages,
|
||||
anthropic_messages_provider_config=anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
return initial_response
|
||||
else:
|
||||
initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response(
|
||||
model=model,
|
||||
|
|
@ -2055,7 +2075,7 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# Call agentic completion hooks
|
||||
# Call agentic completion hooks (non-streaming path only)
|
||||
final_response = await self._call_agentic_completion_hooks(
|
||||
response=initial_response,
|
||||
model=model,
|
||||
|
|
@ -2063,7 +2083,7 @@ class BaseLLMHTTPHandler:
|
|||
anthropic_messages_provider_config=anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream or False,
|
||||
stream=False,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
|
@ -4516,6 +4536,167 @@ class BaseLLMHTTPHandler:
|
|||
return stream, data
|
||||
return stream, data
|
||||
|
||||
@staticmethod
|
||||
def _get_agentic_loop_settings(kwargs: Dict) -> Tuple[int, int, List[str]]:
|
||||
depth = int(kwargs.get("_agentic_loop_depth", 0) or 0)
|
||||
max_loops = int(kwargs.get("max_agentic_loops", 3) or 3)
|
||||
fingerprints = list(kwargs.get("_agentic_loop_fingerprints", []) or [])
|
||||
return depth, max(max_loops, 1), fingerprints
|
||||
|
||||
@staticmethod
|
||||
def _check_agentic_loop_safety(
|
||||
tool_calls: Any,
|
||||
fingerprints: List[str],
|
||||
depth: int,
|
||||
max_loops: int,
|
||||
model: str,
|
||||
) -> str:
|
||||
"""
|
||||
Evaluate agentic-loop safety guards (fingerprint cycle / max depth).
|
||||
|
||||
Raises ValueError on abort. Returns the current fingerprint on success.
|
||||
|
||||
These checks must not be swallowed by the per-callback ``except Exception``
|
||||
block that wraps callback dispatch — they are bounded-loop / cycle-break
|
||||
safety rails and must abort the agentic dispatch when they trip.
|
||||
"""
|
||||
fingerprint = BaseLLMHTTPHandler._fingerprint_agentic_tools(tool_calls)
|
||||
if fingerprint in fingerprints:
|
||||
raise ValueError(
|
||||
"Agentic loop detected repeated tool-call fingerprint; aborting rerun"
|
||||
)
|
||||
if depth >= max_loops:
|
||||
raise ValueError(
|
||||
f"Exceeded max_agentic_loops={max_loops} for model={model}"
|
||||
)
|
||||
return fingerprint
|
||||
|
||||
@staticmethod
|
||||
def _fingerprint_agentic_tools(tools: Dict) -> str:
|
||||
try:
|
||||
return json.dumps(tools, sort_keys=True, default=str)
|
||||
except Exception:
|
||||
return str(tools)
|
||||
|
||||
async def _execute_anthropic_agentic_plan(
|
||||
self,
|
||||
plan: AgenticLoopPlan,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
anthropic_messages_optional_request_params: Dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
kwargs: Dict,
|
||||
depth: int,
|
||||
max_loops: int,
|
||||
fingerprints: List[str],
|
||||
fingerprint: str,
|
||||
stream: bool = False,
|
||||
) -> Any:
|
||||
from litellm.anthropic_interface import messages as anthropic_messages
|
||||
|
||||
patch = plan.request_patch or AgenticLoopRequestPatch()
|
||||
if patch.messages is None:
|
||||
raise ValueError("Agentic loop plan missing patched messages")
|
||||
|
||||
full_model_name = model
|
||||
if logging_obj is not None:
|
||||
agentic_params = logging_obj.model_call_details.get(
|
||||
"agentic_loop_params", {}
|
||||
)
|
||||
full_model_name = cast(str, agentic_params.get("model", model))
|
||||
|
||||
optional_params = dict(anthropic_messages_optional_request_params)
|
||||
optional_params.update(patch.optional_params)
|
||||
if patch.tools is not None:
|
||||
optional_params["tools"] = patch.tools
|
||||
|
||||
max_tokens = patch.max_tokens
|
||||
if max_tokens is None:
|
||||
max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None))
|
||||
else:
|
||||
optional_params.pop("max_tokens", None)
|
||||
if max_tokens is None:
|
||||
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
|
||||
|
||||
internal_keys = {"litellm_logging_obj"}
|
||||
kwargs_for_followup = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if not k.startswith("_websearch_interception")
|
||||
and not k.startswith("_compression_interception")
|
||||
and k not in internal_keys
|
||||
and k not in optional_params
|
||||
}
|
||||
kwargs_for_followup.update(patch.kwargs)
|
||||
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
|
||||
kwargs_for_followup["max_agentic_loops"] = max_loops
|
||||
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
|
||||
|
||||
return await anthropic_messages.acreate(
|
||||
**{
|
||||
"max_tokens": max_tokens,
|
||||
"messages": patch.messages,
|
||||
"model": patch.model or full_model_name,
|
||||
"stream": stream,
|
||||
**optional_params,
|
||||
**kwargs_for_followup,
|
||||
}
|
||||
)
|
||||
|
||||
async def _execute_chat_completion_agentic_plan(
|
||||
self,
|
||||
plan: AgenticLoopPlan,
|
||||
model: str,
|
||||
messages: List[Dict],
|
||||
optional_params: Dict,
|
||||
kwargs: Dict,
|
||||
custom_llm_provider: str,
|
||||
depth: int,
|
||||
max_loops: int,
|
||||
fingerprints: List[str],
|
||||
fingerprint: str,
|
||||
) -> Any:
|
||||
patch = plan.request_patch or AgenticLoopRequestPatch()
|
||||
if patch.messages is None:
|
||||
raise ValueError("Agentic loop plan missing patched messages")
|
||||
|
||||
full_model_name = patch.model or model
|
||||
if "/" not in full_model_name:
|
||||
full_model_name = f"{custom_llm_provider}/{full_model_name}"
|
||||
|
||||
optional_params_for_followup = dict(optional_params)
|
||||
optional_params_for_followup.update(patch.optional_params)
|
||||
if patch.tools is not None:
|
||||
optional_params_for_followup["tools"] = patch.tools
|
||||
|
||||
internal_params = {
|
||||
"_websearch_interception",
|
||||
"acompletion",
|
||||
"litellm_logging_obj",
|
||||
"custom_llm_provider",
|
||||
"model_alias_map",
|
||||
"stream_response",
|
||||
"custom_prompt_dict",
|
||||
}
|
||||
kwargs_for_followup = {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if not k.startswith("_websearch_interception")
|
||||
and not k.startswith("_compression_interception")
|
||||
and k not in internal_params
|
||||
}
|
||||
kwargs_for_followup.update(patch.kwargs)
|
||||
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
|
||||
kwargs_for_followup["max_agentic_loops"] = max_loops
|
||||
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
|
||||
|
||||
return await litellm.acompletion(
|
||||
model=full_model_name,
|
||||
messages=patch.messages,
|
||||
**optional_params_for_followup,
|
||||
**kwargs_for_followup,
|
||||
)
|
||||
|
||||
async def _call_agentic_completion_hooks(
|
||||
self,
|
||||
response: Any,
|
||||
|
|
@ -4541,45 +4722,111 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
|
||||
tools = anthropic_messages_optional_request_params.get("tools", [])
|
||||
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
|
||||
|
||||
for callback in callbacks:
|
||||
if not isinstance(callback, CustomLogger):
|
||||
continue
|
||||
|
||||
should_run: bool = False
|
||||
tool_calls: Any = None
|
||||
try:
|
||||
if isinstance(callback, CustomLogger):
|
||||
# First: Check if agentic loop should run
|
||||
(
|
||||
should_run,
|
||||
tool_calls,
|
||||
) = await callback.async_should_run_agentic_loop(
|
||||
response=response,
|
||||
# First: Check if agentic loop should run. Wrap in try/except
|
||||
# to shield from buggy user callbacks — a callback crash should
|
||||
# not abort the whole request.
|
||||
(
|
||||
should_run,
|
||||
tool_calls,
|
||||
) = await callback.async_should_run_agentic_loop(
|
||||
response=response,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in "
|
||||
"async_should_run_agentic_loop [call_id=%s model=%s]: %s",
|
||||
_call_id,
|
||||
model,
|
||||
str(e),
|
||||
)
|
||||
continue
|
||||
|
||||
if not should_run:
|
||||
continue
|
||||
|
||||
# Safety guards must run OUTSIDE the callback try/except — they are
|
||||
# bounded-loop / cycle-break rails that must propagate to the caller.
|
||||
fingerprint = self._check_agentic_loop_safety(
|
||||
tool_calls=tool_calls,
|
||||
fingerprints=fingerprints,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
model=model,
|
||||
)
|
||||
|
||||
try:
|
||||
kwargs_with_provider = kwargs.copy() if kwargs else {}
|
||||
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
|
||||
build_plan_overridden = (
|
||||
callback.__class__.async_build_agentic_loop_plan
|
||||
is not CustomLogger.async_build_agentic_loop_plan
|
||||
)
|
||||
if not build_plan_overridden:
|
||||
return await callback.async_run_agentic_loop(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
response=response,
|
||||
anthropic_messages_provider_config=anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
kwargs=kwargs_with_provider,
|
||||
)
|
||||
|
||||
if should_run:
|
||||
# Second: Execute agentic loop
|
||||
# Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name
|
||||
kwargs_with_provider = kwargs.copy() if kwargs else {}
|
||||
kwargs_with_provider["custom_llm_provider"] = (
|
||||
custom_llm_provider
|
||||
)
|
||||
agentic_response = await callback.async_run_agentic_loop(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
anthropic_messages_provider_config=anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=kwargs_with_provider,
|
||||
)
|
||||
# First hook that runs agentic loop wins
|
||||
return agentic_response
|
||||
plan = await callback.async_build_agentic_loop_plan(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
anthropic_messages_provider_config=anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=kwargs_with_provider,
|
||||
)
|
||||
|
||||
if plan.response_override is not None:
|
||||
return plan.response_override
|
||||
if plan.terminate:
|
||||
verbose_logger.debug(
|
||||
"Agentic loop terminated by callback=%s reason=%s",
|
||||
callback.__class__.__name__,
|
||||
plan.stop_reason,
|
||||
)
|
||||
return response
|
||||
if not plan.run_agentic_loop:
|
||||
continue
|
||||
|
||||
return await self._execute_anthropic_agentic_plan(
|
||||
plan=plan,
|
||||
model=model,
|
||||
messages=messages,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs_with_provider,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
fingerprints=fingerprints,
|
||||
fingerprint=fingerprint,
|
||||
stream=stream,
|
||||
)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
verbose_logger.exception(
|
||||
|
|
@ -4653,52 +4900,104 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
|
||||
tools = optional_params.get("tools", [])
|
||||
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
|
||||
|
||||
for callback in callbacks:
|
||||
try:
|
||||
if isinstance(callback, CustomLogger):
|
||||
# Check if callback has the chat completion agentic loop method
|
||||
if not hasattr(
|
||||
callback, "async_should_run_chat_completion_agentic_loop"
|
||||
):
|
||||
continue
|
||||
if not isinstance(callback, CustomLogger):
|
||||
continue
|
||||
if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"):
|
||||
continue
|
||||
|
||||
# First: Check if agentic loop should run
|
||||
(
|
||||
should_run,
|
||||
tool_calls,
|
||||
) = await callback.async_should_run_chat_completion_agentic_loop(
|
||||
response=response,
|
||||
should_run: bool = False
|
||||
tool_calls: Any = None
|
||||
try:
|
||||
(
|
||||
should_run,
|
||||
tool_calls,
|
||||
) = await callback.async_should_run_chat_completion_agentic_loop(
|
||||
response=response,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in "
|
||||
"async_should_run_chat_completion_agentic_loop: %s",
|
||||
str(e),
|
||||
)
|
||||
continue
|
||||
|
||||
if not should_run:
|
||||
continue
|
||||
|
||||
# Safety guards must run OUTSIDE the callback try/except — they are
|
||||
# bounded-loop / cycle-break rails that must propagate to the caller.
|
||||
fingerprint = self._check_agentic_loop_safety(
|
||||
tool_calls=tool_calls,
|
||||
fingerprints=fingerprints,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
model=model,
|
||||
)
|
||||
|
||||
try:
|
||||
kwargs_with_provider = kwargs.copy() if kwargs else {}
|
||||
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
|
||||
build_plan_overridden = (
|
||||
callback.__class__.async_build_chat_completion_agentic_loop_plan
|
||||
is not CustomLogger.async_build_chat_completion_agentic_loop_plan
|
||||
)
|
||||
if not build_plan_overridden:
|
||||
return await callback.async_run_chat_completion_agentic_loop(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
response=response,
|
||||
optional_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs,
|
||||
kwargs=kwargs_with_provider,
|
||||
)
|
||||
|
||||
if should_run:
|
||||
# Second: Execute agentic loop
|
||||
# Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name
|
||||
kwargs_with_provider = kwargs.copy() if kwargs else {}
|
||||
kwargs_with_provider["custom_llm_provider"] = (
|
||||
custom_llm_provider
|
||||
)
|
||||
agentic_response = (
|
||||
await callback.async_run_chat_completion_agentic_loop(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
optional_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=kwargs_with_provider,
|
||||
)
|
||||
)
|
||||
# First hook that runs agentic loop wins
|
||||
return agentic_response
|
||||
plan = await callback.async_build_chat_completion_agentic_loop_plan(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
optional_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=kwargs_with_provider,
|
||||
)
|
||||
|
||||
if plan.response_override is not None:
|
||||
return plan.response_override
|
||||
if plan.terminate:
|
||||
verbose_logger.debug(
|
||||
"Agentic chat loop terminated by callback=%s reason=%s",
|
||||
callback.__class__.__name__,
|
||||
plan.stop_reason,
|
||||
)
|
||||
return response
|
||||
if not plan.run_agentic_loop:
|
||||
continue
|
||||
|
||||
return await self._execute_chat_completion_agentic_plan(
|
||||
plan=plan,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs_with_provider,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
fingerprints=fingerprints,
|
||||
fingerprint=fingerprint,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {str(e)}"
|
||||
|
|
@ -5216,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:
|
||||
|
|
@ -5312,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:
|
||||
|
|
|
|||
|
|
@ -294,9 +294,7 @@ class Authenticator:
|
|||
access_token_url = os.getenv(
|
||||
"GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL
|
||||
)
|
||||
client_id = os.getenv(
|
||||
"GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID
|
||||
)
|
||||
client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -48,6 +48,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
Methods can be overridden to customize behavior for different message formats.
|
||||
"""
|
||||
|
||||
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
|
||||
"""
|
||||
Convert chat completions request data to OpenAI-spec structured messages.
|
||||
|
||||
Messages are already in OpenAI format, so this is a simple extraction.
|
||||
"""
|
||||
messages = data.get("messages")
|
||||
if messages is None:
|
||||
return None
|
||||
return cast(List[AllMessageValues], messages)
|
||||
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -68,9 +79,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
tool_calls_to_check: List[ChatCompletionToolParam] = []
|
||||
text_task_mappings: List[Tuple[int, Optional[int]]] = []
|
||||
tool_call_task_mappings: List[Tuple[int, int]] = []
|
||||
# text_task_mappings: Track (message_index, content_index) for each text
|
||||
# content_index is None for string content, int for list content
|
||||
# tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call
|
||||
|
||||
# Step 1: Extract all text content, images, and tool calls
|
||||
for msg_idx, message in enumerate(messages):
|
||||
|
|
@ -92,12 +100,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
inputs["images"] = images_to_check
|
||||
if tool_calls_to_check:
|
||||
inputs["tool_calls"] = tool_calls_to_check # type: ignore
|
||||
if messages:
|
||||
msg_list = cast(List[AllMessageValues], messages)
|
||||
structured_messages = self.get_structured_messages(data)
|
||||
if structured_messages:
|
||||
inputs["structured_messages"] = (
|
||||
openai_messages_without_system(msg_list)
|
||||
openai_messages_without_system(structured_messages)
|
||||
if skip_system
|
||||
else msg_list
|
||||
else structured_messages
|
||||
)
|
||||
# Pass tools (function definitions) to the guardrail
|
||||
tools = data.get("tools")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from litellm.responses.litellm_completion_transformation.transformation import (
|
|||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolParam,
|
||||
)
|
||||
|
|
@ -70,6 +71,24 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
Methods can be overridden to customize behavior for different message formats.
|
||||
"""
|
||||
|
||||
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
|
||||
"""
|
||||
Convert Responses API request data to OpenAI-spec structured messages.
|
||||
|
||||
Transforms `input` (string or ResponseInputParam) and optional
|
||||
`instructions` into chat completion messages.
|
||||
"""
|
||||
input_data = data.get("input")
|
||||
if input_data is None:
|
||||
return None
|
||||
messages = (
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=input_data,
|
||||
responses_api_request=data,
|
||||
)
|
||||
)
|
||||
return cast(List[AllMessageValues], messages) if messages else None
|
||||
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -86,12 +105,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if input_data is None:
|
||||
return data
|
||||
|
||||
structured_messages = (
|
||||
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
|
||||
input=input_data,
|
||||
responses_api_request=data,
|
||||
)
|
||||
)
|
||||
structured_messages = self.get_structured_messages(data)
|
||||
|
||||
# Handle simple string input
|
||||
if isinstance(input_data, str):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
158
litellm/llms/scaleway/audio_transcription/transformation.py
Normal file
158
litellm/llms/scaleway/audio_transcription/transformation.py
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
"""
|
||||
Support for Scaleway's OpenAI-compatible `/v1/audio/transcriptions` endpoint.
|
||||
|
||||
API reference: https://www.scaleway.com/en/developers/api/generative-apis/#path-audio-create-an-audio-transcription
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
|
||||
from litellm.llms.base_llm.audio_transcription.transformation import (
|
||||
AudioTranscriptionRequestData,
|
||||
BaseAudioTranscriptionConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
OpenAIAudioTranscriptionOptionalParams,
|
||||
)
|
||||
from litellm.types.utils import FileTypes, TranscriptionResponse
|
||||
|
||||
|
||||
class ScalewayAudioTranscriptionException(BaseLLMException):
|
||||
pass
|
||||
|
||||
|
||||
class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
|
||||
def get_supported_openai_params(
|
||||
self, model: str
|
||||
) -> List[OpenAIAudioTranscriptionOptionalParams]:
|
||||
return [
|
||||
"language",
|
||||
"prompt",
|
||||
"response_format",
|
||||
"temperature",
|
||||
"timestamp_granularities",
|
||||
]
|
||||
|
||||
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)
|
||||
for k, v in non_default_params.items():
|
||||
if k in supported_params:
|
||||
optional_params[k] = v
|
||||
return optional_params
|
||||
|
||||
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:
|
||||
api_base = (
|
||||
"https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/")
|
||||
)
|
||||
return f"{api_base}/audio/transcriptions"
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
|
||||
) -> BaseLLMException:
|
||||
return ScalewayAudioTranscriptionException(
|
||||
message=error_message,
|
||||
status_code=status_code,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
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:
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("SCW_SECRET_KEY")
|
||||
|
||||
if not api_key:
|
||||
raise ScalewayAudioTranscriptionException(
|
||||
message=(
|
||||
"Scaleway API key not found. Pass `api_key=...` or set the "
|
||||
"SCW_SECRET_KEY environment variable."
|
||||
),
|
||||
status_code=401,
|
||||
headers={},
|
||||
)
|
||||
|
||||
default_headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"accept": "application/json",
|
||||
}
|
||||
default_headers.update(headers or {})
|
||||
return default_headers
|
||||
|
||||
def transform_audio_transcription_request(
|
||||
self,
|
||||
model: str,
|
||||
audio_file: FileTypes,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> AudioTranscriptionRequestData:
|
||||
processed_audio = process_audio_file(audio_file)
|
||||
|
||||
form_fields: dict = {"model": model}
|
||||
for key in self.get_supported_openai_params(model):
|
||||
value = optional_params.get(key)
|
||||
if value is not None:
|
||||
form_fields[key] = value
|
||||
|
||||
files = {
|
||||
"file": (
|
||||
processed_audio.filename,
|
||||
processed_audio.file_content,
|
||||
processed_audio.content_type,
|
||||
)
|
||||
}
|
||||
|
||||
return AudioTranscriptionRequestData(data=form_fields, files=files)
|
||||
|
||||
def transform_audio_transcription_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
) -> TranscriptionResponse:
|
||||
content_type = (raw_response.headers.get("content-type") or "").lower()
|
||||
if "application/json" not in content_type:
|
||||
return TranscriptionResponse(text=raw_response.text)
|
||||
|
||||
try:
|
||||
response_json = raw_response.json()
|
||||
except Exception:
|
||||
raise ScalewayAudioTranscriptionException(
|
||||
message=raw_response.text,
|
||||
status_code=raw_response.status_code,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
|
||||
text = response_json.get("text") or ""
|
||||
response = TranscriptionResponse(text=text)
|
||||
|
||||
if "segments" in response_json:
|
||||
response["segments"] = response_json["segments"]
|
||||
if "language" in response_json:
|
||||
response["language"] = response_json["language"]
|
||||
|
||||
response._hidden_params = response_json
|
||||
return response
|
||||
|
|
@ -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("/")
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue