Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_decrease_anys_opus5

# Conflicts:
#	basedpyright-code-budget.json
#	litellm/caching/valkey_semantic_cache.py
#	litellm/integrations/compression_interception/handler.py
#	litellm/integrations/custom_logger.py
#	litellm/llms/custom_httpx/container_handler.py
#	litellm/llms/infinity/rerank/transformation.py
#	litellm/proxy/agent_endpoints/agent_registry.py
#	litellm/repositories/base_repository.py
#	litellm/repositories/credentials_repository.py
#	litellm/repositories/team_repository.py
#	ruff-strict-budget.json
#	type-discipline-budget.json
This commit is contained in:
mateo-berri 2026-08-29 06:03:33 -07:00
commit 8d4620649f
3582 changed files with 206739 additions and 44627 deletions

View file

@ -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 \
@ -1025,12 +905,12 @@ jobs:
name: Run tests
command: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/agent_tests/**/test_*.py" | grep -v "^tests/agent_tests/local_only_agent_tests/")
TEST_FILES=$(circleci tests glob "tests/agent_tests/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 \
-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"
@ -2575,45 +2421,6 @@ jobs:
- wait_for_service:
url: http://localhost:4000
timeout: "300"
# Add Ruby installation and testing before the existing Node.js and Python tests
- run:
name: Install Ruby and Bundler
command: |
# Clone RVM at pinned tag and verify the commit SHA matches the
# published tag before running its install script.
RVM_VERSION="1.29.12"
RVM_EXPECTED_SHA="6bfc9213c9d6914fe756f524eb034a403d51db81"
git clone --depth 1 --branch "$RVM_VERSION" https://github.com/rvm/rvm.git /tmp/rvm
RVM_ACTUAL_SHA="$(git -C /tmp/rvm rev-parse HEAD)"
if [ "$RVM_ACTUAL_SHA" != "$RVM_EXPECTED_SHA" ]; then
echo "RVM tag $RVM_VERSION resolved to $RVM_ACTUAL_SHA; expected $RVM_EXPECTED_SHA" >&2
exit 1
fi
# Import RVM signing keys (used by `rvm install` to verify Ruby tarballs)
gpg --keyserver hkp://keyserver.ubuntu.com --recv-keys 409B6B1796C275462A1703113804BB82D39DC0E3 7D2BAF1CF37B13E2069D6956105BD0E739499BDB
# Install RVM from the verified checkout. The install script
# sources `scripts/functions/installer` using paths relative to
# its own working directory, so it must be run from /tmp/rvm.
(cd /tmp/rvm && ./install --path "$HOME/.rvm")
source "$HOME/.rvm/scripts/rvm"
# Install Ruby 3.2.2 (RVM verifies the tarball PGP signature)
rvm install 3.2.2
rvm use 3.2.2 --default
# Install latest Bundler
gem install bundler
- run:
name: Run Ruby tests
command: |
source $HOME/.rvm/scripts/rvm
cd tests/pass_through_tests/ruby_passthrough_tests
bundle install
bundle exec rspec
no_output_timeout: 30m
# Install Node.js directly from nodejs.org with SHA256 verification,
# instead of piping NodeSource's setup_24.x apt-repo installer into
# sudo bash (which runs a mutable upstream script unattended).
@ -3105,12 +2912,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 +2968,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 +2990,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
View file

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

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

View file

@ -4,6 +4,22 @@ description: >-
by a job nor listed here, so every entry below is a decision on the record.
test_paths:
- reason: >-
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_caching.py
- tests/local_testing/test_disk_cache_unit_tests.py
- tests/local_testing/test_gcs_cache_unit_tests.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
@ -21,72 +37,14 @@ test_paths:
- tests/documentation_tests/test_requests_lib_usage.py
- tests/documentation_tests/test_standard_logging_payload.py
- reason: >-
Sibling files here are executed by name from the code-quality workflow; this one is referenced
by no job
Named like a test but shaped like a benchmark: it fetches live image URLs, times aiohttp
against httpx, prints the ratio, and asserts nothing, so pytest cannot collect it (its
functions take arguments, not fixtures) and running it beside its siblings in the
code-quality workflow would add a network dependency for a number nothing reads. Exempt
as a script rather than as an unresolved gap; revisit by deleting it once the aiohttp
choice it informed is settled
paths:
- tests/code_coverage_tests/test_aio_http_image_conversion.py
- reason: >-
A second mirror of the package tree living beside tests/test_litellm, which is the mirror the
repo convention names; only test_no_hardcoded_secrets.py is invoked, from the linting
workflow, and whether this directory should exist at all is unresolved
paths:
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
- tests/litellm/integrations/helicone/test_helicone_gemini.py
- tests/litellm/litellm_core_utils/test_json_schema_validation.py
- tests/litellm/llms/anthropic/test_anthropic_reasoning_effort.py
- tests/litellm/llms/anthropic/test_anthropic_schema_filter.py
- tests/litellm/llms/azure/test_azure_embedding.py
- tests/litellm/llms/bedrock/embed/test_embedding.py
- tests/litellm/llms/bedrock/test_nova_imported_models.py
- tests/litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py
- tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py
- tests/litellm/llms/oci/chat/test_oci_chat_transformation.py
- tests/litellm/llms/openai_like/test_abliteration_provider.py
- tests/litellm/llms/openai_like/test_assemblyai_provider.py
- tests/litellm/llms/openai_like/test_empiriolabs_provider.py
- tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py
- tests/litellm/llms/vertex_ai/gemini/test_transformation.py
- tests/litellm/llms/vertex_ai/text_to_speech/test_transformation.py
- tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
- tests/litellm/proxy/agent_endpoints/test_agent_rbac.py
- tests/litellm/proxy/common_utils/test_rbac_utils.py
- tests/litellm/proxy/management_endpoints/test_common_utils.py
- tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py
- tests/litellm/proxy/test_claude_code_marketplace.py
- tests/litellm/proxy/test_init_litellm_callbacks.py
- tests/litellm/proxy/test_prisma_engine_watchdog.py
- tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py
- tests/litellm/test_bedrock_extended_beta_models.py
- tests/litellm/test_bedrock_nemotron_super.py
- tests/litellm/test_proxy_auth.py
- tests/litellm/test_router_retry_backoff_headers.py
- tests/litellm/test_sambanova_model_metadata.py
- tests/litellm/test_stream_chunk_builder_images.py
- reason: >-
Legacy proxy suite superseded by the proxy shards; no job invokes it and whether it still
describes supported behaviour is unresolved
paths:
- tests/old_proxy_tests/tests/test_anthropic_context_caching.py
- tests/old_proxy_tests/tests/test_anthropic_sdk.py
- tests/old_proxy_tests/tests/test_async.py
- tests/old_proxy_tests/tests/test_gemini_context_caching.py
- tests/old_proxy_tests/tests/test_langchain_embedding.py
- tests/old_proxy_tests/tests/test_langchain_request.py
- tests/old_proxy_tests/tests/test_llamaindex.py
- tests/old_proxy_tests/tests/test_mistral_sdk.py
- tests/old_proxy_tests/tests/test_openai_embedding.py
- tests/old_proxy_tests/tests/test_openai_exception_request.py
- tests/old_proxy_tests/tests/test_openai_request.py
- tests/old_proxy_tests/tests/test_openai_request_with_traceparent.py
- tests/old_proxy_tests/tests/test_openai_simple_embedding.py
- tests/old_proxy_tests/tests/test_openai_tts_request.py
- tests/old_proxy_tests/tests/test_pass_through_langfuse.py
- tests/old_proxy_tests/tests/test_q.py
- tests/old_proxy_tests/tests/test_simple_traceparent_openai.py
- tests/old_proxy_tests/tests/test_vertex_sdk_forward_headers.py
- tests/old_proxy_tests/tests/test_vtx_embedding.py
- tests/old_proxy_tests/tests/test_vtx_sdk_embedding.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
@ -116,6 +74,14 @@ test_paths:
- tests/load_tests/test_otel_load_test.py
- tests/load_tests/test_vertex_embeddings_load_test.py
- tests/load_tests/test_vertex_load_tests.py
- reason: >-
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a
pull request job. Until 2026-08-20 the CircleCI agent job hid them behind a grep -v
that this census could not see; the glob now excludes them structurally and this entry
is the decision on the record. Revisit when the A2A bridge gets a recorded-wire fixture
paths:
- tests/agent_tests/local_only_agent_tests
- reason: >-
Third-party integration tests that skip themselves without OCI configuration or sandbox
credentials, neither of which a pull request job holds
@ -123,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: >-
Two prompt-factory tests sitting at the top level of tests/ instead of under the
tests/test_litellm mirror the shards enumerate; they need moving rather than a shard entry
paths:
- tests/litellm_core_utils/test_anthropic_dedup_factory.py
- tests/litellm_core_utils/test_bedrock_converse_dedup_factory.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
paths:
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
dockerfiles:
- reason: >-

5
.github/mutmut-coverage.rc vendored Normal file
View 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

View file

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

View file

@ -1,10 +1,14 @@
from __future__ import annotations
import ast
import operator
import pathlib
import re
import sys
from collections.abc import Iterable, Mapping, Sequence
import warnings
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final
import yaml
@ -25,6 +29,15 @@ DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
GLOB_CHARS = frozenset("*?")
# Trees whose jobs are sharded with no catch-all bucket, so every child that holds
# tests has to be named by some shard or it runs nowhere. A child listed here is
# itself decomposed one level deeper and is checked through its own entry.
SHARDED_ROOTS: tuple[str, ...] = (
"tests/proxy_unit_tests",
"tests/test_litellm",
"tests/test_litellm/proxy",
)
@dataclass(frozen=True, slots=True)
class AllowEntry:
@ -44,6 +57,14 @@ class Allowlist:
return any(relative_path == path for entry in self.dockerfiles for path in entry.paths)
@dataclass(frozen=True, slots=True)
class Section:
name: str
entries: tuple[AllowEntry, ...]
candidates: tuple[str, ...]
matches: Callable[[str, str], bool]
@dataclass(frozen=True, slots=True)
class Scalar:
key: str
@ -107,20 +128,34 @@ def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
)
def _glob_to_regex(token: str) -> re.Pattern[str]:
parts = re.split(r"(\*\*/|\*\*|\*|\?)", token)
def _glob_to_regex(token: str, *, subtree: bool) -> re.Pattern[str]:
parts = re.split(r"(\*\*/|\*\*|\*|\?|\[[^\]]*\])", token)
translated = "".join(
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part, re.escape(part)) for part in parts
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part)
or (part if part.startswith("[") and part.endswith("]") else re.escape(part))
for part in parts
)
return re.compile(rf"{translated}(?:/.*)?$")
return re.compile(rf"{translated}(?:/.*)?$" if subtree else rf"{translated}$")
def _token_covers(token: str, relative_path: str) -> bool:
if GLOB_CHARS & set(token):
return _glob_to_regex(token).match(relative_path) is not None
return _glob_to_regex(token, subtree=True).match(relative_path) is not None
return relative_path == token or relative_path.startswith(f"{token}/")
def _token_names(token: str, relative_path: str) -> bool:
"""Whether the token names this path itself, rather than merely containing it.
A sharded tree has no catch-all bucket, so the ancestor token the census is happy
with (`tests/x` standing in for everything below it) is exactly what would let a
newly added child ride along without a shard.
"""
if GLOB_CHARS & set(token):
return _glob_to_regex(token, subtree=False).match(relative_path) is not None
return token == relative_path
def _test_files() -> tuple[str, ...]:
return tuple(
sorted(
@ -166,6 +201,190 @@ def _describe(paths: tuple[str, ...]) -> str:
return f"{len(paths)} test file(s) invoked by no job: {names}{suffix}"
GLOB_CALL_RE = re.compile(r'circleci tests glob "([^"]+)"')
KEYWORD_RE = re.compile(r"-k\s+\\?[\"']([^\"'\\]+)")
@dataclass(frozen=True, slots=True)
class Slice:
"""One job's selection: the files it globs, narrowed by its `-k` expression."""
job: str
globs: tuple[str, ...]
named: frozenset[str]
required: tuple[str, ...]
excluded: tuple[str, ...]
understood: bool
def claims(self, relative_path: str, inner_names: frozenset[str]) -> bool:
"""Whether this job runs any test in the file.
The question is deliberately per-file, not per-test. An excluded term is only
honoured when it appears in the path, because that is the case where it takes
the whole module with it; a term matching one function inside drops that test
and leaves the file claimed. Losing a whole file is the failure worth a gate,
and answering per-test would mean a baseline of test ids that churns on every
rename.
"""
if relative_path in self.named:
return True
if not any(_token_covers(glob, relative_path) for glob in self.globs):
return False
if not self.understood:
return True # a `-k` this parser cannot model is assumed to claim everything
if any(term.lower() in relative_path.lower() for term in self.excluded):
return False
return not self.required or any(
term.lower() in name.lower() for term in self.required for name in inner_names
)
def _strings(node: object) -> Iterable[str]:
if isinstance(node, str):
yield node
elif isinstance(node, dict):
for value in node.values():
yield from _strings(value)
elif isinstance(node, list):
for value in node:
yield from _strings(value)
def _keyword_terms(
expressions: Sequence[str], *, attributable: bool = True
) -> tuple[tuple[str, ...], tuple[str, ...], bool]:
"""A `-k` expression as (required, excluded, understood).
Only flat `and` chains of bare terms are modelled. Anything with `or`, parentheses
or negation of a group is left unmodelled, and its job is then treated as claiming
every file it globs, so an unparsed selector can never raise a false alarm.
`attributable` is False when a job runs several pytest commands, since a selector
read out of the job's text cannot then be tied to the glob it belongs to, and
pairing one command's exclusion with another's glob would invent a gap.
"""
terms: Final = tuple(part.strip() for expression in expressions for part in expression.split(" and "))
if not attributable and terms:
return (), (), False
if any(("or " in term) or ("(" in term) or (term.startswith("not ") and " " in term[4:]) for term in terms):
return (), (), False
return (
tuple(term for term in terms if term and not term.startswith("not ")),
tuple(term[4:].strip() for term in terms if term.startswith("not ")),
True,
)
def _slices() -> tuple[Slice, ...]:
if not CIRCLECI_CONFIG.exists():
return ()
jobs: Final = yaml.safe_load(CIRCLECI_CONFIG.read_text()).get("jobs", {})
return tuple(
Slice(job=job, globs=globs, named=named, required=required, excluded=excluded, understood=understood)
for job, body in jobs.items()
for text in ("\n".join(_strings(body)),)
if "pytest" in text
for globs in (tuple(GLOB_CALL_RE.findall(text)),)
for named in (frozenset(TEST_TOKEN_RE.findall(text)) & frozenset(_test_files()),)
for required, excluded, understood in (
_keyword_terms(tuple(KEYWORD_RE.findall(text)), attributable=len(globs) < 2),
)
if globs or named
)
def _matchable_names(relative_path: str) -> frozenset[str]:
"""Every name a `-k` term can match for this file: its path, plus the names inside it.
pytest matches a keyword against an item's own name and each of its parents', so a
positive term hits a file when it appears in the path or in a class or function name.
"""
try:
with warnings.catch_warnings():
warnings.simplefilter("ignore") # test files carry stray escapes; their names still parse
tree: Final = ast.parse((REPO_ROOT / relative_path).read_text())
except (OSError, SyntaxError):
return frozenset({relative_path})
return frozenset({relative_path}) | frozenset(
node.name
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
)
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()
if any(_token_covers(glob, path) for slice_ in slices for glob in slice_.globs)
)
return tuple(
Finding(
subject=path,
detail="globbed by a job, then deselected by every one of their -k expressions",
)
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)
)
def _holds_tests(directory: pathlib.Path) -> bool:
return any(directory.rglob("test_*.py"))
def _shard_children(root: str, repo_root: pathlib.Path = REPO_ROOT) -> tuple[str, ...]:
"""Children of a sharded root that carry tests, so each one needs its own shard.
A directory earns an entry by containing a test file rather than by being named
`test_*`, which is what keeps fixture directories (`test_configs`, `expected_*`)
out without a hand-maintained list of exceptions.
"""
return tuple(
sorted(
child.relative_to(repo_root).as_posix()
for child in (repo_root / root).iterdir()
if not child.name.startswith(".")
and (
_holds_tests(child)
if child.is_dir()
else child.name.startswith("test_") and child.suffix == ".py"
)
)
)
def _unassigned_shard_children(
tokens: frozenset[str],
roots: tuple[str, ...] = SHARDED_ROOTS,
repo_root: pathlib.Path = REPO_ROOT,
) -> tuple[Finding, ...]:
return tuple(
Finding(subject=child, detail=f"holds tests but no shard of {root} names it")
for root in roots
if (repo_root / root).is_dir()
for child in _shard_children(root, repo_root)
if child not in roots and not any(_token_names(token, child) for token in tokens)
)
def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
return tuple(
Finding(subject=relative_path, detail="built by no job")
@ -174,6 +393,25 @@ def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tupl
)
def _stale_allowlist_paths(
allowlist: Allowlist,
*,
test_files: tuple[str, ...],
dockerfiles: tuple[str, ...],
) -> tuple[Finding, ...]:
sections: Final[tuple[Section, ...]] = (
Section("test_paths", allowlist.test_paths, test_files, _token_covers),
Section("dockerfiles", allowlist.dockerfiles, dockerfiles, operator.eq),
)
return tuple(
Finding(subject=path, detail=f"listed under '{section.name}' but matches no file the census looks at")
for section in sections
for entry in section.entries
for path in entry.paths
if not any(section.matches(path, candidate) for candidate in section.candidates)
)
def _parse_entry(item: object, section: str) -> AllowEntry:
if not isinstance(item, dict):
raise SystemExit(f"{ALLOWLIST_FILE.name}: '{section}' entries must be mappings")
@ -229,13 +467,56 @@ def _report(title: str, findings: tuple[Finding, ...], remedy: str) -> None:
_write("")
def _check_slices() -> int:
findings: Final = _deselected_everywhere(_load_allowlist())
if findings:
_report(
"test files a -k expression removes from every job that globs them",
findings,
"Give each one a job whose -k keeps it, or list it in "
".github/ci-coverage-allowlist.yml with the reason it may stay unrun.",
)
return 1
_write(f"OK: no test file is globbed by a job and then deselected by every -k across {len(_slices())} slices.")
return 0
def _check_shards() -> int:
findings = _unassigned_shard_children(_invoked_test_tokens(_all_scalars()))
if findings:
_report(
"test directories and files that no shard claims",
findings,
"Add each to the shard it belongs to. A directory that is itself split across "
"several shards belongs in SHARDED_ROOTS instead, so its own children get checked.",
)
return 1
counted = sum(len(_shard_children(root)) for root in SHARDED_ROOTS if (REPO_ROOT / root).is_dir())
_write(f"OK: all {counted} test children across {len(SHARDED_ROOTS)} sharded trees are assigned to a shard.")
return 0
def main() -> int:
if "--shards" in sys.argv[1:]:
return _check_shards()
if "--slices" in sys.argv[1:]:
return _check_slices()
allowlist = _load_allowlist()
scalars = _all_scalars()
test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars))
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
if stale_findings:
_report(
"allowlist entries that exempt nothing",
stale_findings,
"Delete each from .github/ci-coverage-allowlist.yml; the file it named is gone or was renamed.",
)
if test_findings:
_report(
"test files that no CI job invokes",
@ -248,7 +529,7 @@ def main() -> int:
dockerfile_findings,
"Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if test_findings or dockerfile_findings:
if stale_findings or test_findings or dockerfile_findings:
return 1
_write(

View file

@ -0,0 +1,149 @@
#!/usr/bin/env python3
"""Three invariants about what lives in .github/workflows/ and what its names mean.
`.github/workflows/` is a directory GitHub reads, not a place to keep things. Every
file at its top level is parsed as a workflow, so a script or a data file parked there
is either an invalid workflow or an orphan nobody can find. A subdirectory is not read
at all, so helper files may live in one. GitHub accepts both `.yml` and `.yaml`, and
this repo spells them `.yml`, which is a naming rule rather than a validity one and is
reported separately. And the `_` prefix is the repo's only signal that a workflow is a
reusable building block rather than something that runs on its own, which is worth
nothing unless it is true both ways.
WF001 a top-level file in .github/workflows/ that is not a workflow at all
WF002 a workflow whose only trigger is `workflow_call` but is not `_`-prefixed
WF003 a `_`-prefixed workflow that no other workflow can call
WF004 a real workflow spelled `.yaml` where this directory spells them `.yml`
A workflow with `workflow_call` alongside a human trigger is deliberately dual-mode
and belongs under its plain name, so only the call-only ones are held to WF002.
Usage
-----
python assert_workflow_dir_hygiene.py
Exit code 1 if any violation is found.
"""
from __future__ import annotations
import pathlib
import sys
from dataclasses import dataclass
from typing import Final
import yaml
REPO_ROOT: Final = pathlib.Path(__file__).resolve().parents[2]
WORKFLOW_DIR: Final = REPO_ROOT / ".github" / "workflows"
SCRIPT_HOME: Final = ".github/scripts/"
REUSABLE_PREFIX: Final = "_"
CALL_TRIGGER: Final = "workflow_call"
CANONICAL_SUFFIX: Final = ".yml"
WORKFLOW_SUFFIXES: Final = frozenset((CANONICAL_SUFFIX, ".yaml"))
@dataclass(frozen=True, slots=True)
class Finding:
subject: str
code: str
detail: str
def render(self) -> str:
return f" - {self.subject}: {self.code} {self.detail}"
def _triggers(document: object) -> frozenset[str]:
if not isinstance(document, dict):
return frozenset()
raw: Final = document.get("on", document.get(True))
if isinstance(raw, str):
return frozenset({raw})
if isinstance(raw, dict):
return frozenset(str(key) for key in raw)
if isinstance(raw, list):
return frozenset(str(item) for item in raw)
return frozenset()
def _workflows(directory: pathlib.Path) -> tuple[pathlib.Path, ...]:
return tuple(
path
for path in sorted(directory.iterdir())
if path.is_file() and path.suffix in WORKFLOW_SUFFIXES
)
def _strays(directory: pathlib.Path) -> tuple[Finding, ...]:
return tuple(
Finding(
path.name,
"WF001",
f"is not a workflow, and GitHub parses every top-level file here as one; "
f"move it to {SCRIPT_HOME} or into a subdirectory, which GitHub does not read",
)
for path in sorted(directory.iterdir())
if path.is_file() and path.suffix not in WORKFLOW_SUFFIXES
)
def _misspelled(directory: pathlib.Path) -> tuple[Finding, ...]:
return tuple(
Finding(
path.name,
"WF004",
f"is a real workflow and GitHub reads it, but this directory spells them "
f"{CANONICAL_SUFFIX}; rename it to {path.stem}{CANONICAL_SUFFIX}",
)
for path in _workflows(directory)
if path.suffix != CANONICAL_SUFFIX
)
def _misnamed(directory: pathlib.Path) -> tuple[Finding, ...]:
return tuple(
finding
for path in _workflows(directory)
for finding in _naming_findings(path, _triggers(yaml.safe_load(path.read_text(encoding="utf-8"))))
)
def _naming_findings(path: pathlib.Path, triggers: frozenset[str]) -> tuple[Finding, ...]:
underscored: Final = path.name.startswith(REUSABLE_PREFIX)
if triggers == frozenset({CALL_TRIGGER}) and not underscored:
return (
Finding(
path.name,
"WF002",
f"is only callable by another workflow, so name it {REUSABLE_PREFIX}{path.name}",
),
)
if underscored and CALL_TRIGGER not in triggers:
return (
Finding(
path.name,
"WF003",
f"is named as a reusable workflow but has no {CALL_TRIGGER} trigger; "
"add one or drop the prefix",
),
)
return ()
def main() -> int:
findings: Final = _strays(WORKFLOW_DIR) + _misspelled(WORKFLOW_DIR) + _misnamed(WORKFLOW_DIR)
if not findings:
total: Final = len(_workflows(WORKFLOW_DIR))
sys.stdout.write(
f"OK: {total} workflows, every file in .github/workflows/ is one, and the "
f"{REUSABLE_PREFIX} prefix means callable in both directions.\n"
)
return 0
sys.stdout.write("ERROR: .github/workflows/ holds files that break its own conventions\n")
for finding in findings:
sys.stdout.write(f"{finding.render()}\n")
return 1
if __name__ == "__main__":
sys.exit(main())

198
.github/scripts/e2e_egress_sentinel.py vendored Executable file
View 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
View 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
View 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"

View file

@ -1,557 +0,0 @@
#!/usr/bin/env python3
"""One-shot 7-day heads-up sweep for the Agent Shin rollout.
Posts a friendly "the OSS triage bot kicks in next Monday" comment on every
open external PR/issue that currently *would* fail the new rubric i.e.,
every PR/issue Agent Shin would close once the rollout completes. The point
is to give contributors a full week to fix their description before the bot
ever takes a destructive action, so nobody is surprised by an auto-close.
The script is designed to run **exactly once** at rollout, fired by a manual
``workflow_dispatch`` (``dry_run=false``) on the heads-up workflow. Re-runs
are safe: every comment is stamped with the hidden ``HEADS_UP_MARKER`` and
PRs/issues that already carry the marker are skipped.
Dry-run vs. real run
--------------------
Defaults to dry-run. Passing ``--close`` flips into real mode. Every GitHub
mutation goes through ``_agent_shin_actions``, which has a one-line
``if dry_run: log else: do_it`` per call, so the only difference between a
dry-run preview and the real run is the call site that actually hits the
GitHub API.
Local preview::
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm
Real run (the manual rollout dispatch uses this)::
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm --close
"""
from __future__ import annotations
import argparse
import datetime as dt
import json
import os
import sys
from pathlib import Path
from typing import Any
# Make the sibling triage_with_llm + _agent_shin_actions importable when this
# script is invoked directly (the GitHub workflow does `python3 .github/scripts/...`).
_SCRIPTS_DIR = Path(__file__).resolve().parent
if str(_SCRIPTS_DIR) not in sys.path:
sys.path.insert(0, str(_SCRIPTS_DIR))
from _agent_shin_actions import maybe_post_comment # noqa: E402
from agent_shin_shared import ( # noqa: E402
AGENT_SHIN_DEFAULT_BOT_LOGIN,
ALLOWLIST_LOGINS,
list_open_items,
)
from triage_with_llm import ( # noqa: E402
DEFAULT_MODEL,
call_llm_judge,
fetch_issue,
fetch_pr,
gh,
is_internal_contributor,
review_gate,
triage,
)
# Hidden marker so re-runs skip PRs/issues we've already notified. Distinct from
# the within-grace / ready / regressed markers so it can't be confused with the
# steady-state lifecycle comments.
HEADS_UP_MARKER = "<!-- agent-shin:rollout-heads-up -->"
# Placeholder until the litellm-docs PR ships. The rollout blog post explains
# the new rubric, the 7-day grace, and how to recover after an auto-close.
# TODO(docs): replace with the canonical URL once the litellm-docs PR merges.
ROLLOUT_BLOG_URL = "https://docs.litellm.ai/docs/agent_shin_triage_rollout"
# Default cutoff is one week from "now". Computed at runtime so the wording
# stays correct even if the rollout is merged later than planned. The user can
# override with --close-on YYYY-MM-DD when running the script manually.
DEFAULT_GRACE_DAYS = 7
# The daily auto-close sweeps (close_low_quality_prs.yml at 09:00 UTC and
# review_gate.yml at 09:30 UTC) are what actually close a still-failing item,
# so the deadline we promise contributors has to name that wall-clock moment.
ACTIVATION_TIME_UTC = "09:00 UTC"
def _format_cutoff(cutoff: dt.date) -> str:
"""Human-readable, timezone-explicit cutoff, e.g. ``Monday, June 1, 2026
(09:00 UTC)`` the moment a still-failing PR/issue gets closed."""
return (
f"{cutoff.strftime('%A, %B')} {cutoff.day}, {cutoff.year} "
f"({ACTIVATION_TIME_UTC})"
)
def _rubric_section_pr() -> str:
return (
"**Going forward, every external PR needs ONE of:**\n"
"\n"
"- A linked GitHub issue using a closing keyword: "
"`Fixes #1234`, `Closes #1234`, or `Resolves #1234`, OR\n"
"- All three of: a clear **problem description**, **expected vs. "
"actual behavior**, and **end-to-end QA proof** (at least one of a "
"short screen recording / video, before/after screenshots, or the "
"exact commands you ran with their real output; mocked or stubbed "
"runs don't count).\n"
"\n"
"PRs also need a **Greptile confidence score of 4/5 or higher** before "
"the bot will tag them `ready for review`. You can `@greptileai` to "
"request a fresh review at any time, including after the PR is closed."
)
def _rubric_section_issue() -> str:
return (
"**Going forward, every external issue needs:**\n"
"\n"
"- For **bug reports**: end-to-end evidence of the bug (at least one "
"of a screen recording / video, a screenshot, or the exact commands "
"you ran with their real output / traceback) plus expected vs. actual "
"behavior. Written steps with no run output don't count, and mocked "
"or stubbed runs don't count.\n"
"- For **feature requests**: a clear description of the proposed "
"feature plus a use case + concrete example (config, API call, UI "
"flow, or scenario showing what's blocked today)."
)
def _description_only_note(kind: str) -> str:
noun = "PR" if kind == "pr" else "issue"
return (
f"⚠️ **The requirements must live in the {noun} *description*, not in "
"comments.** Some PRs/issues collect 100+ comments from humans and "
"bots; reading the entire thread on every triage run would balloon "
"GitHub API usage (we'd start getting 429'd) and blow out the LLM "
"judge's context. The bot only reads the description, so anything "
"you add as a comment will be invisible to it."
)
def _missing_section(verdict: dict, greptile_score: int | None) -> str:
"""Bullet list of what's currently missing on this PR/issue.
Combines the LLM judge's `missing` list (rubric items) with a Greptile
shortfall (for PRs) so the contributor sees one list of things to fix.
"""
missing = list(verdict.get("missing") or [])
if greptile_score is not None and greptile_score < 4:
missing.insert(
0,
f"Greptile's most recent review scored this PR {greptile_score}/5 "
"(below the 4/5 bar Agent Shin will require).",
)
if not missing:
return (
"_The bot couldn't articulate a specific missing piece; see the "
"rubric link above and double-check the description includes all "
"of it before the rollout._"
)
bullets = "\n".join(f"- {m}" for m in missing)
return f"**What this one is currently missing:**\n\n{bullets}"
def _recovery_section(kind: str) -> str:
if kind == "pr":
return (
"**If the bot closes this PR after the rollout:** update the "
"description with the missing pieces, then either open a fresh "
"PR or comment `@agent-shin reconsider` on the closed PR. If "
"Greptile re-scores you at 4/5 or higher I'll reopen and tag "
"the PR `ready for review`. (`@greptileai` works on closed PRs "
"too; a fresh review is one of the signals that lifts you back "
"into the queue.) This is **not** us losing interest in your "
"change; far from it. We just need open PRs to be a list of "
"things a maintainer can act on, so we can get to yours faster."
)
return (
"**If the bot closes this issue after the rollout:** edit the issue "
"description to add the missing pieces, then comment `@agent-shin "
"reconsider` on the closed issue. I'll re-evaluate and, if the rubric "
"is met, reopen it. (GitHub doesn't let external authors reopen an "
"issue a maintainer or bot closed, so the comment is the reliable "
"path.) This is **not** us saying the bug isn't real or the request "
"isn't useful; it's so the remaining open issues are a list of things "
"a maintainer can act on."
)
def format_heads_up_comment(
*, kind: str, verdict: dict, greptile_score: int | None, cutoff: dt.date
) -> str:
"""Compose the friendly 7-day heads-up comment posted on a failing PR/issue."""
noun = "PR" if kind == "pr" else "issue"
rubric = _rubric_section_pr() if kind == "pr" else _rubric_section_issue()
cutoff_str = _format_cutoff(cutoff)
explanation = (verdict.get("explanation") or "").strip()
explanation_block = (
f"> _(The judge's note for this one: {explanation})_\n\n" if explanation else ""
)
return (
"🚅 **Heads-up: we're turning on the OSS triage bot in "
f"{DEFAULT_GRACE_DAYS} days, on {cutoff_str}.**\n"
"\n"
"We're rolling out **Agent Shin**, an LLM-as-judge triage bot for "
f"external {noun}s. Once it's live, the bot reads each open "
f"{noun}'s description, scores it against a small rubric, and "
f"auto-closes any {noun} that's missing the basics, with a single "
f"comment explaining what's missing and how to recover. Full "
f"context: [Agent Shin rollout blog post]({ROLLOUT_BLOG_URL}).\n"
"\n"
f"{rubric}\n"
"\n"
f"{_description_only_note(kind)}\n"
"\n"
f"{_missing_section(verdict, greptile_score)}\n"
"\n"
f"{explanation_block}"
"**Timeline (you have a week):**\n"
"\n"
f"- We turn the bot on in {DEFAULT_GRACE_DAYS} days, on "
f"**{cutoff_str}**. You have until then to update this {noun}'s "
"description with the missing pieces above.\n"
f"- If this {noun} still fails the rubric at **{cutoff_str}**, "
"we'll close it.\n"
f"- From then on the bot runs daily, and every {noun} that fails "
"the rubric gets a **2-hour lifetime**: one warning comment, then "
"auto-close 2 hours later.\n"
"\n"
f"{_recovery_section(kind)}\n"
"\n"
f"{HEADS_UP_MARKER}"
)
def _list_open_numbers(repo: str, kind: str) -> list[int]:
"""Return every open PR or issue number in ``repo``.
Delegates to ``list_open_items`` so the full backlog is fetched (no cap)
and the `gh {pr,issue} list` invocation stays in one shared place. ``gh
issue list`` would include PRs, but ``list_open_items`` uses the dedicated
command per kind, so the two never mix.
"""
return [
item["number"] for item in list_open_items(kind, repo=repo, fields="number")
]
def _has_heads_up_marker(item: dict) -> bool:
"""Cheap fast-path: check the PR/issue body itself for the marker.
The marker is appended to the *comment* we post, not the body, so this
will only fire if the body literally contains the marker text. We still
do the comment-marker check separately below; this body check just lets
us short-circuit for PRs/issues that quote the marker for any reason.
"""
body = item.get("body") or ""
return HEADS_UP_MARKER in body
def _comments_have_marker(repo: str, number: int) -> bool:
"""True if the bot already posted a comment carrying the marker.
Used for idempotency: a re-run skips items the previous run notified.
Filters by author (matching the sibling marker-checks in
``triage_with_llm._has_marker`` and
``agent_shin_shared.seconds_since_latest_marker_comment``) so a
contributor who quotes the heads-up via GitHub's "Quote reply" — which
preserves HTML comments in the raw markdown can't trick the
idempotency check into silently skipping a real heads-up.
Comments live on the unified issues endpoint regardless of whether the
item is a PR or an issue, so no ``kind`` argument is required here.
"""
expected_login = (
os.environ.get("AGENT_SHIN_BOT_LOGIN") or AGENT_SHIN_DEFAULT_BOT_LOGIN
).lower()
raw = gh(
"api",
"--paginate",
f"repos/{repo}/issues/{number}/comments?per_page=100",
)
for line in raw.splitlines():
line = line.strip()
if not line:
continue
try:
payload = json.loads(line)
except json.JSONDecodeError:
continue
comments = payload if isinstance(payload, list) else [payload]
for comment in comments:
author = ((comment.get("user") or {}).get("login") or "").lower()
if author != expected_login:
continue
if HEADS_UP_MARKER in (comment.get("body") or ""):
return True
return False
def _evaluate_pr(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
"""Run the future PR rubric (review_gate) in dry-run and return the result."""
return review_gate(
repo=repo,
number=number,
close=False, # we only want the verdict, never act here
model=model,
judge=judge,
)
def _evaluate_issue(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
"""Run the future issue rubric (triage kind='issue') in dry-run."""
return triage(
repo=repo,
kind="issue",
number=number,
close=False,
model=model,
judge=judge,
)
def _would_be_closed(kind: str, result: dict) -> bool:
"""True if the future triage would auto-close this PR/issue based on the
rubric (regardless of grace-period gating).
For PRs we trust ``review_gate``'s ``passing`` field — it combines the LLM
verdict and the Greptile score. For issues we read the LLM verdict
directly. Both fields are ``None``/missing on skip paths
(skip-internal-author, skip-llm-error, etc.) where the future bot would
NOT close the item those return False.
"""
if kind == "pr":
passing = result.get("passing")
if passing is None:
return False # skipped — nothing for the heads-up to warn about
return passing is False
verdict = result.get("verdict") or {}
return (verdict.get("verdict") or "").lower() == "fail"
def _process_one(
*,
repo: str,
kind: str,
number: int,
model: str,
cutoff: dt.date,
dry_run: bool,
judge: Any = None,
skip_marker_check: bool = False,
allowlist: frozenset[str] = ALLOWLIST_LOGINS,
) -> dict:
"""Evaluate one PR/issue and post a heads-up if it would be auto-closed.
Returns a per-item dict for the summary table.
"""
base = {"kind": kind, "number": number}
fetcher = fetch_pr if kind == "pr" else fetch_issue
item = fetcher(repo, number)
if (item.get("state") or "") != "open":
return {**base, "action": "skip-not-open"}
if allowlist:
login = (item.get("user") or {}).get("login") or ""
if login.lower() not in allowlist:
return {**base, "action": "skip-not-allowlisted"}
elif is_internal_contributor(item):
return {**base, "action": "skip-internal-author"}
if not skip_marker_check and _has_heads_up_marker(item):
return {**base, "action": "skip-already-marked-in-body"}
if not skip_marker_check and _comments_have_marker(repo, number):
return {**base, "action": "skip-already-notified"}
if kind == "pr":
result = _evaluate_pr(repo=repo, number=number, model=model, judge=judge)
else:
result = _evaluate_issue(repo=repo, number=number, model=model, judge=judge)
if not _would_be_closed(kind, result):
return {**base, "action": "skip-passing", "evaluator": result.get("action")}
verdict = result.get("verdict") or {}
greptile_score = result.get("greptile_score") if kind == "pr" else None
comment = format_heads_up_comment(
kind=kind, verdict=verdict, greptile_score=greptile_score, cutoff=cutoff
)
maybe_post_comment(repo, number, comment, dry_run=dry_run)
return {
**base,
"action": "heads-up-posted" if not dry_run else "would-post-heads-up",
"verdict": (verdict.get("verdict") or "").lower(),
"greptile_score": greptile_score,
}
def _print_summary(results: list[dict]) -> None:
"""Tally per-action counts so a dry-run preview tells you at a glance how
many comments the real run would post."""
counts: dict[str, int] = {}
for r in results:
counts[r["action"]] = counts.get(r["action"], 0) + 1
print("\n=== rollout heads-up summary ===")
for action in sorted(counts):
print(f" {action:35s} {counts[action]}")
print(f" total {len(results)}")
def run(
*,
repo: str,
close: bool,
cutoff: dt.date,
model: str,
kinds: tuple[str, ...] = ("pr", "issue"),
judge: Any = None,
only_numbers: dict[str, list[int]] | None = None,
skip_marker_check: bool = False,
) -> list[dict]:
"""Sweep ``repo`` and post heads-up comments. Returns the per-item results."""
dry_run = not close
if dry_run:
print(
f"[DRY RUN] sweeping {repo}; --close not passed, no comments will be posted."
)
else:
print(f"[REAL RUN] sweeping {repo}; comments WILL be posted.")
print(f"Cutoff date in comment body: {cutoff.isoformat()}")
results: list[dict] = []
for kind in kinds:
if only_numbers and kind in only_numbers:
numbers = list(only_numbers[kind])
else:
numbers = _list_open_numbers(repo, kind)
print(f"\n--- {kind}s: {len(numbers)} open ---")
for n in numbers:
try:
result = _process_one(
repo=repo,
kind=kind,
number=n,
model=model,
cutoff=cutoff,
dry_run=dry_run,
judge=judge,
skip_marker_check=skip_marker_check,
)
except (
Exception
) as exc: # noqa: BLE001 - per-item errors don't abort the sweep
result = {
"kind": kind,
"number": n,
"action": "error",
"error": str(exc),
}
print(f"!! {kind}#{n}: {exc}", file=sys.stderr)
print(f" {kind}#{n}: {result['action']}")
results.append(result)
_print_summary(results)
return results
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--repo", required=True, help="owner/repo")
parser.add_argument(
"--close",
action="store_true",
help=(
"Actually post comments. Without this flag the script is in "
"dry-run mode and only logs what it would do."
),
)
parser.add_argument(
"--close-on",
type=dt.date.fromisoformat,
default=None,
help=(
"Cutoff date shown in the heads-up comment as the rollout date "
f"(default: today + {DEFAULT_GRACE_DAYS} days)."
),
)
parser.add_argument(
"--model",
default=os.environ.get("TRIAGE_MODEL") or DEFAULT_MODEL,
help=f"Model for the rubric LLM judge (default: {DEFAULT_MODEL}).",
)
parser.add_argument(
"--kind",
choices=("pr", "issue", "both"),
default="both",
help="Restrict the sweep to PRs or issues only (default: both).",
)
parser.add_argument(
"--only-pr",
type=int,
action="append",
default=[],
help="Limit the PR sweep to these PR numbers (repeat for several).",
)
parser.add_argument(
"--only-issue",
type=int,
action="append",
default=[],
help="Limit the issue sweep to these issue numbers (repeat for several).",
)
parser.add_argument(
"--ignore-existing-marker",
action="store_true",
help=(
"Re-post on PRs/issues that already carry the heads-up marker. "
"Useful for testing the comment wording on a known PR."
),
)
args = parser.parse_args()
cutoff = args.close_on or (
dt.datetime.now(dt.timezone.utc).date() + dt.timedelta(days=DEFAULT_GRACE_DAYS)
)
kinds: tuple[str, ...]
if args.kind == "pr":
kinds = ("pr",)
elif args.kind == "issue":
kinds = ("issue",)
else:
kinds = ("pr", "issue")
only: dict[str, list[int]] = {}
if args.only_pr:
only["pr"] = args.only_pr
if args.only_issue:
only["issue"] = args.only_issue
# The script must NOT hit the LLM in dry-run if no key is set — we still
# want a useful preview that says "skip-no-llm-key" for items that would
# have been judged. Production runs require OPENAI_API_KEY.
if args.close and not os.environ.get("OPENAI_API_KEY"):
parser.error("OPENAI_API_KEY must be set for --close (real-run) mode.")
run(
repo=args.repo,
close=args.close,
cutoff=cutoff,
model=args.model,
kinds=kinds,
only_numbers=only or None,
skip_marker_check=args.ignore_existing_marker,
)
return 0
if __name__ == "__main__":
sys.exit(main())

View file

@ -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
@ -129,6 +134,13 @@ jobs:
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DIST: ${{ inputs.dist }}
# coverage.py's sys.monitoring backend (PEP 669), the cheapest core it has.
# It is only the default from Python 3.14, and these shards run 3.12, so it
# has to be asked for. Coverage refuses it when branch measurement is on
# (`branch_right_left` needs > 3.14.0a5) and falls back to the slow core with
# a `no-sysmon` warning, so turning on `branch = true` here means giving this
# back until the runners move to 3.14.
COVERAGE_CORE: sysmon
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \
@ -137,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
@ -149,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

View file

@ -23,7 +23,7 @@ jobs:
version: "0.10.9"
- name: Update JSON Data
run: |
uv run --frozen --with 'aiohttp==3.13.3' python ".github/workflows/auto_update_price_and_context_window_file.py"
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
- name: Regenerate JSON Schema
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py

View file

@ -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
@ -79,6 +83,24 @@ jobs:
if: steps.changes.outputs.relevant == 'true'
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Regenerate the lazy OpenAPI snapshot
if: steps.changes.outputs.relevant == 'true'
run: uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
- name: Fail if the lazy OpenAPI snapshot is stale
if: steps.changes.outputs.relevant == 'true'
run: |
if ! git diff --exit-code -- litellm/proxy/_lazy_openapi_snapshot.json; then
echo "::error file=litellm/proxy/_lazy_openapi_snapshot.json::The lazy OpenAPI snapshot is out of sync with the lazily loaded routes."
echo ""
echo "A lazily loaded route or model changed without regenerating the snapshot that /openapi.json serves for unloaded features."
echo "To fix, run from the repo root:"
echo " uv run python -m litellm.proxy._lazy_openapi_snapshot"
echo "then run npm run gen:api from ui/litellm-dashboard and commit both files."
exit 1
fi
echo "_lazy_openapi_snapshot.json is in sync with the lazily loaded routes."
- name: Set up Node.js
if: steps.changes.outputs.relevant == 'true'
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0

View file

@ -40,3 +40,12 @@ jobs:
run: |
python -m pip install "pyyaml==6.0.3"
python .github/scripts/assert_ci_coverage.py
# The census asks whether a job names a file; this asks whether that job's -k
# then throws it back out. A file both globbed and deselected everywhere runs
# nowhere while counting as covered, which is how the caching suite went unrun.
- name: Assert no -k expression deselects a file from every job that globs it
run: python .github/scripts/assert_ci_coverage.py --slices
- name: Assert .github/workflows/ holds only workflows, correctly named
run: python .github/scripts/assert_workflow_dir_hygiene.py

237
.github/workflows/e2e_record_replay.yml vendored Normal file
View 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

View file

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

View file

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

View file

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

View file

@ -0,0 +1,68 @@
name: Sync Together AI model registry
on:
schedule:
- cron: "30 6 * * *"
workflow_dispatch:
permissions:
contents: write
pull-requests: write
jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: litellm_internal_staging
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Look for an already-open sync PR
id: existing
run: |
open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \
--jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')"
echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT"
if [ -n "$open_pr" ]; then
echo "An open sync PR already exists on branch $open_pr; skipping this run."
fi
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
- name: Run the sync
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Regenerate the JSON schema
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
if git diff --quiet; then
echo "Registry already in sync; no PR needed."
exit 0
fi
branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')"
git config user.name "github-actions[bot]"
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json \
model_prices_and_context_window.schema.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base litellm_internal_staging
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

View file

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

View file

@ -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: |
@ -122,6 +137,11 @@ jobs:
uv run --no-sync ruff check .
cd ..
- name: Run Ruff linting (test tree)
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync ruff check --config ruff-tests.toml tests
- name: Check strict-rule budget (delta vs base)
if: steps.changes.outputs.decision != 'skip'
run: |
@ -132,6 +152,11 @@ jobs:
run: |
uv run --no-sync python scripts/type_discipline_gate.py --base "$GATE_BASE_SHA"
- name: Check test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes, litellm global mutation, credential-gated skips, conftest snapshot inventory, delta vs base)
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync python scripts/test_quality_gate.py --base "$GATE_BASE_SHA"
- name: Print OpenAI version
if: steps.changes.outputs.decision != 'skip'
run: |
@ -223,7 +248,7 @@ jobs:
- name: Run secret scan test
run: |
uv run --no-project --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
uv run --no-project --with 'pytest==9.0.2' pytest tests/code_coverage_tests/test_no_hardcoded_secrets.py -v
- name: Run ggshield secret scan
env:

View file

@ -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
View 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

View file

@ -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
@ -111,4 +114,4 @@ jobs:
- name: Audit provider endpoints against the schema
working-directory: terraform/provider
run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json"
run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json" -coverage-allowlist ./tools/endpointaudit/coverage_allowlist.txt

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Core Utilities"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
core-utils:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/litellm_core_utils"
workers: 2
reruns: 1
artifact-name: core-utils

View file

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

View file

@ -1,35 +0,0 @@
name: "Unit Tests: Enterprise, Google GenAI & Routing"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
enterprise-routing:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/enterprise
tests/test_litellm/google_genai
tests/test_litellm/router_utils
tests/test_litellm/router_strategy
workers: 2
reruns: 2
artifact-name: enterprise-routing

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Integrations (Callbacks & Logging)"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
integrations:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/integrations"
workers: 2
reruns: 3
artifact-name: integrations

View file

@ -1,47 +0,0 @@
name: "Unit Tests: LLM Provider Transformations"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
vertex-ai:
name: Vertex AI
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/llms/vertex_ai"
workers: 1
reruns: 2
artifact-name: llm-vertex-ai
other-providers:
name: All Other Providers
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai"
workers: 2
reruns: 2
artifact-name: llm-other-providers

View file

@ -1,53 +0,0 @@
name: "Unit Tests: MCP, Secrets, Containers & Misc"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
misc:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
tests/test_litellm/experimental_mcp_client
tests/test_litellm/models
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
tests/test_litellm/realtime_api
tests/test_litellm/rerank_api
tests/test_litellm/sandbox
tests/test_litellm/test_router
tests/test_litellm/vector_stores
tests/test_litellm/videos
tests/test_litellm/test_*.py
workers: 2
reruns: 2
artifact-name: misc

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Proxy Auth & Key Management"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-auth:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/proxy/auth tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine tests/test_litellm/proxy/client"
workers: 2
reruns: 2
artifact-name: proxy-auth

View file

@ -42,11 +42,10 @@ concurrency:
# pinning the whole file to one worker (the default --dist=loadscope
# behavior for single-file targets).
jobs:
# Fast guard — fails the workflow if a test_*.py file under
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
# The semantic-shard design (no catch-all "remaining" bucket) relies on
# every test file being explicitly assigned; this guard prevents a new
# file from silently dropping out of CI.
# Fast guard — fails the workflow when a test directory or file inside a sharded
# tree is claimed by no shard. The semantic-shard design has no catch-all bucket,
# so an unassigned child runs nowhere; assert_ci_coverage.py holds the tree list
# and reads the same test-path keys the coverage census does.
assert-shard-coverage:
runs-on: ubuntu-latest
timeout-minutes: 2
@ -56,31 +55,8 @@ jobs:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Assert every test_*.py is in a matrix shard
run: |
python3 - <<'PY'
import pathlib, sys, yaml
wf = yaml.safe_load(open(".github/workflows/test-unit-proxy-db.yml"))
matrix = wf["jobs"]["proxy-db"]["strategy"]["matrix"]["include"]
referenced = set()
for entry in matrix:
for token in entry["test-path"].split():
if token.startswith("tests/proxy_unit_tests/"):
referenced.add(pathlib.PurePosixPath(token).name)
actual = {p.name for p in pathlib.Path("tests/proxy_unit_tests").iterdir()
if p.name.startswith("test_") and (p.suffix == ".py" or p.is_dir())
and p.name != "test_configs"}
orphans = sorted(actual - referenced)
if orphans:
print("ERROR: the following files/dirs under tests/proxy_unit_tests/")
print(" are not assigned to any shard in test-unit-proxy-db.yml:")
for o in orphans:
print(f" - {o}")
print()
print("Add each to whichever semantic shard it belongs to.")
sys.exit(1)
print(f"OK: all {len(actual)} files assigned to a shard.")
PY
- name: Assert every test directory and file is claimed by a shard
run: python3 .github/scripts/assert_ci_coverage.py --shards
proxy-db:
needs: assert-shard-coverage

View file

@ -1,81 +0,0 @@
name: "Unit Tests: Proxy API Endpoints"
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:
proxy-endpoints:
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers
tests/test_litellm/proxy/anthropic_endpoints
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/ocr_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/credential_endpoints
tests/test_litellm/proxy/discovery_endpoints
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/shutdown
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
tests/test_litellm/proxy/utils
workers: 2
reruns: 2
artifact-name: proxy-endpoints
# Behavior-pinning tests for litellm/proxy/proxy_server.py. Owns its
# own job (not a path on the proxy-endpoints job above) so its budget
# is independent and its coverage artifact is uploaded separately.
# See: https://www.notion.so/36c43b8acdab81ee845fd5365128a2fc
proxy-server:
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: tests/test_litellm/proxy/proxy_server
workers: 4
reruns: 2
timeout-minutes: 60
job-timeout-minutes: 95
artifact-name: proxy-server

View file

@ -1,42 +0,0 @@
name: "Unit Tests: Proxy Infrastructure"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-infra:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/proxy/db
tests/test_litellm/proxy/middleware
tests/test_litellm/proxy/spend_tracking
tests/test_litellm/proxy/pass_through_endpoints
tests/test_litellm/proxy/_experimental
tests/test_litellm/proxy/experimental
tests/test_litellm/proxy/common_utils
tests/test_litellm/proxy/enterprise_billing
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
workers: 2
reruns: 2
artifact-name: proxy-infra

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Responses, Caching & Types"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
responses-caching-types:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/responses tests/test_litellm/caching tests/test_litellm/types"
workers: 2
reruns: 2
artifact-name: responses-caching-types

251
.github/workflows/test-unit.yml vendored Normal file
View file

@ -0,0 +1,251 @@
name: "Unit 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' }}
# One caller for every tests/test_litellm shard, replacing the nine thin workflow
# files that each wrapped a single call to _test-unit-base.yml. Adding a shard is
# now one matrix entry rather than a new file.
#
# `name` is the shard id and nothing else, so each check reports as
# "<shard> / Run tests" exactly as it did when the shard had its own file. Those
# strings are the branch ruleset's required contexts, so they are load-bearing:
# renaming an entry renames a required check and the ruleset stops matching it.
#
# Every entry states its timeouts even when they equal the base workflow's
# defaults. An absent matrix key renders as an empty string, which is not a
# number, so a partially-specified entry would fail the call rather than fall
# back to the default.
#
# tests/proxy_unit_tests keeps its own caller (test-unit-proxy-db.yml): it is
# already a matrix and carries a shard-coverage guard that reads that file by
# name. Folding it in here is a follow-up, together with generalising that guard
# into assert_ci_coverage.py.
jobs:
unit:
name: ${{ matrix.shard }}
permissions:
contents: read
id-token: write
pull-requests: write
strategy:
fail-fast: false
matrix:
include:
- shard: core-utils
artifact-name: core-utils
test-path: "tests/test_litellm/litellm_core_utils"
workers: 2
reruns: 1
timeout-minutes: 20
job-timeout-minutes: 60
- shard: enterprise-routing
artifact-name: enterprise-routing
test-path: >-
tests/test_litellm/enterprise
tests/test_litellm/google_genai
tests/test_litellm/router_utils
tests/test_litellm/router_strategy
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: integrations
artifact-name: integrations
test-path: "tests/test_litellm/integrations"
workers: 2
reruns: 3
timeout-minutes: 20
job-timeout-minutes: 60
- shard: Vertex AI
artifact-name: llm-vertex-ai
test-path: "tests/test_litellm/llms/vertex_ai"
workers: 1
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: All Other Providers
artifact-name: llm-other-providers
test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai"
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: misc
artifact-name: misc
test-path: >-
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
tests/test_litellm/experimental_mcp_client
tests/test_litellm/models
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/ocr
tests/test_litellm/passthrough
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
tests/test_litellm/videos
tests/test_litellm/test_*.py
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: proxy-auth
artifact-name: proxy-auth
test-path: >-
tests/test_litellm/proxy/auth
tests/test_litellm/proxy/hooks
tests/test_litellm/proxy/policy_engine
tests/test_litellm/proxy/client
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: proxy-endpoints
artifact-name: proxy-endpoints
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/list_api
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers
tests/test_litellm/proxy/anthropic_endpoints
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/ocr_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/credential_endpoints
tests/test_litellm/proxy/discovery_endpoints
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/shutdown
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: 4
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: proxy-server
artifact-name: proxy-server
test-path: "tests/test_litellm/proxy/proxy_server"
workers: 4
reruns: 2
timeout-minutes: 60
job-timeout-minutes: 100
- shard: proxy-infra
artifact-name: proxy-infra
test-path: >-
tests/test_litellm/proxy/db
tests/test_litellm/proxy/middleware
tests/test_litellm/proxy/spend_tracking
tests/test_litellm/proxy/pass_through_endpoints
tests/test_litellm/proxy/_experimental
tests/test_litellm/proxy/experimental
tests/test_litellm/proxy/common_utils
tests/test_litellm/proxy/enterprise_billing
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: 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
test-path: >-
tests/test_litellm/responses
tests/test_litellm/caching
tests/test_litellm/types
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: ${{ matrix.test-path }}
workers: ${{ matrix.workers }}
reruns: ${{ matrix.reruns }}
timeout-minutes: ${{ matrix.timeout-minutes }}
job-timeout-minutes: ${{ matrix.job-timeout-minutes }}
artifact-name: ${{ matrix.artifact-name }}

View file

@ -1,92 +0,0 @@
name: Agent Shin — rollout heads-up (one-shot)
# Fires the 7-day heads-up comment on every open external PR/issue that the
# new triage bot would auto-close. The real sweep is a deliberate one-shot:
# trigger it at rollout via a manual `workflow_dispatch` with `dry_run=false`.
# The script is idempotent (skips items that already carry the
# `<!-- agent-shin:rollout-heads-up -->` marker), so a re-run is harmless.
#
# The automatic push trigger runs DRY-RUN only, so merging the script to
# `litellm_internal_staging` never posts a comment; it just confirms the
# workflow is wired up. Posting real comments requires the manual dispatch,
# which is also the only trigger that exposes `OPENAI_API_KEY`. The heads-up
# is intentionally NOT gated on `AGENT_SHIN_ENABLED`: it has to warn
# contributors while that flag is still off, ahead of the flip that turns on
# auto-closing.
#
# The workflow is a thin shell over `.github/scripts/triage_rollout_heads_up.py`.
# Dry-run vs. real run differ in EXACTLY one CLI flag (`--close`), added only
# on a manual dispatch with `dry_run=false`.
on:
push:
branches:
- litellm_internal_staging
paths:
# The presence of this script on staging IS the rollout merge marker.
# Editing the file later would re-fire the workflow; that's safe because
# the script skips PRs/issues that already have the heads-up marker.
- ".github/scripts/triage_rollout_heads_up.py"
workflow_dispatch:
inputs:
dry_run:
description: "Dry run (true = preview only, false = actually post comments)."
required: false
default: "true"
type: choice
options:
- "true"
- "false"
permissions:
contents: read
issues: write
pull-requests: write
jobs:
heads-up:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- name: Checkout triage scripts
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: .github/scripts
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Install LLM client
run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt
- name: Run heads-up sweep
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Only the manual dispatch (the real-run trigger) needs the LLM key.
# The automatic push trigger runs dry-run and never posts, so it gets
# no key. Mirrors the sibling triage workflows, which expose the key
# only on an enabled/dispatched run rather than unconditionally.
OPENAI_API_KEY: ${{ github.event_name == 'workflow_dispatch' && secrets.OPENAI_API_KEY || '' }}
OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }}
TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }}
# The real run is a deliberate manual dispatch with dry_run=false.
# Use the EXACT "false" comparison so any unexpected input value
# fail-closes to dry-run (mirrors the AGENT_SHIN_ENABLED pattern in
# the sibling workflows). The automatic push trigger always stays
# dry-run, so merging the script never posts.
DRY_RUN_INPUT: ${{ github.event.inputs.dry_run }}
run: |
set -euo pipefail
ARGS=(--repo "${{ github.repository }}")
if [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${DRY_RUN_INPUT:-true}" = "false" ]; then
ARGS+=(--close)
echo "::notice::Manual rollout dispatch with dry_run=false -> heads-up comments WILL be posted."
elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ]; then
echo "::notice::Manual dispatch in dry-run mode -> previewing only, no comments will be posted."
else
echo "::notice::Automatic push trigger -> dry-run preview only. Fire the real rollout sweep with a manual workflow_dispatch (dry_run=false)."
fi
python3 .github/scripts/triage_rollout_heads_up.py "${ARGS[@]}"

View file

@ -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
View file

@ -3,6 +3,8 @@
tests/e2e/.fixtures/
.venv-typecheck
.venv_policy_test
.venv-mutmut
mutants/
.env
.claude
CLAUDE.local.md

View file

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

View file

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

View file

@ -7,6 +7,7 @@
info lint lint-inner lint-dev lint-checks format \
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
lint-test-quality lint-test-quality-budget-update \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \
lint-install lint-fetch-base bootstrap
@ -35,7 +36,8 @@ help:
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
@echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + basedpyright)"
@echo " make lint-test-quality - Gate the test suite against test-quality-budget.json"
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)"
@echo " make check-circular-imports - Check for circular imports"
@echo " make check-import-safety - Check import safety"
@echo " make test - Run all tests"
@ -142,11 +144,13 @@ lint-install:
$(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
# Diff-scoped format check, identical to test-linting.yml's "Check ruff format" step:
# Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step:
# only the litellm Python files changed vs the base are checked, so a pre-existing
# format issue elsewhere doesn't block an unrelated commit.
# format issue elsewhere doesn't block an unrelated commit. Git pathspecs match
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
@files=$$(git diff --name-only origin/litellm_internal_staging...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' || true); \
@files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \
if [ -z "$$files" ]; then \
echo "No changed litellm Python files to format-check."; \
else \
@ -156,6 +160,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
# Linting targets
lint-ruff: $(LINT_DEP_INSTALL)
cd litellm && $(UV_RUN) ruff check . && cd ..
$(UV_RUN) ruff check --config ruff-tests.toml tests
# faster linter for developing ...
# inspiration from:
@ -200,6 +205,12 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging
# Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes,
# litellm module-global mutation, credential-gated skips, conftest snapshot
# inventory), counted across tests/ the same delta-vs-base way.
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging
# --update lowers each limit by what this branch fixed since its branch point, so
# it needs the base ref fetched to resolve the merge-base.
lint-basedpyright-budget-update: install-dev lint-fetch-base
@ -221,8 +232,11 @@ lint-ruff-budget-update: install-dev lint-fetch-base
lint-type-discipline-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_discipline_gate.py --update
# Ratchet all budgets in one shot (ruff strict + type-discipline + basedpyright)
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-basedpyright-budget-update
lint-test-quality-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/test_quality_gate.py --update
# Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update
check-circular-imports: $(LINT_DEP_INSTALL)
cd litellm && $(UV_RUN) python ../tests/documentation_tests/test_circular_imports.py && cd ..
@ -244,7 +258,7 @@ lint:
lint-inner: lint-install lint-fetch-base
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed check-circular-imports check-import-safety
@ -314,7 +328,7 @@ test-unit-helm: install-helm-unittest
# LLM Translation testing targets
test-llm-translation: install-test-deps
@echo "Running LLM translation tests..."
@python .github/workflows/run_llm_translation_tests.py
@python .github/scripts/run_llm_translation_tests.py
test-llm-translation-single: install-test-deps
@echo "Running single LLM translation test file..."

View file

@ -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) | ✅ | ✅ | ✅ | ✅ | | | | | | |

View file

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

View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 18773
"limit": 15800
},
"reportArgumentType": {
"limit": 2566
"limit": 2262
},
"reportAssignmentType": {
"limit": 320
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 5774
"limit": 5174
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5640
"limit": 5590
},
"reportMissingTypeArgument": {
"limit": 15498
"limit": 15327
},
"reportMissingTypeStubs": {
"limit": 40
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1061
"limit": 0
},
"reportOptionalOperand": {
"limit": 0
@ -90,7 +90,7 @@
"limit": 8
},
"reportReturnType": {
"limit": 213
"limit": 181
},
"reportTypedDictNotRequiredAccess": {
"limit": 26
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44589
"limit": 44375
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38882
"limit": 38487
},
"reportUnknownParameterType": {
"limit": 19806
"limit": 19590
},
"reportUnknownVariableType": {
"limit": 30456
"limit": 30126
},
"reportUnnecessaryCast": {
"limit": 116
"limit": 113
},
"reportUnnecessaryComparison": {
"limit": 698
"limit": 697
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 835
"limit": 834
},
"reportUntypedBaseClass": {
"limit": 0
@ -138,9 +138,9 @@
"limit": 139
},
"reportUnusedImport": {
"limit": 544
"limit": 541
},
"reportUnusedVariable": {
"limit": 142
"limit": 133
}
}

View file

@ -27,6 +27,7 @@ EXTRA_BOOLEAN_KEYS = frozenset(
"uses_embed_content",
"use_openai_responses_path",
"bedrock_converse_supports_strict_tools",
"thinking_always_on",
}
)
@ -72,6 +73,11 @@ ARRAY_KEYS: dict[str, JsonSchema] = {
"description": "Output modalities the model can produce.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
},
"reasoning_effort_levels": {
"type": "array",
"description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
"items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
},
"supported_regions": {
"type": "array",
"description": "Cloud regions the model is available in ('global' or region ids).",
@ -156,6 +162,9 @@ COST_DESCRIPTIONS: dict[str, str] = {
"input_cost_per_token": "USD per prompt token.",
"output_cost_per_token": "USD per generated token.",
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
"google_maps_grounding_cost_per_query": (
"USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
),
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
@ -211,6 +220,15 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
"enum": ["low", "medium", "high", "max", "xhigh"],
},
"default_reasoning_effort": {
"type": "string",
"description": (
"Reasoning effort the provider applies when the request omits reasoning_effort. "
"Gates whether a non-default temperature or the top_p/logprobs sampling params are "
"accepted, which hold only when the effort resolves to 'none'."
),
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"comment": STRING,
"audio_transcription_config": STRING,
}

View file

@ -60,4 +60,4 @@ if __name__ == "__main__":
print("\n💡 Tips:")
print("1. Run 'litellm-proxy login' to authenticate first")
print("2. Replace 'https://your-proxy.com' with your actual proxy URL")
print("3. The token is stored locally at ~/.litellm/token.json")
print("3. The token is stored in your OS keychain, or in ~/.litellm/token.json when there is none")

View file

@ -10,6 +10,11 @@
-- partitioned, so existing installs are unaffected until you run this.
--
-- IMPORTANT
-- * After partitioning, `prisma db push` (including the proxy's
-- --use_prisma_db_push startup mode) is NOT supported: it tries to rewrite
-- the primary key back to ("request_id"), which Postgres rejects on a
-- partitioned table. The proxy detects this and exits with guidance.
-- Use the default startup path (`prisma migrate deploy`) instead.
-- * Test on a staging copy first and take a backup.
-- * Postgres cannot convert a populated table to partitioned in place, so this
-- renames the old table aside and creates a fresh partitioned table.

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -1,6 +1,8 @@
"""
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled automatically by the get-responses call.
Cost tracking is handled by the get-responses call, which prices normally only because the
poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
same route are non-inference and free.
"""
from datetime import datetime, timedelta, timezone
@ -9,12 +11,14 @@ from typing import TYPE_CHECKING, Dict, Optional, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
STALE_OBJECT_CLEANUP_BATCH_SIZE,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -113,7 +117,8 @@ class CheckResponsesCost:
Check if background responses are complete and track their cost.
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
- Query the provider to check if response is complete
- Cost is automatically tracked by the get-responses call
- Cost is tracked by the get-responses call, billed because the poll is stamped
with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
- Mark responses in a terminal state as complete in the database
"""
try:
@ -153,6 +158,7 @@ class CheckResponsesCost:
# Prepare metadata with model information for cost tracking
litellm_metadata = {
"user_api_key_user_id": job.created_by or "default-user-id",
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
}
# Add model information if available

View file

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

View file

@ -12,7 +12,7 @@ Endpoints for /project operations
import json
from collections.abc import Sequence
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Request
@ -35,6 +35,8 @@ if TYPE_CHECKING:
LiteLLM_VerificationTokenActions,
)
from litellm import Router
router = APIRouter()
@ -205,6 +207,114 @@ def _check_team_project_limits(
)
def _project_models_missing_positive_quota(
models: list[str] | None,
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
) -> list[str]:
"""Return the models that lack a positive `rpm` AND `tpm` quota.
A valid quota is a positive integer; null, zero, and negative are rejected
because downstream rate limiters treat a non-positive limit as immediately
exhausted (every request blocked).
"""
def _is_positive(value: object) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and value > 0
rpm = rpm_limits or {}
tpm = tpm_limits or {}
return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))]
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
def _project_models_expanding_at_request_time(
models: Sequence[str] | None, access_group_names: frozenset[str]
) -> tuple[str, ...]:
"""Entries project auth expands to many concrete models (`all-proxy-models`, `*` patterns,
access groups). The rate limiter looks quotas up by the exact requested model name, so a
quota keyed on one of these entries is never applied."""
return tuple(
model
for model in (models or ())
if model == SpecialModelNames.all_proxy_models.value or "*" in model or model in access_group_names
)
def _raise_on_project_models_expanding_at_request_time(
models: Sequence[str] | None, access_group_names: frozenset[str]
) -> None:
expanding: Final = _project_models_expanding_at_request_time(models, access_group_names)
if not expanding:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead."
},
)
def _raise_on_missing_project_model_quota(
data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset()
) -> None:
"""Require a positive `rpm`/`tpm` quota for every model on project CREATE.
`model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request
model's `set_model_info` validator, so they are read from there.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
metadata = data.metadata or {}
missing = _project_models_missing_positive_quota(
data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit")
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
},
)
def _raise_on_missing_project_model_quota_on_update(
data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset()
) -> None:
"""Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE.
`/project/update` replaces `models` and `metadata` when they are provided, so the
check runs on what the project WILL look like: a partial update that doesn't touch
models/quota keeps the existing values, while one that adds a model or clears a
model's quota must leave every resulting model with a positive limit.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or [])
resulting_metadata = (
data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {})
)
_raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names)
missing = _project_models_missing_positive_quota(
resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit")
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
},
)
async def _create_budget_for_project(
data: NewProjectRequest,
user_id: str | None,
@ -352,7 +462,9 @@ async def new_project(
```
"""
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
)
@ -399,6 +511,10 @@ async def new_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model added to the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router))
# Check if user has permission to create projects for this team
# only team admins can create projects for their team
has_permission = await _check_user_permission_for_project(
@ -538,7 +654,9 @@ async def update_project(
```
"""
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
user_api_key_cache,
@ -642,6 +760,12 @@ async def update_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model the update would leave on the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota_on_update(
data, existing_project, _router_access_group_names(llm_router)
)
# Prepare update data
update_data = _jsonified(prisma_client, data.model_dump(exclude_none=True, exclude={"project_id"}))
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.57"
version = "0.1.61"
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.57"
version = "0.1.61"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

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

View file

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

View file

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

View file

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

View file

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

View 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

View file

@ -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: {}

View file

@ -1,4 +0,0 @@
UPDATE "LiteLLM_SpendLogs"
SET "created_at" = "endTime",
"updated_at" = "endTime"
WHERE "created_at" > "endTime" + interval '1 hour';

View file

@ -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;

View file

@ -0,0 +1,18 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

View file

@ -0,0 +1,22 @@
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_cost" DOUBLE PRECISION;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "shadow_classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_cache_hit" BOOLEAN NOT NULL DEFAULT false;
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_ShadowEvalFunnel" (
"job_id" TEXT NOT NULL,
"not_sampled" INTEGER NOT NULL DEFAULT 0,
"unjudgeable" INTEGER NOT NULL DEFAULT 0,
"shed" INTEGER NOT NULL DEFAULT 0,
"withheld" INTEGER NOT NULL DEFAULT 0,
CONSTRAINT "LiteLLM_ShadowEvalFunnel_pkey" PRIMARY KEY ("job_id")
);

View file

@ -754,6 +754,7 @@ model LiteLLM_DailyUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -789,6 +790,7 @@ model LiteLLM_DailyOrganizationSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -824,6 +826,7 @@ model LiteLLM_DailyEndUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -858,6 +861,7 @@ model LiteLLM_DailyAgentSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -892,6 +896,7 @@ model LiteLLM_DailyTeamSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -929,6 +934,7 @@ model LiteLLM_DailyTagSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -1502,7 +1508,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,12 +1532,28 @@ model LiteLLM_ShadowEvalAttempt {
shadow_model String?
confidence Float?
judge_cost Float @default(0)
shadow_cost Float @default(0)
real_cost Float? // NULL = row predates cost measurement; comparisons read only measured rows
real_classifier_cost Float @default(0)
shadow_classifier_cost Float @default(0)
real_cache_hit Boolean @default(false)
error String?
created_at DateTime @default(now())
@@index([job_id])
}
// Per-leg sampling funnel counters the attempt rows cannot derive: requests an
// admitting job saw but did not judge. attempted = the leg's attempt rows; the
// leg's eligible traffic = not_sampled + unjudgeable + shed + withheld + attempted.
model LiteLLM_ShadowEvalFunnel {
job_id String @id
not_sampled Int @default(0)
unjudgeable Int @default(0)
shed Int @default(0)
withheld Int @default(0)
}
// ---------------------------------------------------------------------------
// Workflow Run Tracking
//

View file

@ -40,6 +40,65 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
)
_SPEND_LOGS_PK_CLAUSE_RE = re.compile(
r'^(?:DROP\s+CONSTRAINT\s+"[^"]*_pkey"'
r'|ADD\s+(?:CONSTRAINT\s+"[^"]*"\s+)?PRIMARY\s+KEY\s*\([^)]*\))$',
re.IGNORECASE,
)
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
"reconciles the database against schema.prisma, which declares the unpartitioned "
"primary key (\"request_id\"), and Postgres rejects that rewrite with: unique "
"constraint on partitioned table must include all partitioning columns. Start the "
"proxy without --use_prisma_db_push so it uses `prisma migrate deploy`, which only "
"applies shipped migrations and leaves the partitioned primary key alone."
)
def _without_sql_comments(statement: str) -> str:
return "\n".join(
line
for line in statement.splitlines()
if line.strip() and not line.strip().startswith("--")
).strip()
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
if not prefix_match:
return statement
kept = tuple(
clause.strip()
for clause in statement[prefix_match.end():].split(",\n")
if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause.strip())
)
if not kept:
return None
return statement[: prefix_match.end()] + ",\n".join(kept)
def filter_partitioned_spend_logs_diff(diff_sql: str) -> str:
"""Drop statements from a `prisma migrate diff` script that fight the
SpendLogs partitioning runbook (db_scripts/partition_spend_logs.sql): the
primary-key rewrite on "LiteLLM_SpendLogs", which Postgres rejects on a
partitioned table, and drops of runbook artifacts such as
"LiteLLM_SpendLogs_legacy"."""
kept = tuple(
filtered
for statement in diff_sql.split(";")
for bare in (_without_sql_comments(statement),)
if bare and not _SPEND_LOGS_ARTIFACT_DROP_RE.match(bare)
for filtered in (_without_spend_logs_pk_clauses(bare),)
if filtered is not None
)
return "".join(f"{statement};\n\n" for statement in kept)
def _migration_timestamp(name: str) -> int:
"""Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name.
@ -355,7 +414,24 @@ class ProxyExtrasDBManager:
return
logger.info(f"Migration diff created at {diff_sql_path}")
if ProxyExtrasDBManager.spend_logs_is_partitioned():
filtered_sql = filter_partitioned_spend_logs_diff(
diff_sql_path.read_text()
)
diff_sql_path.write_text(filtered_sql)
logger.info(
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
"rewrite and partitioning artifacts from the drift script"
)
if not filtered_sql.strip():
logger.info("Drift script is empty after filtering; nothing to apply")
if not mark_all_applied:
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
return
# 2. Run prisma db execute to apply the migration
applied_ok = False
try:
logger.info("Running prisma db execute to apply the migration diff...")
result = subprocess.run(
@ -376,6 +452,7 @@ class ProxyExtrasDBManager:
)
logger.info(f"prisma db execute stdout: {result.stdout}")
logger.info("✅ Migration diff applied successfully")
applied_ok = True
except subprocess.CalledProcessError as e:
logger.warning(f"Failed to apply migration diff: {e.stderr}")
except subprocess.TimeoutExpired:
@ -384,6 +461,16 @@ class ProxyExtrasDBManager:
# 3. Mark all migrations as applied
if not mark_all_applied:
return
if not applied_ok:
logger.warning(
"Drift script failed to apply; NOT marking migrations as "
"applied so a later migration run can retry them"
)
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
@staticmethod
def _mark_migrations_applied(migrations_dir: str) -> None:
migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
logger.info(f"Resolving {len(migration_names)} migrations")
for migration_name in migration_names:
@ -410,6 +497,55 @@ class ProxyExtrasDBManager:
f"Failed to resolve migration {migration_name}: {e.stderr}"
)
@staticmethod
def spend_logs_is_partitioned() -> bool:
"""True when the connected database's LiteLLM_SpendLogs is a
partitioned table in Prisma's target schema (the `schema` URL param,
falling back to Prisma's default target, public), i.e. the operator
ran db_scripts/partition_spend_logs.sql. Returns False when psycopg is
unavailable or the database cannot be reached, preserving the
pre-existing behavior in those cases."""
database_url = os.getenv("DATABASE_URL")
if not database_url:
return False
try:
import psycopg
except ImportError:
return False
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
try:
with psycopg.connect(
cleaned_url, connect_timeout=10, autocommit=True
) as conn:
row = conn.execute(
"SELECT 1 "
"FROM pg_partitioned_table pt "
"JOIN pg_class c ON c.oid = pt.partrelid "
"JOIN pg_namespace n ON n.oid = c.relnamespace "
"WHERE c.relname = 'LiteLLM_SpendLogs' "
" AND n.nspname = %s",
(
ProxyExtrasDBManager._prisma_schema_param(database_url)
or "public",
),
).fetchone()
except (psycopg.OperationalError, psycopg.DatabaseError):
return False
return row is not None
@staticmethod
def _prisma_schema_param(url: str) -> Optional[str]:
"""The `schema` query param Prisma uses to pick its target schema,
or None when the URL does not set one."""
from urllib.parse import urlparse, parse_qsl
return next(
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
None,
)
@staticmethod
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
@ -528,7 +664,8 @@ class ProxyExtrasDBManager:
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
if not use_migrate:
# Preserve `prisma db push` path unchanged.
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
@ -972,6 +1109,8 @@ class ProxyExtrasDBManager:
)
raise
else:
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
# Use prisma db push with increased timeout
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.87"
version = "0.4.90"
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.87"
version = "0.4.90"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

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

View file

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

View file

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

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

View file

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

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

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

View 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;

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

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

View 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, &params)
}
#[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, .. }
));
}
}

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

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

View file

@ -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]";

View file

@ -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 {

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

View file

@ -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;

View file

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

View file

@ -0,0 +1 @@
pub mod transformation;

View file

@ -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), &params(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"),
]
);
}

View file

@ -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;

View file

@ -1 +1,2 @@
pub mod chat_completions;
pub mod messages;

View file

@ -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::*;

View file

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

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