mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into pr37724
# Conflicts: # tests/test_litellm/test_cost_calculator.py
This commit is contained in:
commit
885d95d71b
2909 changed files with 104892 additions and 29295 deletions
|
|
@ -430,7 +430,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise \
|
||||
--cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=20 \
|
||||
|
|
@ -504,7 +504,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise \
|
||||
--cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=20 \
|
||||
|
|
@ -631,7 +631,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 2"
|
||||
|
|
@ -651,126 +651,6 @@ jobs:
|
|||
- auth_ui_unit_tests_coverage.xml
|
||||
- auth_ui_unit_tests_coverage
|
||||
|
||||
proxy_behavior_tests:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: litellm_test
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Seed DB schema via prisma db push
|
||||
command: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
- run:
|
||||
name: Generate Prisma Client
|
||||
command: uv run --no-sync python -m prisma generate
|
||||
- run:
|
||||
name: Run proxy management behavior tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
uv run --no-sync python -m pytest tests/proxy_behavior \
|
||||
-v --junitxml=test-results/junit.xml --durations=10
|
||||
no_output_timeout: 15m
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
|
||||
proxy_security_tests:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: litellm_test
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Seed DB schema via prisma db push
|
||||
command: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
- run:
|
||||
name: Generate Prisma Client
|
||||
command: uv run --no-sync python -m prisma generate
|
||||
- run:
|
||||
name: Run proxy security tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
uv run --no-sync python -m pytest tests/proxy_security_tests \
|
||||
-v --junitxml=test-results/junit.xml --durations=10
|
||||
no_output_timeout: 15m
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
|
||||
schema_migration_check:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: litellm_test
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
# An empty database; the test applies every committed migration itself.
|
||||
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Generate Prisma Client
|
||||
command: uv run --no-sync python -m prisma generate
|
||||
- run:
|
||||
name: Check schema.prisma is in sync with committed migrations
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
uv run --no-sync python -m pytest tests/proxy_migration_tests \
|
||||
-v --junitxml=test-results/junit.xml --durations=10
|
||||
no_output_timeout: 15m
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
|
||||
litellm_router_testing: # Runs all tests with the "router" keyword
|
||||
docker:
|
||||
- *python312_image
|
||||
|
|
@ -858,7 +738,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 4"
|
||||
|
|
@ -985,7 +865,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=20 \
|
||||
-n 4 \
|
||||
|
|
@ -1030,7 +910,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
no_output_timeout: 15m
|
||||
|
|
@ -1074,7 +954,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 2 \
|
||||
|
|
@ -1120,7 +1000,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
--retries 3 --retry-delay 5"
|
||||
|
|
@ -1211,7 +1091,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 4"
|
||||
|
|
@ -1255,7 +1135,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 4"
|
||||
|
|
@ -1274,40 +1154,6 @@ jobs:
|
|||
paths:
|
||||
- search_coverage.xml
|
||||
- search_coverage
|
||||
litellm_mapped_enterprise_tests:
|
||||
docker:
|
||||
- *python312_image
|
||||
working_directory: ~/project
|
||||
resource_class: large
|
||||
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- setup_litellm_enterprise_pip
|
||||
- run:
|
||||
name: Run enterprise tests
|
||||
command: |
|
||||
uv run --no-sync python -m prisma generate
|
||||
mkdir -p test-results
|
||||
TEST_FILES=$(circleci tests glob "tests/enterprise/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit-enterprise.xml \
|
||||
--durations=10 \
|
||||
-n 4"
|
||||
no_output_timeout: 15m
|
||||
# Store test results
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
batches_testing:
|
||||
docker:
|
||||
- *python312_image
|
||||
|
|
@ -1333,7 +1179,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 2"
|
||||
|
|
@ -1377,7 +1223,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 2"
|
||||
|
|
@ -1422,7 +1268,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
-n 4"
|
||||
|
|
@ -1501,7 +1347,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
-n 4 \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
|
|
@ -1546,7 +1392,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
no_output_timeout: 15m
|
||||
|
|
@ -1599,7 +1445,7 @@ jobs:
|
|||
--verbose \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 -n 2 \
|
||||
--reruns 2 --reruns-delay 1"
|
||||
|
|
@ -3105,12 +2951,6 @@ workflows:
|
|||
filters: *main_branches
|
||||
- auth_ui_unit_tests:
|
||||
filters: *main_branches
|
||||
- proxy_behavior_tests:
|
||||
filters: *main_branches
|
||||
- proxy_security_tests:
|
||||
filters: *main_branches
|
||||
- schema_migration_check:
|
||||
filters: *main_branches
|
||||
- build_docker_database_image:
|
||||
filters: *main_branches
|
||||
- e2e_ui_testing:
|
||||
|
|
@ -3167,8 +3007,6 @@ workflows:
|
|||
filters: *main_branches
|
||||
- search_testing:
|
||||
filters: *main_branches
|
||||
- litellm_mapped_enterprise_tests:
|
||||
filters: *main_branches
|
||||
- batches_testing:
|
||||
filters: *main_branches
|
||||
- litellm_utils_testing:
|
||||
|
|
@ -3191,7 +3029,6 @@ workflows:
|
|||
- guardrails_testing
|
||||
- ocr_testing
|
||||
- search_testing
|
||||
- litellm_mapped_enterprise_tests
|
||||
- batches_testing
|
||||
- litellm_utils_testing
|
||||
- pass_through_unit_testing
|
||||
|
|
|
|||
9
.github/CODEOWNERS
vendored
9
.github/CODEOWNERS
vendored
|
|
@ -1,5 +1,10 @@
|
|||
/ui/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/litellm/proxy/_experimental/out/ @yuneng-jiang @ryan-crabbe-berri
|
||||
/ui/ @yuneng-berri @ryan-crabbe-berri
|
||||
/litellm/proxy/_experimental/out/ @yuneng-berri @ryan-crabbe-berri
|
||||
/ui/Dockerfile
|
||||
/ui/nginx.conf
|
||||
/ui/litellm-dashboard/src/lib/http/schema.d.ts
|
||||
/ui/litellm-dashboard/tsconfig.tsbuildinfo
|
||||
/model_prices_and_context_window.json @mateo-berri
|
||||
/litellm/model_prices_and_context_window_backup.json @mateo-berri
|
||||
/litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri
|
||||
/.github/CODEOWNERS @yuneng-berri
|
||||
|
|
|
|||
31
.github/actions/cache-cargo-build/action.yml
vendored
Normal file
31
.github/actions/cache-cargo-build/action.yml
vendored
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
name: "Cache the Rust build"
|
||||
description: >-
|
||||
Cache the Cargo registry and target directory the root package's build needs,
|
||||
so only the first job on a given Cargo.lock compiles the bridge from scratch.
|
||||
|
||||
litellm builds through maturin, which compiles litellm-rust/crates/python-bridge
|
||||
in release mode before it can produce a wheel. `uv sync` therefore pays a full
|
||||
build in every job that installs the workspace: measured at 2m40s per unit shard
|
||||
on 2026-08-21, more than the whole unit tier spends running tests. Nothing caught
|
||||
it, because the uv cache holds wheels uv downloads rather than wheels it builds,
|
||||
and a path dependency whose source moves every commit could never hit that cache
|
||||
anyway. Cargo rebuilds only what changed when its target directory survives, so a
|
||||
warm job pays for the bridge crate alone.
|
||||
|
||||
The key namespace is separate from test-rust.yml's. Both cache the same directory,
|
||||
but that workflow fills it with debug and clippy artifacts, which a release build
|
||||
cannot reuse, and a shared key would let whichever ran first deny the other a save.
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- name: Restore the Cargo registry and target directory
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
litellm-rust/target
|
||||
key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-release-
|
||||
46
.github/ci-coverage-allowlist.yml
vendored
46
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -5,24 +5,21 @@ description: >-
|
|||
|
||||
test_paths:
|
||||
- reason: >-
|
||||
The caching suite in tests/local_testing, which runs nowhere. Every job that globs that
|
||||
directory either deselects it (local_testing_part1 and part2 carry `-k "... and not caching
|
||||
and not cache"`) or keeps only another keyword (langfuse, router, assistants), and no job
|
||||
names these files the way redis_caching_unit_tests names test_dual_cache.py. Measured
|
||||
2026-08-20 by collecting the directory under each job's own selector: 118 tests across
|
||||
these eight files are selected by none of them. Listed so the gap is a decision rather
|
||||
than an accident, and so the --slices guard has a baseline to ratchet down from. Revisit
|
||||
when tests/local_testing is ported off CircleCI, where the keyless part of this suite
|
||||
belongs in a real job
|
||||
What is left of the caching suite in tests/local_testing that runs nowhere. Every job that
|
||||
globs that directory either deselects it (local_testing_part1 and part2 carry `-k "... and
|
||||
not caching and not cache"`) or keeps only another keyword (langfuse, router, assistants),
|
||||
and no job names these files the way redis_caching_unit_tests names test_dual_cache.py.
|
||||
The gap was eight files and 118 tests when measured 2026-08-20; the five keyless ones now
|
||||
run in the caching-local shard, leaving these three. Measured 2026-08-21 with no provider
|
||||
credentials and no Redis: test_caching.py needs both (37 of 65 fail without them),
|
||||
test_disk_cache_unit_tests.py needs OPENAI_API_KEY for 2 of its 4, and
|
||||
test_gcs_cache_unit_tests.py needs GCS credentials for all 4. They want the keyless/live
|
||||
split that porting tests/local_testing off CircleCI will force, not a job that is red by
|
||||
construction
|
||||
paths:
|
||||
- tests/local_testing/test_cache_preset_key.py
|
||||
- tests/local_testing/test_caching.py
|
||||
- tests/local_testing/test_caching_handler.py
|
||||
- tests/local_testing/test_disk_cache_unit_tests.py
|
||||
- tests/local_testing/test_gcs_cache_unit_tests.py
|
||||
- tests/local_testing/test_prompt_caching.py
|
||||
- tests/local_testing/test_responses_stream_cache_keys.py
|
||||
- tests/local_testing/test_unit_test_caching.py
|
||||
- reason: >-
|
||||
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
|
||||
from a pull request; it needs a live gateway and provider credentials no PR job holds
|
||||
|
|
@ -48,16 +45,6 @@ test_paths:
|
|||
choice it informed is settled
|
||||
paths:
|
||||
- tests/code_coverage_tests/test_aio_http_image_conversion.py
|
||||
- reason: >-
|
||||
The last file of a second mirror that sat beside tests/test_litellm and ran nowhere. Its
|
||||
other 33 files landed in the real mirror during August 2026, 30 as moves and 3 by merging
|
||||
their bodies into the live file of the same name. This one cannot follow either route yet:
|
||||
its live twin was rewritten from 1268 lines to 9434, and of the 19 tests here 5 have no
|
||||
counterpart while 25 assertions fail against today's code, so what survives that rewrite
|
||||
is a judgement about the endpoints, not a merge. Revisit by deciding which of the five
|
||||
behaviours still hold
|
||||
paths:
|
||||
- tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
|
||||
- reason: >-
|
||||
No job invokes this suite and its files mix pure transformation tests with ones driving live
|
||||
vendor vector stores, so assigning them needs a per-file decision
|
||||
|
|
@ -102,17 +89,6 @@ test_paths:
|
|||
- tests/integration/sandbox/test_e2b_sandbox.py
|
||||
- tests/integration/test_oci_integration.py
|
||||
- tests/integration/test_oci_proxy_integration.py
|
||||
- reason: >-
|
||||
A unit test for the proxy-extras package that no job invokes, while the package's other
|
||||
tests live under tests/proxy_migration_tests. Measured 2026-08-20: 24 of its 28 tests pass
|
||||
and the 4 in TestMigrationSQLIdempotency fail, because 13 migrations from 2026-03 onward use
|
||||
bare CREATE TABLE, ADD COLUMN, CREATE INDEX and ADD CONSTRAINT rather than the guarded forms
|
||||
this file requires. It also matches those keywords inside SQL comments, so two further
|
||||
migrations are reported that are in fact fine. Wiring it up means deciding what to do about
|
||||
the 13 first, and they cannot simply be edited: Prisma checksums an applied migration, so a
|
||||
changed one breaks migrate deploy for existing installs
|
||||
paths:
|
||||
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
|
||||
|
||||
dockerfiles:
|
||||
- reason: >-
|
||||
|
|
|
|||
5
.github/mutmut-coverage.rc
vendored
Normal file
5
.github/mutmut-coverage.rc
vendored
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
# mutmut's gather_coverage() looks covered lines up by absolute path, so the
|
||||
# repo's `relative_files = true` makes every lookup miss and mutmut generates
|
||||
# zero mutants. Point COVERAGE_RCFILE here for mutation runs only.
|
||||
[run]
|
||||
relative_files = false
|
||||
23
.github/pull_request_template.md
vendored
23
.github/pull_request_template.md
vendored
|
|
@ -1,7 +1,10 @@
|
|||
<!-- The whole description's target audience is humans, not AI agents: write it in plain, simple,
|
||||
everyday engineering language, extremely parsable and readable at a glance. This goes double for
|
||||
the TLDR, User Flow, and Caveats sections -->
|
||||
|
||||
## TLDR
|
||||
|
||||
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
|
||||
This section must be extremely human parsable, comprehensible, and readable: its target audience is humans, not AI agents -->
|
||||
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max -->
|
||||
|
||||
Problem this solves:
|
||||
|
||||
|
|
@ -110,8 +113,20 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
|
||||
## Caveats (if any)
|
||||
|
||||
<!-- Short bullet points, just like the TLDR: one line per bullet, roughly 10 words max
|
||||
<!-- Group caveats under severity subheadings (### Severe, ### High, ### Medium, ### Low), with
|
||||
short bullet points inside each, just like the TLDR: one line per bullet, roughly 10 words max
|
||||
Call out known limitations, follow-up work, or anything a reviewer should watch out for
|
||||
Include only the tiers that have caveats; drop the empty ones
|
||||
- Severe: inherent to what the PR deliberately ships, there even when the code works as intended:
|
||||
it can degrade or take down a running deployment (e.g. a slow or table-locking boot migration),
|
||||
rewrite data by design, break an existing workflow on purpose, or change auth behavior. An
|
||||
operator must plan around it before rollout
|
||||
- High: an unintended hole: a correctness, security, data-loss, or backward-compatibility bug,
|
||||
unsafe to ship as is
|
||||
- Medium: a real gap someone can hit, but with a workaround or a narrow blast radius
|
||||
- Low: anything else worth noting: naming, cleanup, an edge case nobody hits
|
||||
Nest bullets as deep as helps: hierarchy beats one long line when it makes things clearer to a
|
||||
human reader
|
||||
Leave this section empty if there are none -->
|
||||
|
||||
## QA runbook
|
||||
|
|
@ -134,6 +149,6 @@ Example checklists:
|
|||
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
|
||||
-->
|
||||
|
||||
### Final Attestation
|
||||
## Final Attestation
|
||||
|
||||
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
|
||||
|
|
|
|||
16
.github/scripts/assert_ci_coverage.py
vendored
16
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -312,8 +312,23 @@ def _matchable_names(relative_path: str) -> frozenset[str]:
|
|||
)
|
||||
|
||||
|
||||
def _workflow_named_tokens() -> frozenset[str]:
|
||||
"""Test tokens a GitHub Actions job names directly.
|
||||
|
||||
A CircleCI `-k` that deselects a file no longer means the file runs nowhere once a
|
||||
workflow names it, so the slice check has to credit those the same way the census does.
|
||||
"""
|
||||
return _invoked_test_tokens(
|
||||
scalar
|
||||
for path in _config_files()
|
||||
if path != CIRCLECI_CONFIG
|
||||
for scalar in _scalars(yaml.safe_load(path.read_text(encoding="utf-8")), path.name)
|
||||
)
|
||||
|
||||
|
||||
def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
|
||||
slices: Final = _slices()
|
||||
named_by_workflow: Final = _workflow_named_tokens()
|
||||
globbed: Final = tuple(
|
||||
path
|
||||
for path in _test_files()
|
||||
|
|
@ -326,6 +341,7 @@ def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
|
|||
)
|
||||
for path in globbed
|
||||
if not allowlist.covers_test(path)
|
||||
and not any(_token_covers(token, path) for token in named_by_workflow)
|
||||
and not any(slice_.claims(path, _matchable_names(path)) for slice_ in slices)
|
||||
)
|
||||
|
||||
|
|
|
|||
198
.github/scripts/e2e_egress_sentinel.py
vendored
Executable file
198
.github/scripts/e2e_egress_sentinel.py
vendored
Executable file
|
|
@ -0,0 +1,198 @@
|
|||
"""Prove an e2e replay run makes zero outbound provider calls, by counting them.
|
||||
|
||||
`serve` pins each provider host (`--host`) to a local sink address in the hosts
|
||||
file and binds a counting listener on that address, so any connection the proxy
|
||||
or the record/replay edge opens to a real provider is redirected to the sink,
|
||||
recorded as one line in `--hits-file`, and never leaves the box. The record and
|
||||
replay edge only ever dials `127.0.0.1:<edge-port>` (a different host than the
|
||||
pinned provider names), so in a clean replay the sink sees nothing; a single hit
|
||||
means a provider call escaped the bundle. `assert-empty` turns that hit file into
|
||||
the pass/fail check.
|
||||
|
||||
Stdlib only, so CI runs it under the system interpreter as root (binding :443 and
|
||||
editing the hosts file both need root); `--sink-address`, `--port`, and
|
||||
`--hosts-file` are injectable so it runs unprivileged against a temp hosts file on
|
||||
a high port under test.
|
||||
"""
|
||||
|
||||
# ruff: noqa: T201 # CLI script: its stdout/stderr progress and results are the interface
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import FrameType
|
||||
from typing import Final
|
||||
|
||||
_BLOCK_BEGIN: Final = "# BEGIN e2e-egress-sentinel"
|
||||
_BLOCK_END: Final = "# END e2e-egress-sentinel"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ServeConfig:
|
||||
hosts: tuple[str, ...]
|
||||
sink_address: str
|
||||
ports: tuple[int, ...]
|
||||
hits_file: Path
|
||||
hosts_file: Path
|
||||
ready_file: Path | None
|
||||
pid_file: Path | None
|
||||
|
||||
|
||||
def _pin_block(sink_address: str, hosts: tuple[str, ...]) -> str:
|
||||
lines = "\n".join(f"{sink_address}\t{host}" for host in hosts)
|
||||
return f"\n{_BLOCK_BEGIN}\n{lines}\n{_BLOCK_END}\n"
|
||||
|
||||
|
||||
def _install_pins(hosts_file: Path, sink_address: str, hosts: tuple[str, ...]) -> bytes:
|
||||
original = hosts_file.read_bytes() if hosts_file.exists() else b""
|
||||
hosts_file.write_bytes(original + _pin_block(sink_address, hosts).encode())
|
||||
return original
|
||||
|
||||
|
||||
def _restore_pins(hosts_file: Path, original: bytes) -> None:
|
||||
hosts_file.write_bytes(original)
|
||||
|
||||
|
||||
def _bind(sink_address: str, port: int) -> socket.socket:
|
||||
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
listener.bind((sink_address, port))
|
||||
listener.listen(128)
|
||||
return listener
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _HitLog:
|
||||
path: Path
|
||||
_lock: threading.Lock
|
||||
|
||||
def record(self, *, port: int, peer: tuple[str, int]) -> None:
|
||||
entry = json.dumps({"ts": time.time(), "port": port, "peer": list(peer)})
|
||||
with self._lock:
|
||||
with self.path.open("a", encoding="utf-8") as handle:
|
||||
handle.write(entry + "\n")
|
||||
|
||||
|
||||
def _serve_socket(listener: socket.socket, port: int, hits: _HitLog, stop: threading.Event) -> None:
|
||||
while not stop.is_set():
|
||||
try:
|
||||
conn, peer = listener.accept()
|
||||
except OSError:
|
||||
return
|
||||
hits.record(port=port, peer=(peer[0], peer[1]))
|
||||
try:
|
||||
conn.close()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def serve(config: ServeConfig) -> int:
|
||||
config.hits_file.write_text("", encoding="utf-8")
|
||||
original_hosts = _install_pins(config.hosts_file, config.sink_address, config.hosts)
|
||||
try:
|
||||
listeners = tuple(_bind(config.sink_address, port) for port in config.ports)
|
||||
except OSError as exc:
|
||||
_restore_pins(config.hosts_file, original_hosts)
|
||||
print(f"egress sentinel could not bind a sink: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
stop = threading.Event()
|
||||
hits = _HitLog(path=config.hits_file, _lock=threading.Lock())
|
||||
threads = tuple(
|
||||
threading.Thread(target=_serve_socket, args=(listener, port, hits, stop), daemon=True)
|
||||
for listener, port in zip(listeners, config.ports)
|
||||
)
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
|
||||
def _handle(_signum: int, _frame: FrameType | None) -> None:
|
||||
stop.set()
|
||||
for listener in listeners:
|
||||
try:
|
||||
listener.close()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
signal.signal(signal.SIGTERM, _handle)
|
||||
signal.signal(signal.SIGINT, _handle)
|
||||
|
||||
if config.pid_file is not None:
|
||||
config.pid_file.write_text(str(os.getpid()), encoding="utf-8")
|
||||
if config.ready_file is not None:
|
||||
config.ready_file.write_text("ready", encoding="utf-8")
|
||||
print(
|
||||
f"egress sentinel up: pinned {', '.join(config.hosts)} to {config.sink_address} "
|
||||
f"on port(s) {', '.join(str(p) for p in config.ports)}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
stop.wait()
|
||||
_restore_pins(config.hosts_file, original_hosts)
|
||||
if config.ready_file is not None and config.ready_file.exists():
|
||||
config.ready_file.unlink()
|
||||
if config.pid_file is not None and config.pid_file.exists():
|
||||
config.pid_file.unlink()
|
||||
return 0
|
||||
|
||||
|
||||
def assert_empty(hits_file: Path) -> int:
|
||||
if not hits_file.exists():
|
||||
print(f"egress sentinel recorded no provider calls ({hits_file} absent): zero egress")
|
||||
return 0
|
||||
hits = [line for line in hits_file.read_text(encoding="utf-8").splitlines() if line.strip()]
|
||||
if not hits:
|
||||
print("egress sentinel recorded no provider calls: zero egress")
|
||||
return 0
|
||||
print(f"egress sentinel recorded {len(hits)} provider call(s); replay was not hermetic:", file=sys.stderr)
|
||||
for line in hits:
|
||||
print(f" {line}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
def _serve_from_args(args: argparse.Namespace) -> int:
|
||||
config = ServeConfig(
|
||||
hosts=tuple(args.host),
|
||||
sink_address=args.sink_address,
|
||||
ports=tuple(args.port),
|
||||
hits_file=Path(args.hits_file),
|
||||
hosts_file=Path(args.hosts_file),
|
||||
ready_file=Path(args.ready_file) if args.ready_file else None,
|
||||
pid_file=Path(args.pid_file) if args.pid_file else None,
|
||||
)
|
||||
return serve(config)
|
||||
|
||||
|
||||
def main(argv: tuple[str, ...]) -> int:
|
||||
parser = argparse.ArgumentParser(description="count outbound provider calls during an e2e replay")
|
||||
sub = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
serve_parser = sub.add_parser("serve", help="pin provider hosts and count connection attempts")
|
||||
serve_parser.add_argument("--host", action="append", required=True, help="provider host to pin and watch")
|
||||
serve_parser.add_argument("--sink-address", default="127.0.0.1")
|
||||
serve_parser.add_argument("--port", action="append", type=int, default=None)
|
||||
serve_parser.add_argument("--hits-file", required=True)
|
||||
serve_parser.add_argument("--hosts-file", default="/etc/hosts")
|
||||
serve_parser.add_argument("--ready-file", default=None)
|
||||
serve_parser.add_argument("--pid-file", default=None)
|
||||
|
||||
assert_parser = sub.add_parser("assert-empty", help="exit non-zero if any provider call was recorded")
|
||||
assert_parser.add_argument("--hits-file", required=True)
|
||||
|
||||
args = parser.parse_args(argv)
|
||||
if args.command == "serve":
|
||||
if args.port is None:
|
||||
args.port = [443]
|
||||
return _serve_from_args(args)
|
||||
return assert_empty(Path(args.hits_file))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main(tuple(sys.argv[1:])))
|
||||
55
.github/scripts/e2e_fetch_fixture_bundle.sh
vendored
Executable file
55
.github/scripts/e2e_fetch_fixture_bundle.sh
vendored
Executable file
|
|
@ -0,0 +1,55 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
REPO="${1:-${GITHUB_REPOSITORY:?REPO required}}"
|
||||
ARTIFACT_NAME="${2:-e2e-fixtures-bundle}"
|
||||
BASE_BRANCH="${3:?base branch required}"
|
||||
DEST_DIR="${4:?destination bundle dir required}"
|
||||
|
||||
: "${GH_TOKEN:?GH_TOKEN required to query and download artifacts}"
|
||||
|
||||
WORKDIR="$(mktemp -d)"
|
||||
trap 'rm -rf "${WORKDIR}"' EXIT
|
||||
|
||||
echo "resolving newest non-expired '${ARTIFACT_NAME}' artifact on ${REPO}@${BASE_BRANCH}"
|
||||
|
||||
SELECTED="$(
|
||||
gh api "repos/${REPO}/actions/artifacts" -X GET -f per_page=100 --paginate \
|
||||
--jq ".artifacts[] | select(.name == \"${ARTIFACT_NAME}\" and .expired == false and .workflow_run.head_branch == \"${BASE_BRANCH}\") | {id, digest, created_at, run_id: .workflow_run.id, run_number: .workflow_run.run_number}" \
|
||||
| jq -s 'sort_by(.created_at) | reverse | .[0] // empty'
|
||||
)"
|
||||
|
||||
if [[ -z "${SELECTED}" ]]; then
|
||||
echo "no usable '${ARTIFACT_NAME}' artifact on ${BASE_BRANCH}: the last record run produced none (a red Saturday), so there is nothing fresh to replay; failing loudly instead of replaying a stale bundle" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
RUN_ID="$(echo "${SELECTED}" | jq -r '.run_id')"
|
||||
RUN_NUMBER="$(echo "${SELECTED}" | jq -r '.run_number')"
|
||||
ARTIFACT_ID="$(echo "${SELECTED}" | jq -r '.id')"
|
||||
GH_DIGEST="$(echo "${SELECTED}" | jq -r '.digest // "unknown"')"
|
||||
CREATED_AT="$(echo "${SELECTED}" | jq -r '.created_at')"
|
||||
|
||||
echo "pinned bundle: run #${RUN_NUMBER} (run_id=${RUN_ID}, artifact_id=${ARTIFACT_ID}), recorded ${CREATED_AT}, github digest ${GH_DIGEST}"
|
||||
|
||||
gh run download "${RUN_ID}" --repo "${REPO}" -n "${ARTIFACT_NAME}" -D "${WORKDIR}"
|
||||
|
||||
TARBALL="$(find "${WORKDIR}" -name '*.tar.gz' -type f | head -n 1)"
|
||||
if [[ -z "${TARBALL}" ]]; then
|
||||
echo "downloaded artifact contained no tarball" >&2
|
||||
exit 1
|
||||
fi
|
||||
SIDECAR="${TARBALL}.sha256"
|
||||
if [[ ! -f "${SIDECAR}" ]]; then
|
||||
echo "downloaded artifact has no ${SIDECAR}: cannot verify the bundle digest" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "verifying bundle against its recorded sha256 digest"
|
||||
( cd "$(dirname "${TARBALL}")" && sha256sum -c "$(basename "${SIDECAR}")" )
|
||||
|
||||
mkdir -p "${DEST_DIR}"
|
||||
tar xzf "${TARBALL}" -C "${DEST_DIR}"
|
||||
|
||||
echo "extracted bundle into ${DEST_DIR}"
|
||||
python3 -c "import json,sys; m=json.load(open(sys.argv[1])); print(' recorded_at', m['recorded_at'], 'harness', m['harness_version'], 'format_version', m['format_version'])" "${DEST_DIR}/manifest.json"
|
||||
36
.github/scripts/e2e_pack_fixture_bundle.sh
vendored
Executable file
36
.github/scripts/e2e_pack_fixture_bundle.sh
vendored
Executable file
|
|
@ -0,0 +1,36 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -ne 2 ]]; then
|
||||
echo "usage: $0 <bundle-dir> <out-tarball>" >&2
|
||||
exit 2
|
||||
fi
|
||||
|
||||
BUNDLE_DIR="$1"
|
||||
OUT_TARBALL="$2"
|
||||
|
||||
MANIFEST="${BUNDLE_DIR}/manifest.json"
|
||||
if [[ ! -f "${MANIFEST}" ]]; then
|
||||
echo "no ${MANIFEST}: refusing to publish a bundle with no manifest (record produced nothing)" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "packing fixture bundle from ${BUNDLE_DIR}"
|
||||
python3 -c "import json,sys; m=json.load(open(sys.argv[1])); print(' format_version', m['format_version'], 'recorded_at', m['recorded_at'], 'harness', m['harness_version'])" "${MANIFEST}"
|
||||
|
||||
TEST_DIRS=$(find "${BUNDLE_DIR}" -mindepth 1 -maxdepth 1 -type d | wc -l | tr -d ' ')
|
||||
if [[ "${TEST_DIRS}" -eq 0 ]]; then
|
||||
echo "bundle at ${BUNDLE_DIR} has a manifest but no recorded interactions; refusing to publish an empty bundle" >&2
|
||||
exit 1
|
||||
fi
|
||||
echo " ${TEST_DIRS} recorded test director(ies)"
|
||||
|
||||
mkdir -p "$(dirname "${OUT_TARBALL}")"
|
||||
tar czf "${OUT_TARBALL}" -C "${BUNDLE_DIR}" .
|
||||
|
||||
OUT_DIR="$(cd "$(dirname "${OUT_TARBALL}")" && pwd)"
|
||||
OUT_BASE="$(basename "${OUT_TARBALL}")"
|
||||
( cd "${OUT_DIR}" && sha256sum "${OUT_BASE}" > "${OUT_BASE}.sha256" )
|
||||
|
||||
echo "wrote ${OUT_TARBALL} ($(du -h "${OUT_TARBALL}" | cut -f1)) and ${OUT_BASE}.sha256"
|
||||
cat "${OUT_DIR}/${OUT_BASE}.sha256"
|
||||
13
.github/workflows/_test-unit-base.yml
vendored
13
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -27,7 +27,7 @@ on:
|
|||
default: 20
|
||||
job-timeout-minutes:
|
||||
description: >-
|
||||
Backstop for the whole job. Keep it >= `timeout-minutes` plus 35: 30 for
|
||||
Backstop for the whole job. Keep it >= `timeout-minutes` plus 40: 35 for
|
||||
the per-step ceilings on the setup steps below, and 5 for the runner
|
||||
overhead the job clock charges but no step owns (job init, step
|
||||
transitions, post-job cleanup). That headroom is what makes the test
|
||||
|
|
@ -36,7 +36,7 @@ on:
|
|||
arithmetic, so the sum is passed in rather than computed.
|
||||
required: false
|
||||
type: number
|
||||
default: 55
|
||||
default: 60
|
||||
max-failures:
|
||||
description: "Stop after this many failures"
|
||||
required: false
|
||||
|
|
@ -103,6 +103,11 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 8
|
||||
|
|
@ -144,7 +149,7 @@ jobs:
|
|||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
else
|
||||
|
|
@ -156,7 +161,7 @@ jobs:
|
|||
--reruns-delay 1 \
|
||||
--dist="${DIST}" \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov=./litellm --cov=./enterprise/litellm_enterprise \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
fi
|
||||
|
|
|
|||
4
.github/workflows/check-ui-api-types.yml
vendored
4
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -67,6 +67,10 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install backend dependencies
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
|
|
|||
237
.github/workflows/e2e_record_replay.yml
vendored
Normal file
237
.github/workflows/e2e_record_replay.yml
vendored
Normal file
|
|
@ -0,0 +1,237 @@
|
|||
name: "E2E Record and Replay"
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 8 * * 6"
|
||||
- cron: "0 8 * * 1-5"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
mode:
|
||||
description: "record (hits real providers and publishes a fresh bundle) or replay (bundle only, zero provider egress)"
|
||||
type: choice
|
||||
options:
|
||||
- record
|
||||
- replay
|
||||
default: record
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
record:
|
||||
name: "Record the e2e suite against real providers"
|
||||
if: >-
|
||||
(github.event_name != 'schedule' || github.repository == 'BerriAI/litellm') &&
|
||||
(github.event.schedule == '0 8 * * 6' ||
|
||||
(github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'record'))
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16.6
|
||||
env:
|
||||
POSTGRES_USER: llmproxy
|
||||
POSTGRES_PASSWORD: dbpassword9090
|
||||
POSTGRES_DB: litellm
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U llmproxy"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
env:
|
||||
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
|
||||
LITELLM_MASTER_KEY: sk-e2e-record-replay
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
|
||||
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
|
||||
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: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Start the proxy
|
||||
run: |
|
||||
nohup uv run --no-sync litellm --config tests/e2e/gateway/record_replay_ci_config.yml --port 4000 > proxy.log 2>&1 &
|
||||
for _ in $(seq 1 90); do
|
||||
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "proxy never became live"
|
||||
tail -n 100 proxy.log
|
||||
exit 1
|
||||
|
||||
- name: Record the replayable e2e lane
|
||||
env:
|
||||
E2E_FIXTURE_MODE: record
|
||||
run: |
|
||||
uv run --no-sync pytest tests/e2e -m replayable --reruns 0 -v --tb=short -rA
|
||||
|
||||
- name: Pack the fixture bundle
|
||||
run: |
|
||||
.github/scripts/e2e_pack_fixture_bundle.sh tests/e2e/.fixtures "${RUNNER_TEMP}/bundle/e2e-fixtures.tar.gz"
|
||||
|
||||
- name: Publish the fixture bundle
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: e2e-fixtures-bundle
|
||||
path: |
|
||||
${{ runner.temp }}/bundle/e2e-fixtures.tar.gz
|
||||
${{ runner.temp }}/bundle/e2e-fixtures.tar.gz.sha256
|
||||
if-no-files-found: error
|
||||
retention-days: 30
|
||||
|
||||
- name: Show proxy log on failure
|
||||
if: failure()
|
||||
run: tail -n 300 proxy.log
|
||||
|
||||
replay:
|
||||
name: "Replay the e2e suite from the pinned bundle with zero egress"
|
||||
if: >-
|
||||
(github.event_name != 'schedule' || github.repository == 'BerriAI/litellm') &&
|
||||
(github.event.schedule == '0 8 * * 1-5' ||
|
||||
(github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'replay'))
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 45
|
||||
permissions:
|
||||
contents: read
|
||||
actions: read
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16.6
|
||||
env:
|
||||
POSTGRES_USER: llmproxy
|
||||
POSTGRES_PASSWORD: dbpassword9090
|
||||
POSTGRES_DB: litellm
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U llmproxy"
|
||||
--health-interval 5s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
env:
|
||||
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
|
||||
LITELLM_MASTER_KEY: sk-e2e-record-replay
|
||||
LITELLM_LOCAL_MODEL_COST_MAP: "True"
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
OPENAI_API_KEY: sk-replay-must-never-reach-a-provider
|
||||
ANTHROPIC_API_KEY: sk-ant-replay-must-never-reach-a-provider
|
||||
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: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Fetch the pinned fixture bundle by digest
|
||||
env:
|
||||
BASE_BRANCH: ${{ github.ref_name }}
|
||||
run: |
|
||||
.github/scripts/e2e_fetch_fixture_bundle.sh \
|
||||
"${GITHUB_REPOSITORY}" \
|
||||
e2e-fixtures-bundle \
|
||||
"${BASE_BRANCH}" \
|
||||
tests/e2e/.fixtures
|
||||
|
||||
- name: Start the proxy
|
||||
run: |
|
||||
nohup uv run --no-sync litellm --config tests/e2e/gateway/record_replay_ci_config.yml --port 4000 > proxy.log 2>&1 &
|
||||
for _ in $(seq 1 90); do
|
||||
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "proxy never became live"
|
||||
tail -n 100 proxy.log
|
||||
exit 1
|
||||
|
||||
- name: Start the egress sentinel
|
||||
run: |
|
||||
# shellcheck disable=SC2024 # the log redirect is deliberately the runner user's, so a later non-sudo cat can read it
|
||||
sudo python3 .github/scripts/e2e_egress_sentinel.py serve \
|
||||
--host api.openai.com \
|
||||
--host api.anthropic.com \
|
||||
--hits-file "${RUNNER_TEMP}/egress-hits.jsonl" \
|
||||
--ready-file "${RUNNER_TEMP}/egress-ready" \
|
||||
--pid-file "${RUNNER_TEMP}/egress.pid" \
|
||||
> "${RUNNER_TEMP}/egress-sentinel.log" 2>&1 &
|
||||
for _ in $(seq 1 30); do
|
||||
if [[ -f "${RUNNER_TEMP}/egress-ready" ]]; then
|
||||
cat "${RUNNER_TEMP}/egress-sentinel.log"
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "egress sentinel never became ready"
|
||||
cat "${RUNNER_TEMP}/egress-sentinel.log"
|
||||
exit 1
|
||||
|
||||
- name: Replay the replayable e2e lane
|
||||
env:
|
||||
E2E_FIXTURE_MODE: replay
|
||||
run: |
|
||||
uv run --no-sync pytest tests/e2e -m replayable --reruns 0 -v --tb=short -rA
|
||||
|
||||
- name: Stop the egress sentinel and assert zero provider egress
|
||||
if: always()
|
||||
run: |
|
||||
if [[ -f "${RUNNER_TEMP}/egress.pid" ]]; then
|
||||
sudo kill -TERM "$(cat "${RUNNER_TEMP}/egress.pid")" 2>/dev/null || true
|
||||
sleep 2
|
||||
fi
|
||||
python3 .github/scripts/e2e_egress_sentinel.py assert-empty --hits-file "${RUNNER_TEMP}/egress-hits.jsonl"
|
||||
|
||||
- name: Show proxy log on failure
|
||||
if: failure()
|
||||
run: tail -n 300 proxy.log
|
||||
33
.github/workflows/image-scan.yml
vendored
33
.github/workflows/image-scan.yml
vendored
|
|
@ -17,10 +17,14 @@ on:
|
|||
- backend/Dockerfile
|
||||
- backend/main.py
|
||||
- docker/component_entrypoint.sh
|
||||
- docker/entrypoint.sh
|
||||
- litellm/proxy/prisma_migration.py
|
||||
- litellm-proxy-extras/**
|
||||
- tests/proxy_migration_tests/**
|
||||
- uv.lock
|
||||
- ui/litellm-dashboard/package-lock.json
|
||||
- ui/Dockerfile
|
||||
- ui/nginx.conf
|
||||
- .github/workflows/image-scan.yml
|
||||
schedule:
|
||||
- cron: "41 6 * * *"
|
||||
|
|
@ -183,6 +187,35 @@ jobs:
|
|||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
|
||||
|
||||
ui-image:
|
||||
name: ui-image
|
||||
runs-on: ubuntu-latest
|
||||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Build UI image
|
||||
run: docker build -f ui/Dockerfile -t litellm-ui-scan:${{ github.sha }} .
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Verify the UI serves offline as an arbitrary uid with a read-only root fs
|
||||
env:
|
||||
LITELLM_IMAGE: litellm-ui-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_ui_image_serves_offline.py -v
|
||||
|
||||
backend-image:
|
||||
name: backend-image
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
13
.github/workflows/mutation-test.yml
vendored
13
.github/workflows/mutation-test.yml
vendored
|
|
@ -53,6 +53,9 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
|
||||
|
|
@ -84,11 +87,20 @@ jobs:
|
|||
run: |
|
||||
uv pip uninstall pytest-retry || true
|
||||
|
||||
# Ends before the job's own deadline so a run that outlasts the budget is
|
||||
# still followed by the report and upload steps. mutmut saves after every
|
||||
# mutant result, to mutants/<source path>.meta, so an interrupted run
|
||||
# still scores the mutants it finished and export-cicd-stats can read
|
||||
# them; a cancelled job skips those steps and publishes nothing at all.
|
||||
- name: Run mutmut
|
||||
timeout-minutes: 300
|
||||
env:
|
||||
# Make the mutants/ sandbox win over site-packages on sys.path so the
|
||||
# trampolined files are imported instead of the installed copy.
|
||||
PYTHONPATH: ${{ github.workspace }}/mutants
|
||||
# Without this mutmut finds no covered lines and generates 0 mutants.
|
||||
# See the file itself for why.
|
||||
COVERAGE_RCFILE: ${{ github.workspace }}/.github/mutmut-coverage.rc
|
||||
run: |
|
||||
set -o pipefail
|
||||
mkdir -p mutants
|
||||
|
|
@ -127,6 +139,7 @@ jobs:
|
|||
mutmut-run.log
|
||||
mutants/mutmut-stats.json
|
||||
mutants/mutmut-cicd-stats.json
|
||||
mutants/**/*.meta
|
||||
mutants/litellm/proxy/management_endpoints/**/*.py
|
||||
if-no-files-found: warn
|
||||
retention-days: 14
|
||||
|
|
|
|||
|
|
@ -43,6 +43,9 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
|
|
|
|||
6
.github/workflows/test-code-quality.yml
vendored
6
.github/workflows/test-code-quality.yml
vendored
|
|
@ -56,6 +56,9 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: uv sync --frozen --all-groups --all-extras
|
||||
|
||||
|
|
@ -128,6 +131,9 @@ jobs:
|
|||
- name: check_e2e_no_raw_requests
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
|
||||
|
||||
- name: check_migrations_no_data_rewrites
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
|
||||
|
||||
- name: memory_test
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py
|
||||
|
||||
|
|
|
|||
15
.github/workflows/test-linting.yml
vendored
15
.github/workflows/test-linting.yml
vendored
|
|
@ -67,6 +67,17 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-lint-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-lint-
|
||||
|
||||
- name: Clean Python cache
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
|
|
@ -78,6 +89,10 @@ jobs:
|
|||
run: |
|
||||
uv lock --check || (echo "❌ uv.lock is out of sync with pyproject.toml. Run 'uv lock' locally and commit the result." && exit 1)
|
||||
|
||||
- name: Cache the Rust build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
|
|
|
|||
6
.github/workflows/test-mcp.yml
vendored
6
.github/workflows/test-mcp.yml
vendored
|
|
@ -47,6 +47,10 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
|
|
@ -56,4 +60,4 @@ jobs:
|
|||
- name: Run MCP tests
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov-report=xml --durations=5
|
||||
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml --durations=5
|
||||
|
|
|
|||
145
.github/workflows/test-postgres.yml
vendored
Normal file
145
.github/workflows/test-postgres.yml
vendored
Normal file
|
|
@ -0,0 +1,145 @@
|
|||
name: "Postgres Tests"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
postgres:
|
||||
name: ${{ matrix.shard }}
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: ${{ matrix.job-timeout-minutes }}
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
env:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: litellm_test
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd pg_isready
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 10
|
||||
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- shard: proxy-behavior
|
||||
test-path: "tests/proxy_behavior"
|
||||
seed: db-push
|
||||
workers: 0
|
||||
timeout-minutes: 25
|
||||
job-timeout-minutes: 50
|
||||
|
||||
- shard: proxy-security
|
||||
test-path: "tests/proxy_security_tests"
|
||||
seed: db-push
|
||||
workers: 0
|
||||
timeout-minutes: 15
|
||||
job-timeout-minutes: 40
|
||||
|
||||
- shard: schema-migration
|
||||
test-path: "tests/proxy_migration_tests"
|
||||
seed: none
|
||||
workers: 0
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 45
|
||||
|
||||
env:
|
||||
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
timeout-minutes: 3
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect relevant changes
|
||||
id: changes
|
||||
timeout-minutes: 2
|
||||
uses: ./.github/actions/detect-changes
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-postgres-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-postgres-
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 12
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 5
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Seed database schema
|
||||
if: steps.changes.outputs.decision != 'skip' && matrix.seed != 'none'
|
||||
timeout-minutes: 10
|
||||
run: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
|
||||
- name: Run tests
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: ${{ matrix.timeout-minutes }}
|
||||
env:
|
||||
TEST_PATH: ${{ matrix.test-path }}
|
||||
WORKERS: ${{ matrix.workers }}
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10
|
||||
else
|
||||
uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10 -n "${WORKERS}"
|
||||
fi
|
||||
|
|
@ -88,6 +88,9 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
|
|
|||
|
|
@ -67,6 +67,10 @@ jobs:
|
|||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Cache the Rust build
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: |
|
||||
|
|
|
|||
55
.github/workflows/test-unit.yml
vendored
55
.github/workflows/test-unit.yml
vendored
|
|
@ -55,7 +55,7 @@ jobs:
|
|||
workers: 2
|
||||
reruns: 1
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: enterprise-routing
|
||||
artifact-name: enterprise-routing
|
||||
|
|
@ -67,7 +67,7 @@ jobs:
|
|||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: integrations
|
||||
artifact-name: integrations
|
||||
|
|
@ -75,7 +75,7 @@ jobs:
|
|||
workers: 2
|
||||
reruns: 3
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: Vertex AI
|
||||
artifact-name: llm-vertex-ai
|
||||
|
|
@ -83,7 +83,7 @@ jobs:
|
|||
workers: 1
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: All Other Providers
|
||||
artifact-name: llm-other-providers
|
||||
|
|
@ -91,7 +91,7 @@ jobs:
|
|||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: misc
|
||||
artifact-name: misc
|
||||
|
|
@ -113,6 +113,7 @@ jobs:
|
|||
tests/test_litellm/rag
|
||||
tests/test_litellm/realtime_api
|
||||
tests/test_litellm/rerank_api
|
||||
tests/test_litellm/rust_bridge
|
||||
tests/test_litellm/sandbox
|
||||
tests/test_litellm/test_router
|
||||
tests/test_litellm/vector_stores
|
||||
|
|
@ -121,7 +122,7 @@ jobs:
|
|||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: proxy-auth
|
||||
artifact-name: proxy-auth
|
||||
|
|
@ -133,7 +134,7 @@ jobs:
|
|||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: proxy-endpoints
|
||||
artifact-name: proxy-endpoints
|
||||
|
|
@ -163,14 +164,15 @@ jobs:
|
|||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/rerank_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
tests/test_litellm/proxy/config_resolvers
|
||||
tests/test_litellm/proxy/utils
|
||||
workers: 2
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: proxy-server
|
||||
artifact-name: proxy-server
|
||||
|
|
@ -178,7 +180,7 @@ jobs:
|
|||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 60
|
||||
job-timeout-minutes: 95
|
||||
job-timeout-minutes: 100
|
||||
|
||||
- shard: proxy-infra
|
||||
artifact-name: proxy-infra
|
||||
|
|
@ -194,10 +196,39 @@ jobs:
|
|||
tests/test_litellm/proxy/types_utils
|
||||
tests/test_litellm/proxy/logging_endpoints
|
||||
tests/test_litellm/proxy/test_*.py
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: caching-local
|
||||
artifact-name: caching-local
|
||||
test-path: >-
|
||||
tests/local_testing/test_cache_preset_key.py
|
||||
tests/local_testing/test_caching_handler.py
|
||||
tests/local_testing/test_prompt_caching.py
|
||||
tests/local_testing/test_responses_stream_cache_keys.py
|
||||
tests/local_testing/test_unit_test_caching.py
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: proxy-extras
|
||||
artifact-name: proxy-extras
|
||||
test-path: "tests/litellm-proxy-extras"
|
||||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: enterprise-package
|
||||
artifact-name: enterprise-package
|
||||
test-path: "tests/enterprise"
|
||||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 60
|
||||
|
||||
- shard: responses-caching-types
|
||||
artifact-name: responses-caching-types
|
||||
|
|
@ -208,7 +239,7 @@ jobs:
|
|||
workers: 2
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
job-timeout-minutes: 55
|
||||
job-timeout-minutes: 60
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: ${{ matrix.test-path }}
|
||||
|
|
|
|||
3
.github/workflows/weekly_load_anomaly.yml
vendored
3
.github/workflows/weekly_load_anomaly.yml
vendored
|
|
@ -47,6 +47,9 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache the Rust build
|
||||
uses: ./.github/actions/cache-cargo-build
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
|
|
|||
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -3,6 +3,8 @@
|
|||
tests/e2e/.fixtures/
|
||||
.venv-typecheck
|
||||
.venv_policy_test
|
||||
.venv-mutmut
|
||||
mutants/
|
||||
.env
|
||||
.claude
|
||||
CLAUDE.local.md
|
||||
|
|
|
|||
|
|
@ -37,13 +37,14 @@ If you're resolving a linear ticket, in the "## Linear ticket" section of the PR
|
|||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
|
||||
If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
|
||||
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
|
||||
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
|
||||
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
|
||||
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
|
||||
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
|
||||
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
|
||||
|
||||
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
|
||||
|
||||
|
|
@ -65,6 +66,8 @@ Commit and push your work when you're done without asking
|
|||
|
||||
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
|
||||
|
||||
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
|
||||
|
||||
If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
|
||||
|
||||
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
|
||||
|
|
@ -79,6 +82,8 @@ Do not put names of customers or customer company names in code, PR descriptions
|
|||
|
||||
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
|
||||
|
||||
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
|
||||
|
||||
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
|
||||
|
||||
- Composition over inheritance
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
|
|
|
|||
|
|
@ -292,6 +292,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
|
|||
| [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Codestral (`codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Cognition (`cognition`)](https://docs.litellm.ai/docs/providers/cognition) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Cohere (`cohere`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ |
|
||||
| [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | |
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
|
|
@ -1,18 +1,18 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 19955
|
||||
"limit": 18483
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2566
|
||||
"limit": 2564
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 320
|
||||
},
|
||||
"reportAttributeAccessIssue": {
|
||||
"limit": 488
|
||||
"limit": 483
|
||||
},
|
||||
"reportCallIssue": {
|
||||
"limit": 114
|
||||
"limit": 113
|
||||
},
|
||||
"reportConstantRedefinition": {
|
||||
"limit": 40
|
||||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 6049
|
||||
"limit": 5960
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
|
|
@ -45,7 +45,7 @@
|
|||
"limit": 35
|
||||
},
|
||||
"reportInvalidTypeForm": {
|
||||
"limit": 35
|
||||
"limit": 34
|
||||
},
|
||||
"reportInvalidTypeVarUse": {
|
||||
"limit": 2
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5663
|
||||
"limit": 5659
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15555
|
||||
"limit": 15484
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -72,7 +72,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"limit": 1061
|
||||
"limit": 1058
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"limit": 0
|
||||
|
|
@ -84,7 +84,7 @@
|
|||
"limit": 56
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 1823
|
||||
"limit": 1810
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 8
|
||||
|
|
@ -99,31 +99,31 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44655
|
||||
"limit": 44530
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 39017
|
||||
"limit": 38808
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19885
|
||||
"limit": 19829
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30572
|
||||
"limit": 30356
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 699
|
||||
"limit": 697
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 836
|
||||
"limit": 833
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 0
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ EXTRA_BOOLEAN_KEYS = frozenset(
|
|||
"uses_embed_content",
|
||||
"use_openai_responses_path",
|
||||
"bedrock_converse_supports_strict_tools",
|
||||
"thinking_always_on",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base images
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG PROXY_EXTRAS_SOURCE=published
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from fastapi import HTTPException
|
|||
|
||||
|
||||
class _ENTERPRISE_BannedKeywords(CustomLogger):
|
||||
enforces_request_content: bool = True
|
||||
# Class variables or attributes
|
||||
def __init__(self):
|
||||
banned_keywords_list = litellm.banned_keywords_list
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from fastapi import HTTPException
|
|||
|
||||
|
||||
class _ENTERPRISE_BlockedUserList(CustomLogger):
|
||||
enforces_request_content: bool = True
|
||||
# Class variables or attributes
|
||||
def __init__(self, prisma_client: Optional[PrismaClient]):
|
||||
self.prisma_client = prisma_client
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ from litellm.constants import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import LiteLLM_ManagedObjectTable
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
|
@ -255,6 +257,52 @@ class CheckBatchCost:
|
|||
"so it will no longer be polled"
|
||||
)
|
||||
|
||||
async def _claim_job_for_costing(self, job: "LiteLLM_ManagedObjectTable") -> bool:
|
||||
"""
|
||||
Atomically flip batch_processed from false to true, returning whether this pod won
|
||||
the row. Every pod and uvicorn worker schedules its own poller against the shared
|
||||
table, so without this compare-and-swap two of them can select the same completed
|
||||
batch in one window and both emit an aretrieve_batch spend log for it. Schemas
|
||||
without the column can't be claimed, so they keep the pre-existing behavior.
|
||||
|
||||
Called immediately before the spend log is written rather than before the results
|
||||
fetch, because batch_processed is also what holds off deletion of the files that
|
||||
fetch reads and what keeps an unbilled row selectable by the next poll cycle.
|
||||
"""
|
||||
if not self._has_batch_processed_column:
|
||||
return True
|
||||
try:
|
||||
claimed: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={"id": job.id, "batch_processed": False},
|
||||
data={"batch_processed": True},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to claim job {job.id} for cost tracking: {db_err}"
|
||||
)
|
||||
return False
|
||||
return claimed > 0
|
||||
|
||||
async def _release_job_claim(self, job: "LiteLLM_ManagedObjectTable") -> None:
|
||||
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
|
||||
|
||||
Safe to match on batch_processed=True: while this poller is active the retrieve
|
||||
path leaves the column alone (batch_cost_poller_is_active), so a true value here
|
||||
is always this pod's own claim.
|
||||
"""
|
||||
if not self._has_batch_processed_column:
|
||||
return
|
||||
try:
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={"id": job.id, "batch_processed": True},
|
||||
data={"batch_processed": False},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to release the claim on job {job.id}, "
|
||||
f"so its cost will not be retried: {db_err}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
|
||||
"""A unified id that decodes but carries no model_id can never be routed."""
|
||||
|
|
@ -305,7 +353,7 @@ class CheckBatchCost:
|
|||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
|
||||
|
||||
async def _finalize_unbilled_terminal_job(
|
||||
self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
|
||||
self, job: "prisma_models.LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
|
||||
) -> None:
|
||||
"""Persist a terminal batch that has nothing billable, converting any raw
|
||||
provider file ids to managed ids, and take it out of the poll page."""
|
||||
|
|
@ -572,9 +620,10 @@ class CheckBatchCost:
|
|||
"""
|
||||
Fetch a completed batch's results, compute cost/usage, and emit the
|
||||
aretrieve_batch spend log. Returns (model_name, llm_provider) on
|
||||
success, None when the job can't be routed to a deployment. Raises on
|
||||
results-fetch or cost-computation failures so the caller can leave the
|
||||
job unprocessed and retry it on a later poll.
|
||||
success, None when the job can't be routed to a deployment or when
|
||||
another pod claimed it. Raises on results-fetch or cost-computation
|
||||
failures so the caller can leave the job unprocessed and retry it on a
|
||||
later poll.
|
||||
"""
|
||||
from litellm.batches.batch_utils import (
|
||||
_get_file_content_as_dictionary,
|
||||
|
|
@ -743,12 +792,23 @@ class CheckBatchCost:
|
|||
optional_params={},
|
||||
)
|
||||
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_cost,
|
||||
batch_usage=batch_usage,
|
||||
batch_models=batch_models,
|
||||
)
|
||||
if not await self._claim_job_for_costing(job):
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
|
||||
"in this window, so its cost is already being tracked there"
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_cost,
|
||||
batch_usage=batch_usage,
|
||||
batch_models=batch_models,
|
||||
)
|
||||
except Exception:
|
||||
await self._release_job_claim(job)
|
||||
raise
|
||||
|
||||
# Record batch duration (completed_at - created_at)
|
||||
if prom_logger and response.completed_at and response.created_at:
|
||||
|
|
@ -908,6 +968,16 @@ class CheckBatchCost:
|
|||
)
|
||||
|
||||
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_completed_batch_safe_to_retire,
|
||||
)
|
||||
|
||||
if response.status in ("completed", "complete") and not _completed_batch_safe_to_retire(response):
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: batch {batch_id} is completed but its output file id "
|
||||
f"has not appeared yet; leaving job {job.id} for the next poll cycle"
|
||||
)
|
||||
continue
|
||||
await self._finalize_unbilled_terminal_job(job, response)
|
||||
|
||||
# Record polling run metrics (always, even if nothing was processed)
|
||||
|
|
|
|||
|
|
@ -45,6 +45,8 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
FILE_LIST_CONTINUATION_CHUNK_SIZE,
|
||||
MAX_FILE_LIST_LIMIT,
|
||||
_is_base64_encoded_unified_file_id,
|
||||
apply_unified_file_ids,
|
||||
ensure_batch_response_managed_file_ids,
|
||||
|
|
@ -54,6 +56,8 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
map_raw_file_ids_to_unified,
|
||||
normalize_mime_type_for_provider,
|
||||
resolve_managed_output_file_model_name,
|
||||
validate_file_list_limit,
|
||||
validate_file_list_purpose,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
|
||||
request_tags_from_metadata,
|
||||
|
|
@ -63,9 +67,9 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess
|
|||
AsyncCursorPage,
|
||||
ChatCompletionFileObject,
|
||||
CreateFileRequest,
|
||||
FileListPage,
|
||||
FileObject,
|
||||
OpenAIFileObject,
|
||||
OpenAIFilesPurpose,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -144,7 +148,14 @@ class _ManagedFileRow(Protocol):
|
|||
class _ManagedFileTableActions(Protocol):
|
||||
async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ...
|
||||
|
||||
async def find_many(self, where: Mapping[str, object]) -> Sequence[_ManagedFileRow]: ...
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
take: int = ...,
|
||||
order: Union[Mapping[str, str], Sequence[Mapping[str, str]]] = ...,
|
||||
cursor: Mapping[str, str] = ...,
|
||||
skip: int = ...,
|
||||
) -> Sequence[_ManagedFileRow]: ...
|
||||
|
||||
async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ...
|
||||
|
||||
|
|
@ -1365,12 +1376,76 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
|
||||
async def afile_list(
|
||||
self,
|
||||
purpose: Optional[OpenAIFilesPurpose],
|
||||
purpose: Optional[str],
|
||||
litellm_parent_otel_span: Optional[Span],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
limit: Optional[int] = None,
|
||||
after: Optional[str] = None,
|
||||
**data: Dict,
|
||||
) -> List[OpenAIFileObject]:
|
||||
"""Handled in files_endpoints.py"""
|
||||
return []
|
||||
) -> FileListPage:
|
||||
"""List the managed files the caller owns, newest first.
|
||||
|
||||
Pagination is keyset based on ``unified_file_id`` so a key that owns
|
||||
every file on the proxy still reads one bounded page at a time.
|
||||
``purpose`` is applied after parsing, because the managed file table
|
||||
keeps it inside the ``file_object`` blob instead of a column, and rows
|
||||
whose blob will not parse drop out there too, so a chunk of rows can
|
||||
yield fewer matches than the page holds. Successive chunks are read
|
||||
until the page is full or the caller's rows run out, which keeps
|
||||
``data`` non-empty while matches remain and its last id usable as the
|
||||
next cursor. A first chunk that fills the page costs one query; once a
|
||||
scan has to continue past it, the chunk widens to
|
||||
``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per
|
||||
that many rows instead of one per page. That bound is per query, not
|
||||
per request: the work is still linear in the rows the caller owns, and
|
||||
a filter matching nothing reads every one of them, with no index
|
||||
covering either the owner filter or the sort.
|
||||
"""
|
||||
validate_file_list_limit(limit)
|
||||
validate_file_list_purpose(purpose)
|
||||
|
||||
owner_filter: Final = build_owner_filter(user_api_key_dict)
|
||||
if owner_filter is None:
|
||||
return FileListPage(**build_list_page([]))
|
||||
|
||||
if after:
|
||||
cursor_row = await _managed_file_table(self.prisma_client).find_first(
|
||||
where={**owner_filter, "unified_file_id": after}
|
||||
)
|
||||
if cursor_row is None:
|
||||
raise ProxyException(
|
||||
message=f"Invalid 'after' cursor: no file found with id '{after}'.",
|
||||
type="invalid_request_error",
|
||||
param="after",
|
||||
code=400,
|
||||
openai_code="invalid_value",
|
||||
)
|
||||
|
||||
page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT)
|
||||
matches: Final[List[OpenAIFileObject]] = []
|
||||
cursor_id = after
|
||||
chunk_size = page_size + 1
|
||||
|
||||
while len(matches) <= page_size:
|
||||
cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {}
|
||||
chunk = await _managed_file_table(self.prisma_client).find_many(
|
||||
where=owner_filter,
|
||||
take=chunk_size,
|
||||
order=[{"created_at": "desc"}, {"unified_file_id": "desc"}],
|
||||
**cursor_args,
|
||||
)
|
||||
matches.extend(
|
||||
parsed_file_object.model_copy(update={"id": row.unified_file_id})
|
||||
for row in chunk
|
||||
if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None
|
||||
and (purpose is None or parsed_file_object.purpose == purpose)
|
||||
)
|
||||
if len(chunk) < chunk_size:
|
||||
break
|
||||
cursor_id = chunk[-1].unified_file_id
|
||||
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
|
||||
|
||||
return FileListPage(**build_list_page(matches[:page_size], has_more=len(matches) > page_size))
|
||||
|
||||
def _is_batch_polling_enabled(self) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.58"
|
||||
version = "0.1.60"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.58"
|
||||
version = "0.1.60"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
|
|
@ -100,6 +100,13 @@ spec:
|
|||
- name: DATABASE_URL
|
||||
value: {{ .Values.db.url | quote }}
|
||||
{{- end }}
|
||||
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
|
||||
- name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Values.db.secret.name }}
|
||||
key: {{ .Values.db.secret.readReplicaEndpointKey }}
|
||||
{{- end }}
|
||||
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
|
||||
- name: DATABASE_URL_READ_REPLICA
|
||||
valueFrom:
|
||||
|
|
|
|||
|
|
@ -80,6 +80,96 @@ tests:
|
|||
secretKeyRef:
|
||||
name: my-secret
|
||||
key: my-key
|
||||
- it: should inject DATABASE_READER_HOST from readReplicaEndpointKey before DATABASE_URL_READ_REPLICA
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db:
|
||||
deployStandalone: false
|
||||
useExisting: true
|
||||
secret:
|
||||
name: postgres
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
readReplicaEndpointKey: reader-host
|
||||
readReplicaUrl: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-host
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require
|
||||
# $(VAR) interpolation only resolves vars defined EARLIER in the env
|
||||
# array, so the reader host must precede the composed URL
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[7].name
|
||||
value: DATABASE_READER_HOST
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[8].name
|
||||
value: DATABASE_URL_READ_REPLICA
|
||||
- it: should omit reader host when readReplicaUrl is unset
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db:
|
||||
deployStandalone: false
|
||||
useExisting: true
|
||||
secret:
|
||||
name: postgres
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
readReplicaEndpointKey: reader-host
|
||||
asserts:
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-host
|
||||
- it: should prefer readReplicaUrlKey over readReplicaEndpointKey composition
|
||||
template: deployment.yaml
|
||||
set:
|
||||
db:
|
||||
useExisting: true
|
||||
secret:
|
||||
name: postgres
|
||||
usernameKey: username
|
||||
passwordKey: password
|
||||
readReplicaUrlKey: reader-url
|
||||
readReplicaEndpointKey: reader-host
|
||||
readReplicaUrl: postgresql://ignored
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-url
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_URL_READ_REPLICA
|
||||
value: postgresql://ignored
|
||||
# the unused reader-host secret ref must be suppressed so a missing
|
||||
# key can't fail pod creation
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_READER_HOST
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: postgres
|
||||
key: reader-host
|
||||
- it: should work with extraEnvVars
|
||||
template: deployment.yaml
|
||||
set:
|
||||
|
|
|
|||
|
|
@ -277,6 +277,14 @@ db:
|
|||
# written to db.readReplicaUrl ends up visible in the rendered pod spec
|
||||
# and the Helm release secret.
|
||||
readReplicaUrlKey: ""
|
||||
# Optional: when set, a DATABASE_READER_HOST env var is sourced from this
|
||||
# secret key, so db.readReplicaUrl can compose the reader URL from
|
||||
# individual secret components, e.g.
|
||||
# postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)
|
||||
# Use this when your secret store holds the bare reader hostname rather
|
||||
# than a full connection URL. Only takes effect when readReplicaUrl is
|
||||
# set; ignored when readReplicaUrlKey is set.
|
||||
readReplicaEndpointKey: ""
|
||||
|
||||
# Optional read-replica routing. When set, the proxy sends read-only
|
||||
# queries (find_*, count, group_by, query_raw/_first) to this URL while
|
||||
|
|
|
|||
|
|
@ -213,18 +213,21 @@ whenever the password contains a URL-reserved character (@, /, ?, %, +,
|
|||
|
||||
When `database.writer.useIAMAuth: true`, the chart injects
|
||||
IAM_TOKEN_DB_AUTH=true and omits DATABASE_PASSWORD — the entrypoint mints
|
||||
the URL from DATABASE_HOST/PORT/USER/NAME plus a short-lived IAM token
|
||||
instead of a static password.
|
||||
the URL from DATABASE_HOST/PORT/USER/NAME plus a short-lived AWS RDS IAM
|
||||
token instead of a static password. `database.writer.useAzureEntraAuth: true`
|
||||
does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token,
|
||||
for Azure Database for PostgreSQL. The two are mutually exclusive.
|
||||
|
||||
The read replica is opt-in via `database.reader.host`. The chart emits
|
||||
DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA /
|
||||
DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth
|
||||
modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for
|
||||
password auth. When `database.reader.useIAMAuth: true` it omits
|
||||
password auth. When `database.reader.useIAMAuth: true` (or
|
||||
`database.reader.useAzureEntraAuth: true`) it omits
|
||||
DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the
|
||||
same way. Reader IAM only takes effect when the writer also uses IAM auth
|
||||
(the proxy gates URL minting on IAM_TOKEN_DB_AUTH, which only the writer
|
||||
sets).
|
||||
same way. Reader token auth only takes effect when the writer uses the same
|
||||
token source, since the proxy gates URL minting on the single global
|
||||
IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
|
||||
*/}}
|
||||
{{- define "litellm.serverEnv" -}}
|
||||
{{- $root := .root -}}
|
||||
|
|
@ -254,9 +257,15 @@ sets).
|
|||
- name: DATABASE_SCHEMA
|
||||
value: {{ .schema | quote }}
|
||||
{{- end }}
|
||||
{{- if and .useIAMAuth .useAzureEntraAuth }}
|
||||
{{- fail "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" }}
|
||||
{{- end }}
|
||||
{{- if .useIAMAuth }}
|
||||
- name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
{{- else if .useAzureEntraAuth }}
|
||||
- name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
{{- else }}
|
||||
- name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
|
|
@ -270,6 +279,9 @@ sets).
|
|||
{{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }}
|
||||
{{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }}
|
||||
{{- end }}
|
||||
{{- if and .useAzureEntraAuth (not $root.Values.database.writer.useAzureEntraAuth) }}
|
||||
{{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }}
|
||||
{{- end }}
|
||||
- name: DATABASE_HOST_READ_REPLICA
|
||||
value: {{ .host | quote }}
|
||||
- name: DATABASE_PORT_READ_REPLICA
|
||||
|
|
@ -280,7 +292,7 @@ sets).
|
|||
- name: DATABASE_SCHEMA_READ_REPLICA
|
||||
value: {{ .schema | quote }}
|
||||
{{- end }}
|
||||
{{- if .useIAMAuth }}
|
||||
{{- if or .useIAMAuth .useAzureEntraAuth }}
|
||||
{{- if .passwordSecret.name }}
|
||||
- name: DATABASE_USER_READ_REPLICA
|
||||
valueFrom:
|
||||
|
|
|
|||
116
helm/litellm/tests/database_auth_tests.yaml
Normal file
116
helm/litellm/tests/database_auth_tests.yaml
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
suite: test database token auth env vars
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- backend/configmap.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: writer emits DATABASE_PASSWORD and no token toggle by default
|
||||
template: gateway/deployment.yaml
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-writer-secret
|
||||
key: password
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
|
||||
- it: writer emits AZURE_POSTGRESQL_AUTH and omits DATABASE_PASSWORD under Entra auth
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: IAM_TOKEN_DB_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
|
||||
- it: backend gets the same Entra toggle as the gateway
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
database.writer.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: AZURE_POSTGRESQL_AUTH
|
||||
value: "true"
|
||||
any: true
|
||||
|
||||
- it: writer rejects both token sources at once
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useIAMAuth: true
|
||||
database.writer.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source"
|
||||
|
||||
- it: reader Entra auth without writer Entra auth is rejected
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.reader.host: reader.example.com
|
||||
database.reader.dbname: litellm
|
||||
database.reader.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)"
|
||||
|
||||
- it: reader under Entra auth omits DATABASE_PASSWORD_READ_REPLICA
|
||||
template: gateway/deployment.yaml
|
||||
set:
|
||||
database.writer.useAzureEntraAuth: true
|
||||
database.reader.host: reader.example.com
|
||||
database.reader.dbname: litellm
|
||||
database.reader.useAzureEntraAuth: true
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_HOST_READ_REPLICA
|
||||
value: reader.example.com
|
||||
any: true
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_USER_READ_REPLICA
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: litellm-reader-secret
|
||||
key: username
|
||||
any: true
|
||||
- notContains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: DATABASE_PASSWORD_READ_REPLICA
|
||||
any: true
|
||||
|
|
@ -145,6 +145,8 @@ database:
|
|||
dbname: ""
|
||||
schema: ""
|
||||
useIAMAuth: false
|
||||
# Azure Database for PostgreSQL with a Microsoft Entra ID token; mutually exclusive with useIAMAuth
|
||||
useAzureEntraAuth: false
|
||||
passwordSecret:
|
||||
name: litellm-writer-secret
|
||||
usernameKey: username
|
||||
|
|
@ -159,6 +161,8 @@ database:
|
|||
dbname: ""
|
||||
schema: ""
|
||||
useIAMAuth: false
|
||||
# Azure Database for PostgreSQL with a Microsoft Entra ID token; mutually exclusive with useIAMAuth
|
||||
useAzureEntraAuth: false
|
||||
passwordSecret:
|
||||
name: litellm-reader-secret
|
||||
usernameKey: username
|
||||
|
|
@ -424,9 +428,11 @@ ui:
|
|||
maxUnavailable: ""
|
||||
podAnnotations: {}
|
||||
# Same shape as the gateway blocks of the same name. The nginx runtime
|
||||
# writes its pid, cache, and proxy temp files under the image's root
|
||||
# filesystem, so `securityContext.readOnlyRootFilesystem: true` here needs
|
||||
# emptyDir volumes mounted over those paths.
|
||||
# writes its pid, cache, and proxy temp files under /tmp, so it boots as
|
||||
# any (arbitrary, non-root) uid; `securityContext.readOnlyRootFilesystem:
|
||||
# true` here needs an emptyDir volume mounted over /tmp. Images before
|
||||
# the /tmp move instead need emptyDirs over /var/cache/nginx and /run to
|
||||
# run as a non-root uid at all.
|
||||
podLabels: {}
|
||||
podSecurityContext: {}
|
||||
securityContext: {}
|
||||
|
|
|
|||
|
|
@ -1,4 +0,0 @@
|
|||
UPDATE "LiteLLM_SpendLogs"
|
||||
SET "created_at" = "endTime",
|
||||
"updated_at" = "endTime"
|
||||
WHERE "created_at" > "endTime" + interval '1 hour';
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN "max_budget" DOUBLE PRECISION;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN "shadow_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
|
||||
|
|
@ -1502,7 +1502,8 @@ model LiteLLM_ShadowEvalJob {
|
|||
baseline_model String? // reverse only: the fixed model the router is judged against
|
||||
judge_model String
|
||||
shadow_percentage Float
|
||||
max_turns Int // this key's sample budget: judge at most this many turns
|
||||
max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
|
||||
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
|
|
@ -1525,6 +1526,7 @@ model LiteLLM_ShadowEvalAttempt {
|
|||
shadow_model String?
|
||||
confidence Float?
|
||||
judge_cost Float @default(0)
|
||||
shadow_cost Float @default(0)
|
||||
error String?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.88"
|
||||
version = "0.4.89"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.88"
|
||||
version = "0.4.89"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -295,6 +295,8 @@ fn core_error_kind(error: &CoreError) -> &'static str {
|
|||
CoreError::Http { .. } => "HttpError",
|
||||
CoreError::InvalidResponse(_) => "InvalidResponse",
|
||||
CoreError::Network(_) => "NetworkError",
|
||||
CoreError::Connect(_) => "ConnectError",
|
||||
CoreError::Routing(_) => "RoutingError",
|
||||
CoreError::Unsupported(_) => "UnsupportedRequest",
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -324,6 +324,8 @@ fn core_error_kind(error: &CoreError) -> &'static str {
|
|||
CoreError::Http { .. } => "HttpError",
|
||||
CoreError::InvalidResponse(_) => "InvalidResponse",
|
||||
CoreError::Network(_) => "NetworkError",
|
||||
CoreError::Connect(_) => "ConnectError",
|
||||
CoreError::Routing(_) => "RoutingError",
|
||||
CoreError::Unsupported(_) => "UnsupportedRequest",
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -105,12 +105,20 @@ impl IntoResponse for MessagesRouteError {
|
|||
),
|
||||
CoreError::Http { .. }
|
||||
| CoreError::Network(_)
|
||||
| CoreError::Connect(_)
|
||||
| CoreError::InvalidResponse(_)
|
||||
| CoreError::InvalidType { .. }
|
||||
| CoreError::MissingField(_) => (
|
||||
StatusCode::BAD_GATEWAY,
|
||||
"messages provider request failed".to_string(),
|
||||
),
|
||||
// The gateway has no Python implementation to decline to, so a
|
||||
// request the core cannot serve is reported to the caller. The
|
||||
// reason is a fixed internal string, never provider content.
|
||||
CoreError::Unsupported(reason) => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("messages request is not supported: {reason}"),
|
||||
),
|
||||
};
|
||||
(
|
||||
status,
|
||||
|
|
|
|||
15
litellm-rust/crates/core/src/chat_completions/client.rs
Normal file
15
litellm-rust/crates/core/src/chat_completions/client.rs
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS};
|
||||
|
||||
pub(super) fn http_client() -> &'static reqwest::Client {
|
||||
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS))
|
||||
.connect_timeout(Duration::from_secs(CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new())
|
||||
})
|
||||
}
|
||||
|
|
@ -0,0 +1,28 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
use crate::http_utils::string_headers as shared_string_headers;
|
||||
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
|
||||
|
||||
use super::transformation::ChatCompletionsProviderConfig;
|
||||
|
||||
const HEADER_CONTEXT: &str = "chat completions";
|
||||
|
||||
pub(super) fn chat_completions_provider_config(
|
||||
provider: &str,
|
||||
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
|
||||
match provider {
|
||||
"anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG),
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
"bedrock" => Some(
|
||||
&crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
|
||||
),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers)
|
||||
}
|
||||
254
litellm-rust/crates/core/src/chat_completions/conversation.rs
Normal file
254
litellm-rust/crates/core/src/chat_completions/conversation.rs
Normal file
|
|
@ -0,0 +1,254 @@
|
|||
//! Provider-neutral conversation shape.
|
||||
//!
|
||||
//! Both Anthropic Messages and Bedrock Converse want the same thing out of an
|
||||
//! OpenAI message list: the system prompt lifted out, consecutive same-role
|
||||
//! turns merged, and text blocks that are never empty. That normalization is
|
||||
//! shared here so a provider config only renders the result into its own wire
|
||||
//! shape.
|
||||
//!
|
||||
//! Mirrors Python's `anthropic_messages_pt` /
|
||||
//! `_bedrock_converse_messages_pt` for the text-only surface this route
|
||||
//! accepts; anything richer is declined upstream by the capability gate.
|
||||
|
||||
use crate::constants::EMPTY_TEXT_PLACEHOLDER;
|
||||
|
||||
use super::types::{ChatMessage, ChatMessageContent};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum TurnRole {
|
||||
User,
|
||||
Assistant,
|
||||
}
|
||||
|
||||
impl TurnRole {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::User => "user",
|
||||
Self::Assistant => "assistant",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Turn {
|
||||
pub role: TurnRole,
|
||||
pub texts: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct Conversation {
|
||||
pub system: Vec<String>,
|
||||
pub turns: Vec<Turn>,
|
||||
}
|
||||
|
||||
/// True when the conversation can be sent as-is.
|
||||
///
|
||||
/// Python inserts a placeholder first user turn only under
|
||||
/// `litellm.modify_params`, which the core cannot see, so a conversation that
|
||||
/// does not open on a user turn is declined rather than guessed at.
|
||||
impl Conversation {
|
||||
pub fn opens_on_user_turn(&self) -> bool {
|
||||
self.turns
|
||||
.first()
|
||||
.is_some_and(|turn| turn.role == TurnRole::User)
|
||||
}
|
||||
}
|
||||
|
||||
fn message_texts(content: &ChatMessageContent) -> Vec<String> {
|
||||
match content {
|
||||
ChatMessageContent::Text(text) => vec![text.clone()],
|
||||
ChatMessageContent::Parts(parts) => parts
|
||||
.iter()
|
||||
.filter_map(|part| part.get("text").and_then(|text| text.as_str()))
|
||||
.map(str::to_string)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Python rewrites empty or whitespace-only text rather than dropping it, so an
|
||||
/// entirely empty content list never reaches a provider that rejects one.
|
||||
fn sanitize(text: String) -> String {
|
||||
if text.trim().is_empty() {
|
||||
return EMPTY_TEXT_PLACEHOLDER.to_string();
|
||||
}
|
||||
text
|
||||
}
|
||||
|
||||
pub fn build_conversation(messages: &[ChatMessage]) -> Conversation {
|
||||
let system = messages
|
||||
.iter()
|
||||
.filter(|message| message.role == "system")
|
||||
.filter_map(|message| message.content.as_ref())
|
||||
.flat_map(message_texts)
|
||||
.filter(|text| !text.is_empty())
|
||||
.collect();
|
||||
|
||||
let turns = messages
|
||||
.iter()
|
||||
.filter(|message| message.role != "system")
|
||||
.fold(Vec::<Turn>::new(), |mut turns, message| {
|
||||
let role = if message.role == "assistant" {
|
||||
TurnRole::Assistant
|
||||
} else {
|
||||
TurnRole::User
|
||||
};
|
||||
let texts = message
|
||||
.content
|
||||
.as_ref()
|
||||
.map(message_texts)
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(sanitize);
|
||||
match turns.last_mut() {
|
||||
Some(last) if last.role == role => last.texts.extend(texts),
|
||||
_ => turns.push(Turn {
|
||||
role,
|
||||
texts: texts.collect(),
|
||||
}),
|
||||
}
|
||||
turns
|
||||
});
|
||||
|
||||
// Anthropic and Bedrock both reject trailing whitespace on the final
|
||||
// assistant turn, so Python right-strips it there; mirror that exactly.
|
||||
let turns = match turns.split_last() {
|
||||
Some((last, rest)) if last.role == TurnRole::Assistant => rest
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain([Turn {
|
||||
role: last.role,
|
||||
texts: last
|
||||
.texts
|
||||
.iter()
|
||||
.map(|text| text.trim_end().to_string())
|
||||
.collect(),
|
||||
}])
|
||||
.collect(),
|
||||
_ => turns,
|
||||
};
|
||||
|
||||
Conversation { system, turns }
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: serde_json::Value) -> Vec<ChatMessage> {
|
||||
serde_json::from_value(value).expect("valid messages")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifts_system_messages_out_of_the_turn_list() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
])));
|
||||
assert_eq!(conversation.system, vec!["be terse".to_string()]);
|
||||
assert_eq!(
|
||||
conversation.turns,
|
||||
vec![Turn {
|
||||
role: TurnRole::User,
|
||||
texts: vec!["hi".to_string()]
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_consecutive_same_role_turns() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": "one"},
|
||||
{"role": "user", "content": "two"},
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "three"}
|
||||
])));
|
||||
assert_eq!(
|
||||
conversation.turns,
|
||||
vec![
|
||||
Turn {
|
||||
role: TurnRole::User,
|
||||
texts: vec!["one".to_string(), "two".to_string()]
|
||||
},
|
||||
Turn {
|
||||
role: TurnRole::Assistant,
|
||||
texts: vec!["ack".to_string()]
|
||||
},
|
||||
Turn {
|
||||
role: TurnRole::User,
|
||||
texts: vec!["three".to_string()]
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flattens_text_parts_in_order() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "text", "text": "second"}
|
||||
]}
|
||||
])));
|
||||
assert_eq!(
|
||||
conversation.turns[0].texts,
|
||||
vec!["first".to_string(), "second".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rewrites_empty_and_whitespace_only_text_to_the_python_placeholder() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": ""},
|
||||
{"role": "assistant", "content": " "},
|
||||
{"role": "user", "content": "real"}
|
||||
])));
|
||||
assert_eq!(conversation.turns[0].texts, vec![EMPTY_TEXT_PLACEHOLDER]);
|
||||
assert_eq!(conversation.turns[1].texts, vec![EMPTY_TEXT_PLACEHOLDER]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn right_strips_only_the_final_assistant_turn() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "kept "},
|
||||
{"role": "user", "content": "more"},
|
||||
{"role": "assistant", "content": "stripped "}
|
||||
])));
|
||||
assert_eq!(conversation.turns[1].texts, vec!["kept ".to_string()]);
|
||||
assert_eq!(conversation.turns[3].texts, vec!["stripped".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_strip_when_the_last_turn_is_a_user_turn() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "assistant", "content": "kept "},
|
||||
{"role": "user", "content": "hi "}
|
||||
])));
|
||||
assert_eq!(conversation.turns[0].texts, vec!["kept ".to_string()]);
|
||||
assert_eq!(conversation.turns[1].texts, vec!["hi ".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_whether_the_conversation_opens_on_a_user_turn() {
|
||||
assert!(
|
||||
build_conversation(&messages(json!([{"role": "user", "content": "hi"}])))
|
||||
.opens_on_user_turn()
|
||||
);
|
||||
assert!(
|
||||
!build_conversation(&messages(json!([{"role": "assistant", "content": "hi"}])))
|
||||
.opens_on_user_turn()
|
||||
);
|
||||
assert!(!Conversation::default().opens_on_user_turn());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drops_empty_system_text_the_way_python_skips_empty_system_blocks() {
|
||||
let conversation = build_conversation(&messages(json!([
|
||||
{"role": "system", "content": ""},
|
||||
{"role": "system", "content": "kept"},
|
||||
{"role": "user", "content": "hi"}
|
||||
])));
|
||||
assert_eq!(conversation.system, vec!["kept".to_string()]);
|
||||
}
|
||||
}
|
||||
147
litellm-rust/crates/core/src/chat_completions/handler.rs
Normal file
147
litellm-rust/crates/core/src/chat_completions/handler.rs
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::http_utils::truncate_error_body;
|
||||
|
||||
use super::client::http_client;
|
||||
use super::transformation::ChatCompletionsAuth;
|
||||
use super::types::{
|
||||
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
|
||||
};
|
||||
|
||||
pub(super) async fn execute_chat_completions_provider_call(
|
||||
request: ProviderChatCompletionsRequest,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = serde_json::to_vec(&request.body).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"failed to serialize chat completions request: {err}"
|
||||
))
|
||||
})?;
|
||||
let headers = signed_headers(&request, &body).await?;
|
||||
|
||||
let mut request_builder = http_client().post(&request.url).body(body);
|
||||
for (key, value) in &headers {
|
||||
request_builder = request_builder.header(key, value);
|
||||
}
|
||||
if let Some(duration) = request.timeout {
|
||||
request_builder = request_builder.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request_builder.send().await.map_err(|err| {
|
||||
// Failing to establish the connection means the request never went out,
|
||||
// so the host can still serve it. Everything else here, a timeout
|
||||
// above all, may have reached the provider and been answered.
|
||||
if err.is_connect() || err.is_builder() {
|
||||
CoreError::Connect(err.to_string())
|
||||
} else {
|
||||
CoreError::Network(err.to_string())
|
||||
}
|
||||
})?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let body: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
|
||||
})?;
|
||||
request
|
||||
.config
|
||||
.transform_response(&request.model, ProviderChatResponseData { body })
|
||||
.map_err(as_response_error)
|
||||
}
|
||||
|
||||
/// Re-tag an error raised while normalizing a response the provider already
|
||||
/// returned.
|
||||
///
|
||||
/// A config reports the same variants on either side of the call: a missing
|
||||
/// field or an unsupported block can mean "this request cannot be translated"
|
||||
/// during prepare and "this response cannot be normalized" here. Only the
|
||||
/// second kind has already been billed, and a host that keeps a reference
|
||||
/// implementation must not retry those, so collapse them to one variant that
|
||||
/// can only mean the provider was already called.
|
||||
pub(super) fn as_response_error(err: CoreError) -> CoreError {
|
||||
match err {
|
||||
already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already,
|
||||
other => CoreError::InvalidResponse(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
pub(super) async fn signed_headers(
|
||||
request: &ProviderChatCompletionsRequest,
|
||||
body: &[u8],
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use crate::providers::bedrock::aws_base::{
|
||||
aws_auth_config, aws_signature_headers, host_supplied_credentials,
|
||||
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
|
||||
};
|
||||
|
||||
let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else {
|
||||
return Ok(request.upstream_headers.clone());
|
||||
};
|
||||
// Reattaching a header the signer also emits would put both copies on the
|
||||
// wire, and Bedrock rejects that pair. Python instead drops the caller's
|
||||
// copy and prefers a forwarded Authorization over the signature, so leave
|
||||
// the request to Python rather than serving it a different way here.
|
||||
if request
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, _)| is_sigv4_computed_header(name))
|
||||
{
|
||||
return Err(CoreError::Unsupported(
|
||||
"request forwards a header AWS SigV4 computes",
|
||||
));
|
||||
}
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
|
||||
// A host with its own resolution chain hands the result down; only fall
|
||||
// back to deriving credentials here when it supplied none.
|
||||
let credentials = match host_supplied_credentials(&request.optional_params) {
|
||||
Some(credentials) => credentials,
|
||||
None => {
|
||||
resolve_credentials(
|
||||
aws_auth_config(&request.optional_params, &env_lookup),
|
||||
&env_lookup,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
let signature = sign_bedrock_post(
|
||||
&request.url,
|
||||
body,
|
||||
&aws_signature_headers(&unsigned),
|
||||
region,
|
||||
&credentials,
|
||||
SystemTime::now(),
|
||||
)?;
|
||||
// Every original header goes back on the wire alongside the computed ones,
|
||||
// as Python reattaches them. The guard above already rejected the names
|
||||
// that would collide, so no name appears twice.
|
||||
Ok(unsigned.into_iter().chain(signature).collect())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "bedrock-auth"))]
|
||||
pub(super) async fn signed_headers(
|
||||
request: &ProviderChatCompletionsRequest,
|
||||
_body: &[u8],
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
match &request.auth {
|
||||
ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported(
|
||||
"AWS SigV4 requires the bedrock-auth feature",
|
||||
)),
|
||||
_ => Ok(request.upstream_headers.clone()),
|
||||
}
|
||||
}
|
||||
59
litellm-rust/crates/core/src/chat_completions/mod.rs
Normal file
59
litellm-rust/crates/core/src/chat_completions/mod.rs
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
//! The `/chat/completions` call, the Rust equivalent of Python's
|
||||
//! `litellm.completion()`.
|
||||
//!
|
||||
//! [`chat_completions`] is the top-level entrypoint: give it a model, the
|
||||
//! OpenAI-shaped message list, the provider-mapped optional params, and
|
||||
//! credentials, and it resolves the provider, translates the conversation,
|
||||
//! calls the provider, and returns a typed OpenAI-shaped response.
|
||||
|
||||
mod client;
|
||||
mod common_utils;
|
||||
pub mod conversation;
|
||||
pub(crate) mod handler;
|
||||
mod prepare;
|
||||
pub mod response_utils;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
|
||||
use handler::execute_chat_completions_provider_call;
|
||||
use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config};
|
||||
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
|
||||
|
||||
pub async fn chat_completions(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await
|
||||
}
|
||||
|
||||
/// Whether the core would accept this request, without resolving credentials or
|
||||
/// touching the network.
|
||||
///
|
||||
/// A host that keeps the Python implementation asks this first so it can emit
|
||||
/// its pre-call logging exactly once, on whichever path is about to run.
|
||||
/// Returns the decline reason, or `None` when the request is accepted.
|
||||
pub fn chat_completions_decline_reason(
|
||||
model: &str,
|
||||
custom_llm_provider: Option<&str>,
|
||||
messages: Value,
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<&'static str> {
|
||||
let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else {
|
||||
return Some("provider is not on the rust chat completions path");
|
||||
};
|
||||
let Ok(messages) = parse_messages(messages) else {
|
||||
return Some("unreadable message list");
|
||||
};
|
||||
if messages.is_empty() {
|
||||
return Some("empty message list");
|
||||
}
|
||||
config
|
||||
.unsupported_reason(&messages, optional_params)
|
||||
.map(|reason| reason.0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
118
litellm-rust/crates/core/src/chat_completions/prepare.rs
Normal file
118
litellm-rust/crates/core/src/chat_completions/prepare.rs
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
use serde_json::Value;
|
||||
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::http_utils::has_header;
|
||||
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
|
||||
|
||||
use super::common_utils::{chat_completions_provider_config, string_headers};
|
||||
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
|
||||
use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest};
|
||||
|
||||
pub(super) fn resolve_provider_config<'a>(
|
||||
model: &'a str,
|
||||
custom_llm_provider: Option<&'a str>,
|
||||
) -> CoreResult<(String, &'static dyn ChatCompletionsProviderConfig)> {
|
||||
let provider_info = get_custom_llm_provider(model, custom_llm_provider)
|
||||
.or_else(|| {
|
||||
custom_llm_provider.map(|provider| CustomLlmProvider {
|
||||
model,
|
||||
custom_llm_provider: provider,
|
||||
})
|
||||
})
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidProvider(
|
||||
"unable to resolve custom_llm_provider for chat completions request".to_string(),
|
||||
)
|
||||
})?;
|
||||
let config = chat_completions_provider_config(provider_info.custom_llm_provider)
|
||||
.ok_or_else(|| CoreError::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
|
||||
Ok((provider_info.model.to_string(), config))
|
||||
}
|
||||
|
||||
pub(super) fn parse_messages(messages: Value) -> CoreResult<Vec<ChatMessage>> {
|
||||
serde_json::from_value(messages).map_err(|err| {
|
||||
CoreError::InvalidRequest(format!("invalid chat completions messages: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn prepare_chat_completions_call(
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
) -> CoreResult<ProviderChatCompletionsRequest> {
|
||||
let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?;
|
||||
let env_lookup = |key: &str| std::env::var(key).ok();
|
||||
|
||||
let messages = parse_messages(request.messages)?;
|
||||
if messages.is_empty() {
|
||||
return Err(CoreError::InvalidRequest(
|
||||
"chat completions requires at least one message".to_string(),
|
||||
));
|
||||
}
|
||||
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
|
||||
return Err(CoreError::Unsupported(reason.0));
|
||||
}
|
||||
|
||||
let mut headers = string_headers(request.extra_headers)?;
|
||||
let auth = config.auth(
|
||||
request.api_key,
|
||||
&model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
match &auth {
|
||||
ChatCompletionsAuth::Header { name, value } => {
|
||||
// The deployment's credential replaces whatever the caller forwarded
|
||||
// under the same name, mirroring Python's
|
||||
// `{**headers, **anthropic_headers}`: letting a request header win
|
||||
// would let its sender choose the principal the call bills to.
|
||||
//
|
||||
// The exception is a scheme the provider hands off to entirely, such
|
||||
// as an Anthropic OAuth bearer, where Python drops `x-api-key`
|
||||
// instead of resolving one. Re-adding it there would put the
|
||||
// credential into a header the host removed on purpose.
|
||||
if !config.defers_to_forwarded_auth(&headers) {
|
||||
headers.retain(|(header, _)| !header.eq_ignore_ascii_case(name));
|
||||
headers.push(((*name).to_string(), value.clone()));
|
||||
}
|
||||
}
|
||||
ChatCompletionsAuth::Bearer { token } => {
|
||||
// Bedrock's `get_request_headers` assigns `headers["Authorization"]`
|
||||
// unconditionally once a bearer token resolves, so the deployment's
|
||||
// identity outranks whatever the caller forwarded. Keeping the
|
||||
// caller's would bill and authorize the call as a different
|
||||
// principal than the same deployment uses on Python.
|
||||
//
|
||||
// The `Header` arm below keeps the opposite precedence on purpose:
|
||||
// Anthropic's transform honours a forwarded OAuth bearer.
|
||||
headers.retain(|(name, _)| !name.eq_ignore_ascii_case("authorization"));
|
||||
headers.push(("authorization".to_string(), format!("Bearer {token}")));
|
||||
}
|
||||
// SigV4 signs the serialized body, so the handler adds its headers.
|
||||
ChatCompletionsAuth::AwsSigV4 { .. } => {}
|
||||
}
|
||||
|
||||
for (name, value) in config.default_headers() {
|
||||
if !has_header(&headers, name) {
|
||||
headers.push(((*name).to_string(), (*value).to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
let url = config.complete_url(
|
||||
request.api_base,
|
||||
&model,
|
||||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let transformed =
|
||||
config.transform_request(&model, messages, request.optional_params.clone())?;
|
||||
|
||||
Ok(ProviderChatCompletionsRequest {
|
||||
model,
|
||||
config,
|
||||
url,
|
||||
body: transformed.body,
|
||||
upstream_headers: headers,
|
||||
auth,
|
||||
optional_params: request.optional_params,
|
||||
timeout: request.timeout,
|
||||
})
|
||||
}
|
||||
101
litellm-rust/crates/core/src/chat_completions/response_utils.rs
Normal file
101
litellm-rust/crates/core/src/chat_completions/response_utils.rs
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
//! Response normalization shared by every chat completions provider config.
|
||||
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use super::types::{ChatCompletionsUsage, PromptTokensDetails};
|
||||
|
||||
/// OpenAI finish reasons, mirroring Python's `_FINISH_REASON_MAP` for the
|
||||
/// reasons the providers on this route can emit. Python warns and falls back to
|
||||
/// `stop` for anything unmapped, so do the same.
|
||||
const FINISH_REASONS: &[(&str, &str)] = &[
|
||||
("end_turn", "stop"),
|
||||
("stop_sequence", "stop"),
|
||||
("max_tokens", "length"),
|
||||
("refusal", "content_filter"),
|
||||
("compaction", "length"),
|
||||
("guardrail_intervened", "content_filter"),
|
||||
("content_filtered", "content_filter"),
|
||||
("content_filter", "content_filter"),
|
||||
("stop", "stop"),
|
||||
("length", "length"),
|
||||
];
|
||||
|
||||
pub fn finish_reason_for(provider_reason: &str) -> &'static str {
|
||||
FINISH_REASONS
|
||||
.iter()
|
||||
.find(|(reason, _)| *reason == provider_reason)
|
||||
.map_or("stop", |(_, mapped)| *mapped)
|
||||
}
|
||||
|
||||
/// Python folds cache tokens into `prompt_tokens` and reports the split under
|
||||
/// `prompt_tokens_details`; mirror that so cost tracking agrees on both paths.
|
||||
pub fn usage_from_parts(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cache_read_tokens: u64,
|
||||
cache_creation_tokens: u64,
|
||||
) -> ChatCompletionsUsage {
|
||||
let prompt_tokens = input_tokens + cache_read_tokens + cache_creation_tokens;
|
||||
ChatCompletionsUsage {
|
||||
prompt_tokens,
|
||||
completion_tokens: output_tokens,
|
||||
total_tokens: prompt_tokens + output_tokens,
|
||||
prompt_tokens_details: PromptTokensDetails {
|
||||
cached_tokens: cache_read_tokens,
|
||||
cache_creation_tokens,
|
||||
text_tokens: input_tokens,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn unix_now() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_or(0, |elapsed| elapsed.as_secs())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn maps_every_reason_the_route_can_observe() {
|
||||
assert_eq!(finish_reason_for("end_turn"), "stop");
|
||||
assert_eq!(finish_reason_for("stop_sequence"), "stop");
|
||||
assert_eq!(finish_reason_for("max_tokens"), "length");
|
||||
assert_eq!(finish_reason_for("refusal"), "content_filter");
|
||||
assert_eq!(finish_reason_for("guardrail_intervened"), "content_filter");
|
||||
// Converse emits these two, and folding them into `stop` would report a
|
||||
// filtered completion as a normal one.
|
||||
assert_eq!(finish_reason_for("content_filtered"), "content_filter");
|
||||
assert_eq!(finish_reason_for("content_filter"), "content_filter");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_an_unmapped_reason_to_stop_like_python() {
|
||||
// Python warns and falls back to `stop` for a reason its own map does
|
||||
// not carry, so only a reason absent from `_FINISH_REASON_MAP` belongs
|
||||
// here.
|
||||
assert_eq!(finish_reason_for("something_new"), "stop");
|
||||
assert_eq!(finish_reason_for(""), "stop");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn folds_cache_tokens_into_prompt_tokens() {
|
||||
let usage = usage_from_parts(10, 4, 7, 3);
|
||||
assert_eq!(usage.prompt_tokens, 20);
|
||||
assert_eq!(usage.completion_tokens, 4);
|
||||
assert_eq!(usage.total_tokens, 24);
|
||||
assert_eq!(usage.prompt_tokens_details.cached_tokens, 7);
|
||||
assert_eq!(usage.prompt_tokens_details.cache_creation_tokens, 3);
|
||||
assert_eq!(usage.prompt_tokens_details.text_tokens, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_raw_input_tokens_when_no_cache_is_involved() {
|
||||
let usage = usage_from_parts(12, 5, 0, 0);
|
||||
assert_eq!(usage.prompt_tokens, 12);
|
||||
assert_eq!(usage.total_tokens, 17);
|
||||
assert_eq!(usage.prompt_tokens_details.text_tokens, 12);
|
||||
}
|
||||
}
|
||||
820
litellm-rust/crates/core/src/chat_completions/tests.rs
Normal file
820
litellm-rust/crates/core/src/chat_completions/tests.rs
Normal file
|
|
@ -0,0 +1,820 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::error::CoreError;
|
||||
|
||||
use super::prepare::prepare_chat_completions_call;
|
||||
use super::transformation::ChatCompletionsAuth;
|
||||
use super::types::ChatCompletionsRequest;
|
||||
|
||||
fn request<'a>(
|
||||
model: &'a str,
|
||||
provider: Option<&'a str>,
|
||||
messages: Value,
|
||||
optional_params: Value,
|
||||
) -> ChatCompletionsRequest<'a> {
|
||||
ChatCompletionsRequest {
|
||||
model,
|
||||
messages,
|
||||
optional_params: match optional_params {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
},
|
||||
api_key: Some("sk-test"),
|
||||
api_base: None,
|
||||
custom_llm_provider: provider,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers
|
||||
/// carry resolved credentials), so unwrap the failure case by hand.
|
||||
fn decline(request: ChatCompletionsRequest<'_>) -> CoreError {
|
||||
match prepare_chat_completions_call(request) {
|
||||
Err(error) => error,
|
||||
Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_the_provider_from_the_model_prefix() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(prepared.model, "claude-sonnet-4-5");
|
||||
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
|
||||
assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strips_an_explicit_provider_prefix_from_the_model() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(prepared.model, "claude-sonnet-4-5");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn adds_the_auth_and_default_headers() {
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.contains(&("x-api-key".to_string(), "sk-test".to_string()))
|
||||
);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.contains(&("anthropic-version".to_string(), "2023-06-01".to_string()))
|
||||
);
|
||||
assert!(matches!(
|
||||
prepared.auth,
|
||||
ChatCompletionsAuth::Header {
|
||||
name: "x-api-key",
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
|
||||
// Python builds `{**headers, **anthropic_headers}`, so the deployment's key
|
||||
// overwrites a forwarded one. Honouring the caller's would let whoever sends
|
||||
// the request choose the Anthropic principal it bills to.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"X-Api-Key".to_string(),
|
||||
json!("sk-caller"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let keys: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.collect();
|
||||
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
|
||||
assert_eq!(keys[0].1, "sk-test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
|
||||
// Anthropic's `validate_environment` pops `x-api-key` and sets `authorization`
|
||||
// for an OAuth token, so re-adding the key here would put the credential into
|
||||
// a header the host removed on purpose.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([
|
||||
(
|
||||
"Authorization".to_string(),
|
||||
json!("Bearer sk-ant-oat01-token"),
|
||||
),
|
||||
("X-Api-Key".to_string(), json!("sk-caller")),
|
||||
]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
assert!(
|
||||
!prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"),
|
||||
"the resolved key must not be applied over an OAuth bearer, got {:?}",
|
||||
prepared.upstream_headers
|
||||
);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-ant-oat01-token")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
|
||||
// Only an OAuth bearer replaces the credential. Python sends the deployment's
|
||||
// `x-api-key` alongside any other forwarded `authorization`, so deferring on
|
||||
// the mere presence of that header would drop the deployment's auth.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([
|
||||
("Authorization".to_string(), json!("Bearer unrelated")),
|
||||
("X-Api-Key".to_string(), json!("sk-caller")),
|
||||
]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let keys: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.collect();
|
||||
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
|
||||
assert_eq!(keys[0].1, "sk-test");
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer unrelated"),
|
||||
"the unrelated authorization must survive, got {:?}",
|
||||
prepared.upstream_headers
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_an_unsupported_request_before_resolving_credentials() {
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true}),
|
||||
);
|
||||
call.api_key = None;
|
||||
// No api_key is set and no env is consulted: the gate must run first, so the
|
||||
// error is the decline rather than a missing-credential error.
|
||||
assert_eq!(decline(call), CoreError::Unsupported("streaming"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_an_unknown_provider() {
|
||||
assert_eq!(
|
||||
decline(request(
|
||||
"openai/gpt-4o",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidProvider("openai".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_a_model_with_no_resolvable_provider() {
|
||||
assert!(matches!(
|
||||
decline(request(
|
||||
"claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidProvider(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_an_empty_or_malformed_message_list() {
|
||||
assert_eq!(
|
||||
decline(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([]),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidRequest("chat completions requires at least one message".to_string())
|
||||
);
|
||||
assert!(matches!(
|
||||
decline(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!("not a list"),
|
||||
json!({}),
|
||||
)),
|
||||
CoreError::InvalidRequest(_)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_string_extra_headers() {
|
||||
let mut call = request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
|
||||
assert_eq!(
|
||||
decline(call),
|
||||
CoreError::InvalidRequest(
|
||||
"chat completions extra_headers.x-trace must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[test]
|
||||
fn prepares_a_bedrock_call_without_resolving_credentials() {
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
);
|
||||
call.api_key = None;
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
assert_eq!(
|
||||
prepared.url,
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
assert_eq!(
|
||||
prepared.auth,
|
||||
ChatCompletionsAuth::AwsSigV4 {
|
||||
region: "us-east-1".to_string()
|
||||
}
|
||||
);
|
||||
// SigV4 signs the serialized body, so prepare must not have added an
|
||||
// Authorization header; the handler does it.
|
||||
assert!(
|
||||
!prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
);
|
||||
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
|
||||
// Python signs only the AWS header set and reattaches the rest, so a header
|
||||
// the caller forwarded rides along without joining the canonical request.
|
||||
// Signing it makes Converse 403 on a deployment that works on Python.
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
);
|
||||
// A key would resolve to a bearer token and never reach the signer.
|
||||
call.api_key = None;
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"x-request-id".to_string(),
|
||||
json!("abc-123"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
|
||||
.await
|
||||
.expect("signs");
|
||||
|
||||
let authorization = signed
|
||||
.iter()
|
||||
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.clone())
|
||||
.expect("carries an authorization header");
|
||||
assert!(
|
||||
authorization.starts_with("AWS4-HMAC-SHA256"),
|
||||
"expected a SigV4 signature, got {authorization}"
|
||||
);
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
// It still goes on the wire, it is just not part of the signature.
|
||||
assert!(
|
||||
signed
|
||||
.iter()
|
||||
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
|
||||
"forwarded header was dropped instead of reattached"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[tokio::test]
|
||||
async fn a_forwarded_header_the_signer_computes_declines_to_python() {
|
||||
// Reattaching the caller's copy next to the computed one puts the name on
|
||||
// the wire twice and Bedrock rejects the pair, so a request carrying one
|
||||
// has to go to Python instead of being signed here.
|
||||
for forwarded in [
|
||||
"Authorization",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"Date",
|
||||
] {
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIDEXAMPLE",
|
||||
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
|
||||
}),
|
||||
);
|
||||
call.api_key = None;
|
||||
call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
|
||||
.await
|
||||
.expect_err("{forwarded} should decline instead of being signed");
|
||||
assert!(
|
||||
matches!(error, CoreError::Unsupported(_)),
|
||||
"{forwarded} declined as {error:?}, which the host would not fall back on"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[test]
|
||||
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
|
||||
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
|
||||
// once a bearer token resolves, so the deployment's identity wins on
|
||||
// Python. Keeping the caller's would authorize and bill the call as a
|
||||
// different principal, and only when the deployment carries `rust: true`.
|
||||
let mut call = request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"Authorization".to_string(),
|
||||
json!("Bearer caller-supplied"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let authorizations: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
.collect();
|
||||
assert_eq!(
|
||||
authorizations,
|
||||
vec!["Bearer sk-test"],
|
||||
"the deployment token must be the only authorization on the wire"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
|
||||
// The opposite precedence, and deliberate: Anthropic's own transform
|
||||
// honours a forwarded OAuth bearer, so the Bedrock fix above must not be
|
||||
// generalized into a rule that the configured key always wins.
|
||||
//
|
||||
// An OAuth bearer is the whole of that exception. This forwarded a plain
|
||||
// `x-api-key` until round 17, which read as the same claim and was not:
|
||||
// Python overwrites a forwarded `x-api-key` with the deployment's.
|
||||
let mut call = request(
|
||||
"claude-sonnet-4-5",
|
||||
Some("anthropic"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
);
|
||||
call.extra_headers = Some(Map::from_iter([(
|
||||
"authorization".to_string(),
|
||||
json!("Bearer sk-ant-oat01-forwarded"),
|
||||
)]));
|
||||
let prepared = prepare_chat_completions_call(call).expect("prepares");
|
||||
let keys: Vec<_> = prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
|
||||
.map(|(_, value)| value.as_str())
|
||||
.collect();
|
||||
assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-ant-oat01-forwarded")
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "bedrock-auth")]
|
||||
#[test]
|
||||
fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
|
||||
// The configured bearer identity has its own account and quota boundary,
|
||||
// so a request carrying one must not be signed as whatever principal the
|
||||
// host's AWS credentials resolve to.
|
||||
let prepared = prepare_chat_completions_call(request(
|
||||
"bedrock/us-east-1/anthropic.claude-v2",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"maxTokens": 16}),
|
||||
))
|
||||
.expect("prepares");
|
||||
assert_eq!(
|
||||
prepared.auth,
|
||||
ChatCompletionsAuth::Bearer {
|
||||
token: "sk-test".to_string()
|
||||
}
|
||||
);
|
||||
assert!(
|
||||
prepared
|
||||
.upstream_headers
|
||||
.iter()
|
||||
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
|
||||
&& value == "Bearer sk-test"),
|
||||
"prepare did not carry the bearer token"
|
||||
);
|
||||
}
|
||||
|
||||
fn decline_reason(
|
||||
model: &str,
|
||||
provider: Option<&str>,
|
||||
messages: Value,
|
||||
params: Value,
|
||||
) -> Option<&'static str> {
|
||||
let params = match params {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
};
|
||||
super::chat_completions_decline_reason(model, provider, messages, ¶ms)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_gate_accepts_what_prepare_accepts() {
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_gate_declines_without_resolving_credentials_or_calling_out() {
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true}),
|
||||
),
|
||||
Some("streaming")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"openai/gpt-4o",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
),
|
||||
Some("provider is not on the rust chat completions path")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"claude-sonnet-4-5",
|
||||
None,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({}),
|
||||
),
|
||||
Some("provider is not on the rust chat completions path")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
json!("nope"),
|
||||
json!({})
|
||||
),
|
||||
Some("unreadable message list")
|
||||
);
|
||||
assert_eq!(
|
||||
decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})),
|
||||
Some("empty message list")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
|
||||
// A gate that accepts what prepare then declines would make the host emit
|
||||
// its pre-call logging on a path that falls back, so pin the agreement.
|
||||
for (messages, params) in [
|
||||
(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 8}),
|
||||
),
|
||||
(
|
||||
json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]),
|
||||
json!({"temperature": 0.1}),
|
||||
),
|
||||
(
|
||||
json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]),
|
||||
json!({}),
|
||||
),
|
||||
] {
|
||||
assert_eq!(
|
||||
decline_reason(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
messages.clone(),
|
||||
params.clone()
|
||||
),
|
||||
None,
|
||||
"gate declined {messages}"
|
||||
);
|
||||
prepare_chat_completions_call(request(
|
||||
"anthropic/claude-sonnet-4-5",
|
||||
None,
|
||||
messages.clone(),
|
||||
params,
|
||||
))
|
||||
.unwrap_or_else(|error| panic!("prepare declined {messages}: {error}"));
|
||||
}
|
||||
}
|
||||
|
||||
mod round_trip {
|
||||
use super::*;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use crate::chat_completions::chat_completions;
|
||||
|
||||
async fn read_http_request(socket: &mut TcpStream) -> String {
|
||||
let mut request = Vec::new();
|
||||
let mut buffer = [0_u8; 1024];
|
||||
let header_end = loop {
|
||||
let n = socket.read(&mut buffer).await.expect("reads request");
|
||||
if n == 0 {
|
||||
break request.len();
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
|
||||
break position + 4;
|
||||
}
|
||||
};
|
||||
let headers = String::from_utf8_lossy(&request[..header_end]);
|
||||
let content_length = headers
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().ok())
|
||||
.flatten()
|
||||
})
|
||||
.unwrap_or(0);
|
||||
while request.len().saturating_sub(header_end) < content_length {
|
||||
let n = socket.read(&mut buffer).await.expect("reads body");
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&buffer[..n]);
|
||||
}
|
||||
String::from_utf8(request).expect("request is utf8")
|
||||
}
|
||||
|
||||
fn http_response(status: &str, body: &str) -> String {
|
||||
format!(
|
||||
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
body.len(),
|
||||
body
|
||||
)
|
||||
}
|
||||
|
||||
/// Serve one request from a stub upstream and hand back what it received.
|
||||
async fn serve_once(
|
||||
status: &'static str,
|
||||
body: &'static str,
|
||||
) -> (String, tokio::task::JoinHandle<String>) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
let port = listener.local_addr().expect("addr").port();
|
||||
let handle = tokio::spawn(async move {
|
||||
let (mut socket, _) = listener.accept().await.expect("accepts");
|
||||
let received = read_http_request(&mut socket).await;
|
||||
socket
|
||||
.write_all(http_response(status, body).as_bytes())
|
||||
.await
|
||||
.expect("writes response");
|
||||
socket.flush().await.expect("flushes");
|
||||
received
|
||||
});
|
||||
(format!("http://127.0.0.1:{port}/v1/messages"), handle)
|
||||
}
|
||||
|
||||
fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> {
|
||||
ChatCompletionsRequest {
|
||||
model: "anthropic/claude-sonnet-4-5",
|
||||
messages,
|
||||
optional_params: match params {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
},
|
||||
api_key: Some("sk-test"),
|
||||
api_base: Some(api_base),
|
||||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: Some(std::time::Duration::from_secs(10)),
|
||||
}
|
||||
}
|
||||
|
||||
const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn round_trip_sends_the_translated_body_and_normalizes_the_response() {
|
||||
let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await;
|
||||
let response = chat_completions(call(
|
||||
&api_base,
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect("call succeeds");
|
||||
|
||||
let received = handle.await.expect("server task");
|
||||
let sent: Value = serde_json::from_str(
|
||||
received
|
||||
.split_once("\r\n\r\n")
|
||||
.expect("request has a body")
|
||||
.1,
|
||||
)
|
||||
.expect("body is json");
|
||||
assert_eq!(
|
||||
sent["messages"],
|
||||
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
|
||||
);
|
||||
assert_eq!(
|
||||
sent["system"],
|
||||
json!([{"type": "text", "text": "be terse"}])
|
||||
);
|
||||
assert_eq!(sent["max_tokens"], json!(16));
|
||||
assert!(received.to_lowercase().contains("x-api-key: sk-test"));
|
||||
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("hello")
|
||||
);
|
||||
assert_eq!(response.usage.total_tokens, 15);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_response_it_cannot_normalize_is_reported_as_already_sent() {
|
||||
// The provider was called and billed, so the host must not retry this
|
||||
// on its own path. `MissingField` here would read as a pre-send
|
||||
// decline and be retried; `InvalidResponse` cannot.
|
||||
const NO_USAGE: &str =
|
||||
r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#;
|
||||
let (api_base, handle) = serve_once("200 OK", NO_USAGE).await;
|
||||
let err = chat_completions(call(
|
||||
&api_base,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("response cannot be normalized");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::InvalidResponse(_)),
|
||||
"expected a post-send error, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() {
|
||||
const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#;
|
||||
let (api_base, handle) = serve_once("200 OK", TOOL_USE).await;
|
||||
let err = chat_completions(call(
|
||||
&api_base,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("response cannot be normalized");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::InvalidResponse(_)),
|
||||
"expected a post-send error, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_upstream_error_status_keeps_its_code() {
|
||||
let (api_base, handle) =
|
||||
serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await;
|
||||
let err = chat_completions(call(
|
||||
&api_base,
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("upstream rejects");
|
||||
handle.await.expect("server task");
|
||||
assert!(
|
||||
matches!(err, CoreError::Http { status: 429, .. }),
|
||||
"expected a 429, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_connection_that_is_never_established_declines_instead_of_failing() {
|
||||
// Nothing was sent, so nothing was billed and the host can still serve
|
||||
// the request. Classing this with the post-send failures would turn a
|
||||
// recoverable fallback into a user-facing error on exactly the
|
||||
// deployments whose transport is configured only on the Python client.
|
||||
let port = {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
|
||||
listener.local_addr().expect("has an address").port()
|
||||
// Dropped here, so the port is closed and the connect is refused.
|
||||
};
|
||||
let err = chat_completions(call(
|
||||
&format!("http://127.0.0.1:{port}/v1/messages"),
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
))
|
||||
.await
|
||||
.expect_err("nothing is listening");
|
||||
assert!(
|
||||
matches!(err, CoreError::Connect(_)),
|
||||
"expected a pre-send connect failure, got {err:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() {
|
||||
use crate::chat_completions::handler::as_response_error;
|
||||
|
||||
for original in [
|
||||
CoreError::MissingField("usage"),
|
||||
CoreError::Unsupported("non-text response content block"),
|
||||
CoreError::InvalidRequest("whatever".to_string()),
|
||||
CoreError::Auth("whatever".to_string()),
|
||||
] {
|
||||
let label = format!("{original:?}");
|
||||
assert!(
|
||||
matches!(as_response_error(original), CoreError::InvalidResponse(_)),
|
||||
"{label} must not stay retryable once the provider has answered"
|
||||
);
|
||||
}
|
||||
// An upstream status is already unambiguous, so it survives intact.
|
||||
assert!(matches!(
|
||||
as_response_error(CoreError::Http {
|
||||
status: 500,
|
||||
body: "boom".to_string()
|
||||
}),
|
||||
CoreError::Http { status: 500, .. }
|
||||
));
|
||||
}
|
||||
}
|
||||
155
litellm-rust/crates/core/src/chat_completions/transformation.rs
Normal file
155
litellm-rust/crates/core/src/chat_completions/transformation.rs
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::error::CoreResult;
|
||||
|
||||
use super::types::{
|
||||
ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData,
|
||||
ProviderChatResponseData,
|
||||
};
|
||||
|
||||
/// How the upstream call is authenticated. API-key strategies are resolved in
|
||||
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ChatCompletionsAuth {
|
||||
Header { name: &'static str, value: String },
|
||||
Bearer { token: String },
|
||||
AwsSigV4 { region: String },
|
||||
}
|
||||
|
||||
/// Why a request cannot be served by the Rust path.
|
||||
///
|
||||
/// The core declines rather than guessing: the host turns this into a
|
||||
/// transparent fallback to the Python implementation, which covers the full
|
||||
/// surface. Acceptance is an allowlist, so a parameter or message shape the
|
||||
/// core has never seen declines by construction instead of being translated
|
||||
/// wrong.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct Unsupported(pub &'static str);
|
||||
|
||||
pub const STREAM_PARAM: &str = "stream";
|
||||
|
||||
/// Message fields that carry no meaning for the upstream body, so their
|
||||
/// presence does not make a request untranslatable.
|
||||
const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"];
|
||||
|
||||
pub trait ChatCompletionsProviderConfig: Sync {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String>;
|
||||
|
||||
fn auth(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth>;
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[("content-type", "application/json")]
|
||||
}
|
||||
|
||||
/// Whether an auth header the caller already supplied is the credential this
|
||||
/// request should authenticate with, so the resolved one is not applied.
|
||||
///
|
||||
/// Defaults to false: the deployment's credential outranks anything
|
||||
/// forwarded, which is what every provider wants for its own auth header.
|
||||
/// A provider overrides this only for a scheme it hands off to entirely.
|
||||
fn defers_to_forwarded_auth(&self, _headers: &[(String, String)]) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Provider parameter names (post-mapping) the Rust path knows how to place
|
||||
/// in the upstream body. Anything outside this set declines the request.
|
||||
fn supported_params(&self) -> &'static [&'static str];
|
||||
|
||||
/// Parameters consumed as call configuration (credentials, endpoints)
|
||||
/// rather than placed in the body. Accepted, never serialized.
|
||||
fn config_params(&self) -> &'static [&'static str] {
|
||||
&[]
|
||||
}
|
||||
|
||||
fn unsupported_reason(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
unsupported_param(
|
||||
self.supported_params(),
|
||||
self.config_params(),
|
||||
optional_params,
|
||||
)
|
||||
.or_else(|| messages.iter().find_map(unsupported_message))
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData>;
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse>;
|
||||
}
|
||||
|
||||
pub fn unsupported_param(
|
||||
supported: &'static [&'static str],
|
||||
config: &'static [&'static str],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
if optional_params
|
||||
.get(STREAM_PARAM)
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
{
|
||||
return Some(Unsupported("streaming"));
|
||||
}
|
||||
optional_params
|
||||
.keys()
|
||||
.any(|key| {
|
||||
key != STREAM_PARAM
|
||||
&& !supported.contains(&key.as_str())
|
||||
&& !config.contains(&key.as_str())
|
||||
})
|
||||
.then_some(Unsupported("unrecognized request parameter"))
|
||||
}
|
||||
|
||||
/// Message shapes the core can translate faithfully: text content, either a
|
||||
/// plain string or a non-empty list of parts that are all
|
||||
/// `{"type": "text", "text": ...}`. Tool calls, tool results, and multimodal
|
||||
/// parts decline so Python's fuller translation handles them.
|
||||
pub fn unsupported_message(message: &ChatMessage) -> Option<Unsupported> {
|
||||
if message
|
||||
.extra
|
||||
.keys()
|
||||
.any(|key| !IGNORABLE_MESSAGE_FIELDS.contains(&key.as_str()))
|
||||
{
|
||||
return Some(Unsupported("unrecognized message field"));
|
||||
}
|
||||
if !matches!(message.role.as_str(), "system" | "user" | "assistant") {
|
||||
return Some(Unsupported("unrecognized message role"));
|
||||
}
|
||||
match &message.content {
|
||||
None => Some(Unsupported("message without content")),
|
||||
Some(ChatMessageContent::Text(_)) => None,
|
||||
Some(ChatMessageContent::Parts(parts)) if parts.is_empty() => {
|
||||
Some(Unsupported("message without content"))
|
||||
}
|
||||
Some(ChatMessageContent::Parts(parts)) => parts
|
||||
.iter()
|
||||
.any(|part| {
|
||||
part.get("type").and_then(Value::as_str) != Some("text")
|
||||
|| part.get("text").and_then(Value::as_str).is_none()
|
||||
|| part.as_object().is_some_and(|object| object.len() != 2)
|
||||
})
|
||||
.then_some(Unsupported("non-text message content")),
|
||||
}
|
||||
}
|
||||
112
litellm-rust/crates/core/src/chat_completions/types.rs
Normal file
112
litellm-rust/crates/core/src/chat_completions/types.rs
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
|
||||
|
||||
/// A `/chat/completions` call as it crosses into the core.
|
||||
///
|
||||
/// `optional_params` arrives already mapped to the provider's own parameter
|
||||
/// names by the host, exactly as the messages route receives an already
|
||||
/// Anthropic-shaped body. The core owns the conversation translation, the
|
||||
/// provider call, and the response normalization.
|
||||
pub struct ChatCompletionsRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub messages: Value,
|
||||
pub optional_params: Map<String, Value>,
|
||||
pub api_key: Option<&'a str>,
|
||||
pub api_base: Option<&'a str>,
|
||||
pub custom_llm_provider: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
pub(super) struct ProviderChatCompletionsRequest {
|
||||
pub(super) model: String,
|
||||
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
|
||||
pub(super) url: String,
|
||||
pub(super) body: Value,
|
||||
pub(super) upstream_headers: Vec<(String, String)>,
|
||||
pub(super) auth: ChatCompletionsAuth,
|
||||
#[cfg_attr(not(feature = "bedrock-auth"), allow(dead_code))]
|
||||
pub(super) optional_params: Map<String, Value>,
|
||||
pub(super) timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
/// The provider-shaped request body a config produces. Named rather than a bare
|
||||
/// `Value` so the transform contract stays a typed one, mirroring
|
||||
/// [`crate::audio_transcription::types::AudioTranscriptionRequestData`].
|
||||
pub struct ProviderChatRequestData {
|
||||
pub body: Value,
|
||||
}
|
||||
|
||||
/// The raw provider response body handed back to a config for normalization.
|
||||
pub struct ProviderChatResponseData {
|
||||
pub body: Value,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum ChatMessageContent {
|
||||
Text(String),
|
||||
Parts(Vec<Value>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatMessage {
|
||||
pub role: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<ChatMessageContent>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
/// OpenAI `usage`, including the `prompt_tokens_details` split LiteLLM's Python
|
||||
/// path reports so cost tracking sees the same numbers on either path.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct PromptTokensDetails {
|
||||
pub cached_tokens: u64,
|
||||
pub cache_creation_tokens: u64,
|
||||
pub text_tokens: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsUsage {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub prompt_tokens_details: PromptTokensDetails,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsChoiceMessage {
|
||||
pub role: String,
|
||||
// Whether an empty turn is `None` or `""` is the provider's choice, not a
|
||||
// shared invariant: Anthropic's transform ends on `merged_text or None`
|
||||
// while Converse assigns the joined string unconditionally. Each config
|
||||
// mirrors its own, so keep this optional and serialize it even when None.
|
||||
pub content: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsChoice {
|
||||
pub index: u64,
|
||||
pub message: ChatCompletionsChoiceMessage,
|
||||
pub finish_reason: String,
|
||||
}
|
||||
|
||||
/// The normalized response handed back to the host.
|
||||
///
|
||||
/// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the
|
||||
/// `ModelResponse` it already created, and echoing the provider's own id here
|
||||
/// would change it. Pinned by `response_carries_no_id` in `tests.rs`.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionsResponse {
|
||||
pub created: u64,
|
||||
pub model: String,
|
||||
pub choices: Vec<ChatCompletionsChoice>,
|
||||
pub usage: ChatCompletionsUsage,
|
||||
}
|
||||
|
|
@ -12,8 +12,30 @@ pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
|
|||
|
||||
/// Max characters of an upstream error body echoed across the call boundary
|
||||
/// before truncation, so provider bodies are bounded and data-minimized.
|
||||
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
pub(crate) const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
/// Provider name used for Anthropic Messages when a deployment's provider model
|
||||
/// does not carry an explicit provider prefix.
|
||||
pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
|
||||
|
||||
/// Prefix identifying an Anthropic OAuth token. Mirrors Python's
|
||||
/// `ANTHROPIC_OAUTH_TOKEN_PREFIX`, which is what makes `validate_environment`
|
||||
/// authenticate with `authorization` and drop `x-api-key` entirely.
|
||||
pub(crate) const ANTHROPIC_OAUTH_TOKEN_PREFIX: &str = "sk-ant-oat";
|
||||
|
||||
/// Full-request timeout ceiling for chat completions provider calls, in
|
||||
/// seconds. Mirrors the Python chat completions default.
|
||||
pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Connect timeout for chat completions provider calls, in seconds.
|
||||
pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10;
|
||||
|
||||
/// `object` field every non-streaming chat completion response carries.
|
||||
pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
|
||||
|
||||
/// Placeholder Python substitutes for empty or whitespace-only message text,
|
||||
/// which Anthropic and Bedrock both reject. Must match
|
||||
/// `_EMPTY_TEXT_PLACEHOLDER` in
|
||||
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
|
||||
pub const EMPTY_TEXT_PLACEHOLDER: &str =
|
||||
"[System: Empty message content sanitised to satisfy protocol]";
|
||||
|
|
|
|||
|
|
@ -23,8 +23,19 @@ pub enum CoreError {
|
|||
Http { status: u16, body: String },
|
||||
#[error("upstream network error: {0}")]
|
||||
Network(String),
|
||||
/// The provider was never reached: DNS, TCP, TLS or proxy setup failed
|
||||
/// before any byte of the request went out. Nothing was billed, so a host
|
||||
/// that keeps a reference implementation can serve the request itself.
|
||||
/// A timeout is deliberately not this, since the provider may have received
|
||||
/// and answered the request already.
|
||||
#[error("could not reach the provider: {0}")]
|
||||
Connect(String),
|
||||
#[error("routing error: {0}")]
|
||||
Routing(String),
|
||||
/// The request is outside the surface this route covers in Rust. Hosts that
|
||||
/// keep a reference implementation treat this as "fall back", not "fail".
|
||||
#[error("unsupported by the rust path: {0}")]
|
||||
Unsupported(&'static str),
|
||||
}
|
||||
|
||||
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
|
||||
|
|
|
|||
112
litellm-rust/crates/core/src/http_utils.rs
Normal file
112
litellm-rust/crates/core/src/http_utils.rs
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
//! Header and upstream-body helpers shared by every route module.
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
|
||||
/// Bound an upstream error body before it crosses a host boundary, so provider
|
||||
/// bodies stay data-minimized.
|
||||
pub fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(UPSTREAM_ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
|
||||
pub fn string_headers(
|
||||
context: &'static str,
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"{context} extra_headers.{key} must be a string, got {}",
|
||||
json_type_name(&value)
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn has_header(headers: &[(String, String)], name: &str) -> bool {
|
||||
headers
|
||||
.iter()
|
||||
.any(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
|
||||
headers.iter().any(|(name, value)| {
|
||||
if !name.eq_ignore_ascii_case("authorization") {
|
||||
return false;
|
||||
}
|
||||
let value = value.trim();
|
||||
value.len() > 7
|
||||
&& value[..7].eq_ignore_ascii_case("bearer ")
|
||||
&& !value[7..].trim().is_empty()
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn truncate_leaves_short_bodies_untouched() {
|
||||
assert_eq!(truncate_error_body("short"), "short");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_bounds_long_bodies_by_characters() {
|
||||
let body = "\u{00e9}".repeat(UPSTREAM_ERROR_BODY_MAX_CHARS + 10);
|
||||
let truncated = truncate_error_body(&body);
|
||||
assert!(truncated.ends_with("... (truncated)"));
|
||||
assert_eq!(
|
||||
truncated.chars().count(),
|
||||
UPSTREAM_ERROR_BODY_MAX_CHARS + "... (truncated)".chars().count()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn string_headers_rejects_non_string_values() {
|
||||
let headers = Map::from_iter([("x-trace".to_string(), json!(7))]);
|
||||
let err = string_headers("chat completions", Some(headers)).expect_err("non-string value");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidRequest(
|
||||
"chat completions extra_headers.x-trace must be a string, got number".to_string()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_lookup_is_case_insensitive() {
|
||||
let headers = vec![("X-Api-Key".to_string(), "k".to_string())];
|
||||
assert!(has_header(&headers, "x-api-key"));
|
||||
assert!(!has_header(&headers, "authorization"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bearer_detection_requires_a_non_empty_token() {
|
||||
assert!(has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Bearer abc".to_string()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Bearer ".to_string()
|
||||
)]));
|
||||
assert!(!has_bearer_auth(&[(
|
||||
"Authorization".to_string(),
|
||||
"Basic abc".to_string()
|
||||
)]));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,8 +1,10 @@
|
|||
pub mod audio_transcription;
|
||||
pub mod caching;
|
||||
pub mod call_lifecycle;
|
||||
pub mod chat_completions;
|
||||
pub mod constants;
|
||||
pub mod error;
|
||||
pub mod http_utils;
|
||||
pub mod messages;
|
||||
pub mod ocr;
|
||||
pub mod providers;
|
||||
|
|
|
|||
|
|
@ -1,19 +1,15 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
use crate::error::CoreResult;
|
||||
use crate::http_utils::string_headers as shared_string_headers;
|
||||
use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
|
||||
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
|
||||
|
||||
use super::transformation::AnthropicMessagesProviderConfig;
|
||||
|
||||
pub(super) fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body};
|
||||
|
||||
const HEADER_CONTEXT: &str = "messages";
|
||||
|
||||
pub(super) fn messages_provider_config(
|
||||
provider: &str,
|
||||
|
|
@ -28,37 +24,5 @@ pub(super) fn messages_provider_config(
|
|||
pub(super) fn string_headers(
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> CoreResult<Vec<(String, String)>> {
|
||||
extra_headers
|
||||
.unwrap_or_default()
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
value
|
||||
.as_str()
|
||||
.map(|value| (key.clone(), value.to_string()))
|
||||
.ok_or_else(|| {
|
||||
CoreError::InvalidRequest(format!(
|
||||
"messages extra_headers.{key} must be a string, got {}",
|
||||
json_type_name(&value)
|
||||
))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
|
||||
headers
|
||||
.iter()
|
||||
.any(|(key, _)| key.eq_ignore_ascii_case(name))
|
||||
}
|
||||
|
||||
pub(super) fn has_bearer_auth(headers: &[(String, String)]) -> bool {
|
||||
headers.iter().any(|(name, value)| {
|
||||
if !name.eq_ignore_ascii_case("authorization") {
|
||||
return false;
|
||||
}
|
||||
let value = value.trim();
|
||||
value.len() > 7
|
||||
&& value[..7].eq_ignore_ascii_case("bearer ")
|
||||
&& !value[7..].trim().is_empty()
|
||||
})
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,444 @@
|
|||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: Value) -> Vec<ChatMessage> {
|
||||
serde_json::from_value(value).expect("valid messages")
|
||||
}
|
||||
|
||||
fn params(value: Value) -> Map<String, Value> {
|
||||
match value {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn transform(model: &str, msgs: Value, opts: Value) -> Value {
|
||||
ANTHROPIC_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_request(model, messages(msgs), params(opts))
|
||||
.expect("request transforms")
|
||||
.body
|
||||
}
|
||||
|
||||
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
|
||||
ANTHROPIC_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_response("claude-sonnet-4-5", ProviderChatResponseData { body })
|
||||
}
|
||||
|
||||
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
|
||||
ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_the_messages_body_python_builds() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({"max_tokens": 128, "temperature": 0.2}),
|
||||
);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi"}]}
|
||||
],
|
||||
"system": [{"type": "text", "text": "be terse"}],
|
||||
"max_tokens": 128,
|
||||
"temperature": 0.2
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn omits_system_when_no_system_message_is_present() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"max_tokens": 16}),
|
||||
);
|
||||
assert!(body.get("system").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_consecutive_turns_and_wraps_every_text_in_a_block() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([
|
||||
{"role": "user", "content": "one"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "two"}]},
|
||||
{"role": "assistant", "content": "ack"}
|
||||
]),
|
||||
json!({"max_tokens": 16}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": [
|
||||
{"type": "text", "text": "one"},
|
||||
{"type": "text", "text": "two"}
|
||||
]},
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "ack"}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn right_strips_a_trailing_assistant_prefill_like_python() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "Argentina "}
|
||||
]),
|
||||
json!({"max_tokens": 16}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["messages"][1]["content"][0]["text"],
|
||||
json!("Argentina")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passes_every_supported_param_through_untouched() {
|
||||
let body = transform(
|
||||
"claude-sonnet-4-5",
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"max_tokens": 64,
|
||||
"temperature": 0.1,
|
||||
"top_p": 0.9,
|
||||
"stop_sequences": ["STOP"]
|
||||
}),
|
||||
);
|
||||
assert_eq!(body["max_tokens"], json!(64));
|
||||
assert_eq!(body["temperature"], json!(0.1));
|
||||
assert_eq!(body["top_p"], json!(0.9));
|
||||
assert_eq!(body["stop_sequences"], json!(["STOP"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_top_k_because_python_gates_it_by_model_below_this_point() {
|
||||
// `temperature` and `top_p` arrive already resolved, because
|
||||
// `map_openai_params` applies `_apply_sampling_param` to them before the
|
||||
// gate runs. `top_k` bypasses that and is gated inside `transform_request`,
|
||||
// the function this route replaces, so forwarding it would send `top_k` to
|
||||
// a model that removed sampling params and take a 400 after the call, where
|
||||
// Python drops it and succeeds.
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"top_k": 40})
|
||||
),
|
||||
Some(Unsupported("unrecognized request parameter"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_streaming_before_anything_else() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true, "max_tokens": 16})
|
||||
),
|
||||
Some(Unsupported("streaming"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_an_explicit_stream_false() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": false, "max_tokens": 16})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_any_param_outside_the_allowlist() {
|
||||
for param in [
|
||||
json!({"tools": []}),
|
||||
json!({"tool_choice": {"type": "auto"}}),
|
||||
json!({"thinking": {"type": "enabled"}}),
|
||||
json!({"system": "injected"}),
|
||||
json!({"metadata": {"user_id": "u1"}}),
|
||||
json!({"output_config": {"effort": "high"}}),
|
||||
] {
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": "hi"}]), param.clone()),
|
||||
Some(Unsupported("unrecognized request parameter")),
|
||||
"expected {param} to decline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_tool_calls_tool_results_and_multimodal_content() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": null, "tool_calls": [
|
||||
{"id": "c1", "type": "function",
|
||||
"function": {"name": "f", "arguments": "{}"}}
|
||||
]}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("unrecognized message field"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "ok"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("unrecognized message field"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": "https://x/y.png"}}
|
||||
]}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("non-text message content"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [
|
||||
{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}
|
||||
]}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("non-text message content"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_message_whose_content_list_is_empty() {
|
||||
// An empty list passes every per-part check, so without this it would reach
|
||||
// the provider as an empty `content` array and fail after the call rather
|
||||
// than declining to Python before it.
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": []}]), json!({})),
|
||||
Some(Unsupported("message without content"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]),
|
||||
json!({})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_conversation_that_does_not_open_on_a_user_turn() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "assistant", "content": "prefill"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("conversation does not open on a user turn"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_a_plain_text_conversation() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "again"}]}
|
||||
]),
|
||||
json!({"max_tokens": 16, "temperature": 0.5})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalizes_a_text_response_into_openai_shape() {
|
||||
let response = transform_response(json!({
|
||||
"id": "msg_123",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5-20260101",
|
||||
"content": [{"type": "text", "text": "hello"}, {"type": "text", "text": " there"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": null,
|
||||
"usage": {"input_tokens": 11, "output_tokens": 4}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(response.model, "claude-sonnet-4-5-20260101");
|
||||
assert_eq!(response.choices.len(), 1);
|
||||
assert_eq!(response.choices[0].index, 0);
|
||||
assert_eq!(response.choices[0].message.role, "assistant");
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("hello there")
|
||||
);
|
||||
assert_eq!(response.choices[0].finish_reason, "stop");
|
||||
assert_eq!(response.usage.prompt_tokens, 11);
|
||||
assert_eq!(response.usage.completion_tokens, 4);
|
||||
assert_eq!(response.usage.total_tokens, 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn folds_cache_tokens_into_prompt_tokens_like_python() {
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 2,
|
||||
"cache_read_input_tokens": 5,
|
||||
"cache_creation_input_tokens": 3
|
||||
}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.usage.prompt_tokens, 18);
|
||||
assert_eq!(response.usage.total_tokens, 20);
|
||||
assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5);
|
||||
assert_eq!(
|
||||
response.usage.prompt_tokens_details.cache_creation_tokens,
|
||||
3
|
||||
);
|
||||
assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_max_tokens_stop_reason_to_length() {
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "max_tokens",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.choices[0].finish_reason, "length");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_refusal_returns_the_completion_python_returns() {
|
||||
// `refusal` is a stop_reason, not a content block type, so the content is
|
||||
// ordinary text and this normalizes rather than declining. Python maps it
|
||||
// to content_filter in _FINISH_REASON_MAP and returns the completion.
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "I can't help with that."}],
|
||||
"stop_reason": "refusal",
|
||||
"usage": {"input_tokens": 9, "output_tokens": 6}
|
||||
}))
|
||||
.expect("a refusal still transforms");
|
||||
assert_eq!(response.choices[0].finish_reason, "content_filter");
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("I can't help with that.")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_no_content_rather_than_an_empty_string() {
|
||||
let response = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.choices[0].message.content, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn response_carries_no_id_so_python_keeps_its_chatcmpl_id() {
|
||||
let response = transform_response(json!({
|
||||
"id": "msg_should_not_leak",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "text", "text": "hi"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
let value = serde_json::to_value(response).expect("serializable");
|
||||
assert!(
|
||||
value.get("id").is_none(),
|
||||
"the rust response must not carry an id, got {value}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_response_carrying_a_non_text_block() {
|
||||
let err = transform_response(json!({
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}],
|
||||
"stop_reason": "tool_use",
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1}
|
||||
}))
|
||||
.expect_err("non-text block");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::Unsupported("non-text response content block")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_on_a_response_missing_required_fields() {
|
||||
assert_eq!(
|
||||
transform_response(json!("nope")).expect_err("not an object"),
|
||||
CoreError::InvalidResponse("messages response is not an object".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"model": "m", "usage": {}})).expect_err("no content"),
|
||||
CoreError::MissingField("content")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"model": "m", "content": []})).expect_err("no usage"),
|
||||
CoreError::MissingField("usage")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"content": [], "usage": {}})).expect_err("no model"),
|
||||
CoreError::MissingField("model")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_the_messages_url_and_x_api_key_auth() {
|
||||
let config = &ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None)
|
||||
.expect("url builds"),
|
||||
"https://api.anthropic.com/v1/messages"
|
||||
);
|
||||
assert_eq!(
|
||||
config
|
||||
.auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None)
|
||||
.expect("auth resolves"),
|
||||
ChatCompletionsAuth::Header {
|
||||
name: "x-api-key",
|
||||
value: "sk-x".to_string()
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
config.default_headers(),
|
||||
&[
|
||||
("anthropic-version", "2023-06-01"),
|
||||
("content-type", "application/json"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,211 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::chat_completions::conversation::{Conversation, build_conversation};
|
||||
use crate::chat_completions::transformation::{
|
||||
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
|
||||
unsupported_param,
|
||||
};
|
||||
use crate::chat_completions::types::{
|
||||
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatMessage,
|
||||
ProviderChatRequestData, ProviderChatResponseData,
|
||||
};
|
||||
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
use crate::providers::anthropic::messages::transformation::{
|
||||
complete_anthropic_url, resolve_anthropic_api_key,
|
||||
};
|
||||
|
||||
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
|
||||
|
||||
/// Anthropic parameter names, post `map_openai_params`, that the Rust path can
|
||||
/// place verbatim in the Messages body.
|
||||
///
|
||||
/// `top_k` is deliberately absent even though the Messages API takes it.
|
||||
/// `temperature` and `top_p` reach this gate already resolved, because
|
||||
/// `map_openai_params` runs first and applies `_apply_sampling_param` to them.
|
||||
/// `top_k` bypasses `map_openai_params` entirely, so Python applies that same
|
||||
/// per-model gate inside `transform_request`, the function this route replaces.
|
||||
/// Forwarding it would send `top_k` to a model that removed sampling params and
|
||||
/// take a 400 after the call, where Python drops it and succeeds.
|
||||
const SUPPORTED_PARAMS: &[&str] = &["max_tokens", "temperature", "top_p", "stop_sequences"];
|
||||
|
||||
pub struct AnthropicChatCompletionsConfig;
|
||||
|
||||
pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicChatCompletionsConfig =
|
||||
AnthropicChatCompletionsConfig;
|
||||
|
||||
fn text_block(text: &str) -> Value {
|
||||
json!({"type": "text", "text": text})
|
||||
}
|
||||
|
||||
fn anthropic_body(model: &str, conversation: &Conversation, params: Map<String, Value>) -> Value {
|
||||
let messages: Vec<Value> = conversation
|
||||
.turns
|
||||
.iter()
|
||||
.map(|turn| {
|
||||
json!({
|
||||
"role": turn.role.as_str(),
|
||||
"content": turn.texts.iter().map(|text| text_block(text)).collect::<Vec<_>>(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let system: Vec<Value> = conversation.system.iter().map(|s| text_block(s)).collect();
|
||||
|
||||
let body = Map::from_iter(
|
||||
[
|
||||
("model".to_string(), json!(model)),
|
||||
("messages".to_string(), json!(messages)),
|
||||
]
|
||||
.into_iter()
|
||||
// Python builds `{"model", "messages", **optional_params}` with
|
||||
// `system` already folded into optional_params, so a caller-supplied
|
||||
// key of the same name wins here too.
|
||||
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system))))
|
||||
.chain(params),
|
||||
);
|
||||
Value::Object(body)
|
||||
}
|
||||
|
||||
impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
Ok(complete_anthropic_url(api_base, env_lookup))
|
||||
}
|
||||
|
||||
fn auth(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth> {
|
||||
Ok(ChatCompletionsAuth::Header {
|
||||
name: "x-api-key",
|
||||
value: resolve_anthropic_api_key(api_key, env_lookup)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[
|
||||
("anthropic-version", "2023-06-01"),
|
||||
("content-type", "application/json"),
|
||||
]
|
||||
}
|
||||
|
||||
/// An OAuth bearer is the whole credential: Python's `validate_environment`
|
||||
/// authenticates with it and drops `x-api-key` rather than resolving one, so
|
||||
/// the resolved key must not be applied over the top. Any other forwarded
|
||||
/// `authorization` is unrelated to this header and does not defer, which is
|
||||
/// also what Python does: it sends the deployment's `x-api-key` alongside.
|
||||
fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool {
|
||||
headers.iter().any(|(name, value)| {
|
||||
name.eq_ignore_ascii_case("authorization")
|
||||
&& value
|
||||
.strip_prefix("Bearer ")
|
||||
.is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX))
|
||||
})
|
||||
}
|
||||
|
||||
fn supported_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_PARAMS
|
||||
}
|
||||
|
||||
fn unsupported_reason(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
unsupported_param(SUPPORTED_PARAMS, &[], optional_params)
|
||||
.or_else(|| messages.iter().find_map(unsupported_message))
|
||||
// Anthropic rejects a request whose first turn is not a user turn.
|
||||
// Python only repairs that under `litellm.modify_params`, which the
|
||||
// core cannot observe, so decline instead of guessing.
|
||||
.or_else(|| {
|
||||
(!build_conversation(messages).opens_on_user_turn())
|
||||
.then_some(Unsupported("conversation does not open on a user turn"))
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData> {
|
||||
Ok(ProviderChatRequestData {
|
||||
body: anthropic_body(model, &build_conversation(&messages), optional_params),
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
_model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = response.body.as_object().ok_or_else(|| {
|
||||
CoreError::InvalidResponse("messages response is not an object".into())
|
||||
})?;
|
||||
|
||||
let content = body
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.ok_or(CoreError::MissingField("content"))?;
|
||||
// The route declines tool and thinking requests, so a non-text block
|
||||
// means the response carries something this path never asked for.
|
||||
// Decline rather than silently dropping it; the host falls back.
|
||||
if content
|
||||
.iter()
|
||||
.any(|block| block.get("type").and_then(Value::as_str) != Some("text"))
|
||||
{
|
||||
return Err(CoreError::Unsupported("non-text response content block"));
|
||||
}
|
||||
let text: String = content
|
||||
.iter()
|
||||
.filter_map(|block| block.get("text").and_then(Value::as_str))
|
||||
.collect();
|
||||
|
||||
let usage = body
|
||||
.get("usage")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or(CoreError::MissingField("usage"))?;
|
||||
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
|
||||
|
||||
Ok(ChatCompletionsResponse {
|
||||
created: unix_now(),
|
||||
model: body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.ok_or(CoreError::MissingField("model"))?
|
||||
.to_string(),
|
||||
choices: vec![ChatCompletionsChoice {
|
||||
index: 0,
|
||||
message: ChatCompletionsChoiceMessage {
|
||||
role: "assistant".to_string(),
|
||||
content: (!text.is_empty()).then_some(text),
|
||||
},
|
||||
finish_reason: finish_reason_for(
|
||||
body.get("stop_reason")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(""),
|
||||
)
|
||||
.to_string(),
|
||||
}],
|
||||
usage: usage_from_parts(
|
||||
field("input_tokens"),
|
||||
field("output_tokens"),
|
||||
field("cache_read_input_tokens"),
|
||||
field("cache_creation_input_tokens"),
|
||||
),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests.rs"]
|
||||
mod tests;
|
||||
|
|
@ -1 +1,2 @@
|
|||
pub mod chat_completions;
|
||||
pub mod messages;
|
||||
|
|
|
|||
|
|
@ -8,11 +8,8 @@ use crate::audio_transcription::types::{
|
|||
};
|
||||
use crate::error::{CoreError, CoreResult, json_type_name};
|
||||
|
||||
use super::aws_base::AwsAuthConfig;
|
||||
use super::constants::{
|
||||
AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE,
|
||||
DEFAULT_BEDROCK_REGION,
|
||||
};
|
||||
pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region};
|
||||
use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};
|
||||
|
||||
const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"];
|
||||
|
||||
|
|
@ -21,64 +18,6 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig =
|
|||
|
||||
pub struct BedrockAudioTranscriptionConfig;
|
||||
|
||||
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
||||
let mut stripped = model;
|
||||
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let mut region = None;
|
||||
if let Some((candidate, remainder)) = stripped.split_once('/')
|
||||
&& is_bedrock_region(candidate)
|
||||
{
|
||||
region = Some(candidate.to_string());
|
||||
stripped = remainder;
|
||||
}
|
||||
for prefix in ["nova-2/", "nova/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if region.is_none() {
|
||||
region = stripped
|
||||
.strip_prefix("arn:")
|
||||
.and_then(|value| value.split(':').nth(3))
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
}
|
||||
(stripped.to_string(), region)
|
||||
}
|
||||
|
||||
fn is_bedrock_region(value: &str) -> bool {
|
||||
value.len() > 3
|
||||
&& value.contains('-')
|
||||
&& value
|
||||
.chars()
|
||||
.all(|char| char.is_ascii_alphanumeric() || char == '-')
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_region(
|
||||
model_region: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
if let Some(region) = optional_params
|
||||
.get("aws_region_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return region.to_string();
|
||||
}
|
||||
if let Some(region) = model_region {
|
||||
return region.to_string();
|
||||
}
|
||||
env_lookup(AWS_REGION_NAME)
|
||||
.or_else(|| env_lookup(AWS_REGION))
|
||||
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
|
||||
}
|
||||
|
||||
fn audio_fields(audio: Value) -> CoreResult<(String, String)> {
|
||||
let object = audio.as_object().ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
|
|
@ -203,32 +142,6 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn aws_auth_config(
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> AwsAuthConfig {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
};
|
||||
let env = |key: &str| env_lookup(key);
|
||||
AwsAuthConfig {
|
||||
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
|
||||
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
|
||||
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
|
||||
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
|
||||
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
|
||||
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
|
||||
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
|
||||
web_identity_token: value("aws_web_identity_token")
|
||||
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
|
||||
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
|
||||
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
|
|||
|
|
@ -12,13 +12,15 @@ use aws_sigv4::http_request::{
|
|||
};
|
||||
use aws_sigv4::sign::v4;
|
||||
use aws_smithy_runtime_api::client::identity::Identity;
|
||||
use serde_json::{Map, Value};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use super::constants::{
|
||||
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
|
||||
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_STS_ENDPOINT,
|
||||
AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, BEDROCK_SERVICE,
|
||||
DEFAULT_SESSION_NAME_PREFIX,
|
||||
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
|
||||
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
|
||||
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
|
||||
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
|
||||
SIGV4_COMPUTED_HEADER_NAMES,
|
||||
};
|
||||
|
||||
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
|
||||
|
|
@ -401,6 +403,33 @@ fn default_session_name() -> String {
|
|||
format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}")
|
||||
}
|
||||
|
||||
/// The subset of `headers` SigV4 should cover.
|
||||
///
|
||||
/// Python signs only these and reattaches the rest afterwards, so a forwarded
|
||||
/// client header cannot change the canonical request and invalidate the
|
||||
/// signature. Signing everything instead makes the request 403 on a header the
|
||||
/// caller supplied, on a deployment that works on the Python path.
|
||||
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
let name = name.to_ascii_lowercase();
|
||||
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
|
||||
|| name.starts_with("x-amz-")
|
||||
|| name.starts_with("x-amzn-")
|
||||
})
|
||||
.map(|(name, value)| (name.clone(), value.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Whether the signer produces `name` itself.
|
||||
///
|
||||
/// Python's reattach loop skips these, so a caller-supplied copy never reaches
|
||||
/// the wire next to the computed one.
|
||||
pub fn is_sigv4_computed_header(name: &str) -> bool {
|
||||
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
|
||||
}
|
||||
|
||||
pub fn sign_bedrock_post(
|
||||
url: &str,
|
||||
body: &[u8],
|
||||
|
|
@ -441,6 +470,121 @@ pub fn sign_bedrock_post(
|
|||
.collect())
|
||||
}
|
||||
|
||||
/// Model-id and region parsing shared by every Bedrock route.
|
||||
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
||||
let mut stripped = model;
|
||||
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
let mut region = None;
|
||||
if let Some((candidate, remainder)) = stripped.split_once('/')
|
||||
&& is_bedrock_region(candidate)
|
||||
{
|
||||
region = Some(candidate.to_string());
|
||||
stripped = remainder;
|
||||
}
|
||||
for prefix in ["nova-2/", "nova/"] {
|
||||
if let Some(value) = stripped.strip_prefix(prefix) {
|
||||
stripped = value;
|
||||
break;
|
||||
}
|
||||
}
|
||||
if region.is_none() {
|
||||
// Python splits the whole ARN and takes field 3, the region. Stripping
|
||||
// `arn:` first shifts every field down one, so the region is field 2
|
||||
// here; field 3 is the account id.
|
||||
region = stripped
|
||||
.strip_prefix("arn:")
|
||||
.and_then(|value| value.split(':').nth(2))
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
}
|
||||
(stripped.to_string(), region)
|
||||
}
|
||||
|
||||
fn is_bedrock_region(value: &str) -> bool {
|
||||
value.len() > 3
|
||||
&& value.contains('-')
|
||||
&& value
|
||||
.chars()
|
||||
.all(|char| char.is_ascii_alphanumeric() || char == '-')
|
||||
}
|
||||
|
||||
pub fn resolve_bedrock_region(
|
||||
model_region: Option<&str>,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
if let Some(region) = optional_params
|
||||
.get("aws_region_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return region.to_string();
|
||||
}
|
||||
if let Some(region) = model_region {
|
||||
return region.to_string();
|
||||
}
|
||||
env_lookup(AWS_REGION_NAME)
|
||||
.or_else(|| env_lookup(AWS_REGION))
|
||||
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
|
||||
}
|
||||
|
||||
pub fn aws_auth_config(
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> AwsAuthConfig {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
};
|
||||
let env = |key: &str| env_lookup(key);
|
||||
AwsAuthConfig {
|
||||
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
|
||||
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
|
||||
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
|
||||
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
|
||||
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
|
||||
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
|
||||
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
|
||||
web_identity_token: value("aws_web_identity_token")
|
||||
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
|
||||
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
|
||||
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Credentials a host resolved through its own chain and handed down verbatim.
|
||||
///
|
||||
/// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads
|
||||
/// profiles, STS and boto sessions) passes the result here so the core signs
|
||||
/// with exactly those. Without this the core would re-derive from ambient
|
||||
/// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the
|
||||
/// environment outranks explicit keys in [`classify_auth`] and the two sides
|
||||
/// would sign as different principals.
|
||||
pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option<Credentials> {
|
||||
let value = |key: &str| {
|
||||
optional_params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
};
|
||||
let access_key_id = value("aws_access_key_id")?;
|
||||
let secret_access_key = value("aws_secret_access_key")?;
|
||||
Some(Credentials::new(
|
||||
access_key_id,
|
||||
secret_access_key,
|
||||
value("aws_session_token").map(str::to_string),
|
||||
None,
|
||||
"litellm-host-supplied",
|
||||
))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
@ -458,6 +602,18 @@ mod tests {
|
|||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_the_region_field_of_a_model_arn_not_the_account_id() {
|
||||
// Python's `_get_aws_region_from_model_arn` splits the whole ARN and
|
||||
// takes field 3. Stripping `arn:` first shifts every field down one, so
|
||||
// the region is field 2 here. Taking field 3 after the strip returns
|
||||
// the account id, which is not a region at all.
|
||||
let (_, region) = bedrock_model_id_and_region(
|
||||
"bedrock/arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2",
|
||||
);
|
||||
assert_eq!(region.as_deref(), Some("us-west-2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classification_preserves_python_precedence() {
|
||||
let config = AwsAuthConfig {
|
||||
|
|
@ -610,6 +766,52 @@ mod tests {
|
|||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_forwarded_client_header_is_not_folded_into_the_signature() {
|
||||
// Python signs only the AWS header set, so a header a caller forwarded
|
||||
// cannot change the canonical request. Signing it instead makes the
|
||||
// request 403 the moment anything on the wire rewrites or drops it.
|
||||
let (url, body, mut headers) = parity_inputs();
|
||||
headers.insert("x-request-id".to_string(), "abc-123".to_string());
|
||||
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
|
||||
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
|
||||
let signable = aws_signature_headers(&headers);
|
||||
|
||||
assert!(!signable.contains_key("x-request-id"));
|
||||
assert!(!signable.contains_key("Accept-Encoding"));
|
||||
// The AWS-prefixed one is genuinely part of the signature.
|
||||
assert!(signable.contains_key("x-amzn-trace-id"));
|
||||
assert!(signable.contains_key("Content-Type"));
|
||||
|
||||
let credentials = Credentials::new(
|
||||
"AKIDEXAMPLE",
|
||||
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
|
||||
None,
|
||||
None,
|
||||
"test",
|
||||
);
|
||||
let signed = sign_bedrock_post(
|
||||
&url,
|
||||
&body,
|
||||
&signable,
|
||||
"us-east-1",
|
||||
&credentials,
|
||||
SystemTime::UNIX_EPOCH,
|
||||
)
|
||||
.expect("signs");
|
||||
let authorization = signed
|
||||
.get("Authorization")
|
||||
.expect("carries an authorization header");
|
||||
assert!(
|
||||
!authorization.contains("x-request-id"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
assert!(
|
||||
!authorization.contains("accept-encoding"),
|
||||
"forwarded header reached SignedHeaders: {authorization}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signing_matches_botocore_golden_vector() {
|
||||
let (url, body, headers) = parity_inputs();
|
||||
|
|
|
|||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,580 @@
|
|||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn messages(value: Value) -> Vec<ChatMessage> {
|
||||
serde_json::from_value(value).expect("valid messages")
|
||||
}
|
||||
|
||||
fn params(value: Value) -> Map<String, Value> {
|
||||
match value {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("params must be an object, got {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn transform(msgs: Value, opts: Value) -> Value {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_request(
|
||||
"anthropic.claude-sonnet-4-5-v1:0",
|
||||
messages(msgs),
|
||||
params(opts),
|
||||
)
|
||||
.expect("request transforms")
|
||||
.body
|
||||
}
|
||||
|
||||
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG.transform_response(
|
||||
"anthropic.claude-sonnet-4-5-v1:0",
|
||||
ProviderChatResponseData { body },
|
||||
)
|
||||
}
|
||||
|
||||
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), ¶ms(opts))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_the_converse_body_python_builds() {
|
||||
let body = transform(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({"maxTokens": 128, "temperature": 0.2}),
|
||||
);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"inferenceConfig": {"maxTokens": 128, "temperature": 0.2},
|
||||
"messages": [{"role": "user", "content": [{"text": "hi"}]}],
|
||||
"system": [{"text": "be terse"}]
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn always_emits_inference_config_even_when_empty() {
|
||||
let body = transform(json!([{"role": "user", "content": "hi"}]), json!({}));
|
||||
assert_eq!(body["inferenceConfig"], json!({}));
|
||||
assert!(body.get("system").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn places_only_inference_params_in_inference_config() {
|
||||
let body = transform(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({
|
||||
"maxTokens": 64,
|
||||
"temperature": 0.1,
|
||||
"topP": 0.9,
|
||||
"stopSequences": ["STOP"]
|
||||
}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["inferenceConfig"],
|
||||
json!({"maxTokens": 64, "temperature": 0.1, "topP": 0.9, "stopSequences": ["STOP"]})
|
||||
);
|
||||
assert!(body.get("additionalModelRequestFields").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_consecutive_user_turns_into_one_message() {
|
||||
let body = transform(
|
||||
json!([
|
||||
{"role": "user", "content": "one"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "two"}]},
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "three"}
|
||||
]),
|
||||
json!({}),
|
||||
);
|
||||
assert_eq!(
|
||||
body["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": [{"text": "one"}, {"text": "two"}]},
|
||||
{"role": "assistant", "content": [{"text": "ack"}]},
|
||||
{"role": "user", "content": [{"text": "three"}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_streaming() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"stream": true})
|
||||
),
|
||||
Some(Unsupported("streaming"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_top_k_because_python_routes_it_by_base_model() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
json!({"topK": 40})
|
||||
),
|
||||
Some(Unsupported("unrecognized request parameter"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_tools_and_other_params_outside_the_allowlist() {
|
||||
for param in [
|
||||
json!({"tools": []}),
|
||||
json!({"tool_choice": {"auto": {}}}),
|
||||
json!({"thinking": {"type": "enabled"}}),
|
||||
json!({"requestMetadata": {"k": "v"}}),
|
||||
json!({"outputConfig": {}}),
|
||||
json!({"_parallel_tool_use_config": {}}),
|
||||
] {
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": "hi"}]), param.clone()),
|
||||
Some(Unsupported("unrecognized request parameter")),
|
||||
"expected {param} to decline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_blank_text_rather_than_substituting_the_anthropic_placeholder() {
|
||||
for content in [
|
||||
json!(""),
|
||||
json!(" "),
|
||||
json!([{"type": "text", "text": " "}]),
|
||||
] {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": content}, {"role": "user", "content": "hi"}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("blank message text")),
|
||||
"expected blank content {content} to decline"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_message_whose_content_list_is_empty() {
|
||||
// The blank-text check scans parts, so an empty list clears it; Converse
|
||||
// rejects an empty `content` array, which is a decline the core owes the
|
||||
// host before the call rather than an error after it.
|
||||
assert_eq!(
|
||||
reason(json!([{"role": "user", "content": []}]), json!({})),
|
||||
Some(Unsupported("message without content"))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]),
|
||||
json!({})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_conversation_that_opens_or_closes_on_an_assistant_turn() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "assistant", "content": "prefill"},
|
||||
{"role": "user", "content": "hi"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported(
|
||||
"conversation does not run user turn to user turn"
|
||||
))
|
||||
);
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "prefill"}
|
||||
]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported(
|
||||
"conversation does not run user turn to user turn"
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_a_user_to_user_text_conversation() {
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([
|
||||
{"role": "system", "content": "be terse"},
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
{"role": "user", "content": "again"}
|
||||
]),
|
||||
json!({"maxTokens": 16})
|
||||
),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_the_converse_url_from_the_region_in_the_model_id() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| {
|
||||
None
|
||||
})
|
||||
.expect("url builds"),
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_the_region_env_then_the_default_region() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
let with_env = |key: &str| (key == "AWS_REGION_NAME").then(|| "eu-west-1".to_string());
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env)
|
||||
.expect("url builds"),
|
||||
"https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None)
|
||||
.expect("url builds"),
|
||||
"https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefers_an_explicit_runtime_endpoint_over_the_api_base() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
let overrides = params(json!({"aws_bedrock_runtime_endpoint": "https://vpce.internal/"}));
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(
|
||||
Some("https://ignored.example"),
|
||||
"anthropic.claude-v2",
|
||||
&overrides,
|
||||
&|_| None
|
||||
)
|
||||
.expect("url builds"),
|
||||
"https://vpce.internal/model/anthropic.claude-v2/converse"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn signs_with_sigv4_in_the_resolved_region() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
assert_eq!(
|
||||
config
|
||||
.auth(
|
||||
None,
|
||||
"eu-central-1/anthropic.claude-v2",
|
||||
&Map::new(),
|
||||
&|_| None
|
||||
)
|
||||
.expect("auth resolves"),
|
||||
ChatCompletionsAuth::AwsSigV4 {
|
||||
region: "eu-central-1".to_string()
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() {
|
||||
// Python's get_request_headers reads `api_key` as the Bedrock bearer token
|
||||
// and only falls back to the env when the caller passed none, so each case
|
||||
// pins one of its precedence rules. Signing as the host principal when a
|
||||
// bearer identity is configured would cross an account and quota boundary.
|
||||
let bedrock_env =
|
||||
|key: &str| (key == "AWS_BEARER_TOKEN_BEDROCK").then(|| "from-env".to_string());
|
||||
let no_env = |_: &str| None;
|
||||
let resolve = |api_key, env: &dyn Fn(&str) -> Option<String>| {
|
||||
BEDROCK_CHAT_COMPLETIONS_CONFIG
|
||||
.auth(
|
||||
api_key,
|
||||
"eu-central-1/anthropic.claude-v2",
|
||||
&Map::new(),
|
||||
env,
|
||||
)
|
||||
.expect("auth resolves")
|
||||
};
|
||||
let bearer = |token: &str| ChatCompletionsAuth::Bearer {
|
||||
token: token.to_string(),
|
||||
};
|
||||
let sigv4 = ChatCompletionsAuth::AwsSigV4 {
|
||||
region: "eu-central-1".to_string(),
|
||||
};
|
||||
|
||||
// A caller-supplied key is the bearer token, and outranks the env.
|
||||
assert_eq!(
|
||||
resolve(Some("bedrock-api-key"), &bedrock_env),
|
||||
bearer("bedrock-api-key")
|
||||
);
|
||||
// No key, so the env supplies it.
|
||||
assert_eq!(resolve(None, &bedrock_env), bearer("from-env"));
|
||||
// An empty key is not a bearer token, and deliberately does NOT reach for
|
||||
// the env, which is what Python's `is not None` check does.
|
||||
assert_eq!(resolve(Some(""), &bedrock_env), sigv4);
|
||||
// Whitespace is truthy in Python, so it stays a bearer token rather than
|
||||
// silently becoming a host-credentialed SigV4 request.
|
||||
assert_eq!(resolve(Some(" "), &no_env), bearer(" "));
|
||||
// Neither present, so SigV4 as before.
|
||||
assert_eq!(resolve(None, &no_env), sigv4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalizes_a_converse_response_into_openai_shape() {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"role": "assistant", "content": [
|
||||
{"text": "hello"}, {"text": " there"}
|
||||
]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
|
||||
assert_eq!(response.model, "anthropic.claude-sonnet-4-5-v1:0");
|
||||
assert_eq!(
|
||||
response.choices[0].message.content.as_deref(),
|
||||
Some("hello there")
|
||||
);
|
||||
assert_eq!(response.choices[0].finish_reason, "stop");
|
||||
assert_eq!(response.usage.prompt_tokens, 11);
|
||||
assert_eq!(response.usage.completion_tokens, 4);
|
||||
assert_eq!(response.usage.total_tokens, 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_converse_stop_reasons_python_maps() {
|
||||
for (provider_reason, expected) in [
|
||||
("end_turn", "stop"),
|
||||
("stop_sequence", "stop"),
|
||||
("max_tokens", "length"),
|
||||
("guardrail_intervened", "content_filter"),
|
||||
// Converse emits this one, and Python's `_FINISH_REASON_MAP` carries
|
||||
// it. Folding it into `stop` reports a filtered completion as a normal
|
||||
// one to anything keying on the finish reason.
|
||||
("content_filtered", "content_filter"),
|
||||
("content_filter", "content_filter"),
|
||||
] {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": provider_reason,
|
||||
"usage": {"inputTokens": 1, "outputTokens": 1}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(
|
||||
response.choices[0].finish_reason, expected,
|
||||
"stopReason {provider_reason}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_an_empty_converse_answer_as_an_empty_string_not_null() {
|
||||
// Converse assigns the joined text unconditionally
|
||||
// (`chat_completion_message["content"] = content_str`), unlike Anthropic's
|
||||
// `merged_text or None`, so an empty answer is `""` on both paths. A caller
|
||||
// calling `.strip()` on it would break on the Rust path alone. Reachable
|
||||
// through a filtered or guardrail-intervened response.
|
||||
for content in [json!([]), json!([{"text": ""}])] {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": content}},
|
||||
"stopReason": "content_filtered",
|
||||
"usage": {"inputTokens": 1, "outputTokens": 0}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.choices[0].message.content, Some(String::new()));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_the_total_tokens_converse_sent_rather_than_recomputing_them() {
|
||||
// Python reads `usage["totalTokens"]` straight through here, where Anthropic
|
||||
// has no such field and adds the two counts instead. The two agree while the
|
||||
// gate declines every cache_control request, so this is what keeps them
|
||||
// agreeing if that ever widens.
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 4, "cacheReadInputTokens": 7, "totalTokens": 14}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(
|
||||
response.usage.total_tokens, 14,
|
||||
"provider total was recomputed"
|
||||
);
|
||||
assert_eq!(response.usage.prompt_tokens, 17);
|
||||
assert_eq!(response.usage.completion_tokens, 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn falls_back_to_the_computed_total_when_converse_omits_it() {
|
||||
// Python raises a KeyError on a body with no `totalTokens`. Reporting a zero
|
||||
// instead would be a worse divergence than the one above, so the computed
|
||||
// total stands in.
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 4}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.usage.total_tokens, 14);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_cache_control_message_so_widening_the_gate_is_a_red_test() {
|
||||
// Converse only reports cache token counts when the request carries a
|
||||
// cachePoint block, which is why the provider total and the computed one
|
||||
// cannot disagree today. This is the tripwire: whoever widens the gate to
|
||||
// admit prompt caching has to come back and re-check the usage mapping
|
||||
// rather than discovering a silent number change in production.
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": [
|
||||
{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}
|
||||
]}]),
|
||||
json!({})
|
||||
),
|
||||
Some(Unsupported("non-text message content"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn folds_converse_cache_tokens_into_prompt_tokens() {
|
||||
let response = transform_response(json!({
|
||||
"output": {"message": {"content": [{"text": "x"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {
|
||||
"inputTokens": 10,
|
||||
"outputTokens": 2,
|
||||
"cacheReadInputTokens": 5,
|
||||
"cacheWriteInputTokens": 3
|
||||
}
|
||||
}))
|
||||
.expect("response transforms");
|
||||
assert_eq!(response.usage.prompt_tokens, 18);
|
||||
assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5);
|
||||
assert_eq!(
|
||||
response.usage.prompt_tokens_details.cache_creation_tokens,
|
||||
3
|
||||
);
|
||||
assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn declines_a_response_carrying_a_tool_use_block() {
|
||||
let err = transform_response(json!({
|
||||
"output": {"message": {"content": [
|
||||
{"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}}
|
||||
]}},
|
||||
"stopReason": "tool_use",
|
||||
"usage": {"inputTokens": 1, "outputTokens": 1}
|
||||
}))
|
||||
.expect_err("tool use block");
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::Unsupported("non-text response content block")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn errors_on_a_response_missing_required_fields() {
|
||||
assert_eq!(
|
||||
transform_response(json!("nope")).expect_err("not an object"),
|
||||
CoreError::InvalidResponse("converse response is not an object".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"usage": {}})).expect_err("no output"),
|
||||
CoreError::MissingField("output.message.content")
|
||||
);
|
||||
assert_eq!(
|
||||
transform_response(json!({"output": {"message": {"content": []}}})).expect_err("no usage"),
|
||||
CoreError::MissingField("usage")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_aws_call_configuration_without_serializing_it() {
|
||||
let call_config = json!({
|
||||
"maxTokens": 16,
|
||||
"aws_access_key_id": "AKIA",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_session_token": "token",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_profile_name": "litellm-stage",
|
||||
"aws_role_name": "role",
|
||||
"aws_session_name": "session",
|
||||
"aws_web_identity_token": "wit",
|
||||
"aws_sts_endpoint": "https://sts.example",
|
||||
"aws_external_id": "ext",
|
||||
"aws_bedrock_runtime_endpoint": "https://vpce.internal"
|
||||
});
|
||||
assert_eq!(
|
||||
reason(
|
||||
json!([{"role": "user", "content": "hi"}]),
|
||||
call_config.clone()
|
||||
),
|
||||
None
|
||||
);
|
||||
let body = transform(json!([{"role": "user", "content": "hi"}]), call_config);
|
||||
assert_eq!(
|
||||
body,
|
||||
json!({
|
||||
"inferenceConfig": {"maxTokens": 16},
|
||||
"messages": [{"role": "user", "content": [{"text": "hi"}]}]
|
||||
}),
|
||||
"aws call configuration must not reach the Converse body"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_a_complete_converse_url_untouched() {
|
||||
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
|
||||
let already_built =
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-v2%3A0/converse";
|
||||
assert_eq!(
|
||||
config
|
||||
.complete_url(
|
||||
Some(already_built),
|
||||
"anthropic.claude-v2",
|
||||
&Map::new(),
|
||||
&|_| None
|
||||
)
|
||||
.expect("url builds"),
|
||||
already_built,
|
||||
"a host that encoded the model id itself must not have it re-derived"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn host_supplied_credentials_outrank_ambient_profile_and_role_state() {
|
||||
use crate::providers::bedrock::aws_base::host_supplied_credentials;
|
||||
|
||||
let supplied = params(json!({
|
||||
"aws_access_key_id": "AKIAHOST",
|
||||
"aws_secret_access_key": "hostsecret",
|
||||
"aws_session_token": "hosttoken"
|
||||
}));
|
||||
let credentials = host_supplied_credentials(&supplied).expect("host credentials");
|
||||
assert_eq!(credentials.access_key_id(), "AKIAHOST");
|
||||
assert_eq!(credentials.secret_access_key(), "hostsecret");
|
||||
assert_eq!(credentials.session_token(), Some("hosttoken"));
|
||||
|
||||
// Without a full static pair there is nothing to honor, so the core falls
|
||||
// back to deriving credentials itself.
|
||||
assert!(host_supplied_credentials(¶ms(json!({"aws_access_key_id": "AKIA"}))).is_none());
|
||||
assert!(
|
||||
host_supplied_credentials(¶ms(
|
||||
json!({"aws_access_key_id": " ", "aws_secret_access_key": "s"})
|
||||
))
|
||||
.is_none()
|
||||
);
|
||||
assert!(host_supplied_credentials(&Map::new()).is_none());
|
||||
}
|
||||
|
|
@ -0,0 +1,297 @@
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation};
|
||||
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
|
||||
use crate::chat_completions::transformation::{
|
||||
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
|
||||
unsupported_param,
|
||||
};
|
||||
use crate::chat_completions::types::{
|
||||
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse,
|
||||
ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData,
|
||||
ProviderChatResponseData,
|
||||
};
|
||||
use crate::error::{CoreError, CoreResult};
|
||||
|
||||
use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region};
|
||||
use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE};
|
||||
|
||||
/// Converse parameter names, post `map_openai_params`, that the Rust path can
|
||||
/// place verbatim in `inferenceConfig`.
|
||||
///
|
||||
/// `topK` is deliberately absent: Python routes it to
|
||||
/// `additionalModelRequestFields` for Anthropic base models and to
|
||||
/// `inferenceConfig` otherwise, and that branch reads the model catalog the
|
||||
/// core cannot see.
|
||||
const SUPPORTED_PARAMS: &[&str] = &["maxTokens", "temperature", "topP", "stopSequences"];
|
||||
|
||||
/// Params that belong in `inferenceConfig`, in the order Python's
|
||||
/// `AmazonConverseConfig` declares them, so bodies compare cleanly.
|
||||
const INFERENCE_CONFIG_PARAMS: &[&str] = SUPPORTED_PARAMS;
|
||||
|
||||
const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "aws_bedrock_runtime_endpoint";
|
||||
|
||||
/// AWS call configuration a host passes down: consumed for signing and endpoint
|
||||
/// resolution, never serialized into the Converse body.
|
||||
const CONFIG_PARAMS: &[&str] = &[
|
||||
"aws_access_key_id",
|
||||
"aws_secret_access_key",
|
||||
"aws_session_token",
|
||||
"aws_region_name",
|
||||
"aws_session_name",
|
||||
"aws_profile_name",
|
||||
"aws_role_name",
|
||||
"aws_web_identity_token",
|
||||
"aws_sts_endpoint",
|
||||
"aws_external_id",
|
||||
AWS_BEDROCK_RUNTIME_ENDPOINT,
|
||||
];
|
||||
|
||||
const CONVERSE_PATH_SUFFIX: &str = "/converse";
|
||||
|
||||
pub struct BedrockChatCompletionsConfig;
|
||||
|
||||
pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: BedrockChatCompletionsConfig =
|
||||
BedrockChatCompletionsConfig;
|
||||
|
||||
fn converse_body(conversation: &Conversation, params: &Map<String, Value>) -> Value {
|
||||
let messages: Vec<Value> = conversation
|
||||
.turns
|
||||
.iter()
|
||||
.map(|turn| {
|
||||
json!({
|
||||
"role": turn.role.as_str(),
|
||||
"content": turn.texts.iter().map(|text| json!({"text": text})).collect::<Vec<_>>(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
let inference_config = Map::from_iter(INFERENCE_CONFIG_PARAMS.iter().filter_map(|name| {
|
||||
params
|
||||
.get(*name)
|
||||
.map(|value| ((*name).to_string(), value.clone()))
|
||||
}));
|
||||
|
||||
let system: Vec<Value> = conversation
|
||||
.system
|
||||
.iter()
|
||||
.map(|text| json!({"text": text}))
|
||||
.collect();
|
||||
|
||||
Value::Object(Map::from_iter(
|
||||
[
|
||||
(
|
||||
"inferenceConfig".to_string(),
|
||||
Value::Object(inference_config),
|
||||
),
|
||||
("messages".to_string(), json!(messages)),
|
||||
]
|
||||
.into_iter()
|
||||
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))),
|
||||
))
|
||||
}
|
||||
|
||||
fn has_blank_text(message: &ChatMessage) -> bool {
|
||||
match &message.content {
|
||||
None => false,
|
||||
Some(ChatMessageContent::Text(text)) => text.trim().is_empty(),
|
||||
Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| {
|
||||
part.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.is_none_or(|text| text.trim().is_empty())
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
let (model_id, model_region) = bedrock_model_id_and_region(model);
|
||||
let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup);
|
||||
let endpoint = optional_params
|
||||
.get(AWS_BEDROCK_RUNTIME_ENDPOINT)
|
||||
.and_then(Value::as_str)
|
||||
.or(api_base)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", ®ion));
|
||||
let endpoint = endpoint.trim_end_matches('/');
|
||||
// A host that already built the full Converse URL (LiteLLM's Python
|
||||
// path encodes the model id itself) passes it through untouched, the
|
||||
// way the Anthropic config leaves a complete `/v1/messages` URL alone.
|
||||
if endpoint.ends_with(CONVERSE_PATH_SUFFIX) {
|
||||
return Ok(endpoint.to_string());
|
||||
}
|
||||
Ok(format!("{endpoint}/model/{model_id}{CONVERSE_PATH_SUFFIX}"))
|
||||
}
|
||||
|
||||
fn auth(
|
||||
&self,
|
||||
api_key: Option<&str>,
|
||||
model: &str,
|
||||
optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<ChatCompletionsAuth> {
|
||||
// Python reads `api_key` as the Bedrock bearer token and consults the
|
||||
// env only when the caller passed none, so a caller-supplied empty key
|
||||
// falls through to SigV4 without reaching for the environment. An
|
||||
// all-whitespace token stays a bearer token here because Python sends
|
||||
// it too: treating it as absent would sign as the host principal
|
||||
// instead, which is the identity swap this branch exists to prevent.
|
||||
let bearer = match api_key {
|
||||
Some(key) => Some(key.to_string()),
|
||||
None => env_lookup(AWS_BEARER_TOKEN_BEDROCK),
|
||||
}
|
||||
.filter(|token| !token.is_empty());
|
||||
if let Some(token) = bearer {
|
||||
return Ok(ChatCompletionsAuth::Bearer { token });
|
||||
}
|
||||
let (_, model_region) = bedrock_model_id_and_region(model);
|
||||
Ok(ChatCompletionsAuth::AwsSigV4 {
|
||||
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
|
||||
})
|
||||
}
|
||||
|
||||
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[("Content-Type", "application/json")]
|
||||
}
|
||||
|
||||
fn supported_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_PARAMS
|
||||
}
|
||||
|
||||
fn config_params(&self) -> &'static [&'static str] {
|
||||
CONFIG_PARAMS
|
||||
}
|
||||
|
||||
fn unsupported_reason(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
unsupported_param(SUPPORTED_PARAMS, CONFIG_PARAMS, optional_params)
|
||||
.or_else(|| messages.iter().find_map(unsupported_message))
|
||||
// Python's Converse translation drops blank text blocks instead of
|
||||
// substituting the placeholder the shared conversation builder
|
||||
// applies, so decline blank text rather than diverge.
|
||||
.or_else(|| {
|
||||
messages
|
||||
.iter()
|
||||
.any(has_blank_text)
|
||||
.then_some(Unsupported("blank message text"))
|
||||
})
|
||||
// Converse has no assistant prefill: Python inserts a continue turn
|
||||
// when a conversation opens or closes on an assistant message, and
|
||||
// only under `litellm.modify_params`, which the core cannot see.
|
||||
// Declining both ends also keeps the shared builder's final
|
||||
// assistant right-strip (an Anthropic rule) unreachable here.
|
||||
.or_else(|| {
|
||||
let conversation = build_conversation(messages);
|
||||
let ends_on_assistant = conversation
|
||||
.turns
|
||||
.last()
|
||||
.is_some_and(|turn| turn.role == TurnRole::Assistant);
|
||||
(!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported(
|
||||
"conversation does not run user turn to user turn",
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
_model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<ProviderChatRequestData> {
|
||||
Ok(ProviderChatRequestData {
|
||||
body: converse_body(&build_conversation(&messages), &optional_params),
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
) -> CoreResult<ChatCompletionsResponse> {
|
||||
let body = response.body.as_object().ok_or_else(|| {
|
||||
CoreError::InvalidResponse("converse response is not an object".into())
|
||||
})?;
|
||||
|
||||
let content = body
|
||||
.get("output")
|
||||
.and_then(|output| output.get("message"))
|
||||
.and_then(|message| message.get("content"))
|
||||
.and_then(Value::as_array)
|
||||
.ok_or(CoreError::MissingField("output.message.content"))?;
|
||||
// The route declines tool requests, so anything other than a text block
|
||||
// is something this path never asked for. Decline; the host falls back.
|
||||
if content.iter().any(|block| {
|
||||
block
|
||||
.as_object()
|
||||
.is_none_or(|block| block.len() != 1 || !block.contains_key("text"))
|
||||
}) {
|
||||
return Err(CoreError::Unsupported("non-text response content block"));
|
||||
}
|
||||
let text: String = content
|
||||
.iter()
|
||||
.filter_map(|block| block.get("text").and_then(Value::as_str))
|
||||
.collect();
|
||||
|
||||
let usage = body
|
||||
.get("usage")
|
||||
.and_then(Value::as_object)
|
||||
.ok_or(CoreError::MissingField("usage"))?;
|
||||
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
|
||||
let computed = usage_from_parts(
|
||||
field("inputTokens"),
|
||||
field("outputTokens"),
|
||||
field("cacheReadInputTokens"),
|
||||
field("cacheWriteInputTokens"),
|
||||
);
|
||||
// Converse reports `totalTokens` and Python passes it straight through,
|
||||
// where Anthropic has no such field and Python adds the two counts
|
||||
// instead, so only this provider overrides the computed total. Python
|
||||
// does a bare `usage["totalTokens"]` lookup, so a body without the key
|
||||
// raises there rather than reporting a zero; fall back to the computed
|
||||
// total, which is the closest thing to that without failing the call.
|
||||
let usage = ChatCompletionsUsage {
|
||||
total_tokens: usage
|
||||
.get("totalTokens")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(computed.total_tokens),
|
||||
..computed
|
||||
};
|
||||
|
||||
Ok(ChatCompletionsResponse {
|
||||
created: unix_now(),
|
||||
// Converse echoes no model id, so Python reports the requested one.
|
||||
model: model.to_string(),
|
||||
choices: vec![ChatCompletionsChoice {
|
||||
index: 0,
|
||||
message: ChatCompletionsChoiceMessage {
|
||||
role: "assistant".to_string(),
|
||||
// Converse assigns the joined string unconditionally, so an
|
||||
// empty response is `""` here and not `None` as it is on
|
||||
// Anthropic. A caller calling `.strip()` on it would break
|
||||
// on this path alone.
|
||||
content: Some(text),
|
||||
},
|
||||
finish_reason: finish_reason_for(
|
||||
body.get("stopReason").and_then(Value::as_str).unwrap_or(""),
|
||||
)
|
||||
.to_string(),
|
||||
}],
|
||||
usage,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests.rs"]
|
||||
mod tests;
|
||||
|
|
@ -11,6 +11,31 @@ pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
|
|||
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
|
||||
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
|
||||
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
|
||||
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
|
||||
|
||||
/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors
|
||||
/// Python's `_filter_headers_for_aws_signature` allowlist.
|
||||
pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[
|
||||
"host",
|
||||
"content-type",
|
||||
"date",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"x-amz-content-sha256",
|
||||
"x-amz-algorithm",
|
||||
"x-amz-credential",
|
||||
"x-amz-signedheaders",
|
||||
"x-amz-signature",
|
||||
];
|
||||
/// Headers the signer emits itself. Mirrors Python's `SIGV4_COMPUTED_HEADERS`,
|
||||
/// which the reattach loop skips so a caller's copy cannot ride alongside the
|
||||
/// computed one.
|
||||
pub const SIGV4_COMPUTED_HEADER_NAMES: &[&str] = &[
|
||||
"authorization",
|
||||
"x-amz-date",
|
||||
"x-amz-security-token",
|
||||
"date",
|
||||
];
|
||||
pub const BEDROCK_SERVICE: &str = "bedrock";
|
||||
pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session";
|
||||
pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2";
|
||||
|
|
|
|||
|
|
@ -5,4 +5,5 @@
|
|||
#[cfg(feature = "bedrock-auth")]
|
||||
pub mod audio_transcription;
|
||||
pub mod aws_base;
|
||||
pub mod chat_completions;
|
||||
mod constants;
|
||||
|
|
|
|||
|
|
@ -6,6 +6,10 @@ use litellm_ai_gateway::io::audio_transcription::{
|
|||
};
|
||||
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
|
||||
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
|
||||
use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
|
||||
use litellm_core::chat_completions::{
|
||||
chat_completions as run_chat_completions, chat_completions_decline_reason,
|
||||
};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::messages::messages as run_messages;
|
||||
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
|
||||
|
|
@ -16,6 +20,20 @@ use serde_json::{Map, Value};
|
|||
|
||||
mod gil;
|
||||
|
||||
pyo3::create_exception!(
|
||||
_native,
|
||||
RustBridgeDeclined,
|
||||
pyo3::exceptions::PyException,
|
||||
"The route declined before calling the provider, so the host may retry on its own path."
|
||||
);
|
||||
|
||||
pyo3::create_exception!(
|
||||
_native,
|
||||
RustUpstreamError,
|
||||
pyo3::exceptions::PyException,
|
||||
"The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response."
|
||||
);
|
||||
|
||||
type MarshaledOcrInputs = (
|
||||
Value,
|
||||
Option<Map<String, Value>>,
|
||||
|
|
@ -45,6 +63,15 @@ fn messages_response_to_py(
|
|||
json_to_py(py, value)
|
||||
}
|
||||
|
||||
fn chat_completions_response_to_py(
|
||||
py: Python<'_>,
|
||||
response: ChatCompletionsResponse,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let value =
|
||||
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
json_to_py(py, value)
|
||||
}
|
||||
|
||||
fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Auth(message) => PyValueError::new_err(message),
|
||||
|
|
@ -56,6 +83,33 @@ fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
|||
}
|
||||
}
|
||||
|
||||
/// Map a core error for a route whose host keeps a Python implementation.
|
||||
///
|
||||
/// The distinction the host needs is whether the provider was already called.
|
||||
/// Everything raised before the request goes out is safe for the host to retry
|
||||
/// on its own path; anything after it is not, because the provider has already
|
||||
/// done the work and billed for it.
|
||||
fn chat_completions_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Unsupported(_)
|
||||
| CoreError::Auth(_)
|
||||
| CoreError::InvalidProvider(_)
|
||||
| CoreError::InvalidRequest(_)
|
||||
| CoreError::InvalidType { .. }
|
||||
| CoreError::MissingField(_)
|
||||
| CoreError::Routing(_)
|
||||
// Nothing reached the provider, so serving it on Python cannot double
|
||||
// bill and is the only way the caller gets an answer at all.
|
||||
| CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()),
|
||||
CoreError::Http { status, body } => {
|
||||
RustUpstreamError::new_err((status, format!("{status}: {body}")))
|
||||
}
|
||||
CoreError::Network(message) | CoreError::InvalidResponse(message) => {
|
||||
RustUpstreamError::new_err((0u16, message))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn optional_object_to_map(
|
||||
py: Python<'_>,
|
||||
name: &'static str,
|
||||
|
|
@ -430,6 +484,143 @@ fn amessages(
|
|||
})
|
||||
}
|
||||
|
||||
type MarshaledChatCompletionsInputs = (
|
||||
Value,
|
||||
Map<String, Value>,
|
||||
Option<Map<String, Value>>,
|
||||
Option<Duration>,
|
||||
);
|
||||
|
||||
fn marshal_chat_completions_inputs(
|
||||
py: Python<'_>,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<MarshaledChatCompletionsInputs> {
|
||||
let messages = py_to_json(py, messages.bind(py))?;
|
||||
if !messages.is_array() {
|
||||
return Err(PyValueError::new_err("messages must be a list"));
|
||||
}
|
||||
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
|
||||
let extra_headers = match extra_headers {
|
||||
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
|
||||
None => None,
|
||||
};
|
||||
Ok((
|
||||
messages,
|
||||
optional_params,
|
||||
extra_headers,
|
||||
optional_timeout(timeout_seconds),
|
||||
))
|
||||
}
|
||||
|
||||
/// The decline reason for this request, or `None` when the Rust path accepts
|
||||
/// it. Resolves no credentials and performs no I/O, so a host can ask before
|
||||
/// committing to either path.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))]
|
||||
fn chat_completions_decline(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
custom_llm_provider: Option<String>,
|
||||
) -> PyResult<Option<String>> {
|
||||
let messages = py_to_json(py, messages.bind(py))?;
|
||||
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
|
||||
Ok(chat_completions_decline_reason(
|
||||
&model,
|
||||
custom_llm_provider.as_deref(),
|
||||
messages,
|
||||
&optional_params,
|
||||
)
|
||||
.map(str::to_string))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn chat_completions(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
custom_llm_provider: Option<String>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
|
||||
py,
|
||||
messages,
|
||||
optional_params,
|
||||
extra_headers,
|
||||
timeout_seconds,
|
||||
)?;
|
||||
|
||||
let result = gil::release_gil(py, || {
|
||||
pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions(
|
||||
ChatCompletionsRequest {
|
||||
model: &model,
|
||||
messages,
|
||||
optional_params,
|
||||
api_key: api_key.as_deref(),
|
||||
api_base: api_base.as_deref(),
|
||||
custom_llm_provider: custom_llm_provider.as_deref(),
|
||||
extra_headers,
|
||||
timeout,
|
||||
},
|
||||
))
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(response) => chat_completions_response_to_py(py, response),
|
||||
Err(err) => Err(chat_completions_error_to_pyerr(err)),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn achat_completions(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
messages: Py<PyAny>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
custom_llm_provider: Option<String>,
|
||||
extra_headers: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Bound<'_, PyAny>> {
|
||||
let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
|
||||
py,
|
||||
messages,
|
||||
optional_params,
|
||||
extra_headers,
|
||||
timeout_seconds,
|
||||
)?;
|
||||
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
let response = run_chat_completions(ChatCompletionsRequest {
|
||||
model: &model,
|
||||
messages,
|
||||
optional_params,
|
||||
api_key: api_key.as_deref(),
|
||||
api_base: api_base.as_deref(),
|
||||
custom_llm_provider: custom_llm_provider.as_deref(),
|
||||
extra_headers,
|
||||
timeout,
|
||||
})
|
||||
.await
|
||||
.map_err(chat_completions_error_to_pyerr)?;
|
||||
|
||||
Python::attach(|py| chat_completions_response_to_py(py, response))
|
||||
})
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let stats = PyDict::new(py);
|
||||
|
|
@ -439,12 +630,18 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
|||
|
||||
#[pymodule]
|
||||
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let py = module.py();
|
||||
module.add_function(wrap_pyfunction!(ocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(aocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(transcription, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(atranscription, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(messages, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(amessages, module)?)?;
|
||||
module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
|
||||
module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())?;
|
||||
module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(chat_completions, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(achat_completions, module)?)?;
|
||||
module.add_class::<ResponsesWebSocketConnection>()?;
|
||||
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
|
||||
Ok(())
|
||||
|
|
|
|||
|
|
@ -199,6 +199,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = (
|
|||
None # Fields to exclude from StandardLoggingPayload before callbacks receive it
|
||||
)
|
||||
log_raw_request_response: bool = False
|
||||
log_client_error_tracebacks: bool = False
|
||||
request_correlation_in_logs: bool = False
|
||||
redact_messages_in_exceptions: Optional[bool] = False
|
||||
redact_user_api_key_info: Optional[bool] = False
|
||||
|
|
@ -453,6 +454,7 @@ max_end_user_budget_id: Optional[str] = None
|
|||
# backwards compatibility — arbitrary client-supplied identifiers still
|
||||
# pass through unchanged.
|
||||
validate_end_user_id_in_db: bool = False
|
||||
block_requests_for_models_without_pricing: bool = False
|
||||
disable_end_user_cost_tracking: Optional[bool] = None
|
||||
disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
|
||||
enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
|
||||
|
|
@ -462,6 +464,11 @@ prometheus_metrics_config: Optional[List] = None
|
|||
prometheus_exclude_metrics: Optional[List[str]] = None
|
||||
prometheus_exclude_labels: Optional[List[str]] = None
|
||||
prometheus_emit_stream_label: bool = False
|
||||
prometheus_deployment_and_latency_caller_identity: Literal[
|
||||
"api_key_alias",
|
||||
"user_email",
|
||||
"both",
|
||||
] = "api_key_alias"
|
||||
# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on
|
||||
# `litellm_proxy_failed_requests_metric`. Off by default to preserve the
|
||||
# pre-unification label set so existing dashboards / recording rules keyed on
|
||||
|
|
@ -1627,6 +1634,9 @@ if TYPE_CHECKING:
|
|||
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
|
||||
)
|
||||
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
|
||||
from .llms.together_ai.chat.transformation import (
|
||||
TogetherAIChatConfig as TogetherAIChatConfig,
|
||||
)
|
||||
from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig
|
||||
from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig as VertexGeminiConfig,
|
||||
|
|
@ -1800,6 +1810,9 @@ if TYPE_CHECKING:
|
|||
from .llms.gemini.interactions.transformation import (
|
||||
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
|
||||
)
|
||||
from .llms.vertex_ai.interactions.transformation import (
|
||||
VertexAIInteractionsConfig as VertexAIInteractionsConfig,
|
||||
)
|
||||
from .llms.openai.chat.o_series_transformation import (
|
||||
OpenAIOSeriesConfig as OpenAIOSeriesConfig,
|
||||
OpenAIOSeriesConfig as OpenAIO1Config,
|
||||
|
|
|
|||
|
|
@ -177,6 +177,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"AmazonAnthropicClaudeMessagesConfig",
|
||||
"AmazonMantleMessagesConfig",
|
||||
"TogetherAIConfig",
|
||||
"TogetherAIChatConfig",
|
||||
"NLPCloudConfig",
|
||||
"VertexGeminiConfig",
|
||||
"GoogleAIStudioGeminiConfig",
|
||||
|
|
@ -242,6 +243,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"OpenRouterResponsesAPIConfig",
|
||||
"BedrockMantleResponsesAPIConfig",
|
||||
"GoogleAIStudioInteractionsConfig",
|
||||
"VertexAIInteractionsConfig",
|
||||
"OpenAIOSeriesConfig",
|
||||
"AnthropicSkillsConfig",
|
||||
"BaseSkillsAPIConfig",
|
||||
|
|
@ -740,6 +742,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
"AmazonMantleMessagesConfig",
|
||||
),
|
||||
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
|
||||
"TogetherAIChatConfig": (
|
||||
".llms.together_ai.chat.transformation",
|
||||
"TogetherAIChatConfig",
|
||||
),
|
||||
"NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"),
|
||||
"VertexGeminiConfig": (
|
||||
".llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini",
|
||||
|
|
@ -977,6 +983,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.gemini.interactions.transformation",
|
||||
"GoogleAIStudioInteractionsConfig",
|
||||
),
|
||||
"VertexAIInteractionsConfig": (
|
||||
".llms.vertex_ai.interactions.transformation",
|
||||
"VertexAIInteractionsConfig",
|
||||
),
|
||||
"OpenAIOSeriesConfig": (
|
||||
".llms.openai.chat.o_series_transformation",
|
||||
"OpenAIOSeriesConfig",
|
||||
|
|
|
|||
|
|
@ -8,6 +8,12 @@ from logging import Formatter
|
|||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE,
|
||||
MAX_STRING_LENGTH_STDOUT_LOG,
|
||||
)
|
||||
from litellm.litellm_core_utils.env_utils import get_env_int
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string, redact_structured_value
|
||||
|
|
@ -82,6 +88,24 @@ def redact_secrets(value: str) -> str:
|
|||
return _redact_string(value)
|
||||
|
||||
|
||||
def _substituted_color_message(record: logging.LogRecord) -> str | None:
|
||||
"""Render a record's ``color_message`` against its args, or None if absent.
|
||||
|
||||
uvicorn's colorized formatter re-renders `color_message` against
|
||||
record.args at emit time (see uvicorn.logging.ColourizedFormatter) instead
|
||||
of using the already-formatted record.msg, so it has to be substituted
|
||||
before args are cleared or it is later formatted with no args and prints
|
||||
the raw "%s://%s:%d" placeholders instead of the URL.
|
||||
"""
|
||||
color_message: Final = record.__dict__.get("color_message")
|
||||
if not isinstance(color_message, str) or not record.args:
|
||||
return None
|
||||
try:
|
||||
return color_message % record.args
|
||||
except TypeError:
|
||||
return color_message
|
||||
|
||||
|
||||
class SecretRedactionFilter(logging.Filter):
|
||||
"""Scrubs known secret/credential patterns from log records."""
|
||||
|
||||
|
|
@ -91,6 +115,12 @@ class SecretRedactionFilter(logging.Filter):
|
|||
if not _ENABLE_SECRET_REDACTION:
|
||||
return True
|
||||
|
||||
# Runs before args are cleared, and before the extra-field loop below
|
||||
# that redacts the substituted result.
|
||||
substituted_color_message: Final = _substituted_color_message(record)
|
||||
if substituted_color_message is not None:
|
||||
record.color_message = substituted_color_message # rebind-ok: a Filter scrubs records in place
|
||||
|
||||
try:
|
||||
record.msg = _redact_string(record.getMessage())
|
||||
record.args = None
|
||||
|
|
@ -101,7 +131,7 @@ class SecretRedactionFilter(logging.Filter):
|
|||
# Redact exception tracebacks
|
||||
if record.exc_info and record.exc_info[1] is not None:
|
||||
try:
|
||||
record.exc_text = _redact_string(self._formatter.formatException(record.exc_info))
|
||||
record.exc_text = _redact_string(record.exc_text or self._formatter.formatException(record.exc_info))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
|
@ -116,6 +146,72 @@ class SecretRedactionFilter(logging.Filter):
|
|||
_secret_filter: Final = SecretRedactionFilter()
|
||||
|
||||
|
||||
def _get_max_string_length_stdout_log() -> int:
|
||||
"""Read the limit per record so a value loaded later via proxy config
|
||||
environment_variables is honored."""
|
||||
return get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", MAX_STRING_LENGTH_STDOUT_LOG)
|
||||
|
||||
|
||||
def _stdout_truncation_marker(skipped_chars: int) -> str:
|
||||
return (
|
||||
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. "
|
||||
f"{LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE}) ..."
|
||||
)
|
||||
|
||||
|
||||
def _truncate_for_stdout_log(text: str, limit: int) -> str:
|
||||
kept_chars: Final = limit - len(_stdout_truncation_marker(len(text)))
|
||||
if kept_chars <= 0:
|
||||
return text[:limit]
|
||||
head_chars: Final = kept_chars // 2
|
||||
tail_chars: Final = kept_chars - head_chars
|
||||
return f"{text[:head_chars]}{_stdout_truncation_marker(len(text) - kept_chars)}{text[-tail_chars:]}"
|
||||
|
||||
|
||||
class StdoutLogTruncationFilter(logging.Filter):
|
||||
"""Bounds how much of an oversized log line reaches stdout.
|
||||
|
||||
A provider error string can echo the whole request payload, so one failed agentic
|
||||
request writes hundreds of KB to stdout, repeatedly as the exception propagates from
|
||||
the router to the proxy handler and into its traceback, all inline on the event loop.
|
||||
|
||||
DEBUG records pass through untouched, since dumping full payloads is the point of
|
||||
`--detailed_debug`, and logging callbacks (OTEL, Datadog, etc.) don't run through
|
||||
logging filters at all, so they still get the untruncated error.
|
||||
"""
|
||||
|
||||
_formatter = logging.Formatter()
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if record.levelno < logging.INFO:
|
||||
return True
|
||||
|
||||
limit: Final = _get_max_string_length_stdout_log()
|
||||
if limit <= 0:
|
||||
return True
|
||||
|
||||
try:
|
||||
message: Final = record.getMessage()
|
||||
except (TypeError, ValueError):
|
||||
return True
|
||||
|
||||
if len(message) > limit:
|
||||
record.msg = _truncate_for_stdout_log(message, limit) # rebind-ok: the Filter interface mutates the record
|
||||
record.args = None # rebind-ok: args are consumed by the truncated message above
|
||||
|
||||
if isinstance(record.exc_info, tuple):
|
||||
exc_text: Final = record.exc_text or self._formatter.formatException(record.exc_info)
|
||||
if len(exc_text) > limit:
|
||||
record.exc_text = _truncate_for_stdout_log( # rebind-ok: the Filter interface mutates the record
|
||||
exc_text, limit
|
||||
)
|
||||
|
||||
return True
|
||||
|
||||
|
||||
_stdout_truncation_filter: Final = StdoutLogTruncationFilter()
|
||||
|
||||
|
||||
class CorrelationContextFilter(logging.Filter):
|
||||
"""Stamps each log record with the current request's trace_id and session_id from contextvars.
|
||||
|
||||
|
|
@ -301,6 +397,7 @@ def _setup_json_exception_handlers(formatter):
|
|||
error_handler: Final = logging.StreamHandler()
|
||||
error_handler.setFormatter(formatter)
|
||||
error_handler.addFilter(_secret_filter)
|
||||
error_handler.addFilter(_stdout_truncation_filter)
|
||||
error_handler.addFilter(_correlation_filter)
|
||||
|
||||
# Setup excepthook for uncaught exceptions
|
||||
|
|
@ -365,6 +462,12 @@ verbose_router_logger.addHandler(handler)
|
|||
verbose_proxy_logger.addHandler(handler)
|
||||
verbose_logger.addHandler(handler)
|
||||
|
||||
# Filters attached to the logger, not the handler, survive callers swapping in their own
|
||||
# handlers (JSON mode, uvicorn log config, a host app's root handler).
|
||||
verbose_router_logger.addFilter(_stdout_truncation_filter)
|
||||
verbose_proxy_logger.addFilter(_stdout_truncation_filter)
|
||||
verbose_logger.addFilter(_stdout_truncation_filter)
|
||||
|
||||
|
||||
def _suppress_loggers():
|
||||
"""Suppress noisy loggers at INFO level"""
|
||||
|
|
|
|||
|
|
@ -12,11 +12,13 @@ import json
|
|||
|
||||
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
import redis
|
||||
import redis.asyncio as async_redis
|
||||
from redis.credentials import CredentialProvider
|
||||
|
||||
from litellm import get_secret, get_secret_str
|
||||
from litellm._redis_credential_provider import (
|
||||
|
|
@ -49,6 +51,7 @@ def _get_redis_kwargs():
|
|||
include_args: Final = {
|
||||
"url",
|
||||
"redis_connect_func",
|
||||
"credential_provider",
|
||||
"gcp_service_account",
|
||||
"gcp_ssl_ca_certs",
|
||||
"azure_redis_ad_token",
|
||||
|
|
@ -134,6 +137,7 @@ def _get_redis_cluster_kwargs(client=None):
|
|||
"ssl_check_hostname",
|
||||
"ssl_ca_certs",
|
||||
"redis_connect_func", # Needed for sync clusters and IAM detection
|
||||
"credential_provider",
|
||||
"gcp_service_account",
|
||||
"gcp_ssl_ca_certs",
|
||||
"azure_redis_ad_token",
|
||||
|
|
@ -153,7 +157,8 @@ def _get_redis_cluster_kwargs(client=None):
|
|||
def _get_redis_env_kwarg_mapping():
|
||||
PREFIX: Final = "REDIS_"
|
||||
|
||||
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs()}
|
||||
exclude_from_environment: Final = frozenset({"credential_provider"})
|
||||
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs() if x not in exclude_from_environment}
|
||||
|
||||
|
||||
def _redis_kwargs_from_environment():
|
||||
|
|
@ -351,6 +356,12 @@ def get_redis_url_from_environment():
|
|||
return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
|
||||
|
||||
|
||||
def _url_without_userinfo(url: str) -> str:
|
||||
parts: Final = urlsplit(url)
|
||||
netloc: Final = parts.netloc.rsplit("@", 1)[-1]
|
||||
return urlunsplit((parts.scheme, netloc, parts.path, parts.query, parts.fragment))
|
||||
|
||||
|
||||
def _get_redis_client_logic(**env_overrides):
|
||||
"""
|
||||
Common functionality across sync + async redis client implementations
|
||||
|
|
@ -408,54 +419,58 @@ def _get_redis_client_logic(**env_overrides):
|
|||
if _service_name is not None:
|
||||
redis_kwargs["service_name"] = _service_name
|
||||
|
||||
# Handle GCP IAM authentication
|
||||
_gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
|
||||
_gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
|
||||
|
||||
if _gcp_service_account is not None:
|
||||
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
|
||||
redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func(
|
||||
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
|
||||
if redis_kwargs.get("credential_provider") is None:
|
||||
# Handle GCP IAM authentication
|
||||
_gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str(
|
||||
"REDIS_GCP_SERVICE_ACCOUNT"
|
||||
)
|
||||
# Store GCP service account in redis_connect_func for async cluster access
|
||||
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
|
||||
_gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
|
||||
|
||||
# Remove GCP-specific kwargs that shouldn't be passed to Redis client
|
||||
redis_kwargs.pop("gcp_service_account", None)
|
||||
redis_kwargs.pop("gcp_ssl_ca_certs", None)
|
||||
if _gcp_service_account is not None:
|
||||
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
|
||||
redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func(
|
||||
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
|
||||
)
|
||||
# Store GCP service account in redis_connect_func for async cluster access
|
||||
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
|
||||
|
||||
# Only enable SSL if explicitly requested AND SSL CA certs are provided
|
||||
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
|
||||
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
|
||||
# Only enable SSL if explicitly requested AND SSL CA certs are provided
|
||||
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
|
||||
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
|
||||
|
||||
# Handle Azure AD authentication (after GCP IAM block)
|
||||
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
|
||||
# Handle Azure AD authentication (after GCP IAM block)
|
||||
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
|
||||
|
||||
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
|
||||
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
|
||||
|
||||
if _azure_ad_enabled and _gcp_service_account is not None:
|
||||
verbose_logger.warning(
|
||||
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
|
||||
"Using GCP IAM. Remove one to avoid misconfiguration."
|
||||
)
|
||||
if _azure_ad_enabled and _gcp_service_account is not None:
|
||||
verbose_logger.warning(
|
||||
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
|
||||
"Using GCP IAM. Remove one to avoid misconfiguration."
|
||||
)
|
||||
|
||||
if _azure_ad_enabled and _gcp_service_account is None:
|
||||
_azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
|
||||
_azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
|
||||
_azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET")
|
||||
if _azure_ad_enabled and _gcp_service_account is None:
|
||||
_azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
|
||||
_azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
|
||||
_azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str(
|
||||
"AZURE_CLIENT_SECRET"
|
||||
)
|
||||
|
||||
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
|
||||
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
|
||||
azure_client_id=_azure_client_id,
|
||||
azure_tenant_id=_azure_tenant_id,
|
||||
azure_client_secret=_azure_client_secret,
|
||||
)
|
||||
# Marker for async paths to detect Azure AD auth. The live credential
|
||||
# object is attached separately as `_azure_credential` by
|
||||
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
|
||||
# are intentionally NOT exposed on the function to avoid leaking
|
||||
# credentials via inspection or logging.
|
||||
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
|
||||
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
|
||||
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
|
||||
azure_client_id=_azure_client_id,
|
||||
azure_tenant_id=_azure_tenant_id,
|
||||
azure_client_secret=_azure_client_secret,
|
||||
)
|
||||
# Marker for async paths to detect Azure AD auth. The live credential
|
||||
# object is attached separately as `_azure_credential` by
|
||||
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
|
||||
# are intentionally NOT exposed on the function to avoid leaking
|
||||
# credentials via inspection or logging.
|
||||
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
|
||||
|
||||
redis_kwargs.pop("gcp_service_account", None)
|
||||
redis_kwargs.pop("gcp_ssl_ca_certs", None)
|
||||
|
||||
# Always remove Azure-specific kwargs that shouldn't be passed to Redis client
|
||||
redis_kwargs.pop("azure_redis_ad_token", None)
|
||||
|
|
@ -463,6 +478,13 @@ def _get_redis_client_logic(**env_overrides):
|
|||
redis_kwargs.pop("azure_tenant_id", None)
|
||||
redis_kwargs.pop("azure_client_secret", None)
|
||||
|
||||
if redis_kwargs.get("credential_provider") is not None:
|
||||
redis_kwargs.pop("redis_connect_func", None)
|
||||
redis_kwargs.pop("username", None)
|
||||
redis_kwargs.pop("password", None)
|
||||
if redis_kwargs.get("url") is not None:
|
||||
redis_kwargs["url"] = _url_without_userinfo(redis_kwargs["url"])
|
||||
|
||||
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
|
||||
# Only strip host/port/db/password when not routing to a cluster.
|
||||
# When startup_nodes is also present the cluster path takes priority and
|
||||
|
|
@ -530,8 +552,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
|
|||
service_name: Final = redis_kwargs.get("service_name")
|
||||
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
|
||||
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
|
||||
sentinel_kwargs: Final = dict(connection_kwargs)
|
||||
sentinel_kwargs["password"] = sentinel_password
|
||||
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
|
||||
|
||||
if not sentinel_nodes or not service_name:
|
||||
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
|
||||
|
|
@ -549,14 +570,22 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
|
|||
return sentinel.master_for(service_name, **connection_kwargs)
|
||||
|
||||
|
||||
def _sentinel_auth_kwargs(connection_kwargs: dict, sentinel_password: str | None) -> dict:
|
||||
"""The Sentinel monitors are separate servers that authenticate with their own password, so the
|
||||
data node's credential provider never belongs on them: leaving it there makes redis-py send the
|
||||
data node's token to a monitor, which fails whether the monitor is unauthenticated or has its
|
||||
own password."""
|
||||
kept: Final = ((k, v) for k, v in connection_kwargs.items() if k != "credential_provider")
|
||||
return dict(kept, password=sentinel_password)
|
||||
|
||||
|
||||
def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
|
||||
sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes")
|
||||
sentinel_password: Final = redis_kwargs.get("sentinel_password")
|
||||
service_name: Final = redis_kwargs.get("service_name")
|
||||
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
|
||||
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
|
||||
sentinel_kwargs: Final = dict(connection_kwargs)
|
||||
sentinel_kwargs["password"] = sentinel_password
|
||||
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
|
||||
|
||||
if not sentinel_nodes or not service_name:
|
||||
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
|
||||
|
|
@ -574,6 +603,41 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
|
|||
return sentinel.master_for(service_name, **connection_kwargs)
|
||||
|
||||
|
||||
def _async_credential_provider(redis_connect_func: object | None) -> CredentialProvider | None:
|
||||
"""The Azure AD and GCP IAM connect funcs run their AUTH exchange with the blocking client
|
||||
API, so on an async connection their ``send_command``/``read_response`` calls return
|
||||
coroutines nobody awaits and every connect fails. Async paths authenticate through a
|
||||
``CredentialProvider`` instead, which redis-py consults per connection so the token stays
|
||||
fresh. Any other ``redis_connect_func`` is left where it is, since redis-py awaits it
|
||||
itself when it is a coroutine function."""
|
||||
gcp_service_account: Final = getattr(redis_connect_func, "_gcp_service_account", None)
|
||||
if gcp_service_account is not None:
|
||||
return GCPIAMCredentialProvider(gcp_service_account)
|
||||
|
||||
azure_credential: Final = getattr(redis_connect_func, "_azure_credential", None)
|
||||
if azure_credential is not None:
|
||||
return AzureADCredentialProvider(azure_credential, username=os.environ.get("REDIS_USERNAME") or None)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _async_auth_kwargs(redis_kwargs: dict) -> dict:
|
||||
"""Swaps a connect func an async path cannot run for the equivalent credential provider,
|
||||
which supersedes any static username or password redis-py would otherwise reject it with."""
|
||||
explicit_provider: Final = redis_kwargs.get("credential_provider")
|
||||
credential_provider: Final = (
|
||||
explicit_provider
|
||||
if explicit_provider is not None
|
||||
else _async_credential_provider(redis_kwargs.get("redis_connect_func"))
|
||||
)
|
||||
if credential_provider is None:
|
||||
return redis_kwargs
|
||||
|
||||
superseded: Final = frozenset({"redis_connect_func", "username", "password"})
|
||||
kept: Final = ((k, v) for k, v in redis_kwargs.items() if k not in superseded)
|
||||
return dict(kept, credential_provider=credential_provider) # mutable-ok: the branches below mutate these kwargs
|
||||
|
||||
|
||||
def get_redis_client(**env_overrides):
|
||||
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
|
||||
|
||||
|
|
@ -600,7 +664,7 @@ def get_redis_async_client(
|
|||
connection_pool: async_redis.BlockingConnectionPool | None = None,
|
||||
**env_overrides,
|
||||
) -> async_redis.Redis | async_redis.RedisCluster:
|
||||
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
|
||||
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
from redis.cluster import ClusterNode
|
||||
|
|
@ -611,28 +675,12 @@ def get_redis_async_client(
|
|||
if arg in args:
|
||||
cluster_kwargs[arg] = redis_kwargs[arg]
|
||||
|
||||
# Handle GCP IAM authentication for async clusters
|
||||
redis_connect_func = cluster_kwargs.pop("redis_connect_func", None)
|
||||
|
||||
# Use a CredentialProvider so the IAM token is regenerated on every new
|
||||
# connection — mirrors the sync path where redis_connect_func is invoked
|
||||
# per connection. Without this, the token would expire after ~1 hour.
|
||||
if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
|
||||
# Handle Azure AD authentication for async clusters via CredentialProvider
|
||||
# so the credential's internal cache + silent refresh runs per connection
|
||||
# (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry).
|
||||
elif redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
|
||||
cluster_kwargs["credential_provider"] = AzureADCredentialProvider(
|
||||
redis_connect_func._azure_credential,
|
||||
username=os.environ.get("REDIS_USERNAME") or None,
|
||||
)
|
||||
|
||||
new_startup_nodes: Final[list[ClusterNode]] = []
|
||||
|
||||
for item in redis_kwargs["startup_nodes"]:
|
||||
new_startup_nodes.append(ClusterNode(**item))
|
||||
cluster_kwargs.pop("startup_nodes", None)
|
||||
cluster_kwargs.pop("redis_connect_func", None)
|
||||
|
||||
# Default to a periodic health check + TCP keepalive so a connection silently dropped
|
||||
# by a cluster restart (e.g. ElastiCache Serverless maintenance) is revalidated and
|
||||
|
|
@ -641,8 +689,16 @@ def get_redis_async_client(
|
|||
cluster_kwargs.setdefault("health_check_interval", REDIS_CLUSTER_HEALTH_CHECK_INTERVAL)
|
||||
cluster_kwargs.setdefault("socket_keepalive", True)
|
||||
|
||||
# A single node's client-side timeout must reset only that node's connections,
|
||||
# not tear down the whole cluster client for every concurrent caller.
|
||||
from litellm.caching.redis_cluster_node_isolation import (
|
||||
get_litellm_async_redis_cluster_class,
|
||||
)
|
||||
|
||||
async_redis_cluster_class: Final = get_litellm_async_redis_cluster_class()
|
||||
|
||||
# Create async RedisCluster with IAM token as password if available
|
||||
cluster_client: Final = async_redis.RedisCluster(
|
||||
cluster_client: Final = async_redis_cluster_class(
|
||||
startup_nodes=new_startup_nodes,
|
||||
**cluster_kwargs,
|
||||
)
|
||||
|
|
@ -667,19 +723,6 @@ def get_redis_async_client(
|
|||
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
|
||||
return _init_async_redis_sentinel(redis_kwargs)
|
||||
|
||||
# Wrap GCP / Azure AD auth in a CredentialProvider for the standard async
|
||||
# Redis client. The async client doesn't support redis_connect_func, but it
|
||||
# does honour credential_provider — which is called per connection, so the
|
||||
# underlying SDK can refresh tokens silently before they expire.
|
||||
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
|
||||
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
|
||||
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
|
||||
redis_connect_func._azure_credential,
|
||||
username=os.environ.get("REDIS_USERNAME") or None,
|
||||
)
|
||||
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
|
||||
|
||||
_pretty_print_redis_config(redis_kwargs=redis_kwargs)
|
||||
|
||||
if connection_pool is not None:
|
||||
|
|
@ -693,7 +736,7 @@ def get_redis_async_client(
|
|||
def get_redis_connection_pool(
|
||||
**env_overrides,
|
||||
) -> async_redis.BlockingConnectionPool | None:
|
||||
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
|
||||
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
|
||||
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
|
||||
|
||||
if "startup_nodes" in redis_kwargs:
|
||||
|
|
@ -714,25 +757,25 @@ def get_redis_connection_pool(
|
|||
)
|
||||
return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
|
||||
|
||||
# Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed
|
||||
# connections re-fetch tokens via the SDK's internal cache + silent refresh
|
||||
# rather than reusing a single token captured at pool creation.
|
||||
redis_connect_func: Final = redis_kwargs.pop("redis_connect_func", None)
|
||||
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
|
||||
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
|
||||
redis_connect_func._azure_credential,
|
||||
username=os.environ.get("REDIS_USERNAME") or None,
|
||||
)
|
||||
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
|
||||
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
|
||||
|
||||
if redis_kwargs.pop("ssl", None):
|
||||
redis_kwargs["connection_class"] = async_redis.SSLConnection
|
||||
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
|
||||
|
||||
|
||||
def _redis_kwargs_for_logging(redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {
|
||||
key: "<credential provider>"
|
||||
if key == "credential_provider" and value is not None
|
||||
else "<redis connect function>"
|
||||
if key == "redis_connect_func" and value is not None
|
||||
else value
|
||||
for key, value in redis_kwargs.items()
|
||||
}
|
||||
|
||||
|
||||
def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
||||
"""Pretty print the Redis configuration using rich with sensitive data masking"""
|
||||
redis_kwargs_for_logging: Final = _redis_kwargs_for_logging(redis_kwargs)
|
||||
try:
|
||||
import logging
|
||||
|
||||
|
|
@ -750,7 +793,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
masker = SensitiveDataMasker()
|
||||
|
||||
# Mask sensitive data in redis_kwargs
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
|
||||
|
||||
# Create main panel title
|
||||
title: Final = Text("Redis Configuration", style="bold blue")
|
||||
|
|
@ -813,7 +856,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
|
|||
except ImportError:
|
||||
# Fallback to simple logging if rich is not available
|
||||
masker = SensitiveDataMasker()
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
|
||||
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
|
||||
verbose_logger.info("Redis configuration: %s", masked_redis_kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.error("Error pretty printing Redis configuration: %s", e)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ Custom A2A Card Resolver for LiteLLM.
|
|||
Extends the A2A SDK's card resolver to support multiple well-known paths.
|
||||
"""
|
||||
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -48,6 +49,43 @@ def is_localhost_or_internal_url(url: str | None) -> bool:
|
|||
return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS)
|
||||
|
||||
|
||||
_CANONICAL_PROTOCOL_BINDINGS: Final = MappingProxyType(
|
||||
{
|
||||
"jsonrpc": "JSONRPC",
|
||||
"http+json": "HTTP+JSON",
|
||||
"grpc": "GRPC",
|
||||
}
|
||||
)
|
||||
|
||||
_LEGACY_PROTOCOL_VERSION: Final = "0.3"
|
||||
|
||||
|
||||
def normalize_agent_card_interfaces(agent_card: "AgentCard") -> "AgentCard":
|
||||
"""
|
||||
Canonicalize the supported interfaces of spec-adjacent agent cards.
|
||||
|
||||
Some A2A servers (e.g. LangGraph Platform) serve agent cards with lowercase
|
||||
bindings like "jsonrpc", but a2a-sdk's ClientFactory matches bindings
|
||||
case-sensitively against its uppercase TransportProtocol constants and fails
|
||||
with "no compatible transports found." for spec-adjacent casings.
|
||||
|
||||
The same servers also speak the A2A 0.3 JSON dialect ("kind"-discriminated
|
||||
payloads) while declaring protocolVersion "1.0", which a2a-sdk's strict v1
|
||||
proto parsing rejects. A mis-cased binding fingerprints such a server, so its
|
||||
declared version is downgraded to 0.3 to route the SDK's ClientFactory onto
|
||||
its v0.3 compat transport, which speaks that dialect.
|
||||
"""
|
||||
normalized: Final = type(agent_card)()
|
||||
normalized.CopyFrom(agent_card)
|
||||
for interface in normalized.supported_interfaces:
|
||||
canonical: str | None = _CANONICAL_PROTOCOL_BINDINGS.get(interface.protocol_binding.lower())
|
||||
if canonical is None or canonical == interface.protocol_binding:
|
||||
continue
|
||||
interface.protocol_binding = canonical
|
||||
interface.protocol_version = _LEGACY_PROTOCOL_VERSION
|
||||
return normalized
|
||||
|
||||
|
||||
def get_agent_card_url(agent_card: "AgentCard") -> str | None:
|
||||
"""Return the agent endpoint URL from the resolved SDK card."""
|
||||
url: Final = getattr(agent_card, "url", None)
|
||||
|
|
|
|||
|
|
@ -73,6 +73,7 @@ except ImportError:
|
|||
from litellm.a2a_protocol.card_resolver import (
|
||||
LiteLLMA2ACardResolver,
|
||||
get_agent_card_url,
|
||||
normalize_agent_card_interfaces,
|
||||
)
|
||||
from litellm.a2a_protocol.exception_mapping_utils import (
|
||||
handle_a2a_localhost_retry,
|
||||
|
|
@ -782,13 +783,17 @@ async def create_a2a_client(
|
|||
if extra_headers:
|
||||
verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
|
||||
|
||||
resolver: Final = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
|
||||
agent_card: Final = normalize_agent_card_interfaces(
|
||||
await resolver.get_agent_card(http_kwargs={"headers": extra_headers} if extra_headers else None)
|
||||
)
|
||||
|
||||
a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
|
||||
base_url,
|
||||
agent_card,
|
||||
client_config=ClientConfig( # pyright: ignore[reportOptionalCall]
|
||||
httpx_client=httpx_client,
|
||||
streaming=streaming,
|
||||
),
|
||||
resolver_http_kwargs={"headers": extra_headers} if extra_headers else None,
|
||||
)
|
||||
# Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse
|
||||
# the configured httpx client and this agent's headers without excavating
|
||||
|
|
@ -799,9 +804,7 @@ async def create_a2a_client(
|
|||
if extra_headers
|
||||
else None
|
||||
)
|
||||
agent_card: Final = getattr(a2a_client, "_card", None)
|
||||
if agent_card is not None:
|
||||
a2a_client._litellm_agent_card = agent_card
|
||||
a2a_client._litellm_agent_card = agent_card
|
||||
|
||||
verbose_logger.info("A2A client created for %s", base_url)
|
||||
|
||||
|
|
|
|||
|
|
@ -296,6 +296,32 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
)
|
||||
|
||||
|
||||
def _provider_output_file_id(output_file_id: str) -> str:
|
||||
"""
|
||||
Resolve the file id the provider actually knows: unified ids yield their embedded
|
||||
llm_output_file_id, model-encoded ids decode to the raw provider id, raw ids pass through.
|
||||
"""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
get_original_file_id,
|
||||
)
|
||||
|
||||
unified_file_id: Final = _is_base64_encoded_unified_file_id(output_file_id)
|
||||
if not unified_file_id:
|
||||
return get_original_file_id(output_file_id)
|
||||
try:
|
||||
extracted: Final = unified_file_id.split("llm_output_file_id,")[1].split(";")[0]
|
||||
except (IndexError, AttributeError) as e:
|
||||
verbose_logger.error(
|
||||
"Failed to extract LLM output file ID from unified file ID: %s, error: %s",
|
||||
output_file_id,
|
||||
e,
|
||||
)
|
||||
return output_file_id
|
||||
verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", extracted)
|
||||
return extracted
|
||||
|
||||
|
||||
async def _fetch_batch_output_file_content(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
|
|
@ -311,23 +337,11 @@ async def _fetch_batch_output_file_content(
|
|||
Required for Azure and other providers that need authentication
|
||||
"""
|
||||
from litellm.files.main import afile_content
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
)
|
||||
|
||||
if batch.output_file_id is None:
|
||||
raise ValueError("Output file id is None cannot retrieve file content")
|
||||
|
||||
file_id = batch.output_file_id
|
||||
is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id)
|
||||
if is_base64_unified_file_id:
|
||||
try:
|
||||
file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0]
|
||||
verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id)
|
||||
except (IndexError, AttributeError) as e:
|
||||
verbose_logger.error(
|
||||
"Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e
|
||||
)
|
||||
file_id: Final = _provider_output_file_id(batch.output_file_id)
|
||||
|
||||
# Build kwargs for afile_content with credentials from litellm_params
|
||||
file_content_kwargs: Final = {
|
||||
|
|
@ -537,7 +551,7 @@ def _get_batch_job_usage_from_response_body(
|
|||
return usage
|
||||
|
||||
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict:
|
||||
def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> Mapping[str, Any]:
|
||||
"""
|
||||
Get the ``result`` object from a line of an Anthropic message batch results JSONL file.
|
||||
|
||||
|
|
@ -549,7 +563,7 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[st
|
|||
|
||||
def _get_response_from_batch_job_output_file(
|
||||
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
|
||||
) -> Any:
|
||||
) -> Mapping[str, Any]:
|
||||
"""
|
||||
Get the response from the batch job output file
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from collections.abc import Sequence
|
|||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
|
@ -60,6 +61,13 @@ def resolve_embedding_max_input_tokens(
|
|||
return deployment_max_input_tokens
|
||||
|
||||
|
||||
def resolve_embedding_timeout(configured_timeout: float | None) -> float:
|
||||
"""Explicit cache setting first, else the short semantic-cache default."""
|
||||
if configured_timeout is not None:
|
||||
return configured_timeout
|
||||
return SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
|
||||
|
||||
def truncate_embedding_input(prompt: str, embedding_model: str, max_input_tokens: int | None) -> str:
|
||||
"""Keep only the first ``max_input_tokens`` tokens of ``prompt`` for the embedding call."""
|
||||
if max_input_tokens is None:
|
||||
|
|
|
|||
|
|
@ -98,6 +98,7 @@ class Cache:
|
|||
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
|
||||
qdrant_semantic_cache_vector_size: int | None = None,
|
||||
semantic_cache_embedding_max_input_tokens: int | None = None,
|
||||
semantic_cache_embedding_timeout: float | None = None,
|
||||
# GCP IAM authentication parameters
|
||||
gcp_service_account: str | None = None,
|
||||
gcp_ssl_ca_certs: str | None = None,
|
||||
|
|
@ -124,6 +125,7 @@ class Cache:
|
|||
qdrant_collection_name (str, optional): The name for your qdrant collection. Required if type is "qdrant-semantic".
|
||||
similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic".
|
||||
semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens.
|
||||
semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS.
|
||||
|
||||
# Disk Cache Args
|
||||
disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
|
||||
|
|
@ -195,6 +197,7 @@ class Cache:
|
|||
embedding_model=redis_semantic_cache_embedding_model,
|
||||
index_name=redis_semantic_cache_index_name,
|
||||
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
|
||||
embedding_timeout=semantic_cache_embedding_timeout,
|
||||
**kwargs,
|
||||
)
|
||||
elif type == LiteLLMCacheType.VALKEY_SEMANTIC:
|
||||
|
|
@ -211,6 +214,7 @@ class Cache:
|
|||
index_name=valkey_semantic_cache_index_name,
|
||||
startup_nodes=redis_startup_nodes,
|
||||
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
|
||||
embedding_timeout=semantic_cache_embedding_timeout,
|
||||
**kwargs,
|
||||
)
|
||||
elif type == LiteLLMCacheType.QDRANT_SEMANTIC:
|
||||
|
|
@ -223,6 +227,7 @@ class Cache:
|
|||
embedding_model=qdrant_semantic_cache_embedding_model,
|
||||
vector_size=qdrant_semantic_cache_vector_size,
|
||||
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
|
||||
embedding_timeout=semantic_cache_embedding_timeout,
|
||||
)
|
||||
elif type == LiteLLMCacheType.LOCAL:
|
||||
self.cache = InMemoryCache()
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import asyncio
|
|||
import datetime
|
||||
import inspect
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -27,6 +27,7 @@ import litellm
|
|||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.caching import InMemoryCache
|
||||
from litellm.caching.caching import S3Cache
|
||||
from litellm.constants import CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
||||
update_response_metadata,
|
||||
)
|
||||
|
|
@ -124,6 +125,29 @@ def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") ->
|
|||
return details.model_dump(exclude_none=True) if hasattr(details, "model_dump") else {}
|
||||
|
||||
|
||||
_PENDING_CACHE_WRITES: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs to pending write tasks
|
||||
|
||||
|
||||
async def _complete_cache_write_despite_cancellation(write_factory: Callable[[], Awaitable[None]]) -> None:
|
||||
try:
|
||||
await write_factory()
|
||||
except asyncio.CancelledError:
|
||||
try:
|
||||
await asyncio.wait_for(write_factory(), timeout=CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS)
|
||||
except Exception as flush_error: # noqa: BLE001 # shutdown flush failures are logged, never raised
|
||||
verbose_logger.warning(
|
||||
"LiteLLM Cache: pending cache write failed during event loop shutdown: %s", flush_error
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "asyncio.Task[None]":
|
||||
task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory))
|
||||
_PENDING_CACHE_WRITES.add(task)
|
||||
task.add_done_callback(_PENDING_CACHE_WRITES.discard)
|
||||
return task
|
||||
|
||||
|
||||
def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
|
||||
"""Read the caller-supplied ``cache_key`` off the request kwargs."""
|
||||
return request_kwargs.get("cache_key", None)
|
||||
|
|
@ -983,6 +1007,7 @@ class LLMCachingHandler:
|
|||
|
||||
if litellm.cache is None:
|
||||
return
|
||||
cache: Final = litellm.cache
|
||||
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs.update(
|
||||
|
|
@ -1004,24 +1029,24 @@ class LLMCachingHandler:
|
|||
):
|
||||
if (
|
||||
isinstance(result, EmbeddingResponse)
|
||||
and litellm.cache is not None
|
||||
and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
|
||||
and not isinstance(cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
|
||||
):
|
||||
asyncio.create_task(
|
||||
litellm.cache.async_add_cache_pipeline(
|
||||
create_cache_write_task(
|
||||
lambda: cache.async_add_cache_pipeline(
|
||||
result, dynamic_cache_object=self.dual_cache, **new_kwargs
|
||||
)
|
||||
)
|
||||
else:
|
||||
asyncio.create_task(
|
||||
litellm.cache.async_add_cache(
|
||||
result.model_dump_json(),
|
||||
result_json: Final = result.model_dump_json()
|
||||
create_cache_write_task(
|
||||
lambda: cache.async_add_cache(
|
||||
result_json,
|
||||
dynamic_cache_object=self.dual_cache,
|
||||
**new_kwargs,
|
||||
)
|
||||
)
|
||||
else:
|
||||
asyncio.create_task(litellm.cache.async_add_cache(result, **new_kwargs))
|
||||
create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs))
|
||||
|
||||
def sync_set_cache(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -16,7 +16,11 @@ from typing import TYPE_CHECKING, Any, Final, cast
|
|||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose
|
||||
from litellm.constants import QDRANT_SCALAR_QUANTILE, QDRANT_VECTOR_SIZE
|
||||
from litellm.constants import (
|
||||
QDRANT_SCALAR_QUANTILE,
|
||||
QDRANT_VECTOR_SIZE,
|
||||
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_str_from_messages,
|
||||
)
|
||||
|
|
@ -26,6 +30,7 @@ from ._embedding_router import (
|
|||
build_router_embedding_metadata,
|
||||
resolve_embedding_max_input_tokens,
|
||||
resolve_embedding_router,
|
||||
resolve_embedding_timeout,
|
||||
truncate_embedding_input,
|
||||
)
|
||||
from .base_cache import BaseCache
|
||||
|
|
@ -37,6 +42,7 @@ if TYPE_CHECKING:
|
|||
class QdrantSemanticCache(BaseCache):
|
||||
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
|
||||
embedding_max_input_tokens: int | None = None
|
||||
embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -49,6 +55,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
host_type=None,
|
||||
vector_size=None,
|
||||
embedding_max_input_tokens: int | None = None,
|
||||
embedding_timeout: float | None = None,
|
||||
):
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_httpx_client,
|
||||
|
|
@ -68,6 +75,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
self.similarity_threshold = similarity_threshold
|
||||
self.embedding_model = embedding_model
|
||||
self.embedding_max_input_tokens = embedding_max_input_tokens
|
||||
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
|
||||
self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
|
||||
headers = {}
|
||||
|
||||
|
|
@ -222,11 +230,15 @@ class QdrantSemanticCache(BaseCache):
|
|||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
metadata=build_router_embedding_metadata(metadata),
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
)
|
||||
return litellm.embedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
|
||||
|
|
@ -238,19 +250,25 @@ class QdrantSemanticCache(BaseCache):
|
|||
|
||||
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
|
||||
embedding_input: Final = self._embedding_input(prompt, router)
|
||||
if router is not None:
|
||||
return await router.aembedding(
|
||||
embedding_call: Final = (
|
||||
router.aembedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
metadata=build_router_embedding_metadata(metadata),
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
)
|
||||
if router is not None
|
||||
else litellm.aembedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
return await litellm.aembedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
)
|
||||
return await asyncio.wait_for(embedding_call, self.embedding_timeout)
|
||||
|
||||
def set_cache(self, key, value, **kwargs):
|
||||
print_verbose(f"qdrant semantic-cache set_cache, kwargs: {kwargs}")
|
||||
|
|
|
|||
|
|
@ -175,6 +175,10 @@ _RedisCallResult = TypeVar("_RedisCallResult")
|
|||
_swallowed_redis_failures: Final[ContextVar[int]] = ContextVar("litellm_swallowed_redis_failures", default=0)
|
||||
|
||||
|
||||
def _opaque_kwarg_key(value: object) -> str:
|
||||
return f"{type(value).__name__}-{id(value)}"
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _redis_health_error_types() -> tuple[type, ...]:
|
||||
"""Exception types that mean the Redis backend itself is unhealthy.
|
||||
|
|
@ -399,10 +403,9 @@ class RedisCache(BaseCache):
|
|||
Generate a cache key for the async Redis client based on connection parameters.
|
||||
This ensures different Redis configurations use different cached clients.
|
||||
"""
|
||||
# Create a stable representation of redis_kwargs for hashing
|
||||
# Sort keys to ensure consistent hash regardless of parameter order
|
||||
sorted_kwargs: Final = sorted(self.redis_kwargs.items())
|
||||
kwargs_str: Final = json.dumps(sorted_kwargs, sort_keys=True)
|
||||
kwargs_str: Final = json.dumps(sorted_kwargs, sort_keys=True, default=_opaque_kwarg_key)
|
||||
kwargs_hash: Final = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16]
|
||||
return f"async-redis-client-{kwargs_hash}"
|
||||
|
||||
|
|
@ -432,7 +435,7 @@ class RedisCache(BaseCache):
|
|||
"""
|
||||
if key is None:
|
||||
return key
|
||||
if self.namespace is not None and not key.startswith(self.namespace):
|
||||
if self.namespace and not key.startswith(self.namespace + ":"):
|
||||
key = self.namespace + ":" + key
|
||||
|
||||
return key
|
||||
|
|
@ -1384,10 +1387,10 @@ class RedisCache(BaseCache):
|
|||
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
|
||||
"""
|
||||
try:
|
||||
import redis.asyncio as redis_async
|
||||
from .._redis import get_redis_async_client
|
||||
|
||||
# Create a fresh Redis client with current settings
|
||||
redis_client: Final = redis_async.Redis(**self.redis_kwargs)
|
||||
redis_client: Final = get_redis_async_client(**self.redis_kwargs)
|
||||
|
||||
# Test the connection
|
||||
ping_result: Final = await redis_client.ping()
|
||||
|
|
|
|||
|
|
@ -64,22 +64,9 @@ class RedisClusterCache(RedisCache):
|
|||
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
|
||||
"""
|
||||
try:
|
||||
import redis.asyncio as redis_async
|
||||
from redis.cluster import ClusterNode
|
||||
from .._redis import get_redis_async_client
|
||||
|
||||
# Create ClusterNode objects from startup_nodes
|
||||
cluster_kwargs: Final = self.redis_kwargs.copy()
|
||||
startup_nodes: Final = cluster_kwargs.pop("startup_nodes", [])
|
||||
|
||||
new_startup_nodes: Final[list[ClusterNode]] = []
|
||||
for item in startup_nodes:
|
||||
new_startup_nodes.append(ClusterNode(**item))
|
||||
|
||||
# Create a fresh Redis Cluster client with current settings
|
||||
redis_client: Final = redis_async.RedisCluster(
|
||||
startup_nodes=new_startup_nodes,
|
||||
**cluster_kwargs,
|
||||
)
|
||||
redis_client: Final = get_redis_async_client(**self.redis_kwargs)
|
||||
|
||||
# Test the connection
|
||||
ping_result: Final = await redis_client.ping()
|
||||
|
|
|
|||
201
litellm/caching/redis_cluster_node_isolation.py
Normal file
201
litellm/caching/redis_cluster_node_isolation.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
"""Bounds the blast radius of a single node's transient connection error on the async
|
||||
Redis Cluster client.
|
||||
|
||||
redis-py's ``RedisCluster._execute_command`` responds to a ``ConnectionError`` or
|
||||
``TimeoutError`` on ANY one node by tearing down every node's connections and flipping
|
||||
the client into "needs reinitialization", which forces every other concurrent caller
|
||||
sharing this client through one reinit lock until the whole cluster topology is
|
||||
re-walked. Under real proxy load, a client-side socket timeout on a single node is a
|
||||
routine event (the event loop was too busy to read the response before ``socket_timeout``
|
||||
elapsed) and does not mean the cluster's topology moved, so treating it as a full-cluster
|
||||
event turns one slow node into a proxy-wide latency spike while Redis itself stays
|
||||
healthy -- confirmed live: pausing one of three local cluster nodes made every concurrent
|
||||
command against the other two, untouched nodes stall for the full pause duration too.
|
||||
|
||||
``get_litellm_async_redis_cluster_class`` returns a ``RedisCluster`` subclass that resets
|
||||
only the node that actually failed (mirroring what a plain, non-cluster Redis client
|
||||
already does when one of its pooled connections errors), leaving every other node's
|
||||
connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-covered,
|
||||
retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
|
||||
topology changed.
|
||||
|
||||
redis-py 8.x fixed this upstream with gentler machinery than this override's
|
||||
``node.disconnect()`` (which also kills connections other coroutines are mid-operation
|
||||
on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
|
||||
killed connection): it marks in-use connections for reconnect only after their current
|
||||
operation completes, disconnects only the idle pooled ones, and defers reinitialization
|
||||
to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
|
||||
recovery API, the factory returns the base ``RedisCluster`` unmodified.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from redis.asyncio.cluster import RedisCluster as _AsyncRedisClusterType
|
||||
|
||||
|
||||
class _ClusterNodeAttrs(Protocol):
|
||||
"""The subset of ``redis.asyncio.cluster.ClusterNode`` this override reads. redis-py
|
||||
ships no resolvable stub for these members under the repo's current types-redis pin,
|
||||
so a plain attribute access resolves every downstream use to ``Unknown`` under strict
|
||||
mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's
|
||||
own logic fully typed without a banned ``typing.cast``."""
|
||||
|
||||
async def execute_command(
|
||||
self,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: mirrors redis-py's own ClusterNode.execute_command signature, a raw command dispatch with no fixed keyword contract
|
||||
) -> object: ...
|
||||
async def disconnect(self) -> None: ...
|
||||
|
||||
|
||||
class _NodesManagerAttrs(Protocol):
|
||||
_moved_exception: object
|
||||
|
||||
def get_node_from_slot(
|
||||
self, slot: int, read_from_replicas: bool, load_balancing_strategy: object
|
||||
) -> _ClusterNodeAttrs: ...
|
||||
|
||||
|
||||
class _ClusterAttrs(Protocol):
|
||||
RedisClusterRequestTTL: int
|
||||
reinitialize_counter: int
|
||||
reinitialize_steps: int
|
||||
read_from_replicas: bool
|
||||
load_balancing_strategy: object
|
||||
nodes_manager: _NodesManagerAttrs
|
||||
|
||||
def get_node(self, node_name: str) -> _ClusterNodeAttrs: ...
|
||||
async def _determine_slot(self, *args: object) -> int: ...
|
||||
async def aclose(self) -> None: ...
|
||||
|
||||
|
||||
#: redis-py versions this override's copied ``_execute_command`` body has been verified
|
||||
#: against. A version outside this set may have changed the method's structure in a way
|
||||
#: this override can't see (Python won't error -- it'll just run our now-stale copy), so
|
||||
#: construction logs a loud warning rather than silently trusting an unverified copy.
|
||||
_VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
|
||||
|
||||
|
||||
def get_litellm_async_redis_cluster_class(
|
||||
cluster_node_class: type | None = None,
|
||||
) -> type["_AsyncRedisClusterType"]:
|
||||
"""Returns the base ``RedisCluster`` when the installed redis-py already recovers a
|
||||
node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
|
||||
subclass with the per-node isolation fix for older versions whose upstream branch
|
||||
tears down the whole cluster client.
|
||||
|
||||
``cluster_node_class`` exists for dependency injection in tests; production callers
|
||||
leave it unset and the installed ``ClusterNode`` is used.
|
||||
|
||||
Imported lazily because this module is reachable from a base ``import litellm`` while
|
||||
redis is not a base dependency. Cheap to call repeatedly: the underlying redis
|
||||
submodules are cached in ``sys.modules`` after the first import.
|
||||
"""
|
||||
import redis
|
||||
from redis.asyncio.cluster import (
|
||||
ClusterNode as _AsyncClusterNode, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin
|
||||
)
|
||||
from redis.asyncio.cluster import (
|
||||
RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # same stale-stub gap as the import above
|
||||
)
|
||||
from redis.cluster import get_node_name
|
||||
from redis.commands import READ_COMMANDS
|
||||
from redis.exceptions import (
|
||||
AskError,
|
||||
BusyLoadingError,
|
||||
ClusterDownError,
|
||||
ClusterError,
|
||||
MaxConnectionsError,
|
||||
MovedError,
|
||||
SlotNotCoveredError,
|
||||
TryAgainError,
|
||||
)
|
||||
from redis.exceptions import ConnectionError as _RedisConnectionError
|
||||
from redis.exceptions import TimeoutError as _RedisTimeoutError
|
||||
|
||||
node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
|
||||
if hasattr(node_class, "update_active_connections_for_reconnect"):
|
||||
verbose_logger.debug(
|
||||
"redis-py %s recovers a node-level connection error per-connection upstream; "
|
||||
"using the base RedisCluster without litellm's node-isolation override.",
|
||||
redis.__version__,
|
||||
)
|
||||
return _BaseAsyncRedisCluster
|
||||
|
||||
if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
|
||||
verbose_logger.warning(
|
||||
"redis-py %s is not in the set this cluster-teardown-storm fix was verified "
|
||||
"against (%s). The per-node-isolation override may not match the installed library's "
|
||||
"real _execute_command behavior.",
|
||||
redis.__version__,
|
||||
sorted(_VERIFIED_REDIS_VERSIONS),
|
||||
)
|
||||
|
||||
class LiteLLMAsyncRedisCluster(
|
||||
_BaseAsyncRedisCluster # pyright: ignore[reportUntypedBaseClass] # same stale-stub gap as the import above; the base class itself is unresolvable, not this subclass's own code
|
||||
):
|
||||
async def _execute_command(
|
||||
self,
|
||||
target_node: _ClusterNodeAttrs,
|
||||
*args: object,
|
||||
**kwargs: object, # kwargs-ok: overrides redis-py's own **kwargs signature; the keyword contract is defined by the Redis command being dispatched, not by this method
|
||||
) -> object:
|
||||
cluster: _ClusterAttrs = self
|
||||
node = target_node
|
||||
|
||||
asking = moved = False
|
||||
redirect_addr: str | None = None
|
||||
ttl = cluster.RedisClusterRequestTTL
|
||||
|
||||
while ttl > 0:
|
||||
ttl -= 1
|
||||
try:
|
||||
if asking:
|
||||
assert redirect_addr is not None
|
||||
node = cluster.get_node(node_name=redirect_addr)
|
||||
await node.execute_command("ASKING")
|
||||
asking = False
|
||||
elif moved:
|
||||
slot = await cluster._determine_slot(*args) # pyright: ignore[reportPrivateUsage] # mirrors upstream's own un-overridden branch, which makes this identical private call from the same subclass
|
||||
node = cluster.nodes_manager.get_node_from_slot(
|
||||
slot,
|
||||
cluster.read_from_replicas and args[0] in READ_COMMANDS,
|
||||
(cluster.load_balancing_strategy if args[0] in READ_COMMANDS else None),
|
||||
)
|
||||
moved = False
|
||||
|
||||
return await node.execute_command(*args, **kwargs)
|
||||
except (BusyLoadingError, MaxConnectionsError):
|
||||
raise
|
||||
except (_RedisConnectionError, _RedisTimeoutError):
|
||||
# Reset only the node that actually failed instead of the upstream
|
||||
# default (`await self.aclose()`, a full-cluster teardown that forces
|
||||
# every other concurrent caller through the shared reinit lock).
|
||||
await node.disconnect()
|
||||
raise
|
||||
except (ClusterDownError, SlotNotCoveredError):
|
||||
await cluster.aclose()
|
||||
await asyncio.sleep(0.25)
|
||||
raise
|
||||
except MovedError as e:
|
||||
cluster.reinitialize_counter += 1
|
||||
if cluster.reinitialize_steps and cluster.reinitialize_counter % cluster.reinitialize_steps == 0:
|
||||
await cluster.aclose()
|
||||
cluster.reinitialize_counter = 0
|
||||
else:
|
||||
cluster.nodes_manager._moved_exception = e # pyright: ignore[reportPrivateUsage] # mirrors upstream's own un-overridden branch; redis-py exposes no public setter for this
|
||||
moved = True
|
||||
except AskError as e:
|
||||
redirect_addr = get_node_name(host=e.host, port=e.port)
|
||||
asking = True
|
||||
except TryAgainError:
|
||||
if ttl < cluster.RedisClusterRequestTTL / 2:
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
raise ClusterError("TTL exhausted.")
|
||||
|
||||
return LiteLLMAsyncRedisCluster
|
||||
|
|
@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
|
|||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_str_from_messages,
|
||||
)
|
||||
|
|
@ -27,6 +28,7 @@ from ._embedding_router import (
|
|||
build_router_embedding_metadata,
|
||||
resolve_embedding_max_input_tokens,
|
||||
resolve_embedding_router,
|
||||
resolve_embedding_timeout,
|
||||
truncate_embedding_input,
|
||||
)
|
||||
from .base_cache import BaseCache
|
||||
|
|
@ -47,6 +49,7 @@ class RedisSemanticCache(BaseCache):
|
|||
DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index"
|
||||
CACHE_KEY_FIELD_NAME: str = "litellm_cache_key"
|
||||
embedding_max_input_tokens: int | None = None
|
||||
embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -58,6 +61,7 @@ class RedisSemanticCache(BaseCache):
|
|||
embedding_model: str = "text-embedding-ada-002",
|
||||
index_name: str | None = None,
|
||||
embedding_max_input_tokens: int | None = None,
|
||||
embedding_timeout: float | None = None,
|
||||
**kwargs: object,
|
||||
):
|
||||
"""
|
||||
|
|
@ -74,6 +78,8 @@ class RedisSemanticCache(BaseCache):
|
|||
index_name: Name for the Redis index
|
||||
embedding_max_input_tokens: Truncate prompts to this many tokens before
|
||||
embedding; defaults to the Router deployment's configured max_input_tokens
|
||||
embedding_timeout: Seconds a cache lookup may spend embedding the prompt before it
|
||||
gives up and lets the request continue to the LLM
|
||||
ttl: Default time-to-live for cache entries in seconds
|
||||
**kwargs: Additional arguments passed to the Redis client
|
||||
|
||||
|
|
@ -99,6 +105,7 @@ class RedisSemanticCache(BaseCache):
|
|||
self.distance_threshold = 1 - similarity_threshold
|
||||
self.embedding_model = embedding_model
|
||||
self.embedding_max_input_tokens = embedding_max_input_tokens
|
||||
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
|
||||
|
||||
# Set up Redis connection
|
||||
if redis_url is None:
|
||||
|
|
@ -349,6 +356,8 @@ class RedisSemanticCache(BaseCache):
|
|||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
metadata=build_router_embedding_metadata(metadata),
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
),
|
||||
)
|
||||
else:
|
||||
|
|
@ -358,6 +367,8 @@ class RedisSemanticCache(BaseCache):
|
|||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
),
|
||||
)
|
||||
return embedding_response["data"][0]["embedding"]
|
||||
|
|
@ -512,20 +523,26 @@ class RedisSemanticCache(BaseCache):
|
|||
|
||||
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
|
||||
embedding_input: Final = self._embedding_input(prompt, router)
|
||||
embedding_call: Final = (
|
||||
router.aembedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
metadata=build_router_embedding_metadata(metadata),
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
)
|
||||
if router is not None
|
||||
else litellm.aembedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
timeout=self.embedding_timeout,
|
||||
num_retries=0,
|
||||
)
|
||||
)
|
||||
try:
|
||||
if router is not None:
|
||||
embedding_response = await router.aembedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
metadata=build_router_embedding_metadata(metadata),
|
||||
)
|
||||
else:
|
||||
embedding_response = await litellm.aembedding(
|
||||
model=self.embedding_model,
|
||||
input=embedding_input,
|
||||
cache={"no-store": True, "no-cache": True},
|
||||
)
|
||||
embedding_response: Final = await asyncio.wait_for(embedding_call, self.embedding_timeout)
|
||||
return embedding_response["data"][0]["embedding"]
|
||||
except Exception as e:
|
||||
print_verbose(f"Error generating async embedding: {e}")
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from litellm._logging import print_verbose
|
|||
from litellm._uuid import uuid
|
||||
from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
|
||||
|
||||
from ._embedding_router import resolve_embedding_timeout
|
||||
from .redis_semantic_cache import RedisSemanticCache
|
||||
|
||||
|
||||
|
|
@ -62,6 +63,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
sync_client: Redis | None = None,
|
||||
async_client: AsyncRedis | None = None,
|
||||
embedding_max_input_tokens: int | None = None,
|
||||
embedding_timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
if similarity_threshold is None:
|
||||
|
|
@ -80,6 +82,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
self.similarity_threshold = similarity_threshold
|
||||
self.embedding_model = embedding_model
|
||||
self.embedding_max_input_tokens = embedding_max_input_tokens
|
||||
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
|
||||
self.index_name = index_name or self.DEFAULT_VALKEY_INDEX_NAME
|
||||
self.key_prefix = f"{self.index_name}:"
|
||||
self._index_dim: int | None = None
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req
|
|||
import json
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
|
||||
|
||||
from openai.types.responses.custom_tool_param import CustomToolParam
|
||||
from openai.types.responses.response_input_param import (
|
||||
|
|
@ -21,6 +21,9 @@ from pydantic import BaseModel
|
|||
import litellm
|
||||
from litellm import ModelResponse
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
responses_reasoning_item_from_thinking_blocks,
|
||||
)
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
from litellm.llms.base_llm.bridges.completion_transformation import (
|
||||
CompletionTransformationBridge,
|
||||
|
|
@ -32,6 +35,7 @@ from litellm.responses.sse_output_recovery import (
|
|||
)
|
||||
from litellm.responses.utils import normalize_responses_api_stream_options
|
||||
from litellm.types.llms.openai import (
|
||||
REASONING_EFFORT,
|
||||
ChatCompletionAnnotation,
|
||||
ChatCompletionReasoningItem,
|
||||
ChatCompletionToolCallChunk,
|
||||
|
|
@ -85,6 +89,22 @@ def _get_reasoning_items(
|
|||
return []
|
||||
|
||||
|
||||
def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: # mutable-ok: API message payload
|
||||
"""Reasoning input items for an assistant message.
|
||||
|
||||
Stored reasoning items win because they carry an id the Responses API minted; thinking
|
||||
blocks are the fallback for turns that arrived over another API surface.
|
||||
"""
|
||||
items: Final = _get_reasoning_items(msg)
|
||||
stored: Final = [_reasoning_item_to_response_input(item) for item in items] # mutable-ok: API message payload
|
||||
if stored:
|
||||
return stored
|
||||
raw_blocks: Final = msg.get("thinking_blocks") or ()
|
||||
blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json
|
||||
from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks)
|
||||
return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload
|
||||
|
||||
|
||||
def _build_reasoning_item(
|
||||
item_id: str,
|
||||
encrypted_content: str | None,
|
||||
|
|
@ -113,6 +133,58 @@ def _build_reasoning_item(
|
|||
}
|
||||
|
||||
|
||||
def _reasoning_item_from_output_item(item: object) -> _BuiltReasoningItem | None:
|
||||
from openai.types.responses import ResponseReasoningItem
|
||||
|
||||
if isinstance(item, ResponseReasoningItem):
|
||||
return _build_reasoning_item(
|
||||
item_id=item.id,
|
||||
encrypted_content=getattr(item, "encrypted_content", None),
|
||||
summary_raw=item.summary,
|
||||
)
|
||||
if isinstance(item, dict) and item.get("type") == "reasoning":
|
||||
return _build_reasoning_item(
|
||||
item_id=item.get("id", ""),
|
||||
encrypted_content=item.get("encrypted_content"),
|
||||
summary_raw=item.get("summary"),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[_BuiltReasoningItem, ...]:
|
||||
return tuple(
|
||||
reasoning_item
|
||||
for reasoning_item in (_reasoning_item_from_output_item(item) for item in output_items)
|
||||
if reasoning_item is not None
|
||||
)
|
||||
|
||||
|
||||
def _as_chat_reasoning_items(
|
||||
reasoning_items: Sequence[_BuiltReasoningItem],
|
||||
) -> list[ChatCompletionReasoningItem] | None:
|
||||
if not reasoning_items:
|
||||
return None
|
||||
# cast-ok: _BuiltReasoningItem is the structural shape ChatCompletionReasoningItem
|
||||
# describes, and TypedDict invariance is what stops the two from unifying here.
|
||||
return cast(list[ChatCompletionReasoningItem], list(reasoning_items))
|
||||
|
||||
|
||||
def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Literal["length", "content_filter"]:
|
||||
if incomplete_reason == "content_filter":
|
||||
return "content_filter"
|
||||
return "length"
|
||||
|
||||
|
||||
def _incomplete_reason_from_response_payload(response_payload: object) -> str | None:
|
||||
if not isinstance(response_payload, Mapping):
|
||||
return None
|
||||
incomplete_details: Final = response_payload.get("incomplete_details")
|
||||
if not isinstance(incomplete_details, Mapping):
|
||||
return None
|
||||
reason: Final = incomplete_details.get("reason")
|
||||
return reason if isinstance(reason, str) else None
|
||||
|
||||
|
||||
class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
|
||||
provider_specific_fields: Mapping[str, object]
|
||||
|
||||
|
|
@ -320,8 +392,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
)
|
||||
)
|
||||
elif role == "assistant" and tool_calls and isinstance(tool_calls, list):
|
||||
for r_item in _get_reasoning_items(msg):
|
||||
input_items.append(_reasoning_item_to_response_input(r_item))
|
||||
input_items.extend(_reasoning_input_items(msg))
|
||||
if content:
|
||||
input_items.append(
|
||||
{ # mutable-ok: API message payload
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": self._convert_content_to_responses_format(content, "assistant"),
|
||||
}
|
||||
)
|
||||
for tool_call in tool_calls:
|
||||
function = tool_call.get("function")
|
||||
custom = tool_call.get("custom")
|
||||
|
|
@ -348,15 +427,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
raise ValueError(f"tool call not supported: {tool_call}")
|
||||
elif content is not None:
|
||||
if role == "assistant":
|
||||
for r_item in _get_reasoning_items(msg):
|
||||
input_items.append(_reasoning_item_to_response_input(r_item))
|
||||
input_items.extend(_reasoning_input_items(msg))
|
||||
input_items.append(
|
||||
{
|
||||
{ # mutable-ok: API message payload
|
||||
"type": "message",
|
||||
"role": role,
|
||||
"content": self._convert_content_to_responses_format(content, cast(str, role)),
|
||||
}
|
||||
)
|
||||
elif role == "assistant":
|
||||
input_items.extend(_reasoning_input_items(msg))
|
||||
|
||||
return input_items, instructions
|
||||
|
||||
|
|
@ -657,6 +737,27 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
|
||||
return choices
|
||||
|
||||
@staticmethod
|
||||
def _build_empty_incomplete_choice(
|
||||
output_items: Sequence[object],
|
||||
finish_reason: Literal["length", "content_filter"],
|
||||
) -> "Choices":
|
||||
from litellm.types.utils import Choices, Message
|
||||
|
||||
reasoning_items: Final = _reasoning_items_from_output_items(output_items)
|
||||
reasoning_content: Final = " ".join(
|
||||
summary_block["text"]
|
||||
for reasoning_item in reasoning_items
|
||||
for summary_block in reasoning_item["summary"]
|
||||
if summary_block.get("text")
|
||||
)
|
||||
message: Final = Message(
|
||||
content="",
|
||||
reasoning_content=reasoning_content if reasoning_content else None,
|
||||
reasoning_items=_as_chat_reasoning_items(reasoning_items),
|
||||
)
|
||||
return Choices(message=message, finish_reason=finish_reason, index=0)
|
||||
|
||||
@classmethod
|
||||
def _extract_output_from_completed_event(cls, parsed_chunk: Mapping[str, object]) -> list[dict[str, object]] | None:
|
||||
response_payload: Final = parsed_chunk.get("response")
|
||||
|
|
@ -763,11 +864,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
handle_raw_dict_callback=self._handle_raw_dict_response_item,
|
||||
)
|
||||
|
||||
if len(choices) == 0:
|
||||
if raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None:
|
||||
raise ValueError(f"{model} unable to complete request: {raw_response.incomplete_details.reason}")
|
||||
response_is_incomplete: Final = raw_response.status == "incomplete" or (
|
||||
raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None
|
||||
)
|
||||
|
||||
if len(choices) == 0 and not response_is_incomplete:
|
||||
raise ValueError(f"Unknown items in responses API response: {output_items}")
|
||||
|
||||
if response_is_incomplete:
|
||||
incomplete_finish_reason: Final = _map_incomplete_reason_to_finish_reason(
|
||||
raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None
|
||||
)
|
||||
if len(choices) == 0:
|
||||
choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason))
|
||||
else:
|
||||
raise ValueError(f"Unknown items in responses API response: {output_items}")
|
||||
for choice in choices:
|
||||
choice.finish_reason = incomplete_finish_reason
|
||||
|
||||
setattr(model_response, "choices", choices)
|
||||
|
||||
|
|
@ -1002,22 +1114,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
|
||||
)
|
||||
|
||||
# If string is passed, map with optional summary based on flag/env var
|
||||
if reasoning_effort == "none":
|
||||
return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none")
|
||||
elif reasoning_effort == "high":
|
||||
return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high")
|
||||
elif reasoning_effort == "xhigh":
|
||||
return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh")
|
||||
elif reasoning_effort == "medium":
|
||||
if reasoning_effort in get_args(REASONING_EFFORT):
|
||||
return (
|
||||
Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium")
|
||||
)
|
||||
elif reasoning_effort == "low":
|
||||
return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low")
|
||||
elif reasoning_effort == "minimal":
|
||||
return (
|
||||
Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal")
|
||||
Reasoning(effort=reasoning_effort, summary="detailed")
|
||||
if auto_summary_enabled
|
||||
else Reasoning(effort=reasoning_effort)
|
||||
)
|
||||
return None
|
||||
|
||||
|
|
@ -1392,12 +1493,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
)
|
||||
]
|
||||
)
|
||||
elif event_type == "response.completed":
|
||||
# Response is fully complete - now we can signal is_finished=True
|
||||
# This ensures we don't prematurely end the stream before tool_calls arrive
|
||||
|
||||
# Check if response contains function_call items in output
|
||||
# to determine correct finish_reason
|
||||
elif event_type in ("response.completed", "response.incomplete"):
|
||||
response_data: Final = parsed_chunk.get("response", {})
|
||||
output_items: Final = response_data.get("output", []) if response_data else []
|
||||
|
||||
|
|
@ -1407,25 +1503,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
if isinstance(item, dict)
|
||||
)
|
||||
|
||||
finish_reason: Final = "tool_calls" if has_function_calls else "stop"
|
||||
finish_reason: Final = (
|
||||
_map_incomplete_reason_to_finish_reason(_incomplete_reason_from_response_payload(response_data))
|
||||
if event_type == "response.incomplete"
|
||||
else ("tool_calls" if has_function_calls else "stop")
|
||||
)
|
||||
|
||||
# Extract reasoning items with encrypted_content for round-tripping
|
||||
completed_reasoning_items: list[_BuiltReasoningItem] | None = None
|
||||
for item in output_items:
|
||||
if not isinstance(item, dict) or item.get("type") != "reasoning":
|
||||
continue
|
||||
if completed_reasoning_items is None:
|
||||
completed_reasoning_items = []
|
||||
completed_reasoning_items.append(
|
||||
_build_reasoning_item(
|
||||
item_id=item.get("id", ""),
|
||||
encrypted_content=item.get("encrypted_content"),
|
||||
summary_raw=item.get("summary"),
|
||||
)
|
||||
)
|
||||
completed_reasoning_items_typed: Final = cast(
|
||||
list[ChatCompletionReasoningItem] | None,
|
||||
completed_reasoning_items,
|
||||
terminal_reasoning_items_typed: Final = _as_chat_reasoning_items(
|
||||
_reasoning_items_from_output_items(output_items)
|
||||
)
|
||||
|
||||
usage = None
|
||||
|
|
@ -1439,7 +1524,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
index=0,
|
||||
delta=Delta(
|
||||
content="",
|
||||
reasoning_items=completed_reasoning_items_typed,
|
||||
reasoning_items=terminal_reasoning_items_typed,
|
||||
),
|
||||
finish_reason=finish_reason,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import sys
|
|||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none
|
||||
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_range, get_env_int_or_none
|
||||
|
||||
DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
|
||||
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
|
||||
|
|
@ -48,6 +48,11 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = (
|
|||
# Data URIs exceeding this are replaced with a size placeholder.
|
||||
# Set to 0 to disable truncation.
|
||||
MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
|
||||
REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
|
||||
# in-memory stand-in handed to provider converters for redacted arguments; never stored
|
||||
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}"
|
||||
|
||||
MAX_STRING_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", 4096)
|
||||
|
||||
# When true, adds detailed per-phase timing breakdown headers to responses.
|
||||
# Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms
|
||||
|
|
@ -144,6 +149,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
|
|||
"x-litellm-adaptive-router-model",
|
||||
"x-litellm-applied-guardrails",
|
||||
"x-litellm-guardrail-scan-id",
|
||||
"x-litellm-cache-key",
|
||||
]
|
||||
|
||||
# Gemini model-specific minimal thinking budget constants
|
||||
|
|
@ -323,6 +329,17 @@ DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RE
|
|||
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20))
|
||||
MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768))
|
||||
MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000))
|
||||
# tiktoken's BPE merge loop is quadratic in the length of a single regex piece, so a long run of one
|
||||
# repeated character (dot leaders, whitespace, zero-padded base64) can take minutes on a multi-MB payload.
|
||||
# Encoding in chunks makes the cost linear, at a drift of at most ~1 token per chunk boundary. The upper
|
||||
# bound keeps a misconfigured chunk size from restoring the quadratic cost this exists to remove.
|
||||
TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS: Final = 4096
|
||||
TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
|
||||
"TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS",
|
||||
default=1024,
|
||||
minimum=1,
|
||||
maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
|
||||
)
|
||||
MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
|
||||
MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
|
||||
OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))
|
||||
|
|
@ -364,6 +381,7 @@ AZURE_OPERATION_POLLING_TIMEOUT: Final = int(os.getenv("AZURE_OPERATION_POLLING_
|
|||
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30"))
|
||||
AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96))
|
||||
REDIS_SOCKET_TIMEOUT: Final = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1))
|
||||
CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS: Final[float] = 5.0
|
||||
REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5))
|
||||
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5))
|
||||
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60))
|
||||
|
|
@ -423,6 +441,9 @@ DEFAULT_REQUEST_TIMEOUT_SECONDS: Final[float] = 6000.0
|
|||
# deadline and connect handshake (see ``http_handler`` cached handler paths).
|
||||
COMPLETION_HTTP_FALLBACK_SECONDS: Final[float] = 600.0
|
||||
HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: Final[float] = 5.0
|
||||
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = float(
|
||||
os.getenv("SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS", "5.0")
|
||||
)
|
||||
request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS))))
|
||||
request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ
|
||||
DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes
|
||||
|
|
@ -444,6 +465,8 @@ CONNECTION_ERROR_PATTERNS: Final[list[str]] = [
|
|||
]
|
||||
STREAM_SSE_DONE_STRING: Final[str] = "[DONE]"
|
||||
STREAM_SSE_DATA_PREFIX: Final[str] = "data: "
|
||||
STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n'
|
||||
STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8")
|
||||
### SPEND TRACKING ###
|
||||
DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = float(
|
||||
os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400)
|
||||
|
|
@ -467,6 +490,9 @@ MAX_TIME_TO_CLEAR_QUEUE: Final = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0)
|
|||
LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
|
||||
os.getenv("LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS", 0.5)
|
||||
) # Cooldown time in seconds before allowing another aggressive clear (default: 0.5s)
|
||||
LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
|
||||
LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
|
||||
LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0
|
||||
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
|
||||
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
|
||||
)
|
||||
|
|
@ -730,6 +756,7 @@ openai_compatible_endpoints: Final[list] = [
|
|||
"api.groq.com/openai/v1",
|
||||
"https://integrate.api.nvidia.com/v1",
|
||||
"api.deepseek.com/v1",
|
||||
"api.together.ai/v1",
|
||||
"api.together.xyz/v1",
|
||||
"app.empower.dev/api/v1",
|
||||
"https://api.friendli.ai/serverless/v1",
|
||||
|
|
@ -763,6 +790,8 @@ openai_compatible_endpoints: Final[list] = [
|
|||
"https://api.libertai.io/v1",
|
||||
"https://pinstripes.io/v1",
|
||||
"https://api.meta.ai/v1",
|
||||
"https://api.cognition.ai/v1",
|
||||
"https://api.scx.ai/v1",
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -830,6 +859,8 @@ openai_compatible_providers: Final[list] = [
|
|||
"pinstripes", # Pinstripes - JSON-configured provider
|
||||
"darkbloom",
|
||||
"meta", # Meta Model API (Muse Spark) - JSON-configured provider
|
||||
"cognition",
|
||||
"scx-ai",
|
||||
]
|
||||
openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions`
|
||||
"together_ai",
|
||||
|
|
@ -1333,6 +1364,8 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
|
|||
LITELLM_METADATA_FIELD: Final = "litellm_metadata"
|
||||
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
|
||||
AUTO_ROUTED_REQUEST_METADATA_KEY: Final = "_auto_routed_request"
|
||||
ROUTER_MODEL_NAME_RESPONSE_FIELD: Final = "router_model_name"
|
||||
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
|
||||
|
|
@ -1342,6 +1375,11 @@ LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
|
|||
"Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.). "
|
||||
"To increase the truncation limit, set `MAX_STRING_LENGTH_PROMPT_IN_DB` in your env."
|
||||
)
|
||||
LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE: Final = (
|
||||
"Truncation is a stdout logging safeguard. "
|
||||
"Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.) and at DEBUG level. "
|
||||
"To increase the truncation limit, set `MAX_STRING_LENGTH_STDOUT_LOG` in your env."
|
||||
)
|
||||
|
||||
########################### LiteLLM Proxy Specific Constants ###########################
|
||||
########################################################################################
|
||||
|
|
@ -1508,14 +1546,22 @@ TOOL_SPEND_TOP_TOOLS: Final = 100
|
|||
SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
|
||||
SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
|
||||
SPEND_LOG_WRITE_BATCH_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_BYTES", 2_000_000)))
|
||||
SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_ROWS", "100")))
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
|
||||
SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYTES", "64000000")))
|
||||
SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
|
||||
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3")))
|
||||
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2"))
|
||||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
|
||||
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
|
||||
RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500")))
|
||||
RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100")))
|
||||
RESET_BUDGET_JOB_NAME: Final = "reset_budget_job"
|
||||
# Comfortably longer than one PROXY_BUDGET_RESCHEDULER_MIN_TIME tick, so a healthy
|
||||
# leader keeps the lease across its own run, and a crashed one strands the sweep for
|
||||
# at most a single tick.
|
||||
RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
|
||||
PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
|
||||
MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)))
|
||||
|
|
@ -1525,6 +1571,19 @@ STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, int(os.getenv("STALE_OBJECT_CLEA
|
|||
# installations with large numbers of stale managed objects).
|
||||
_batch_polling_env: Final = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
|
||||
PROXY_BATCH_POLLING_ENABLED: Final = _batch_polling_env == "true"
|
||||
BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS: Final = float(
|
||||
os.getenv("BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS", "5")
|
||||
)
|
||||
BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS: Final = float(
|
||||
os.getenv("BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS", "60")
|
||||
)
|
||||
BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS: Final = float(
|
||||
os.getenv("BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", "3600")
|
||||
)
|
||||
_background_interaction_cost_polling_env: Final = os.getenv(
|
||||
"BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true"
|
||||
).lower()
|
||||
BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true"
|
||||
PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605))
|
||||
PROXY_BATCH_WRITE_AT: Final = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10
|
||||
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)
|
||||
|
|
@ -1580,6 +1639,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
|
|||
"public_model_groups_links",
|
||||
"cost_discount_config",
|
||||
"cost_margin_config",
|
||||
"block_requests_for_models_without_pricing",
|
||||
"budget_exceeded_throttle_percentage",
|
||||
# Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS)
|
||||
# must be listed here so a DB write from one worker overrides the live litellm attribute on
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
|
|||
StandardBuiltInToolCostTracking,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
|
||||
InteractionsUsageObjectTransformation,
|
||||
TranscriptionUsageObjectTransformation,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
|
|
@ -150,6 +151,7 @@ _VIDEO_CALL_TYPES: Final = frozenset(
|
|||
}
|
||||
)
|
||||
|
||||
|
||||
_SPEECH_CALL_TYPES: Final = frozenset(
|
||||
{
|
||||
CallTypes.speech.value,
|
||||
|
|
@ -793,14 +795,27 @@ def _select_model_name_for_cost_calc(
|
|||
and custom_llm_provider is not None
|
||||
and not _model_contains_known_llm_provider(return_model)
|
||||
): # add provider prefix if not already present, to match model_cost
|
||||
if region_name is not None:
|
||||
return_model = f"{custom_llm_provider}/{region_name}/{return_model}"
|
||||
else:
|
||||
return_model = f"{custom_llm_provider}/{return_model}"
|
||||
provider_prefix: Final = custom_llm_provider if region_name is None else f"{custom_llm_provider}/{region_name}"
|
||||
return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", region_name)
|
||||
|
||||
return return_model
|
||||
|
||||
|
||||
def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str:
|
||||
"""Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the
|
||||
registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already
|
||||
resolves downstream (custom-priced router ids) or no stripped candidate is registered (#38069)."""
|
||||
segments: Final = model.split("/")
|
||||
if "/".join(segments[1:]) in litellm.model_cost:
|
||||
return model
|
||||
head_len: Final = 2 if region_name is not None and len(segments) > 2 and segments[1] == region_name else 1
|
||||
head: Final = "/".join(segments[:head_len])
|
||||
tail: Final = segments[head_len:]
|
||||
strippable: Final = next((index for index, segment in enumerate(tail) if segment in LlmProvidersSet), len(tail))
|
||||
candidates: Final = (f"{head}/{'/'.join(tail[start:])}" for start in range(min(strippable, len(tail) - 1) + 1))
|
||||
return next((candidate for candidate in candidates if candidate in litellm.model_cost), model)
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _model_contains_known_llm_provider(model: str) -> bool:
|
||||
"""
|
||||
|
|
@ -915,6 +930,8 @@ def _get_usage_object(
|
|||
usage_obj,
|
||||
)
|
||||
)
|
||||
elif isinstance(usage_obj, dict) and InteractionsUsageObjectTransformation.is_interactions_usage_object(usage_obj):
|
||||
return InteractionsUsageObjectTransformation.transform_interactions_usage_object(usage_obj)
|
||||
elif isinstance(usage_obj, dict):
|
||||
return Usage(**usage_obj)
|
||||
elif isinstance(usage_obj, BaseModel):
|
||||
|
|
@ -1291,6 +1308,10 @@ def completion_cost(
|
|||
)
|
||||
if tr_usage is not None:
|
||||
_usage = tr_usage.model_dump()
|
||||
elif InteractionsUsageObjectTransformation.is_interactions_usage_object(_usage):
|
||||
_usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
|
||||
_usage
|
||||
).model_dump()
|
||||
else:
|
||||
_usage = _usage
|
||||
|
||||
|
|
@ -1375,23 +1396,36 @@ def completion_cost(
|
|||
if custom_pricing and litellm_logging_obj is not None:
|
||||
_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
|
||||
if _litellm_params is not None:
|
||||
_metadata = _litellm_params.get("metadata", {}) or {}
|
||||
_video_model_info = _metadata.get("model_info", None)
|
||||
_video_model_info = next(
|
||||
(
|
||||
model_info
|
||||
for _metadata_key in ("metadata", "litellm_metadata")
|
||||
if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info"))
|
||||
is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
usage_obj = getattr(completion_response, "usage", None)
|
||||
duration_seconds: float | None = None
|
||||
video_resolution: str | None = None
|
||||
provider_reported_cost: float | None = None
|
||||
if completion_response is not None and usage_obj:
|
||||
# Handle both dict and Pydantic Usage object
|
||||
if isinstance(usage_obj, dict):
|
||||
duration_seconds = usage_obj.get("duration_seconds", None)
|
||||
_vr = usage_obj.get("video_resolution", None)
|
||||
provider_reported_cost = usage_obj.get("provider_reported_cost_usd", None)
|
||||
else:
|
||||
duration_seconds = getattr(usage_obj, "duration_seconds", None)
|
||||
_vr = getattr(usage_obj, "video_resolution", None)
|
||||
provider_reported_cost = getattr(usage_obj, "provider_reported_cost_usd", None)
|
||||
if _vr is not None:
|
||||
video_resolution = str(_vr).strip().lower()
|
||||
|
||||
if _video_model_info is None and provider_reported_cost is not None:
|
||||
return float(provider_reported_cost)
|
||||
|
||||
if duration_seconds is not None:
|
||||
# Calculate cost based on video duration using video-specific cost calculation
|
||||
from litellm.llms.openai.cost_calculation import (
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
|
|||
from litellm.exceptions import LiteLLMUnknownProvider
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_request_utils import flatten_form_field_values
|
||||
from litellm.litellm_core_utils.mock_functions import mock_image_generation
|
||||
from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
|
|
@ -422,24 +423,32 @@ def image_generation(
|
|||
aimg_generation=aimg_generation,
|
||||
)
|
||||
elif custom_llm_provider == "azure_ai":
|
||||
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AzureFoundryModelInfo,
|
||||
get_azure_ai_auth_headers,
|
||||
)
|
||||
|
||||
api_base = AzureFoundryModelInfo.get_api_base(api_base)
|
||||
api_key = AzureFoundryModelInfo.get_api_key(api_key)
|
||||
if extra_headers is not None:
|
||||
optional_params["extra_headers"] = extra_headers
|
||||
|
||||
default_headers = {
|
||||
caller_header_names = frozenset(name.lower() for name in headers)
|
||||
caller_set_auth = "api-key" in caller_header_names or "authorization" in caller_header_names
|
||||
auth_headers = (
|
||||
headers
|
||||
if caller_set_auth
|
||||
else get_azure_ai_auth_headers(
|
||||
api_key=api_key,
|
||||
litellm_params=litellm_params_dict,
|
||||
api_key_header="api-key",
|
||||
)
|
||||
)
|
||||
request_headers: Final = {
|
||||
"Content-Type": "application/json",
|
||||
**auth_headers,
|
||||
**headers,
|
||||
}
|
||||
# Only add api-key header if api_key is not None
|
||||
# Azure AD authentication will use Authorization header instead
|
||||
if api_key is not None:
|
||||
default_headers["api-key"] = api_key
|
||||
|
||||
for k, v in default_headers.items():
|
||||
if k not in headers:
|
||||
headers[k] = v
|
||||
|
||||
model_response = azure_chat_completions.image_generation(
|
||||
model=model,
|
||||
|
|
@ -455,7 +464,7 @@ def image_generation(
|
|||
api_version=api_version,
|
||||
aimg_generation=aimg_generation,
|
||||
client=client,
|
||||
headers=headers,
|
||||
headers=request_headers,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
elif (
|
||||
|
|
@ -846,6 +855,18 @@ def image_edit(
|
|||
additional_drop_params=kwargs.get("additional_drop_params"),
|
||||
)
|
||||
|
||||
if (
|
||||
custom_llm_provider == "openai"
|
||||
or custom_llm_provider == "azure"
|
||||
or custom_llm_provider in litellm.openai_compatible_providers
|
||||
):
|
||||
image_edit_request_params.update(
|
||||
flatten_form_field_values(
|
||||
non_default_params,
|
||||
extra_body if isinstance(extra_body, dict) else None,
|
||||
)
|
||||
)
|
||||
|
||||
# Pre Call logging
|
||||
litellm_logging_obj.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
|
|
@ -995,6 +1016,9 @@ async def aimage_edit(
|
|||
response_format=response_format,
|
||||
size=size,
|
||||
user=user,
|
||||
extra_headers=extra_headers,
|
||||
extra_query=extra_query,
|
||||
extra_body=extra_body,
|
||||
timeout=timeout,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
**kwargs,
|
||||
|
|
|
|||
|
|
@ -27,11 +27,15 @@ from litellm.types.integrations.anthropic_cache_control_hook import (
|
|||
CacheControlInjectionPoint,
|
||||
CacheControlMessageInjectionPoint,
|
||||
)
|
||||
from litellm.types.llms.anthropic import AnthropicSystemMessageContent
|
||||
from litellm.types.llms.anthropic import (
|
||||
AllAnthropicToolsValues,
|
||||
AnthropicSystemMessageContent,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionCachedContent,
|
||||
ChatCompletionTextObject,
|
||||
ChatCompletionToolParam,
|
||||
PromptCacheBreakpoint,
|
||||
PromptCacheOptions,
|
||||
)
|
||||
|
|
@ -57,6 +61,8 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
|
|||
OPENAI_API_HOST: Final = "api.openai.com"
|
||||
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
|
||||
|
||||
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
|
||||
|
||||
|
||||
def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
|
||||
model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
|
||||
|
|
@ -625,6 +631,50 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
]
|
||||
return points
|
||||
|
||||
@staticmethod
|
||||
def messages_with_default_injections(
|
||||
messages: list[AllMessageValues],
|
||||
models: Iterable[str],
|
||||
tools: list[AllToolParamValues] | None = None,
|
||||
enable_prompt_caching: bool | None = None,
|
||||
) -> list[AllMessageValues]:
|
||||
"""Return the messages auto prompt caching will send, default breakpoints included.
|
||||
|
||||
Router cache affinity depends on this. Deployment selection runs before the injection in
|
||||
`litellm.acompletion`, so it has to reproduce the markers to derive the same cache key the
|
||||
success event later writes from the sent messages. `models` is every candidate model of the
|
||||
group: the first that would auto-inject decides, since the default breakpoints (system
|
||||
prompt and trailing turn) do not depend on which deployment serves the call. Returns the
|
||||
input list itself when auto-injection would not apply
|
||||
"""
|
||||
points: Final = next(
|
||||
(
|
||||
candidate
|
||||
for candidate in (
|
||||
AnthropicCacheControlHook.get_default_injection_points(
|
||||
messages=messages,
|
||||
system=None,
|
||||
model=model,
|
||||
custom_llm_provider=None,
|
||||
tools=tools,
|
||||
enable_prompt_caching=enable_prompt_caching,
|
||||
)
|
||||
for model in models
|
||||
)
|
||||
if candidate
|
||||
),
|
||||
None,
|
||||
)
|
||||
if not points:
|
||||
return messages
|
||||
return AnthropicCacheControlHook._apply_message_injections(
|
||||
points=cast( # cast-ok: the default points are all message-location points
|
||||
list[CacheControlMessageInjectionPoint], points
|
||||
),
|
||||
messages=copy.deepcopy(messages),
|
||||
max_blocks=MAX_CACHE_CONTROL_BLOCKS,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def maybe_seed_default_injection_points(
|
||||
non_default_params: dict[str, Any],
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue