Merge remote-tracking branch 'origin/main' into fix/trim-messages-response-tokens-shape

This commit is contained in:
taltal-beep 2026-10-15 09:22:07 +03:00
commit 1e3571271d
2170 changed files with 272373 additions and 43205 deletions

View file

@ -38,7 +38,7 @@ commands:
parameters:
category:
type: enum
enum: ["backend", "client", "provider-harness", "redis-compat"]
enum: ["backend", "client"]
default: "backend"
steps:
- run:
@ -339,8 +339,15 @@ commands:
Add-Content -Path $PROFILE -Value "`$env:Path = `"$cargoBin;`$env:Path`""
}
setup_litellm_test_deps:
parameters:
checkout:
type: boolean
default: true
steps:
- checkout
- when:
condition: << parameters.checkout >>
steps:
- checkout
- setup_google_dns
- install_uv
- install_rust
@ -376,13 +383,28 @@ jobs:
shell: bash.exe
command: bash .circleci/scripts/path_filter.sh windows-release
- run:
name: Install Python
name: Install Python 3.11.9 from python.org
no_output_timeout: 6m
command: |
choco install python --version=3.11.0 -y --no-progress --force
refreshenv
python --version
environment:
CHOCOLATEY_CONFIRM_ALL: "true"
$ProgressPreference = "SilentlyContinue"
$installer = Join-Path $env:TEMP "python-3.11.9-amd64.exe"
Invoke-WebRequest -Uri https://www.python.org/ftp/python/3.11.9/python-3.11.9-amd64.exe -OutFile $installer -TimeoutSec 120
$expected = "5ee42c4eee1e6b4464bb23722f90b45303f79442df63083f05322f1785f5fdde"
$actual = (Get-FileHash -Path $installer -Algorithm SHA256).Hash.ToLower()
if ($actual -ne $expected) {
throw "Python installer hash mismatch: expected $expected got $actual"
}
$arguments = "/quiet InstallAllUsers=1 PrependPath=1 Include_test=0 Include_doc=0 Include_launcher=0 TargetDir=C:\Python311"
$process = Start-Process -FilePath $installer -ArgumentList $arguments -PassThru
if (-not $process.WaitForExit(300000)) {
$process.Kill()
throw "Python installer did not finish within 5 minutes"
}
if ($process.ExitCode -ne 0) {
throw "Python installer exited with $($process.ExitCode)"
}
Remove-Item $installer
& C:\Python311\python.exe --version
- install_windows_toolchain
- run:
name: Record the Rust build environment for the release cargo cache key
@ -405,7 +427,7 @@ jobs:
environment:
UV_HTTP_TIMEOUT: "300"
command: |
$env:Path = "$HOME\.cargo\bin;$HOME\.local\bin;$env:Path"
$env:Path = "C:\Python311;C:\Python311\Scripts;$HOME\.cargo\bin;$HOME\.local\bin;$env:Path"
uv build --wheel --out-dir dist
if ($LASTEXITCODE -ne 0) {
exit $LASTEXITCODE
@ -512,11 +534,12 @@ jobs:
--set image.repository=litellm-ci \
--set image.tag=${IMAGE_TAG} \
--set image.pullPolicy=Never
# Wait for pod to be ready
echo "Waiting 30 seconds for pod to be ready..."
sleep 30
if ! kubectl rollout status deployment/litellm --timeout=180s; then
kubectl describe pods -l app.kubernetes.io/name=litellm
kubectl logs -l app.kubernetes.io/name=litellm --tail=200
exit 1
fi
# Print pod logs before running tests
echo "Printing pod logs..."
kubectl logs $(kubectl get pods -l app.kubernetes.io/name=litellm -o jsonpath="{.items[0].metadata.name}")
@ -619,7 +642,7 @@ jobs:
name: Run Basic Proxy Startup Tests (Health Readiness and Chat Completion)
command: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
TEST_FILES=$(circleci tests glob "tests/integration/configuration/test_basic_proxy_startup.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 --tb=short \
@ -696,7 +719,7 @@ jobs:
name: Run second round of tests
command: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
TEST_FILES=$(circleci tests glob "tests/integration/configuration/test_basic_proxy_startup.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 --tb=short \
@ -708,84 +731,46 @@ jobs:
# Store test results
- store_test_results:
path: test-results
proxy_pass_through_endpoint_tests:
machine:
image: ubuntu-2204:2024.04.1
e2e_ui_build:
docker:
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: large
working_directory: ~/project
steps:
- checkout
- run:
name: Generate LiteLLM master key
command: |
key="$(openssl rand -hex 16)"
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- start_postgres
- attach_workspace:
at: ~/project
- run:
name: Load Docker Database Image
command: |
zstd -d litellm-docker-database.tar.zst --stdout | docker load
docker images | grep litellm-docker-database
- run:
name: Run Docker container
command: |
docker run -d \
-p 4000:4000 \
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" \
-e OPENAI_API_KEY=$OPENAI_API_KEY \
-e GEMINI_API_KEY=$GEMINI_API_KEY \
-e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
-e ASSEMBLYAI_API_KEY=$ASSEMBLYAI_API_KEY \
-e AZURE_API_KEY=$AZURE_API_KEY \
-e AZURE_API_BASE=$AZURE_API_BASE \
-e USE_DDTRACE=True \
-e DD_API_KEY=$DD_API_KEY \
-e DD_SITE=$DD_SITE \
-e LITELLM_LICENSE=$LITELLM_LICENSE \
-e LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES=true \
-e LITELLM_LOG=ERROR \
--add-host host.docker.internal:host-gateway \
--name my-app \
-v $(pwd)/litellm/proxy/example_config_yaml/pass_through_config.yaml:/app/config.yaml \
-v $(pwd)/litellm/proxy/example_config_yaml/custom_auth_basic.py:/app/custom_auth_basic.py \
litellm-docker-database:ci \
--config /app/config.yaml \
--port 4000
- run:
name: Start outputting logs
command: docker logs -f my-app
background: true
- wait_for_service:
url: http://localhost:4000
timeout: "300"
# 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).
- install_node
- restore_cache:
keys:
- ui-dashboard-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- run:
name: Install Node.js test dependencies
name: Install dashboard dependencies
command: |
cd tests/pass_through_tests
cd ui/litellm-dashboard
npm ci
- save_cache:
key: ui-dashboard-node-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- run:
name: Run Vertex AI, Google AI Studio Node.js tests
name: Build UI bundle
command: |
cd tests/pass_through_tests
NODE_OPTIONS=--experimental-vm-modules npx jest . --verbose
no_output_timeout: 30m
cd ui/litellm-dashboard
npm run build
mkdir -p ../../ui-bundle
mv out ../../ui-bundle/out
find ../../ui-bundle/out -name '*.html' ! -name 'index.html' | while read -r f; do
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
done
- persist_to_workspace:
root: ~/project
paths:
- ui-bundle/out
e2e_ui_testing:
docker:
@ -798,7 +783,7 @@ jobs:
POSTGRES_USER: e2euser
POSTGRES_PASSWORD: e2epassword
POSTGRES_DB: litellm_e2e
resource_class: large
resource_class: xlarge
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e"
@ -817,6 +802,8 @@ jobs:
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes:
category: client
- attach_workspace:
at: ~/project
- setup_google_dns
- install_node
- install_uv
@ -835,7 +822,7 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
# The cimg/python:3.12-browsers image already ships the Chromium system
@ -844,36 +831,19 @@ jobs:
# 5-10 minutes to the job and frequently stalls on flaky Ubuntu mirrors,
# so we install just the browser binary.
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
cd tests/e2e/ui
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
key: ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
- ~/.cache/ms-playwright
- run:
name: Build UI from source
# Prior version used `cp -r out/ ../../litellm/proxy/_experimental/out/`.
# GNU cp (used on CircleCI's Ubuntu image) interprets that as "copy the
# source directory as a child of the destination" when the destination
# already exists — silently creating `_experimental/out/out/` instead of
# replacing the served bundle. The proxy continued serving whatever was
# checked into `_experimental/out/*`, so this job was effectively testing
# the pre-build bundle on every run. Replace-and-move guarantees the
# freshly built bundle is what the proxy actually serves.
name: Use built UI bundle
command: |
cd ui/litellm-dashboard
npm run build
rm -rf ../../litellm/proxy/_experimental/out
mv out ../../litellm/proxy/_experimental/out
# Restructure HTML so extensionless routes work (login.html -> login/index.html)
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
done
rm -rf litellm/proxy/_experimental/out
cp -r ui-bundle/out litellm/proxy/_experimental/out
- wait_for_service:
url: tcp://localhost:5432
timeout: "30"
@ -940,7 +910,7 @@ jobs:
command: |
cd tests/e2e/ui
LITELLM_LICENSE="$LITELLM_LICENSE" \
npx playwright test --config playwright.config.ts
npx playwright test --config playwright.config.ts --workers=3 --retries=0
no_output_timeout: 10m
- store_artifacts:
path: tests/e2e/ui/test-results
@ -978,6 +948,8 @@ jobs:
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes:
category: client
- attach_workspace:
at: ~/project
- setup_google_dns
- install_node
- install_uv
@ -996,31 +968,23 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
cd tests/e2e/ui
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
key: ui-playwright-node-deps-v1-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
- ~/.cache/ms-playwright
- run:
name: Build UI from source
name: Use built UI bundle
command: |
cd ui/litellm-dashboard
npm run build
rm -rf ../../litellm/proxy/_experimental/out
mv out ../../litellm/proxy/_experimental/out
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
done
rm -rf litellm/proxy/_experimental/out
cp -r ui-bundle/out litellm/proxy/_experimental/out
- wait_for_service:
url: tcp://localhost:5432
timeout: "30"
@ -1271,45 +1235,6 @@ jobs:
! docker exec my-app true 2>/dev/null || fail "container is still running"
echo "Proxy refused to serve (exit $exit_code) and never reached startup. Test passed."
provider_replay_harness:
docker:
- &python312_image
image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
- image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
working_directory: ~/project
resource_class: medium
environment:
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
E2E_PROVIDER_CACHE: "0"
E2E_FIXTURE_MODE: live
steps:
- checkout
- run:
name: Generate LiteLLM master key
command: |
key="$(openssl rand -hex 16)"
printf 'export LITELLM_MASTER_KEY=sk-%s\n' "$key" >> "$BASH_ENV"
- skip_if_unrelated_changes:
category: provider-harness
- setup_litellm_test_deps
- wait_for_service:
url: tcp://localhost:6379
- run:
name: Test provider capture and replay harness
command: |
mkdir -p test-results/provider-replay-harness
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
--junitxml=test-results/provider-replay-harness/junit.xml \
tests/e2e_harness/test_provider_edge.py tests/e2e_harness/test_fixture_bundle.py \
tests/e2e_harness/test_fixture_canonical.py tests/e2e_harness/test_fixture_mode.py \
tests/code_coverage_tests/test_provider_replay_harness.py \
tests/code_coverage_tests/test_provider_cache.py
- store_test_results:
path: test-results/provider-replay-harness
integration_contracts:
parameters:
suite:
@ -1358,9 +1283,19 @@ jobs:
- run:
name: Build the candidate dashboard
command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build
- when:
condition:
equal: [mcp, << parameters.suite >>]
steps:
- run:
name: Install MCP SDK1 peer
command: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
server_args: "-c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000"
- start_redis
- run:
name: Run owned integration contracts
@ -1413,12 +1348,27 @@ jobs:
default: standard
shards:
type: integer
machine:
image: ubuntu-2204:2024.04.1
resource_class: medium
docker:
- image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
resource_class: small
working_directory: ~/project
environment:
UV_PYTHON_DOWNLOADS: never
steps:
- setup_litellm_test_deps
- checkout
- install_uv
- run:
name: Install the locked shard verifier and coverage dependencies
command: |
uv export --frozen --all-groups --all-extras --no-emit-project --no-emit-workspace --no-header --no-annotate \
| awk '/^[^ ]/ { keep = $0 ~ /^(annotated-types|coverage|pydantic|pydantic-core|typing-extensions|typing-inspection)==/ } keep' \
> /tmp/fan-in-requirements.txt
test "$(grep -c '^[^ ]' /tmp/fan-in-requirements.txt)" -eq 6
uv venv .venv --python 3.12
uv pip install --no-config --python .venv --require-hashes -r /tmp/fan-in-requirements.txt
- attach_workspace:
at: /tmp/integration-shards
- run:
@ -1534,7 +1484,8 @@ jobs:
steps:
- checkout
- skip_if_unrelated_changes
- setup_litellm_test_deps
- setup_litellm_test_deps:
checkout: false
- start_postgres:
db_name: litellm_test
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
@ -1579,112 +1530,15 @@ jobs:
- store_artifacts:
path: test-results
mcp_integration:
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
working_directory: ~/project
environment:
COVERAGE_CORE: sysmon
LITELLM_LOCAL_MODEL_COST_MAP: "True"
steps:
- checkout
- skip_if_unrelated_changes
- setup_litellm_test_deps
- run:
name: Generate Prisma client
command: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- run:
name: Install MCP SDK1 peer
command: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV"
- run:
name: Run MCP integration tests
command: |
mkdir -p test-results
env -u OPENAI_API_KEY -u ANTHROPIC_API_KEY \
timeout --signal=TERM 20m uv run --no-sync pytest \
tests/mcp_tests tests/unit/experimental_mcp_client tests/unit/proxy/_experimental/mcp_server \
tests/unit/responses/mcp --tb=short -vv --maxfail=10 -n 2 --dist=loadscope --reruns 0 \
--reruns-delay 1 --timeout=120 --rerun-except "from pytest-timeout" --durations=20 \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml:coverage.xml \
--cov-config=pyproject.toml -o junit_family=xunit1 --junitxml=test-results/junit.xml
- install_codecov_cli
- run:
name: Upload MCP integration coverage
when: always
command: |
[ -f coverage.xml ] || { echo "no coverage.xml produced"; exit 1; }
codecov upload-process --disable-search --fail-on-error -f coverage.xml \
-F mcp-integration -C "$CIRCLE_SHA1" \
-n "mcp-integration-${CIRCLE_BUILD_NUM}" --git-service github
- store_test_results:
path: test-results
- store_artifacts:
path: test-results
redis_compat:
parameters:
redis_py:
type: string
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
working_directory: ~/project
steps:
- checkout
- skip_if_unrelated_changes:
category: redis-compat
- setup_litellm_test_deps
- run:
name: Pin redis-py version
command: |
uv pip install "redis==<< parameters.redis_py >>"
uv run --no-sync python -c "import redis; assert redis.__version__ == '<< parameters.redis_py >>', redis.__version__; print('redis-py', redis.__version__)"
- run:
name: Install redis-server 7.2.16
command: |
mkdir -p "$HOME/.local/bin"
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
docker cp "${cid}":/usr/local/bin/redis-server "$HOME/.local/bin/redis-server"
docker rm "$cid"
redis-server --version | grep -q 'v=7.2.16'
- run:
name: Run Redis compatibility tests
command: |
mkdir -p test-results
env -u CASSETTE_REDIS_URL -u AZURE_CLIENT_ID -u AZURE_CLIENT_SECRET -u AZURE_TENANT_ID \
timeout --signal=TERM 15m uv run --no-sync pytest \
tests/unit/test_redis.py tests/unit/caching/test_redis_connection_pool.py \
tests/unit/caching/test_redis_cluster_cache.py tests/unit/caching/test_evicted_client_closer.py \
tests/integration/sdk/test_redis_cluster_iam_auth.py::test_sync_cluster_authenticates_with_azure_credentials \
tests/integration/sdk/test_redis_cluster_iam_auth.py::test_sync_cluster_authenticates_with_gcp_credentials \
--tb=short -vv --reruns 2 --reruns-delay 1 --durations=20 --cov=./litellm \
--cov-report=xml:coverage.xml -o junit_family=xunit1 --junitxml=test-results/junit.xml
- when:
condition:
equal: ["5.3.1", << parameters.redis_py >>]
steps:
- install_codecov_cli
- run:
name: Upload Redis compatibility coverage
when: always
command: |
[ -f coverage.xml ] || { echo "no coverage.xml produced"; exit 1; }
codecov upload-process --disable-search --fail-on-error -f coverage.xml \
-F redis-compat -C "$CIRCLE_SHA1" \
-n "redis-compat-${CIRCLE_BUILD_NUM}" --git-service github
- store_test_results:
path: test-results
- store_artifacts:
path: test-results
workflows:
migration_startup:
when: << pipeline.parameters.run_migration_tests >>
jobs: &migration_jobs
migration_startup_scheduled:
triggers:
- schedule:
cron: "17 0,6,12,18 * * *"
filters:
branches:
only: main
jobs:
- build_docker_database_image:
migration_qualification: true
- migration_startup_tests:
@ -1709,33 +1563,11 @@ workflows:
suite: shaped
baseline: true
requires: [build_docker_database_image]
migration_startup_scheduled:
triggers:
- schedule:
cron: "17 0,6,12,18 * * *"
filters:
branches:
only: main
jobs: *migration_jobs
routing_parity:
when:
not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- routing_parity:
name: routing-parity-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, cost, mcp]
integration:
unless:
or:
- << pipeline.parameters.run_migration_tests >>
- not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- integration_contracts:
name: integration-<< matrix.suite >>
filters: &integration_default_run not pipeline.parameters.run_migration_tests and pipeline.parameters.routing_parity_base == ""
matrix:
parameters:
suite: [sdk, cost, browser]
@ -1745,150 +1577,178 @@ workflows:
suite: security
parallelism: 1
resource_class: xlarge
filters: *integration_default_run
- integration_contracts:
name: integration-management
suite: management
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-management-coverage
suite: management
shards: 4
requires: [integration-management]
filters: *integration_default_run
- integration_contracts:
name: integration-accounting
suite: accounting
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-accounting-coverage
suite: accounting
shards: 4
requires: [integration-accounting]
filters: *integration_default_run
- integration_contracts:
name: integration-providers
suite: providers
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-providers-coverage
suite: providers
shards: 4
requires:
- integration-providers
filters: *integration_default_run
- integration_contracts:
name: integration-database
suite: database
parallelism: 2
filters: *integration_default_run
- integration_shard_coverage:
name: integration-database-coverage
suite: database
shards: 2
requires: [integration-database]
filters: *integration_default_run
- integration_contracts:
name: integration-mcp
suite: mcp
resource_class: xlarge
parallelism: 2
filters: *integration_default_run
- integration_shard_coverage:
name: integration-mcp-coverage
suite: mcp
shards: 2
requires:
- integration-mcp
filters: *integration_default_run
- integration_contracts:
name: integration-extensions
suite: extensions
parallelism: 8
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-extensions-coverage
suite: extensions
shards: 8
requires:
- integration-extensions
filters: *integration_default_run
- integration_contracts:
name: integration-management-replica
suite: management
mode: replica
parallelism: 4
resource_class: xlarge
filters: *integration_default_run
- integration_shard_coverage:
name: integration-management-replica-coverage
suite: management
mode: replica
shards: 4
requires: [integration-management-replica]
filters: *integration_default_run
- integration_contracts:
name: integration-database-replica
suite: database
mode: replica
parallelism: 2
filters: *integration_default_run
- integration_shard_coverage:
name: integration-database-replica-coverage
suite: database
mode: replica
shards: 2
requires: [integration-database-replica]
- postgres_suite:
name: proxy-behavior
test_path: tests/proxy_behavior
seed: true
coverage_flag: lens-postgres
timeout_minutes: 25
- postgres_suite:
name: proxy-security
test_path: tests/proxy_security_tests
seed: true
timeout_minutes: 15
- postgres_suite:
name: schema-migration
test_path: tests/proxy_migration_tests
seed: false
timeout_minutes: 20
workers: 2
filters: *integration_default_run
- postgres_suite:
name: roi-database
test_path: tests/integration/database/test_roi_observed.py
seed: false
coverage_flag: roi-postgres
timeout_minutes: 10
- mcp_integration:
name: mcp-integration
- redis_compat:
name: redis-compat-<< matrix.redis_py >>
matrix:
parameters:
redis_py:
- "5.3.1"
- "6.4.0"
- "7.4.1"
- "8.0.1"
- windows_release_wheel
- base_sdk_install
build_and_test:
unless:
or:
- << pipeline.parameters.run_migration_tests >>
- not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- provider_replay_harness
- build_docker_database_image
- e2e_ui_testing
- e2e_ui_testing_server_root_path
filters: *integration_default_run
- windows_release_wheel:
filters: *integration_default_run
- base_sdk_install:
filters: *integration_default_run
- build_docker_database_image:
filters: *integration_default_run
- proxy_logging_guardrails_model_info_tests:
requires:
- build_docker_database_image
- proxy_pass_through_endpoint_tests:
requires:
- build_docker_database_image
filters: *integration_default_run
- db_migration_disable_update_check:
requires:
- build_docker_database_image
- helm_chart_testing:
requires:
- build_docker_database_image
filters: *integration_default_run
- test_bad_database_url:
requires:
- build_docker_database_image
filters: *integration_default_run
- helm_chart_testing:
requires:
- build_docker_database_image
filters: *integration_default_run
- e2e_ui_build:
filters: *integration_default_run
- e2e_ui_testing:
requires: [e2e_ui_build]
filters: *integration_default_run
- e2e_ui_testing_server_root_path:
requires: [e2e_ui_build]
filters: *integration_default_run
- build_docker_database_image:
name: build_docker_database_image-migration
migration_qualification: true
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-startup
suite: startup
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-recovery
suite: recovery
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-legacy-and-pooling
suite: legacy
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-upgrade
suite: upgrade
baseline: true
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- migration_startup_tests:
name: migration-upgrade-shaped
suite: shaped
baseline: true
requires: [build_docker_database_image-migration]
filters: pipeline.parameters.run_migration_tests
- routing_parity:
name: routing-parity-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, cost, mcp]
filters: pipeline.parameters.routing_parity_base != ""

View file

@ -1,12 +1,11 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only|mcp-dependencies|windows-release|redis-compat>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|cost-map-only|mcp-dependencies|windows-release|redis-compat>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
has_cost_map=false
has_mcp_dependencies=false
has_windows_release=false
@ -15,7 +14,7 @@ outside_cost_map_set=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .circleci/config.yml | .circleci/scripts/classify_changes.sh | .circleci/scripts/path_filter.sh | pyproject.toml | uv.lock)
litellm/_redis.py | litellm/_redis_credential_provider.py | litellm/caching/redis_cache.py | litellm/caching/evicted_client_closer.py | tests/unit/test_redis.py | tests/integration/sdk/test_redis_cluster_iam_auth.py | tests/unit/caching/test_redis_connection_pool.py | tests/unit/caching/test_redis_cluster_cache.py | tests/unit/caching/test_evicted_client_closer.py | .github/workflows/test-unit.yml | .github/actions/detect-changes/* | .github/scripts/detect_changes.sh | .circleci/scripts/classify_changes.sh | tests/integration/run.py | pyproject.toml | uv.lock)
has_redis_compat=true ;;
esac
case "$file" in
@ -23,11 +22,6 @@ while IFS= read -r file || [ -n "$file" ]; do
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/unit/test_circleci_path_filter.py | tests/unit/test_detect_changes.py)
has_mcp_dependencies=true ;;
esac
case "$file" in
tests/e2e/*/*.py | tests/e2e_harness/*/*.py) : ;;
tests/e2e/*.py | tests/e2e_harness/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/unit/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
has_provider_harness=true ;;
esac
case "$file" in
litellm-rust/* | litellm/rust_bridge/* | rust-toolchain.toml | pyproject.toml | uv.lock | tests/windows_tests/* | .circleci/*)
has_windows_release=true ;;
@ -41,7 +35,7 @@ while IFS= read -r file || [ -n "$file" ]; do
case "$file" in
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
has_cost_map=true ;;
tests/test_litellm/* | tests/proxy_unit_tests/* | tests/unit/proxy/*) : ;;
tests/unit/proxy/*) : ;;
*) outside_cost_map_set=true ;;
esac
done
@ -53,9 +47,6 @@ case "$category" in
cost-map-only)
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
;;
provider-harness)
[ "$has_provider_harness" = true ] && echo run || echo skip
;;
windows-release)
[ "$has_windows_release" = true ] && echo run || echo skip
;;

View file

@ -1,7 +1,7 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: path_filter.sh <backend|client|provider-harness|redis-compat>}"
category="${1:?usage: path_filter.sh <backend|client|windows-release>}"
here="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
run_full() {

View file

@ -128,20 +128,14 @@ fi
sudo iptables -L integration_only -n -v -x > "$results/egress-guard.txt"
awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/egress-guard.txt"
isolated_workers=false
if [ "$mode" = parity ]; then
export INTEGRATION_WORKERS=1
else
case "$suite" in
management|accounting) export INTEGRATION_WORKERS=2; isolated_workers=true ;;
providers|extensions) export INTEGRATION_WORKERS=4; isolated_workers=true ;;
management|accounting) export INTEGRATION_WORKERS=2 ;;
providers|extensions) export INTEGRATION_WORKERS=4 ;;
esac
fi
if [ "$isolated_workers" = false ]; then
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
.venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 &
upstream_pid=$!
fi
if [ "$suite" = cost ]; then
export INTEGRATION_WORKERS=8
fi
@ -149,7 +143,16 @@ if [ "$suite" = security ]; then
export INTEGRATION_WORKERS=4
fi
if [ "$suite" = mcp ]; then
export INTEGRATION_WORKERS=4 INTEGRATION_COVERAGE=1
export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1
fi
isolated_workers=false
if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then
isolated_workers=true
fi
if [ "$isolated_workers" = false ]; then
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
.venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 &
upstream_pid=$!
fi
coverage_data="$PWD/$results/coverage/data"
proxy_command=(.venv/bin/python -m integration._support.proxy)
@ -191,9 +194,6 @@ start_proxy() {
> "$results/$log_name" 2>&1 &
launched_pid=$!
}
if { [ "$suite" = management ] || [ "$suite" = mcp ]; } && [ "$isolated_workers" = true ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
fi
if [ "$isolated_workers" = false ]; then
start_proxy 4000 proxy.log
proxy_pid="$launched_pid"
@ -201,7 +201,7 @@ if [ "$isolated_workers" = false ]; then
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
if [ "$suite" = management ] || [ "$suite" = mcp ]; then
if [ "$suite" = management ] || [ "$suite" = mcp ] || [ "$suite" = accounting ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log
peer_pid="$launched_pid"
@ -239,13 +239,24 @@ if [ "${CIRCLE_NODE_TOTAL:-1}" -gt 1 ]; then
shard_arguments=(--shard-files "$results/node-files.txt")
fi
shared_services=()
worker_coverage=()
if [ "$isolated_workers" = false ]; then
shared_services=(
"REDIS_HOST=$REDIS_HOST" "REDIS_PORT=$REDIS_PORT"
"INTEGRATION_PROXY_URL=$INTEGRATION_PROXY_URL" "INTEGRATION_PEER_URL=$INTEGRATION_PEER_URL"
"INTEGRATION_UPSTREAM_URL=$INTEGRATION_UPSTREAM_URL" "INTEGRATION_MASTER_KEY=$INTEGRATION_MASTER_KEY"
)
elif [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then
worker_coverage=("INTEGRATION_COVERAGE_FILE=$coverage_data")
fi
env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
DATABASE_URL="$DATABASE_URL" \
"${shared_services[@]}" "${worker_coverage[@]}" \
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
LITELLM_MODE=PRODUCTION \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
@ -254,6 +265,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_PROXY_DATABASE_URL="$INTEGRATION_PROXY_DATABASE_URL" \
INTEGRATION_PROXY_READ_REPLICA_URL="$INTEGRATION_PROXY_READ_REPLICA_URL" \
INTEGRATION_ROUTING="$INTEGRATION_ROUTING" \
MCP_TEST_PEER_PYTHON="${MCP_TEST_PEER_PYTHON:-}" \
.venv/bin/python tests/integration/run.py "$suite" --results "$results" \
"${shard_arguments[@]}"

View file

@ -10,7 +10,7 @@ from typing import Final
from xml.etree import ElementTree
SUITES: Final = {
"startup": (("test_startup.py",), 12),
"startup": (("test_startup.py",), 13),
"recovery": (("test_recovery.py",), 15),
"legacy": (("test_legacy.py", "test_pooling.py"), 11),
"upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5),

4
.github/CODEOWNERS vendored
View file

@ -1,2 +1,2 @@
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerrylu-berri
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri

View file

@ -68,7 +68,7 @@ After: the same request comes back with real token counts, so the dashboard show
- [ ] The handful of test files covering my change pass locally, e.g. `uv run pytest tests/unit/<your_test_file>.py -v`. Leave the suites (`make test-unit-*`, `make test-unit`) to CI: it finishes in ~15 minutes where a laptop takes an hour or more
- [ ] My PR passes all required CI/CD checks (e.g., lint, schema.d.ts sync check, etc.)
- [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem
- [ ] I have received a Greptile **Confidence Score of at least 4/5** before requesting a maintainer review (Greptile reviews automatically once the PR is opened; only comment `@greptileai` to re-request a review after pushing changes)
- [ ] Every automated review on my PR's latest commit has finished, and I have addressed or replied to each of its findings before requesting a maintainer review
## Delays in PR merge?

View file

@ -23,18 +23,6 @@ test_paths:
and that pull request gives it a new home
paths:
- tests/test_rust_python_harness.py
- reason: >-
Live-provider caching cases in tests/local_testing that remain outside CI. No CircleCI job
globs that directory anymore, and the legacy jobs that did always deselected these files.
The two Redis Cluster IAM authentication tests moved to tests/integration/sdk and remain
selected by node ID in redis_compat, where redis-server is installed.
The five keyless files now run in the caching-local shard, leaving live cases in these two.
Measured 2026-08-21 with no provider credentials: test_disk_cache_unit_tests.py needs
OPENAI_API_KEY for 2 of its 4 tests, and test_gcs_cache_unit_tests.py needs GCS credentials
for all 4. Keep these live cases outside PR CI rather than making a job red by construction
paths:
- 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
@ -60,48 +48,6 @@ test_paths:
choice it informed is settled
paths:
- tests/code_coverage_tests/test_aio_http_image_conversion.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
paths:
- tests/vector_store_tests/rag/test_rag_bedrock.py
- tests/vector_store_tests/rag/test_rag_openai.py
- tests/vector_store_tests/rag/test_rag_s3_vectors.py
- tests/vector_store_tests/rag/test_rag_vertex_ai.py
- tests/vector_store_tests/test_azure_ai_vector_store.py
- tests/vector_store_tests/test_azure_vector_store.py
- tests/vector_store_tests/test_bedrock_vector_store.py
- tests/vector_store_tests/test_gemini_vector_store.py
- tests/vector_store_tests/test_milvus_vector_store.py
- tests/vector_store_tests/test_openai_vector_store.py
- tests/vector_store_tests/test_ragflow_vector_store.py
- tests/vector_store_tests/test_s3_vectors_vector_store.py
- tests/vector_store_tests/test_vertex_ai_search_api_vector_store.py
- tests/vector_store_tests/test_vertex_ai_vector_store.py
- reason: >-
Throughput and memory-growth measurements whose runtime and variance make them unsuitable for
a per-pull-request job
paths:
- tests/load_tests/test_datadog_load_test.py
- tests/load_tests/test_langsmith_load_test.py
- tests/load_tests/test_linear_memory_growth.py
- tests/load_tests/test_memory_usage.py
- 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: >-
Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
locally rather than in pull-request jobs
paths:
- tests/load_tests/test_granian_admission_saturation.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

View file

@ -37,7 +37,6 @@ GLOB_CHARS = frozenset("*?")
# 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/test_litellm",
"tests/unit",
"tests/unit/enterprise",
"tests/unit/enterprise/enterprise_callbacks",

View file

@ -0,0 +1,30 @@
name: Deploy main to gateway-dev
on:
push:
branches: [main]
permissions: {}
concurrency:
group: deploy-dev-dispatch
cancel-in-progress: true
jobs:
dispatch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/create-github-app-token@fee1f7d63c2ff003460e3d139729b119787bc349 # v2.2.2
id: token
with:
app-id: ${{ secrets.LITELLM_OPS_DISPATCH_APP_ID }}
private-key: ${{ secrets.LITELLM_OPS_DISPATCH_APP_PRIVATE_KEY }}
owner: BerriAI
repositories: litellm-ops
permission-actions: write
- name: Dispatch the litellm-ops dev deploy
env:
GH_TOKEN: ${{ steps.token.outputs.token }}
run: gh api --method POST repos/BerriAI/litellm-ops/actions/workflows/deploy-dev.yml/dispatches -f ref=main

View file

@ -27,7 +27,7 @@ on:
- backend/routes/allowlist.py
- pyproject.toml
- litellm-proxy-extras/**
- tests/proxy_migration_tests/**
- tests/integration/database/image/**
- uv.lock
- ui/litellm-dashboard/package-lock.json
- ui/Dockerfile
@ -92,7 +92,7 @@ jobs:
LITELLM_ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v
# Scans the whole shipped artifact: OS/apk plus every language package
# baked into the image, including ones no lockfile declares (e.g. prisma's
@ -145,7 +145,7 @@ jobs:
LITELLM_ADMIN_MCP_SCHEMA_IMAGE: litellm-admin-mcp-schema:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v
migrations-image:
@ -177,7 +177,7 @@ jobs:
LITELLM_MIGRATION_SCRIPT: /app/run.py
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
python -m pytest --noconftest tests/integration/database/image/test_offline_image_migration.py -v
gateway-image:
name: gateway-image
@ -208,7 +208,7 @@ jobs:
LITELLM_IMAGE_COMPONENT: gateway
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
python -m pytest --noconftest tests/integration/database/image/test_component_image_serves_offline.py tests/integration/database/image/test_image_bedrock_realtime_extra.py tests/integration/database/image/test_image_admin_mcp.py -v
ui-image:
name: ui-image
@ -237,7 +237,7 @@ jobs:
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
python -m pytest --noconftest tests/integration/database/image/test_ui_image_serves_offline.py -v
backend-image:
name: backend-image
@ -272,4 +272,4 @@ jobs:
LITELLM_IMAGE_COMPONENT: backend
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_admin_mcp.py -v
python -m pytest --noconftest tests/integration/database/image/test_component_image_serves_offline.py tests/integration/database/image/test_image_admin_mcp.py -v

View file

@ -91,12 +91,12 @@ jobs:
- name: Run the Redis chaos load test
env:
E2E_REDIS_CHAOS: "1"
E2E_PERF_TESTS: "1"
LITELLM_PROXY_URL: http://localhost:4000
REDIS_HOST: 127.0.0.1
REDIS_PORT: "6379"
run: |
uv run --no-sync pytest tests/e2e/load/test_redis_chaos_e2e.py -v --tb=short -rA -s
uv run --no-sync pytest tests/e2e/perf/test_redis_chaos_e2e.py -v --tb=short -rA -s
- name: Show proxy log on failure
if: failure()

View file

@ -22,6 +22,16 @@ jobs:
actions: read
runs-on: ubuntu-latest
timeout-minutes: 15
services:
redis:
image: redis@sha256:e2debfb7956fa12c7ddc79d7e645c8cf26b30c99a6e9161ea9bf4171e1668a5f
ports:
- 6379:6379
options: >-
--health-cmd "redis-cli ping"
--health-interval 10s
--health-timeout 5s
--health-retries 5
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
@ -189,8 +199,9 @@ jobs:
if: steps.changes.outputs.decision != 'skip'
env:
LITELLM_MASTER_KEY: sk-e2e-harness-tests-reach-no-proxy
E2E_CACHE_TEST_REDIS_URL: redis://127.0.0.1:6379/0
run: |
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
if ! git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- tests/e2e tests/e2e_harness ':(exclude)tests/e2e/ui' tests/code_coverage_tests/test_provider_replay_harness.py tests/code_coverage_tests/test_provider_cache.py pyproject.toml uv.lock .github/workflows/test-linting.yml | grep -q .; then
echo "No changed e2e harness files; skipping."
exit 0
fi
@ -198,6 +209,9 @@ jobs:
CLAUDE_VERSION="$(retry uv run --no-sync python tests/e2e/claude_code/pr_gate_version_resolver.py)"
tests/e2e/claude_code/cron_vm/install_claude_code.sh "$CLAUDE_VERSION" "$RUNNER_TEMP/claude-cli"
PATH="$RUNNER_TEMP/claude-cli:$PATH" uv run --no-sync pytest -q tests/e2e_harness
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o "pythonpath=tests/e2e tests/e2e_harness" -p no:rerunfailures \
tests/code_coverage_tests/test_provider_replay_harness.py \
tests/code_coverage_tests/test_provider_cache.py
- name: Check for circular imports
if: steps.changes.outputs.decision != 'skip'
@ -586,6 +600,9 @@ jobs:
- name: check_unsafe_enterprise_import
run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
- name: check_telemetry_imports
run: uv run --no-sync python ./tests/code_coverage_tests/check_telemetry_imports.py
- name: ban_copy_deepcopy_kwargs
run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
@ -601,11 +618,8 @@ jobs:
- name: check_migrations_no_data_rewrites
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
- name: check_greptile_files_json
run: uv run --no-sync python ./tests/code_coverage_tests/check_greptile_files_json.py
- name: test_check_greptile_files_json
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_check_greptile_files_json.py
- name: check_migrations_lock_timeout
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_lock_timeout.py
- name: check_no_publicly_known_master_key
run: uv run --no-sync python ./tests/code_coverage_tests/check_no_publicly_known_master_key.py

View file

@ -234,7 +234,30 @@ jobs:
exit 1
- name: Run the startup tests
run: .venv/bin/python -m pytest --tb=short -vv tests/basic_proxy_startup_tests
shell: bash
run: |
set -euo pipefail
base_url=http://localhost:4000
readiness_response="$(curl -sS --fail-with-body --write-out $'\n%{http_code}' "${base_url}/health/readiness")"
readiness_status="${readiness_response##*$'\n'}"
readiness_body="${readiness_response%$'\n'*}"
[[ "${readiness_status}" == "200" ]]
jq -e '.status == "healthy" or .status == "connected"' <<< "${readiness_body}" > /dev/null
liveness_status="$(curl -sS --fail-with-body --output /dev/null --write-out '%{http_code}' "${base_url}/health/liveness")"
[[ "${liveness_status}" == "200" ]]
completion_response="$(curl -sS --fail-with-body --write-out $'\n%{http_code}' \
--header "Authorization: Bearer ${LITELLM_MASTER_KEY}" \
--header 'Content-Type: application/json' \
--data '{"model":"gpt-4","messages":[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Hello!"}]}' \
"${base_url}/chat/completions")"
completion_status="${completion_response##*$'\n'}"
completion_body="${completion_response%$'\n'*}"
[[ "${completion_status}" == "200" ]]
jq -e '
(.choices | length) == 1 and
.choices[0].message.role == "assistant" and
.choices[0].message.content == "Hello! This is a mock response from the fake OpenAI endpoint."
' <<< "${completion_body}" > /dev/null
- name: Show the proxy logs
if: always()

View file

@ -18,7 +18,6 @@ on:
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"
@ -49,7 +48,6 @@ on:
- "litellm/ocr/**"
- "litellm/llms/base_llm/ocr/**"
- "litellm/llms/custom_httpx/llm_http_handler.py"
- "tests/test_litellm/conftest.py"
- "Makefile"
- ".cargo/**"
- "pyproject.toml"

View file

@ -174,8 +174,6 @@ jobs:
- shard: integrations
test-path: |-
tests/test_litellm/integrations
tests/test_litellm/tracing
tests/unit/integrations
python-version: "3.12"
workers: 4
@ -235,6 +233,7 @@ jobs:
tests/unit/llms/base_llm
tests/unit/llms/baseten
tests/unit/llms/black_forest_labs
tests/unit/llms/brave
tests/unit/llms/bytez
tests/unit/llms/cerebras
tests/unit/llms/chat
@ -270,6 +269,7 @@ jobs:
tests/unit/llms/gdc
tests/unit/llms/gemini
tests/unit/llms/gigachat
tests/unit/llms/google_pse
tests/unit/llms/github_copilot
tests/unit/llms/gradient_ai
tests/unit/llms/groq
@ -292,6 +292,7 @@ jobs:
tests/unit/llms/manus
tests/unit/llms/meta_llama
tests/unit/llms/microsoft_365_copilot
tests/unit/llms/milvus
tests/unit/llms/minimax
tests/unit/llms/mistral
tests/unit/llms/modelscope
@ -301,6 +302,7 @@ jobs:
tests/unit/llms/nebius
tests/unit/llms/neosantara
tests/unit/llms/nimble
tests/unit/llms/nlp_cloud
tests/unit/llms/novita
tests/unit/llms/nscale
tests/unit/llms/nvidia_nim
@ -311,6 +313,7 @@ jobs:
tests/unit/llms/oobabooga
tests/unit/llms/openai_like
tests/unit/llms/opencode
tests/unit/llms/opensandbox
tests/unit/llms/openrouter
tests/unit/llms/ovhcloud
tests/unit/llms/parallel_ai
@ -347,6 +350,7 @@ jobs:
tests/unit/llms/v0
tests/unit/llms/valkey
tests/unit/llms/vercel_ai_gateway
tests/unit/llms/vllm
tests/unit/llms/volcengine
tests/unit/llms/voyage
tests/unit/llms/wandb
@ -376,7 +380,6 @@ jobs:
- shard: root
test-path: >-
tests/test_litellm/test_*.py
tests/unit/test_*.py
python-version: "3.12"
workers: 4
@ -395,16 +398,6 @@ jobs:
job-timeout-minutes: 60
dist: loadscope
- shard: router-unit-tests
test-path: >-
tests/router_unit_tests
python-version: "3.12"
workers: 0
reruns: 0
timeout-minutes: 10
job-timeout-minutes: 60
dist: loadscope
- shard: endpoints
test-path: >-
tests/unit/a2a_protocol
@ -424,6 +417,7 @@ jobs:
tests/unit/rag
tests/unit/realtime_api
tests/unit/rerank_api
tests/unit/search
tests/unit/vector_stores
tests/unit/videos
python-version: "3.12"
@ -458,6 +452,7 @@ jobs:
tests/unit/repositories
tests/unit/secret_managers
tests/unit/skills/test_skills_main.py
tests/unit/telemetry tests/unit/proxy/telemetry
tests/unit/tracing
python-version: "3.12"
workers: 2
@ -775,7 +770,7 @@ jobs:
- shard: proxy-server-core
test-path: |-
tests/proxy_unit_tests/test_proxy_server_gemini_pass_through.py
tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py
tests/unit/proxy/test__lazy_features.py
tests/unit/proxy/test_aproxy_startup.py
tests/unit/proxy/test_proxy_server.py
@ -806,6 +801,10 @@ jobs:
tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_oauth_passthrough_tools.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_max_concurrent_requests.py
tests/unit/proxy/_experimental/mcp_server/faults/test_list_outcomes.py
tests/unit/proxy/_experimental/mcp_server/outbound_credentials
python-version: "3.12"
workers: 4
@ -814,9 +813,31 @@ jobs:
job-timeout-minutes: 60
dist: loadscope
- shard: mcp-server
test-path: |-
tests/unit/experimental_mcp_client
tests/unit/proxy/_experimental/mcp_server
tests/unit/responses/mcp
--ignore=tests/unit/experimental_mcp_client/test_mcp_client.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_capabilities.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_interactions.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_elicitation_handler.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_operations.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
--ignore=tests/unit/proxy/_experimental/mcp_server/test_db_credentials.py
--ignore=tests/unit/proxy/_experimental/mcp_server/outbound_credentials
python-version: "3.12"
workers: 4
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
dist: loadscope
- shard: custom-logging
test-path: |-
tests/proxy_unit_tests/test_proxy_custom_logger.py
tests/unit/proxy/test_proxy_server_custom_logger.py
tests/unit/proxy/test_custom_callback_input.py
tests/unit/proxy/test_custom_logger_s3_gcs.py
python-version: "3.12"
@ -880,7 +901,7 @@ jobs:
- shard: endpoints-and-responses
test-path: |-
tests/proxy_unit_tests/test_proxy_exception_mapping.py
tests/unit/proxy/test_proxy_server_exception_mapping.py
tests/unit/proxy/lens
tests/unit/proxy/auth/test_models_fallback_endpoint.py
tests/unit/proxy/common_utils/test_check_batch_cost.py
@ -1336,6 +1357,115 @@ jobs:
- name: Check consumer types with the installed SDK
if: steps.changes.outputs.decision != 'skip'
run: uv run --no-sync basedpyright --pythonpath .venv/bin/python --project tests/typing/pyrightconfig.json
redis-compat:
name: redis-py ${{ matrix.redis-py }} compatibility
needs: rust-bridge
if: ${{ !cancelled() }}
permissions:
contents: read
id-token: write
pull-requests: read
runs-on: ubuntu-latest
timeout-minutes: 25
strategy:
fail-fast: false
matrix:
redis-py: ["5.3.1", "6.4.0", "7.4.1", "8.0.1"]
env:
UV_PYTHON: "3.12"
LITELLM_LOCAL_MODEL_COST_MAP: "True"
REDIS_PY_VERSION: ${{ matrix.redis-py }}
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
with:
category: redis-compat
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: ${{ env.UV_PYTHON }}
- 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: ./.github/actions/cache-uv-downloads
- name: Set up the Rust build
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
uses: ./.github/actions/rust-bridge
with:
artifact: ${{ needs.rust-bridge.outputs.artifact }}
- name: Install dependencies with redis-py ${{ matrix.redis-py }}
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
env:
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
run: |
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra caching --extra extra_proxy --extra proxy-runtime --extra utils
else
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra caching --extra extra_proxy --extra proxy-runtime --extra utils --no-install-project
uv pip install --no-deps --python .venv/bin/python rust-bridge-dist/*.whl
cp rust-bridge-dist/litellm/rust_bridge/_native.abi3.so litellm/rust_bridge/_native.abi3.so
fi
uv pip install --python .venv/bin/python "redis==${REDIS_PY_VERSION:?}"
uv run --no-sync python -c 'import os, redis; assert redis.__version__ == os.environ["REDIS_PY_VERSION"], redis.__version__'
- name: Install redis-server 7.2.16
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
run: |
mkdir -p "$RUNNER_TEMP/redis-bin"
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server"
docker rm "$cid"
"$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16'
echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH"
- name: Run Redis compatibility tests
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 15
env:
COVERAGE_CORE: sysmon
run: |
uv run --no-sync pytest \
tests/unit/test_redis.py \
tests/unit/caching/test_redis_connection_pool.py \
tests/unit/caching/test_redis_cluster_cache.py \
tests/unit/caching/test_evicted_client_closer.py \
tests/integration/sdk/test_redis_cluster_iam_auth.py \
--tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \
--cov=./litellm --cov-report=xml:coverage.xml
- name: Upload redis-py ${{ matrix.redis-py }} coverage to Codecov
if: always() && steps.changes.outputs.decision != 'skip' && matrix.redis-py == '5.3.1'
continue-on-error: true
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
with:
use_oidc: true
files: coverage.xml
root_dir: ${{ github.workspace }}
flags: redis-compat
fail_ci_if_error: false
coverage:
name: coverage
needs: unit
@ -1381,10 +1511,20 @@ jobs:
root_dir: ${{ github.workspace }}
flags: unit
fail_ci_if_error: false
- name: Upload MCP server coverage under its own flag
continue-on-error: true
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
with:
use_oidc: true
files: coverage-reports/coverage-mcp-server-${{ github.run_id }}-${{ github.run_attempt }}/coverage.xml
root_dir: ${{ github.workspace }}
flags: mcp-integration
fail_ci_if_error: false
unit-passed:
name: unit passed
permissions: {}
needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat]
needs: [rust-bridge, assert-shard-coverage, unit, ui-unit, docs, helm, openai-sdk-compat, redis-compat]
if: always()
runs-on: ubuntu-latest
timeout-minutes: 2

View file

@ -63,7 +63,7 @@ jobs:
- name: Start the proxy
run: |
nohup uv run --no-sync litellm --config tests/e2e/load/weekly_anomaly_config.yml --port 4000 > proxy.log 2>&1 &
nohup uv run --no-sync litellm --config tests/e2e/perf/weekly_anomaly_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
@ -78,7 +78,7 @@ jobs:
env:
E2E_WEEKLY_ANOMALY: "1"
run: |
uv run --no-sync pytest tests/e2e/load/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
uv run --no-sync pytest tests/e2e/perf/test_weekly_session_anomaly_e2e.py -v --tb=short -rA
- name: Show proxy log on failure
if: failure()

View file

@ -1,613 +0,0 @@
{
"files": [
{
"path": "litellm/_v2/AGENTS.md",
"description": "Conventions for code under litellm/_v2/",
"scope": [
"litellm/_v2/**"
]
},
{
"path": "litellm/_v2/cache/AGENTS.md",
"description": "Conventions for code under litellm/_v2/cache/",
"scope": [
"litellm/_v2/cache/**"
]
},
{
"path": "litellm/proxy/_experimental/mcp_server/AGENTS.md",
"description": "Conventions for code under litellm/proxy/_experimental/mcp_server/",
"scope": [
"litellm/proxy/_experimental/mcp_server/**"
]
},
{
"path": "litellm/rust_bridge/AGENTS.md",
"description": "Conventions for code under litellm/rust_bridge/",
"scope": [
"litellm/rust_bridge/**"
]
},
{
"path": "litellm/tracing/AGENTS.md",
"description": "Conventions for code under litellm/tracing/",
"scope": [
"litellm/tracing/**"
]
},
{
"path": "litellm-rust/AGENTS.md",
"description": "Conventions for code under litellm-rust/",
"scope": [
"litellm-rust/**"
]
},
{
"path": "litellm-rust/crates/cache-response/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/cache-response/",
"scope": [
"litellm-rust/crates/cache-response/**"
]
},
{
"path": "litellm-rust/crates/callbacks-legacy-python/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/callbacks-legacy-python/",
"scope": [
"litellm-rust/crates/callbacks-legacy-python/**"
]
},
{
"path": "litellm-rust/crates/coroutine/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/coroutine/",
"scope": [
"litellm-rust/crates/coroutine/**"
]
},
{
"path": "litellm-rust/crates/gateway/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/gateway/",
"scope": [
"litellm-rust/crates/gateway/**"
]
},
{
"path": "litellm-rust/crates/gateway-auth/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/gateway-auth/",
"scope": [
"litellm-rust/crates/gateway-auth/**"
]
},
{
"path": "litellm-rust/crates/gateway-inference/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/gateway-inference/",
"scope": [
"litellm-rust/crates/gateway-inference/**"
]
},
{
"path": "litellm-rust/crates/gateway-mcp/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/gateway-mcp/",
"scope": [
"litellm-rust/crates/gateway-mcp/**"
]
},
{
"path": "litellm-rust/crates/host/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/host/",
"scope": [
"litellm-rust/crates/host/**"
]
},
{
"path": "litellm-rust/crates/host/src/machine/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/host/src/machine/",
"scope": [
"litellm-rust/crates/host/src/machine/**"
]
},
{
"path": "litellm-rust/crates/host-http/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/host-http/",
"scope": [
"litellm-rust/crates/host-http/**"
]
},
{
"path": "litellm-rust/crates/host-native/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/host-native/",
"scope": [
"litellm-rust/crates/host-native/**"
]
},
{
"path": "litellm-rust/crates/host-python/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/host-python/",
"scope": [
"litellm-rust/crates/host-python/**"
]
},
{
"path": "litellm-rust/crates/http/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/http/",
"scope": [
"litellm-rust/crates/http/**"
]
},
{
"path": "litellm-rust/crates/inference/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/inference/",
"scope": [
"litellm-rust/crates/inference/**"
]
},
{
"path": "litellm-rust/crates/inference-chat/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/inference-chat/",
"scope": [
"litellm-rust/crates/inference-chat/**"
]
},
{
"path": "litellm-rust/crates/inference-messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/inference-messages/",
"scope": [
"litellm-rust/crates/inference-messages/**"
]
},
{
"path": "litellm-rust/crates/inference-ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/inference-ocr/",
"scope": [
"litellm-rust/crates/inference-ocr/**"
]
},
{
"path": "litellm-rust/crates/inference-responses/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/inference-responses/",
"scope": [
"litellm-rust/crates/inference-responses/**"
]
},
{
"path": "litellm-rust/crates/inference-transcription/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/inference-transcription/",
"scope": [
"litellm-rust/crates/inference-transcription/**"
]
},
{
"path": "litellm-rust/crates/llms/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/",
"scope": [
"litellm-rust/crates/llms/**"
]
},
{
"path": "litellm-rust/crates/llms/src/anthropic/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/",
"scope": [
"litellm-rust/crates/llms/src/anthropic/**"
]
},
{
"path": "litellm-rust/crates/llms/src/anthropic/batches/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/batches/",
"scope": [
"litellm-rust/crates/llms/src/anthropic/batches/**"
]
},
{
"path": "litellm-rust/crates/llms/src/anthropic/chat/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/chat/",
"scope": [
"litellm-rust/crates/llms/src/anthropic/chat/**"
]
},
{
"path": "litellm-rust/crates/llms/src/anthropic/count_tokens/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/count_tokens/",
"scope": [
"litellm-rust/crates/llms/src/anthropic/count_tokens/**"
]
},
{
"path": "litellm-rust/crates/llms/src/anthropic/messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/anthropic/messages/",
"scope": [
"litellm-rust/crates/llms/src/anthropic/messages/**"
]
},
{
"path": "litellm-rust/crates/llms/src/aws_textract/ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/aws_textract/ocr/",
"scope": [
"litellm-rust/crates/llms/src/aws_textract/ocr/**"
]
},
{
"path": "litellm-rust/crates/llms/src/azure_ai/messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/azure_ai/messages/",
"scope": [
"litellm-rust/crates/llms/src/azure_ai/messages/**"
]
},
{
"path": "litellm-rust/crates/llms/src/azure_ai/ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/azure_ai/ocr/",
"scope": [
"litellm-rust/crates/llms/src/azure_ai/ocr/**"
]
},
{
"path": "litellm-rust/crates/llms/src/base_llm/messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/base_llm/messages/",
"scope": [
"litellm-rust/crates/llms/src/base_llm/messages/**"
]
},
{
"path": "litellm-rust/crates/llms/src/bedrock/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/",
"scope": [
"litellm-rust/crates/llms/src/bedrock/**"
]
},
{
"path": "litellm-rust/crates/llms/src/bedrock/audio_transcription/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/audio_transcription/",
"scope": [
"litellm-rust/crates/llms/src/bedrock/audio_transcription/**"
]
},
{
"path": "litellm-rust/crates/llms/src/bedrock/chat/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/chat/",
"scope": [
"litellm-rust/crates/llms/src/bedrock/chat/**"
]
},
{
"path": "litellm-rust/crates/llms/src/bedrock/messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/bedrock/messages/",
"scope": [
"litellm-rust/crates/llms/src/bedrock/messages/**"
]
},
{
"path": "litellm-rust/crates/llms/src/cohere/ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/cohere/ocr/",
"scope": [
"litellm-rust/crates/llms/src/cohere/ocr/**"
]
},
{
"path": "litellm-rust/crates/llms/src/deepseek/messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/deepseek/messages/",
"scope": [
"litellm-rust/crates/llms/src/deepseek/messages/**"
]
},
{
"path": "litellm-rust/crates/llms/src/mistral/ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/mistral/ocr/",
"scope": [
"litellm-rust/crates/llms/src/mistral/ocr/**"
]
},
{
"path": "litellm-rust/crates/llms/src/openai/responses/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/openai/responses/",
"scope": [
"litellm-rust/crates/llms/src/openai/responses/**"
]
},
{
"path": "litellm-rust/crates/llms/src/openai_like/chat/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/openai_like/chat/",
"scope": [
"litellm-rust/crates/llms/src/openai_like/chat/**"
]
},
{
"path": "litellm-rust/crates/llms/src/reducto/ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/reducto/ocr/",
"scope": [
"litellm-rust/crates/llms/src/reducto/ocr/**"
]
},
{
"path": "litellm-rust/crates/llms/src/vertex_ai/messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/vertex_ai/messages/",
"scope": [
"litellm-rust/crates/llms/src/vertex_ai/messages/**"
]
},
{
"path": "litellm-rust/crates/llms/src/vertex_ai/ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms/src/vertex_ai/ocr/",
"scope": [
"litellm-rust/crates/llms/src/vertex_ai/ocr/**"
]
},
{
"path": "litellm-rust/crates/llms-types/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/",
"scope": [
"litellm-rust/crates/llms-types/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/formats/audio_transcription/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/audio_transcription/",
"scope": [
"litellm-rust/crates/llms-types/src/formats/audio_transcription/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/formats/batches/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/batches/",
"scope": [
"litellm-rust/crates/llms-types/src/formats/batches/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/formats/chat_completions/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/chat_completions/",
"scope": [
"litellm-rust/crates/llms-types/src/formats/chat_completions/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/formats/messages/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/messages/",
"scope": [
"litellm-rust/crates/llms-types/src/formats/messages/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/formats/ocr/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/ocr/",
"scope": [
"litellm-rust/crates/llms-types/src/formats/ocr/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/formats/responses/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/formats/responses/",
"scope": [
"litellm-rust/crates/llms-types/src/formats/responses/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/providers/anthropic/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/anthropic/",
"scope": [
"litellm-rust/crates/llms-types/src/providers/anthropic/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/providers/bedrock/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/bedrock/",
"scope": [
"litellm-rust/crates/llms-types/src/providers/bedrock/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/providers/minimax/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/minimax/",
"scope": [
"litellm-rust/crates/llms-types/src/providers/minimax/**"
]
},
{
"path": "litellm-rust/crates/llms-types/src/providers/vertex_ai/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/llms-types/src/providers/vertex_ai/",
"scope": [
"litellm-rust/crates/llms-types/src/providers/vertex_ai/**"
]
},
{
"path": "litellm-rust/crates/model-catalog/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/model-catalog/",
"scope": [
"litellm-rust/crates/model-catalog/**"
]
},
{
"path": "litellm-rust/crates/python-bridge/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/python-bridge/",
"scope": [
"litellm-rust/crates/python-bridge/**"
]
},
{
"path": "litellm-rust/crates/python-bridge/src/cache/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/",
"scope": [
"litellm-rust/crates/python-bridge/src/cache/**"
]
},
{
"path": "litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/native/",
"scope": [
"litellm-rust/crates/python-bridge/src/cache/native/**"
]
},
{
"path": "litellm-rust/crates/python-bridge/src/cache/python/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/cache/python/",
"scope": [
"litellm-rust/crates/python-bridge/src/cache/python/**"
]
},
{
"path": "litellm-rust/crates/python-bridge/src/routes/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/python-bridge/src/routes/",
"scope": [
"litellm-rust/crates/python-bridge/src/routes/**"
]
},
{
"path": "litellm-rust/crates/python-compat/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/python-compat/",
"scope": [
"litellm-rust/crates/python-compat/**"
]
},
{
"path": "litellm-rust/crates/secrets-aws/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/secrets-aws/",
"scope": [
"litellm-rust/crates/secrets-aws/**"
]
},
{
"path": "litellm-rust/crates/secrets-azure/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/secrets-azure/",
"scope": [
"litellm-rust/crates/secrets-azure/**"
]
},
{
"path": "litellm-rust/crates/secrets-cyberark/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/secrets-cyberark/",
"scope": [
"litellm-rust/crates/secrets-cyberark/**"
]
},
{
"path": "litellm-rust/crates/secrets-google/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/secrets-google/",
"scope": [
"litellm-rust/crates/secrets-google/**"
]
},
{
"path": "litellm-rust/crates/secrets-hashicorp/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/secrets-hashicorp/",
"scope": [
"litellm-rust/crates/secrets-hashicorp/**"
]
},
{
"path": "litellm-rust/crates/storage-clickhouse/AGENTS.md",
"description": "Conventions for code under litellm-rust/crates/storage-clickhouse/",
"scope": [
"litellm-rust/crates/storage-clickhouse/**"
]
},
{
"path": "litellm-rust/docs/adrs/AGENTS.md",
"description": "Conventions for code under litellm-rust/docs/adrs/",
"scope": [
"litellm-rust/docs/adrs/**"
]
},
{
"path": "tests/AGENTS.md",
"description": "Conventions for code under tests/",
"scope": [
"tests/**"
]
},
{
"path": "tests/e2e/AGENTS.md",
"description": "Conventions for code under tests/e2e/",
"scope": [
"tests/e2e/**"
]
},
{
"path": "tests/e2e_harness/AGENTS.md",
"description": "Conventions for code under tests/e2e_harness/",
"scope": [
"tests/e2e_harness/**"
]
},
{
"path": "tests/integration/AGENTS.md",
"description": "Conventions for code under tests/integration/",
"scope": [
"tests/integration/**"
]
},
{
"path": "tests/integration/translation/AGENTS.md",
"description": "Conventions for code under tests/integration/translation/",
"scope": [
"tests/integration/translation/**"
]
},
{
"path": "tests/rust-python-harness/AGENTS.md",
"description": "Conventions for code under tests/rust-python-harness/",
"scope": [
"tests/rust-python-harness/**"
]
},
{
"path": "tests/rust-python-harness/strategies/e2e_parity/AGENTS.md",
"description": "Conventions for code under tests/rust-python-harness/strategies/e2e_parity/",
"scope": [
"tests/rust-python-harness/strategies/e2e_parity/**"
]
},
{
"path": "tests/rust-python-harness/strategies/trace_parity/AGENTS.md",
"description": "Conventions for code under tests/rust-python-harness/strategies/trace_parity/",
"scope": [
"tests/rust-python-harness/strategies/trace_parity/**"
]
},
{
"path": "tests/rust-python-harness/strategies/unit_tests_parity/AGENTS.md",
"description": "Conventions for code under tests/rust-python-harness/strategies/unit_tests_parity/",
"scope": [
"tests/rust-python-harness/strategies/unit_tests_parity/**"
]
},
{
"path": "tests/rust-python-harness/strategies/unit_tests_rust/AGENTS.md",
"description": "Conventions for code under tests/rust-python-harness/strategies/unit_tests_rust/",
"scope": [
"tests/rust-python-harness/strategies/unit_tests_rust/**"
]
},
{
"path": "tests/test_litellm_rust/AGENTS.md",
"description": "Conventions for code under tests/test_litellm_rust/",
"scope": [
"tests/test_litellm_rust/**"
]
},
{
"path": "tests/unit/AGENTS.md",
"description": "Conventions for code under tests/unit/",
"scope": [
"tests/unit/**"
]
},
{
"path": "tests/unit/rust_bridge/AGENTS.md",
"description": "Conventions for code under tests/unit/rust_bridge/",
"scope": [
"tests/unit/rust_bridge/**"
]
},
{
"path": "ui/litellm-dashboard/AGENTS.md",
"description": "Conventions for code under ui/litellm-dashboard/",
"scope": [
"ui/litellm-dashboard/**"
]
},
{
"path": "ui/litellm-dashboard/src/components/chat/AGENTS.md",
"description": "Conventions for code under ui/litellm-dashboard/src/components/chat/",
"scope": [
"ui/litellm-dashboard/src/components/chat/**"
]
}
]
}

View file

@ -23,7 +23,7 @@ When adding new features, add meaningful tests. Don't add tests that don't check
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
Never test structure or implementation details of code, only function of it. A test should catch a change that breaks behavior by accident. If the only way it can fail is someone deliberately undoing the change, don't write it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent". Assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
@ -130,4 +130,4 @@ Before implementing:
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review. Never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage and has acceptable reviews from the review bots on the PR, such as Veria and Bugbot. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review. Never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or raise `ANY_CAPS` to satisfy a check

View file

@ -39,7 +39,7 @@ help:
@echo " make check-circular-imports - Check for circular imports"
@echo " make check-import-safety - Check import safety"
@echo " make test - Run all tests"
@echo " make test-unit - Run unit tests (tests/unit and tests/test_litellm)"
@echo " make test-unit - Run unit tests (tests/unit)"
@echo " make test-unit-llms - Run LLM provider tests (~225 files)"
@echo " make test-unit-proxy-guardrails - Run proxy guardrails+mgmt tests (~51 files)"
@echo " make test-unit-proxy-core - Run proxy auth+client+db+hooks tests (~52 files)"
@ -300,7 +300,7 @@ test: install-test-deps
$(UV_RUN) pytest tests/
test-unit: install-test-deps
$(UV_RUN) pytest tests/unit tests/test_litellm -x -vv -n 4
$(UV_RUN) pytest tests/unit -x -vv -n 4
# Matrix test targets (matching CI workflow groups)
test-unit-llms: install-test-deps
@ -328,7 +328,7 @@ test-unit-other: install-test-deps
$(UV_RUN) pytest tests/unit/caching tests/unit/responses tests/unit/secret_managers tests/unit/vector_stores tests/unit/a2a_protocol tests/unit/completion_extras tests/unit/containers tests/unit/enterprise tests/unit/experimental_mcp_client tests/unit/google_genai tests/unit/images tests/unit/interactions tests/unit/router_strategy tests/unit/router_utils tests/unit/types --tb=short -vv -n 4 --durations=20
test-unit-root: install-test-deps
$(UV_RUN) pytest tests/unit/test_*.py tests/test_litellm/test_*.py --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/test_*.py --tb=short -vv -n 4 --durations=20
# Proxy unit tests (tests/unit/proxy split alphabetically)
test-proxy-unit-a: install-test-deps
@ -338,7 +338,7 @@ test-proxy-unit-b: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/test_[p-z]*.py tests/unit/skills --tb=short -vv -n 2 --durations=20
test-integration: install-test-deps
$(UV_RUN) pytest tests/ -k "not test_litellm"
$(UV_RUN) pytest tests/
test-unit-helm: install-helm-unittest
helm unittest -f 'tests/*.yaml' helm/litellm-helm

View file

@ -9,7 +9,7 @@
<a href="https://render.com/deploy?repo=https://github.com/BerriAI/litellm" target="_blank" rel="nofollow"><img src="https://render.com/images/deploy-to-render-button.svg" alt="Deploy to Render" height="40"></a>
<a href="https://railway.com/deploy/RhvhdC?referralCode=7mRv9K&utm_medium=integration&utm_source=template&utm_campaign=generic"><img src="https://railway.com/button.svg" alt="Deploy on Railway" height="40"></a>
<a href="https://console.aws.amazon.com/cloudshell/home" target="_blank" rel="nofollow"><img src="./.github/deploy-on-aws.png" alt="Deploy on AWS" height="40"></a>
<a href="https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true" target="_blank" rel="nofollow"><img src="./.github/deploy-on-gcp.png" alt="Deploy on GCP" height="40"></a>
<a href="https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Fterraform-google-litellm&cloudshell_git_branch=v1.105.0&cloudshell_workspace=examples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true" target="_blank" rel="nofollow"><img src="./.github/deploy-on-gcp.png" alt="Deploy on GCP" height="40"></a>
</p>
</p>
<h4 align="center"><a href="https://docs.litellm.ai/docs/simple_proxy" target="_blank">LiteLLM Proxy Server (AI Gateway)</a> | <a href="https://docs.litellm.ai/docs/enterprise#hosted-litellm-proxy" target="_blank"> Hosted Proxy</a> | <a href="https://litellm.ai/enterprise"target="_blank">Enterprise Tier</a> | <a href="https://www.litellm.ai/ai-gateway" target="_blank">Website</a></h4>
@ -342,6 +342,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | |
| [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | |
| [CompactifAI (`compactifai`)](https://docs.litellm.ai/docs/providers/compactifai) | ✅ | ✅ | ✅ | | | | | | | |
| [CoralBricks (`coralbricks`)](https://docs.litellm.ai/docs/providers/coralbricks) | ✅ | ✅ | ✅ | | | | | | | |
| [Custom (`custom`)](https://docs.litellm.ai/docs/providers/custom_llm_server) | ✅ | ✅ | ✅ | | | | | | | |
| [Custom OpenAI (`custom_openai`)](https://docs.litellm.ai/docs/providers/openai_compatible) | ✅ | ✅ | ✅ | | | ✅ | ✅ | ✅ | ✅ | |
| [Dashscope (`dashscope`)](https://docs.litellm.ai/docs/providers/dashscope) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ |
@ -526,9 +527,9 @@ Provider API keys live in AWS Secrets Manager; reference ARNs via `gateway_extra
#### GCP — Cloud Run + Cloud SQL + Memorystore + HTTPS LB
[![Open in Cloud Shell](https://gstatic.com/cloudssh/images/open-btn.png)](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true)
[![Open in Cloud Shell](https://gstatic.com/cloudssh/images/open-btn.png)](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Fterraform-google-litellm&cloudshell_git_branch=v1.105.0&cloudshell_workspace=examples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true)
Real 1-click. Opens Cloud Shell, clones this repo, and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md) — pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply.
Real 1-click. Opens Cloud Shell, clones the [terraform-google-litellm](https://github.com/BerriAI/terraform-google-litellm) module mirror at the latest stable release (a small download instead of the full litellm repo), and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md). Pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply
[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/google/latest)

View file

@ -32,6 +32,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/token",
"/onboarding/",
"/audit",
"/telemetry/",
"/oauth/",
"/invitation/",
"/jwt/",

View file

@ -56,7 +56,7 @@ def _rendered_schema(cost_map: CostMap) -> str:
def _evaluation_output_price_failure(
key: str, entry: object, evaluation_output_price_sources: Mapping[str, str]
) -> str | None:
if not isinstance(entry, dict) or entry.get("mode") != "evaluation":
if not isinstance(entry, dict) or entry.get("mode") not in ("evaluation", "decisions"):
return None
output_cost: Final = entry.get("output_cost_per_token")
if (
@ -67,7 +67,7 @@ def _evaluation_output_price_failure(
):
return None
return (
f"{key}: mode evaluation rows bill input only; output_cost_per_token must be 0 unless the vendor "
f"{key}: mode evaluation and decisions rows bill input only; output_cost_per_token must be 0 unless the vendor "
"prices output tokens for this model (add the source to EVALUATION_OUTPUT_PRICE_SOURCES)"
)

View file

@ -322,7 +322,12 @@ def create_migration(
# Write the SQL to migration.sql
migration_file = migration_dir / "migration.sql"
migration_file.write_text(diff_sql)
migration_sql: Final = (
diff_sql
if "CONCURRENTLY" in diff_sql.upper()
else "SET LOCAL lock_timeout = '10s';\n\n" + diff_sql
)
migration_file.write_text(migration_sql)
print(f"Created migration in {migration_dir}")
return True

View file

@ -71,6 +71,8 @@ flag_management:
carryforward: false
- name: lens-python-310
carryforward: false
- name: lens-postgres
carryforward: false
- name: proxy-db-key-generation
carryforward: false
- name: proxy-db-auth-checks

View file

@ -55,6 +55,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_litellm_call_id_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED

View file

@ -24,7 +24,7 @@ from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
from pydantic import BaseModel, ValidationError
from typing_extensions import ReadOnly, Unpack
import litellm
@ -54,6 +54,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
can_access_resource,
resolve_resource_owner_id,
)
from litellm.llms.xai.batches.transformation import is_xai_batch_results_id
from litellm.proxy._types import (
CallTypes,
LiteLLM_ManagedFileTable,
@ -102,6 +103,7 @@ from litellm.types.utils import (
CallTypesLiteral,
LiteLLMBatch,
LiteLLMFineTuningJob,
LlmProviders,
LLMResponseTypes,
SpecialEnums,
)
@ -194,6 +196,32 @@ def _is_transient_file_retrieve_error(error: Exception) -> bool:
return isinstance(error, (httpx.TransportError, APIConnectionError, asyncio.TimeoutError))
_MEASURED_SIZE_SAVE_ATTEMPTS: Final = 3
def _stored_json_text(raw_file_object: object) -> str:
"""JSON text of a stored file object exactly as read, for a write that must not race a newer one."""
if isinstance(raw_file_object, BaseModel):
return raw_file_object.model_dump_json()
return json.dumps(raw_file_object)
def _with_measured_size(file_object: OpenAIFileObject, size_bytes: int | None) -> OpenAIFileObject:
"""Fill a file object whose provider reported no size (0 bytes) with a size measured elsewhere."""
if not size_bytes or file_object.bytes:
return file_object
return file_object.model_copy(update={"bytes": size_bytes})
def _over_stored_object(file_object: OpenAIFileObject, stored_object: OpenAIFileObject | None) -> OpenAIFileObject:
"""Refreshed details for a stored row, keeping its size when none is reported and its stored created_at."""
if stored_object is None:
return file_object
return _with_measured_size(file_object, stored_object.bytes).model_copy(
update={"created_at": stored_object.created_at}
)
def _proxy_llm_router() -> Router | None:
import litellm.proxy.proxy_server as proxy_server_module
@ -252,6 +280,23 @@ def _has_provider_file_retrieve_route(
return bool(router_credentials) or _model_name_file_retrieve_provider(model_name) is not None
def _is_xai_batch_output(
*,
llm_router: Router | None,
model_id: str | None,
model_name: str | None,
provider_file_id: str,
) -> bool:
"""Whether a batch output row holds an xAI batch's results, from its batch id and its deployment's provider."""
if not is_xai_batch_results_id(provider_file_id):
return False
router_credentials: Final = _provider_file_retrieve_credentials(llm_router=llm_router, model_id=model_id)
provider: Final = (router_credentials or {}).get("custom_llm_provider") or _model_name_file_retrieve_provider(
model_name
)
return provider == LlmProviders.XAI.value
class _ManagedFileRow(Protocol):
unified_file_id: str
file_object: OpenAIFileObject
@ -845,6 +890,57 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
litellm_parent_otel_span=None,
)
async def _save_xai_refreshed_file_object(
self,
stored: LiteLLM_ManagedFileTable,
file_object: OpenAIFileObject,
) -> OpenAIFileObject | None:
"""Save an xAI batch output's refreshed details, keeping a known size, only over the row as read; drop the cached copy if a newer write won."""
stored_object: Final = _parse_managed_file_object(stored.file_object, stored.unified_file_id)
sized_file_object: Final = _over_stored_object(file_object, stored_object)
if not await ManagedFileRepository(self.prisma_client).update_file_object(
stored.unified_file_id,
sized_file_object,
if_file_object=_stored_json_text(stored.file_object) if stored.file_object is not None else None,
):
await self.internal_usage_cache.async_set_cache(
key=stored.unified_file_id, value=None, litellm_parent_otel_span=None
)
return None
await self.internal_usage_cache.async_set_cache(
key=stored.unified_file_id,
value=stored.model_copy(update={"file_object": sized_file_object}).model_dump(),
litellm_parent_otel_span=None,
)
return sized_file_object
async def _save_batch_output_file_object(
self,
stored: LiteLLM_ManagedFileTable,
file_object: OpenAIFileObject,
size_bytes: int | None,
) -> None:
"""Save an xAI batch output's details; after losing to a refresh, size the refreshed details instead of replacing them."""
current_row: LiteLLM_ManagedFileTable | None = stored # rebind-ok: the row as last read
details: OpenAIFileObject = file_object # rebind-ok: a concurrent refresh's details replace the job's copy
for _ in range(_MEASURED_SIZE_SAVE_ATTEMPTS):
if current_row is None:
return
saved = await self._save_xai_refreshed_file_object( # rebind-ok: per attempt
current_row, _with_measured_size(details, size_bytes)
)
if saved is not None or not size_bytes:
return
latest_row = await ManagedFileRepository( # rebind-ok: per attempt
self.prisma_client, use_writer=True
).table.find_first(where={"unified_file_id": stored.unified_file_id})
current_row = (
None if latest_row is None else LiteLLM_ManagedFileTable.model_validate(latest_row.model_dump())
)
latest_object = None if current_row is None else current_row.file_object # rebind-ok: per attempt
if latest_object is not None and not latest_object.litellm_details_fallback:
details = latest_object
async def store_batch_output_file(
self,
*,
@ -885,6 +981,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
return
if stored_object is not None and not stored_object.litellm_details_fallback:
if (
stored_file is not None
and size_bytes
and not stored_object.bytes
and _is_xai_batch_output(
llm_router=_proxy_llm_router(),
model_id=model_id,
model_name=model_name,
provider_file_id=provider_file_id,
)
):
await self._save_batch_output_file_object(stored_file, stored_object, size_bytes)
return
fallback_written_recently: Final = (
@ -924,13 +1032,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
)
xai_batch_output: Final = _is_xai_batch_output(
llm_router=_proxy_llm_router(),
model_id=model_id,
model_name=model_name,
provider_file_id=provider_file_id,
)
if stored_file is not None and xai_batch_output:
await self._save_batch_output_file_object(stored_file, file_object, size_bytes)
return
if stored_file is not None:
await self._save_refreshed_file_object(stored_file, file_object)
return
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=file_object,
file_object=_with_measured_size(file_object, size_bytes) if xai_batch_output else file_object,
litellm_parent_otel_span=litellm_parent_otel_span,
model_mappings={model_id: provider_file_id} if model_id else {},
user_api_key_dict=owner,
@ -1162,7 +1279,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
return False
async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None:
async def check_file_ids_access(self, file_ids: Sequence[str], user_api_key_dict: UserAPIKeyAuth) -> None:
"""
Check if the user has access to a list of file IDs.
Only checks managed (unified) file IDs.
@ -1269,6 +1386,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
elif call_type == CallTypes.acreate_batch.value:
input_file_id = cast(Optional[str], data.get("input_file_id"))
if input_file_id:
await self.check_file_ids_access((input_file_id,), user_api_key_dict)
model_file_id_mapping = await self.get_model_file_id_mapping(
[input_file_id], user_api_key_dict.parent_otel_span
)
@ -1315,11 +1433,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
elif retrieve_object_id and accessor_key == "batch_id":
await self.enforce_batch_object_access(retrieve_object_id, user_api_key_dict)
elif call_type == CallTypes.acreate_fine_tuning_job.value:
input_file_id = cast(Optional[str], data.get("training_file"))
if input_file_id:
model_file_id_mapping = await self.get_model_file_id_mapping(
[input_file_id], user_api_key_dict.parent_otel_span
)
fine_tuning_file_ids: Final = tuple(
file_id
for file_id in (data.get("training_file"), data.get("validation_file"))
if isinstance(file_id, str) and file_id
)
await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict)
return data
@ -1941,8 +2060,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
if refreshed_file_object is None:
return _public_file_object(file_object, file_id)
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
return _public_file_object(refreshed_file_object, file_id)
if not _is_xai_batch_output(
llm_router=llm_router if llm_router is not None else _proxy_llm_router(),
model_id=model_id,
model_name=model_id,
provider_file_id=provider_file_id,
):
await self._save_refreshed_file_object(stored_file_object, refreshed_file_object)
return _public_file_object(refreshed_file_object, file_id)
saved_file_object: Final = await self._save_xai_refreshed_file_object(
stored_file_object, refreshed_file_object
)
return _public_file_object(
saved_file_object or _over_stored_object(refreshed_file_object, file_object), file_id
)
except Exception as error:
verbose_logger.warning(
"Failed to refresh batch file object for "

View file

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

View file

@ -121,6 +121,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
# Realtime / streaming
"/v1/realtime",
"/realtime",
"/v1/live/sessions",
"/live/sessions",
# Health & ops
"/health",
"/metrics",

View file

@ -440,10 +440,10 @@ db:
# Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks
# the pod NotReady and blocks only until in-flight requests actually finish
# (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always
# waiting the worst-case duration. The drain runs once (the preStop hook and
# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few
# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before
# SIGKILL.
# waiting the worst-case duration. After SIGTERM the server waits up to
# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels
# them, so set terminationGracePeriodSeconds a few seconds above twice
# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL.
#
# /health/drain is off by default; enable it with
# general_settings.enable_drain_endpoint: true. The kubelet calls preStop

View file

@ -5,3 +5,21 @@ This chart deploys separate gateway, backend and UI services. Configure Lens thr
Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for the complete values and verification flow. It also links the GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
For an ingress with separate component hostnames, set all three `ingress.componentHosts` values, leave `ingress.host` and `ingress.extraPaths` empty, configure `ui.backendUrl` as the backend's public origin and `backend.gatewayUrl` as the gateway's public origin, and set `ui.publicUrl` to the browser-facing UI origin. The UI Deployment receives `LITELLM_BACKEND_URL` from `ui.backendUrl` only when `ingress.componentHosts` is configured.
```yaml
ingress:
enabled: true
host: ""
componentHosts:
ui: ui.example.com
backend: backend.example.com
gateway: gateway.example.com
extraPaths: []
ui:
publicUrl: https://ui.example.com
backendUrl: https://backend.example.com
backend:
gatewayUrl: https://gateway.example.com
```

View file

@ -22,12 +22,14 @@ Reminders:
- database.writer.{host,port,dbname} (writer connection pieces)
- database.writer.passwordSecret.{name,usernameKey,passwordKey}
(Secret holding the writer DB username + password)
- database.writer.passwordSecret.{hostKey,urlKey} (optional, read writer host or full URL from that Secret)
- database.writer.useIAMAuth: true (optional — chart sets IAM_TOKEN_DB_AUTH=true and
omits DATABASE_PASSWORD / DATABASE_URL so the proxy
mints the URL from an IAM token at startup)
- database.reader.host (optional — enables read-replica routing; reader
.passwordSecret.name is required when set, unless
.useIAMAuth is true)
- database.reader.host / database.reader.passwordSecret.{hostKey,urlKey}
(optional, enables read-replica routing. The Secret
name is required for hostKey, urlKey, and for host
with password auth)
- database.reader.useIAMAuth: true (optional, requires database.writer.useIAMAuth: true —
chart emits DATABASE_*_READ_REPLICA env vars and
omits DATABASE_PASSWORD_READ_REPLICA /
@ -35,6 +37,8 @@ Reminders:
reader URL from an IAM token at startup)
- redis.passwordSecret.name (optional — set when redis.host is provided and the
cache requires auth)
- redis.passwordSecret.{hostKey,urlKey} (optional, read Redis host or full URL from that Secret)
(the Secret name is required when either key is set)
- redis.cluster: true (optional — chart sets REDIS_CLUSTER_NODES from
redis.host / redis.port so the proxy's Cache()
constructs a RedisClusterCache; the cluster client

View file

@ -27,6 +27,30 @@ Common naming + label helpers shared by gateway, backend, and ui templates.
{{- printf "%s-ui" (include "litellm.fullname" .) | trunc 63 | trimSuffix "-" -}}
{{- end -}}
{{- define "litellm.ui.publicOrigin" -}}
{{- $origin := .Values.ui.publicUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "ui.publicUrl must be an origin like https://ui.example.com, without a path" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.ui.backendOrigin" -}}
{{- $origin := .Values.ui.backendUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.backend.gatewayOrigin" -}}
{{- $origin := .Values.backend.gatewayUrl | trim | trimSuffix "/" -}}
{{- if not (regexMatch "^https?://[^/\\s]+$" $origin) -}}
{{- fail "backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set" -}}
{{- end -}}
{{- $origin -}}
{{- end -}}
{{- define "litellm.commonLabels" -}}
app.kubernetes.io/name: {{ include "litellm.name" . }}
app.kubernetes.io/instance: {{ .Release.Name }}
@ -202,7 +226,9 @@ the chart context (needed for .Values), `component` selects which component's
Sensitive values (master key, DB username + password, Redis password) come
only from referenced Secrets; the chart never accepts inline values for them.
The chart never assembles DATABASE_URL itself. It emits only the discrete
The chart never assembles DATABASE_URL itself. When
database.writer.passwordSecret.urlKey is set, the chart passes the operator's
URL through from the Secret. Otherwise it emits only the discrete
DATABASE_HOST/PORT/USER/NAME/SCHEMA (+ DATABASE_PASSWORD for password auth)
vars; the proxy's entrypoint (DatabaseURLSettings in
litellm/proxy/db/db_url_settings.py) builds the URL from them and
@ -218,11 +244,11 @@ 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` (or
The read replica is opt-in via `database.reader.host`, `hostKey`, or `urlKey`.
The chart emits DATABASE_URL_READ_REPLICA alone for urlKey, or discrete
DATABASE_*_READ_REPLICA vars for host and hostKey configurations, plus
DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for 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 token auth only takes effect when the writer uses the same
@ -242,8 +268,36 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
value: {{ $component.logLevel | quote }}
{{- end }}
{{- with $root.Values.database.writer }}
{{- $hostKey := .passwordSecret.hostKey }}
{{- $urlKey := .passwordSecret.urlKey }}
{{- if and .host $hostKey }}
{{- fail "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or .host $hostKey .dbname .schema .sslMode .sslRootCert .useIAMAuth .useAzureEntraAuth) }}
{{- fail "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth" }}
{{- end }}
{{- if and (not .host) (not $hostKey) (not $urlKey) }}
{{- fail "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not .passwordSecret.name) }}
{{- fail "database.writer.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $urlKey }}
- name: DATABASE_URL
valueFrom:
secretKeyRef:
name: {{ .passwordSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else }}
- name: DATABASE_HOST
value: {{ required "database.writer.host is required" .host | quote }}
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ .passwordSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ .host | quote }}
{{- end }}
- name: DATABASE_PORT
value: {{ .port | default 5432 | quote }}
- name: DATABASE_USER
@ -282,8 +336,29 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
key: {{ .passwordSecret.passwordKey | default "password" }}
{{- end }}
{{- end }}
{{- end }}
{{- with $root.Values.database.reader }}
{{- if .host }}
{{- $readerSecret := .passwordSecret | default dict }}
{{- $hostKey := $readerSecret.hostKey }}
{{- $urlKey := $readerSecret.urlKey }}
{{- $readerEnabled := or .host $hostKey $urlKey }}
{{- if and .host $hostKey }}
{{- fail "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or .host $hostKey .dbname .schema .useIAMAuth .useAzureEntraAuth) }}
{{- fail "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not $readerSecret.name) }}
{{- fail "database.reader.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $readerEnabled }}
{{- if $urlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: {{ $readerSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else }}
{{- 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 }}
@ -291,11 +366,18 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
{{- 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
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ $readerSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ .host | quote }}
{{- end }}
- name: DATABASE_PORT_READ_REPLICA
value: {{ .port | default 5432 | quote }}
- name: DATABASE_NAME_READ_REPLICA
value: {{ required "database.reader.dbname is required when database.reader.host is set" .dbname | quote }}
value: {{ required "database.reader.dbname is required when database.reader.host, hostKey, or urlKey is set" .dbname | quote }}
{{- if .schema }}
- name: DATABASE_SCHEMA_READ_REPLICA
value: {{ .schema | quote }}
@ -325,6 +407,7 @@ IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
{{- end }}
{{- end }}
{{- end }}
{{- end }}
{{/*
The migrations Job (helm.sh/hook: pre-upgrade) is the single owner of
`prisma migrate deploy`. Without this, every gateway/backend pod also runs
@ -341,9 +424,34 @@ harmless no-op for the Job and authoritative for the app pods.
tracking, pod lock manager) via its REDIS_* env fallback. An explicit
`general_settings.coordination_redis` block in proxy_config takes
precedence over anything emitted here. */}}
{{- if $root.Values.redis.host }}
{{- $redisSecret := $root.Values.redis.passwordSecret | default dict }}
{{- $hostKey := $redisSecret.hostKey }}
{{- $urlKey := $redisSecret.urlKey }}
{{- if and $root.Values.redis.host $hostKey }}
{{- fail "redis.host and redis.passwordSecret.hostKey cannot both be set" }}
{{- end }}
{{- if and $urlKey (or $root.Values.redis.host $hostKey $root.Values.redis.cluster) }}
{{- fail "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster" }}
{{- end }}
{{- if and (or $hostKey $urlKey) (not $redisSecret.name) }}
{{- fail "redis.passwordSecret.name is required when hostKey or urlKey is set" }}
{{- end }}
{{- if $urlKey }}
- name: REDIS_URL
valueFrom:
secretKeyRef:
name: {{ $redisSecret.name | quote }}
key: {{ $urlKey | quote }}
{{- else if or $root.Values.redis.host $hostKey }}
- name: REDIS_HOST
{{- if $hostKey }}
valueFrom:
secretKeyRef:
name: {{ $redisSecret.name | quote }}
key: {{ $hostKey | quote }}
{{- else }}
value: {{ $root.Values.redis.host | quote }}
{{- end }}
- name: REDIS_PORT
value: {{ $root.Values.redis.port | quote }}
{{- if $root.Values.redis.passwordSecret.name }}
@ -360,7 +468,11 @@ harmless no-op for the Job and authoritative for the app pods.
configured endpoint; the cluster client discovers the remaining nodes from
CLUSTER SLOTS at startup. */}}
- name: REDIS_CLUSTER_NODES
{{- if $hostKey }}
value: {{ printf "[{\"host\":\"$(REDIS_HOST)\",\"port\":%v}]" (int $root.Values.redis.port) | quote }}
{{- else }}
value: {{ printf "[{\"host\":%q,\"port\":%v}]" $root.Values.redis.host (int $root.Values.redis.port) | quote }}
{{- end }}
{{- end }}
{{- end }}
{{- with $component.extraEnv }}

View file

@ -1,4 +1,6 @@
{{- if .Values.backend.enabled }}
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
apiVersion: apps/v1
kind: Deployment
metadata:
@ -58,7 +60,17 @@ spec:
protocol: TCP
env:
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
{{- if .Values.ui.publicUrl }}
- name: LITELLM_CORS_ORIGINS
value: {{ include "litellm.ui.publicOrigin" . | quote }}
- name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
{{- end }}
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
{{- if .Values.backend.gatewayUrl }}
- name: LITELLM_GATEWAY_URL
value: {{ if $hasComponentHosts }}{{ include "litellm.backend.gatewayOrigin" . | quote }}{{ else }}{{ .Values.backend.gatewayUrl | quote }}{{ end }}
{{- end }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml

View file

@ -56,6 +56,10 @@ spec:
protocol: TCP
env:
{{- include "litellm.lensConnectionEnv" . | nindent 12 }}
{{- if .Values.ui.publicUrl }}
- name: LITELLM_CORS_ORIGINS
value: {{ include "litellm.ui.publicOrigin" . | quote }}
{{- end }}
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.gateway) | nindent 12 }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH

View file

@ -1,3 +1,36 @@
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
{{- if $hasComponentHosts }}
{{- if not (and $componentHosts.ui $componentHosts.backend $componentHosts.gateway) }}
{{- fail "ingress.componentHosts.ui, backend, and gateway must all be set together" }}
{{- end }}
{{- if not .Values.gateway.enabled }}
{{- fail "ingress.componentHosts requires gateway.enabled=true" }}
{{- end }}
{{- if not .Values.backend.enabled }}
{{- fail "ingress.componentHosts requires backend.enabled=true" }}
{{- end }}
{{- if not .Values.ui.enabled }}
{{- fail "ingress.componentHosts requires ui.enabled=true" }}
{{- end }}
{{- if .Values.ingress.host }}
{{- fail "ingress.host must be empty when ingress.componentHosts is configured" }}
{{- end }}
{{- if .Values.ingress.extraPaths }}
{{- fail "ingress.extraPaths must be empty when ingress.componentHosts is configured" }}
{{- end }}
{{- if not (trim .Values.ui.backendUrl) }}
{{- fail "ui.backendUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- $_ := include "litellm.ui.backendOrigin" . }}
{{- if not (trim .Values.backend.gatewayUrl) }}
{{- fail "backend.gatewayUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- $_ := include "litellm.backend.gatewayOrigin" . }}
{{- if not (trim .Values.ui.publicUrl) }}
{{- fail "ui.publicUrl is required when ingress.componentHosts is configured" }}
{{- end }}
{{- end }}
{{- if .Values.ingress.enabled -}}
{{- $gatewayName := include "litellm.gateway.fullname" . -}}
{{- $backendName := include "litellm.backend.fullname" . -}}
@ -90,6 +123,66 @@
adding to it.
*/}}
{{- $builtinPathKeys := list "/test|Exact" "/debug/memory/summary|Exact" "/|Prefix" -}}
{{- if $hasComponentHosts }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
name: {{ include "litellm.fullname" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
{{- with .Values.ingress.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
spec:
{{- with .Values.ingress.className }}
ingressClassName: {{ . | quote }}
{{- end }}
{{- with .Values.ingress.tls }}
tls:
{{- toYaml . | nindent 4 }}
{{- end }}
rules:
- host: {{ $componentHosts.ui | quote }}
http:
paths:
- path: /
pathType: Prefix
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
- host: {{ $componentHosts.backend | quote }}
http:
paths:
- path: /
pathType: Prefix
backend:
service:
name: {{ $backendName }}
port:
number: {{ $backendPort }}
- host: {{ $componentHosts.gateway | quote }}
http:
paths:
{{- if or (eq (include "litellm.lens.mode" .) "bundled") (and (eq (include "litellm.lens.mode" .) "external") .Values.lensWorker.externalServiceName) }}
- path: /lens-ingest
pathType: Prefix
backend:
service:
name: {{ .Values.lensWorker.externalServiceName | default (printf "%s-lens-worker" (include "litellm.fullname" .)) }}
port:
number: {{ .Values.lensWorker.service.port }}
{{- end }}
- path: /
pathType: Prefix
backend:
service:
name: {{ $gatewayName }}
port:
number: {{ $gatewayPort }}
{{- else }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
@ -217,3 +310,4 @@ spec:
port:
number: {{ $backendPort }}
{{- end }}
{{- end }}

View file

@ -1,3 +1,5 @@
{{- $componentHosts := .Values.ingress.componentHosts | default dict -}}
{{- $hasComponentHosts := or $componentHosts.ui $componentHosts.backend $componentHosts.gateway -}}
{{- if .Values.ui.enabled }}
apiVersion: apps/v1
kind: Deployment
@ -56,9 +58,9 @@ spec:
- name: LITELLM_LOG
value: {{ .Values.ui.logLevel | quote }}
{{- end }}
{{- if .Values.ui.backendUrl }}
{{- if and .Values.ui.backendUrl $hasComponentHosts }}
- name: LITELLM_BACKEND_URL
value: {{ .Values.ui.backendUrl | quote }}
value: {{ include "litellm.ui.backendOrigin" . | quote }}
{{- end }}
{{- with .Values.ui.extraEnv }}
{{- toYaml . | nindent 12 }}

View file

@ -0,0 +1,441 @@
suite: ingress component hosts
templates:
- ingress.yaml
- backend/deployment.yaml
- gateway/deployment.yaml
- gateway/configmap.yaml
- ui/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: routes each hostname to its component and injects the external URLs
template: ingress.yaml
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
lensWorker.mode: disabled
ingress.tls:
- hosts:
- ui.example.com
- backend.example.com
- gateway.example.com
secretName: component-hosts
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- equal:
path: spec.rules[0].host
value: ui.example.com
- equal:
path: spec.tls[0]
value:
hosts:
- ui.example.com
- backend.example.com
- gateway.example.com
secretName: component-hosts
- equal:
path: spec.rules[0].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- equal:
path: spec.rules[1].host
value: backend.example.com
- equal:
path: spec.rules[1].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-backend
port:
number: 4001
- equal:
path: spec.rules[2].host
value: gateway.example.com
- equal:
path: spec.rules[2].http.paths
value:
- path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- it: keeps Lens ingestion ahead of the gateway catch-all
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
lensWorker.mode: external
lensWorker.externalServiceName: existing-lens
lensWorker.externalUrl: https://lens.example.com
lensWorker.publicUrl: https://lens.example.com
lensWorker.serviceTokenSecret.name: existing-lens-signing
lensWorker.gateway.secretName: existing-lens-signing
asserts:
- equal:
path: spec.rules[2].http.paths[0].path
value: /lens-ingest
- equal:
path: spec.rules[2].http.paths[0].backend.service.name
value: existing-lens
- equal:
path: spec.rules[2].http.paths[1].path
value: /
template: ingress.yaml
- it: rejects a partial component host configuration
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts.ui, backend, and gateway must all be set together
template: ingress.yaml
- it: rejects component hosts when the gateway is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
gateway.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires gateway.enabled=true
template: ingress.yaml
- it: rejects component hosts when the backend is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
backend.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires backend.enabled=true
template: ingress.yaml
- it: rejects component hosts when the UI is disabled
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
ui.enabled: false
asserts:
- failedTemplate:
errorMessage: ingress.componentHosts requires ui.enabled=true
template: ingress.yaml
- it: rejects a legacy host alongside component hosts
set:
ingress.enabled: true
ingress.host: legacy.example.com
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
asserts:
- failedTemplate:
errorMessage: ingress.host must be empty when ingress.componentHosts is configured
template: ingress.yaml
- it: rejects extra paths alongside component hosts
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ingress.extraPaths:
- path: /custom
asserts:
- failedTemplate:
errorMessage: ingress.extraPaths must be empty when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the UI backend URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.backendUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the backend gateway URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: requires the public UI origin in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.publicUrl is required when ingress.componentHosts is configured
template: ingress.yaml
- it: exposes the gateway URL in the backend Deployment
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: " https://gateway.example.com/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_GATEWAY_URL
value: https://gateway.example.com
- it: preserves the backend gateway URL verbatim in single-host mode
set:
ingress.enabled: true
ingress.host: litellm.example.com
backend.gatewayUrl: " https://gateway.example.com/api/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_GATEWAY_URL
value: " https://gateway.example.com/api/ "
- it: exposes the backend URL in the UI Deployment
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: " https://backend.example.com/ "
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
template: ui/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BACKEND_URL
value: https://backend.example.com
- it: omits the backend URL from the UI Deployment in single-host mode
set:
ingress.enabled: true
ingress.host: litellm.example.com
ui.backendUrl: https://litellm.example.com/api
template: ui/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_BACKEND_URL
any: true
- it: enables credentialed CORS from the public UI origin on the backend
set:
ui.publicUrl: https://ui.example.com
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
- it: lets backend extraEnv override generated CORS values
set:
ui.publicUrl: https://ui.example.com
backend.extraEnv:
- name: LITELLM_CORS_ORIGINS
value: https://backend-override.example.com
- name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "false"
template: backend/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].env[0].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[0].value
value: https://ui.example.com
- equal:
path: spec.template.spec.containers[0].env[1].name
value: LITELLM_CORS_ALLOW_CREDENTIALS
- equal:
path: spec.template.spec.containers[0].env[1].value
value: "true"
- equal:
path: spec.template.spec.containers[0].env[10].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[10].value
value: https://backend-override.example.com
- equal:
path: spec.template.spec.containers[0].env[11].name
value: LITELLM_CORS_ALLOW_CREDENTIALS
- equal:
path: spec.template.spec.containers[0].env[11].value
value: "false"
- it: exposes the public UI origin to the gateway without enabling credentials
set:
ui.publicUrl: https://ui.example.com
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ALLOW_CREDENTIALS
value: "true"
- it: trims the public UI origin before setting backend CORS
set:
ui.publicUrl: " https://ui.example.com/ "
template: backend/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- it: trims the public UI origin before setting gateway CORS
set:
ui.publicUrl: " https://ui.example.com/ "
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_CORS_ORIGINS
value: https://ui.example.com
- it: rejects a UI public URL with a path
set:
ui.publicUrl: https://ui.example.com/app
template: gateway/deployment.yaml
asserts:
- failedTemplate:
errorMessage: ui.publicUrl must be an origin like https://ui.example.com, without a path
- it: lets gateway extraEnv override the generated CORS origin
set:
ui.publicUrl: https://ui.example.com
gateway.extraEnv:
- name: LITELLM_CORS_ORIGINS
value: https://gateway-override.example.com
template: gateway/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].env[0].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[0].value
value: https://ui.example.com
- equal:
path: spec.template.spec.containers[0].env[9].name
value: LITELLM_CORS_ORIGINS
- equal:
path: spec.template.spec.containers[0].env[9].value
value: https://gateway-override.example.com
---
suite: component host UI backend origin validation
templates:
- ingress.yaml
values:
- ./values/required.yaml
tests:
- it: rejects a pathful UI backend URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com/api
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com
asserts:
- failedTemplate:
errorMessage: ui.backendUrl must be an origin like https://backend.example.com, without a path, when ingress.componentHosts is set
- it: rejects a pathful backend gateway URL in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: https://gateway.example.com/v1
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set
- it: rejects a backend gateway URL without a scheme in component host mode
set:
ingress.enabled: true
ingress.componentHosts.ui: ui.example.com
ingress.componentHosts.backend: backend.example.com
ingress.componentHosts.gateway: gateway.example.com
ui.backendUrl: https://backend.example.com
ui.publicUrl: https://ui.example.com
backend.gatewayUrl: gateway.example.com
asserts:
- failedTemplate:
errorMessage: backend.gatewayUrl must be an origin like https://gateway.example.com, without a path, when ingress.componentHosts is set

View file

@ -191,3 +191,400 @@ tests:
content:
name: DATABASE_PASSWORD_READ_REPLICA
any: true
- it: writer hostKey reads the host from the Secret and supports IAM auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.hostKey: hostname
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- it: writer urlKey supplies only DATABASE_URL to each database client container and keeps pool env
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.schema: ""
database.writer.sslMode: ""
database.writer.sslRootCert: ""
database.writer.passwordSecret.urlKey: connection-url
database.connectionPool.enabled: true
gateway.collector.enabled: true
gateway.metricsServer.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLMODE
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SSLROOTCERT
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: gateway/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: gateway/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[2].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_USER
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_PASSWORD
any: true
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[2].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: backend/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: backend/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: backend/deployment.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: backend/deployment.yaml
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: connection-url
template: migrations-job.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER
any: true
template: migrations-job.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
template: migrations-job.yaml
- notExists:
path: >-
spec.template.spec.containers[0].env[?(@.name == "DATABASE_HOST" ||
@.name == "DATABASE_PORT" || @.name == "DATABASE_USER" ||
@.name == "DATABASE_NAME" || @.name == "DATABASE_SCHEMA" ||
@.name == "DATABASE_SSLMODE" || @.name == "DATABASE_SSLROOTCERT" ||
@.name == "DATABASE_PASSWORD")]
template: migrations-job.yaml
- it: writer rejects host together with hostKey
template: gateway/deployment.yaml
set:
database.writer.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "database.writer.host and database.writer.passwordSecret.hostKey cannot both be set"
- it: writer rejects urlKey with discrete connection settings
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.schema: public
database.writer.sslMode: require
database.writer.sslRootCert: /etc/ssl/ca.crt
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer rejects urlKey with IAM auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer rejects urlKey with Entra auth
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.dbname: ""
database.writer.passwordSecret.urlKey: connection-url
database.writer.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, sslMode, sslRootCert, useIAMAuth, or useAzureEntraAuth"
- it: writer requires an endpoint
template: gateway/deployment.yaml
set:
database.writer.host: ""
asserts:
- failedTemplate:
errorMessage: "one of database.writer.host, database.writer.passwordSecret.hostKey, or database.writer.passwordSecret.urlKey is required"
- it: writer hostKey requires a Secret name
template: gateway/deployment.yaml
set:
database.writer.host: ""
database.writer.passwordSecret.hostKey: hostname
database.writer.passwordSecret.name: ""
asserts:
- failedTemplate:
errorMessage: "database.writer.passwordSecret.name is required when hostKey or urlKey is set"
- it: reader renders without replica env vars when passwordSecret is null
template: gateway/deployment.yaml
set:
database.reader.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
any: true
- it: reader hostKey enables replica routing and reads its host from the Secret
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
database.reader.passwordSecret.hostKey: hostname
database.reader.dbname: litellm
database.reader.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
valueFrom:
secretKeyRef:
name: litellm-reader-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME_READ_REPLICA
value: litellm
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- it: reader urlKey emits only the replica URL
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: litellm-reader-secret
key: connection-url
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PORT_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_NAME_READ_REPLICA
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_SCHEMA_READ_REPLICA
any: true
- it: reader rejects host together with hostKey
template: gateway/deployment.yaml
set:
database.reader.host: reader.example.com
database.reader.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "database.reader.host and database.reader.passwordSecret.hostKey cannot both be set"
- it: reader rejects urlKey with discrete connection settings
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
database.reader.dbname: litellm
database.reader.schema: public
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
- it: reader rejects urlKey with token auth
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
database.reader.passwordSecret.urlKey: connection-url
database.reader.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.urlKey cannot be combined with host, hostKey, dbname, schema, useIAMAuth, or useAzureEntraAuth"
- it: reader urlKey requires a Secret name
template: gateway/deployment.yaml
set:
database.reader.passwordSecret.urlKey: connection-url
database.reader.passwordSecret.name: ""
asserts:
- failedTemplate:
errorMessage: "database.reader.passwordSecret.name is required when hostKey or urlKey is set"

View file

@ -21,6 +21,27 @@ tests:
name: REDIS_CLUSTER_NODES
any: true
- it: gateway renders without redis env vars when passwordSecret is null
template: gateway/deployment.yaml
set:
redis.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- it: gateway emits host, port and password when redis is configured
template: gateway/deployment.yaml
set:
@ -68,6 +89,27 @@ tests:
name: redis-secret
key: redis-password
- it: backend renders without redis env vars when passwordSecret is null
template: backend/deployment.yaml
set:
redis.passwordSecret: null
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- it: gateway omits REDIS_PASSWORD for an auth-less redis
template: gateway/deployment.yaml
set:
@ -107,3 +149,174 @@ tests:
content:
name: REDIS_CLUSTER_NODES
any: true
- it: gateway gets redis host from hostKey and keeps port and password
template: gateway/deployment.yaml
set:
redis.host: ""
redis.port: 6380
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.passwordKey: redis-password
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
valueFrom:
secretKeyRef:
name: redis-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6380"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: redis-password
- it: backend gets redis host from hostKey with the same coordination env
template: backend/deployment.yaml
set:
redis.host: ""
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
valueFrom:
secretKeyRef:
name: redis-secret
key: hostname
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
value: "6379"
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
valueFrom:
secretKeyRef:
name: redis-secret
key: password
- it: gateway expands the redis cluster seed from hostKey at runtime
template: gateway/deployment.yaml
set:
redis.host: ""
redis.passwordSecret.name: redis-secret
redis.passwordSecret.hostKey: hostname
redis.cluster: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
value: '[{"host":"$(REDIS_HOST)","port":6379}]'
- equal:
path: spec.template.spec.containers[0].env[8].name
value: REDIS_HOST
- equal:
path: spec.template.spec.containers[0].env[11].name
value: REDIS_CLUSTER_NODES
- it: gateway gets REDIS_URL from urlKey without discrete Redis env
template: gateway/deployment.yaml
set:
redis.passwordSecret.name: redis-secret
redis.passwordSecret.urlKey: connection-url
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_URL
valueFrom:
secretKeyRef:
name: redis-secret
key: connection-url
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_HOST
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PORT
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_PASSWORD
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: REDIS_CLUSTER_NODES
any: true
- it: redis rejects host together with hostKey
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.host and redis.passwordSecret.hostKey cannot both be set"
- it: redis urlKey rejects host
template: gateway/deployment.yaml
set:
redis.host: redis.example.com
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis urlKey rejects hostKey
template: gateway/deployment.yaml
set:
redis.passwordSecret.hostKey: hostname
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis urlKey rejects cluster mode
template: gateway/deployment.yaml
set:
redis.cluster: true
redis.passwordSecret.urlKey: connection-url
redis.passwordSecret.name: redis-secret
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.urlKey cannot be combined with host, hostKey, or cluster"
- it: redis hostKey requires a Secret name
template: gateway/deployment.yaml
set:
redis.passwordSecret.hostKey: hostname
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"
- it: redis urlKey requires a Secret name
template: gateway/deployment.yaml
set:
redis.passwordSecret.urlKey: connection-url
asserts:
- failedTemplate:
errorMessage: "redis.passwordSecret.name is required when hostKey or urlKey is set"

View file

@ -23,7 +23,11 @@ ingress:
# wildcard pathType, so that rule could never match there.
controller: alb
annotations: {}
host: "" # optional; if set, becomes the rule's host
host: "" # optional; if set, becomes the single ingress rule's host
componentHosts:
ui: ""
backend: ""
gateway: ""
tls: []
# Extra HTTP paths appended to the ingress rule. Additive: every built-in
# UI / gateway / backend path is still rendered, these entries are placed
@ -208,6 +212,8 @@ database:
name: litellm-writer-secret
usernameKey: username
passwordKey: password
hostKey: "" # Secret key containing the writer database host
urlKey: "" # Secret key containing the complete writer database URL
# libpq sslmode / sslrootcert applied to the writer and reader URLs (Prisma and the
# in-container PgBouncer); e.g. verify-full with /etc/ssl/certs/ca-certificates.crt for AWS RDS.
# sslRootCert on its own implies sslMode verify-full
@ -229,6 +235,8 @@ database:
name: litellm-reader-secret
usernameKey: username
passwordKey: password
hostKey: "" # Secret key containing the reader database host
urlKey: "" # Secret key containing the complete reader database URL
# In-container connection pool (PgBouncer, transaction mode) shared by every
# gateway worker in the pod. Without it each of the `gateway.numWorkers`
@ -274,6 +282,8 @@ redis:
passwordSecret:
name: "" # Leave empty for auth-less Redis
passwordKey: password
hostKey: "" # Secret key containing the Redis host
urlKey: "" # Secret key containing the complete Redis URL
# ---------- gateway (LLM data plane) ----------
gateway:
@ -493,6 +503,7 @@ gateway:
backend:
enabled: true
logLevel: INFO
gatewayUrl: ""
extraEnv: []
envConfigMaps: []
envSecrets: []
@ -575,9 +586,10 @@ ui:
service:
type: ClusterIP
port: 3000
# The dashboard expects to know where to reach the backend API. Set this to
# the externally-routable URL (typically the ingress host + /api or similar).
# The backend's public origin, passed to the UI only when ingress.componentHosts is configured.
backendUrl: ""
# Public origin of the UI, used for backend and gateway CORS.
publicUrl: ""
resources:
requests:
cpu: 500m

View file

@ -1,5 +1,4 @@
import random
import re
import time
from collections.abc import Generator, Mapping
from contextlib import contextmanager
@ -10,12 +9,11 @@ from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
migration_ddl_lock_timeout,
migration_lock_timeout,
)
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
_LOCK_TIMEOUT_PINNED_RE: Final = re.compile(r"(?:-c\s*|--)lock_timeout=")
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
if TYPE_CHECKING:
import psycopg
@ -30,29 +28,12 @@ def _migration_url(database_url: str, direct_url: "str | None") -> str:
return urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema)))))
def _with_ddl_lock_timeout(url: str) -> str:
split: Final = urlsplit(url)
pairs: Final = tuple(parse_qsl(split.query))
if any(key == "pgbouncer" and value == "true" for key, value in pairs):
return url
existing: Final = next((value for key, value in pairs if key == "options"), "")
if _LOCK_TIMEOUT_PINNED_RE.search(existing):
return url
timeout_ms: Final = max(1, int(migration_ddl_lock_timeout() * 1000))
options: Final = f"{existing} -c lock_timeout={timeout_ms}".strip()
rewritten: Final = tuple((key, value) for key, value in pairs if key != "options") + (("options", options),)
return urlunsplit(split._replace(query=urlencode(rewritten)))
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
database_url: Final = environment.get("DATABASE_URL")
if not database_url:
return environment
migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
return {
**environment,
"DATABASE_URL": _with_ddl_lock_timeout(migration_url),
}
return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
@dataclass(frozen=True, slots=True)
@ -75,10 +56,15 @@ class MigrationCoordinator:
def check_connection(self) -> None:
self.connection.execute("SELECT 1")
def release_statement_snapshot(self) -> None:
# a parameterized statement leaves its portal snapshot registered on the backend until the
# next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away
self.connection.execute("SELECT 1")
def acquire_prisma_lock(self) -> None:
deadline: Final = time.monotonic() + migration_lock_timeout()
while time.monotonic() < deadline:
if _try_lock(self.connection, 72707369):
if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY):
return
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
raise RuntimeError(
@ -115,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
)
@contextmanager
def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
"""Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one
pinned connection while a Prisma command runs with its session advisory lock disabled, so a
transaction pooler can never strand that lock on an idle backend after the command exits."""
with migration_lock(database_url) as coordinator:
coordinator.acquire_prisma_lock()
coordinator.release_statement_snapshot()
yield
@contextmanager
def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
"""A session-level, non-blocking hold of the migration coordinator lock on an autocommit

View file

@ -0,0 +1,14 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_TelemetryReport" (
"id" TEXT NOT NULL,
"window_start" DOUBLE PRECISION NOT NULL,
"window_end" DOUBLE PRECISION NOT NULL,
"report" JSONB NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_TelemetryReport_pkey" PRIMARY KEY ("id")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_window_end_idx" ON "LiteLLM_TelemetryReport"("window_end");
CREATE INDEX IF NOT EXISTS "LiteLLM_TelemetryReport_created_at_idx" ON "LiteLLM_TelemetryReport"("created_at");

View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "mcp_tool_permissions" JSONB;

View file

@ -0,0 +1,4 @@
SET LOCAL lock_timeout = '10s';
-- AlterTable
ALTER TABLE "LiteLLM_AccessGroupTable" ADD COLUMN IF NOT EXISTS "access_skill_names" TEXT[] DEFAULT ARRAY[]::TEXT[];

View file

@ -60,13 +60,11 @@ PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_DDL_LOCK_TIMEOUT"
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT = 10.0
BOOTSTRAP_ARG = "--version"
PRISMA_CONSOLE_SCRIPT = "prisma"
@ -113,11 +111,6 @@ def migration_lock_timeout() -> float:
return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
def migration_ddl_lock_timeout() -> float:
"""Seconds one migration statement may wait on a table lock before Postgres cancels it."""
return _timeout_from_env(MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR, DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT)
def prisma_bootstrap_timeout() -> float:
"""Seconds the one-time Node toolchain install may run for."""
return _timeout_from_env(

View file

@ -1588,7 +1588,9 @@ model LiteLLM_AccessGroupTable {
// Resource memberships - explicit arrays per type
access_model_names String[] @default([])
access_mcp_server_ids String[] @default([])
mcp_tool_permissions Json?
access_agent_ids String[] @default([])
access_skill_names String[] @default([])
assigned_team_ids String[] @default([])
assigned_key_ids String[] @default([])
@ -2045,3 +2047,14 @@ model LiteLLM_LensTraceSignal {
@@id([trace_id, trace_ref])
}
model LiteLLM_TelemetryReport {
id String @id
window_start Float
window_end Float
report Json
created_at DateTime @default(now())
@@index([window_end])
@@index([created_at])
}

View file

@ -9,6 +9,7 @@ import tempfile
import threading
import time
from collections.abc import Callable
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
@ -18,7 +19,6 @@ from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR,
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
ensure_prisma_toolchain,
@ -1022,7 +1022,11 @@ class ProxyExtrasDBManager:
def _setup_database_v2(use_migrate: bool) -> bool:
if not use_migrate:
return ProxyExtrasDBManager._run_database_v2(False)
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
from litellm_proxy_extras.migration_lock import (
migration_environment,
migration_lock,
prisma_migration_lock,
)
from litellm_proxy_extras.migration_recovery import (
baseline_current_schema,
recover_completed_migration,
@ -1056,7 +1060,9 @@ class ProxyExtrasDBManager:
migration_environment(_get_prisma_env()),
)
while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing):
while not ProxyExtrasDBManager._run_database_v2(
True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url)
):
continue
return True
@ -1065,6 +1071,7 @@ class ProxyExtrasDBManager:
use_migrate: bool,
recover_completed: Callable[[str], bool] = lambda name: False,
baseline_existing: "Callable[[str], None] | None" = None,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> bool:
"""
v2 migration resolver (opt-in via --use_v2_migration_resolver).
@ -1127,11 +1134,12 @@ class ProxyExtrasDBManager:
while not budget.exhausted:
attempt_started = time.monotonic()
try:
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=deploy_timeout,
env=migration_env,
)
with deploy_lock():
result = prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "deploy"],
timeout=deploy_timeout,
env=migration_env,
)
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
return True
@ -1156,6 +1164,7 @@ class ProxyExtrasDBManager:
e,
budget,
time.monotonic() - attempt_started,
deploy_lock=deploy_lock,
)
if next_budget.attempts_left < budget.attempts_left:
@ -1165,10 +1174,9 @@ class ProxyExtrasDBManager:
raise RuntimeError(
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
"attempts that made no progress (timeouts, lock timeouts or deadlock retries). "
"Check database connectivity, load, long-running transactions on the migrated tables, "
"and _prisma_migrations ledger state. Raise "
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out, or "
f"{MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if they timed out waiting for a table lock."
"Check database connectivity, load and _prisma_migrations ledger state. "
"If a migration hit SET LOCAL lock_timeout, check for long-running transactions on the migrated tables. "
f"Raise {PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
)
finally:
os.chdir(original_dir)
@ -1189,15 +1197,19 @@ class ProxyExtrasDBManager:
return match.group(1) if match else None
@staticmethod
def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
def _v2_roll_back_migration_best_effort(
migration_name: str,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> None:
from litellm_proxy_extras.migration_lock import migration_environment
try:
prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
timeout=prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
with deploy_lock():
prisma_toolchain.run_prisma(
[_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name],
timeout=prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
pass
@ -1206,6 +1218,7 @@ class ProxyExtrasDBManager:
error: subprocess.CalledProcessError,
budget: "_MigrateAttemptBudget",
attempt_seconds: float = 0.0,
deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext,
) -> "_MigrateAttemptBudget":
"""Recover from one failed `prisma migrate deploy`, and price the pass.
@ -1227,14 +1240,14 @@ class ProxyExtrasDBManager:
"deadlock race, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs:
logger.info(
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
logger.info(
@ -1278,17 +1291,17 @@ class ProxyExtrasDBManager:
"rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
ProxyExtrasDBManager._log_migration_lock_holders(migration_name)
logger.warning(
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying. "
f"Raise {MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if the database needs longer.",
"Migration %s timed out waiting for a table lock after hitting SET LOCAL lock_timeout, rolling "
"its ledger row back and retrying. Check for long-running transactions on the migrated tables.",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock)
return budget.spend()
raise RuntimeError(

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.108"
version = "0.4.109"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -30,7 +30,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.108"
version = "0.4.109"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -3922,13 +3922,16 @@ version = "0.1.0"
dependencies = [
"bytes",
"futures-util",
"litellm-coroutine",
"litellm-host",
"parking_lot",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
@ -4024,6 +4027,7 @@ dependencies = [
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-framer",
"litellm-host",
"litellm-host-native",
"litellm-http",
@ -4107,9 +4111,11 @@ name = "litellm-inference-testing"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-auth",
"litellm-http",
"litellm-inference",
"litellm-secrets",
"tokio",
]
[[package]]

View file

@ -80,6 +80,7 @@ http = "1"
google-cloud-auth = { version = "1.16.0", default-features = false }
jsonwebtoken = { version = "11.1.0", default-features = false }
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
parking_lot = "0.12.5"
proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }

View file

@ -9,7 +9,7 @@ use litellm_host::{
interceptors::{RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing, epoch_seconds},
};
use litellm_host_python::{HookStep, from_py, missing_state, to_py};
use litellm_host_python::{Error as HostError, HookStep, from_py, to_py};
use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
@ -160,7 +160,7 @@ impl LegacyLogging {
self.response
.as_ref()
.map(|response| HookStep::Ready(response.clone_ref(py)))
.ok_or_else(missing_state)
.ok_or_else(|| HostError::MissingState.into())
}
fn dispatch_success(&self, py: Python<'_>) -> PyResult<()> {
@ -209,7 +209,7 @@ impl LegacyLogging {
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.stream_billing().ok_or_else(missing_state)?;
let billing = self.stream_billing().ok_or(HostError::MissingState)?;
let billed = Streaming::Success.call(
py,
(
@ -442,7 +442,7 @@ impl LegacyLogging {
py: Python<'_>,
facts: &litellm_host::interceptors::ExecutionFacts,
) -> PyResult<HookStep<Self, ()>> {
use litellm_host::interceptors::ResultSource;
use litellm_host::interceptors::{Cost, ResultSource};
let logger = self.logger()?.object(py);
let params = logger
@ -462,6 +462,12 @@ impl LegacyLogging {
),
)?;
let details = logger.getattr("model_call_details")?;
let reported_cost = match &facts.cost {
Cost::Reported { amount } => amount.as_f64(),
Cost::CacheHit | Cost::Deferred => None,
};
logger.setattr("_provider_reported_cost", reported_cost)?;
details.call_method1("pop", ("response_cost", py.None()))?;
self.cache_key = match &facts.source {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
@ -519,7 +525,7 @@ impl LegacyLogging {
return self.stream_failure(py);
}
if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() {
let error = self.error.as_ref().ok_or_else(missing_state)?;
let error = self.error.as_ref().ok_or(HostError::MissingState)?;
return Ok(HookStep::Await(
DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type())?,
Self::resume_deployment_failure,
@ -530,7 +536,7 @@ impl LegacyLogging {
pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
if self.stream_billing().is_none() {
return Err(missing_state());
return Err(HostError::MissingState.into());
}
if let Some(key) = &self.cache_key {
head.bind(py).set_item("cache_key", key)?;
@ -545,7 +551,7 @@ impl LegacyLogging {
}
pub(crate) fn stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
let stream = self.stream.as_mut().ok_or_else(missing_state)?;
let stream = self.stream.as_mut().ok_or(HostError::MissingState)?;
if stream.first_chunk.is_none() {
stream.first_chunk = Some(datetime(py, epoch_seconds())?);
}
@ -958,6 +964,81 @@ on_pre_call = lambda additional_args: None
check = lambda: None
";
#[rstest]
#[case::reported(Some(0.37), false)]
#[case::free(Some(0.0), false)]
#[case::deferred(None, false)]
#[case::cached(None, true)]
fn result_ready_hands_only_a_reported_cost_to_the_spend_logging_record(
#[case] reported: Option<f64>,
#[case] cache_hit: bool,
#[values(0.001, 7.0)] existing_cost: f64,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
run(
py,
&locals,
c"
logger.litellm_params = {}
logger.optional_params = {}
logger.model_call_details = {'response_cost': 7.0}
logger._provider_reported_cost = None
",
);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
let source = if cache_hit {
litellm_host::interceptors::ResultSource::Cache {
key: "cached-test".into(),
}
} else {
litellm_host::interceptors::ResultSource::Provider
};
let facts = litellm_host::interceptors::ExecutionFacts {
provider: litellm_host::interceptors::ProviderIdentity {
model: "native-test".into(),
provider: "edenai".into(),
},
cost: litellm_host::interceptors::Cost::settle(
&source,
reported.and_then(serde_json::Number::from_f64),
),
source,
};
let previous = litellm_host::interceptors::ExecutionFacts {
provider: facts.provider.clone(),
cost: litellm_host::interceptors::Cost::Reported {
amount: serde_json::Number::from_f64(9.0).unwrap(),
},
source: litellm_host::interceptors::ResultSource::Provider,
};
logging.result_ready(py, &previous).unwrap();
local(&locals, "logger")
.getattr("model_call_details")
.unwrap()
.set_item("response_cost", existing_cost)
.unwrap();
let step = logging.result_ready(py, &facts).unwrap();
assert!(matches!(step, HookStep::Ready(())));
locals.set_item("reported", reported).unwrap();
locals.set_item("cache_hit", cache_hit).unwrap();
run(
py,
&locals,
c"
assert logger.model_call_details.get('response_cost') is None, logger.model_call_details
assert logger._provider_reported_cost == reported, logger._provider_reported_cost
assert logger.model_call_details['cache_hit'] is cache_hit, logger.model_call_details
assert logger.update['custom_llm_provider'] == 'edenai', logger.update
",
);
});
}
const DOCUMENT: &str = "data:application/pdf;base64,YWJj";
const EDITED: &str = "data:application/pdf;base64,ZWRpdGVk";

View file

@ -1,7 +1,7 @@
use pyo3::prelude::*;
use strum::{IntoStaticStr, VariantArray};
const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python";
const MODULE: &str = "litellm.rust_bridge.host.callbacks_legacy_python";
/// Every litellm Python internal the native call still borrows, grouped by the subsystem it
/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today
@ -9,7 +9,7 @@ const MODULE: &str = "litellm.rust_bridge.callbacks_legacy_python";
/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a
/// user's own callback is not borrowing and does not belong here.
///
/// `litellm/rust_bridge/callbacks_legacy_python.py` is the only Python module behind it, and
/// `litellm/rust_bridge/host/callbacks_legacy_python.py` is the only Python module behind it, and
/// `python_contract.json` pins each function's parameters on both sides.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum LegacyPython {

View file

@ -6,7 +6,7 @@ use pyo3::types::{PyDict, PyTuple};
use crate::{LegacyLogging, PublicCall};
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
/// `tests/unit/rust_bridge/host/test_callbacks_legacy_python.py` pins this file to the Python
/// signatures, and [`namespace`] binds every fake call against it.
pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json");
@ -23,10 +23,10 @@ import sys
import traceback
import types
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.callbacks_legacy_python'):
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.host', 'litellm.rust_bridge.host.callbacks_legacy_python'):
sys.modules.setdefault(name, types.ModuleType(name))
legacy = sys.modules['litellm.rust_bridge.callbacks_legacy_python']
legacy = sys.modules['litellm.rust_bridge.host.callbacks_legacy_python']
CONTRACT = json.loads(python_contract)

View file

@ -16,18 +16,28 @@ pub enum LlmProviders {
AzureAi,
#[strum(serialize = "bedrock")]
Bedrock,
#[strum(serialize = "bedrock_mantle")]
BedrockMantle,
#[strum(serialize = "cohere")]
Cohere,
#[strum(serialize = "deepseek")]
Deepseek,
#[strum(serialize = "edenai")]
Edenai,
#[strum(serialize = "minimax")]
Minimax,
#[strum(serialize = "mistral")]
Mistral,
#[strum(serialize = "openai")]
Openai,
#[strum(serialize = "openai_like")]
OpenaiLike,
#[strum(serialize = "openrouter")]
Openrouter,
#[strum(serialize = "reducto")]
Reducto,
#[strum(serialize = "tencent")]
Tencent,
#[strum(serialize = "vertex_ai")]
VertexAi,
}

View file

@ -3,7 +3,7 @@
Core must pause mid-call to ask the host for things it cannot do itself (Python callbacks, secret and token reads, `before_send` rewrites, stream demand), then continue where it stopped. Any change to this crate must keep every requirement below. The alternatives section says which one each rejected design breaks
- R1 Core never calls the host: it names an op and waits for the answer, so it stays free of PyO3 and of any other host runtime
- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/lifecycle.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
- R2 Async host work is awaited by the host's own driver in the caller's asyncio task (`litellm/rust_bridge/execution.py`), so `contextvars` writes reach the caller. A Rust-side `into_future` would run it in a copied context
- R3 The body awaits real I/O (HTTP, `spawn_blocking`, timers) between yields, so `resume` is itself a future driven by the caller's runtime
- R4 Route code stays straight-line async (`host.route(OcrOp::ReadDocument).await?`) instead of hand-written states
- R5 Each op fixes its answer type at compile time: a host cannot answer `ReadDocument` with a token, and core never matches a result variant it did not ask for

View file

@ -1,6 +1,6 @@
use std::str;
use bytes::{Buf, BufMut, BytesMut};
use bytes::{Buf, BufMut, Bytes, BytesMut};
use tokio_util::codec::{Decoder, Encoder};
use crate::SseError;
@ -29,38 +29,107 @@ impl Decoder for SseCodec {
return Ok(None);
}
while let Some(end) = block_end(src) {
let block = src.split_to(end);
let pending = lines(&block)
.map(|(line, _)| line)
.take_while(|line| !line.is_empty())
.try_fold(Pending::default(), Pending::apply)?;
if let Some(event) = pending.dispatch() {
if let Some(event) = event_in(&src.split_to(end))? {
return Ok(Some(event));
}
}
Ok(None)
}
fn decode_eof(&mut self, _pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
Ok(None)
/// An event the stream ended inside of is dropped, as the SSE spec requires.
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseEvent>, SseError> {
self.decode(pending)
}
}
pub fn decode_all(wire: &[u8]) -> Result<Vec<SseEvent>, SseError> {
let mut codec = SseCodec::default();
let mut buffer = BytesMut::from(wire);
let mut events = Vec::new();
while let Some(event) = codec.decode_eof(&mut buffer)? {
events.push(event);
}
Ok(events)
}
impl SseCodec {
fn skip_bom(&mut self, src: &mut BytesMut) -> bool {
if self.past_bom {
return true;
}
if src.starts_with(BOM) {
src.advance(BOM.len());
} else if BOM.starts_with(src) {
return false;
}
self.past_bom = true;
true
skip_bom(&mut self.past_bom, src)
}
}
fn event_in(block: &[u8]) -> Result<Option<SseEvent>, SseError> {
Ok(lines(block)
.map(|(line, _)| line)
.take_while(|line| !line.is_empty())
.try_fold(Pending::default(), Pending::apply)?
.dispatch())
}
/// One event together with the exact wire bytes that carried it, for relays that forward a
/// provider's bytes untouched but still need to see each event go by.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SseBlock {
pub bytes: Bytes,
pub event: SseEvent,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct SseBlockCodec {
past_bom: bool,
}
impl Decoder for SseBlockCodec {
type Item = SseBlock;
type Error = SseError;
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
self.decode_block(src, false)
}
fn decode_eof(&mut self, pending: &mut BytesMut) -> Result<Option<SseBlock>, SseError> {
self.decode_block(pending, true)
}
}
impl SseBlockCodec {
fn decode_block(
&mut self,
src: &mut BytesMut,
eof: bool,
) -> Result<Option<SseBlock>, SseError> {
if !skip_bom(&mut self.past_bom, src) {
return Ok(None);
}
while let Some(end) = block_end(src) {
if !eof && end == src.len() && src.ends_with(b"\r") {
return Ok(None);
}
let block = src.split_to(end);
if let Some(event) = event_in(&block)? {
return Ok(Some(SseBlock {
bytes: block.freeze(),
event,
}));
}
}
Ok(None)
}
}
fn skip_bom(past_bom: &mut bool, src: &mut BytesMut) -> bool {
if *past_bom {
return true;
}
if src.starts_with(BOM) {
src.advance(BOM.len());
} else if BOM.starts_with(src) {
return false;
}
*past_bom = true;
true
}
fn block_end(bytes: &[u8]) -> Option<usize> {
lines(bytes)
.find(|(line, _)| line.is_empty())
@ -168,3 +237,179 @@ impl Encoder<SseEvent> for SseCodec {
Ok(())
}
}
#[cfg(test)]
mod tests {
use futures_util::{StreamExt, TryStreamExt, stream};
use rstest::rstest;
use super::*;
use crate::frames;
fn pieces(wire: &[u8], size: usize) -> Vec<Bytes> {
wire.chunks(size).map(Bytes::copy_from_slice).collect()
}
async fn blocks(wire: &[u8], size: usize) -> Result<Vec<SseBlock>, SseError> {
frames(
stream::iter(pieces(wire, size).into_iter().map(Ok::<_, std::io::Error>)),
SseBlockCodec::default(),
)
.try_collect()
.await
}
#[rstest]
#[case::lf(
b": keepalive\n\nevent: delta\ndata: a\ndata: b\n\n",
b"event: delta\ndata: a\ndata: b\n\n"
)]
#[case::crlf(
b": keepalive\r\n\r\nevent: delta\r\ndata: a\r\ndata: b\r\n\r\n",
b"event: delta\r\ndata: a\r\ndata: b\r\n\r\n"
)]
#[case::cr(
b": keepalive\r\revent: delta\rdata: a\rdata: b\r\r",
b"event: delta\rdata: a\rdata: b\r\r"
)]
#[case::bom(b"\xEF\xBB\xBFdata: \xe2\x82\xac\n\n", b"data: \xe2\x82\xac\n\n")]
#[case::empty_data(b"event: ignored\n\ndata:\n\n", b"data:\n\n")]
#[case::extensions(
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n",
b"id: 1\nretry: 10\nvendor: x\ndata: a\n\n"
)]
#[tokio::test]
async fn blocks_carry_the_wire_bytes_and_the_event_they_dispatch(
#[case] wire: &[u8],
#[case] expected: &[u8],
#[values(1, 3, 4096)] size: usize,
) {
let blocks = blocks(wire, size).await.unwrap();
let expected_events = decode_all(expected).unwrap();
assert_eq!(expected_events.len(), 1);
assert_eq!(
blocks,
vec![SseBlock {
bytes: Bytes::copy_from_slice(expected),
event: expected_events[0].clone(),
}]
);
assert_eq!(expected_events, decode_all(wire).unwrap());
}
#[rstest]
#[case::empty(b"", vec![])]
#[case::unterminated(b"data: partial\n", vec![])]
#[case::partial_tail(b"data: complete\n\ndata: partial\n", vec![&b"data: complete\n\n"[..]])]
#[case::multiple(b"data: first\n\ndata: second\n\n", vec![&b"data: first\n\n"[..], &b"data: second\n\n"[..]])]
#[case::trailing_comment(b"data: first\n\n: bye\n\n", vec![&b"data: first\n\n"[..]])]
#[tokio::test]
async fn blocks_dispatch_complete_events_and_drop_a_cut_tail(
#[case] wire: &[u8],
#[case] expected: Vec<&[u8]>,
) {
let mut blocks = Box::pin(frames(
stream::iter(pieces(wire, 1).into_iter().map(Ok::<_, std::io::Error>)),
SseBlockCodec::default(),
));
for expected in expected {
assert_eq!(blocks.next().await.unwrap().unwrap().bytes, expected);
}
assert!(blocks.next().await.is_none());
}
#[rstest]
fn bulk_decode_preserves_every_field() {
let events = decode_all(b"id: 1\nretry: 20\nevent: delta\ndata: a\ndata: b\n\n").unwrap();
assert_eq!(
events,
vec![SseEvent {
event: Some("delta".into()),
data: "a\nb".into(),
id: Some("1".into()),
retry: Some(20),
}]
);
}
#[rstest]
#[case::only_a_tail(b"data: partial\n", vec![])]
#[case::tail_after_an_event(b"data: complete\n\ndata: partial\n", vec!["complete"])]
#[case::no_terminator_at_all(b"data: partial", vec![])]
fn bulk_decode_drops_the_event_a_stream_ends_inside(
#[case] wire: &[u8],
#[case] expected: Vec<&str>,
) {
let data: Vec<String> = decode_all(wire)
.unwrap()
.into_iter()
.map(|e| e.data)
.collect();
assert_eq!(data, expected);
}
#[rstest]
#[case::trailing_blank_lines(b"data: x\n\n\n\n")]
#[case::trailing_comment(b"data: x\n\n: bye\n\n")]
#[case::lone_cr_at_eof(b"data: x\r\r")]
fn bulk_decode_accepts_a_stream_that_ends_between_events(#[case] wire: &[u8]) {
assert_eq!(decode_all(wire).unwrap().len(), 1);
}
#[rstest]
fn bulk_decode_rejects_invalid_utf8() {
assert!(matches!(
decode_all(b"data: \xff\n\n"),
Err(SseError::InvalidUtf8(_))
));
}
#[rstest]
#[tokio::test]
async fn block_framing_preserves_transport_errors() {
let input = stream::iter([
Ok(Bytes::from_static(b"data: first\n\n")),
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset",
)),
]);
let mut blocks = Box::pin(frames(input, SseBlockCodec::default()));
assert_eq!(
blocks.next().await.unwrap().unwrap().bytes,
Bytes::from_static(b"data: first\n\n")
);
let Some(Err(SseError::Body(error))) = blocks.next().await else {
panic!("expected body error");
};
let cause = error
.get_ref()
.unwrap()
.downcast_ref::<std::io::Error>()
.unwrap();
assert_eq!(cause.kind(), std::io::ErrorKind::ConnectionReset);
assert!(blocks.next().await.is_none());
}
#[rstest]
#[tokio::test]
async fn invalid_utf8_stops_block_framing_after_prior_events() {
let mut blocks = Box::pin(frames(
stream::iter(
pieces(b"data: first\n\ndata: \xff\n\n", 1)
.into_iter()
.map(Ok::<_, std::io::Error>),
),
SseBlockCodec::default(),
));
assert_eq!(
blocks.next().await.unwrap().unwrap().bytes,
Bytes::from_static(b"data: first\n\n")
);
assert!(matches!(
blocks.next().await,
Some(Err(SseError::InvalidUtf8(_)))
));
assert!(blocks.next().await.is_none());
}
}

View file

@ -97,12 +97,11 @@ async fn dispatch_follows_the_data_buffer(#[case] wire: &[u8], #[case] expected:
}
#[rstest]
#[case::unterminated_single(b"data: partial\n", vec![])]
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
#[case::lone_cr_terminates_at_eof(b"data: x\r\r", vec![event(None, "x")])]
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
#[case::trailing_blank_lines(b"data: x\n\n\n", vec![event(None, "x")])]
#[case::trailing_comment(b"data: x\n\n: bye\n\n", vec![event(None, "x")])]
#[tokio::test]
async fn eof_dispatches_only_terminated_events(
async fn eof_between_events_ends_the_stream_cleanly(
#[case] wire: &[u8],
#[case] expected: Vec<SseEvent>,
) {
@ -112,6 +111,19 @@ async fn eof_dispatches_only_terminated_events(
);
}
#[rstest]
#[case::unterminated_single(b"data: partial\n", vec![])]
#[case::unterminated_tail_after_complete(b"data: complete\n\ndata: unfinished\n", vec![event(None, "complete")])]
#[case::lone_cr_line_then_eof(b"data: x\r", vec![])]
#[case::partial_bom(b"\xEF\xBB", vec![])]
#[tokio::test]
async fn eof_inside_an_event_drops_it_after_the_terminated_events(
#[case] wire: &[u8],
#[case] expected: Vec<SseEvent>,
) {
assert_eq!(collect(every(wire, 1)).await.unwrap(), expected);
}
#[rstest]
#[case::inside_the_first_line(vec![&b"data: a\r"[..], &b"\ndata: b\r\n\r\n"[..]])]
#[case::inside_the_blank_line(vec![&b"data: a\r\ndata: b\r\n\r"[..], &b"\n"[..]])]

View file

@ -19,7 +19,12 @@ use crate::{Deployment, Error, Gateway, JsonObject, RequestId, request};
/// Client headers Python forwards to Anthropic-speaking providers on every call.
const ANTHROPIC_API_HEADERS: [&str; 2] = ["anthropic-version", "anthropic-beta"];
const ANTHROPIC_API_HEADER_PROVIDERS: &str = "anthropic,bedrock,bedrock_mantle,vertex_ai";
const OPENROUTER_API_HEADERS: [&str; 2] = [
"x-openrouter-metadata",
"x-openrouter-experimental-metadata",
];
const ANTHROPIC_API_HEADER_PROVIDERS: &str =
"anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter";
pub async fn create(
State(gateway): State<Arc<Gateway>>,
@ -88,17 +93,28 @@ fn project(
}
fn anthropic_api_headers(headers: &HeaderMap) -> Option<ProviderSpecificHeaders> {
let extra_headers: Map<String, Value> = ANTHROPIC_API_HEADERS
let scopes = [
(
ANTHROPIC_API_HEADER_PROVIDERS,
ANTHROPIC_API_HEADERS.as_slice(),
),
("openrouter", OPENROUTER_API_HEADERS.as_slice()),
];
let forwarded: Vec<ProviderSpecificHeader> = scopes
.into_iter()
.filter_map(|name| {
let value = headers.get(name)?.to_str().ok()?;
Some((name.to_owned(), Value::from(value)))
.filter_map(|(provider, names)| {
let extra_headers: Map<String, Value> = names
.iter()
.filter_map(|name| {
let value = headers.get(*name)?.to_str().ok()?;
Some(((*name).to_owned(), Value::from(value)))
})
.collect();
(!extra_headers.is_empty()).then(|| ProviderSpecificHeader {
custom_llm_provider: provider.into(),
extra_headers,
})
})
.collect();
(!extra_headers.is_empty()).then(|| {
ProviderSpecificHeaders::One(ProviderSpecificHeader {
custom_llm_provider: ANTHROPIC_API_HEADER_PROVIDERS.into(),
extra_headers,
})
})
(!forwarded.is_empty()).then_some(ProviderSpecificHeaders::Many(forwarded))
}

View file

@ -192,3 +192,52 @@ async fn model_permissions_prevent_provider_calls(#[case] path: &str) {
assert_eq!(response.status(), 403);
assert!(upstream.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[case::completion(false)]
#[case::streaming(true)]
#[tokio::test]
async fn openrouter_receives_native_messages_headers_and_its_deployment_bearer(
#[case] streaming: bool,
) {
let upstream = MockServer::start().await;
let message = json!({"id": "msg_test", "type": "message", "role": "assistant", "model": "anthropic/test-model", "content": [{"type": "text", "text": "hello"}], "stop_reason": "end_turn", "stop_sequence": null, "usage": {"input_tokens": 1, "output_tokens": 1, "cost": 0.001, "cost_details": "unexpected"}, "openrouter_metadata": {"provider": "test-provider"}});
let sse = "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"cost\":0.001}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let template = if streaming {
ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")
} else {
ResponseTemplate::new(200).set_body_json(message.clone())
};
Mock::given(method("POST"))
.and(path("/api/v1/messages"))
.and(header("authorization", "Bearer test-key"))
.and(header("anthropic-beta", "test-feature"))
.and(header("anthropic-version", "test-version"))
.and(header("x-openrouter-metadata", "enabled"))
.and(body_json(json!({"model": "anthropic/test-model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "stream": streaming, "provider": {"order": ["test-provider"]}})))
.respond_with(template).expect(1).mount(&upstream).await;
let request = Request::post("/v1/messages")
.header("content-type", "application/json")
.header("anthropic-beta", "test-feature")
.header("anthropic-version", "test-version")
.header("x-openrouter-metadata", "enabled")
.header("x-api-key", "proxy-key")
.header("authorization", "Bearer proxy-key")
.body(Body::from(json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16, "stream": streaming, "provider": {"order": ["test-provider"]}}).to_string())).unwrap();
let response = support::app(
"openrouter/anthropic/test-model",
&format!("{}/api/v1", upstream.uri()),
)
.oneshot(request)
.await
.unwrap();
assert_eq!(response.status(), 200);
if streaming {
assert_eq!(to_bytes(response.into_body(), 4096).await.unwrap(), sse);
} else {
assert_eq!(support::json(response).await, message);
}
let requests = upstream.received_requests().await.unwrap();
assert_eq!(requests.len(), 1);
assert!(!requests[0].headers.contains_key("x-api-key"));
}

View file

@ -24,7 +24,7 @@ From the repository root, this loopback-only example runs the existing Python MC
```sh
cargo run --manifest-path litellm-rust/Cargo.toml -p litellm-gateway-mcp --example stdio_proxy -- \
.venv/bin/python -c 'import runpy; runpy.run_path("tests/mcp_tests/mcp_e2e_upstream_server.py")["mcp"].run(transport="stdio")'
.venv/bin/python -c 'import runpy; runpy.run_path("tests/integration/_support/mcp_e2e_upstream_server.py")["mcp"].run(transport="stdio")'
```
```sh

View file

@ -7,7 +7,7 @@ The binary reads YAML from `LITELLM_CONFIG` and mounts configured MCP servers al
From the repository root, start the Python fixture in one terminal:
```sh
MCP_HOST=127.0.0.1 MCP_PORT=8090 .venv/bin/python tests/mcp_tests/mcp_e2e_upstream_server.py
MCP_HOST=127.0.0.1 MCP_PORT=8090 .venv/bin/python tests/integration/_support/mcp_e2e_upstream_server.py
```
Start the gateway in another terminal with the example configuration:

View file

@ -6,7 +6,7 @@ This crate owns CPython execution mechanics for generic `litellm-host` machines
`src/native.rs` belongs here: it runs a generic machine through the Python runtime and owns its pending execution and abort handle. Keep provider selection, request projection and public exception policy out of it. A rename to `machine_runner.rs` is optional and must not change behavior
The execution handle receives its Python lifecycle binding from its consumer through `PythonLifecycle` rather than import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
`run_call` returns the unstarted `Execution` and never names a Python stream class or driver. The consumer drives it from Python and supplies the step classes through `StepTypes` rather than have this crate import a fixed `litellm.rust_bridge` module. Generic suspension and execution state validation belong here. Public stream wrappers and `_hidden_params` conventions belong to the consumer
Creating a resolved asyncio Future from an already constructed Python value belongs here, alongside runtime waiting, interpreter detachment and panic containment. Choosing which callable exceptions become a public `RuntimeError` belongs to the consumer. `python-bridge::callable::wrap_failure` owns that policy
@ -28,6 +28,7 @@ Terminal notifications share the selected response or exception. An ordinary not
- `PythonBinding::decode_request` receives the keyword view returned by `prepare_arguments` and updated by `arguments_prepared`, not the caller's dict. A binding that decodes from it inherits all composed argument rewrites
- A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the binding's `map_error`. A Python exception raised inside the call, and a failure in `prepare_arguments` or `transform_response`, is raised as is
- A failing `map_error` is raised with the native error's text as its `__context__`, never swallowed
- Consumers import this crate's error as `use litellm_host_python::Error as HostError`, so it never shadows their own `Error`
- Use standard PyO3 ownership and conversion APIs
- Prefer `Bound<'py, T>` for attached operations/results, `Py<T>` for retention. Binding/unbinding does not copy payloads
- Use `pythonize` for selected Serde data, never a JSON-text round trip. Share conversion with `Pythonized<T>`
@ -37,7 +38,7 @@ Terminal notifications share the selected response or exception. An ordinary not
- Keep diagnostic counters in the consumer. Wrapper invocations do not measure every interpreter release
- Release exclusive class borrows/locks before Python calls or decrements that can invoke finalizers. Expose retained Python edges to GC without calling Python during traversal
- Keep coroutine driving in the shared Python driver and the native handle
- Shared driver implementation: `litellm/rust_bridge/lifecycle.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
- Shared driver implementation: `litellm/rust_bridge/execution.py`, handle: `src/handle.rs`, call driver: `src/driver.rs`, native-backed behavior tests: `tests/lifecycle.py`. The consumer supplies the lifecycle binding
- Every lifecycle suspension is awaited inline in the caller's task. `into_future` creates a separate task and cannot satisfy this contract
- References: [ownership](https://pyo3.rs/v0.29.2/types.html), [conversions](https://pyo3.rs/v0.29.2/conversions/traits.html), [pythonize errors](https://docs.rs/pythonize/0.29.0/src/pythonize/error.rs.html)
- [GC](https://pyo3.rs/v0.29.2/class/protocols.html#garbage-collector-integration), [re-entry](https://pyo3.rs/v0.29.2/class/call.html), [parallelism](https://pyo3.rs/v0.29.2/parallelism.html), [async delivery source](https://docs.rs/pyo3-async-runtimes/0.29.0/src/pyo3_async_runtimes/generic.rs.html)

View file

@ -6,11 +6,14 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-coroutine.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
parking_lot.workspace = true
serde.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["rt", "sync"] }
pyo3.workspace = true

View file

@ -0,0 +1,58 @@
use std::ops::ControlFlow;
use litellm_coroutine::{Abandoned, Co, Reply};
use pyo3::prelude::*;
use crate::Error;
use crate::handle::Suspension;
pub(crate) type Answer = PyResult<Py<PyAny>>;
pub(crate) struct Yielded {
pub(crate) suspension: Suspension,
pub(crate) reply: Reply<Answer>,
}
/// Hands suspensions to the Python caller and receives its answers.
#[derive(Clone)]
pub(crate) struct PythonCaller {
co: Co<Yielded>,
}
impl PythonCaller {
/// Connects the call body to its coroutine's yield handle.
pub(crate) fn new(co: Co<Yielded>) -> Self {
Self { co }
}
/// Asks the caller to await Python work and return its value or exception.
pub(crate) async fn await_python(&self, awaitable: Py<PyAny>) -> Answer {
self.suspend(Suspension::Await(awaitable)).await
}
/// Hands the caller a stream head and waits for demand or closure.
pub(crate) async fn open_stream(&self, head: Py<PyAny>) -> ControlFlow<()> {
self.stream_boundary(Suspension::StreamOpen(head)).await
}
/// Delivers a chunk and waits for the caller to request another or close.
pub(crate) async fn yield_chunk(&self, chunk: Py<PyAny>) -> ControlFlow<()> {
self.stream_boundary(Suspension::StreamChunk(chunk)).await
}
/// Maps stream demand to continuation and any error to detachment.
async fn stream_boundary(&self, suspension: Suspension) -> ControlFlow<()> {
match self.suspend(suspension).await {
Ok(_) => ControlFlow::Continue(()),
Err(_) => ControlFlow::Break(()),
}
}
/// Yields a suspension with its reply handle and waits for the answer.
async fn suspend(&self, suspension: Suspension) -> Answer {
match self.co.yield_(|reply| Yielded { suspension, reply }).await {
Ok(answer) => answer,
Err(Abandoned) => Err(Error::MissingState.into()),
}
}
}

File diff suppressed because it is too large Load diff

View file

@ -1,20 +1,58 @@
use pyo3::prelude::*;
use pyo3::{
PyErr,
exceptions::{PyRuntimeError, PyTypeError},
};
/// A failure host-python raises itself, surfacing in Python as the exception `PyErr::from` picks.
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("missing native call state")]
MissingState,
#[error("execution is already running")]
ExecutionRunning,
#[error("execution is closed")]
ExecutionClosed,
#[error("can't send non-None value to a just-started execution")]
SendBeforeStart,
#[error("execution requires start before resume and can only start once")]
ExecutionOutOfOrder,
#[error("execution was closed while running")]
ExecutionClosedWhileRunning,
#[error("native execution closed")]
NativeExecutionClosed,
#[error("synchronous native routes cannot run from a Tokio context; use the async route")]
SyncInsideTokio,
#[error("the context body produced no value")]
ContextBodyProducedNothing,
#[error("file read must return bytes or str, got {got}")]
FileContentType { got: String },
}
impl From<Error> for PyErr {
fn from(error: Error) -> Self {
match error {
Error::SendBeforeStart | Error::FileContentType { .. } => {
PyTypeError::new_err(error.to_string())
}
Error::MissingState
| Error::ExecutionRunning
| Error::ExecutionClosed
| Error::ExecutionOutOfOrder
| Error::ExecutionClosedWhileRunning
| Error::NativeExecutionClosed
| Error::SyncInsideTokio
| Error::ContextBodyProducedNothing => PyRuntimeError::new_err(error.to_string()),
}
}
}
/// Why a custom operation the host answered did not produce a result: the route's own code
/// rejected it, which the route classifies like any other native failure, or Python code
/// raised, which reaches the caller as it was raised.
#[derive(Debug)]
#[derive(Debug, thiserror::Error)]
pub enum InvokeError<E> {
#[error("{0}")]
Native(E),
Python(PyErr),
}
impl<E> From<PyErr> for InvokeError<E> {
fn from(error: PyErr) -> Self {
Self::Python(error)
}
}
pub fn missing_state() -> PyErr {
pyo3::exceptions::PyRuntimeError::new_err("missing native call state")
#[error(transparent)]
Python(#[from] PyErr),
}

View file

@ -3,13 +3,14 @@
use bytes::Bytes;
use pyo3::{
exceptions::PyTypeError,
gc::{PyTraverseError, PyVisit},
prelude::*,
pybacked::PyBackedBytes,
types::{PyBytes, PyString},
};
use crate::Error;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FileContent {
pub bytes: Bytes,
@ -50,10 +51,10 @@ impl PythonFileReader {
} else if value.is_instance_of::<PyBytes>() {
py_bytes(&value)?
} else {
return Err(PyTypeError::new_err(format!(
"file read must return bytes or str, got {}",
value.get_type(),
)));
return Err(Error::FileContentType {
got: value.get_type().to_string(),
}
.into());
};
Ok(FileContent {
bytes,

View file

@ -3,7 +3,9 @@ use std::sync::{
atomic::{AtomicU64, Ordering},
};
use pyo3::{exceptions::PyRuntimeError, prelude::*};
use pyo3::prelude::*;
use crate::Error;
/// The caller's `contextvars` context, captured at the Python entry point so blocking Python
/// work started from Rust sees the same request-local values as the Python caller.
@ -49,7 +51,7 @@ impl PythonContext {
.lock()
.expect("context value slot poisoned")
.take()
.ok_or_else(|| PyRuntimeError::new_err("the context body produced no value"))
.ok_or_else(|| Error::ContextBodyProducedNothing.into())
}
}

View file

@ -1,19 +1,30 @@
use std::panic::{AssertUnwindSafe, catch_unwind};
use crate::panic_to_pyerr;
use pyo3::exceptions::{PyBaseException, PyRuntimeError};
use crate::{Error, panic_to_pyerr};
use pyo3::exceptions::PyBaseException;
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;
pub type PythonLifecycle = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
pub type StepTypes = for<'py> fn(Python<'py>) -> PyResult<Bound<'py, PyModule>>;
pub enum ExecutionStep {
Return(Py<PyAny>),
/// What a suspended execution waits on: an awaitable the caller awaits, or a stream
/// boundary the caller's stream reads past.
pub enum Suspension {
Await(Py<PyAny>),
/// The call streams: the caller gets a stream over this execution carrying this head,
/// and the execution stays suspended until the stream asks for a chunk.
Open(Py<PyAny>),
Yield(Py<PyAny>),
StreamOpen(Py<PyAny>),
StreamChunk(Py<PyAny>),
}
enum Resume {
Send(Py<PyAny>),
Throw(PyErr),
}
pub enum ExecutionStep {
Suspended(Suspension),
Return(Py<PyAny>),
}
pub trait ExecutionBody: Send + Sync {
@ -31,77 +42,59 @@ enum ExecutionState {
#[pyclass]
pub struct Execution {
state: ExecutionState,
lifecycle: PythonLifecycle,
steps: StepTypes,
}
impl Execution {
pub fn new(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
pub fn new(body: impl ExecutionBody + 'static, steps: StepTypes) -> Self {
Self {
state: ExecutionState::Created(Box::new(body)),
lifecycle,
steps,
}
}
pub fn into_coroutine(self, py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
let binding = (self.lifecycle)(py)?;
let execution = Py::new(py, self)?;
binding.getattr("drive")?.call1((execution,))
}
pub(crate) fn into_sync_stream(
self,
py: Python<'_>,
head: Py<PyAny>,
) -> PyResult<Bound<'_, PyAny>> {
(self.lifecycle)(py)?
.getattr("SyncStream")?
.call1((Py::new(py, self)?, head))
}
/// An execution already started elsewhere and now waiting for its next input.
pub fn suspended(body: impl ExecutionBody + 'static, lifecycle: PythonLifecycle) -> Self {
Self {
state: ExecutionState::Suspended(Box::new(body)),
lifecycle,
}
}
fn advance(
slf: &Bound<'_, Self>,
py: Python<'_>,
result: Option<PyResult<Py<PyAny>>>,
) -> PyResult<Py<PyAny>> {
let mut body = {
fn advance(slf: &Bound<'_, Self>, py: Python<'_>, input: Resume) -> PyResult<Py<PyAny>> {
let (mut body, started) = {
let mut execution = slf.borrow_mut();
match (&execution.state, result.is_some()) {
(ExecutionState::Created(_), false) | (ExecutionState::Suspended(_), true) => {}
let started = match (&execution.state, &input) {
(ExecutionState::Created(_), Resume::Send(value)) if value.is_none(py) => false,
(ExecutionState::Created(_), Resume::Send(_)) => {
return Err(Error::SendBeforeStart.into());
}
(ExecutionState::Suspended(_), _) => true,
(ExecutionState::Running, _) => {
return Err(PyRuntimeError::new_err("execution is already running"));
return Err(Error::ExecutionRunning.into());
}
(ExecutionState::Closed, _) => {
return Err(PyRuntimeError::new_err("execution is closed"));
return Err(Error::ExecutionClosed.into());
}
_ => {
return Err(PyRuntimeError::new_err(
"execution requires start before resume and can only start once",
));
(ExecutionState::Created(_), Resume::Throw(_)) => {
return Err(Error::ExecutionOutOfOrder.into());
}
}
};
match std::mem::replace(&mut execution.state, ExecutionState::Running) {
ExecutionState::Created(body) | ExecutionState::Suspended(body) => body,
ExecutionState::Created(body) | ExecutionState::Suspended(body) => (body, started),
_ => unreachable!(),
}
};
let lifecycle = slf.borrow().lifecycle;
let result = match input {
Resume::Send(value) => started.then_some(Ok(value)),
Resume::Throw(error) => Some(Err(error)),
};
let steps = slf.borrow().steps;
let outcome = catch_unwind(AssertUnwindSafe(|| {
let step = body.resume(result)?;
let (tag, value, suspended) = match step {
ExecutionStep::Await(value) => ("Await", value, true),
ExecutionStep::Open(head) => ("Open", head, true),
ExecutionStep::Yield(value) => ("Yield", value, true),
ExecutionStep::Suspended(Suspension::Await(value)) => ("Await", value, true),
ExecutionStep::Suspended(Suspension::StreamOpen(head)) => {
("StreamOpen", head, true)
}
ExecutionStep::Suspended(Suspension::StreamChunk(chunk)) => {
("StreamChunk", chunk, true)
}
ExecutionStep::Return(value) => ("Complete", value, false),
};
let step = lifecycle(py)?.getattr(tag)?.call1((value,))?.unbind();
let step = steps(py)?.getattr(tag)?.call1((value,))?.unbind();
Ok((step, suspended))
}))
.map_err(panic_to_pyerr)
@ -116,9 +109,7 @@ impl Execution {
drop(body);
outcome.and_then(|(step, suspended)| {
if suspended {
Err(PyRuntimeError::new_err(
"execution was closed while running",
))
Err(Error::ExecutionClosedWhileRunning.into())
} else {
Ok(step)
}
@ -130,24 +121,16 @@ impl Execution {
#[pymethods]
impl Execution {
fn start(slf: &Bound<'_, Self>, py: Python<'_>) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, None)
fn send(slf: &Bound<'_, Self>, py: Python<'_>, value: Py<PyAny>) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, Resume::Send(value))
}
fn resume_value(
slf: &Bound<'_, Self>,
py: Python<'_>,
value: Py<PyAny>,
) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, Some(Ok(value)))
}
fn resume_error(
fn throw(
slf: &Bound<'_, Self>,
py: Python<'_>,
error: Bound<'_, PyBaseException>,
) -> PyResult<Py<PyAny>> {
Self::advance(slf, py, Some(Err(PyErr::from_value(error.into_any()))))
Self::advance(slf, py, Resume::Throw(PyErr::from_value(error.into_any())))
}
fn close(slf: &Bound<'_, Self>) {

View file

@ -6,7 +6,7 @@ use pyo3::{
types::PyDict,
};
use crate::{HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, missing_state};
use crate::{Error, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned};
pub(super) enum ChainStep<T> {
Ready(T),
@ -110,7 +110,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Py<PyDict>>> {
let Some(Continuation::Arguments(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Arguments))
@ -132,7 +132,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Box<WireRequest>>> {
let Some(Continuation::Wire(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Wire))
@ -154,7 +154,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<Py<PyAny>>> {
let Some(Continuation::Response(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Response))
@ -171,7 +171,7 @@ impl<H: PythonCallHooks> ChainHooks for HookAdapter<H> {
result: PyResult<Py<PyAny>>,
) -> PyResult<ChainStep<()>> {
let Some(Continuation::Event(resume)) = self.continuation.take() else {
return Err(missing_state());
return Err(Error::MissingState.into());
};
let step = resume(&mut self.hooks, py, result)?;
Ok(self.step(step, Continuation::Event))

View file

@ -10,9 +10,7 @@ use pyo3::{
types::PyDict,
};
use crate::{
HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime, missing_state,
};
use crate::{Error, HookStep, PythonCallEvent, PythonCallHooks, PythonOwned, PythonRuntime};
use super::adapter::{ChainHooks, ChainStep, HookAdapter};
@ -72,7 +70,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, Py<PyDict>>> {
let (index, context) = self.arguments.take().ok_or_else(missing_state)?;
let (index, context) = self.arguments.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let step = self.hooks[index].resume_arguments(py, result)?;
self.arguments_step(py, index, step, context)
@ -113,7 +111,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
let (index, context) = self.wire.take().ok_or_else(missing_state)?;
let (index, context) = self.wire.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let step = self.hooks[index].resume_wire(py, result)?;
self.wire_step(py, index, step, &context)
@ -154,7 +152,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
let (index, context) = self.response.take().ok_or_else(missing_state)?;
let (index, context) = self.response.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let step = self.hooks[index].resume_response(py, result)?;
self.response_step(py, index, step, context)
@ -195,7 +193,7 @@ impl HookChain {
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> PyResult<HookStep<Self, ()>> {
let (index, event) = self.event.take().ok_or_else(missing_state)?;
let (index, event) = self.event.take().ok_or(Error::MissingState)?;
let result = resume_unless_cancelled(py, result)?;
let result = self.hooks[index].resume_event(py, result);
let step = notification_result(py, is_terminal(&event), result)?;

View file

@ -5,6 +5,7 @@
mod argument;
mod binding;
mod caller;
mod conversion_cache;
mod driver;
mod error;
@ -23,11 +24,11 @@ pub use argument::{effective_py_args, lookup, present};
pub use binding::PythonBinding;
pub use conversion_cache::{FromPythonCache, ToPythonCache};
pub use driver::{CallOptions, run_call};
pub use error::{InvokeError, missing_state};
pub use error::{Error, InvokeError};
pub use file_reader::{FileContent, PythonFileReader, py_bytes};
pub use fork_gate::RuntimeAlreadyStarted;
pub use gil::{PythonContext, attach_blocking, release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep, PythonLifecycle};
pub use handle::{Execution, ExecutionBody, ExecutionStep, StepTypes, Suspension};
pub use hooks::{HookChain, HookResume, HookStep, PythonCallEvent, PythonCallHooks, PythonRuntime};
pub use marshal::{
Pythonized, from_py, from_py_argument, json_loads, json_object_field, panic_to_pyerr, to_py,

View file

@ -5,11 +5,10 @@ use futures_util::future::{AbortHandle, Abortable};
use litellm_host::call::HostedCompletion;
use litellm_host::machine::{HostFailure, Machine, MachineStep};
use litellm_host::protocol::Protocol;
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use tokio::sync::Mutex;
use crate::missing_state;
use crate::Error;
use crate::runtime::{poll_async_value, run_async_value, run_sync_value};
type NativeResult<M> = Result<
@ -65,7 +64,7 @@ where
py: Python<'_>,
interruption: Option<HostFailure<<M::Protocol as Protocol>::Error>>,
) -> PyResult<NativePoll<NativeResult<M>>> {
let state = Arc::clone(self.state.as_ref().ok_or_else(missing_state)?);
let state = Arc::clone(self.state.as_ref().ok_or(Error::MissingState)?);
let future = async move {
let mut state = state.lock().await;
let result = match interruption {
@ -90,7 +89,7 @@ where
run_async_value(py, async move {
Abortable::new(future, registration)
.await
.map_err(|_| PyRuntimeError::new_err("native execution closed"))?
.map_err(|_| Error::NativeExecutionClosed)?
})?
.unbind(),
))
@ -103,12 +102,12 @@ where
pub(super) fn take_result(&self) -> PyResult<NativeResult<M>> {
self.state
.as_ref()
.ok_or_else(missing_state)?
.ok_or(Error::MissingState)?
.try_lock()
.map_err(|_| missing_state())?
.map_err(|_| Error::MissingState)?
.result
.take()
.ok_or_else(missing_state)
.ok_or_else(|| Error::MissingState.into())
.map(|result| {
result.map(|step| match step {
MachineStep::Suspended(op) => MachineStep::Suspended(op),

View file

@ -5,7 +5,7 @@ use std::task::{Context, Poll, Waker};
use std::time::Duration;
use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
use crate::{Pythonized, panic_to_pyerr, release_gil};
use crate::{Error, Pythonized, panic_to_pyerr, release_gil};
use futures_util::FutureExt;
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
@ -112,9 +112,7 @@ where
F: Future<Output = PyResult<T>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
return Err(Error::SyncInsideTokio.into());
}
release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?
}
@ -131,9 +129,7 @@ where
F: Future<Output = Result<T, E>> + Send + 'static,
{
if Handle::try_current().is_ok() {
return Err(PyRuntimeError::new_err(
"synchronous native routes cannot run from a Tokio context; use the async route",
));
return Err(Error::SyncInsideTokio.into());
}
let result = release_gil(py, move || runtime.block_on(wait_for_sync_result(future)))?;

View file

@ -0,0 +1,770 @@
mod support;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use litellm_host::{
hooks::CallHooks,
lifecycle::{CallEvent, ExecutionEvent, FailureOrigin},
machine::CallMachine,
observation::observation_channel,
};
use litellm_host_python::{CallOptions, HookChain, PythonOwned, PythonRuntime, run_call};
use pyo3::{
exceptions::{
PyBaseException, PyException, PyLookupError, PyRuntimeError, PyStopIteration, PyTypeError,
PyValueError, asyncio::CancelledError,
},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
};
use rstest::rstest;
use support::{
HookScript, Interpreter, Log, MachineFactory, OpScript, Synthetic, binding, call_options,
failing_machine, handed, hooks, interpreter, lifecycle_module, open_stream, read_all,
run_composed, run_scripted, settled, success_machine,
};
fn observed(capacity: usize) -> (CallOptions, impl FnMut() -> Vec<CallEvent>) {
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
let options = CallOptions {
asynchronous: false,
steps: lifecycle_module,
observers: Some(sender),
};
(options, move || {
std::iter::from_fn(|| receiver.try_recv().ok()).collect()
})
}
#[rstest]
#[case::success(OpScript::AwaitAnswer)]
#[case::failure(OpScript::AwaitFailure)]
fn asynchronous_host_operations_resume_the_same_machine(
#[from(interpreter)] _python: Interpreter,
#[case] op: OpScript,
) {
Python::attach(|py| {
let machine: MachineFactory = Box::new(|_| {
CallMachine::<Synthetic>::new(None, |host| {
Box::pin(async move { host.services.call(|reply| ("read", reply)).await })
})
});
let (result, log) = run_scripted(py, machine, op, HookScript::Plain, true);
match op {
OpScript::AwaitAnswer => {
assert_eq!(result.unwrap().extract::<String>(py).unwrap(), "awaited")
}
OpScript::AwaitFailure => {
assert!(result.unwrap_err().is_instance_of::<PyLookupError>(py));
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
1
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded:")));
}
_ => unreachable!(),
}
});
}
#[rstest]
fn composed_hooks_share_prepared_arguments_and_one_execution(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let chain = HookChain::new()
.with(hooks(&log, HookScript::ObserveArguments))
.with(hooks(&log, HookScript::RewriteArguments))
.with(hooks(&log, HookScript::ReplaceResponse));
let result = handed(
py,
run_call(
py,
move |_, _, request| Ok(success_machine()(request)),
binding(&log, OpScript::Answer),
chain,
PyDict::new(py).unbind(),
call_options(asynchronous),
),
asynchronous,
);
if asynchronous {
assert!(log.entries().is_empty());
}
let response = settled(py, result, asynchronous).unwrap();
assert_eq!(response.extract::<String>(py).unwrap(), "replaced");
let entries = log.entries();
assert!(entries.iter().any(|entry| entry == "adopted:hook"));
assert_eq!(
entries.iter().filter(|entry| *entry == "project").count(),
1
);
assert_eq!(
entries.iter().filter(|entry| *entry == "op:sign").count(),
1
);
assert_eq!(
entries
.iter()
.filter(|entry| *entry == "succeeded:replaced")
.count(),
3
);
assert!(!entries.iter().any(|entry| entry.starts_with("failed:")));
});
}
#[rstest]
#[case::draining(4, false)]
#[case::full(1, false)]
#[case::closed(4, true)]
fn observation_delivery_preserves_callback_order_and_response(
#[from(interpreter)] _python: Interpreter,
#[case] capacity: usize,
#[case] closed: bool,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (expected, expected_log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::ReplaceResponse,
asynchronous,
);
let (sender, mut receiver) = observation_channel(NonZeroUsize::new(capacity).unwrap());
if closed {
receiver.close();
}
let log = Log::default();
let (actual, actual_log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::ReplaceResponse),
std::convert::identity,
CallOptions {
asynchronous,
steps: lifecycle_module,
observers: Some(sender.clone()),
},
);
assert_eq!(
actual.unwrap().extract::<String>(py).unwrap(),
expected.unwrap().extract::<String>(py).unwrap()
);
assert_eq!(actual_log, expected_log);
let events: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()).collect();
assert_eq!(events.len() as u64 + sender.dropped_events(), 3);
if !closed && capacity >= 3 {
let [
CallEvent::Started { start_time },
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }),
CallEvent::Succeeded { timing, .. },
] = events.as_slice()
else {
panic!("expected start, provider response and success: {events:?}");
};
assert_eq!(raw.body, "raw");
assert_eq!(*start_time, timing.start_time);
assert!(timing.end_time >= timing.start_time);
}
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, false, 1)]
#[case::transformation(HookScript::FailAfterSuccess, false, 1)]
#[case::terminal_cancellation(HookScript::CancelTerminal, true, 0)]
#[case::terminal_failure(HookScript::FailTerminal, false, 0)]
fn a_terminal_hook_failure_is_reported_once_and_never_replays_callbacks(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] cancelled: bool,
#[case] failure_callbacks: usize,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (options, mut events) = observed(4);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, script),
std::convert::identity,
CallOptions {
asynchronous,
..options
},
);
let error = result.unwrap_err();
assert_eq!(error.is_instance_of::<CancelledError>(py), cancelled);
let events = events();
assert!(matches!(events.first(), Some(CallEvent::Started { .. })));
assert_eq!(
matches!(events.last(), Some(CallEvent::Cancelled { .. })),
cancelled
);
assert_eq!(
matches!(
events.last(),
Some(CallEvent::Failed {
origin: FailureOrigin::Host,
..
})
),
!cancelled
);
let terminal = events.iter().filter(|event| {
matches!(
event,
CallEvent::Failed { .. }
| CallEvent::Succeeded { .. }
| CallEvent::Cancelled { .. }
)
});
assert_eq!(terminal.count(), 1);
assert_eq!(
log.iter()
.filter(|entry| entry.starts_with("failed:"))
.count(),
failure_callbacks
);
assert!(
!log.iter()
.any(|entry| entry.starts_with("succeeded:") || entry == "cancelled")
);
});
}
#[rstest]
#[case::preparation(HookScript::FailBegin, OpScript::Answer)]
#[case::native_decode(HookScript::Plain, OpScript::RejectRequestNatively)]
#[case::python_decode(HookScript::Plain, OpScript::RaiseRequestPython)]
fn startup_failures_release_the_factory_without_constructing_a_machine(
#[from(interpreter)] _python: Interpreter,
#[case] script: HookScript,
#[case] op: OpScript,
#[values(false, true)] asynchronous: bool,
) {
struct Release(Arc<AtomicBool>);
impl Drop for Release {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
Python::attach(|py| {
let constructed = Arc::new(AtomicBool::new(false));
let did_construct = constructed.clone();
let released = Arc::new(AtomicBool::new(false));
let release = Release(released.clone());
let machine: MachineFactory = Box::new(move |request| {
let _release = release;
did_construct.store(true, Ordering::SeqCst);
success_machine()(request)
});
let (result, log) = run_scripted(py, machine, op, script, asynchronous);
assert!(result.is_err());
assert!(!constructed.load(Ordering::SeqCst));
assert!(released.load(Ordering::SeqCst));
assert_eq!(
log.iter()
.filter(|event| event.starts_with("failed:"))
.count(),
1
);
assert_eq!(
log.iter().filter(|event| *event == "adapter.close").count(),
1
);
assert_eq!(log.iter().filter(|event| *event == "host.close").count(), 1);
});
}
#[rstest]
fn success_runs_every_step_in_order_and_returns_the_public_response(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::Plain,
asynchronous,
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:1|sign|rewritten"
);
assert_eq!(
log,
[
"started",
"begin",
"project",
"op:sign",
"before_provider_request",
"response:raw",
"complete",
"after_success",
"succeeded:project:1|sign|rewritten",
"adapter.close",
"host.close",
]
);
});
}
#[rstest]
fn explicitly_closing_a_stream_dispatches_success_for_delivered_chunks(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (options, mut events) = observed(3);
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
CallOptions {
asynchronous,
..options
},
);
if asynchronous {
for method in ["__anext__", "aclose", "aclose"] {
let stop = stream
.call_method0(method)
.unwrap()
.call_method1("send", (py.None(),))
.unwrap_err();
assert!(stop.is_instance_of::<PyStopIteration>(py));
}
} else {
assert_eq!(
stream
.call_method0("__next__")
.unwrap()
.extract::<String>()
.unwrap(),
"first"
);
stream.call_method0("close").unwrap();
stream.call_method0("close").unwrap();
}
assert!(matches!(
events().as_slice(),
[CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
));
assert_eq!(
log.entries(),
[
"started",
"begin",
"opened",
"delivered",
"succeeded:None",
"adapter.close"
]
);
});
}
#[rstest]
fn a_stream_carries_its_head_before_the_first_chunk(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let stream = open_stream(
py,
hooks(&log, HookScript::Plain),
call_options(asynchronous),
);
let hidden: std::collections::HashMap<String, std::collections::HashMap<String, String>> =
stream.getattr("_hidden_params").unwrap().extract().unwrap();
assert_eq!(
hidden["additional_headers"],
std::collections::HashMap::from([("request-id".to_string(), "req_1".to_string())])
);
assert_eq!(log.entries(), ["started", "begin", "opened"]);
assert_eq!(read_all(py, &stream, asynchronous), ["first", "second"]);
});
}
struct Classification {
provider_fails: bool,
op: OpScript,
classifier_fails: bool,
message: &'static str,
native_context: Option<&'static str>,
steps: &'static [&'static str],
}
#[rstest]
#[case::provider_failure(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: false,
message: "classified: provider exploded",
native_context: None,
steps: &["project", "classify:provider exploded"],
})]
#[case::host_op_rejection(Classification {
provider_fails: false,
op: OpScript::RejectNatively,
classifier_fails: false,
message: "classified: op rejected",
native_context: None,
steps: &["project", "op:sign", "classify:op rejected"],
})]
#[case::host_op_python_exception(Classification {
provider_fails: false,
op: OpScript::RaisePython,
classifier_fails: false,
message: "op failed",
native_context: None,
steps: &["project", "op:sign"],
})]
#[case::failing_classifier(Classification {
provider_fails: true,
op: OpScript::Answer,
classifier_fails: true,
message: "classifier failed",
native_context: Some("provider exploded"),
steps: &["project", "classify:provider exploded"],
})]
fn a_native_failure_is_classified_once_and_a_python_exception_is_raised_as_is(
#[from(interpreter)] _python: Interpreter,
#[case] case: Classification,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let Classification {
provider_fails,
op,
classifier_fails,
message,
native_context,
steps,
} = case;
let log = Log::default();
let binding = if classifier_fails {
binding(&log, op).with_failing_classifier()
} else {
binding(&log, op)
};
let machine = if provider_fails {
failing_machine()
} else {
success_machine()
};
let (result, log) = run_composed(
py,
machine,
binding,
hooks(&log, HookScript::Plain),
std::convert::identity,
call_options(asynchronous),
);
let error = result.unwrap_err();
assert_eq!(error.value(py).to_string(), message);
match native_context {
Some(context) => {
assert!(error.is_instance_of::<PyTypeError>(py));
let chained = error.context(py).unwrap();
assert!(chained.is_instance_of::<PyRuntimeError>(py));
assert_eq!(chained.value(py).to_string(), context);
}
None => assert!(error.is_instance_of::<PyValueError>(py)),
}
let expected: Vec<String> = ["started", "begin"]
.iter()
.chain(steps)
.map(|step| step.to_string())
.chain([
format!("failed:Call:{message}"),
"adapter.close".into(),
"host.close".into(),
])
.collect();
assert_eq!(log, expected);
});
}
#[rstest]
fn begin_failures_are_host_failures_without_provider_mapping(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailBegin,
false,
);
assert_eq!(result.unwrap_err().value(py).to_string(), "begin failed");
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:begin failed",
"adapter.close",
"host.close"
]
);
});
}
enum ArgumentPolicy {
Inherit,
Reject(Py<PyBaseException>),
}
impl CallHooks<PythonRuntime> for ArgumentPolicy {
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
match self {
Self::Inherit => arguments.bind(py).set_item("api_key", "inherited"),
Self::Reject(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())),
}
}
}
impl PythonOwned for ArgumentPolicy {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
match self {
Self::Inherit => Ok(()),
Self::Reject(error) => visit.call(error),
}
}
}
#[rstest]
fn resource_setup_uses_prepared_arguments_and_failures_are_terminal(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
#[values(false, true)] fail_setup: bool,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let error = PyValueError::new_err("setup failed").into_value(py);
let setup_error = error.clone_ref(py);
let arguments = PyDict::new(py);
arguments.set_item("api_key", "original").unwrap();
let result = handed(
py,
run_call(
py,
move |py, prepared, request| {
let source: String = prepared.get_item("prepared")?.unwrap().extract()?;
let key: String = prepared.get_item("api_key")?.unwrap().extract()?;
setup_log.push(format!("setup:{source}:{key}"));
if fail_setup {
return Err(PyErr::from_value(setup_error.into_bound(py).into_any()));
}
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
HookChain::new()
.with(hooks(&log, HookScript::RewriteArguments))
.with(ArgumentPolicy::Inherit),
arguments.clone().unbind(),
call_options(asynchronous),
),
asynchronous,
);
if asynchronous {
assert!(log.entries().is_empty());
}
let settled = settled(py, result, asynchronous);
assert_eq!(
arguments
.get_item("api_key")
.unwrap()
.unwrap()
.extract::<String>()
.unwrap(),
"original"
);
assert_eq!(
&log.entries()[..4],
["started", "begin", "project", "setup:hook:inherited"]
);
if fail_setup {
assert!(settled.unwrap_err().value(py).is(error.bind(py)));
assert_eq!(
log.entries(),
[
"started",
"begin",
"project",
"setup:hook:inherited",
"failed:Host:setup failed",
"adapter.close",
"host.close",
]
);
} else {
assert!(settled.is_ok());
assert!(
log.entries()
.iter()
.any(|entry| entry.starts_with("succeeded:"))
);
}
});
}
#[rstest]
fn closing_an_unstarted_call_never_initializes_resources(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let log = Log::default();
let setup_log = log.clone();
let pending = handed(
py,
run_call(
py,
move |_, _, request| {
setup_log.push("setup");
Ok(success_machine()(request))
},
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
PyDict::new(py).unbind(),
call_options(true),
),
true,
)
.unwrap();
assert!(log.entries().is_empty());
pending.call_method0(py, "close").unwrap();
drop(pending);
assert_eq!(log.entries(), ["adapter.close", "host.close"]);
});
}
#[rstest]
fn an_argument_policy_rejection_is_the_callers_error_and_the_machine_never_starts(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let raised = PyValueError::new_err("over budget").into_value(py);
let log = Log::default();
let (result, log) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| {
HookChain::new()
.with(hooks)
.with(ArgumentPolicy::Reject(raised.clone_ref(py)))
},
call_options(asynchronous),
);
assert!(result.unwrap_err().value(py).is(raised.bind(py)));
assert_eq!(
log,
[
"started",
"begin",
"failed:Host:over budget",
"adapter.close",
"host.close"
]
);
});
}
#[rstest]
fn the_host_projects_from_the_keyword_view_the_argument_policy_rewrote(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let log = Log::default();
let (result, _) = run_composed(
py,
success_machine(),
binding(&log, OpScript::Answer),
hooks(&log, HookScript::Plain),
|hooks| HookChain::new().with(hooks).with(ArgumentPolicy::Inherit),
call_options(asynchronous),
);
assert_eq!(
result.unwrap().extract::<String>(py).unwrap(),
"project:2|sign|rewritten"
);
});
}
#[rstest]
fn a_failure_while_finalizing_fails_the_call_instead_of_succeeding(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::Answer,
HookScript::FailAfterSuccess,
asynchronous,
);
assert_eq!(
result.unwrap_err().value(py).to_string(),
"after_success failed"
);
assert_eq!(
&log[log.len() - 4..],
[
"after_success",
"failed:Host:after_success failed",
"adapter.close",
"host.close"
]
);
assert!(!log.iter().any(|entry| entry.starts_with("succeeded")));
});
}
#[rstest]
fn cancellation_ends_the_call_without_terminal_dispatch(
#[from(interpreter)] _python: Interpreter,
#[values(false, true)] asynchronous: bool,
) {
Python::attach(|py| {
let (result, log) = run_scripted(
py,
success_machine(),
OpScript::CancelRequest,
HookScript::Plain,
asynchronous,
);
let error = result.unwrap_err();
assert!(!error.is_instance_of::<PyException>(py));
assert_eq!(
log,
["started", "begin", "project", "adapter.close", "host.close"]
);
});
}

View file

@ -0,0 +1,343 @@
mod support;
use litellm_host::{hooks::CallHooks, lifecycle::Timing};
use litellm_host_python::{
Execution, ExecutionBody, ExecutionStep, HookStep, PythonCallEvent, PythonOwned, PythonRuntime,
Suspension, run_call,
};
use pyo3::{
exceptions::PyBaseException,
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyCFunction, PyDict},
};
use rstest::rstest;
use support::{
Interpreter, Log, OpScript, binding, call_options, handed, install_lifecycle_module,
interpreter, lifecycle_module, success_machine,
};
struct RetainingHost {
retained: Option<Py<PyAny>>,
}
impl ExecutionBody for RetainingHost {
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| Ok(ExecutionStep::Return(py.None())))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.retained)
}
}
#[pyfunction]
fn retaining_coroutine(py: Python<'_>, retained: Py<PyAny>) -> PyResult<Py<Execution>> {
Py::new(
py,
Execution::new(
RetainingHost {
retained: Some(retained),
},
lifecycle_module,
),
)
}
struct AwaitBody(Option<Py<PyAny>>);
impl ExecutionBody for AwaitBody {
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
match self.0.take() {
Some(awaitable) => Ok(ExecutionStep::Suspended(Suspension::Await(awaitable))),
None => result
.expect("selected await completed")
.map(ExecutionStep::Return),
}
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
}
#[pyfunction]
fn await_execution(awaitable: Py<PyAny>) -> Execution {
Execution::new(AwaitBody(Some(awaitable)), lifecycle_module)
}
struct CallingBody(Py<PyAny>);
impl ExecutionBody for CallingBody {
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| self.0.call0(py).map(ExecutionStep::Return))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
}
#[pyfunction]
fn calling_execution(callback: Py<PyAny>) -> Execution {
Execution::new(CallingBody(callback), lifecycle_module)
}
struct ErrorBody(Option<Py<PyBaseException>>);
impl ExecutionBody for ErrorBody {
fn resume(&mut self, _: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| {
Err(PyErr::from_value(
self.0.take().unwrap().into_bound(py).into_any(),
))
})
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
}
#[pyfunction]
fn error_execution(error: Bound<'_, PyBaseException>) -> Execution {
Execution::new(ErrorBody(Some(error.unbind())), lifecycle_module)
}
/// Where a paused call keeps the test's object: in its hooks, or selected as the call's
/// response or exception, with the call paused in the hook that is told about it.
#[derive(Clone, Copy, Debug)]
enum Holder {
Hooks,
Response,
Exception,
}
struct PausingHooks {
retained: Option<Py<PyAny>>,
pause: Option<Py<PyAny>>,
holder: Holder,
}
impl CallHooks<PythonRuntime> for PausingHooks {
fn transform_response(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
_: Timing,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
match self.holder {
Holder::Hooks => Ok(HookStep::Ready(response)),
Holder::Response => Ok(HookStep::Ready(self.retained.take().unwrap())),
Holder::Exception => Err(PyErr::from_value(
self.retained.take().unwrap().into_bound(py),
)),
}
}
fn on_event(
&mut self,
_: Python<'_>,
event: PythonCallEvent<'_>,
) -> PyResult<HookStep<Self, ()>> {
let pauses_here = match self.holder {
Holder::Hooks => true,
Holder::Response => matches!(event, PythonCallEvent::Succeeded { .. }),
Holder::Exception => matches!(event, PythonCallEvent::Failed { .. }),
};
Ok(match self.pause.take_if(|_| pauses_here) {
Some(pause) => {
HookStep::Await(pause, |_, _, answer| answer.map(|_| HookStep::Ready(())))
}
None => HookStep::Ready(()),
})
}
}
impl PythonOwned for PausingHooks {
fn close(&mut self, _: Python<'_>) {
self.retained = None;
self.pause = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.retained)?;
visit.call(&self.pause)
}
}
fn paused_call(
py: Python<'_>,
retained: Py<PyAny>,
pause: Py<PyAny>,
holder: Holder,
) -> PyResult<Py<PyAny>> {
let machine = success_machine();
handed(
py,
run_call(
py,
move |_, _, request| Ok(machine(request)),
binding(&Log::default(), OpScript::Answer),
PausingHooks {
retained: Some(retained),
pause: Some(pause),
holder,
},
PyDict::new(py).unbind(),
call_options(true),
),
true,
)
}
#[rstest]
fn every_suspension_is_awaited_inline_on_the_callers_task(
#[from(interpreter)] _python: Interpreter,
) {
Python::attach(|py| {
let module = install_lifecycle_module(py);
let locals = PyDict::new(py);
locals
.set_item("drive", module.getattr("drive").unwrap())
.unwrap();
locals
.set_item(
"await_execution",
wrap_pyfunction!(await_execution, py).unwrap(),
)
.unwrap();
locals
.set_item(
"calling_execution",
wrap_pyfunction!(calling_execution, py).unwrap(),
)
.unwrap();
let probe = std::ffi::CString::new(include_str!("lifecycle.py")).unwrap();
py.run(&probe, Some(&locals), Some(&locals)).unwrap();
});
}
#[rstest]
fn retained_exception_frames_are_collectable(#[from(interpreter)] _python: Interpreter) {
Python::attach(|py| {
let locals = PyDict::new(py);
locals
.set_item(
"error_execution",
wrap_pyfunction!(error_execution, py).unwrap(),
)
.unwrap();
py.run(
pyo3::ffi::c_str!(
r#"
import gc
import weakref
class Retained:
pass
def cycle():
retained = Retained()
try:
raise ValueError('retained traceback')
except ValueError as error:
retained.owner = error_execution(error)
return weakref.ref(retained)
reference = cycle()
gc.collect()
assert reference() is None
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
});
}
#[rstest]
fn cycles_through_a_retained_execution_are_collectable(#[from(interpreter)] _python: Interpreter) {
Python::attach(|py| {
let locals = PyDict::new(py);
locals
.set_item(
"retaining_coroutine",
wrap_pyfunction!(retaining_coroutine, py).unwrap(),
)
.unwrap();
py.run(
pyo3::ffi::c_str!(
r#"
import gc
import weakref
class Retained:
pass
def cycle():
retained = Retained()
coroutine = retaining_coroutine(retained)
retained.coroutine = coroutine
return weakref.ref(retained)
retained_ref = cycle()
gc.collect()
assert retained_ref() is None
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
});
}
#[rstest]
#[case::held_by_the_hooks(Holder::Hooks)]
#[case::selected_as_the_response(Holder::Response)]
#[case::selected_as_the_exception(Holder::Exception)]
fn cycles_through_a_call_suspended_in_a_hook_are_collectable(
#[from(interpreter)] _python: Interpreter,
#[case] holder: Holder,
) {
Python::attach(|py| {
let locals = PyDict::new(py);
let call = PyCFunction::new_closure(py, None, None, move |args, _| {
let (retained, pause) = args.extract::<(Py<PyAny>, Py<PyAny>)>()?;
paused_call(args.py(), retained, pause, holder)
})
.unwrap();
locals.set_item("paused_call", call).unwrap();
py.run(
pyo3::ffi::c_str!(
r#"
import gc
import weakref
class Retained(Exception):
pass
class Pause:
def __await__(self):
yield
def cycle():
retained = Retained()
coroutine = paused_call(retained, Pause())
coroutine.send(None)
retained.coroutine = coroutine
return weakref.ref(retained)
reference = cycle()
gc.collect()
assert reference() is None
"#
),
Some(&locals),
Some(&locals),
)
.unwrap();
});
}

View file

@ -504,7 +504,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
}
HookStep::Await(awaitable, resume) => {
self.resume = Some(resume);
Ok(litellm_host_python::ExecutionStep::Await(awaitable))
Ok(litellm_host_python::ExecutionStep::Suspended(
litellm_host_python::Suspension::Await(awaitable),
))
}
}
})
@ -518,9 +520,9 @@ impl litellm_host_python::ExecutionBody for Preparing {
fn lifecycle(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
let source =
std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/lifecycle.py"))
std::ffi::CString::new(include_str!("../../../../litellm/rust_bridge/execution.py"))
.unwrap();
PyModule::from_code(py, &source, c"lifecycle.py", c"hook_chain_test_lifecycle")
PyModule::from_code(py, &source, c"execution.py", c"hook_chain_test_lifecycle")
}
#[rstest]
@ -556,9 +558,16 @@ second = Hooks('b', invoke)
},
lifecycle,
);
locals
.set_item("call", execution.into_coroutine(py).unwrap())
let module = lifecycle(py).unwrap();
let call = module
.getattr("drive")
.unwrap()
.call1((
Py::new(py, execution).unwrap(),
module.getattr("AsyncStream").unwrap(),
))
.unwrap();
locals.set_item("call", call).unwrap();
py.run(
c"
async def exercise():

View file

@ -36,11 +36,17 @@ async def exercise():
releaser = asyncio.create_task(release())
execution = await_execution(CustomAwaitable())
try:
execution.resume_value(None)
execution.throw(ValueError("early"))
except RuntimeError:
pass
else:
raise AssertionError("resumed an unstarted execution")
try:
execution.send(1)
except TypeError:
pass
else:
raise AssertionError("sent a value into an unstarted execution")
wrapped = drive(execution)
try:
wrapped.send(1)
@ -129,7 +135,7 @@ async def exercise():
def reenter():
try:
active.start()
active.send(None)
except RuntimeError as error:
observed.append(str(error))
return result
@ -151,7 +157,7 @@ async def exercise():
execution = calling_execution(callback)
callback.owner = execution
if started:
assert execution.start().value is result
assert execution.send(None).value is result
return weakref.ref(callback)
for started in (False, True):

View file

@ -0,0 +1,684 @@
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
use litellm_host::{
call::{CallOutput, HostedMachine, hosted_call},
hooks::CallHooks,
interceptors::{Interceptors, RawResponse, RequestContext, WireRequest},
lifecycle::{ExecutionEvent, Timing},
machine::{CallMachine, MachineFault},
protocol::{Protocol, Reply},
};
use litellm_host_python::{
CallOptions, HookStep, InvokeError, PythonBinding, PythonCallEvent, PythonCallHooks,
PythonHostCalls, PythonOwned, PythonRuntime, run_call,
};
use pyo3::{
exceptions::{PyTypeError, PyValueError, asyncio::CancelledError},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyString},
};
use rstest::fixture;
static PYTHON_GLOBALS: Mutex<()> = Mutex::new(());
/// One initialized interpreter with the real lifecycle protocol installed, held for the
/// length of a test so tests sharing the interpreter run one at a time.
pub struct Interpreter(MutexGuard<'static, ()>);
#[fixture]
pub fn interpreter() -> Interpreter {
let guard = PYTHON_GLOBALS
.lock()
.unwrap_or_else(PoisonError::into_inner);
Python::initialize();
Python::attach(|py| {
py.import("asyncio").unwrap();
install_lifecycle_module(py);
});
Interpreter(guard)
}
pub fn lifecycle_module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
py.import("host_test_lifecycle")
}
/// Hands a call's execution to its caller the way the product entrypoints do: an async call
/// gets the coroutine from `drive`, a sync call the value or stream `run` settles to.
pub fn handed(
py: Python<'_>,
execution: PyResult<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
let execution = execution?;
let module = lifecycle_module(py)?;
let handed = if asynchronous {
module.getattr("drive")?.call1((execution,))?
} else {
module
.getattr("run")?
.call1((execution, module.getattr("SyncStream")?))?
};
Ok(handed.unbind())
}
pub fn call_options(asynchronous: bool) -> CallOptions {
CallOptions::new(asynchronous, lifecycle_module)
}
pub fn install_lifecycle_module(py: Python<'_>) -> Bound<'_, PyModule> {
let source = std::ffi::CString::new(concat!(
include_str!("../../../../../litellm/rust_bridge/execution.py"),
"\n_drive = drive\n\n\nasync def drive(execution):\n return await _drive(execution, AsyncStream)\n",
))
.unwrap();
PyModule::from_code(
py,
&source,
pyo3::ffi::c_str!("execution.py"),
pyo3::ffi::c_str!("host_test_lifecycle"),
)
.unwrap()
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Error(pub String);
impl std::fmt::Display for Error {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str(&self.0)
}
}
impl From<MachineFault> for Error {
fn from(fault: MachineFault) -> Self {
Self(format!("{fault:?}"))
}
}
pub struct Synthetic;
impl Protocol for Synthetic {
type Response = String;
type Error = Error;
type Request = String;
type HostCall = (&'static str, Reply<String>);
type Chunk = std::convert::Infallible;
type StreamHead = std::convert::Infallible;
}
pub fn wire() -> WireRequest {
WireRequest {
url: "https://example.invalid".into(),
headers: Vec::new(),
body: serde_json::json!({}),
}
}
pub fn context() -> RequestContext {
RequestContext {
model: "model".into(),
custom_llm_provider: "provider".into(),
optional_params: serde_json::json!({}),
secret_fields: Vec::new(),
api_key: None,
}
}
#[derive(Clone, Default)]
pub struct Log(Arc<Mutex<Vec<String>>>);
impl Log {
pub fn push(&self, entry: impl Into<String>) {
self.0.lock().unwrap().push(entry.into());
}
pub fn entries(&self) -> Vec<String> {
self.0.lock().unwrap().clone()
}
}
#[derive(Clone, Copy, Debug)]
pub enum OpScript {
Answer,
RaisePython,
RejectNatively,
RejectRequestNatively,
RaiseRequestPython,
CancelRequest,
AwaitAnswer,
AwaitFailure,
}
/// The fake route's public exception, kept as a value so a test sees what `classify`
/// produced before the driver raises it.
#[derive(Debug, PartialEq, Eq)]
pub struct Classified(pub String);
impl From<Classified> for PyErr {
fn from(classified: Classified) -> Self {
PyValueError::new_err(format!("classified: {}", classified.0))
}
}
pub struct SyntheticBinding {
log: Log,
op: OpScript,
classifier_fails: bool,
pending_reply: Option<Reply<String>>,
}
pub fn binding(log: &Log, op: OpScript) -> SyntheticBinding {
SyntheticBinding {
log: log.clone(),
op,
classifier_fails: false,
pending_reply: None,
}
}
impl SyntheticBinding {
pub fn with_failing_classifier(self) -> Self {
Self {
classifier_fails: true,
..self
}
}
fn answer(&self, value: impl FnOnce() -> String) -> Result<String, InvokeError<Error>> {
match self.op {
OpScript::Answer | OpScript::AwaitAnswer | OpScript::AwaitFailure => Ok(value()),
OpScript::RaisePython | OpScript::RaiseRequestPython => {
Err(PyValueError::new_err("op failed").into())
}
OpScript::RejectNatively | OpScript::RejectRequestNatively => {
Err(InvokeError::Native(Error("op rejected".into())))
}
OpScript::CancelRequest => Err(CancelledError::new_err(()).into()),
}
}
}
impl PythonBinding for SyntheticBinding {
type Protocol = Synthetic;
type Failure = Classified;
fn decode_request(
&mut self,
_: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<String, InvokeError<Error>> {
self.log.push("project");
match self.op {
OpScript::RejectRequestNatively
| OpScript::RaiseRequestPython
| OpScript::CancelRequest => self.answer(String::new),
_ => Ok(format!("project:{}", arguments.len())),
}
}
fn encode_stream_head(
&mut self,
_: Python<'_>,
head: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match head {}
}
fn encode_chunk(
&mut self,
_: Python<'_>,
chunk: std::convert::Infallible,
) -> PyResult<Py<PyAny>> {
match chunk {}
}
fn encode_response(&mut self, py: Python<'_>, response: String) -> PyResult<Py<PyAny>> {
self.log.push("complete");
Ok(PyString::new(py, &response).into_any().unbind())
}
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
self.log.push(format!("classify:{error}"));
if self.classifier_fails {
return Err(PyTypeError::new_err("classifier failed"));
}
Ok(Classified(error.0))
}
fn host_error(error: &PyErr) -> Error {
Error(error.to_string())
}
}
impl PythonHostCalls<Synthetic> for SyntheticBinding {
fn handle_host_call(
&mut self,
_: Python<'_>,
(op, reply): (&'static str, Reply<String>),
) -> Result<(), InvokeError<Error>> {
self.log.push(format!("op:{op}"));
self.answer(|| op.to_string())
.map(|answer| reply.send(answer))
}
fn begin_host_call(
&mut self,
py: Python<'_>,
(op, reply): (&'static str, Reply<String>),
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
if !matches!(self.op, OpScript::AwaitAnswer | OpScript::AwaitFailure) {
return self.handle_host_call(py, (op, reply)).map(|()| None);
}
self.pending_reply = Some(reply);
let module = PyModule::from_code(
py,
pyo3::ffi::c_str!(
"async def answer(fail):\n if fail:\n raise LookupError('async host failed')\n return 'awaited'\n"
),
pyo3::ffi::c_str!("host_op.py"),
pyo3::ffi::c_str!("host_op"),
)?;
Ok(Some(
module
.getattr("answer")?
.call1((matches!(self.op, OpScript::AwaitFailure),))?
.unbind(),
))
}
fn resume_host_call(
&mut self,
py: Python<'_>,
result: PyResult<Py<PyAny>>,
) -> Result<Option<Py<PyAny>>, InvokeError<Error>> {
let answer = result?.extract::<String>(py)?;
self.pending_reply.take().unwrap().send(answer);
Ok(None)
}
}
impl PythonOwned for SyntheticBinding {
fn close(&mut self, _: Python<'_>) {
self.log.push("host.close");
}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
#[derive(Clone, Copy, Debug)]
pub enum HookScript {
Plain,
RewriteArguments,
ObserveArguments,
FailBegin,
ReplaceResponse,
FailAfterSuccess,
CancelTerminal,
FailTerminal,
}
pub struct SyntheticHooks {
log: Log,
script: HookScript,
}
pub fn hooks(log: &Log, script: HookScript) -> SyntheticHooks {
SyntheticHooks {
log: log.clone(),
script,
}
}
impl CallHooks<PythonRuntime> for SyntheticHooks {
fn prepare_arguments(
&mut self,
_: Python<'_>,
arguments: Py<PyDict>,
_: f64,
) -> PyResult<HookStep<Self, Py<PyDict>>> {
self.log.push("begin");
if matches!(self.script, HookScript::FailBegin) {
return Err(PyValueError::new_err("begin failed"));
}
if matches!(self.script, HookScript::RewriteArguments) {
let prepared = Python::attach(|py| -> PyResult<Py<PyDict>> {
let prepared = arguments.bind(py).copy()?;
prepared.set_item("prepared", "hook")?;
prepared.set_item("api_key", "hook-key")?;
Ok(prepared.unbind())
})?;
return Ok(HookStep::Ready(prepared));
}
Ok(HookStep::Ready(arguments))
}
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
if matches!(self.script, HookScript::ObserveArguments) {
let value: String = arguments
.bind(py)
.get_item("prepared")?
.unwrap()
.extract()?;
self.log.push(format!("adopted:{value}"));
}
Ok(())
}
fn before_provider_request(
&mut self,
_: Python<'_>,
wire: Box<WireRequest>,
_: &RequestContext,
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
self.log.push("before_provider_request");
Ok(HookStep::Ready(Box::new(WireRequest {
url: "rewritten".into(),
..*wire
})))
}
fn transform_response(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
_: Timing,
) -> PyResult<HookStep<Self, Py<PyAny>>> {
self.log.push("after_success");
match self.script {
HookScript::ReplaceResponse => Ok(HookStep::Ready(
"replaced".into_pyobject(py)?.into_any().unbind(),
)),
HookScript::FailAfterSuccess => Err(PyValueError::new_err("after_success failed")),
_ => Ok(HookStep::Ready(response)),
}
}
fn on_event(
&mut self,
py: Python<'_>,
event: PythonCallEvent<'_>,
) -> PyResult<HookStep<Self, ()>> {
let terminal = matches!(
event,
PythonCallEvent::Succeeded { .. } | PythonCallEvent::Failed { .. }
);
if terminal && matches!(self.script, HookScript::CancelTerminal) {
return Err(CancelledError::new_err("callback cancelled"));
}
if terminal && matches!(self.script, HookScript::FailTerminal) {
return Err(PyValueError::new_err("callback failed"));
}
self.log.push(match event {
PythonCallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "cache_hit".into(),
PythonCallEvent::Started { .. } => "started".into(),
PythonCallEvent::Cancelled { .. } => "cancelled".into(),
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
format!("response:{}", raw.body)
}
PythonCallEvent::Succeeded { response, .. } => {
format!("succeeded:{}", response.bind(py))
}
PythonCallEvent::Failed { origin, error, .. } => {
format!("failed:{origin:?}:{}", error.value(py))
}
});
Ok(HookStep::Ready(()))
}
fn on_stream_open(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
self.log.push("opened");
Ok(())
}
fn on_stream_chunk(&mut self, _: Python<'_>, _: &Py<PyAny>) -> PyResult<()> {
self.log.push("delivered");
Ok(())
}
}
impl PythonOwned for SyntheticHooks {
fn close(&mut self, _: Python<'_>) {
self.log.push("adapter.close");
}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
pub type MachineFactory = Box<dyn FnOnce(String) -> CallMachine<Synthetic> + Send + Sync>;
/// Answers to projection, to the route op and to `before_provider_request` all reach the
/// response, so a driver that misroutes a reply changes what the call returns.
pub fn success_machine() -> MachineFactory {
Box::new(move |projected| {
CallMachine::new(None, move |host| {
Box::pin(async move {
let signed = host.services.call(|reply| ("sign", reply)).await?;
let wire = host
.interceptors
.before_provider_request(wire(), context())
.await?;
host.interceptors
.after_provider_response(RawResponse { body: "raw".into() })
.await?;
Ok(format!("{projected}|{signed}|{}", wire.url))
})
})
})
}
pub fn failing_machine() -> MachineFactory {
Box::new(move |_| {
CallMachine::new(None, |_| {
Box::pin(async move { Err(Error("provider exploded".into())) })
})
})
}
/// Drives an asynchronous call's coroutine to its settled value from a sync test. A
/// synchronous call is already settled.
pub fn settled(
py: Python<'_>,
result: PyResult<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Py<PyAny>> {
if !asynchronous {
return result;
}
let coroutine = result?;
let completed = coroutine
.call_method1(py, "send", (py.None(),))
.unwrap_err();
if !completed.is_instance_of::<pyo3::exceptions::PyStopIteration>(py) {
return Err(completed);
}
completed.value(py).getattr("value").map(Bound::unbind)
}
pub fn run_scripted(
py: Python<'_>,
machine: MachineFactory,
op: OpScript,
script: HookScript,
asynchronous: bool,
) -> (PyResult<Py<PyAny>>, Vec<String>) {
let log = Log::default();
run_composed(
py,
machine,
binding(&log, op),
hooks(&log, script),
std::convert::identity,
call_options(asynchronous),
)
}
pub fn run_composed<L: PythonCallHooks + 'static>(
py: Python<'_>,
machine: MachineFactory,
binding: SyntheticBinding,
hooks: SyntheticHooks,
compose: impl FnOnce(SyntheticHooks) -> L,
options: CallOptions,
) -> (PyResult<Py<PyAny>>, Vec<String>) {
let asynchronous = options.asynchronous;
let log = binding.log.clone();
let arguments = PyDict::new(py);
arguments.set_item("model", "m").unwrap();
let result = handed(
py,
run_call(
py,
move |_, _, request| Ok(machine(request)),
binding,
compose(hooks),
arguments.unbind(),
options,
),
asynchronous,
);
(settled(py, result, asynchronous), log.entries())
}
pub struct Streaming;
impl Protocol for Streaming {
type Response = ();
type Error = Error;
type Request = ();
type HostCall = std::convert::Infallible;
type Chunk = &'static str;
type StreamHead = Vec<(&'static str, &'static str)>;
}
pub struct StreamingBinding;
impl PythonBinding for StreamingBinding {
type Protocol = Streaming;
type Failure = Classified;
fn decode_request(
&mut self,
_: Python<'_>,
_: &Bound<'_, PyDict>,
) -> Result<(), InvokeError<Error>> {
Ok(())
}
fn encode_stream_head(
&mut self,
py: Python<'_>,
head: Vec<(&'static str, &'static str)>,
) -> PyResult<Py<PyAny>> {
let headers = PyDict::new(py);
for (name, value) in head {
headers.set_item(name, value)?;
}
let hidden = PyDict::new(py);
hidden.set_item("additional_headers", headers)?;
Ok(hidden.into_any().unbind())
}
fn encode_chunk(&mut self, py: Python<'_>, chunk: &'static str) -> PyResult<Py<PyAny>> {
Ok(PyString::new(py, chunk).into_any().unbind())
}
fn encode_response(&mut self, py: Python<'_>, (): ()) -> PyResult<Py<PyAny>> {
Ok(py.None())
}
fn map_error(&self, _: Python<'_>, error: Error) -> PyResult<Classified> {
Ok(Classified(error.0))
}
fn host_error(error: &PyErr) -> Error {
Error(error.to_string())
}
}
impl PythonHostCalls<Streaming> for StreamingBinding {
fn handle_host_call(
&mut self,
_: Python<'_>,
op: std::convert::Infallible,
) -> Result<(), InvokeError<Error>> {
match op {}
}
}
impl PythonOwned for StreamingBinding {
fn close(&mut self, _: Python<'_>) {}
fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(())
}
}
pub fn streaming_machine() -> impl FnOnce(()) -> HostedMachine<Streaming> + Send + Sync {
move |()| {
hosted_call((), None, |(), _, _, _observations| async {
Ok(CallOutput::Stream {
head: vec![("request-id", "req_1")],
chunks: Box::pin(futures_util::stream::iter([Ok("first"), Ok("second")])),
})
})
}
}
/// Starts a streaming call and returns the caller's stream object.
pub fn open_stream(
py: Python<'_>,
hooks: SyntheticHooks,
options: CallOptions,
) -> Bound<'_, PyAny> {
let asynchronous = options.asynchronous;
let handed = handed(
py,
run_call(
py,
|_, _, request| Ok(streaming_machine()(request)),
StreamingBinding,
hooks,
PyDict::new(py).unbind(),
options,
),
asynchronous,
)
.unwrap();
if asynchronous {
let stop = handed.call_method1(py, "send", (py.None(),)).unwrap_err();
stop.value(py).getattr("value").unwrap()
} else {
handed.into_bound(py)
}
}
/// Drives an `AsyncStream` or `SyncStream` to completion from a sync test.
pub fn read_all(py: Python<'_>, stream: &Bound<'_, PyAny>, asynchronous: bool) -> Vec<String> {
if !asynchronous {
return stream
.try_iter()
.unwrap()
.map(|chunk| chunk.unwrap().extract().unwrap())
.collect();
}
std::iter::from_fn(|| {
let stop = stream
.call_method0("__anext__")
.unwrap()
.call_method1("send", (py.None(),))
.unwrap_err();
if stop.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py) {
return None;
}
assert!(stop.is_instance_of::<pyo3::exceptions::PyStopIteration>(py));
Some(stop.value(py).getattr("value").unwrap().extract().unwrap())
})
.collect()
}

View file

@ -42,10 +42,32 @@ pub enum ResultSource {
Cache { key: String },
}
/// What the call cost, settled once by the route. A host records it; it never re-derives it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Cost {
/// The amount the provider itself billed, taken from its response.
Reported { amount: serde_json::Number },
/// Served from the cache, so nothing was billed.
CacheHit,
/// No provider figure: the host estimates from usage and its price table.
Deferred,
}
impl Cost {
pub fn settle(source: &ResultSource, reported: Option<serde_json::Number>) -> Self {
match (source, reported) {
(ResultSource::Cache { .. }, _) => Self::CacheHit,
(ResultSource::Provider, Some(amount)) => Self::Reported { amount },
(ResultSource::Provider, None) => Self::Deferred,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ExecutionFacts {
pub provider: ProviderIdentity,
pub source: ResultSource,
pub cost: Cost,
}
pub trait Interceptors<E>: Send + Sync {

View file

@ -138,7 +138,9 @@ pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
}
let value = value.trim();
value.len() > 7
&& value[..7].eq_ignore_ascii_case("bearer ")
&& value
.get(..7)
.is_some_and(|prefix| prefix.eq_ignore_ascii_case("bearer "))
&& !value[7..].trim().is_empty()
})
}
@ -279,31 +281,19 @@ mod tests {
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 abc".to_string()
)]));
assert!(!has_bearer_auth(&[(
"Authorization".to_string(),
"Bearer ".to_string()
)]));
assert!(!has_bearer_auth(&[(
"authorization".to_string(),
String::new()
)]));
assert!(!has_bearer_auth(&[(
"Authorization".to_string(),
"Basic abc".to_string()
)]));
assert!(!has_bearer_auth(&[(
"x-api-key".to_string(),
"abc".to_string()
)]));
#[rstest::rstest]
#[case::bearer("Authorization", "Bearer abc", true)]
#[case::lowercase("authorization", "bearer abc", true)]
#[case::blank_token("Authorization", "Bearer ", false)]
#[case::empty("authorization", "", false)]
#[case::basic("Authorization", "Basic abc", false)]
#[case::different_header("x-api-key", "abc", false)]
#[case::non_ascii("authorization", "invali\u{00e9}auth", false)]
fn bearer_detection_requires_a_non_empty_token(
#[case] name: &str,
#[case] value: &str,
#[case] expected: bool,
) {
assert_eq!(has_bearer_auth(&[(name.into(), value.into())]), expected);
}
}

View file

@ -140,7 +140,7 @@ impl Default for HttpSettings {
proxies: EnvironmentProxies::default(),
connect_timeout: Duration::from_secs(10),
tcp_keepalive: None,
pool_idle_timeout: Duration::from_secs(120),
pool_idle_timeout: Duration::from_secs(15),
}
}
}
@ -330,7 +330,7 @@ mod tests {
&[("AIOHTTP_TCP_KEEPIDLE", "45")],
))]);
assert_eq!(settings.tcp_keepalive, None);
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120));
assert_eq!(settings.pool_idle_timeout, Duration::from_secs(15));
}
fn proxy_flags(

View file

@ -26,6 +26,7 @@ serde_json = { workspace = true, features = ["preserve_order"] }
tracing.workspace = true
[dev-dependencies]
litellm-framer = { workspace = true, features = ["sse"] }
litellm-cache-memory.workspace = true
litellm-host-native.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }

View file

@ -1,14 +1,21 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_http::request::string_headers as shared_string_headers;
pub(super) use litellm_http::request::truncate_error_body;
use litellm_llms::tencent::messages::transformation::TENCENT_MESSAGES_CONFIG;
use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
base_llm::messages::transformation::BaseMessagesConfig,
bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
deepseek::messages::transformation::DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG,
edenai::messages::transformation::EDENAI_MESSAGES_CONFIG,
minimax::messages::transformation::MINIMAX_MESSAGES_CONFIG,
vertex_ai::messages::transformation::VERTEX_ANTHROPIC_MESSAGES_CONFIG,
};
use litellm_llms::{
bedrock::claude_platform::messages::transformation::BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG,
bedrock_mantle::messages::transformation::BEDROCK_MANTLE_MESSAGES_CONFIG,
};
use serde_json::{Map, Value};
use super::Error;
@ -34,14 +41,25 @@ impl MessagesProvider {
/// Python's `get_provider_anthropic_messages_config`: Vertex AI serves only its Claude
/// partner models on this route.
pub(crate) fn messages_provider(provider: LlmProviders, model: &str) -> Option<MessagesProvider> {
use litellm_llms::openrouter::messages::transformation::OPENROUTER_MESSAGES_CONFIG;
let config: &'static dyn BaseMessagesConfig = match provider {
LlmProviders::Tencent => &TENCENT_MESSAGES_CONFIG,
LlmProviders::Anthropic => &ANTHROPIC_MESSAGES_CONFIG,
LlmProviders::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG,
LlmProviders::Bedrock if model.starts_with("claude_platform/") => {
&BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
}
LlmProviders::Bedrock if model.starts_with("mantle/") => &BEDROCK_MANTLE_MESSAGES_CONFIG,
LlmProviders::Bedrock => &BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
LlmProviders::BedrockMantle => &BEDROCK_MANTLE_MESSAGES_CONFIG,
LlmProviders::Deepseek => &DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG,
LlmProviders::Edenai => &EDENAI_MESSAGES_CONFIG,
LlmProviders::Minimax => &MINIMAX_MESSAGES_CONFIG,
LlmProviders::VertexAi if model.to_ascii_lowercase().contains("claude") => {
&VERTEX_ANTHROPIC_MESSAGES_CONFIG
}
LlmProviders::Openrouter => &OPENROUTER_MESSAGES_CONFIG,
_ => return None,
};
Some(MessagesProvider { provider, config })
@ -64,10 +82,13 @@ mod tests {
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
#[case::tencent("tencent")]
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[case::bedrock("bedrock")]
#[case::deepseek("deepseek")]
#[case::edenai("edenai")]
#[case::minimax("minimax")]
#[case::vertex_ai("vertex_ai")]
fn provider_keeps_its_python_name(#[case] name: &str) {
let provider =

View file

@ -7,7 +7,7 @@ use litellm_http::transport::Error as TransportError;
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
messages::{
streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
streaming::{self, ByteStream, StreamDecoder},
transformation::BaseMessagesConfig,
},
};
@ -25,7 +25,7 @@ use litellm_inference::{context::CallContext, outbound::outbound_request};
pub(super) struct ProviderCall {
pub identity: ProviderIdentity,
pub wire: WireRequest,
provider: super::common_utils::MessagesProvider,
pub(super) provider: super::common_utils::MessagesProvider,
signer: Option<litellm_auth_aws::SigV4Signer>,
timeout: Option<Duration>,
stream: bool,
@ -245,7 +245,7 @@ fn decoded_chunks(
let Some(event) = events.try_next().await? else {
return Ok(None);
};
let chunk = encode_anthropic_sse(&event)?;
let chunk = streaming::encode(&event)?;
log_chunk(provider, "client_response", &chunk);
Ok(Some((chunk, events)))
})
@ -271,7 +271,6 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
#[cfg(test)]
mod tests {
use litellm_llms::base_llm::messages::streaming::anthropic_sse_event_stream;
use rstest::rstest;
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
@ -300,7 +299,7 @@ mod tests {
.await
.unwrap();
let MessagesCallResponse::Stream { mut chunks, .. } =
streaming_response(response, Some(anthropic_sse_event_stream), "test")
streaming_response(response, Some(streaming::event_stream), "test")
else {
panic!("a streaming response returns chunks");
};

View file

@ -7,7 +7,7 @@ mod types;
use futures_util::FutureExt;
use litellm_auth::AuthServices;
use litellm_host::interceptors::{ExecutionFacts, Interceptors, ResultSource};
use litellm_host::interceptors::{Cost, ExecutionFacts, Interceptors, ResultSource};
use litellm_inference::{caching::CallCache, context::CallContext};
use litellm_secrets::source::SecretSource;
@ -87,6 +87,7 @@ impl MessagesRoute {
&request.wire,
);
let identity = request.identity.clone();
let config = request.provider.config();
let (output, source) = match cache.lookup().await {
Some(hit) => hit,
None => (
@ -94,9 +95,14 @@ impl MessagesRoute {
ResultSource::Provider,
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => config.reported_cost(response),
MessagesCallResponse::Stream { .. } => None,
};
context
.result_ready(ExecutionFacts {
provider: identity,
cost: Cost::settle(&source, reported),
source: source.clone(),
})
.await?;

View file

@ -76,6 +76,10 @@ fn prepare_provider_request(
shaping,
..
} = call;
let connection_params = litellm_router_types::LitellmParams {
api_base: api_base.clone().or(litellm_params.api_base.clone()),
..litellm_params
};
let config = provider.config();
let env_lookup = |key: &str| secrets.get(key);
@ -95,25 +99,22 @@ fn prepare_provider_request(
let forwarded = string_headers(Some(
extra_headers.into_iter().flatten().chain(scoped).collect(),
))?;
let validated = config.validate_environment(
let forwarded = config.request_headers(
with_default_headers(forwarded, config.default_headers()),
&transformed,
);
let environment = config.validate_environment(
forwarded,
api_key.as_deref(),
&transformed.model,
&litellm_params,
&connection_params,
&env_lookup,
)?;
let environment = ValidatedEnvironment {
headers: config.request_headers(
with_default_headers(validated.headers, config.default_headers()),
&transformed,
),
auth: validated.auth,
};
let url = config.get_complete_url(
api_base.as_deref(),
&transformed.model,
&litellm_params,
&connection_params,
transformed.params.stream == Some(true),
&env_lookup,
)?;
@ -475,4 +476,44 @@ mod tests {
))
);
}
#[rstest]
#[case::mantle(
"bedrock_mantle/claude-test",
"https://bedrock-mantle.us-west-2.api.aws",
"/anthropic/v1/messages",
"bedrock"
)]
#[case::platform(
"bedrock/claude_platform/claude-test",
"https://aws-external-anthropic.us-west-2.api.aws",
"/v1/messages",
"aws-external-anthropic"
)]
fn explicit_endpoint_region_is_used_by_the_signer(
shaping: MessagesShaping,
#[case] model: &str,
#[case] base: &str,
#[case] path: &str,
#[case] service: &str,
) {
let prepared = prepare_with_secrets(MessagesCall {
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
api_key: None,
api_base: Some(base.into()),
custom_llm_provider: None,
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
}, &|name: &str| match name {
"AWS_REGION_NAME" => Some("us-east-1".into()),
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
_ => None,
}).unwrap();
assert_eq!(prepared.url, format!("{base}{path}"));
assert!(
matches!(prepared.environment.auth, litellm_llms::base_llm::auth::AuthScheme::AwsSigV4 { region, service: actual, .. } if region == "us-west-2" && actual == service)
);
}
}

View file

@ -0,0 +1,323 @@
use std::{sync::Arc, time::Duration};
use futures_util::TryStreamExt;
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use wiremock::ResponseTemplate;
#[path = "../support/mod.rs"]
mod support;
use support::*;
const MODEL: &str = "claude-test";
#[fixture]
fn call() -> MessagesCall {
MessagesCall {
body: serde_json::from_value(json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})).unwrap(),
api_key: None,
api_base: None,
custom_llm_provider: None,
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
}
}
fn with_fields(call: MessagesCall, fields: Value) -> MessagesCall {
let body = serde_json::to_value(call.body).unwrap();
MessagesCall {
body: serde_json::from_value(Value::Object(
body.as_object()
.unwrap()
.clone()
.into_iter()
.chain(fields.as_object().unwrap().clone())
.collect(),
))
.unwrap(),
..call
}
}
fn with_model(call: MessagesCall, model: &str) -> MessagesCall {
MessagesCall {
body: MessagesRequest {
model: model.into(),
..call.body
},
..call
}
}
fn message_response() -> ResponseTemplate {
json_response(
json!({"id": "msg_1", "type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}], "model": MODEL, "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 2}}),
)
}
async fn run_with(
secrets: Arc<RecordingSecrets>,
call: MessagesCall,
) -> Result<MessagesCallResponse, Error> {
messages_route(secrets).execute(call, &(), None).await
}
async fn run(call: MessagesCall) -> Result<MessagesCallResponse, Error> {
run_with(Arc::new(RecordingSecrets::empty()), call).await
}
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.unwrap() {
MessagesCallResponse::Complete(message) => *message,
MessagesCallResponse::Stream { .. } => panic!("expected a completed message"),
}
}
#[rstest]
#[case::platform(
"bedrock/claude_platform/claude-test",
"/v1/messages",
"x-api-key",
"key"
)]
#[case::mantle(
"bedrock_mantle/anthropic.claude-test",
"/anthropic/v1/messages",
"authorization",
"Bearer key"
)]
#[case::mantle_alias(
"bedrock/mantle/anthropic.claude-test",
"/anthropic/v1/messages",
"authorization",
"Bearer key"
)]
#[tokio::test]
async fn prefixes_select_native_messages_and_strip_connection_fields(
call: MessagesCall,
#[case] model: &str,
#[case] path: &str,
#[case] header: &str,
#[case] credential: &str,
) {
let upstream = upstream([message_response()]).await;
let message = run_message(MessagesCall {
api_key: Some("key".into()),
api_base: Some(upstream.uri()),
custom_llm_provider: None,
..with_fields(
call,
json!({"model": model, "workspace_id": "request-workspace"}),
)
})
.await;
assert_eq!(message.model, MODEL);
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), path);
assert_eq!(request.header(header), Some(credential));
assert_eq!(request.json()["model"], model.rsplit('/').next().unwrap());
if path == "/v1/messages" {
assert_eq!(
request.header("anthropic-workspace-id"),
Some("request-workspace")
);
assert_eq!(request.json().get("workspace_id"), None);
}
}
#[rstest]
#[case::platform(
"bedrock/claude_platform/claude-test",
"ANTHROPIC_AWS_API_KEY",
"ANTHROPIC_AWS_BASE_URL",
"/v1/messages",
"x-api-key",
"managed"
)]
#[case::mantle(
"bedrock_mantle/claude-test",
"AWS_BEARER_TOKEN_BEDROCK",
"BEDROCK_MANTLE_API_BASE",
"/anthropic/v1/messages",
"authorization",
"Bearer managed"
)]
#[tokio::test]
async fn managed_keys_bases_and_workspaces_reach_the_upstream(
call: MessagesCall,
#[case] model: &str,
#[case] key_name: &str,
#[case] base_name: &str,
#[case] path: &str,
#[case] header: &str,
#[case] credential: &str,
) {
let upstream = upstream([message_response()]).await;
let base = upstream.uri();
let secrets = Arc::new(RecordingSecrets::new([
(key_name, "managed"),
(base_name, base.as_str()),
("ANTHROPIC_AWS_WORKSPACE_ID", "managed-workspace"),
]));
let output = run_with(
secrets,
MessagesCall {
custom_llm_provider: None,
..with_model(call, model)
},
)
.await
.unwrap();
assert!(matches!(output, MessagesCallResponse::Complete(_)));
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), path);
assert_eq!(request.header(header), Some(credential));
if path == "/v1/messages" {
assert_eq!(
request.header("anthropic-workspace-id"),
Some("managed-workspace")
);
}
}
#[rstest]
#[case::platform("bedrock/claude_platform/claude-test", "aws-external-anthropic")]
#[case::mantle("bedrock_mantle/claude-test", "bedrock")]
#[tokio::test]
async fn managed_aws_credentials_sign_the_exact_native_request(
call: MessagesCall,
#[case] model: &str,
#[case] service: &str,
) {
let upstream = upstream([message_response()]).await;
let secrets = Arc::new(RecordingSecrets::new([
("AWS_ACCESS_KEY_ID", "test-access"),
("AWS_SECRET_ACCESS_KEY", "test-secret"),
("AWS_SESSION_TOKEN", "test-session"),
("AWS_REGION_NAME", "us-west-2"),
("ANTHROPIC_AWS_WORKSPACE_ID", "managed-workspace"),
]));
run_with(
secrets,
MessagesCall {
api_base: Some(upstream.uri()),
custom_llm_provider: None,
..with_model(call, model)
},
)
.await
.unwrap();
let request = only_request(&upstream).await;
let auth = request.header("authorization").unwrap();
assert!(auth.starts_with("AWS4-HMAC-SHA256 Credential=test-access/"));
assert!(auth.contains(&format!("/us-west-2/{service}/aws4_request")));
assert_eq!(request.header("x-amz-security-token"), Some("test-session"));
assert_eq!(request.header("x-api-key"), None);
assert_eq!(request.json()["model"], "claude-test");
}
#[rstest]
#[tokio::test]
async fn private_mantle_endpoint_is_used_and_signed(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let params = litellm_router_types::LitellmParams {
aws: litellm_auth::AwsParams {
aws_bedrock_runtime_endpoint: Some(upstream.uri()),
aws_access_key_id: Some("private-access".into()),
aws_secret_access_key: Some("private-secret".into()),
aws_region_name: Some("us-west-2".into()),
..Default::default()
},
..Default::default()
};
run(MessagesCall {
custom_llm_provider: None,
litellm_params: params,
..with_model(call, "bedrock/mantle/claude-test")
})
.await
.unwrap();
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), "/anthropic/v1/messages");
assert!(
request
.header("authorization")
.unwrap()
.contains("/us-west-2/bedrock/aws4_request")
);
}
#[rstest]
#[case::platform("bedrock/claude_platform/claude-test")]
#[case::mantle("bedrock_mantle/claude-test")]
#[tokio::test]
async fn native_streams_relay_every_sse_byte(call: MessagesCall, #[case] model: &str) {
let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let upstream =
upstream([ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")]).await;
let response = messages_route(no_secrets())
.execute(
MessagesCall {
api_key: Some("key".into()),
api_base: Some(upstream.uri()),
custom_llm_provider: None,
..with_fields(
call,
json!({"model": model, "stream": true, "workspace_id": "workspace"}),
)
},
&(),
None,
)
.await
.unwrap();
let MessagesCallResponse::Stream { chunks, .. } = response else {
panic!("expected SSE");
};
assert_eq!(
chunks.try_collect::<Vec<_>>().await.unwrap().concat(),
sse.as_bytes()
);
assert_eq!(only_request(&upstream).await.json()["stream"], true);
}
#[rstest]
#[case::missing(false)]
#[case::configured(true)]
#[tokio::test]
async fn deployment_workspace_is_required_and_sent_as_a_header(
call: MessagesCall,
#[case] configured: bool,
) {
let upstream = upstream([message_response()]).await;
let params = if configured {
serde_json::from_value(json!({"model": "bedrock/claude_platform/claude-test", "workspace_id": "deployment-workspace"})).unwrap()
} else {
Default::default()
};
let result = run(MessagesCall {
api_key: Some("key".into()),
api_base: Some(upstream.uri()),
litellm_params: params,
..with_model(call, "bedrock/claude_platform/claude-test")
})
.await;
if !configured {
assert!(matches!(result, Err(Error::Auth(_))));
assert!(received(&upstream).await.is_empty());
return;
}
assert!(matches!(result.unwrap(), MessagesCallResponse::Complete(_)));
let sent = only_request(&upstream).await;
assert_eq!(
sent.header("anthropic-workspace-id"),
Some("deployment-workspace")
);
assert_eq!(sent.json().get("workspace_id"), None);
}

View file

@ -13,8 +13,8 @@ use litellm_cache_response::{
CacheOptions, CacheScope, ResponseCache, ResponseCacheConfig, ResponseCacheService, ScopedCache,
};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext, ResultSource,
WireRequest,
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RawResponse, RequestContext,
ResultSource, WireRequest,
};
use litellm_inference::{
RouteError,
@ -241,3 +241,55 @@ async fn messages_cache_identity_follows_resolved_configuration_and_request_call
first.verify().await;
second.verify().await;
}
#[rstest]
#[case::reported(json!(0.37), Some(0.37))]
#[case::free(json!(0), Some(0.0))]
#[case::malformed(json!("invalid"), None)]
#[case::missing(Value::Null, None)]
#[tokio::test]
async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_hit(
cache: Arc<dyn ResponseCacheService>,
#[case] cost: Value,
#[case] expected: Option<f64>,
) {
use litellm_inference_messages::{MessagesCall, MessagesCallResponse};
let response = json!({"id": "msg_cost", "type": "message", "role": "assistant", "model": "native-test",
"content": [{"type": "text", "text": "answer"}], "stop_reason": "end_turn", "stop_sequence": null,
"usage": {"input_tokens": 3, "output_tokens": 2}, "cost": cost});
let upstream = support::upstream([support::json_response(response)]).await;
let hooks = ChangingHooks::default();
let route = support::messages_route(litellm_inference_testing::no_secrets())
.with_cache(ScopedCache::new(cache, CacheScope::Shared));
let request = || {
MessagesCall {
body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})).unwrap(),
api_key: Some("eden-key".into()), api_base: Some(upstream.uri()), custom_llm_provider: Some("edenai".into()),
litellm_params: Default::default(), extra_headers: None, provider_specific_header: None,
timeout: None, shaping: Default::default(),
}
};
let first = route.execute(request(), &hooks, None).await.unwrap();
let second = route.execute(request(), &hooks, None).await.unwrap();
let (MessagesCallResponse::Complete(first), MessagesCallResponse::Complete(second)) =
(first, second)
else {
panic!("expected Messages responses");
};
assert_eq!(first, second);
assert_eq!(first.extra.get("cost"), Some(&cost));
assert_eq!(support::received(&upstream).await.len(), 1);
let facts = hooks.facts.lock().unwrap();
let [provider, cached] = facts.as_slice() else {
panic!("expected provider and cache results");
};
assert_eq!(provider.source, ResultSource::Provider);
assert_eq!(
provider.cost,
expected.map_or(Cost::Deferred, |amount| Cost::Reported {
amount: serde_json::Number::from_f64(amount).unwrap()
})
);
assert!(matches!(cached.source, ResultSource::Cache { .. }));
assert_eq!(cached.cost, Cost::CacheHit);
}

View file

@ -0,0 +1,8 @@
#[path = "live/support.rs"]
mod support;
#[path = "live/edenai.rs"]
mod edenai;
#[path = "live/vertex_ai.rs"]
mod vertex_ai;

View file

@ -0,0 +1,92 @@
use litellm_host::interceptors::Cost;
use litellm_inference_messages::MessagesRoute;
use litellm_inference_testing::live::within_deadline;
use rstest::rstest;
use serde_json::{Number, Value, json};
use super::support::{LiveCall, assert_text, call, complete, request, route, stream_events};
const PROVIDER: &str = "edenai";
#[rstest]
#[case::basic(json!({"messages": [{"role": "user", "content": "Say hello."}]}))]
#[case::system_text(json!({
"system": "You are a helpful assistant.",
"messages": [{"role": "user", "content": "Say hello."}]
}))]
#[case::portable_cache_hint(json!({
"system": [{"type": "text", "text": "You are a helpful assistant.", "cache_control": {"type": "ephemeral", "ttl": "1h"}}],
"messages": [{"role": "user", "content": "Say hello."}]
}))]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn completion_and_cost_handoff(route: MessagesRoute, #[case] payload: Value) {
within_deadline(async {
let host = LiveCall::new(PROVIDER);
let message = complete(request(&route, call(PROVIDER, payload), &host).await);
assert_text(&message);
host.assert_provider_result();
let expected_cost = message
.extra
.get("cost")
.and_then(|value| match value {
Value::Number(number) => number.as_f64(),
Value::String(text) => text.parse::<f64>().ok(),
_ => None,
})
.filter(|cost| *cost >= 0.0)
.and_then(Number::from_f64)
.map_or(Cost::Deferred, |amount| Cost::Reported { amount });
assert_eq!(host.facts().cost, expected_cost);
println!("{}", serde_json::to_string(&message).unwrap());
})
.await;
}
#[rstest]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn streaming(route: MessagesRoute) {
within_deadline(async {
let host = LiveCall::new(PROVIDER);
let events = stream_events(
request(
&route,
call(
PROVIDER,
json!({
"stream": true, "messages": [{"role": "user", "content": "Say hello."}]
}),
),
&host,
)
.await,
)
.await;
assert!(events.iter().any(|event| {
event["delta"]["text"]
.as_str()
.is_some_and(|text| !text.is_empty())
}));
host.assert_provider_result();
assert_eq!(host.facts().cost, Cost::Deferred);
})
.await;
}
#[rstest]
#[case::complete(false)]
#[case::streaming(true)]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn tool_round_trip(route: MessagesRoute, #[case] stream: bool) {
super::support::tool_round_trip(
route,
PROVIDER,
stream,
json!({
"max_tokens": 512, "tool_choice": {"type": "auto"}
}),
)
.await;
}

View file

@ -0,0 +1,265 @@
use std::{sync::OnceLock, time::Duration};
use futures_util::{FutureExt, TryStreamExt};
use litellm_host::interceptors::{
ExecutionFacts, Interceptors, RawResponse, RequestContext, ResultSource, WireRequest,
};
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesRoute};
use litellm_inference_testing::live::{LiveResources, model, within_deadline};
use litellm_llms_types::formats::messages::MessagesResponse;
use serde_json::{Value, json};
pub(super) struct LiveCall {
provider: &'static str,
facts: OnceLock<ExecutionFacts>,
}
impl LiveCall {
pub(super) fn new(provider: &'static str) -> Self {
Self {
provider,
facts: OnceLock::new(),
}
}
pub(super) fn facts(&self) -> &ExecutionFacts {
self.facts
.get()
.expect("the route must deliver execution facts")
}
pub(super) fn assert_provider_result(&self) {
let facts = self.facts();
assert_eq!(facts.provider.provider, self.provider);
assert_eq!(facts.source, ResultSource::Provider);
println!(
"{}",
json!({"route": "messages", "provider": self.provider, "source": "provider"})
);
}
}
impl Interceptors<Error> for LiveCall {
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
assert!(
self.facts.set(facts).is_ok(),
"result must be accepted exactly once"
);
Ok(())
}
async fn before_provider_request(
&self,
wire: WireRequest,
context: RequestContext,
) -> Result<WireRequest, Error> {
assert_eq!(context.custom_llm_provider, self.provider);
println!(
"{}",
json!({"route": "messages", "provider": self.provider, "model": context.model})
);
Ok(wire)
}
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Ok(())
}
}
#[rstest::fixture]
pub(super) fn route() -> MessagesRoute {
let resources = LiveResources::default();
MessagesRoute::new(resources.http, resources.auth, resources.secrets)
}
pub(super) async fn request(
route: &MessagesRoute,
call: MessagesCall,
host: &LiveCall,
) -> MessagesCallResponse {
route
.execute(call, host, None)
.boxed()
.await
.expect("live request failed")
}
pub(super) fn call(provider: &str, payload: Value) -> MessagesCall {
let Value::Object(fields) = payload else {
panic!("live payload must be an object")
};
let body = serde_json::from_value(Value::Object(
[("max_tokens".into(), json!(128))]
.into_iter()
.chain(fields)
.chain([("model".into(), json!(model("messages", provider)))])
.collect(),
))
.expect("live payload must be a valid Messages request");
MessagesCall {
body,
api_key: None,
api_base: None,
custom_llm_provider: Some(provider.into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(60)),
shaping: Default::default(),
}
}
pub(super) fn complete(response: MessagesCallResponse) -> Box<MessagesResponse> {
match response {
MessagesCallResponse::Complete(message) => message,
MessagesCallResponse::Stream { .. } => panic!("a complete request returned a stream"),
}
}
pub(super) fn assert_text(message: &MessagesResponse) {
assert!(!message.id.is_empty());
assert_eq!(message.role, "assistant");
assert!(message.stop_reason.is_some());
assert!(
message
.content
.iter()
.any(|block| block["text"].as_str().is_some_and(|text| !text.is_empty()))
);
assert!(
message
.usage
.as_ref()
.and_then(|usage| usage["output_tokens"].as_u64())
.is_some_and(|count| count > 0)
);
}
pub(super) async fn stream_events(response: MessagesCallResponse) -> Vec<Value> {
let MessagesCallResponse::Stream { head, chunks } = response else {
panic!("stream request returned a complete response")
};
let content_types: Vec<_> = head
.headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("content-type"))
.map(|(_, value)| value.split(';').next().unwrap().trim())
.collect();
assert_eq!(content_types, ["text/event-stream"]);
let frames = litellm_framer::frames(chunks, litellm_framer::sse::SseCodec::default())
.try_collect::<Vec<_>>()
.await
.expect("live SSE framing failed");
let data_frames = match frames.last() {
Some(frame) if frame.data == "[DONE]" => &frames[..frames.len() - 1],
_ => frames.as_slice(),
};
let events: Vec<Value> = data_frames
.iter()
.map(|frame| serde_json::from_str(&frame.data).expect("SSE data must be JSON"))
.collect();
assert_eq!(
events.first().map(|event| event["type"].as_str()),
Some(Some("message_start"))
);
assert_eq!(
events.last().map(|event| event["type"].as_str()),
Some(Some("message_stop"))
);
assert!(!events.iter().any(|event| event["type"] == "error"));
assert!(events.iter().any(|event| {
event["usage"]["output_tokens"]
.as_u64()
.is_some_and(|count| count > 0)
}));
println!("{}", json!({"events": events}));
events
}
fn prompt() -> Value {
json!({"role": "user", "content": "Use the echo tool with value live-check, then tell me the result."})
}
fn tools() -> Value {
json!([{"name": "echo", "description": "Return a value", "input_schema": {
"type": "object", "properties": {"value": {"type": "string", "enum": ["live-check"]}}, "required": ["value"]
}}])
}
fn fragments(events: &[Value], index: u64, field: &str) -> String {
events
.iter()
.filter(|event| event["index"].as_u64() == Some(index))
.filter_map(|event| event["delta"][field].as_str())
.collect()
}
fn streamed_block(start: &Value, events: &[Value]) -> Value {
let index = start["index"]
.as_u64()
.expect("content block must have an index");
let block = &start["content_block"];
match block["type"].as_str() {
Some("tool_use") => {
let partial = fragments(events, index, "partial_json");
let input = match partial.is_empty() {
true => block["input"].clone(),
false => {
serde_json::from_str(&partial).expect("streamed tool arguments must be JSON")
}
};
json!({"type": "tool_use", "id": block["id"], "name": block["name"], "input": input})
}
Some("text") => {
json!({"type": "text", "text": format!("{}{}", block["text"].as_str().unwrap_or_default(), fragments(events, index, "text"))})
}
Some("thinking") => {
json!({"type": "thinking", "thinking": format!("{}{}", block["thinking"].as_str().unwrap_or_default(), fragments(events, index, "thinking")),
"signature": format!("{}{}", block["signature"].as_str().unwrap_or_default(), fragments(events, index, "signature"))})
}
_ => block.clone(),
}
}
pub(super) async fn tool_round_trip(
route: MessagesRoute,
provider: &'static str,
stream: bool,
tool_params: Value,
) {
within_deadline(async {
let first_host = LiveCall::new(provider);
let payload = json!({"messages": [prompt()], "tools": tools(), "stream": stream});
let fields = payload.as_object().expect("tool payload must be an object").clone();
let params = tool_params.as_object().expect("tool parameters must be an object").clone();
let first = request(&route, call(provider, Value::Object(fields.into_iter().chain(params).collect())), &first_host).await;
let content = match stream {
false => {
let message = complete(first);
assert_eq!(message.stop_reason.as_deref(), Some("tool_use"));
message.content
}
true => {
let events = stream_events(first).await;
assert!(events.iter().any(|event| event["delta"]["stop_reason"] == "tool_use"));
events.iter().filter(|event| event["type"] == "content_block_start")
.map(|event| streamed_block(event, &events)).collect()
}
};
let tool = content.iter().find(|block| block["type"] == "tool_use").expect("requested tool must be returned");
assert_eq!(tool["name"], "echo");
assert_eq!(tool["input"]["value"], "live-check");
let tool_id = tool["id"].as_str().filter(|id| !id.is_empty()).expect("tool must have an id");
first_host.assert_provider_result();
let second_host = LiveCall::new(provider);
let second = complete(request(&route, call(provider, json!({
"tools": tools(), "messages": [prompt(),
{"role": "assistant", "content": content},
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": tool_id, "content": "live-check"}]}]
})), &second_host).await);
assert_text(&second);
second_host.assert_provider_result();
println!("{}", serde_json::to_string(&second).unwrap());
})
.await;
}

View file

@ -0,0 +1,68 @@
use litellm_inference_messages::MessagesRoute;
use litellm_inference_testing::live::within_deadline;
use rstest::rstest;
use serde_json::{Value, json};
use super::support::{LiveCall, assert_text, call, complete, request, route, stream_events};
const PROVIDER: &str = "vertex_ai";
#[rstest]
#[case::basic(json!({"messages": [{"role": "user", "content": "Say hello."}]}))]
#[case::system_text(json!({
"system": "You are a helpful assistant.",
"messages": [{"role": "user", "content": "Say hello."}]
}))]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn completion(route: MessagesRoute, #[case] payload: Value) {
within_deadline(async {
let host = LiveCall::new(PROVIDER);
let message = complete(request(&route, call(PROVIDER, payload), &host).await);
assert_text(&message);
host.assert_provider_result();
println!("{}", serde_json::to_string(&message).unwrap());
})
.await;
}
#[rstest]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn streaming(route: MessagesRoute) {
within_deadline(async {
let host = LiveCall::new(PROVIDER);
let response = request(
&route,
call(
PROVIDER,
json!({"stream": true, "messages": [{"role": "user", "content": "Say hello."}]}),
),
&host,
)
.await;
let events = stream_events(response).await;
assert!(events.iter().any(|event| {
event["delta"]["text"]
.as_str()
.is_some_and(|text| !text.is_empty())
}));
host.assert_provider_result();
})
.await;
}
#[rstest]
#[case::complete(false)]
#[case::streaming(true)]
#[ignore = "calls a real provider, requires credentials and a live model"]
#[tokio::test]
async fn tool_round_trip(route: MessagesRoute, #[case] stream: bool) {
super::support::tool_round_trip(
route,
PROVIDER,
stream,
json!({"tool_choice": {"type": "auto"}}),
)
.await;
}

View file

@ -24,6 +24,7 @@ mod request;
mod response;
mod secrets;
mod stream;
mod tencent;
const MODEL: &str = "claude-sonnet-4-5";

View file

@ -11,6 +11,9 @@ use super::*;
#[rstest]
#[case::anthropic_key("anthropic", Some("sk-ant"), &[], ("x-api-key", "sk-ant"), &["authorization"])]
#[case::minimax_key("minimax", Some("sk-test"), &[], ("x-api-key", "sk-test"), &["authorization"])]
#[case::minimax_forwarded_key("minimax", Some("deployment"), &[("X-Api-Key", "caller")], ("x-api-key", "caller"), &["authorization"])]
#[case::minimax_forwarded_bearer("minimax", None, &[("Authorization", "Bearer caller")], ("authorization", "Bearer caller"), &["x-api-key"])]
#[case::azure_key("azure_ai", Some("sk-azure"), &[], ("x-api-key", "sk-azure"), &["authorization"])]
#[case::deepseek_key("deepseek", Some("sk-deepseek"), &[], ("x-api-key", "sk-deepseek"), &["authorization"])]
#[case::deepseek_forwards_caller_authorization(
@ -88,6 +91,7 @@ async fn credentials_become_exactly_one_auth_header(
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[case::deepseek("deepseek")]
#[case::minimax("minimax")]
#[tokio::test]
async fn a_call_without_credentials_fails_before_sending(
call: MessagesCall,
@ -123,6 +127,9 @@ async fn a_call_without_credentials_fails_before_sending(
"/v1/messages",
"/v1/messages"
)]
#[case::minimax(MODEL, Some("minimax"), "", "/v1/messages")]
#[case::minimax_prefix("minimax/claude-sonnet-4-5", None, "/v1/messages", "/v1/messages")]
#[case::minimax_endpoint_slash(MODEL, Some("minimax"), "/v1/messages/", "/v1/messages")]
#[case::azure_ai(MODEL, Some("azure_ai"), "", "/anthropic/v1/messages")]
#[case::deepseek(MODEL, Some("deepseek"), "", "/anthropic/v1/messages")]
#[case::deepseek_openai_compatible_base(MODEL, Some("deepseek"), "/beta", "/anthropic/v1/messages")]
@ -301,6 +308,7 @@ fn sent_betas(request: &wiremock::Request) -> BetaSet {
#[tokio::test]
async fn feature_betas_join_the_callers_betas_in_one_sorted_header(
call: MessagesCall,
#[values("anthropic", "minimax")] provider: &str,
#[case] fields: Value,
#[case] features: &[AnthropicBeta],
) {
@ -312,6 +320,7 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header(
run_message(with_fields(
MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]),
@ -362,6 +371,7 @@ async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: M
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[case::deepseek("deepseek")]
#[case::minimax("minimax")]
#[tokio::test]
async fn caller_protocol_headers_win_over_the_defaults(call: MessagesCall, #[case] provider: &str) {
let upstream = upstream([message_response()]).await;
@ -452,6 +462,7 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i
#[tokio::test]
async fn reasoning_auto_summary_marks_active_thinking_on_the_wire(
call: MessagesCall,
#[values("anthropic", "minimax")] provider: &str,
#[case] thinking: Value,
#[case] expected: Value,
) {
@ -459,6 +470,7 @@ async fn reasoning_auto_summary_marks_active_thinking_on_the_wire(
run_message(with_fields(
MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
@ -517,6 +529,7 @@ async fn reasoning_auto_summary_marks_active_thinking_on_the_wire(
#[tokio::test]
async fn reasoning_is_translated_by_the_model_capabilities(
call: MessagesCall,
#[values("anthropic", "minimax")] provider: &str,
#[case] capabilities: MessagesModelCapabilities,
#[case] fields: Value,
#[case] expected: Value,
@ -525,6 +538,7 @@ async fn reasoning_is_translated_by_the_model_capabilities(
run_message(with_fields(
MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
@ -592,6 +606,7 @@ async fn replayed_history_is_cleaned_before_sending(
#[case::anthropic("anthropic")]
#[case::azure("azure_ai")]
#[case::deepseek("deepseek")]
#[case::minimax("minimax")]
#[tokio::test]
async fn metadata_is_reduced_to_the_user_id(call: MessagesCall, #[case] provider: &str) {
let upstream = upstream([message_response()]).await;
@ -703,11 +718,16 @@ async fn vertex_ai_addresses_the_model_in_the_url_and_not_in_the_body(
#[case::numeric_user_id(json!({"metadata": {"user_id": 7}}))]
#[case::missing_max_tokens(json!({"max_tokens": null}))]
#[tokio::test]
async fn an_invalid_request_fails_before_sending(call: MessagesCall, #[case] fields: Value) {
async fn an_invalid_request_fails_before_sending(
call: MessagesCall,
#[values("anthropic", "minimax")] provider: &str,
#[case] fields: Value,
) {
let upstream = upstream([message_response()]).await;
let error = run(with_fields(
MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
@ -721,6 +741,45 @@ async fn an_invalid_request_fails_before_sending(call: MessagesCall, #[case] fie
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::billing_string(json!("x-anthropic-billing-header: private"), None)]
#[case::billing_blocks(json!([{ "type": "text", "text": "x-anthropic-billing-header: private" }]), None)]
#[case::mixed_blocks(
json!([{ "type": "text", "text": "x-anthropic-billing-header: private" }, { "type": "text", "text": "keep" }]),
Some(json!([{ "type": "text", "text": "keep" }]))
)]
#[case::ordinary_system(json!("keep"), Some(json!("keep")))]
#[tokio::test]
async fn minimax_strips_only_billing_system_blocks(
call: MessagesCall,
#[case] system: Value,
#[case] expected: Option<Value>,
) {
let upstream = upstream([message_response()]).await;
let fields = json!({
"system": system,
"tools": [{"name": "echo", "input_schema": {"type": "object"}}],
"messages": [{"role": "user", "content": [{
"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}
}]}]
});
run_message(with_fields(
MessagesCall {
custom_llm_provider: Some("minimax".into()),
api_key: Some("sk".into()),
api_base: Some(upstream.uri()),
..call
},
fields.clone(),
))
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent.get("system"), expected.as_ref());
assert_eq!(sent["tools"], fields["tools"]);
assert_eq!(sent["messages"], fields["messages"]);
}
#[rstest]
#[case::azure("azure_ai", &[], json!([
{"type": "text", "text": "top level"},
@ -835,3 +894,38 @@ async fn provider_validation_runs_before_caller_parameter_removal(
);
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn edenai_preserves_native_thinking_and_strips_billing_and_cache_extensions(
call: MessagesCall,
) {
let upstream = upstream([message_response()]).await;
let thinking = json!({"type": "enabled", "budget_tokens": 512, "native_field": true});
run_message(MessagesCall {
custom_llm_provider: Some("edenai".into()),
api_base: Some(format!("{}/v3", upstream.uri())),
api_key: Some("eden-key".into()),
extra_headers: headers([("Anthropic-Beta", "unregistered-native-beta")]),
..with_fields(call, json!({
"thinking": thinking, "temperature": 0.25,
"system": [{"type": "text", "text": "x-anthropic-billing-header: hidden"},
{"type": "text", "text": "keep", "cache_control": {"type": "ephemeral", "ttl": "1h", "scope": "global"}}],
}))
}).await;
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/v3/v1/messages");
assert_eq!(sent.header("authorization"), Some("Bearer eden-key"));
assert_eq!(
sent.header("anthropic-beta"),
Some("unregistered-native-beta")
);
assert_eq!(
sent.json(),
json!({
"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}],
"thinking": thinking, "temperature": 0.25,
"system": [{"type": "text", "text": "keep", "cache_control": {"type": "ephemeral"}}],
})
);
}

View file

@ -25,6 +25,13 @@ use super::*;
"/anthropic/v1/messages",
&["DEEPSEEK_API_KEY", "DEEPSEEK_ANTHROPIC_API_BASE", "DEEPSEEK_API_BASE"]
)]
#[case::minimax(
"minimax",
"MINIMAX_API_KEY",
"MINIMAX_API_BASE",
"/v1/messages",
&["MINIMAX_API_KEY", "MINIMAX_API_BASE"]
)]
#[tokio::test]
async fn the_credential_and_base_come_from_the_secret_source(
call: MessagesCall,

View file

@ -150,9 +150,21 @@ async fn stream_through(host: &RecordingStreamHost) -> Result<MessagesOutput, Er
#[rstest]
#[tokio::test]
async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(call: MessagesCall) {
async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(
call: MessagesCall,
#[values("anthropic", "minimax")] provider: &str,
) {
let upstream = upstream([sse_response()]).await;
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
let host = RecordingStreamHost::new(
streaming(
MessagesCall {
custom_llm_provider: Some(provider.into()),
..call
},
upstream.uri(),
),
usize::MAX,
);
let outcome = stream_through(&host).await.expect("streamed call succeeds");
@ -336,6 +348,7 @@ async fn the_timeout_covers_a_stalled_stream_body(call: MessagesCall) {
#[rstest]
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[case::edenai("edenai")]
#[tokio::test]
async fn the_sdk_returns_stream_headers_and_every_sse_byte(
call: MessagesCall,
@ -346,6 +359,7 @@ async fn the_sdk_returns_stream_headers_and_every_sse_byte(
.execute(
MessagesCall {
custom_llm_provider: Some(provider.into()),
api_key: Some("stream-key".into()),
litellm_params: Default::default(),
..streaming(call, upstream.uri())
},

View file

@ -0,0 +1,260 @@
use std::ops::ControlFlow;
use bytes::Bytes;
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities;
use litellm_llms_types::providers::anthropic::DEFAULT_HEADERS;
use rstest::{fixture, rstest};
use super::*;
#[fixture]
fn tencent_call(call: MessagesCall) -> MessagesCall {
MessagesCall {
custom_llm_provider: Some("tencent".into()),
api_key: Some("deployment-key".into()),
..call
}
}
#[rstest]
#[case::deployment_key(&[], "x-api-key", "deployment-key")]
#[case::forwarded_key(&[("X-Api-Key", "caller-key")], "x-api-key", "caller-key")]
#[case::forwarded_bearer(&[("Authorization", "Bearer caller")], "authorization", "Bearer caller")]
#[tokio::test]
async fn credentials_reach_the_selected_provider(
tencent_call: MessagesCall,
#[case] forwarded: &[(&str, &str)],
#[case] header: &str,
#[case] credential: &str,
) {
let upstream = upstream([message_response()]).await;
run_message(MessagesCall {
api_base: Some(upstream.uri()),
extra_headers: headers(forwarded.iter().copied()),
..tencent_call
})
.await;
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values(header), [credential]);
let version = DEFAULT_HEADERS
.iter()
.find_map(|(name, value)| (*name == "anthropic-version").then_some(*value))
.unwrap();
assert_eq!(sent.header_values("anthropic-version"), [version]);
assert_eq!(sent.json()["model"], MODEL);
}
#[rstest]
#[tokio::test]
async fn credentials_and_base_resolve_from_the_secret_source(tencent_call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let base = format!("{}/v1/chat/completions/", upstream.uri());
let secrets = Arc::new(RecordingSecrets::new([
("TENCENT_API_KEY", "secret-key"),
("TENCENT_ANTHROPIC_API_BASE", base.as_str()),
("TENCENT_API_BASE", UNREACHABLE_BASE),
]));
run_with(
secrets,
MessagesCall {
api_key: None,
..tencent_call
},
)
.await
.unwrap();
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/v1/messages");
assert_eq!(sent.header("x-api-key"), Some("secret-key"));
}
#[rstest]
#[tokio::test]
async fn missing_credentials_fail_before_sending(tencent_call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let error = run(MessagesCall {
api_key: None,
api_base: Some(upstream.uri()),
..tencent_call
})
.await
.unwrap_err();
assert!(matches!(
error,
Error::Auth(litellm_auth::Error::MissingApiKey {
environment_variable: "TENCENT_API_KEY",
..
})
));
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::missing_max_tokens(json!({"max_tokens": null}))]
#[case::invalid_metadata(json!({"metadata": {"user_id": 7}}))]
#[case::invalid_reasoning_effort(json!({"reasoning_effort": "unknown"}))]
#[tokio::test]
async fn malformed_requests_fail_before_sending(tencent_call: MessagesCall, #[case] fields: Value) {
let upstream = upstream([message_response()]).await;
let error = run(with_fields(
MessagesCall {
api_base: Some(upstream.uri()),
..tencent_call
},
fields,
))
.await
.unwrap_err();
assert!(error.is_request(), "{error:?}");
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::legacy(false, json!({"type": "enabled", "budget_tokens": 4095}))]
#[case::adaptive(true, json!({"type": "adaptive", "display": "summarized"}))]
#[tokio::test]
async fn reasoning_effort_becomes_model_appropriate_thinking(
tencent_call: MessagesCall,
#[case] adaptive: bool,
#[case] thinking: Value,
) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
api_base: Some(upstream.uri()),
shaping: MessagesShaping {
capabilities: MessagesModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: adaptive,
..Default::default()
},
..Default::default()
},
..tencent_call
},
json!({"max_tokens": 4096, "reasoning_effort": "high"}),
))
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent["thinking"], thinking);
assert_eq!(sent.get("reasoning_effort"), None);
assert_eq!(
sent.get("output_config"),
adaptive.then(|| json!({"effort": "high"})).as_ref()
);
}
#[rstest]
#[case::encrypted_reasoning(json!({"type": "thinking", "thinking": "private", "signature": "litellm_encrypted_reasoning:private"}))]
#[case::orphaned_advisor(json!({"type": "server_tool_use", "name": "advisor", "id": "srvtoolu_advisor", "input": {}}))]
#[tokio::test]
async fn replayed_history_is_cleaned_without_losing_text(
tencent_call: MessagesCall,
#[case] block: Value,
) {
let upstream = upstream([message_response()]).await;
run_message(with_fields(
MessagesCall {
api_base: Some(upstream.uri()),
..tencent_call
},
json!({"messages": [
{"role": "assistant", "content": [block, {"type": "text", "text": "keep"}]},
{"role": "user", "content": "continue"}
]}),
))
.await;
assert_eq!(
only_request(&upstream).await.json()["messages"],
json!([
{"role": "assistant", "content": [{"type": "text", "text": "keep"}]},
{"role": "user", "content": "continue"}
])
);
}
#[rstest]
#[case::billing_string(json!("x-anthropic-billing-header: private"), None)]
#[case::mixed_blocks(json!([
{"type": "text", "text": "x-anthropic-billing-header: private"},
{"type": "text", "text": "keep"}
]), Some(json!([{ "type": "text", "text": "keep" }])))]
#[case::ordinary_system(json!("keep"), Some(json!("keep")))]
#[tokio::test]
async fn billing_is_removed_and_tools_and_cache_hints_are_preserved(
tencent_call: MessagesCall,
#[case] system: Value,
#[case] expected: Option<Value>,
) {
let upstream = upstream([message_response()]).await;
let fields = json!({
"system": system,
"tools": [{"name": "echo", "input_schema": {"type": "object"}}],
"messages": [{"role": "user", "content": [{
"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}
}]}]
});
run_message(with_fields(
MessagesCall {
api_base: Some(upstream.uri()),
..tencent_call
},
fields.clone(),
))
.await;
let sent = only_request(&upstream).await.json();
assert_eq!(sent.get("system"), expected.as_ref());
assert_eq!(sent["tools"], fields["tools"]);
assert_eq!(sent["messages"], fields["messages"]);
}
#[derive(Default)]
struct StreamSink(Mutex<Vec<u8>>);
impl litellm_host_native::in_process::StreamConsumer<Messages> for StreamSink {
async fn open_stream(
&self,
_: litellm_inference_messages::route::MessagesStreamHead,
) -> Result<ControlFlow<()>, Error> {
Ok(ControlFlow::Continue(()))
}
async fn send_chunk(&self, chunk: Bytes) -> Result<ControlFlow<()>, Error> {
self.0.lock().unwrap().extend_from_slice(&chunk);
Ok(ControlFlow::Continue(()))
}
}
#[rstest]
#[tokio::test]
async fn streaming_relays_provider_events_unchanged(tencent_call: MessagesCall) {
let events = concat!(
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n",
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"OK\"}}\n\n",
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"
);
let upstream =
upstream([ResponseTemplate::new(200).set_body_raw(events, "text/event-stream")]).await;
let host = LocalMessagesHost::new(with_fields(
MessagesCall {
api_base: Some(upstream.uri()),
..tencent_call
},
json!({"stream": true}),
));
let sink = StreamSink::default();
let outcome = litellm_host_native::in_process::run_hosted(
machine(Arc::new(RecordingSecrets::empty()))(host.request().unwrap()),
litellm_host_native::in_process::Host {
services: &(),
interceptors: &host,
stream: &sink,
observers: None,
},
)
.await
.unwrap();
assert!(matches!(outcome, MessagesOutput::StreamEnded));
assert_eq!(sink.0.into_inner().unwrap(), events.as_bytes());
assert_eq!(only_request(&upstream).await.json()["stream"], true);
}

View file

@ -0,0 +1,147 @@
use std::{sync::Arc, time::Duration};
use futures_util::TryStreamExt;
use litellm_inference_messages::{Error, MessagesCall, MessagesCallResponse, MessagesShaping};
use litellm_inference_testing::{RecordingSecrets, no_secrets};
use rstest::{fixture, rstest};
use serde_json::json;
use wiremock::ResponseTemplate;
#[path = "../support/mod.rs"]
mod support;
use support::*;
#[fixture]
fn call() -> MessagesCall {
MessagesCall {
body: serde_json::from_value(json!({"model": "openrouter/anthropic/test-model", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})).unwrap(),
api_key: None,
api_base: None,
custom_llm_provider: None,
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: Some(Duration::from_secs(5)),
shaping: MessagesShaping::default(),
}
}
fn message_response() -> ResponseTemplate {
json_response(
json!({"id": "msg_1", "type": "message", "role": "assistant", "model": "anthropic/test-model", "content": [{"type": "text", "text": "hi"}], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1}}),
)
}
#[rstest]
#[case::explicit(Some("explicit"), None, "Bearer explicit")]
#[case::managed(None, None, "Bearer managed")]
#[case::forwarded(None, Some("bEaReR caller"), "bEaReR caller")]
#[case::blank(Some("explicit"), Some("Bearer "), "Bearer explicit")]
#[case::malformed(Some("explicit"), Some("invali\u{00e9}auth"), "Bearer explicit")]
#[tokio::test]
async fn only_a_valid_bearer_can_override_the_configured_key(
call: MessagesCall,
#[case] api_key: Option<&str>,
#[case] authorization: Option<&str>,
#[case] expected: &str,
) {
let upstream = upstream([message_response()]).await;
let extra_headers = [("x-api-key".to_string(), json!("caller-anthropic-key"))]
.into_iter()
.chain(authorization.map(|value| ("Authorization".to_string(), json!(value))))
.collect();
let secrets = Arc::new(RecordingSecrets::new([("OPENROUTER_API_KEY", "managed")]));
let output = messages_route(secrets)
.execute(
MessagesCall {
api_key: api_key.map(str::to_string),
api_base: Some(format!("{}/api/v1", upstream.uri())),
extra_headers: Some(extra_headers),
..call
},
&(),
None,
)
.await
.unwrap();
assert!(matches!(output, MessagesCallResponse::Complete(_)));
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), "/api/v1/messages");
assert_eq!(request.header_values("authorization"), [expected]);
assert_eq!(request.header("x-api-key"), None);
assert_eq!(request.json()["model"], "anthropic/test-model");
}
#[rstest]
#[tokio::test]
async fn an_anthropic_key_alone_is_not_an_openrouter_credential(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let result = messages_route(no_secrets())
.execute(
MessagesCall {
api_base: Some(upstream.uri()),
extra_headers: Some(
json!({"x-api-key": "caller-key"})
.as_object()
.unwrap()
.clone(),
),
..call
},
&(),
None,
)
.await;
assert!(matches!(
result,
Err(Error::Auth(litellm_auth::Error::MissingApiKey {
provider: "OpenRouter",
..
}))
));
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn the_secret_manager_supplies_the_native_base_and_key(call: MessagesCall) {
let upstream = upstream([message_response()]).await;
let base = format!("{}/api/v1", upstream.uri());
let secrets = Arc::new(RecordingSecrets::new([
("OPENROUTER_API_KEY", "managed"),
("OPENROUTER_API_BASE", base.as_str()),
]));
messages_route(secrets)
.execute(call, &(), None)
.await
.unwrap();
let request = only_request(&upstream).await;
assert_eq!(request.url.path(), "/api/v1/messages");
assert_eq!(request.header("authorization"), Some("Bearer managed"));
}
#[rstest]
#[tokio::test]
async fn streaming_relay_preserves_every_event_and_billed_usage(call: MessagesCall) {
let sse = "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"cost\":\"0.001\",\"cost_details\":\"unexpected\"}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let upstream =
upstream([ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")]).await;
let streamed = MessagesCall {
body: serde_json::from_value(json!({"model": "openrouter/anthropic/test-model", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}], "stream": true})).unwrap(),
api_key: Some("key".into()),
api_base: Some(upstream.uri()),
..call
};
let output = messages_route(no_secrets())
.execute(streamed, &(), None)
.await
.unwrap();
let MessagesCallResponse::Stream { chunks, .. } = output else {
panic!("expected SSE");
};
assert_eq!(
chunks.try_collect::<Vec<_>>().await.unwrap().concat(),
sse.as_bytes()
);
assert_eq!(only_request(&upstream).await.json()["stream"], true);
}

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