Resolve merge conflicts with litellm_internal_staging

This commit is contained in:
Dushyant Acharya 2026-06-13 23:09:28 +05:30
commit 875532bb8b
No known key found for this signature in database
2692 changed files with 288440 additions and 55270 deletions

View file

@ -111,6 +111,28 @@ commands:
- wait_for_service:
url: tcp://localhost:6379
timeout: "60"
start_openai_record_replay_proxy:
description: "Start the record/replay proxy (tests/_openai_record_replay_proxy.py) on host port 8090 and wait until healthy. Models whose api_base points here replay recorded provider responses, so the E2E run neither pays for nor depends on the live provider. The default upstream is OpenAI; a non-OpenAI model must point its api_base at /__recorder_upstream/<host>/ so the recorder forwards there instead of defaulting to OpenAI. Run after uv deps are synced."
steps:
- run:
name: Start record/replay proxy
background: true
command: |
CASSETTE_REDIS_URL="$CASSETTE_REDIS_URL" \
RECORDER_UPSTREAM_BASE_URL="https://api.openai.com" \
uv run --no-sync python tests/_openai_record_replay_proxy.py --host 0.0.0.0 --port 8090
- run:
name: Wait for record/replay proxy
command: |
for i in $(seq 1 30); do
if curl -sf http://localhost:8090/__recorder_health >/dev/null 2>&1; then
echo "record/replay proxy is up"
exit 0
fi
sleep 1
done
echo "record/replay proxy did not become ready" >&2
exit 1
setup_litellm_enterprise_pip:
steps:
- run:
@ -182,7 +204,14 @@ jobs:
- run:
name: Run Windows-specific test
command: |
uv run --no-sync python -m pytest tests/windows_tests/test_litellm_on_windows.py -v
uv run --no-sync python -m pytest tests/windows_tests/ -v
- run:
name: Guard against MAX_PATH-busting packaged wheel paths
environment:
UV_HTTP_TIMEOUT: "300"
command: |
uv build --wheel --out-dir dist
uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py
local_testing_part1:
docker:
@ -228,7 +257,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm \
--cov-report=xml \
@ -242,8 +271,15 @@ jobs:
- run:
name: Rename the coverage files
command: |
mv coverage.xml local_testing_part1_coverage.xml
mv .coverage local_testing_part1_coverage
# When CI reruns only the failed tests, a parallel node can receive
# zero tests and pytest never writes coverage. Emit empty placeholders
# so persist_to_workspace and the downstream coverage combine stay green.
if [ -f coverage.xml ]; then
mv coverage.xml local_testing_part1_coverage.xml
mv .coverage local_testing_part1_coverage
else
touch local_testing_part1_coverage.xml local_testing_part1_coverage
fi
# Store test results
- store_test_results:
@ -293,7 +329,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm \
--cov-report=xml \
@ -307,8 +343,15 @@ jobs:
- run:
name: Rename the coverage files
command: |
mv coverage.xml local_testing_part2_coverage.xml
mv .coverage local_testing_part2_coverage
# When CI reruns only the failed tests, a parallel node can receive
# zero tests and pytest never writes coverage. Emit empty placeholders
# so persist_to_workspace and the downstream coverage combine stay green.
if [ -f coverage.xml ]; then
mv coverage.xml local_testing_part2_coverage.xml
mv .coverage local_testing_part2_coverage
else
touch local_testing_part2_coverage.xml local_testing_part2_coverage
fi
# Store test results
- store_test_results:
@ -356,7 +399,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -409,7 +452,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/proxy_admin_ui_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -431,6 +474,120 @@ jobs:
- auth_ui_unit_tests_coverage.xml
- auth_ui_unit_tests_coverage
proxy_behavior_tests:
docker:
- *python312_image
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- checkout
- setup_google_dns
- install_uv
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Seed DB schema via prisma db push
command: |
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
- run:
name: Generate Prisma Client
command: uv run --no-sync python -m prisma generate
- run:
name: Run proxy management behavior tests
command: |
mkdir -p test-results
uv run --no-sync python -m pytest tests/proxy_behavior \
-v --junitxml=test-results/junit.xml --durations=10
no_output_timeout: 15m
- store_test_results:
path: test-results
proxy_security_tests:
docker:
- *python312_image
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- checkout
- setup_google_dns
- install_uv
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Seed DB schema via prisma db push
command: |
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
- run:
name: Generate Prisma Client
command: uv run --no-sync python -m prisma generate
- run:
name: Run proxy security tests
command: |
mkdir -p test-results
uv run --no-sync python -m pytest tests/proxy_security_tests \
-v --junitxml=test-results/junit.xml --durations=10
no_output_timeout: 15m
- store_test_results:
path: test-results
schema_migration_check:
docker:
- *python312_image
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
working_directory: ~/project
environment:
# An empty database; the test applies every committed migration itself.
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- checkout
- setup_google_dns
- install_uv
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Generate Prisma Client
command: uv run --no-sync python -m prisma generate
- run:
name: Check schema.prisma is in sync with committed migrations
command: |
mkdir -p test-results
uv run --no-sync python -m pytest tests/proxy_migration_tests \
-v --junitxml=test-results/junit.xml --durations=10
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_router_testing: # Runs all tests with the "router" keyword
docker:
- *python312_image
@ -457,12 +614,17 @@ jobs:
- run:
name: Run tests
command: |
# On a "rerun failed tests" build a parallel node can receive no
# tests, so the test command never creates test-results. Pre-create it
# so store_test_results doesn't fail the node on a missing path.
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v \
-k 'router' \
-n 4 \
@ -504,7 +666,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/router_unit_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -547,7 +709,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -589,7 +751,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/llm_translation/**/test_*.py" | grep -v "^tests/llm_translation/realtime/")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v \
--junitxml=test-results/junit.xml \
--durations=20 \
@ -625,7 +787,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/llm_translation/realtime/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -668,7 +830,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/agent_tests/**/test_*.py" | grep -v "^tests/agent_tests/local_only_agent_tests/")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -710,7 +872,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/guardrails_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -754,7 +916,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/unified_google_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -805,7 +967,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/llm_responses_api_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -836,7 +998,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -878,7 +1040,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/search_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -922,7 +1084,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/enterprise/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit-enterprise.xml \
--durations=10 \
@ -952,7 +1114,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/batches_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -994,7 +1156,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/litellm_utils_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1037,7 +1199,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/pass_through_unit_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1080,7 +1242,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/image_gen_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1112,7 +1274,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/logging_callback_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm --cov-report=xml \
-n 4 \
@ -1155,7 +1317,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/audio_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1206,7 +1368,7 @@ jobs:
tests/local_testing/test_router_utils.py)
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1456,7 +1618,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit-2.xml \
--durations=5"
@ -1485,6 +1647,7 @@ jobs:
command: |
zstd -d litellm-docker-database.tar.zst --stdout | docker load
docker tag litellm-docker-database:ci my-app:latest
- start_openai_record_replay_proxy
- run:
name: Run Docker container
command: |
@ -1515,6 +1678,7 @@ jobs:
-e LANGFUSE_PROJECT2_PUBLIC=$LANGFUSE_PROJECT2_PUBLIC \
-e LANGFUSE_PROJECT1_SECRET=$LANGFUSE_PROJECT1_SECRET \
-e LANGFUSE_PROJECT2_SECRET=$LANGFUSE_PROJECT2_SECRET \
-e RECORDER_OPENAI_BASE_URL=http://host.docker.internal:8090/v1 \
--add-host host.docker.internal:host-gateway \
--name my-app \
-v $(pwd)/proxy_server_config.yaml:/app/config.yaml \
@ -1539,7 +1703,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-s -v -x \
--junitxml=test-results/junit.xml \
-n 4 \
@ -1622,7 +1786,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/openai_endpoints_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-s -vv \
--junitxml=test-results/junit.xml \
--durations=5"
@ -1652,6 +1816,7 @@ jobs:
command: |
zstd -d litellm-docker-database.tar.zst --stdout | docker load
docker images | grep litellm-docker-database
- start_openai_record_replay_proxy
- run:
name: Run Docker container
# intentionally give bad redis credentials here
@ -1675,6 +1840,7 @@ jobs:
-e DD_SITE=$DD_SITE \
-e AWS_REGION_NAME=$AWS_REGION_NAME \
-e COHERE_API_KEY=$COHERE_API_KEY \
-e RECORDER_COHERE_BASE_URL=http://host.docker.internal:8090/__recorder_upstream/api.cohere.com \
-e GCS_FLUSH_INTERVAL="1" \
--add-host host.docker.internal:host-gateway \
--name my-app \
@ -1698,7 +1864,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/otel_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v \
--junitxml=test-results/junit.xml \
--durations=5"
@ -1748,7 +1914,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit-2.xml \
--durations=5"
@ -1824,7 +1990,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/spend_tracking_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--junitxml=test-results/junit.xml \
--durations=5"
@ -1922,7 +2088,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/multi_instance_e2e_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--junitxml=test-results/junit.xml \
--durations=5"
@ -1985,7 +2151,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/store_model_in_db_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2065,7 +2231,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--junitxml=test-results/junit-2.xml \
--durations=5"
@ -2209,7 +2375,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/pass_through_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2240,6 +2406,7 @@ jobs:
command: |
zstd -d litellm-docker-database.tar.zst --stdout | docker load
docker images | grep litellm-docker-database
- start_openai_record_replay_proxy
- run:
name: Run Docker container with test config
command: |
@ -2248,6 +2415,7 @@ jobs:
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
-e LITELLM_MASTER_KEY="sk-1234" \
-e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
-e RECORDER_ANTHROPIC_BASE_URL=http://host.docker.internal:8090/__recorder_upstream/api.anthropic.com \
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
-e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
-e AWS_REGION_NAME="us-east-1" \
@ -2275,7 +2443,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/proxy_e2e_anthropic_messages_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2400,6 +2568,11 @@ jobs:
environment:
DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e"
CI: "true"
# Boot the proxy with an external logout URL so proxyLogoutUrl.spec.ts can
# assert the redirect. Set at job level so both the proxy boot step and the
# Playwright step (whose skip guard reads this) see the same value. Safe for
# the rest of the suite: nothing else performs a logout.
PROXY_LOGOUT_URL: "https://www.example.com"
steps:
- checkout
- setup_google_dns
@ -2476,7 +2649,8 @@ jobs:
MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
DISABLE_SCHEMA_UPDATE: "true"
SERVER_ROOT_PATH: ""
PROXY_LOGOUT_URL: ""
# PROXY_LOGOUT_URL is inherited from the job-level environment so the
# proxy and proxyLogoutUrl.spec.ts agree on the logout target.
# LITELLM_LICENSE is forwarded from the project env so premium-gated
# UI flows can be exercised. license.spec.ts asserts the resulting
# JWT carries premium_user=true; if it ever stops being passed, that
@ -2516,6 +2690,122 @@ jobs:
path: ui/litellm-dashboard/playwright-report
destination: e2e-playwright-report
e2e_ui_testing_server_root_path:
docker:
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
environment:
POSTGRES_USER: e2euser
POSTGRES_PASSWORD: e2epassword
POSTGRES_DB: litellm_e2e
resource_class: large
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e"
CI: "true"
# The whole job exercises the proxy mounted under a prefix. SERVER_ROOT_PATH
# is read both by the proxy at boot (to rewrite the built UI bundle in place)
# and by migration.serverRootPath.config.ts, which refuses to run without it.
SERVER_ROOT_PATH: "/litellm"
steps:
- checkout
- setup_google_dns
- install_uv
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- run:
name: Install Python dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
uv run --no-sync python -m prisma generate --schema litellm/proxy/schema.prisma
- save_cache:
key: v1-uv-cache-{{ checksum "uv.lock" }}
paths:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
command: |
cd ui/litellm-dashboard
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- ~/.cache/ms-playwright
- run:
name: Build UI from source
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
- wait_for_service:
url: tcp://localhost:5432
timeout: "30"
- run:
name: Push Prisma schema
command: uv run --no-sync python -m prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
- run:
name: Seed database
command: |
PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \
-f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
- run:
name: Start mock LLM server
command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py
background: true
- run:
name: Start LiteLLM proxy under a server root path
environment:
LITELLM_MASTER_KEY: "sk-1234"
MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
DISABLE_SCHEMA_UPDATE: "true"
# Output flows to this step's own log, so a boot crash is visible here
# rather than swallowed by a downstream readiness probe.
command: |
LITELLM_LICENSE="$LITELLM_LICENSE" \
uv run --no-sync python -m litellm.proxy.proxy_cli \
--config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \
--port 4000
background: true
- run:
name: Wait for prefixed proxy to be ready
command: |
for i in $(seq 1 60); do
HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 -H "Authorization: Bearer sk-1234" http://127.0.0.1:4000/litellm/health 2>/dev/null || true)
if [ "$HTTP_CODE" = "200" ]; then
echo "Prefixed proxy is ready"
exit 0
fi
sleep 2
done
echo "Prefixed proxy failed to start; see the 'Start LiteLLM proxy under a server root path' step for the boot log"
exit 1
- run:
name: Run migration smoke under SERVER_ROOT_PATH
command: |
cd ui/litellm-dashboard
LITELLM_LICENSE="$LITELLM_LICENSE" \
npx playwright test --config e2e_tests/migration.serverRootPath.config.ts
no_output_timeout: 10m
- store_artifacts:
path: ui/litellm-dashboard/test-results
destination: e2e-server-root-path-test-results
- store_artifacts:
path: ui/litellm-dashboard/playwright-report
destination: e2e-server-root-path-playwright-report
build_docker_database_image:
machine:
image: ubuntu-2204:2024.04.1
@ -2611,10 +2901,18 @@ workflows:
filters: *main_branches
- auth_ui_unit_tests:
filters: *main_branches
- proxy_behavior_tests:
filters: *main_branches
- proxy_security_tests:
filters: *main_branches
- schema_migration_check:
filters: *main_branches
- build_docker_database_image:
filters: *main_branches
- e2e_ui_testing:
filters: *main_branches
- e2e_ui_testing_server_root_path:
filters: *main_branches
- build_and_test:
requires:
- build_docker_database_image

View file

@ -8,3 +8,6 @@
# Update pydantic code to fix warnings (GH-3600)
876840e9957bc7e9f7d6a2b58c4d7c53dad16481
# style(ui): run prettier --write across the dashboard (#29622)
7edf3a9cb55548b143df1692f4ed7c4681d7fcf7

3
.gitattributes vendored
View file

@ -1 +1,2 @@
*.ipynb linguist-vendored
*.ipynb linguist-vendored
ui/litellm-dashboard/src/lib/http/schema.d.ts linguist-generated

75
.githooks/commit-msg Executable file
View file

@ -0,0 +1,75 @@
#!/usr/bin/env bash
#
# commit-msg — enforce Conventional Commits 1.0.0
# https://www.conventionalcommits.org/en/v1.0.0/
#
# Subject format: <type>(<scope>)!: <description>
# - <type> must be one of the angular types (feat, fix, ...)
# - (<scope>) is optional
# - ! is optional and marks a breaking change
# - <description> is mandatory and must be non-empty
#
# Bypass: commit with --no-verify.
# Merge, revert, fixup!, squash!, and amend! messages are passed through.
set -eu
COMMIT_MSG_FILE="${1:-}"
if [ -z "$COMMIT_MSG_FILE" ] || [ ! -f "$COMMIT_MSG_FILE" ]; then
echo "commit-msg: missing commit message file" >&2
exit 1
fi
# First non-comment, non-empty line is the subject.
subject=""
while IFS= read -r line || [ -n "$line" ]; do
case "$line" in
''|'#'*) continue ;;
esac
subject="$line"
break
done < "$COMMIT_MSG_FILE"
if [ -z "$subject" ]; then
echo "commit-msg: empty commit message" >&2
exit 1
fi
# Pass-through commits generated by git itself.
case "$subject" in
"Merge "*|"Revert \""*|"fixup! "*|"squash! "*|"amend! "*)
exit 0
;;
esac
ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert"
# Description must not start with an uppercase letter — kept in sync with the
# subjectPattern in .github/workflows/conventional-commits.yml so the local
# hook is the strictly tighter of the two gates. (Without this guard, a commit
# like "feat: Add thing" passes locally but fails the PR-title CI check.)
PATTERN="^(${ALLOWED_TYPES})(\([^)]+\))?!?: [^A-Z].*"
if printf '%s' "$subject" | grep -Eq "$PATTERN"; then
exit 0
fi
cat >&2 <<EOF
✗ Commit message does not follow Conventional Commits.
Got: $subject
Expected: <type>(<scope>)!: <description>
(description must start with a lowercase letter)
Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert
Examples:
feat(router): add weighted round-robin strategy
fix(bedrock): decouple STS region from aws_region_name
chore(deps): bump black to 26.3.1
refactor!: drop Python 3.8 support
See https://www.conventionalcommits.org/en/v1.0.0/
To bypass (use sparingly): git commit --no-verify
EOF
exit 1

92
.githooks/pre-push Executable file
View file

@ -0,0 +1,92 @@
#!/usr/bin/env bash
#
# pre-push — enforce Conventional Branches
# https://conventional-branch.github.io/
#
# Branch format: <type>/<description>
# <type> must be one of: feature, bugfix, hotfix, release, chore
#
# Protected branches (always allowed):
# - main
# - litellm_internal_staging
# - dependabot/*
# - gh-readonly-queue/*
#
# Tag pushes and branch deletions are skipped.
# Bypass: git push --no-verify.
set -eu
ZERO_OID="0000000000000000000000000000000000000000"
ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000"
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
PROTECTED_NAMES="main litellm_internal_staging"
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
is_protected() {
branch="$1"
for name in $PROTECTED_NAMES; do
if [ "$branch" = "$name" ]; then
return 0
fi
done
for prefix in $PROTECTED_PREFIXES; do
case "$branch" in "$prefix"*) return 0 ;; esac
done
return 1
}
invalid=""
while read -r local_ref local_oid remote_ref remote_oid; do
# Branch deletion (no local commit being pushed).
if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then
continue
fi
# Only validate branch pushes; ignore tags and other ref namespaces.
case "$remote_ref" in
refs/heads/*) ;;
*) continue ;;
esac
branch="${remote_ref#refs/heads/}"
if is_protected "$branch"; then
continue
fi
if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then
invalid="$invalid $branch"
fi
done
if [ -n "$invalid" ]; then
cat >&2 <<EOF
✗ Branch name does not follow Conventional Branches.
Invalid:$invalid
Expected: <type>/<description>
Allowed types: feature, bugfix, hotfix, release, chore
Examples:
feature/weighted-round-robin
bugfix/streaming-empty-chunks
chore/bump-deps
hotfix/auth-bypass
Protected (always allowed): main, litellm_internal_staging,
dependabot/*, gh-readonly-queue/*.
See https://conventional-branch.github.io/
Rename with: git branch -m <new-name>
To bypass (use sparingly): git push --no-verify
EOF
exit 1
fi
exit 0

View file

@ -10,9 +10,9 @@
**Please complete all items before asking a LiteLLM maintainer to review your PR**
- [ ] I have Added testing in the [`tests/test_litellm/`](https://github.com/BerriAI/litellm/tree/main/tests/test_litellm) directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code)
- [ ] I have added meaningful tests
- [ ] My PR passes all unit tests on [`make test-unit`](https://docs.litellm.ai/docs/extras/contributing_code)
- [ ] My PR's scope is as isolated as possible, it only solves 1 specific problem
- [ ] My PR's scope is as isolated as possible; it only solves 1 specific problem
- [ ] I have requested a Greptile review by commenting `@greptileai` and received a **Confidence Score of at least 4/5** before requesting a maintainer review
## Delays in PR merge?

View file

@ -27,6 +27,11 @@ on:
required: false
type: number
default: 10
dist:
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
required: false
type: string
default: "loadscope"
artifact-name:
description: "Unique name for the coverage artifact (must be unique per run)"
required: true
@ -82,18 +87,31 @@ jobs:
MAX_FAILURES: ${{ inputs.max-failures }}
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DIST: ${{ inputs.dist }}
run: |
uv run --no-sync pytest ${TEST_PATH:?} \
--tb=short -vv \
--maxfail="${MAX_FAILURES}" \
-n "${WORKERS}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--dist=loadscope \
--durations=20 \
--cov=./litellm \
--cov-report=xml:coverage.xml \
--cov-config=pyproject.toml
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \
--tb=short -vv \
--maxfail="${MAX_FAILURES}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--durations=20 \
--cov=./litellm \
--cov-report=xml:coverage.xml \
--cov-config=pyproject.toml
else
uv run --no-sync pytest ${TEST_PATH:?} \
--tb=short -vv \
--maxfail="${MAX_FAILURES}" \
-n "${WORKERS}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--dist="${DIST}" \
--durations=20 \
--cov=./litellm \
--cov-report=xml:coverage.xml \
--cov-config=pyproject.toml
fi
- name: Save coverage report
if: always()

View file

@ -1,190 +0,0 @@
name: _Unit Test Services Base (Reusable)
on:
workflow_call:
inputs:
test-path:
description: "Pytest path(s) to run"
required: true
type: string
workers:
description: "Number of pytest-xdist workers (0 = no parallelism)"
required: false
type: number
default: 2
reruns:
description: "Number of reruns for flaky tests"
required: false
type: number
default: 2
timeout-minutes:
description: "Job timeout in minutes"
required: false
type: number
default: 20
max-failures:
description: "Stop after this many failures"
required: false
type: number
default: 10
enable-postgres:
description: "Start a local Postgres service container and run Prisma migrations"
required: false
type: boolean
default: false
dist:
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
required: false
type: string
default: "loadscope"
artifact-name:
description: "Unique name for the coverage artifact (must be unique per run)"
required: false
type: string
default: "run"
permissions:
contents: read
# The postgres service container below is spawned per-job on localhost and
# destroyed with the job. Nothing outside the runner can reach it. The
# user/password/database here are not secrets — they're bootstrap values
# for a throwaway container — so we hardcode them instead of attaching
# every matrix shard to a GHA environment just to read three "secrets"
# (which also produces a "temporarily deployed to …" notification on the
# PR timeline per shard per push).
jobs:
run:
name: Run tests
runs-on: ubuntu-latest
timeout-minutes: ${{ inputs.timeout-minutes }}
services:
postgres:
image: postgres@sha256:705a5d5b5836f3fcba0d02c4d281e6a7dd9ed2dd4078640f08a1e1e9896e097d # postgres:14
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: litellm
POSTGRES_DB: litellm_test
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready"
--health-interval 10s
--health-timeout 5s
--health-retries 5
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
with:
version: "0.10.9"
- name: Cache uv dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-services-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-services-
- name: Install dependencies
run: |
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Run Prisma migrations
if: ${{ inputs.enable-postgres }}
env:
DATABASE_URL: "postgresql://litellm:litellm@localhost:5432/litellm_test"
run: |
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
- name: Run tests
env:
TEST_PATH: ${{ inputs.test-path }}
MAX_FAILURES: ${{ inputs.max-failures }}
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DIST: ${{ inputs.dist }}
DATABASE_URL: ${{ inputs.enable-postgres && 'postgresql://litellm:litellm@localhost:5432/litellm_test' || '' }}
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \
--tb=short -vv \
--maxfail="${MAX_FAILURES}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--durations=20 \
--cov=./litellm \
--cov-report=xml:coverage.xml \
--cov-config=pyproject.toml
else
uv run --no-sync pytest ${TEST_PATH:?} \
--tb=short -vv \
--maxfail="${MAX_FAILURES}" \
-n "${WORKERS}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--dist="${DIST}" \
--durations=20 \
--cov=./litellm \
--cov-report=xml:coverage.xml \
--cov-config=pyproject.toml
fi
- name: Save coverage report
if: always()
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}
path: coverage.xml
retention-days: 1
upload-coverage:
name: Upload coverage to Codecov
needs: run
if: always()
runs-on: ubuntu-latest
permissions:
contents: read
id-token: write
pull-requests: write
steps:
- name: Checkout code
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Download coverage report
uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e # v4.2.1
with:
pattern: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}
path: coverage-reports
merge-multiple: true
- name: Upload to Codecov
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
with:
use_oidc: true
directory: coverage-reports
root_dir: ${{ github.workspace }}
flags: ${{ inputs.artifact-name }}
fail_ci_if_error: false

View file

@ -0,0 +1,84 @@
name: Check UI API Types Sync
on:
pull_request:
paths:
- "litellm/proxy/**"
- "litellm/types/**"
- "ui/litellm-dashboard/src/lib/http/schema.d.ts"
- "ui/litellm-dashboard/scripts/gen-api-types.mjs"
- "ui/litellm-dashboard/package.json"
- "ui/litellm-dashboard/package-lock.json"
- ".github/workflows/check-ui-api-types.yml"
permissions:
contents: read
jobs:
check-sync:
name: Verify schema.d.ts matches the proxy OpenAPI spec
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
with:
version: "0.10.9"
- name: Cache uv dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Install backend dependencies
run: uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Set up Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
with:
node-version: "20"
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json
- name: Install dashboard dependencies
working-directory: ui/litellm-dashboard
run: npm ci
- name: Regenerate types from the live spec
working-directory: ui/litellm-dashboard
env:
LITELLM_PYTHON: "uv run --no-sync python"
run: npm run gen:api
- name: Fail if types are stale
run: |
if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec."
echo ""
echo "A backend route or model changed without regenerating the dashboard types."
echo "To fix, run from ui/litellm-dashboard:"
echo " npm run gen:api"
echo "then commit the updated src/lib/http/schema.d.ts."
exit 1
fi
echo "schema.d.ts is in sync with the proxy OpenAPI spec."

View file

@ -0,0 +1,46 @@
name: Conventional PR Title
# Squash-merge replaces the merge commit subject with the PR title, so
# enforcing Conventional Commits at the PR-title level is what actually gates
# the commits that land on the default branch. The local commit-msg hook
# (.githooks/commit-msg) is a best-effort assist; this workflow is the gate.
#
# See https://www.conventionalcommits.org/en/v1.0.0/
on:
pull_request:
types: [opened, edited, reopened, synchronize, labeled, unlabeled]
permissions:
pull-requests: read
jobs:
lint-pr-title:
name: Validate PR title
runs-on: ubuntu-latest
steps:
- name: Check title against Conventional Commits
uses: amannn/action-semantic-pull-request@48f256284bd46cdaab1048c3721360e808335d50 # v6.1.1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
# Must mirror the type list in .githooks/commit-msg.
types: |
feat
fix
docs
style
refactor
perf
test
build
ci
chore
revert
requireScope: false
subjectPattern: ^(?![A-Z]).+$
subjectPatternError: |
The subject "{subject}" must start with a lowercase character.
# Allow merges/reverts that GitHub generates automatically.
ignoreLabels: |
ignore-semantic-pull-request

View file

@ -63,3 +63,28 @@ jobs:
sha: commitHash,
});
core.info(`Created branch ${branchName} at ${commitHash}`);
- name: Create stable line branch
env:
TAG: ${{ inputs.tag }}
COMMIT_HASH: ${{ inputs.commit_hash }}
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
with:
script: |
const tag = process.env.TAG;
const commitHash = process.env.COMMIT_HASH;
const match = tag.match(/^v?(\d+)\.(\d+)\.0$/);
if (!match) {
core.info(`Tag ${tag} is not the X.Y.0 stable opener; skipping stable line branch`);
return;
}
const lineBranch = `stable/${match[1]}.${match[2]}.x`;
await github.rest.git.createRef({
owner: context.repo.owner,
repo: context.repo.repo,
ref: `refs/heads/${lineBranch}`,
sha: commitHash,
});
core.info(`Created branch ${lineBranch} at ${commitHash}`);

View file

@ -52,6 +52,22 @@ jobs:
// are stable maintenance releases, not pre-releases.
const isPrerelease = /(?:rc|nightly|alpha|beta|[-.]dev)/i.test(tag);
// A stable release should only claim the repo "latest" badge when its
// version is >= the current latest. Otherwise a backport (e.g. 1.84.6)
// would steal "latest" from a newer line (e.g. 1.88.1).
const versionKey = (rawTag) => {
const m = String(rawTag).match(/^v?(\d+)\.(\d+)\.(\d+)/);
if (!m) return null;
const maintenance = String(rawTag).match(/(?:\.post|\.patch\.)(\d+)/i);
return [Number(m[1]), Number(m[2]), Number(m[3]), maintenance ? Number(maintenance[1]) : 0];
};
const isAtLeast = (a, b) => {
for (let i = 0; i < a.length; i++) {
if (a[i] !== b[i]) return a[i] > b[i];
}
return true;
};
const cosignSection = [
`## Verify Docker Image Signature`,
``,
@ -90,6 +106,22 @@ jobs:
].join('\n');
try {
let makeLatest = "false";
const newVersion = versionKey(tag);
if (!isPrerelease && newVersion) {
let latestVersion = null;
try {
const latest = await github.rest.repos.getLatestRelease({
owner: context.repo.owner,
repo: context.repo.repo,
});
latestVersion = versionKey(latest.data.tag_name);
} catch (error) {
if (error.status !== 404) throw error;
}
makeLatest = (!latestVersion || isAtLeast(newVersion, latestVersion)) ? "true" : "false";
}
const response = await github.rest.repos.createRelease({
draft: true,
generate_release_notes: true,
@ -108,6 +140,7 @@ jobs:
release_id: response.data.id,
body: updatedBody,
draft: false,
make_latest: makeLatest,
});
} catch (error) {

View file

@ -67,6 +67,12 @@ jobs:
uv run --no-sync ruff check .
cd ..
- name: Check strict-rule budget (delta vs base)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
uv run --no-sync python scripts/ruff_strict_gate.py --base "$BASE_SHA"
- name: Print OpenAI version
run: |
uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')"

View file

@ -36,3 +36,79 @@ jobs:
- name: Build
run: npm run build
frontend-lint:
runs-on: ubuntu-latest
timeout-minutes: 8
defaults:
run:
working-directory: ui/litellm-dashboard
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Collect changed files
id: changed
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
: > "$RUNNER_TEMP/prettier_files.txt"
: > "$RUNNER_TEMP/eslint_files.txt"
while IFS= read -r f; do
[ -f "$f" ] || continue
case "$f" in
*.js | *.jsx | *.ts | *.tsx | *.mjs | *.cjs)
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt"
printf '%s\n' "$f" >> "$RUNNER_TEMP/eslint_files.txt" ;;
*.json | *.css | *.scss | *.md | *.mdx | *.yml | *.yaml | *.html)
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt" ;;
esac
done < <(git diff --name-only --diff-filter=ACMR --relative "$BASE_SHA"...HEAD -- .)
if [ -s "$RUNNER_TEMP/prettier_files.txt" ] || [ -s "$RUNNER_TEMP/eslint_files.txt" ]; then
echo "has_files=true" >> "$GITHUB_OUTPUT"
else
echo "has_files=false" >> "$GITHUB_OUTPUT"
echo "No lintable UI files changed in this PR; nothing to check."
fi
- name: Setup Node.js
if: steps.changed.outputs.has_files == 'true'
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
with:
node-version: "20"
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json
- name: Install dependencies
if: steps.changed.outputs.has_files == 'true'
run: npm ci
- name: Lint changed files (prettier + eslint)
if: steps.changed.outputs.has_files == 'true'
run: |
prettier_files=()
eslint_files=()
while IFS= read -r f; do prettier_files+=("$f"); done < "$RUNNER_TEMP/prettier_files.txt"
while IFS= read -r f; do eslint_files+=("$f"); done < "$RUNNER_TEMP/eslint_files.txt"
status=0
if [ ${#prettier_files[@]} -gt 0 ]; then
echo "::group::Prettier (${#prettier_files[@]} files)"
npx prettier --check "${prettier_files[@]}" || { status=1; echo "::error::Unformatted files. Fix with: npm run format"; }
echo "::endgroup::"
fi
if [ ${#eslint_files[@]} -gt 0 ]; then
echo "::group::ESLint (${#eslint_files[@]} files)"
npx eslint --no-warn-ignored --pass-on-unpruned-suppressions "${eslint_files[@]}" || status=1
echo "::endgroup::"
fi
exit $status
- name: Check lint budgets
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
run: |
npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json

View file

@ -28,6 +28,8 @@ jobs:
tests/test_litellm/completion_extras
tests/test_litellm/containers
tests/test_litellm/experimental_mcp_client
tests/test_litellm/models
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/passthrough

View file

@ -1,9 +1,10 @@
name: "Unit Tests: Proxy DB Operations"
# Uses DATABASE_URL secret — only runs on trusted branches, not PRs.
on:
push:
branches: [main, "litellm_**"]
pull_request:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -30,9 +31,6 @@ concurrency:
# xdist balances its 188 parametrized cases across workers instead of
# pinning the whole file to one worker (the default --dist=loadscope
# behavior for single-file targets).
# * test_db_schema_migration.py is isolated because one test in it
# (test_aaaasschema_migration_check) takes ~170s — by itself it
# determines the shard's wall-clock floor.
jobs:
# Fast guard — fails the workflow if a test_*.py file under
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
@ -166,18 +164,6 @@ jobs:
dist: loadscope
timeout: 15
# ---- db-and-spend: isolate the 170s schema-migration test ----
# test_db_schema_migration.py has exactly one test, and that test
# is mostly waiting on `prisma migrate deploy` / `prisma migrate
# diff` subprocesses (~170s). It does no CPU-bound Python work
# inside the test. Running with workers=0 (serial, no xdist)
# skips the 4-worker cold-start cost we'd otherwise pay for a
# single test, saving ~4 minutes of wall-clock.
- test-group: schema-migration
test-path: "tests/proxy_unit_tests/test_db_schema_migration.py"
workers: 0
dist: loadscope
timeout: 15
- test-group: db-and-spend
test-path: >-
tests/proxy_unit_tests/test_prisma_client_backoff_retry.py
@ -232,12 +218,11 @@ jobs:
workers: 4
dist: loadscope
timeout: 15
uses: ./.github/workflows/_test-unit-services-base.yml
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: ${{ matrix.test-path }}
workers: ${{ matrix.workers }}
reruns: 2
timeout-minutes: ${{ matrix.timeout }}
enable-postgres: true
dist: ${{ matrix.dist }}
artifact-name: proxy-db-${{ matrix.test-group }}

View file

@ -33,13 +33,16 @@ jobs:
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/discovery_endpoints
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/shutdown
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/utils
workers: 2
reruns: 2
artifact-name: proxy-endpoints

View file

@ -1,34 +0,0 @@
name: "Unit Tests: Proxy Management-Endpoint Behavior Pinning"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_branch
- "litellm_**"
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
proxy-mgmt-behavior:
uses: ./.github/workflows/_test-unit-services-base.yml
with:
test-path: tests/proxy_behavior
# workers=0 (no xdist): the world seed is a single shared Postgres
# state — two xdist workers both call seed_world() and race on the
# ``behavior-pin-budget`` row, producing UniqueViolation + cascading
# missing-membership FK failures. The whole suite is ~7s sequentially,
# so the cost of disabling parallelism here is negligible.
workers: 0
reruns: 0
enable-postgres: true
artifact-name: proxy-mgmt-behavior
timeout-minutes: 15

View file

@ -1,28 +0,0 @@
name: "Unit Tests: Security"
# Kept push-only (was previously required by DATABASE_URL secret scoping;
# now the postgres credentials are ephemeral localhost values but the
# push-trigger stays to match the proxy-db workflow cadence).
on:
push:
branches: [main, "litellm_**"]
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
security:
uses: ./.github/workflows/_test-unit-services-base.yml
with:
test-path: "tests/proxy_security_tests/"
workers: 1
reruns: 2
timeout-minutes: 20
enable-postgres: true
artifact-name: security

View file

@ -101,6 +101,31 @@ jobs:
docker logs litellm-test
exit 1
- name: Setup Node for Playwright
uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
with:
node-version: "20"
- name: Install UI deps and Chromium
working-directory: ui/litellm-dashboard
run: |
npm ci
npx playwright install --with-deps chromium
- name: Run SERVER_ROOT_PATH redirect e2e
working-directory: ui/litellm-dashboard
env:
SERVER_ROOT_PATH: ${{ matrix.root_path }}
run: npx playwright test --config=e2e_tests/serverRootPath.config.ts
- name: Upload Playwright artifacts on failure
if: failure()
uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
with:
name: playwright-trace-${{ strategy.job-index }}
path: ui/litellm-dashboard/test-results/
retention-days: 7
- name: Cleanup
if: always()
run: |

5
.gitignore vendored
View file

@ -28,6 +28,8 @@ litellm/tests/config_*.yaml
litellm/tests/langfuse.log
langfuse.log
.langfuse.log
.pin_list.txt
.cov_new.xml
litellm/tests/test_custom_logger.py
litellm/tests/langfuse.log
litellm/tests/dynamo*.log
@ -120,4 +122,5 @@ crash.log
crash.*.log
# .terraform.lock.hcl is intentionally NOT ignored — it pins provider versions
# and should be committed.
.vscode
.vscode
.pin_list.txt

294
AGENTS.md
View file

@ -1,293 +1 @@
# INSTRUCTIONS FOR LITELLM
This document provides comprehensive instructions for AI agents working in the LiteLLM repository.
## OVERVIEW
LiteLLM is a unified interface for 100+ LLMs that:
- Translates inputs to provider-specific completion, embedding, and image generation endpoints
- Provides consistent OpenAI-format output across all providers
- Includes retry/fallback logic across multiple deployments (Router)
- Offers a proxy server (LLM Gateway) with budgets, rate limits, and authentication
- Supports advanced features like function calling, streaming, caching, and observability
## REPOSITORY STRUCTURE
### Core Components
- `litellm/` - Main library code
- `llms/` - Provider-specific implementations (OpenAI, Anthropic, Azure, etc.)
- `proxy/` - Proxy server implementation (LLM Gateway)
- `router_utils/` - Load balancing and fallback logic
- `types/` - Type definitions and schemas
- `integrations/` - Third-party integrations (observability, caching, etc.)
### Key Directories
- `tests/` - Comprehensive test suites
- `ui/litellm-dashboard/` - Admin dashboard UI
- `enterprise/` - Enterprise-specific features
Documentation lives in the separate [BerriAI/litellm-docs](https://github.com/BerriAI/litellm-docs) repository and is served at [docs.litellm.ai](https://docs.litellm.ai).
## DEVELOPMENT GUIDELINES
### MAKING CODE CHANGES
1. **Provider Implementations**: When adding/modifying LLM providers:
- Follow existing patterns in `litellm/llms/{provider}/`
- Implement proper transformation classes that inherit from `BaseConfig`
- Support both sync and async operations
- Handle streaming responses appropriately
- Include proper error handling with provider-specific exceptions
2. **Type Safety**:
- Use proper type hints throughout
- Update type definitions in `litellm/types/`
- Ensure compatibility with both Pydantic v1 and v2
3. **Testing**:
- Add tests in appropriate `tests/` subdirectories
- Include both unit tests and integration tests
- Test provider-specific functionality thoroughly
- Consider adding load tests for performance-critical changes
### MAKING CODE CHANGES FOR THE UI (IGNORE FOR BACKEND)
1. **Always use `antd` for new UI components — Tremor is DEPRECATED**
- We are migrating off of `@tremor/react`. Do not introduce new `Badge`, `Text`, `Card`, `Grid`, `Title`, or other imports from `@tremor/react` in any new or modified file.
- Use `antd` equivalents: `Tag` for labels, plain `<span>`/`<div>` with Tailwind classes (or `Typography.Text`) for text, `Card` from `antd`, etc. Note that `antd` has no `"yellow"` Tag color — use `"gold"` for amber/yellow.
- The only exception is the Tremor Table component and its required Tremor Table sub components.
2. **Use Common Components as much as possible**:
- These are usually defined in the `common_components` directory
- Use these components as much as possible and avoid building new components unless needed
3. **Testing**:
- The codebase uses **Vitest** and **React Testing Library**
- **Query Priority Order**: Use query methods in this order: `getByRole`, `getByLabelText`, `getByPlaceholderText`, `getByText`, `getByTestId`
- **Always use `screen`** instead of destructuring from `render()` (e.g., use `screen.getByText()` not `getByText`)
- **Wrap user interactions in `act()`**: Always wrap `fireEvent` calls with `act()` to ensure React state updates are properly handled
- **Use `query` methods for absence checks**: Use `queryBy*` methods (not `getBy*`) when expecting an element to NOT be present
- **Test names must start with "should"**: All test names should follow the pattern `it("should ...")`
- **Mock external dependencies**: Check `setupTests.ts` for global mocks and mock child components/networking calls as needed
- **Structure tests properly**:
- First test should verify the component renders successfully
- Subsequent tests should focus on functionality and user interactions
- Use `waitFor` for async operations that aren't already awaited
- **Avoid using `querySelector`**: Prefer React Testing Library queries over direct DOM manipulation
### IMPORTANT PATTERNS
1. **Function/Tool Calling**:
- LiteLLM standardizes tool calling across providers
- OpenAI format is the standard, with transformations for other providers
- See `litellm/llms/anthropic/chat/transformation.py` for complex tool handling
2. **Streaming**:
- All providers should support streaming where possible
- Use consistent chunk formatting across providers
- Handle both sync and async streaming
3. **Error Handling**:
- Use provider-specific exception classes
- Maintain consistent error formats across providers
- Include proper retry logic and fallback mechanisms
4. **Configuration**:
- Support both environment variables and programmatic configuration
- Use `BaseConfig` classes for provider configurations
- Allow dynamic parameter passing
## PROXY SERVER (LLM GATEWAY)
The proxy server is a critical component that provides:
- Authentication and authorization
- Rate limiting and budget management
- Load balancing across multiple models/deployments
- Observability and logging
- Admin dashboard UI
- Enterprise features
Key files:
- `litellm/proxy/proxy_server.py` - Main server implementation
- `litellm/proxy/auth/` - Authentication logic
- `litellm/proxy/management_endpoints/` - Admin API endpoints
**Database (proxy)**: Use Prisma model methods (`prisma_client.db.<model>.upsert`, `.find_many`, `.find_unique`, etc.), not raw SQL (`execute_raw`/`query_raw`). See COMMON PITFALLS for details.
## MCP (MODEL CONTEXT PROTOCOL) SUPPORT
LiteLLM supports MCP for agent workflows:
- MCP server integration for tool calling
- Transformation between OpenAI and MCP tool formats
- Support for external MCP servers (Zapier, Jira, Linear, etc.)
- See `litellm/experimental_mcp_client/` and `litellm/proxy/_experimental/mcp_server/`
## RUNNING SCRIPTS
Use `uv run python script.py` to run Python scripts in the project environment (for non-test files).
## GITHUB TEMPLATES
When opening issues or pull requests, follow these templates:
### Bug Reports (`.github/ISSUE_TEMPLATE/bug_report.yml`)
- Describe what happened vs. expected behavior
- Include relevant log output
- Specify LiteLLM version
- Indicate if you're part of an ML Ops team (helps with prioritization)
### Feature Requests (`.github/ISSUE_TEMPLATE/feature_request.yml`)
- Clearly describe the feature
- Explain motivation and use case with concrete examples
### Pull Requests (`.github/pull_request_template.md`)
- Add at least 1 test in `tests/litellm/`
- Ensure `make test-unit` passes
## TESTING CONSIDERATIONS
1. **Provider Tests**: Test against real provider APIs when possible
2. **Proxy Tests**: Include authentication, rate limiting, and routing tests
3. **Performance Tests**: Load testing for high-throughput scenarios
4. **Integration Tests**: End-to-end workflows including tool calling
## DOCUMENTATION
- Keep documentation in sync with code changes
- Update provider documentation when adding new providers
- Include code examples for new features
- Update changelog and release notes
## SECURITY CONSIDERATIONS
- Handle API keys securely
- Validate all inputs, especially for proxy endpoints
- Consider rate limiting and abuse prevention
- Follow security best practices for authentication
## ENTERPRISE FEATURES
- Some features are enterprise-only
- Check `enterprise/` directory for enterprise-specific code
- Maintain compatibility between open-source and enterprise versions
## COMMON PITFALLS TO AVOID
1. **Breaking Changes**: LiteLLM has many users - avoid breaking existing APIs
2. **Provider Specifics**: Each provider has unique quirks - handle them properly
3. **Rate Limits**: Respect provider rate limits in tests
4. **Memory Usage**: Be mindful of memory usage in streaming scenarios
5. **Dependencies**: Keep dependencies minimal and well-justified
6. **UI/Backend Contract Mismatch**: When adding a new entity type to the UI, always check whether the backend endpoint accepts a single value or an array. Match the UI control accordingly (single-select vs. multi-select) to avoid silently dropping user selections
7. **Missing Tests for New Entity Types**: When adding a new entity type (e.g., in `EntityUsage`, `UsageViewSelect`), always add corresponding tests in the existing test files and update any icon/component mocks
8. **Raw SQL in proxy DB code**: Do not use `execute_raw` or `query_raw` for proxy database access. Use Prisma model methods (e.g. `prisma_client.db.litellm_tooltable.upsert()`, `.find_many()`, `.find_unique()`) so behavior stays consistent with the schema, the client stays mockable in tests, and you avoid the pitfalls of hand-written SQL (parameter ordering, type casting, schema drift)
8. **Do not hardcode model-specific flags**: Put model-specific capability flags in `model_prices_and_context_window.json` and read them via `get_model_info` (or existing helpers like `supports_reasoning`). This prevents users from needing to upgrade LiteLLM each time a new model supports a feature.
**Example of BAD** (hardcoded model checks):
```python
@staticmethod
def _is_effort_supported_model(model: str) -> bool:
"""Check if the model supports the output_config.effort parameter..."""
model_lower = model.lower()
if AnthropicConfig._is_claude_4_6_model(model):
return True
return any(
v in model_lower for v in ("opus-4-5", "opus_4_5", "opus-4.5", "opus_4.5")
)
```
**Example of GOOD** (config-driven or helper that reads from config):
```python
if (
"claude-3-7-sonnet" in model
or AnthropicConfig._is_claude_4_6_model(model)
or supports_reasoning(
model=model,
custom_llm_provider=self.custom_llm_provider,
)
):
...
```
Using helpers like `supports_reasoning` (which read from `model_prices_and_context_window.json` / `get_model_info`) allows future model updates to "just work" without code changes.
9. **Never close HTTP/SDK clients on cache eviction**: Do not add `close()`, `aclose()`, or `create_task(close_fn())` inside `LLMClientCache._remove_key()` or any cache eviction path. Evicted clients may still be held by in-flight requests; closing them causes `RuntimeError: Cannot send a request, as the client has been closed.` in production after the cache TTL (1 hour) expires. Connection cleanup is handled at shutdown by `close_litellm_async_clients()`. See PR #22247 for the full incident history.
## HELPFUL RESOURCES
- Main documentation: https://docs.litellm.ai/ (source: [BerriAI/litellm-docs](https://github.com/BerriAI/litellm-docs))
- Provider-specific docs: https://docs.litellm.ai/docs/providers/
- Admin UI for testing proxy features
## WHEN IN DOUBT
- Follow existing patterns in the codebase
- Check similar provider implementations
- Ensure comprehensive test coverage
- Update documentation appropriately
- Consider backward compatibility impact
## Cursor Cloud specific instructions
### Environment
- uv is installed in `~/.local/bin`; the update script ensures it is on `PATH`.
- Python 3.12, Node 22 are pre-installed.
- The project virtual environment lives under `.venv/`.
### Running the proxy server
Create a minimal config file and start the proxy:
```yaml
# config.yaml
model_list:
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake-model
api_key: fake-key
api_base: https://fake-api.example.com
general_settings:
master_key: sk-1234
litellm_settings:
drop_params: True
telemetry: False
```
```bash
uv run litellm --config config.yaml --port 4000
```
The proxy takes ~15-20 seconds to fully start (it runs Prisma migrations on boot). Wait for `/health` to return before sending requests. Without a PostgreSQL `DATABASE_URL`, the proxy connects to a default Neon dev database embedded in the `litellm-proxy-extras` package.
### Running tests
See `CLAUDE.md` and the `Makefile` for standard commands. Key notes:
- `uv sync --group proxy-dev --extra proxy` installs the Prisma and proxy-side test dependencies used by the standard local workflow.
- The `--timeout` pytest flag is NOT available; don't pass it.
- Unit tests: `uv run pytest tests/test_litellm/ -x -vv -n 4`
- **Before committing, always run `uv run black .` to format your code.** Black formatting is enforced in CI.
- If `uv sync` fails because the lockfile is outdated, run `uv lock` and retry.
### Lint
```bash
cd litellm && uv run ruff check .
```
Ruff is the primary fast linter. For the full lint suite (including mypy, black, circular imports), run `make lint` per `CLAUDE.md`.
### UI Dashboard development
- The UI is at `ui/litellm-dashboard/`. Run `npm run dev` from that directory for the Next.js dev server on port 3000.
- The proxy at port 4000 serves a **pre-built** static UI from `litellm/proxy/_experimental/out/`. After making UI code changes, you must run `npm run build` in the dashboard directory and copy the output: `cp -r ui/litellm-dashboard/out/* litellm/proxy/_experimental/out/` for the proxy to serve the updated UI.
- SVGs used as provider logos (loaded via `<img>` tags) must NOT use `fill="currentColor"` — replace with an explicit color like `#000000` or use the `-color` variant from lobehub icons, since CSS color inheritance does not work inside `<img>` elements.
- Provider logos live in `ui/litellm-dashboard/public/assets/logos/` (source) and `litellm/proxy/_experimental/out/assets/logos/` (pre-built). Both locations must have the file for it to work in dev and proxy-served modes.
- UI Vitest tests: `cd ui/litellm-dashboard && npx vitest run`
Read @CLAUDE.md for coding guidelines

View file

@ -240,6 +240,24 @@ graph LR
7. `DBSpendUpdateWriter.update_database()` queues spend increments to Redis
8. Background job `update_spend` flushes queued spend to PostgreSQL every 60s
### Data Access Layer (Models & Repositories)
Database entities and the operations on them live in two packages at the root of `litellm/` so both the gateway (`proxy/`) and the SDK can use them without importing proxy internals:
- `litellm/models/` holds the canonical Pydantic definitions for every persisted entity (`LiteLLM_VerificationToken`, `LiteLLM_TeamTable`, `LiteLLM_UserTable`, etc.). `proxy/_types.py` re-exports these for backwards compatibility, so existing imports keep working.
- `litellm/repositories/` holds the data-access layer. `BaseRepository[T]` provides the generic CRUD (`find_by_id`, `find_many`, `create`, `update`, `delete`, `count`, `exists`); entity repositories such as `VerificationTokenRepository`, `TeamRepository`, and `UserRepository` add domain-specific queries and writes on top of it.
Conventions to follow when touching this layer:
| Concern | How it's handled |
|---------|------------------|
| JSON columns | Prisma `Json` columns are stored as JSON strings. Repositories `json.dumps()` on write and `json.loads()` on read (see `_to_model` and the `_build_*_data` helpers). |
| Archive-then-delete | `delete_team` / `delete_token` copy the row into the `LiteLLM_Deleted*` table and delete the original inside a single `prisma_client.db.tx()` transaction. Archive payloads are built explicitly so only columns that exist on the archive table are written. |
| Column vs. field names | Where a model field differs from its DB column (for example `org_id` maps to the `organization_id` column), the repository translates in both directions rather than relying on Pydantic to guess. |
| Array mutations | Adds use Prisma's atomic `push` (`add_member`, `add_admin`, `add_models`) to avoid read-modify-write races. Removals fall back to read-modify-write because Prisma has no atomic array remove. |
To add a new entity, define the model under `litellm/models/`, re-export it from `proxy/_types.py` if existing code imports it from there, and add a repository under `litellm/repositories/` (subclass `BaseRepository` for plain CRUD, or add bespoke methods when the entity needs encryption, archiving, or atomic array updates). Mirror the tests in `tests/test_litellm/repositories/`.
---
## 2. SDK Request Flow

211
CLAUDE.md
View file

@ -1,181 +1,94 @@
# CLAUDE.md
Do not write comments unless they are absolutely necessary to explain some very complex business logic. Please clean up if there are comments that are not absolutely necessary. Do not remove comments that are unrelated to the addition of the code of this PR
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
Explanation: code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive to the reader, while being both easy to maintain and high performance
## Documentation
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
- correct
- secure
- performant
- readable
- easy to maintain/change
- modern
Documentation lives in a separate repository: [BerriAI/litellm-docs](https://github.com/BerriAI/litellm-docs). It is served at [docs.litellm.ai](https://docs.litellm.ai). Do not create or edit documentation files in this repository — open doc PRs against `BerriAI/litellm-docs` instead.
In that order of importance
## Development Commands
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
### Installation
- `make install-dev` - Install core development dependencies
- `make install-proxy-dev` - Install proxy development dependencies with full feature set
- `make install-test-deps` - Install the full local test environment and generate the Prisma client
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)
### Testing
- `make test` - Run all tests
- `make test-unit` - Run unit tests (tests/test_litellm) with 4 parallel workers
- `make test-integration` - Run integration tests (excludes unit tests)
- `pytest tests/` - Direct pytest execution
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
### Code Quality
- `make lint` - Run all linting (Ruff, MyPy, Black, circular imports, import safety)
- `make format` - Apply Black code formatting
- `make lint-ruff` - Run Ruff linting only
- `make lint-mypy` - Run MyPy type checking only
- **Before committing, always run `uv run black .` to format your code.** Black formatting is enforced in CI.
When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose
### Single Test Files
- `uv run pytest tests/path/to/test_file.py -v` - Run specific test file
- `uv run pytest tests/path/to/test_file.py::test_function -v` - Run specific test
Always use @.github/pull_request_template.md as a guide for your PR body
### Running Scripts
- `uv run python script.py` - Run Python scripts (use for non-test files)
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
### GitHub Issue & PR Templates
When contributing to the project, use the appropriate templates:
If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ";", ".", etc.
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file)
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
**Bug Reports** (`.github/ISSUE_TEMPLATE/bug_report.yml`):
- Describe what happened vs. what you expected
- Include relevant log output
- Specify your LiteLLM version
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
**Feature Requests** (`.github/ISSUE_TEMPLATE/feature_request.yml`):
- Describe the feature clearly
- Explain the motivation and use case
Run tests, format your code, and lint your code before each commit
**Pull Requests** (`.github/pull_request_template.md`):
- Add at least 1 test in `tests/litellm/`
- Ensure `make test-unit` passes
When you fix strict-rule violations gated by `ruff-strict-budget.json`, run `make lint-strict-budget-update` and commit the lowered baselines so the ceilings ratchet down instead of leaving stale headroom
## Architecture Overview
Ask to commit and push your work when you're done (or if you're confident that your code is good and works, just do it)
LiteLLM is a unified interface for 100+ LLM providers with two main components:
When you must use real LLM models to, for example, write e2e tests, write a QA runbook, etc., make sure to use the latest models (doesn't have to be smartest, can also be a modern small, fast one. No strong preference for smart vs fast here, just use something modern) as of the year and month of the current date. Do a web search as necessary to figure that out
### Core Library (`litellm/`)
- **Main entry point**: `litellm/main.py` - Contains core completion() function
- **Provider implementations**: `litellm/llms/` - Each provider has its own subdirectory
- **Router system**: `litellm/router.py` + `litellm/router_utils/` - Load balancing and fallback logic
- **Type definitions**: `litellm/types/` - Pydantic models and type hints
- **Integrations**: `litellm/integrations/` - Third-party observability, caching, logging
- **Caching**: `litellm/caching/` - Multiple cache backends (Redis, in-memory, S3, etc.)
If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
### Proxy Server (`litellm/proxy/`)
- **Main server**: `proxy_server.py` - FastAPI application
- **Authentication**: `auth/` - API key management, JWT, OAuth2
- **Database**: `db/` - Prisma ORM with PostgreSQL/SQLite support
- **Management endpoints**: `management_endpoints/` - Admin APIs for keys, teams, models
- **Pass-through endpoints**: `pass_through_endpoints/` - Provider-specific API forwarding
- **Guardrails**: `guardrails/` - Safety and content filtering hooks
- **UI Dashboard**: Served from `_experimental/out/` (Next.js build)
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
## Key Patterns
When working on a PR, keep the PR description in sync with new commits being made
### Provider Implementation
- Providers inherit from base classes in `litellm/llms/base.py`
- Each provider has transformation functions for input/output formatting
- Support both sync and async operations
- Handle streaming responses and function calling
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
### Error Handling
- Provider-specific exceptions mapped to OpenAI-compatible errors
- Fallback logic handled by Router system
- Comprehensive logging through `litellm/_logging.py`
Do not put names of customers or customer company names in code, PRs, and issues. The codebase is public
### Configuration
- YAML config files for proxy server (see `proxy/example_config_yaml/`)
- Environment variables for API keys and settings
- Database schema managed via Prisma (`proxy/schema.prisma`)
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
## Development Notes
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
### Code Style
- Uses Black formatter, Ruff linter, MyPy type checker
- Pydantic v2 for data validation
- Async/await patterns throughout
- Type hints required for all public APIs
- **Avoid imports within methods** — place all imports at the top of the file (module-level). Inline imports inside functions/methods make dependencies harder to trace and hurt readability. The only exception is avoiding circular imports where absolutely necessary.
- **Use dict spread for immutable copies** — prefer `{**original, "key": new_value}` over `dict(obj)` + mutation. The spread produces the final dict in one step and makes intent clear.
- **Guard at resolution time** — when resolving an optional value through a fallback chain (`a or b or ""`), raise immediately if the resolved result being empty is an error. Don't pass empty strings or sentinel values downstream for the callee to deal with.
- **Extract complex comprehensions to named helpers** — a set/dict comprehension that calls into the DB or manager (e.g. "which of these server IDs are OAuth2?") belongs in a named helper function, not inline in the caller.
- **FastAPI parameter declarations** — mark required query/form params with `= Query(...)` / `= Form(...)` explicitly when other params in the same handler are optional. Mixing `str` (required) with `Optional[str] = None` in the same signature causes silent 422s when the required param is missing.
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), etc.
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match
- No monster files or god objects
- No file sprawl: deliberate file and folder structure
- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions
### Testing Strategy
- Unit tests in `tests/test_litellm/`
- Integration tests for each provider in `tests/llm_translation/`
- Proxy tests in `tests/proxy_unit_tests/`
- Load tests in `tests/load_tests/`
- **Always add tests when adding new entity types or features** — if the existing test file covers other entity types, add corresponding tests for the new one
- **Keep monkeypatch stubs in sync with real signatures** — when a function gains a new optional parameter, update every `fake_*` / `stub_*` in tests that patch it to also accept that kwarg (even as `**kwargs`). Stale stubs fail with `unexpected keyword argument` and mask real bugs.
- **Test all branches of name→ID resolution** — when adding server/resource lookup that resolves names to UUIDs, test: (1) name resolves and UUID is allowed, (2) name resolves but UUID is not allowed, (3) name does not resolve at all. The silent-fallback path is where access-control bugs hide.
if you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and bringing it closer to the max, just validate it in the caller (a simple function that returns the typed thing or raises will do) and then pass the now typed variable in
### UI / Backend Consistency
- When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select
Follow conventional commits for commit names and PR titles
### UI Component Library
- **Always use `antd` for new UI components** — we are migrating off of `@tremor/react`. Do not introduce new `Badge`, `Text`, `Card`, `Grid`, `Title`, or other imports from `@tremor/react` in any new or modified file. Use `antd` equivalents: `Tag` for labels, `Typography.Text` / `Typography.Title` / `Typography.Paragraph` for textual content (avoid plain text-only `<span>`, `<p>`, `<h*>` when Typography fits), and `Card` from `antd`. Note that `antd` has no `"yellow"` Tag color — use `"gold"` for amber/yellow.
## Think Before Coding
### MCP OAuth / OpenAPI Transport Mapping
- **`available_on_public_internet: false` with `delegate_auth_to_upstream: true` (oauth2, interactive — not `client_credentials`)** — LiteLLM still allows the anonymous upstream PKCE path (no proxy API key for `/authorize` and matching MCP routes). The internal-only flag mainly affects other surfaces (e.g. IP-based discovery). Rely on the upstream IdP and network policy; the dashboard shows a warning when both are set, and the proxy logs a warning when the server is loaded from config or the database.
- `TRANSPORT.OPENAPI` is a UI-only concept. The backend only accepts `"http"`, `"sse"`, or `"stdio"`. Always map it to `"http"` before any API call (including pre-OAuth temp-session calls).
- FastAPI validation errors return `detail` as an array of `{loc, msg, type}` objects. Error extractors must handle: array (map `.msg`), string, nested `{error: string}`, and fallback.
- When an MCP server already has `authorization_url` stored, skip OAuth discovery (`_discovery_metadata`) — the server URL for OpenAPI MCPs is the spec file, not the API base, and fetching it causes timeouts.
- `client_id` should be optional in the `/authorize` endpoint — if the server has a stored `client_id` in credentials, use that. Never require callers to re-supply it.
**Don't assume. Don't hide confusion. Surface tradeoffs**
### MCP Credential Storage
- OAuth credentials and BYOK credentials share the `litellm_mcpusercredentials` table, distinguished by a `"type"` field in the JSON payload (`"oauth2"` vs plain string).
- When deleting OAuth credentials, check type before deleting to avoid accidentally deleting a BYOK credential for the same `(user_id, server_id)` pair.
- Always pass the raw `expires_at` timestamp to the client — never set it to `None` for expired credentials. Let the frontend compute the "Expired" display state from the timestamp.
- Use `RecordNotFoundError` (not bare `except Exception`) when catching "already deleted" in credential delete endpoints.
Before implementing:
- State your assumptions explicitly. If uncertain, ask
- If multiple interpretations exist, present them. Don't pick silently
- If a simpler approach exists, say so. Push back when warranted
- If something is unclear, stop. Name what's confusing. Ask
### Browser Storage Safety (UI)
- Never write LiteLLM access tokens or API keys to `localStorage` — use `sessionStorage` only. `localStorage` survives browser close and is readable by any injected script (XSS).
- Shared utility functions (e.g. `extractErrorMessage`) belong in `src/utils/` — never define them inline in hooks or duplicate them across files.
## Simplicity First
### Database Migrations
- Prisma handles schema migrations
- Migration files auto-generated with `prisma migrate dev`
- Always test migrations against both PostgreSQL and SQLite
**Minimum code that solves the problem. Nothing speculative**
### Proxy database access
- **Do not write raw SQL** for proxy DB operations. Use Prisma model methods instead of `execute_raw` / `query_raw`.
- Use the generated client: `prisma_client.db.<model>` (e.g. `litellm_tooltable`, `litellm_usertable`) with `.upsert()`, `.find_many()`, `.find_unique()`, `.update()`, `.update_many()` as appropriate. This avoids schema/client drift, keeps code testable with simple mocks, and matches patterns used in spend logs and other proxy code.
- **No N+1 queries.** Never query the DB inside a loop. Batch-fetch with `{"in": ids}` and distribute in-memory.
- **Batch writes.** Use `create_many`/`update_many`/`delete_many` instead of individual calls (these return counts only; `update_many`/`delete_many` no-op silently on missing rows). When multiple separate writes target the same table (e.g. in `batch_()`), order by primary key to avoid deadlocks.
- **Push work to the DB.** Filter, sort, group, and aggregate in SQL, not Python. Verify Prisma generates the expected SQL — e.g. prefer `group_by` over `find_many(distinct=...)` which does client-side processing.
- **Bound large result sets.** Prisma materializes full results in memory. For results over ~10 MB, paginate with `take`/`skip` or `cursor`/`take`, always with an explicit `order`. Prefer cursor-based pagination (`skip` is O(n)). Don't paginate naturally small result sets.
- **Limit fetched columns on wide tables.** Use `select` to fetch only needed fields — returns a partial object, so downstream code must not access unselected fields.
- **Check index coverage.** For new or modified queries, check `schema.prisma` for a supporting index. Prefer extending an existing index (e.g. `@@index([a])` → `@@index([a, b])`) over adding a new one, unless it's a `@@unique`. Only add indexes for large/frequent queries.
- **Keep schema files in sync.** Apply schema changes to all `schema.prisma` copies (`schema.prisma`, `litellm/proxy/`, `litellm-proxy-extras/`) with a migration under `litellm-proxy-extras/litellm_proxy_extras/migrations/`.
- No features beyond what was asked
- No abstractions for single-use code
- No "flexibility" or "configurability" that wasn't requested
- No error handling for impossible scenarios
- If you write 200 lines and it could be 50, rewrite it
### Setup Wizard (`litellm/setup_wizard.py`)
- The wizard is implemented as a single `SetupWizard` class with `@staticmethod` methods — keep it that way. No module-level functions except `run_setup_wizard()` (the public entrypoint) and pure helpers (color, ANSI).
- Use `litellm.utils.check_valid_key(model, api_key)` for credential validation — never roll a custom completion call.
- Do not hardcode provider env-key names or model lists that already exist in the codebase. Add a `test_model` field to each provider entry to drive `check_valid_key`; set it to `None` for providers that can't be validated with a single API key (Azure, Bedrock, Ollama).
### Enterprise Features
- Enterprise-specific code in `enterprise/` directory
- Optional features enabled via environment variables
- Separate licensing and authentication for enterprise features
### CI Supply-Chain Safety
- **Never pipe a remote script into a shell** (`curl ... | bash`, `wget ... | sh`). Download the artifact to a file, verify its SHA-256 checksum, then install.
- **Pin every external tool to a specific version** with a full URL (not `latest` or `stable`). Unversioned downloads silently change under you.
- **Verify checksums for all downloaded binaries.** Use the provider's official `.sha256` / `.sha256sum` sidecar file when available; otherwise compute and hardcode the digest.
- **Prefer reusable CircleCI commands** (`commands:` section) so a tool is installed and verified in exactly one place, then referenced everywhere with `- install_<tool>` or `- wait_for_service`.
- **Don't add tools just because they were there before.** Audit whether an external dependency is still needed. If it can be replaced with a shell one-liner or a tool already in the image, remove it.
- These rules apply to every download in CI: binaries, install scripts, language version managers, package repos. No exceptions.
### HTTP Client Cache Safety
- **Never close HTTP/SDK clients on cache eviction.** `LLMClientCache._remove_key()` must not call `close()`/`aclose()` on evicted clients — they may still be used by in-flight requests. Doing so causes `RuntimeError: Cannot send a request, as the client has been closed.` after the 1-hour TTL expires. Cleanup happens at shutdown via `close_litellm_async_clients()`.
### Troubleshooting: DB schema out of sync after proxy restart
`litellm-proxy-extras` runs `prisma migrate deploy` on startup using **its own** bundled migration files, which may lag behind schema changes in the current worktree. Symptoms: `Unknown column`, `Invalid prisma invocation`, or missing data on new fields.
**Diagnose:** Run `\d "TableName"` in psql and compare against `schema.prisma` — missing columns confirm the issue.
**Fix options:**
1. **Create a Prisma migration** (permanent) — run `prisma migrate dev --name <description>` in the worktree. The generated file will be picked up by `prisma migrate deploy` on next startup.
2. **Apply manually for local dev** — `psql -d litellm -c "ALTER TABLE ... ADD COLUMN IF NOT EXISTS ..."` after each proxy start. Fine for dev, not for production.
3. **Update litellm-proxy-extras** — if the package is installed from PyPI, its migration directory must include the new file. Either update the package or run the migration manually until the next release ships it.
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify

View file

@ -38,18 +38,25 @@ Before contributing code to LiteLLM, you must sign our [Contributor License Agre
git clone https://github.com/YOUR_USERNAME/litellm.git
cd litellm
# Create a new branch for your feature
git checkout -b your-feature-branch
# Create a new branch for your feature (see "Commit and Branch Conventions" below)
git checkout -b feature/your-feature
# Install development dependencies
make install-dev
# Install git hooks that enforce commit + branch conventions (one-time, opt-in)
make install-hooks
# Verify your setup works
make help
```
That's it! Your local development environment is ready.
## Commit and Branch Conventions
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out.
### 2. Development Workflow
Here's the recommended workflow for making changes:
@ -67,12 +74,12 @@ make lint
# Run unit tests to ensure nothing is broken
make test-unit
# Commit your changes
# Commit your changes (must follow Conventional Commits — see above)
git add .
git commit -m "Your descriptive commit message"
git commit -m "feat(scope): your descriptive commit message"
# Push and create a PR
git push origin your-feature-branch
# Push and create a PR (branch must follow Conventional Branches — see above)
git push origin feature/your-feature
```
## Adding Testing

View file

@ -68,22 +68,24 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
npm install -g npm@11.14.0 tar@7.5.11 glob@13.0.6 @isaacs/brace-expansion@5.0.1 brace-expansion@5.0.5 minimatch@10.2.4 diff@8.0.3 picomatch@4.0.4 && \
GLOBAL="$(npm root -g)" && \
for pkg in tar glob @isaacs/brace-expansion brace-expansion minimatch diff picomatch; do \
name="${pkg##*/}"; \
find "$GLOBAL/npm" -type d -name "$name" -path "*/node_modules/$pkg" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/$pkg" "$d"; \
done; \
done && \
npm cache clean --force && \
{ apk del --no-cache npm 2>/dev/null || true; }
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
# Copy only what runtime needs. The application is installed inside the venv;
# the rest of the builder's /app is source and build metadata that must not
# ship (manifest-scanning tools attribute everything in it to this image).
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
COPY --from=builder /app/.venv /app/.venv
COPY --from=builder /app/docker /app/docker
COPY --from=builder /app/schema.prisma /app/schema.prisma
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
# enterprise/ is imported by source path at runtime (proxy_cli puts the
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
# whole /root/.cache drags in the uv build cache (~660 MB, includes a

109
GEMINI.md
View file

@ -1,108 +1 @@
# GEMINI.md
This file provides guidance to Gemini when working with code in this repository.
## Development Commands
### Installation
- `make install-dev` - Install core development dependencies
- `make install-proxy-dev` - Install proxy development dependencies with full feature set
- `make install-test-deps` - Install all test dependencies
### Testing
- `make test` - Run all tests
- `make test-unit` - Run unit tests (tests/test_litellm) with 4 parallel workers
- `make test-integration` - Run integration tests (excludes unit tests)
- `pytest tests/` - Direct pytest execution
### Code Quality
- `make lint` - Run all linting (Ruff, MyPy, Black, circular imports, import safety)
- `make format` - Apply Black code formatting
- `make lint-ruff` - Run Ruff linting only
- `make lint-mypy` - Run MyPy type checking only
### Single Test Files
- `uv run pytest tests/path/to/test_file.py -v` - Run specific test file
- `uv run pytest tests/path/to/test_file.py::test_function -v` - Run specific test
### Running Scripts
- `uv run python script.py` - Run Python scripts (use for non-test files)
### GitHub Issue & PR Templates
When contributing to the project, use the appropriate templates:
**Bug Reports** (`.github/ISSUE_TEMPLATE/bug_report.yml`):
- Describe what happened vs. what you expected
- Include relevant log output
- Specify your LiteLLM version
**Feature Requests** (`.github/ISSUE_TEMPLATE/feature_request.yml`):
- Describe the feature clearly
- Explain the motivation and use case
**Pull Requests** (`.github/pull_request_template.md`):
- Add at least 1 test in `tests/litellm/`
- Ensure `make test-unit` passes
## Architecture Overview
LiteLLM is a unified interface for 100+ LLM providers with two main components:
### Core Library (`litellm/`)
- **Main entry point**: `litellm/main.py` - Contains core completion() function
- **Provider implementations**: `litellm/llms/` - Each provider has its own subdirectory
- **Router system**: `litellm/router.py` + `litellm/router_utils/` - Load balancing and fallback logic
- **Type definitions**: `litellm/types/` - Pydantic models and type hints
- **Integrations**: `litellm/integrations/` - Third-party observability, caching, logging
- **Caching**: `litellm/caching/` - Multiple cache backends (Redis, in-memory, S3, etc.)
### Proxy Server (`litellm/proxy/`)
- **Main server**: `proxy_server.py` - FastAPI application
- **Authentication**: `auth/` - API key management, JWT, OAuth2
- **Database**: `db/` - Prisma ORM with PostgreSQL/SQLite support
- **Management endpoints**: `management_endpoints/` - Admin APIs for keys, teams, models
- **Pass-through endpoints**: `pass_through_endpoints/` - Provider-specific API forwarding
- **Guardrails**: `guardrails/` - Safety and content filtering hooks
- **UI Dashboard**: Served from `_experimental/out/` (Next.js build)
## Key Patterns
### Provider Implementation
- Providers inherit from base classes in `litellm/llms/base.py`
- Each provider has transformation functions for input/output formatting
- Support both sync and async operations
- Handle streaming responses and function calling
### Error Handling
- Provider-specific exceptions mapped to OpenAI-compatible errors
- Fallback logic handled by Router system
- Comprehensive logging through `litellm/_logging.py`
### Configuration
- YAML config files for proxy server (see `proxy/example_config_yaml/`)
- Environment variables for API keys and settings
- Database schema managed via Prisma (`proxy/schema.prisma`)
## Development Notes
### Code Style
- Uses Black formatter, Ruff linter, MyPy type checker
- Pydantic v2 for data validation
- Async/await patterns throughout
- Type hints required for all public APIs
### Testing Strategy
- Unit tests in `tests/test_litellm/`
- Integration tests for each provider in `tests/llm_translation/`
- Proxy tests in `tests/proxy_unit_tests/`
- Load tests in `tests/load_tests/`
### Database Migrations
- Prisma handles schema migrations
- Migration files auto-generated with `prisma migrate dev`
- Always test migrations against both PostgreSQL and SQLite
### Enterprise Features
- Enterprise-specific code in `enterprise/` directory
- Optional features enabled via environment variables
- Separate licensing and authentication for enterprise features
Read @CLAUDE.md for coding guidelines

View file

@ -5,7 +5,8 @@
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
info lint lint-dev format \
install-dev install-proxy-dev install-test-deps \
lint-strict-budget lint-strict-budget-update \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety
# Default target
@ -17,12 +18,15 @@ help:
@echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)"
@echo " make install-test-deps - Install the full local test environment"
@echo " make install-helm-unittest - Install helm unittest plugin"
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
@echo " make format - Apply Black code formatting"
@echo " make format-check - Check Black code formatting (matches CI)"
@echo " make lint - Run all linting (Ruff, MyPy, Black check, circular imports, import safety)"
@echo " make lint-ruff - Run Ruff linting only"
@echo " make lint-mypy - Run MyPy type checking only"
@echo " make lint-black - Check Black formatting (matches CI)"
@echo " make lint-strict-budget - Gate the codebase total of each strict ruff rule against its ceiling"
@echo " make lint-strict-budget-update - Re-capture per-rule baselines in ruff-strict-budget.json (ratchet)"
@echo " make check-circular-imports - Check for circular imports"
@echo " make check-import-safety - Check import safety"
@echo " make test - Run all tests"
@ -68,6 +72,11 @@ install-test-deps: install-proxy-dev
install-helm-unittest:
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4 || echo "ignore error if plugin exists"
# Install git hooks that enforce Conventional Commits and Conventional Branches.
# Opt-in: not chained into install-dev.
install-hooks:
./scripts/install_git_hooks.sh
# Formatting
format: install-dev
cd litellm && $(UV_RUN) black . && cd ..
@ -116,6 +125,12 @@ lint-mypy: install-dev
lint-black: format-check
lint-strict-budget: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py
lint-strict-budget-update: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py --update
check-circular-imports: install-dev
cd litellm && $(UV_RUN) python ../tests/documentation_tests/test_circular_imports.py && cd ..
@ -123,7 +138,7 @@ check-import-safety: install-dev
@$(UV_RUN) python -c "from litellm import *; print('[from litellm import *] OK! no issues!');" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1)
# Combined linting (matches test-linting.yml workflow)
lint: format-check lint-ruff lint-mypy check-circular-imports check-import-safety
lint: format-check lint-ruff lint-mypy check-circular-imports check-import-safety lint-strict-budget
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed lint-mypy check-circular-imports check-import-safety
@ -146,7 +161,7 @@ test-unit-proxy-core: install-test-deps
$(UV_RUN) pytest tests/test_litellm/proxy/auth tests/test_litellm/proxy/client tests/test_litellm/proxy/db tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine --tb=short -vv -n 4 --durations=20
test-unit-proxy-misc: install-test-deps
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20
test-unit-integrations: install-test-deps
$(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20

View file

@ -37,7 +37,7 @@
</a>
</h4>
<img width="2688" height="1600" alt="Group 7154 (1)" src="https://github.com/user-attachments/assets/c5ee0412-6fb5-4fb6-ab5b-bafae4209ca6" />
<img alt="LiteLLM AI Gateway" src="https://github.com/user-attachments/assets/c5ee0412-6fb5-4fb6-ab5b-bafae4209ca6" />
---
@ -407,7 +407,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature
### Run in Developer Mode
#### Services
1. Setup .env file in root
2. Run dependant services `docker-compose up db prometheus`
2. Run dependent services `docker-compose up db prometheus`
#### Backend
1. (In root) create virtual environment `python -m venv .venv`

View file

@ -12,17 +12,27 @@ USER root
COPY --from=uvbin /uv /uvx /usr/local/bin/
RUN apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile
# nodejs/npm so `prisma generate` uses Wolfi's Node via PRISMA_USE_GLOBAL_NODE
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
# UV_COMPILE_BYTECODE=1 precompiles .pyc at install time → faster cold start.
# UV_LINK_MODE=copy avoids hardlink warnings when uv installs from a
# BuildKit cache mount (different filesystem).
# UV_PYTHON_DOWNLOADS=0 force uv to use the apk-installed CPython instead of
# silently pulling a managed interpreter.
# PRISMA_USE_GLOBAL_NODE explicit (matches default) so an env override can't
# silently re-enable nodeenv's Node download.
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
UV_COMPILE_BYTECODE=1 \
UV_PYTHON_DOWNLOADS=0 \
PRISMA_USE_GLOBAL_NODE=true \
PATH="/app/.venv/bin:${PATH}"
# Stage 1 — install dependencies only.
@ -58,7 +68,11 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata python3 libsndfile libatomic
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
# wolfi-base ships an unprivileged `nonroot` account (UID/GID 65532) with
# /home/nonroot. We run the backend as that user

View file

@ -20,7 +20,11 @@ DatabaseURLSettings.from_env().apply_to_env()
from litellm.proxy.proxy_server import app
from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES
from backend.routes.allowlist import (
BACKEND_EXACT_PATHS,
BACKEND_MOUNT_PATHS,
BACKEND_PATH_PREFIXES,
)
def _is_backend_route(route) -> bool:
@ -29,8 +33,9 @@ def _is_backend_route(route) -> bool:
if path is None:
return False
if isinstance(route, Mount):
# Static UI mounts are served by the dedicated UI container, not here.
return False
# The dashboard UI static mounts are served by the dedicated UI container.
# Only Mounts in the backend allowlist (e.g. swagger docs) remain on backend.
return path in BACKEND_MOUNT_PATHS
if path in BACKEND_EXACT_PATHS:
return True
return any(path.startswith(prefix) for prefix in BACKEND_PATH_PREFIXES)

View file

@ -133,3 +133,9 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
"/fallback/login",
}
)
BACKEND_MOUNT_PATHS: frozenset[str] = frozenset(
{
"/swagger", # API documentation static assets belong to the backend
}
)

View file

@ -1,5 +1,5 @@
module github.com/BerriAI/litellm/cookbook/gollem_go_agent_framework
go 1.25.1
go 1.26.3
require github.com/fugue-labs/gollem v0.1.0

View file

@ -0,0 +1,99 @@
-- Converts an existing LiteLLM_SpendLogs table into a native Postgres
-- range-partitioned table keyed on "startTime".
--
-- Why: at high request volume, retention via DELETE leaves dead tuples that
-- autovacuum cannot reclaim quickly enough, so the table keeps growing on disk
-- (seen at 450GB+ after ~1 month). With partitioning, retention drops whole
-- partitions, which is instant and returns disk to the OS immediately.
--
-- This is an opt-in, manual operation. The default LiteLLM schema is NOT
-- partitioned, so existing installs are unaffected until you run this.
--
-- IMPORTANT
-- * Test on a staging copy first and take a backup.
-- * Postgres cannot convert a populated table to partitioned in place, so this
-- renames the old table aside and creates a fresh partitioned table.
-- * The partition key ("startTime") must be part of the primary key, so the
-- PK becomes composite ("request_id", "startTime"). LiteLLM's write path uses
-- INSERT ... ON CONFLICT DO NOTHING, which is compatible with this.
-- * Choose a partition granularity ("day" is the recommended default for
-- high-volume tables) and keep it consistent with SPEND_LOG_PARTITION_INTERVAL.
--
-- After running this, enable the feature and set a retention period in
-- proxy_config.yaml:
-- general_settings:
-- use_spend_logs_partitioning: true
-- maximum_spend_logs_retention_period: "30d"
-- The spend-log cleanup job then verifies the table is partitioned and reclaims
-- disk by dropping expired partitions instead of deleting rows. It also
-- pre-creates upcoming partitions on each run. To roll back, see
-- db_scripts/unpartition_spend_logs.sql.
BEGIN;
ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_legacy";
-- Renaming a table does NOT rename its indexes, and index names are unique per
-- schema. Move the legacy table's indexes aside so the CREATE INDEX statements
-- below actually create indexes on the new partitioned table instead of being
-- silently skipped by IF NOT EXISTS, and so the new PK keeps the canonical
-- name instead of getting a "_pkey1" suffix.
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey"
RENAME TO "LiteLLM_SpendLogs_legacy_pkey";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_startTime_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_startTime_request_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
) PARTITION BY RANGE ("startTime");
ALTER TABLE "LiteLLM_SpendLogs"
ADD PRIMARY KEY ("request_id", "startTime");
-- Recreate every index Prisma defines on the table. LIKE ... INCLUDING DEFAULTS
-- INCLUDING GENERATED copies columns and defaults but NOT indexes, so without
-- these the admin-UI cost-reporting queries that filter by end_user/session_id
-- fall back to sequential scans. On a partitioned parent these propagate to
-- every current and future partition automatically.
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx"
ON "LiteLLM_SpendLogs" ("startTime");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
ON "LiteLLM_SpendLogs" ("startTime", "request_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
ON "LiteLLM_SpendLogs" ("end_user");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
-- Safety net: any row whose startTime has no explicit partition lands here so
-- writes never fail. The cleanup job never drops the DEFAULT partition.
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault"
PARTITION OF "LiteLLM_SpendLogs" DEFAULT;
COMMIT;
-- Backfill (optional). Rows route to the correct partition automatically.
-- For large legacy tables, copy in time-bounded batches during a low-traffic
-- window instead of one statement, or simply keep "LiteLLM_SpendLogs_legacy"
-- read-only until its data ages past your retention, then DROP it.
--
-- Backfilled rows land in the DEFAULT partition until explicit partitions
-- cover their dates. Postgres refuses to create a partition whose range
-- overlaps rows already in DEFAULT, so the cleanup job may log a warning when
-- pre-creating today's partition right after a backfill; it recovers on its
-- own once those dates age out, and future partitions are unaffected because
-- they are always created ahead of writes.
--
-- INSERT INTO "LiteLLM_SpendLogs"
-- SELECT * FROM "LiteLLM_SpendLogs_legacy"
-- WHERE "startTime" >= now() - interval '30 days';
--
-- DROP TABLE "LiteLLM_SpendLogs_legacy";

View file

@ -0,0 +1,69 @@
-- Rolls back db_scripts/partition_spend_logs.sql: converts the native
-- range-partitioned "LiteLLM_SpendLogs" table back into a plain,
-- non-partitioned table matching the default LiteLLM schema.
--
-- When/why: run this if you want to stop using partition-based retention and
-- return to DELETE-based cleanup, or to restore the original single-column
-- primary key ("request_id") that the partitioned layout had to widen to a
-- composite ("request_id", "startTime").
--
-- IMPORTANT
-- * Test on a staging copy first and take a backup.
-- * Postgres cannot convert a partitioned table back in place, so this
-- renames the partitioned table aside and creates a fresh plain table.
-- * The composite PK could in principle hold the same "request_id" in more
-- than one partition, so rows are copied with ON CONFLICT DO NOTHING to
-- restore the single-column PK without failing on such duplicates.
-- * For large tables the INSERT ... SELECT copies every surviving row and may
-- run long; do it during a low-traffic window.
-- * Also remove use_spend_logs_partitioning from proxy_config.yaml (or set it
-- to false) so the cleanup job returns to DELETE-based retention.
BEGIN;
ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_partitioned";
-- Renaming a table does NOT rename its indexes, and index names are unique per
-- schema. Move the partitioned table's indexes aside so the CREATE INDEX
-- statements below actually create indexes on the new plain table instead of
-- being silently skipped by IF NOT EXISTS, and so the new PK keeps the
-- canonical name.
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey"
RENAME TO "LiteLLM_SpendLogs_partitioned_pkey";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey1"
RENAME TO "LiteLLM_SpendLogs_partitioned_pkey1";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_request_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED
);
ALTER TABLE "LiteLLM_SpendLogs"
ADD PRIMARY KEY ("request_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx"
ON "LiteLLM_SpendLogs" ("startTime");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
ON "LiteLLM_SpendLogs" ("startTime", "request_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
ON "LiteLLM_SpendLogs" ("end_user");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
INSERT INTO "LiteLLM_SpendLogs"
SELECT * FROM "LiteLLM_SpendLogs_partitioned"
ON CONFLICT ("request_id") DO NOTHING;
DROP TABLE "LiteLLM_SpendLogs_partitioned";
COMMIT;

View file

@ -30,7 +30,7 @@ spec:
checksum/config: {{ include (print $.Template.BasePath "/configmap-litellm.yaml") . | sha256sum }}
{{- end }}
{{- with .Values.podAnnotations }}
{{- toYaml . | nindent 8 }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
labels:
{{- include "litellm.labels" . | nindent 8 }}

View file

@ -377,3 +377,28 @@ tests:
content:
name: sidecar-tpl
image: "ghcr.io/berriai/litellm-database:test"
- it: should support tpl in podAnnotations
template: deployment.yaml
set:
image:
repository: ghcr.io/berriai/litellm-database
tag: test
# Mirrors the real-world scenario this feature unblocks:
# user disables the built-in ConfigMap (and its built-in checksum/config
# annotation) and re-implements checksum/config themselves via tpl.
proxyConfigMap:
create: false
podAnnotations:
checksum/config: "{{ .Values.image.tag }}"
example.com/some-key: "{{ .Values.image.repository }}"
example.com/literal: "plain-string-value"
asserts:
- equal:
path: spec.template.metadata.annotations["checksum/config"]
value: "test"
- equal:
path: spec.template.metadata.annotations["example.com/some-key"]
value: "ghcr.io/berriai/litellm-database"
- equal:
path: spec.template.metadata.annotations["example.com/literal"]
value: "plain-string-value"

View file

@ -285,11 +285,31 @@ db:
deployStandalone: true
# Lifecycle hooks for the LiteLLM container
#
# 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.
#
# /health/drain is off by default; enable it with
# general_settings.enable_drain_endpoint: true. The kubelet calls preStop
# hooks without proxy credentials, so when the health port is reachable from
# other pods (the common case) also set
# general_settings.drain_endpoint_token (or the DRAIN_ENDPOINT_TOKEN env
# var) and send the same value on the X-Drain-Token header from the hook.
# Calls missing/wrong the token get a 401 and have no side effect.
# Example:
# lifecycle:
# preStop:
# exec:
# command: ["/bin/sh", "-c", "sleep 10"]
# httpGet:
# path: /health/drain
# port: 4000
# httpHeaders:
# - name: X-Drain-Token
# value: <same value as drain_endpoint_token>
lifecycle: {}
# Settings for Bitnami postgresql chart (if db.deployStandalone is true, ignored

View file

@ -66,36 +66,31 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \
GLOBAL="$(npm root -g)" && \
find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
done && \
find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
done && \
npm cache clean --force && \
{ apk del --no-cache npm 2>/dev/null || true; }
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
# Copy only what runtime needs. The application is installed inside the venv;
# the rest of the builder's /app is source and build metadata that must not
# ship (manifest-scanning tools attribute everything in it to this image).
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
COPY --from=builder /app/.venv /app/.venv
COPY --from=builder /app/docker /app/docker
COPY --from=builder /app/schema.prisma /app/schema.prisma
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
# enterprise/ is imported by source path at runtime (proxy_cli puts the
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy them from the builder so they survive
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
COPY --from=builder /root/.cache /root/.cache
# Only the Prisma subdirs: the whole /root/.cache drags in the uv build cache.
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
find /app/.venv -type d -path "*/tornado/test" -delete

View file

@ -95,7 +95,21 @@ RUN for i in 1 2 3; do \
apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
done
COPY --from=builder /app /app
# Copy only what runtime needs. The application is installed inside the venv;
# the rest of the builder's /app is source and build metadata that must not
# ship (manifest-scanning tools attribute everything in it to this image).
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
# Prisma caches live under /app/.cache here (XDG_CACHE_HOME /
# PRISMA_BINARY_CACHE_DIR) so the runtime prisma generate finds them.
COPY --from=builder /app/.venv /app/.venv
COPY --from=builder /app/docker /app/docker
COPY --from=builder /app/schema.prisma /app/schema.prisma
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
# enterprise/ is imported by source path at runtime (proxy_cli puts the
# working directory on sys.path; litellm/proxy/hooks resolves
# enterprise.enterprise_hooks from it)
COPY --from=builder /app/enterprise /app/enterprise
COPY --from=builder /app/.cache /app/.cache
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets

View file

@ -19,12 +19,26 @@ RESEND_API_ENDPOINT = "https://api.resend.com/emails"
class ResendEmailLogger(BaseEmailLogger):
"""
Send emails using Resend's API.
Required env vars:
- RESEND_API_KEY
Optional env vars:
- RESEND_FROM_EMAIL: Override the default sender address. Must be on a
domain verified in your Resend account. When unset, falls back to the
`from_email` argument passed by the caller (which defaults to
`notifications@alerts.litellm.ai` and only works on LiteLLM Cloud).
"""
def __init__(self, internal_usage_cache=None, **kwargs):
super().__init__(internal_usage_cache=internal_usage_cache, **kwargs)
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.resend_api_key = os.getenv("RESEND_API_KEY")
self.resend_from_email = os.getenv("RESEND_FROM_EMAIL")
async def send_email(
self,
@ -33,13 +47,14 @@ class ResendEmailLogger(BaseEmailLogger):
subject: str,
html_body: str,
):
sender_email = self.resend_from_email or from_email
verbose_logger.debug(
f"Sending email from {from_email} to {to_email} with subject {subject}"
f"Sending email from {sender_email} to {to_email} with subject {subject}"
)
response = await self.async_httpx_client.post(
url=RESEND_API_ENDPOINT,
json={
"from": from_email,
"from": sender_email,
"to": to_email,
"subject": subject,
"html": html_body,

View file

@ -504,7 +504,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if retrieve_file_id
else False
)
if potential_file_id:
if potential_file_id and "llm_output_file_id," in potential_file_id:
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
if model_id:
data["model"] = model_id
@ -658,7 +658,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if isinstance(content, str):
continue
for c in content:
if c["type"] == "file":
if c.get("type") == "file":
file_object = cast(ChatCompletionFileObject, c)
file_object_file_field = file_object["file"]
file_id = file_object_file_field.get("file_id")
@ -1058,7 +1058,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return file_id.split("llm_output_file_model_id,")[1].split(";")[0]
def get_output_file_id_from_unified_file_id(self, file_id: str) -> str:
return file_id.split("llm_output_file_id,")[1].split(";")[0]
marker = "llm_output_file_id,"
if marker not in file_id:
raise ValueError(
f"Unified id does not contain {marker!r}: {file_id[:80]!r}"
)
return file_id.split(marker, 1)[1].split(";")[0]
async def async_post_call_success_hook(
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
@ -1099,13 +1104,33 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
for file_attr in ["output_file_id", "error_file_id"]:
file_id_value = getattr(response, file_attr, None)
if file_id_value and model_id:
original_file_id = file_id_value
unified_file_id = self.get_unified_output_file_id(
output_file_id=original_file_id,
model_id=model_id,
model_name=resolved_model_name,
decoded_output_file_id = _is_base64_encoded_unified_file_id(
file_id_value
)
setattr(response, file_attr, unified_file_id)
if (
decoded_output_file_id
and "llm_output_file_id," in decoded_output_file_id
):
provider_file_id = (
self.get_output_file_id_from_unified_file_id(
decoded_output_file_id
)
)
unified_file_id = file_id_value
elif decoded_output_file_id:
verbose_logger.warning(
f"Skipping {file_attr}={file_id_value!r}: "
"unified id is not a managed file output id"
)
continue
else:
provider_file_id = file_id_value
unified_file_id = self.get_unified_output_file_id(
output_file_id=provider_file_id,
model_id=model_id,
model_name=resolved_model_name,
)
setattr(response, file_attr, unified_file_id)
# Use llm_router credentials when available. Without credentials,
# Azure and other auth-required providers return 500/401.
@ -1125,27 +1150,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
or {}
)
file_object = await litellm.afile_retrieve(
file_id=original_file_id,
file_id=provider_file_id,
**_creds,
)
else:
file_object = await litellm.afile_retrieve(
custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", # type: ignore[arg-type]
file_id=original_file_id,
file_id=provider_file_id,
)
verbose_logger.debug(
f"Successfully retrieved file object for {file_attr}={original_file_id}"
f"Successfully retrieved file object for {file_attr}={provider_file_id}"
)
except Exception as e:
verbose_logger.warning(
f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand."
f"Failed to retrieve file object for {file_attr}={provider_file_id}: {str(e)}. Storing with None and will fetch on-demand."
)
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=file_object,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
model_mappings={model_id: original_file_id},
model_mappings={model_id: provider_file_id},
user_api_key_dict=user_api_key_dict,
)
await self.store_unified_object_id(

View file

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

View file

@ -12,17 +12,27 @@ USER root
COPY --from=uvbin /uv /uvx /usr/local/bin/
RUN apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile
# nodejs/npm so `prisma generate` uses Wolfi's Node via PRISMA_USE_GLOBAL_NODE
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
# UV_COMPILE_BYTECODE=1 precompiles .pyc at install time → faster cold start.
# UV_LINK_MODE=copy avoids hardlink warnings when uv installs from a
# BuildKit cache mount (different filesystem).
# UV_PYTHON_DOWNLOADS=0 force uv to use the apk-installed CPython instead of
# silently pulling a managed interpreter.
# PRISMA_USE_GLOBAL_NODE explicit (matches default) so an env override can't
# silently re-enable nodeenv's Node download.
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
UV_COMPILE_BYTECODE=1 \
UV_PYTHON_DOWNLOADS=0 \
PRISMA_USE_GLOBAL_NODE=true \
PATH="/app/.venv/bin:${PATH}"
# Stage 1 — install dependencies only.
@ -58,7 +68,11 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata python3 libsndfile libatomic
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
# wolfi-base ships an unprivileged `nonroot` account (UID/GID 65532) with
# /home/nonroot. We run the proxy as that user.

View file

@ -106,6 +106,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
# Health & ops
"/health",
"/metrics",
"/watsonx"
)
GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(

View file

@ -56,16 +56,34 @@ app.kubernetes.io/component: ui
{{- end -}}
{{/*
Shared ServiceAccount name used by all three component Deployments. When
`serviceAccount.create` is true and `serviceAccount.name` is empty, default
to the chart fullname. When `create` is false, fall back to the provided
name or the namespace's `default` SA.
Per-component ServiceAccount name helpers.
Each component (gateway, backend, ui) has its own SA config under
.Values.serviceAccounts.<component>. When `create` is true and `name` is
empty the chart defaults to "<release>-litellm-<component>". When `create`
is false the chart uses the provided name, or the namespace `default` SA.
*/}}
{{- define "litellm.serviceAccountName" -}}
{{- if .Values.serviceAccount.create -}}
{{ default (include "litellm.fullname" .) .Values.serviceAccount.name }}
{{- define "litellm.gateway.serviceAccountName" -}}
{{- if .Values.serviceAccounts.gateway.create -}}
{{ default (include "litellm.gateway.fullname" .) .Values.serviceAccounts.gateway.name }}
{{- else -}}
{{ default "default" .Values.serviceAccount.name }}
{{ default "default" .Values.serviceAccounts.gateway.name }}
{{- end -}}
{{- end -}}
{{- define "litellm.backend.serviceAccountName" -}}
{{- if .Values.serviceAccounts.backend.create -}}
{{ default (include "litellm.backend.fullname" .) .Values.serviceAccounts.backend.name }}
{{- else -}}
{{ default "default" .Values.serviceAccounts.backend.name }}
{{- end -}}
{{- end -}}
{{- define "litellm.ui.serviceAccountName" -}}
{{- if .Values.serviceAccounts.ui.create -}}
{{ default (include "litellm.ui.fullname" .) .Values.serviceAccounts.ui.name }}
{{- else -}}
{{ default "default" .Values.serviceAccounts.ui.name }}
{{- end -}}
{{- end -}}

View file

@ -12,14 +12,20 @@ spec:
{{- include "litellm.backend.selectorLabels" . | nindent 6 }}
template:
metadata:
{{- with .Values.backend.podAnnotations }}
{{- if or .Values.gateway.config.create .Values.backend.podAnnotations }}
annotations:
{{- if .Values.gateway.config.create }}
checksum/config: {{ include (print $.Template.BasePath "/gateway/configmap.yaml") . | sha256sum }}
{{- end }}
{{- with .Values.backend.podAnnotations }}
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}
labels:
{{- include "litellm.backend.selectorLabels" . | nindent 8 }}
spec:
serviceAccountName: {{ include "litellm.serviceAccountName" . }}
serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }}
automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
@ -34,7 +40,17 @@ spec:
protocol: TCP
env:
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
{{- if .Values.gateway.config.create }}
- name: CONFIG_FILE_PATH
value: /app/config/config.yaml
{{- end }}
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
{{- if .Values.gateway.config.create }}
volumeMounts:
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- with .Values.backend.livenessProbe }}
livenessProbe:
{{- toYaml . | nindent 12 }}
@ -45,6 +61,12 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.backend.resources | nindent 12 }}
{{- if .Values.gateway.config.create }}
volumes:
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- with .Values.backend.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}

View file

@ -22,7 +22,8 @@ spec:
labels:
{{- include "litellm.gateway.selectorLabels" . | nindent 8 }}
spec:
serviceAccountName: {{ include "litellm.serviceAccountName" . }}
serviceAccountName: {{ include "litellm.gateway.serviceAccountName" . }}
automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}

View file

@ -28,7 +28,7 @@ spec:
app.kubernetes.io/component: migrations
spec:
restartPolicy: Never
serviceAccountName: {{ include "litellm.serviceAccountName" . }}
serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}

View file

@ -1,13 +1,51 @@
{{- if .Values.serviceAccount.create -}}
{{- $prev := false -}}
{{- if .Values.serviceAccounts.gateway.create -}}
{{- $prev = true }}
apiVersion: v1
kind: ServiceAccount
metadata:
name: {{ include "litellm.serviceAccountName" . }}
name: {{ include "litellm.gateway.serviceAccountName" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
{{- with .Values.serviceAccount.annotations }}
app.kubernetes.io/component: gateway
{{- with .Values.serviceAccounts.gateway.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
automountServiceAccountToken: {{ .Values.serviceAccount.automount }}
automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }}
{{- end }}
{{- if .Values.serviceAccounts.backend.create }}
{{- if $prev }}
---
{{- end }}
{{- $prev = true }}
apiVersion: v1
kind: ServiceAccount
metadata:
name: {{ include "litellm.backend.serviceAccountName" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: backend
{{- with .Values.serviceAccounts.backend.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }}
{{- end }}
{{- if .Values.serviceAccounts.ui.create }}
{{- if $prev }}
---
{{- end }}
apiVersion: v1
kind: ServiceAccount
metadata:
name: {{ include "litellm.ui.serviceAccountName" . }}
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: ui
{{- with .Values.serviceAccounts.ui.annotations }}
annotations:
{{- toYaml . | nindent 4 }}
{{- end }}
automountServiceAccountToken: {{ .Values.serviceAccounts.ui.automount }}
{{- end }}

View file

@ -19,7 +19,8 @@ spec:
labels:
{{- include "litellm.ui.selectorLabels" . | nindent 8 }}
spec:
serviceAccountName: {{ include "litellm.serviceAccountName" . }}
serviceAccountName: {{ include "litellm.ui.serviceAccountName" . }}
automountServiceAccountToken: {{ .Values.serviceAccounts.ui.automount }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}

View file

@ -14,16 +14,33 @@ ingress:
host: "" # optional; if set, becomes the rule's host
tls: []
# Shared ServiceAccount used by all three component Deployments. Set
# `create: true` to have the chart provision it (e.g. when wiring an EKS
# Pod Identity association by SA name). Set `name` to use an existing SA
# (chart-created or out-of-band). When both are empty / false, pods run
# with the namespace's `default` SA.
serviceAccount:
create: false
automount: true
annotations: {}
name: ""
# Per-component ServiceAccounts for gateway, backend, and ui.
#
# Each section mirrors the old shared serviceAccount shape. Set `create:
# true` to have the chart provision the SA (useful for EKS Pod Identity /
# GKE Workload Identity annotations). Set `name` to bind an existing SA.
# When both are unset the component pod runs with the namespace `default` SA.
#
# The UI SA deliberately defaults to `automount: false` — the static nginx
# container does not need the K8s API and should not carry a projected
# ServiceAccount token that a compromised container could use to call the
# cloud-provider metadata service or the K8s API.
serviceAccounts:
gateway:
create: false
automount: true
annotations: {}
name: ""
backend:
create: false
automount: true
annotations: {}
name: ""
ui:
create: false
automount: false
annotations: {}
name: ""
# Pre-install / pre-upgrade Helm hook that runs `prisma migrate deploy`
# against the writer database, creating the LiteLLM schema (tables that

View file

@ -0,0 +1,23 @@
-- AlterTable: add admin-configured env_vars to MCP server table
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "env_vars" JSONB DEFAULT '[]';
-- CreateTable: per-user env var values for MCP servers
CREATE TABLE IF NOT EXISTS "LiteLLM_MCPUserEnvVars" (
"id" TEXT NOT NULL,
"user_id" TEXT NOT NULL,
"server_id" TEXT NOT NULL,
"values_b64" TEXT NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_MCPUserEnvVars_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_user_id_server_id_key" ON "LiteLLM_MCPUserEnvVars"("user_id", "server_id");
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_user_id_idx" ON "LiteLLM_MCPUserEnvVars"("user_id");
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_server_id_idx" ON "LiteLLM_MCPUserEnvVars"("server_id");

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "oauth_passthrough" BOOLEAN NOT NULL DEFAULT false;

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "oauth2_flow" TEXT;

View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "timeout" DOUBLE PRECISION;

View file

@ -311,6 +311,11 @@ model LiteLLM_MCPServerTable {
tool_name_to_description Json? @default("{}")
extra_headers String[] @default([])
static_headers Json? @default("{}")
// Admin-configured environment variables interpolated into static_headers
// via ${NAME} syntax. Stored as an array of
// {name, value, scope, description}. scope is "global" (value used as-is)
// or "user" (value supplied per-user via LiteLLM_MCPUserEnvVars).
env_vars Json? @default("[]")
// Health check status
status String? @default("unknown")
last_health_check DateTime?
@ -322,13 +327,16 @@ model LiteLLM_MCPServerTable {
authorization_url String?
token_url String?
registration_url String?
oauth2_flow String?
allow_all_keys Boolean @default(false)
available_on_public_internet Boolean @default(true)
delegate_auth_to_upstream Boolean @default(false)
oauth_passthrough Boolean @default(false)
is_byok Boolean @default(false)
byok_description String[] @default([])
byok_api_key_help_url String?
source_url String?
timeout Float?
// BYOM submission lifecycle
approval_status String? @default("active")
submitted_by String?
@ -363,6 +371,21 @@ model LiteLLM_MCPUserCredentials {
@@unique([user_id, server_id])
}
// Per-user environment variable values for MCP servers.
// values_b64 is an encrypted JSON object: {VAR_NAME: "value", ...}.
model LiteLLM_MCPUserEnvVars {
id String @id @default(uuid())
user_id String
server_id String
values_b64 String
created_at DateTime @default(now())
updated_at DateTime @default(now()) @updatedAt
@@unique([user_id, server_id])
@@index([user_id])
@@index([server_id])
}
// Generate Tokens for Proxy
model LiteLLM_VerificationToken {
token String @id

View file

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

View file

@ -16,8 +16,17 @@ import os
# Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available
import dotenv as _dotenv
def _dev_env_hot_reload_enabled() -> bool:
"""The proxy exports this flag when started with ``--reload``. A reloaded
worker is a fresh process that inherits the reloader's environment, so an
edited ``.env`` value stays masked by the stale inherited one unless we
let the file win; overriding makes the edit take effect on reload."""
return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True"
if os.getenv("LITELLM_MODE", "DEV") == "DEV":
_dotenv.load_dotenv()
_dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
from typing import (
Callable,
@ -34,6 +43,7 @@ from typing import (
Type,
)
from litellm.types.integrations.datadog import DatadogInitParams
from litellm.types.integrations.newrelic import NewRelicInitParams
from litellm._logging import (
set_verbose,
_turn_on_debug,
@ -145,10 +155,12 @@ _custom_logger_compatible_callbacks_literal = Literal[
"gitlab",
"cloudzero",
"focus",
"mavvrik",
"vantage",
"posthog",
"levo",
"compression_interception",
"newrelic",
]
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
@ -225,6 +237,11 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
route_all_chat_openai_to_responses: bool = (
os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true"
) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge
# When True, Gemini/Vertex Live setup is deferred until client `session.update`.
# Default False preserves historical behavior (auto-send setup on connect).
gemini_live_defer_setup: bool = (
os.getenv("LITELLM_GEMINI_LIVE_DEFER_SETUP", "false").lower() == "true"
)
use_legacy_interactions_schema: bool = (
os.getenv("LITELLM_USE_LEGACY_INTERACTIONS_SCHEMA", "false").lower() == "true"
) # When True, sends Api-Revision: 2026-05-07 to Google so responses use the legacy `outputs`
@ -235,6 +252,7 @@ api_key: Optional[str] = None
openai_key: Optional[str] = None
groq_key: Optional[str] = None
gigachat_key: Optional[str] = None
xai_key: Optional[str] = None
databricks_key: Optional[str] = None
openai_like_key: Optional[str] = None
azure_key: Optional[str] = None
@ -272,6 +290,7 @@ ovhcloud_key: Optional[str] = None
lemonade_key: Optional[str] = None
sap_service_key: Optional[str] = None
amazon_nova_api_key: Optional[str] = None
inception_key: Optional[str] = None
common_cloud_provider_auth_params: dict = {
"params": ["project", "region_name", "token"],
"providers": ["vertex_ai", "bedrock", "watsonx", "azure", "vertex_ai_beta"],
@ -343,6 +362,9 @@ enable_gemini_default_thinking_level_low: bool = (
####################
logging: bool = True
enable_loadbalancing_on_batch_endpoints: Optional[bool] = None
require_managed_files: bool = (
False # proxy only - require target_model_names on POST /v1/files
)
enable_caching_on_provider_specific_optional_params: bool = (
False # feature-flag for caching on optional params - e.g. 'top_k'
)
@ -396,6 +418,7 @@ s3_callback_params: Optional[Dict] = None
s3_audit_callback_params: Optional[Dict] = None
datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] = None
datadog_params: Optional[Union[DatadogInitParams, Dict]] = None
newrelic_params: Optional[Union[NewRelicInitParams, Dict]] = None
aws_sqs_callback_params: Optional[Dict] = None
generic_logger_headers: Optional[Dict] = None
default_key_generate_params: Optional[Dict] = None
@ -426,6 +449,13 @@ custom_prometheus_metadata_labels: List[str] = []
custom_prometheus_tags: List[str] = []
prometheus_metrics_config: Optional[List] = None
prometheus_emit_stream_label: bool = False
# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on
# `litellm_proxy_failed_requests_metric`. Off by default to preserve the
# pre-unification label set so existing dashboards / recording rules keyed on
# that metric keep matching after upgrade. Enable when downstream consumers
# are ready to split 429s by source (vendor vs. litellm) and dimension
# (RPM/TPM/concurrent/budget).
prometheus_emit_rate_limit_labels: bool = False
prometheus_user_budget_label_include_email_alias: bool = False
prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000
prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0
@ -437,6 +467,7 @@ disable_copilot_system_to_assistant: bool = (
False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior.
)
public_mcp_servers: Optional[List[str]] = None
public_mcp_hub_strict_whitelist: bool = True
public_model_groups: Optional[List[str]] = None
public_agent_groups: Optional[List[str]] = None
# Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
@ -545,6 +576,7 @@ cohere_models: Set = set()
cohere_chat_models: Set = set()
mistral_chat_models: Set = set()
text_completion_codestral_models: Set = set()
text_completion_inception_models: Set = set()
anthropic_models: Set = set()
openrouter_models: Set = set()
datarobot_models: Set = set()
@ -603,6 +635,7 @@ cerebras_models: Set = set()
galadriel_models: Set = set()
nvidia_nim_models: Set = set()
nvidia_riva_models: Set = set()
soniox_models: Set = set()
sambanova_models: Set = set()
sambanova_embedding_models: Set = set()
novita_models: Set = set()
@ -622,6 +655,7 @@ publicai_models: Set = set()
v0_models: Set = set()
morph_models: Set = set()
lambda_ai_models: Set = set()
inception_models: Set = set()
hyperbolic_models: Set = set()
black_forest_labs_models: Set = set()
recraft_models: Set = set()
@ -786,6 +820,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
fireworks_ai_embedding_models.add(key)
elif value.get("litellm_provider") == "text-completion-codestral":
text_completion_codestral_models.add(key)
elif value.get("litellm_provider") == "text-completion-inception":
text_completion_inception_models.add(key)
elif value.get("litellm_provider") == "xai":
xai_models.add(key)
elif value.get("litellm_provider") == "zai":
@ -832,6 +868,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
nvidia_nim_models.add(key)
elif value.get("litellm_provider") == "nvidia_riva":
nvidia_riva_models.add(key)
elif value.get("litellm_provider") == "soniox":
soniox_models.add(key)
elif value.get("litellm_provider") == "sambanova":
sambanova_models.add(key)
elif value.get("litellm_provider") == "sambanova-embedding-models":
@ -872,6 +910,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
morph_models.add(key)
elif value.get("litellm_provider") == "lambda_ai":
lambda_ai_models.add(key)
elif value.get("litellm_provider") == "inception":
inception_models.add(key)
elif value.get("litellm_provider") == "hyperbolic":
hyperbolic_models.add(key)
elif value.get("litellm_provider") == "black_forest_labs":
@ -974,6 +1014,7 @@ model_list = list(
| watsonx_models
| gemini_models
| text_completion_codestral_models
| text_completion_inception_models
| xai_models
| zai_models
| fal_ai_models
@ -994,6 +1035,7 @@ model_list = list(
| galadriel_models
| nvidia_nim_models
| nvidia_riva_models
| soniox_models
| sambanova_models
| azure_text_models
| novita_models
@ -1012,6 +1054,7 @@ model_list = list(
| v0_models
| morph_models
| lambda_ai_models
| inception_models
| black_forest_labs_models
| recraft_models
| cometapi_models
@ -1068,6 +1111,7 @@ models_by_provider: dict = {
"fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models,
"aleph_alpha": aleph_alpha_models,
"text-completion-codestral": text_completion_codestral_models,
"text-completion-inception": text_completion_inception_models,
"xai": xai_models,
"zai": zai_models,
"fal_ai": fal_ai_models,
@ -1092,6 +1136,7 @@ models_by_provider: dict = {
"galadriel": galadriel_models,
"nvidia_nim": nvidia_nim_models,
"nvidia_riva": nvidia_riva_models,
"soniox": soniox_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"novita": novita_models,
"nebius": nebius_models | nebius_embedding_models,
@ -1112,6 +1157,7 @@ models_by_provider: dict = {
"v0": v0_models,
"morph": morph_models,
"lambda_ai": lambda_ai_models,
"inception": inception_models,
"hyperbolic": hyperbolic_models,
"black_forest_labs": black_forest_labs_models,
"recraft": recraft_models,
@ -1271,6 +1317,8 @@ from .exceptions import (
NotFoundError,
PermissionDeniedError,
RateLimitError,
RateLimitErrorCategory,
RateLimitType,
ServiceUnavailableError,
BadGatewayError,
OpenAIError,
@ -1332,6 +1380,7 @@ from .search.main import *
from .realtime_api.main import (
_arealtime,
acreate_realtime_client_secret,
acreate_realtime_transcription_session,
arealtime_calls,
)
from .responses.main import _aresponses_websocket
@ -1682,6 +1731,9 @@ if TYPE_CHECKING:
from .llms.voyage.embedding.transformation_contextual import (
VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig,
)
from .llms.voyage.embedding.transformation_multimodal import (
VoyageMultimodalEmbeddingConfig as VoyageMultimodalEmbeddingConfig,
)
from .llms.infinity.embedding.transformation import (
InfinityEmbeddingConfig as InfinityEmbeddingConfig,
)
@ -1722,6 +1774,9 @@ if TYPE_CHECKING:
from .llms.openrouter.responses.transformation import (
OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig,
)
from .llms.bedrock_mantle.responses.transformation import (
BedrockMantleResponsesAPIConfig as BedrockMantleResponsesAPIConfig,
)
from .llms.gemini.interactions.transformation import (
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
)
@ -1863,6 +1918,9 @@ if TYPE_CHECKING:
from .llms.codestral.completion.transformation import (
CodestralTextCompletionConfig as CodestralTextCompletionConfig,
)
from .llms.inception.completion.transformation import (
InceptionTextCompletionConfig as InceptionTextCompletionConfig,
)
from .llms.azure.azure import (
AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig,
)
@ -1931,6 +1989,9 @@ if TYPE_CHECKING:
from .llms.lambda_ai.chat.transformation import (
LambdaAIChatConfig as LambdaAIChatConfig,
)
from .llms.inception.chat.transformation import (
InceptionChatConfig as InceptionChatConfig,
)
from .llms.hyperbolic.chat.transformation import (
HyperbolicChatConfig as HyperbolicChatConfig,
)

View file

@ -223,6 +223,7 @@ LLM_CONFIG_NAMES = (
"GenAIHubOrchestrationConfig",
"VoyageEmbeddingConfig",
"VoyageContextualEmbeddingConfig",
"VoyageMultimodalEmbeddingConfig",
"InfinityEmbeddingConfig",
"PerplexityEmbeddingConfig",
"AzureAIStudioConfig",
@ -237,6 +238,7 @@ LLM_CONFIG_NAMES = (
"PerplexityResponsesConfig",
"DatabricksResponsesAPIConfig",
"OpenRouterResponsesAPIConfig",
"BedrockMantleResponsesAPIConfig",
"GoogleAIStudioInteractionsConfig",
"OpenAIOSeriesConfig",
"AnthropicSkillsConfig",
@ -267,6 +269,7 @@ LLM_CONFIG_NAMES = (
"AIMLChatConfig",
"VolcEngineChatConfig",
"CodestralTextCompletionConfig",
"InceptionTextCompletionConfig",
"AzureOpenAIAssistantsAPIConfig",
"HerokuChatConfig",
"CometAPIConfig",
@ -310,6 +313,7 @@ LLM_CONFIG_NAMES = (
"MorphChatConfig",
"RAGFlowConfig",
"LambdaAIChatConfig",
"InceptionChatConfig",
"HyperbolicChatConfig",
"VercelAIGatewayConfig",
"OVHCloudChatConfig",
@ -318,6 +322,7 @@ LLM_CONFIG_NAMES = (
"LemonadeChatConfig",
"SnowflakeEmbeddingConfig",
"AmazonNovaChatConfig",
"SonioxAudioTranscriptionConfig",
)
# Types that support lazy loading via _lazy_import_types
@ -899,6 +904,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.voyage.embedding.transformation_contextual",
"VoyageContextualEmbeddingConfig",
),
"VoyageMultimodalEmbeddingConfig": (
".llms.voyage.embedding.transformation_multimodal",
"VoyageMultimodalEmbeddingConfig",
),
"InfinityEmbeddingConfig": (
".llms.infinity.embedding.transformation",
"InfinityEmbeddingConfig",
@ -956,6 +965,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.openrouter.responses.transformation",
"OpenRouterResponsesAPIConfig",
),
"BedrockMantleResponsesAPIConfig": (
".llms.bedrock_mantle.responses.transformation",
"BedrockMantleResponsesAPIConfig",
),
"GoogleAIStudioInteractionsConfig": (
".llms.gemini.interactions.transformation",
"GoogleAIStudioInteractionsConfig",
@ -1040,6 +1053,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.codestral.completion.transformation",
"CodestralTextCompletionConfig",
),
"InceptionTextCompletionConfig": (
".llms.inception.completion.transformation",
"InceptionTextCompletionConfig",
),
"AzureOpenAIAssistantsAPIConfig": (
".llms.azure.azure",
"AzureOpenAIAssistantsAPIConfig",
@ -1154,6 +1171,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
"MorphChatConfig": (".llms.morph.chat.transformation", "MorphChatConfig"),
"RAGFlowConfig": (".llms.ragflow.chat.transformation", "RAGFlowConfig"),
"LambdaAIChatConfig": (".llms.lambda_ai.chat.transformation", "LambdaAIChatConfig"),
"InceptionChatConfig": (
".llms.inception.chat.transformation",
"InceptionChatConfig",
),
"HyperbolicChatConfig": (
".llms.hyperbolic.chat.transformation",
"HyperbolicChatConfig",
@ -1180,6 +1201,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
".llms.amazon_nova.chat.transformation",
"AmazonNovaChatConfig",
),
"SonioxAudioTranscriptionConfig": (
".llms.soniox.audio_transcription.transformation",
"SonioxAudioTranscriptionConfig",
),
}
# Import map for utils module lazy imports

View file

@ -24,6 +24,22 @@ else:
UserAPIKeyAuth = Any
def _get_otel_v2_class() -> Optional[type]:
"""Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent.
Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry
SDK at module scope, so importing it eagerly would break installs without the
SDK. The V2 logger only exists when ``LITELLM_OTEL_V2`` is enabled (which
requires the SDK), so a failed import simply means "no V2 logger in play".
"""
try:
from litellm.integrations.otel.logger import OpenTelemetryV2
return OpenTelemetryV2
except Exception:
return None
class ServiceLogging(CustomLogger):
"""
Separate class used for monitoring health of litellm-adjacent services (redis/postgres).
@ -38,6 +54,37 @@ class ServiceLogging(CustomLogger):
if "prometheus_system" in litellm.service_callback:
self.prometheusServicesLogger = PrometheusServicesLogger()
def _resolve_otel_service_logger(self, callback: Any) -> Optional[Any]:
"""Resolve the OTel logger (legacy or V2) to emit a service span on.
Returns the logger instance whose ``async_service_*_hook`` should fire for
this ``callback``, or ``None`` when ``callback`` is not an OTel callback.
The V2 ``OpenTelemetryV2`` logger is a plain ``CustomLogger`` and is NOT a
subclass of the legacy ``OpenTelemetry``, so the legacy ``isinstance``
check alone misses it — which is why redis/postgres service spans never
showed up under ``LITELLM_OTEL_V2``. Match both the legacy and V2 types,
whether the callback is the logger instance itself or the ``"otel"`` string
(which routes to the proxy's registered ``open_telemetry_logger``).
"""
otel_v2_cls = _get_otel_v2_class()
def _is_otel_logger(obj: Any) -> bool:
if isinstance(obj, OpenTelemetry):
return True
return otel_v2_cls is not None and isinstance(obj, otel_v2_cls)
if _is_otel_logger(callback):
return callback
if callback == "otel":
from litellm.proxy.proxy_server import open_telemetry_logger
if open_telemetry_logger is not None and _is_otel_logger(
open_telemetry_logger
):
return open_telemetry_logger
return None
def service_success_hook(
self,
service: ServiceTypes,
@ -129,6 +176,13 @@ class ServiceLogging(CustomLogger):
event_metadata=event_metadata,
)
# OTel loggers already fired this event. ``service_callback`` can hold more
# than one reference that resolves to the *same* logger — the ``"otel"``
# string AND the registered instance both map to ``open_telemetry_logger``
# (the V2 logger self-registers its instance even when the string is
# present, unlike V1). Without this guard each such reference emits its own
# span, so a single DB call shows up as duplicate ``postgres ...`` spans.
emitted_otel_logger_ids: set = set()
for callback in litellm.service_callback:
if callback == "prometheus_system":
await self.init_prometheus_services_logger_if_none()
@ -144,19 +198,18 @@ class ServiceLogging(CustomLogger):
end_time=end_time,
event_metadata=event_metadata,
)
elif callback == "otel" or isinstance(callback, OpenTelemetry):
_otel_logger_to_use: Optional[OpenTelemetry] = None
if isinstance(callback, OpenTelemetry):
_otel_logger_to_use = callback
else:
from litellm.proxy.proxy_server import open_telemetry_logger
if open_telemetry_logger is not None and isinstance(
open_telemetry_logger, OpenTelemetry
):
_otel_logger_to_use = open_telemetry_logger
if _otel_logger_to_use is not None and parent_otel_span is not None:
else:
_otel_logger_to_use = self._resolve_otel_service_logger(callback)
# No ``parent_otel_span is not None`` gate: a background service
# call (no request on the stack) has no parent, and dropping it
# here is what hid those calls from traces entirely. The OTel
# logger decides what to do with a missing parent — legacy V1
# no-ops, V2 emits a root span (and skips metrics-only pings).
if (
_otel_logger_to_use is not None
and id(_otel_logger_to_use) not in emitted_otel_logger_ids
):
emitted_otel_logger_ids.add(id(_otel_logger_to_use))
await _otel_logger_to_use.async_service_success_hook(
payload=payload,
parent_otel_span=parent_otel_span,
@ -238,6 +291,9 @@ class ServiceLogging(CustomLogger):
event_metadata=event_metadata,
)
# Dedupe OTel loggers per event — see ``async_service_success_hook`` for why
# the same logger can be referenced twice in ``service_callback``.
emitted_otel_logger_ids: set = set()
for callback in litellm.service_callback:
if callback == "prometheus_system":
await self.init_prometheus_services_logger_if_none()
@ -255,22 +311,19 @@ class ServiceLogging(CustomLogger):
end_time=end_time,
event_metadata=event_metadata,
)
elif callback == "otel" or isinstance(callback, OpenTelemetry):
_otel_logger_to_use: Optional[OpenTelemetry] = None
if isinstance(callback, OpenTelemetry):
_otel_logger_to_use = callback
else:
from litellm.proxy.proxy_server import open_telemetry_logger
if open_telemetry_logger is not None and isinstance(
open_telemetry_logger, OpenTelemetry
):
_otel_logger_to_use = open_telemetry_logger
else:
_otel_logger_to_use = self._resolve_otel_service_logger(callback)
if not isinstance(error, str):
error = str(error)
if _otel_logger_to_use is not None and parent_otel_span is not None:
# See the success hook: no parent gate, so background failures
# are traced too. V1 no-ops without a parent; V2 emits a root.
if (
_otel_logger_to_use is not None
and id(_otel_logger_to_use) not in emitted_otel_logger_ids
):
emitted_otel_logger_ids.add(id(_otel_logger_to_use))
await _otel_logger_to_use.async_service_failure_hook(
payload=payload,
error=error,
@ -318,6 +371,8 @@ class ServiceLogging(CustomLogger):
service=ServiceTypes.LITELLM,
duration=_duration,
call_type=kwargs.get("call_type", "unknown"),
start_time=start_time,
end_time=end_time,
)
except Exception as e:
raise e

View file

@ -19,10 +19,22 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
A2AStreamingContext,
)
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
from litellm.interactions.agents.utils import merge_agent_headers
# litellm_params key carrying the authenticated principal (hashed virtual key) so
# A2A provider configs can scope provider-side state (e.g. LangFlow session memory)
# per key instead of trusting the client-supplied A2A contextId.
A2A_USER_API_KEY_HASH_PARAM = "litellm_a2a_user_api_key_hash"
# Agent metadata fields stored in litellm_params that are not valid litellm.acompletion() kwargs
_AGENT_ONLY_PARAMS = frozenset(
{"is_public", "agent_name", "agent_id", "agent_card_params"}
{
"is_public",
"agent_name",
"agent_id",
"agent_card_params",
A2A_USER_API_KEY_HASH_PARAM,
}
)
@ -37,6 +49,9 @@ class A2ACompletionBridgeHandler:
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> Dict[str, Any]:
"""
Handle non-streaming A2A request via litellm.acompletion.
@ -46,29 +61,31 @@ class A2ACompletionBridgeHandler:
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
api_base: API base URL from agent_card_params
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Returns:
A2A SendMessageResponse dict
"""
# Get provider config for custom_llm_provider
custom_llm_provider = litellm_params.get("custom_llm_provider")
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
# If provider config exists, use it
if a2a_provider_config is not None:
verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}")
response_data = await a2a_provider_config.handle_non_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
if not _skip_a2a_provider_routing:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
return response_data
if a2a_provider_config is not None:
verbose_logger.info(
f"A2A: Using provider config for {custom_llm_provider}"
)
return await a2a_provider_config.handle_non_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
agent_extra_headers=agent_extra_headers,
)
# Extract message from params
message = params.get("message", {})
@ -94,7 +111,7 @@ class A2ACompletionBridgeHandler:
)
# Build completion params dict
completion_params = {
completion_params: Dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
@ -107,6 +124,20 @@ class A2ACompletionBridgeHandler:
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so the helper
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
# those keys authoritative over the client-supplied A2A metadata.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
if agent_extra_headers:
completion_params["extra_headers"] = merge_agent_headers(
dynamic_headers=agent_extra_headers,
static_headers=completion_params.get("extra_headers"),
)
# Call litellm.acompletion
response = await litellm.acompletion(**completion_params)
@ -129,6 +160,9 @@ class A2ACompletionBridgeHandler:
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> AsyncIterator[Dict[str, Any]]:
"""
Handle streaming A2A request via litellm.acompletion with stream=True.
@ -144,32 +178,34 @@ class A2ACompletionBridgeHandler:
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
api_base: API base URL from agent_card_params
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Yields:
A2A streaming response events
"""
# Get provider config for custom_llm_provider
custom_llm_provider = litellm_params.get("custom_llm_provider")
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
# If provider config exists, use it
if a2a_provider_config is not None:
verbose_logger.info(
f"A2A: Using provider config for {custom_llm_provider} (streaming)"
if not _skip_a2a_provider_routing:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
async for chunk in a2a_provider_config.handle_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
):
yield chunk
if a2a_provider_config is not None:
verbose_logger.info(
f"A2A: Using provider config for {custom_llm_provider} (streaming)"
)
return
async for chunk in a2a_provider_config.handle_streaming(
request_id=request_id,
params=params,
api_base=api_base,
litellm_params=litellm_params,
agent_extra_headers=agent_extra_headers,
):
yield chunk
return
# Extract message from params
message = params.get("message", {})
@ -201,7 +237,7 @@ class A2ACompletionBridgeHandler:
)
# Build completion params dict
completion_params = {
completion_params: Dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
@ -214,6 +250,20 @@ class A2ACompletionBridgeHandler:
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
}
completion_params.update(litellm_params_to_add)
# Apply forward metadata AFTER the litellm_params merge so the helper
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
# those keys authoritative over the client-supplied A2A metadata.
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
completion_params=completion_params,
a2a_message=message,
params=params,
)
if agent_extra_headers:
completion_params["extra_headers"] = merge_agent_headers(
dynamic_headers=agent_extra_headers,
static_headers=completion_params.get("extra_headers"),
)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)
@ -276,6 +326,7 @@ async def handle_a2a_completion(
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""Convenience function for non-streaming A2A completion."""
return await A2ACompletionBridgeHandler.handle_non_streaming(
@ -283,6 +334,7 @@ async def handle_a2a_completion(
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
)
@ -291,6 +343,7 @@ async def handle_a2a_completion_streaming(
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
"""Convenience function for streaming A2A completion."""
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
@ -298,5 +351,6 @@ async def handle_a2a_completion_streaming(
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
):
yield chunk

View file

@ -45,10 +45,80 @@ class A2ACompletionBridgeTransformation:
Static methods for transforming between A2A and OpenAI message formats.
"""
@staticmethod
def _extract_text_from_a2a_parts(parts: List[Dict[str, Any]]) -> str:
"""Extract text from A2A parts (with or without explicit ``kind``)."""
content_parts: List[str] = []
for part in parts:
if not isinstance(part, dict):
continue
kind = part.get("kind")
text = part.get("text")
if text is None:
continue
if kind in (None, "", "text"):
content_parts.append(str(text))
return "\n".join(content_parts)
@staticmethod
def get_forward_metadata(
a2a_message: Dict[str, Any],
params: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, Any]]:
"""
Merge A2A metadata from MessageSendParams and the message for downstream providers.
Forwarded once on the LangGraph run payload (``metadata``), not duplicated on
each input message — see ``apply_forward_metadata_to_completion_params``.
"""
merged: Dict[str, Any] = {}
if params and isinstance(params.get("metadata"), dict):
merged.update(params["metadata"])
message_metadata = a2a_message.get("metadata")
if isinstance(message_metadata, dict):
merged.update(message_metadata)
return merged or None
@staticmethod
def apply_forward_metadata_to_completion_params(
completion_params: Dict[str, Any],
a2a_message: Dict[str, Any],
params: Optional[Dict[str, Any]] = None,
) -> None:
"""
Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph).
Uses ``extra_body`` so we do not collide with LiteLLM's spend-log ``metadata`` kwarg.
"""
forward_metadata = A2ACompletionBridgeTransformation.get_forward_metadata(
a2a_message=a2a_message,
params=params,
)
if not forward_metadata:
return
extra_body = completion_params.get("extra_body")
if not isinstance(extra_body, dict):
extra_body = {}
# Layer client-supplied A2A metadata under any agent-owner-configured
# ``extra_body.metadata`` so the configured keys remain authoritative
# and an A2A caller cannot overwrite server-set run metadata.
existing_metadata = extra_body.get("metadata")
existing_dict: Dict[str, Any] = (
existing_metadata if isinstance(existing_metadata, dict) else {}
)
merged_metadata: Dict[str, Any] = {**forward_metadata, **existing_dict}
extra_body = {**extra_body, "metadata": merged_metadata}
completion_params["extra_body"] = extra_body
verbose_logger.debug(
f"A2A -> completion forward metadata keys={list(forward_metadata.keys())}"
)
@staticmethod
def a2a_message_to_openai_messages(
a2a_message: Dict[str, Any],
) -> List[Dict[str, str]]:
) -> List[Dict[str, Any]]:
"""
Transform an A2A message to OpenAI message format.
@ -70,21 +140,20 @@ class A2ACompletionBridgeTransformation:
elif role == "system":
openai_role = "system"
# Extract text content from parts
content_parts = []
for part in parts:
kind = part.get("kind", "")
if kind == "text":
text = part.get("text", "")
content_parts.append(text)
if not isinstance(parts, list):
parts = []
content = "\n".join(content_parts) if content_parts else ""
content = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts)
# Do not attach A2A message.metadata here — the completion bridge forwards it
# once at run level via extra_body.metadata (LangGraph POST /runs/wait shape).
openai_message: Dict[str, Any] = {"role": openai_role, "content": content}
verbose_logger.debug(
f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}"
)
return [{"role": openai_role, "content": content}]
return [openai_message]
@staticmethod
def openai_response_to_a2a_response(
@ -110,6 +179,7 @@ class A2ACompletionBridgeTransformation:
# Build A2A message
a2a_message = {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
@ -119,9 +189,7 @@ class A2ACompletionBridgeTransformation:
a2a_response = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"message": a2a_message,
},
"result": a2a_message,
}
verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}")
@ -235,50 +303,3 @@ class A2ACompletionBridgeTransformation:
"taskId": ctx.task_id,
},
}
@staticmethod
def openai_chunk_to_a2a_chunk(
chunk: Any,
request_id: Optional[str] = None,
is_final: bool = False,
) -> Optional[Dict[str, Any]]:
"""
Transform a LiteLLM streaming chunk to A2A streaming format.
NOTE: This method is deprecated for streaming. Use the event-based
methods (create_task_event, create_status_update_event,
create_artifact_update_event) instead for proper A2A streaming.
Args:
chunk: LiteLLM ModelResponse chunk
request_id: Original A2A request ID
is_final: Whether this is the final chunk
Returns:
A2A streaming chunk dict or None if no content
"""
# Extract delta content
content = ""
if chunk is not None and hasattr(chunk, "choices") and chunk.choices:
choice = chunk.choices[0]
if hasattr(choice, "delta") and choice.delta:
content = choice.delta.content or ""
if not content and not is_final:
return None
# Build A2A streaming chunk (legacy format)
a2a_chunk = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"message": {
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
},
"final": is_final,
},
}
return a2a_chunk

View file

@ -132,6 +132,7 @@ async def _send_message_via_completion_bridge(
custom_llm_provider: str,
api_base: Optional[str],
litellm_params: Dict[str, Any],
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> LiteLLMSendMessageResponse:
"""
Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
@ -157,9 +158,12 @@ async def _send_message_via_completion_bridge(
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
)
return LiteLLMSendMessageResponse.from_dict(response_dict)
return LiteLLMSendMessageResponse.from_dict(
response_dict, request_id=str(request.id)
)
async def _execute_a2a_send_with_retry(
@ -281,6 +285,7 @@ async def asend_message(
custom_llm_provider=custom_llm_provider,
api_base=api_base,
litellm_params=litellm_params,
agent_extra_headers=agent_extra_headers,
)
# Standard A2A client flow
@ -317,15 +322,6 @@ async def asend_message(
)
card_url = getattr(agent_card, "url", None) if agent_card else None
context_id = trace_id or str(uuid.uuid4())
message = request.params.message
if isinstance(message, dict):
if message.get("context_id") is None:
message["context_id"] = context_id
else:
if getattr(message, "context_id", None) is None:
message.context_id = context_id
a2a_response = await _execute_a2a_send_with_retry(
a2a_client=a2a_client,
request=request,
@ -338,7 +334,9 @@ async def asend_message(
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
# Wrap in LiteLLM response type for _hidden_params support
response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response)
response = LiteLLMSendMessageResponse.from_a2a_response(
a2a_response, request_id=str(request.id)
)
# Calculate token usage from request and response
response_dict = a2a_response.model_dump(mode="json", exclude_none=True)
@ -514,6 +512,7 @@ async def asend_message_streaming( # noqa: PLR0915
params=params,
litellm_params=litellm_params,
api_base=api_base,
agent_extra_headers=agent_extra_headers,
):
yield chunk
return

View file

@ -37,6 +37,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
request_id=request_id,
params=params,
litellm_params=litellm_params,
agent_extra_headers=kwargs.get("agent_extra_headers"),
)
async def handle_streaming(
@ -57,5 +58,6 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
request_id=request_id,
params=params,
litellm_params=litellm_params,
agent_extra_headers=kwargs.get("agent_extra_headers"),
):
yield chunk

View file

@ -6,7 +6,7 @@ completion bridge that would otherwise strip the envelope.
"""
import json
from typing import Any, AsyncIterator, Dict, cast
from typing import Any, AsyncIterator, Dict, Optional, cast
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
@ -29,6 +29,7 @@ class BedrockAgentCoreA2AHandler:
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""
Handle non-streaming A2A request to AgentCore.
@ -37,6 +38,8 @@ class BedrockAgentCoreA2AHandler:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (model, api_key, etc.)
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Returns:
A2A JSON-RPC response dict from the AgentCore agent
@ -47,6 +50,7 @@ class BedrockAgentCoreA2AHandler:
params=params,
litellm_params=litellm_params,
method="message/send",
agent_extra_headers=agent_extra_headers,
)
)
@ -77,6 +81,7 @@ class BedrockAgentCoreA2AHandler:
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
"""
Handle streaming A2A request to AgentCore.
@ -85,6 +90,8 @@ class BedrockAgentCoreA2AHandler:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (model, api_key, etc.)
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Yields:
A2A streaming response events from the AgentCore agent
@ -96,6 +103,7 @@ class BedrockAgentCoreA2AHandler:
litellm_params=litellm_params,
method="message/send",
stream=True,
agent_extra_headers=agent_extra_headers,
)
)

View file

@ -6,11 +6,66 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT).
"""
import json
from typing import Any, AsyncIterator, Dict, Tuple
from typing import Any, AsyncIterator, Dict, Mapping, Optional, Tuple
from litellm._logging import verbose_logger
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
# Reserved outbound header names that must never be sourced from per-request
# ``agent_extra_headers`` for AgentCore requests. ``agent_extra_headers`` carries
# values rewritten from the client-controlled ``x-a2a-{agent}-*`` convention, so
# allowing these would let any caller with access to the agent spoof the AWS
# request identity / SigV4 metadata by overwriting headers the proxy sets from
# trusted server-side config.
#
# The runtime headers (session / user id) are derived server-side from
# ``runtimeSessionId`` / ``runtimeUserId`` in the agent's ``litellm_params``;
# ``authorization`` is set by the AgentCore signer (JWT or SigV4); ``host`` and
# the ``x-amz-*`` family are owned by SigV4 itself.
_RESERVED_EXACT_HEADERS = frozenset(
{
"authorization",
"host",
}
)
_RESERVED_PREFIX_HEADERS: Tuple[str, ...] = (
"x-amzn-bedrock-agentcore-runtime-",
"x-amz-",
)
def _filter_reserved_headers(
agent_extra_headers: Optional[Mapping[str, str]],
) -> Optional[Dict[str, str]]:
"""
Strip reserved AWS / AgentCore headers from caller-supplied
``agent_extra_headers`` before they are merged into the signed request.
Returns ``None`` if the result is empty.
"""
if not agent_extra_headers:
return None
filtered: Dict[str, str] = {}
dropped: list = []
for k, v in agent_extra_headers.items():
k_lower = k.lower()
if k_lower in _RESERVED_EXACT_HEADERS or any(
k_lower.startswith(prefix) for prefix in _RESERVED_PREFIX_HEADERS
):
dropped.append(k)
continue
filtered[k] = v
if dropped:
verbose_logger.warning(
"BedrockAgentCore A2A: dropping reserved header(s) from "
"agent_extra_headers (not forwarded to AgentCore): %s",
sorted(dropped),
)
return filtered or None
class BedrockAgentCoreA2ATransformation:
"""
@ -27,6 +82,7 @@ class BedrockAgentCoreA2ATransformation:
litellm_params: Dict[str, Any],
method: str = "message/send",
stream: bool = False,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Tuple[str, dict, bytes]:
"""
Build the AgentCore URL, construct a JSON-RPC envelope, and sign the request.
@ -37,6 +93,15 @@ class BedrockAgentCoreA2ATransformation:
litellm_params: Agent's litellm_params (model, api_key, etc.)
method: JSON-RPC method name (default: "message/send")
stream: Whether this is a streaming request
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call. Merged into
the headers dict before signing so SigV4 includes them in the signature.
Reserved AWS / AgentCore identity headers (``authorization``, ``host``,
``x-amzn-bedrock-agentcore-runtime-*``, ``x-amz-*``) are filtered out
here to prevent a caller-controlled ``x-a2a-{agent}-*`` header from
spoofing the AgentCore runtime user id or other SigV4 metadata. Use
``api_key`` / ``runtimeUserId`` / ``runtimeSessionId`` in litellm_params
(not ``agent_extra_headers``) to override those values.
Returns:
Tuple of (url, signed_headers, signed_body_bytes)
@ -85,6 +150,13 @@ class BedrockAgentCoreA2ATransformation:
if runtime_user_id:
headers["X-Amzn-Bedrock-AgentCore-Runtime-User-Id"] = runtime_user_id
# Merge per-request agent headers before signing so SigV4 covers them.
# Reserved headers are stripped first to prevent client-controlled values
# from spoofing the AgentCore runtime identity / SigV4 metadata.
safe_extra_headers = _filter_reserved_headers(agent_extra_headers)
if safe_extra_headers:
headers.update(safe_extra_headers)
# Sign the request (SigV4 or JWT depending on api_key presence)
signed_headers, signed_body = agentcore_config.sign_request(
headers=headers,

View file

@ -48,4 +48,16 @@ class A2AProviderConfigManager:
return BedrockAgentCoreA2AConfig()
if custom_llm_provider == "langflow":
from litellm.a2a_protocol.providers.langflow.config import LangFlowA2AConfig
return LangFlowA2AConfig()
if custom_llm_provider == "watsonx_orchestrate":
from litellm.a2a_protocol.providers.watsonx_orchestrate.config import (
WatsonxOrchestrateA2AConfig,
)
return WatsonxOrchestrateA2AConfig()
return None

View file

@ -0,0 +1,62 @@
from typing import Any, AsyncIterator, Dict, Optional
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2A_USER_API_KEY_HASH_PARAM,
A2ACompletionBridgeHandler,
)
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.llms.langflow.a2a import merge_a2a_session_into_litellm_params
class LangFlowA2AConfig(BaseA2AProviderConfig):
"""A2A bridge for LangFlow: scopes contextId to the authenticated key as the
LangFlow session_id, then uses completion."""
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
**kwargs,
) -> Dict[str, Any]:
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for LangFlowA2AConfig "
"(must contain custom_llm_provider and model)"
)
litellm_params = merge_a2a_session_into_litellm_params(
litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)
)
return await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
api_base=api_base,
_skip_a2a_provider_routing=True,
)
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for LangFlowA2AConfig "
"(must contain custom_llm_provider and model)"
)
litellm_params = merge_a2a_session_into_litellm_params(
litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)
)
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
api_base=api_base,
_skip_a2a_provider_routing=True,
):
yield chunk

View file

@ -1,74 +0,0 @@
# A2A to LiteLLM Completion Bridge
Routes A2A protocol requests through `litellm.acompletion`, enabling any LiteLLM-supported provider to be invoked via A2A.
## Flow
```
A2A Request → Transform → litellm.acompletion → Transform → A2A Response
```
## SDK Usage
Use the existing `asend_message` and `asend_message_streaming` functions with `litellm_params`:
```python
from litellm.a2a_protocol import asend_message, asend_message_streaming
from a2a.types import SendMessageRequest, SendStreamingMessageRequest, MessageSendParams
from uuid import uuid4
# Non-streaming
request = SendMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={"role": "user", "parts": [{"kind": "text", "text": "Hello!"}], "messageId": uuid4().hex}
)
)
response = await asend_message(
request=request,
api_base="http://localhost:2024",
litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
)
# Streaming
stream_request = SendStreamingMessageRequest(
id=str(uuid4()),
params=MessageSendParams(
message={"role": "user", "parts": [{"kind": "text", "text": "Hello!"}], "messageId": uuid4().hex}
)
)
async for chunk in asend_message_streaming(
request=stream_request,
api_base="http://localhost:2024",
litellm_params={"custom_llm_provider": "langgraph", "model": "agent"},
):
print(chunk)
```
## Proxy Usage
Configure an agent with `custom_llm_provider` in `litellm_params`:
```yaml
agents:
- agent_name: my-langgraph-agent
agent_card_params:
name: "LangGraph Agent"
url: "http://localhost:2024" # Used as api_base
litellm_params:
custom_llm_provider: langgraph
model: agent
```
When an A2A request hits `/a2a/{agent_id}/message/send`, the bridge:
1. Detects `custom_llm_provider` in agent's `litellm_params`
2. Transforms A2A message → OpenAI messages
3. Calls `litellm.acompletion(model="langgraph/agent", api_base="http://localhost:2024")`
4. Transforms response → A2A format
## Classes
- `A2ACompletionBridgeTransformation` - Static methods for message format conversion
- `A2ACompletionBridgeHandler` - Static methods for handling requests (streaming/non-streaming)

View file

@ -1,5 +0,0 @@
"""
LiteLLM Completion bridge provider for A2A protocol.
Routes A2A requests through litellm.acompletion based on custom_llm_provider.
"""

View file

@ -1,301 +0,0 @@
"""
Handler for A2A to LiteLLM completion bridge.
Routes A2A requests through litellm.acompletion based on custom_llm_provider.
A2A Streaming Events (in order):
1. Task event (kind: "task") - Initial task creation with status "submitted"
2. Status update (kind: "status-update") - Status change to "working"
3. Artifact update (kind: "artifact-update") - Content/artifact delivery
4. Status update (kind: "status-update") - Final status "completed" with final=true
"""
from typing import Any, AsyncIterator, Dict, Optional
import litellm
from litellm._logging import verbose_logger
from litellm.a2a_protocol.litellm_completion_bridge.pydantic_ai_transformation import (
PydanticAITransformation,
)
from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
A2ACompletionBridgeTransformation,
A2AStreamingContext,
)
class A2ACompletionBridgeHandler:
"""
Static methods for handling A2A requests via LiteLLM completion.
"""
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
) -> Dict[str, Any]:
"""
Handle non-streaming A2A request via litellm.acompletion.
Args:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
api_base: API base URL from agent_card_params
Returns:
A2A SendMessageResponse dict
"""
# Check if this is a Pydantic AI agent request
custom_llm_provider = litellm_params.get("custom_llm_provider")
if custom_llm_provider == "pydantic_ai_agents":
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
verbose_logger.info(
f"Pydantic AI: Routing to Pydantic AI agent at {api_base}"
)
# Send request directly to Pydantic AI agent
response_data = await PydanticAITransformation.send_non_streaming_request(
api_base=api_base,
request_id=request_id,
params=params,
)
return response_data
# Extract message from params
message = params.get("message", {})
# Transform A2A message to OpenAI format
openai_messages = (
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
full_model = f"{custom_llm_provider}/{model}"
else:
full_model = model
verbose_logger.info(
f"A2A completion bridge: model={full_model}, api_base={api_base}"
)
# Build completion params dict
completion_params = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": False,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
# Call litellm.acompletion
response = await litellm.acompletion(**completion_params)
# Transform response to A2A format
a2a_response = (
A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
response=response,
request_id=request_id,
)
)
verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}")
return a2a_response
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
) -> AsyncIterator[Dict[str, Any]]:
"""
Handle streaming A2A request via litellm.acompletion with stream=True.
Emits proper A2A streaming events:
1. Task event (kind: "task") - Initial task with status "submitted"
2. Status update (kind: "status-update") - Status "working"
3. Artifact update (kind: "artifact-update") - Content delivery
4. Status update (kind: "status-update") - Final "completed" status
Args:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
api_base: API base URL from agent_card_params
Yields:
A2A streaming response events
"""
# Check if this is a Pydantic AI agent request
custom_llm_provider = litellm_params.get("custom_llm_provider")
if custom_llm_provider == "pydantic_ai_agents":
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
verbose_logger.info(
f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}"
)
# Get non-streaming response first
response_data = await PydanticAITransformation.send_non_streaming_request(
api_base=api_base,
request_id=request_id,
params=params,
)
# Convert to fake streaming
async for chunk in PydanticAITransformation.fake_streaming_from_response(
response_data=response_data,
request_id=request_id,
):
yield chunk
return
# Extract message from params
message = params.get("message", {})
# Create streaming context
ctx = A2AStreamingContext(
request_id=request_id,
input_message=message,
)
# Transform A2A message to OpenAI format
openai_messages = (
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
full_model = f"{custom_llm_provider}/{model}"
else:
full_model = model
verbose_logger.info(
f"A2A completion bridge streaming: model={full_model}, api_base={api_base}"
)
# Build completion params dict
completion_params = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": True,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)
yield task_event
# 2. Emit status update (kind: "status-update", status: "working")
working_event = A2ACompletionBridgeTransformation.create_status_update_event(
ctx=ctx,
state="working",
final=False,
message_text="Processing request...",
)
yield working_event
# Call litellm.acompletion with streaming
response = await litellm.acompletion(**completion_params)
# 3. Accumulate content and emit artifact update
accumulated_text = ""
chunk_count = 0
async for chunk in response: # type: ignore[union-attr]
chunk_count += 1
# Extract delta content
content = ""
if chunk is not None and hasattr(chunk, "choices") and chunk.choices:
choice = chunk.choices[0]
if hasattr(choice, "delta") and choice.delta:
content = choice.delta.content or ""
if content:
accumulated_text += content
# Emit artifact update with accumulated content
if accumulated_text:
artifact_event = (
A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=accumulated_text,
)
)
yield artifact_event
# 4. Emit final status update (kind: "status-update", status: "completed", final: true)
completed_event = A2ACompletionBridgeTransformation.create_status_update_event(
ctx=ctx,
state="completed",
final=True,
)
yield completed_event
verbose_logger.info(
f"A2A completion bridge streaming completed: request_id={request_id}, chunks={chunk_count}"
)
# Convenience functions that delegate to the class methods
async def handle_a2a_completion(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
) -> Dict[str, Any]:
"""Convenience function for non-streaming A2A completion."""
return await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
api_base=api_base,
)
async def handle_a2a_completion_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
) -> AsyncIterator[Dict[str, Any]]:
"""Convenience function for streaming A2A completion."""
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
api_base=api_base,
):
yield chunk

View file

@ -1,284 +0,0 @@
"""
Transformation utilities for A2A <-> OpenAI message format conversion.
A2A Message Format:
{
"role": "user",
"parts": [{"kind": "text", "text": "Hello!"}],
"messageId": "abc123"
}
OpenAI Message Format:
{"role": "user", "content": "Hello!"}
A2A Streaming Events:
- Task event (kind: "task") - Initial task creation with status "submitted"
- Status update (kind: "status-update") - Status changes (working, completed)
- Artifact update (kind: "artifact-update") - Content/artifact delivery
"""
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from uuid import uuid4
from litellm._logging import verbose_logger
class A2AStreamingContext:
"""
Context holder for A2A streaming state.
Tracks task_id, context_id, and message accumulation.
"""
def __init__(self, request_id: str, input_message: Dict[str, Any]):
self.request_id = request_id
self.task_id = str(uuid4())
self.context_id = str(uuid4())
self.input_message = input_message
self.accumulated_text = ""
self.has_emitted_task = False
self.has_emitted_working = False
class A2ACompletionBridgeTransformation:
"""
Static methods for transforming between A2A and OpenAI message formats.
"""
@staticmethod
def a2a_message_to_openai_messages(
a2a_message: Dict[str, Any],
) -> List[Dict[str, str]]:
"""
Transform an A2A message to OpenAI message format.
Args:
a2a_message: A2A message with role, parts, and messageId
Returns:
List of OpenAI-format messages
"""
role = a2a_message.get("role", "user")
parts = a2a_message.get("parts", [])
# Map A2A roles to OpenAI roles
openai_role = role
if role == "user":
openai_role = "user"
elif role == "assistant":
openai_role = "assistant"
elif role == "system":
openai_role = "system"
# Extract text content from parts
content_parts = []
for part in parts:
kind = part.get("kind", "")
if kind == "text":
text = part.get("text", "")
content_parts.append(text)
content = "\n".join(content_parts) if content_parts else ""
verbose_logger.debug(
f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}"
)
return [{"role": openai_role, "content": content}]
@staticmethod
def openai_response_to_a2a_response(
response: Any,
request_id: Optional[str] = None,
) -> Dict[str, Any]:
"""
Transform a LiteLLM ModelResponse to A2A SendMessageResponse format.
Args:
response: LiteLLM ModelResponse object
request_id: Original A2A request ID
Returns:
A2A SendMessageResponse dict
"""
# Extract content from response
content = ""
if hasattr(response, "choices") and response.choices:
choice = response.choices[0]
if hasattr(choice, "message") and choice.message:
content = choice.message.content or ""
# Build A2A message
a2a_message = {
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
}
# Build A2A response
a2a_response = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"message": a2a_message,
},
}
verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}")
return a2a_response
@staticmethod
def _get_timestamp() -> str:
"""Get current timestamp in ISO format with timezone."""
return datetime.now(timezone.utc).isoformat()
@staticmethod
def create_task_event(
ctx: A2AStreamingContext,
) -> Dict[str, Any]:
"""
Create the initial task event with status 'submitted'.
This is the first event emitted in an A2A streaming response.
"""
return {
"id": ctx.request_id,
"jsonrpc": "2.0",
"result": {
"contextId": ctx.context_id,
"history": [
{
"contextId": ctx.context_id,
"kind": "message",
"messageId": ctx.input_message.get("messageId", uuid4().hex),
"parts": ctx.input_message.get("parts", []),
"role": ctx.input_message.get("role", "user"),
"taskId": ctx.task_id,
}
],
"id": ctx.task_id,
"kind": "task",
"status": {
"state": "submitted",
},
},
}
@staticmethod
def create_status_update_event(
ctx: A2AStreamingContext,
state: str,
final: bool = False,
message_text: Optional[str] = None,
) -> Dict[str, Any]:
"""
Create a status update event.
Args:
ctx: Streaming context
state: Status state ('working', 'completed')
final: Whether this is the final event
message_text: Optional message text for 'working' status
"""
status: Dict[str, Any] = {
"state": state,
"timestamp": A2ACompletionBridgeTransformation._get_timestamp(),
}
# Add message for 'working' status
if state == "working" and message_text:
status["message"] = {
"contextId": ctx.context_id,
"kind": "message",
"messageId": str(uuid4()),
"parts": [{"kind": "text", "text": message_text}],
"role": "agent",
"taskId": ctx.task_id,
}
return {
"id": ctx.request_id,
"jsonrpc": "2.0",
"result": {
"contextId": ctx.context_id,
"final": final,
"kind": "status-update",
"status": status,
"taskId": ctx.task_id,
},
}
@staticmethod
def create_artifact_update_event(
ctx: A2AStreamingContext,
text: str,
) -> Dict[str, Any]:
"""
Create an artifact update event with content.
Args:
ctx: Streaming context
text: The text content for the artifact
"""
return {
"id": ctx.request_id,
"jsonrpc": "2.0",
"result": {
"artifact": {
"artifactId": str(uuid4()),
"name": "response",
"parts": [{"kind": "text", "text": text}],
},
"contextId": ctx.context_id,
"kind": "artifact-update",
"taskId": ctx.task_id,
},
}
@staticmethod
def openai_chunk_to_a2a_chunk(
chunk: Any,
request_id: Optional[str] = None,
is_final: bool = False,
) -> Optional[Dict[str, Any]]:
"""
Transform a LiteLLM streaming chunk to A2A streaming format.
NOTE: This method is deprecated for streaming. Use the event-based
methods (create_task_event, create_status_update_event,
create_artifact_update_event) instead for proper A2A streaming.
Args:
chunk: LiteLLM ModelResponse chunk
request_id: Original A2A request ID
is_final: Whether this is the final chunk
Returns:
A2A streaming chunk dict or None if no content
"""
# Extract delta content
content = ""
if chunk is not None and hasattr(chunk, "choices") and chunk.choices:
choice = chunk.choices[0]
if hasattr(choice, "delta") and choice.delta:
content = choice.delta.content or ""
if not content and not is_final:
return None
# Build A2A streaming chunk (legacy format)
a2a_chunk = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"message": {
"role": "agent",
"parts": [{"kind": "text", "text": content}],
"messageId": uuid4().hex,
},
"final": is_final,
},
}
return a2a_chunk

View file

@ -31,6 +31,7 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
params=params,
api_base=api_base,
timeout=kwargs.get("timeout", 60.0),
agent_extra_headers=kwargs.get("agent_extra_headers"),
)
async def handle_streaming(
@ -50,5 +51,6 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
timeout=kwargs.get("timeout", 60.0),
chunk_size=kwargs.get("chunk_size", 50),
delay_ms=kwargs.get("delay_ms", 10),
agent_extra_headers=kwargs.get("agent_extra_headers"),
):
yield chunk

View file

@ -28,6 +28,7 @@ class PydanticAIHandler:
params: Dict[str, Any],
api_base: Optional[str] = None,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""
Handle non-streaming request to Pydantic AI agent.
@ -37,6 +38,8 @@ class PydanticAIHandler:
params: A2A MessageSendParams containing the message
api_base: Base URL of the Pydantic AI agent
timeout: Request timeout in seconds
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Returns:
A2A SendMessageResponse dict
@ -51,6 +54,7 @@ class PydanticAIHandler:
request_id=request_id,
params=params,
timeout=timeout,
agent_extra_headers=agent_extra_headers,
)
return response_data
@ -63,6 +67,7 @@ class PydanticAIHandler:
timeout: float = 60.0,
chunk_size: int = 50,
delay_ms: int = 10,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
"""
Handle streaming request to Pydantic AI agent with fake streaming.
@ -78,6 +83,8 @@ class PydanticAIHandler:
timeout: Request timeout in seconds
chunk_size: Number of characters per chunk
delay_ms: Delay between chunks in milliseconds
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
admin extra_headers) to forward on the upstream HTTP call.
Yields:
A2A streaming response events
@ -94,6 +101,7 @@ class PydanticAIHandler:
request_id=request_id,
params=params,
timeout=timeout,
agent_extra_headers=agent_extra_headers,
)
# Convert raw task response to fake streaming chunks

View file

@ -6,7 +6,7 @@ This module provides fake streaming by converting non-streaming responses into s
"""
import asyncio
from typing import Any, AsyncIterator, Dict, cast
from typing import Any, AsyncIterator, Dict, Optional, cast
from uuid import uuid4
from litellm._logging import verbose_logger
@ -86,6 +86,7 @@ class PydanticAITransformation:
request_id: str,
max_attempts: int = 30,
poll_interval: float = 0.5,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""
Poll for task completion using tasks/get method.
@ -112,7 +113,10 @@ class PydanticAITransformation:
response = await client.post(
endpoint,
json=poll_request,
headers={"Content-Type": "application/json"},
headers={
**(agent_extra_headers or {}),
"Content-Type": "application/json",
},
)
response.raise_for_status()
poll_data = response.json()
@ -142,6 +146,7 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -189,7 +194,10 @@ class PydanticAITransformation:
response = await client.post(
endpoint,
json=a2a_request,
headers={"Content-Type": "application/json"},
headers={
**(agent_extra_headers or {}),
"Content-Type": "application/json",
},
)
response.raise_for_status()
response_data = response.json()
@ -211,6 +219,7 @@ class PydanticAITransformation:
endpoint=endpoint,
task_id=task_id,
request_id=request_id,
agent_extra_headers=agent_extra_headers,
)
verbose_logger.info(
@ -225,6 +234,7 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""
Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
@ -234,6 +244,7 @@ class PydanticAITransformation:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message (dict or Pydantic model)
timeout: Request timeout in seconds
agent_extra_headers: Per-request headers to forward on the upstream HTTP call.
Returns:
Standard A2A non-streaming response format with message
@ -244,6 +255,7 @@ class PydanticAITransformation:
request_id=request_id,
params=params,
timeout=timeout,
agent_extra_headers=agent_extra_headers,
)
# Transform to standard A2A non-streaming format
@ -258,6 +270,7 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -269,6 +282,7 @@ class PydanticAITransformation:
request_id: A2A JSON-RPC request ID
params: A2A MessageSendParams containing the message
timeout: Request timeout in seconds
agent_extra_headers: Per-request headers to forward on the upstream HTTP call.
Returns:
Raw Pydantic AI task response (with history/artifacts)
@ -278,6 +292,7 @@ class PydanticAITransformation:
request_id=request_id,
params=params,
timeout=timeout,
agent_extra_headers=agent_extra_headers,
)
@staticmethod
@ -289,16 +304,16 @@ class PydanticAITransformation:
Transform Pydantic AI task response to standard A2A non-streaming format.
Pydantic AI returns a task with history/artifacts, but the standard A2A
non-streaming format expects:
non-streaming format expects ``result`` to be the Message directly
(``kind="message"``), per the A2A spec / ``SendMessageResponse``:
{
"jsonrpc": "2.0",
"id": "...",
"result": {
"message": {
"role": "agent",
"parts": [{"kind": "text", "text": "..."}],
"messageId": "..."
}
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": "..."}],
"messageId": "..."
}
}
@ -316,6 +331,7 @@ class PydanticAITransformation:
# Build standard A2A message
a2a_message = {
"kind": "message",
"role": "agent",
"parts": parts if parts else [{"kind": "text", "text": full_text}],
"messageId": message_id,
@ -325,9 +341,7 @@ class PydanticAITransformation:
return {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"message": a2a_message,
},
"result": a2a_message,
}
@staticmethod

View file

@ -0,0 +1,3 @@
"""
IBM watsonx Orchestrate (WXO) A2A provider.
"""

View file

@ -0,0 +1,55 @@
"""
A2A provider configuration for IBM watsonx Orchestrate (WXO).
"""
from typing import Any, AsyncIterator, Dict, Optional
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
WatsonxOrchestrateHandler,
)
class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
"""A2A bridge for IBM watsonx Orchestrate (REST runs API + poll/SSE)."""
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
**kwargs: Any,
) -> Dict[str, Any]:
"""Handle a non-streaming A2A request via WXO runs API."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for WatsonxOrchestrateA2AConfig "
"(must contain cp4d_host, instance_id, wxo_agent_id, api_key)"
)
return await WatsonxOrchestrateHandler.handle_non_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
)
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
**kwargs: Any,
) -> AsyncIterator[Dict[str, Any]]:
"""Handle a streaming A2A request via WXO streaming runs API."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for WatsonxOrchestrateA2AConfig "
"(must contain cp4d_host, instance_id, wxo_agent_id, api_key)"
)
async for chunk in WatsonxOrchestrateHandler.handle_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
):
yield chunk

View file

@ -0,0 +1,373 @@
"""
Handler for IBM watsonx Orchestrate (WXO) agent provider.
"""
import asyncio
import hashlib
import json
import time
from typing import Any, AsyncIterator, Dict, NamedTuple, Optional, Tuple, cast
import httpx
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import (
WatsonxOrchestrateTransformation,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
)
from litellm.types.llms.custom_http import httpxSpecialProvider
_IBM_CLOUD_IAM_URL = "https://iam.cloud.ibm.com/identity/token"
_POLL_INTERVAL_S = 2.0
_MAX_POLL_ATTEMPTS = 90
_TOKEN_CACHE_TTL_BUFFER_S = 60
_token_cache: Dict[str, Tuple[str, float]] = {}
class WXORequestParams(NamedTuple):
cp4d_host: str
instance_id: str
wxo_agent_id: str
api_key: str
username: Optional[str]
auth_mode: str
thread_id: Optional[str]
class WatsonxOrchestrateHandler:
@staticmethod
def _http_client(timeout: float = 90.0) -> AsyncHTTPHandler:
return get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
params={"timeout": timeout},
)
@staticmethod
def _token_cache_key(
auth_mode: str,
cp4d_host: str,
api_key: str,
username: Optional[str],
) -> str:
material = f"{auth_mode}:{cp4d_host}:{username or ''}:{api_key}"
return hashlib.sha256(material.encode()).hexdigest()
@staticmethod
def _cp4d_token_ttl_seconds(
expiration: Any, now_wall: Optional[float] = None
) -> int:
# CP4D returns expiration as absolute Unix epoch seconds, not a duration.
expires_at = int(expiration)
wall = now_wall if now_wall is not None else time.time()
return max(expires_at - int(wall), 0)
@staticmethod
async def _get_bearer_token(
cp4d_host: str,
auth_mode: str,
api_key: str,
username: Optional[str] = None,
client: Optional[AsyncHTTPHandler] = None,
) -> str:
cache_key = WatsonxOrchestrateHandler._token_cache_key(
auth_mode, cp4d_host, api_key, username
)
now = time.monotonic()
cached = _token_cache.get(cache_key)
if cached and cached[1] > now:
return cached[0]
if client is None:
client = WatsonxOrchestrateHandler._http_client(timeout=30.0)
if auth_mode == "ibm_cloud":
response = await client.post(
_IBM_CLOUD_IAM_URL,
data={
"grant_type": "urn:ibm:params:oauth:grant-type:apikey",
"apikey": api_key,
},
headers={"Content-Type": "application/x-www-form-urlencoded"},
)
response.raise_for_status()
payload = response.json()
token = str(payload["access_token"])
ttl_s = int(payload.get("expires_in", 3600))
else:
if not username:
raise ValueError(
"'username' is required in litellm_params when auth_mode='cp4d'"
)
token_url = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize"
response = await client.post(
token_url,
json={"username": username, "api_key": api_key},
headers={"Content-Type": "application/json"},
)
response.raise_for_status()
payload = response.json()
token = str(payload["token"])
expiration = payload.get("expiration")
if expiration is None:
ttl_s = 3600
else:
ttl_s = WatsonxOrchestrateHandler._cp4d_token_ttl_seconds(expiration)
expires_at = now + max(ttl_s - _TOKEN_CACHE_TTL_BUFFER_S, 0)
_token_cache[cache_key] = (token, expires_at)
for stale_key, (_, stale_expires_at) in list(_token_cache.items()):
if stale_expires_at <= now:
del _token_cache[stale_key]
return token
@staticmethod
async def _poll_run(
base_url: str,
run_id: str,
auth_headers: Dict[str, str],
client: AsyncHTTPHandler,
max_attempts: int = _MAX_POLL_ATTEMPTS,
interval_s: float = _POLL_INTERVAL_S,
) -> Dict[str, Any]:
url = f"{base_url}/v1/orchestrate/runs/{run_id}"
for attempt in range(max_attempts):
await asyncio.sleep(interval_s)
response = await client.get(url, headers=auth_headers)
response.raise_for_status()
result: Dict[str, Any] = response.json()
status = result.get("status", "")
verbose_logger.debug(
f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'"
)
if status in WatsonxOrchestrateTransformation.TERMINAL_STATES:
return result
raise asyncio.TimeoutError(
f"WXO run '{run_id}' did not reach a terminal state after "
f"{max_attempts * interval_s:.0f}s"
)
@staticmethod
async def _get_successful_run_data(
run_data: Dict[str, Any],
base_url: str,
auth_headers: Dict[str, str],
client: AsyncHTTPHandler,
) -> Dict[str, Any]:
status = run_data.get("status", "")
if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
run_id = run_data.get("run_id") or run_data.get("id") or ""
if not run_id:
raise ValueError(f"WXO: No run_id in response: {run_data}")
run_data = await WatsonxOrchestrateHandler._poll_run(
base_url=base_url,
run_id=run_id,
auth_headers=auth_headers,
client=client,
)
status = run_data.get("status", "")
if status not in WatsonxOrchestrateTransformation.SUCCESS_STATES:
raise RuntimeError(
f"WXO run ended with non-success status '{status}': {run_data}"
)
return run_data
@staticmethod
async def _accumulate_wxo_sse_text(response: Any) -> str:
accumulated_text = ""
async for line in response.aiter_lines():
if not line.startswith("data:"):
continue
data_str = line[5:].strip()
if not data_str or data_str == "[DONE]":
continue
try:
event = json.loads(data_str)
except json.JSONDecodeError:
continue
chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(
event
)
if chunk_text:
accumulated_text += chunk_text
return accumulated_text
@staticmethod
def _extract_litellm_params(litellm_params: Dict[str, Any]) -> WXORequestParams:
cp4d_host = litellm_params.get("cp4d_host") or ""
instance_id = litellm_params.get("instance_id") or ""
wxo_agent_id = litellm_params.get("wxo_agent_id") or ""
api_key = litellm_params.get("api_key") or ""
if not cp4d_host:
raise ValueError("'cp4d_host' is required in litellm_params for WXO agents")
if not instance_id:
raise ValueError(
"'instance_id' is required in litellm_params for WXO agents"
)
if not wxo_agent_id:
raise ValueError(
"'wxo_agent_id' is required in litellm_params for WXO agents"
)
if not api_key:
raise ValueError("'api_key' is required in litellm_params for WXO agents")
return WXORequestParams(
cp4d_host=cp4d_host,
instance_id=instance_id,
wxo_agent_id=wxo_agent_id,
api_key=api_key,
username=litellm_params.get("username") or None,
auth_mode=litellm_params.get("auth_mode") or "cp4d",
thread_id=litellm_params.get("thread_id") or None,
)
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
) -> Dict[str, Any]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=90.0)
token = await WatsonxOrchestrateHandler._get_bearer_token(
cp4d_host=wxo.cp4d_host,
auth_mode=wxo.auth_mode,
api_key=wxo.api_key,
username=wxo.username,
client=client,
)
base_url = WatsonxOrchestrateTransformation.get_api_base_url(
wxo.cp4d_host, wxo.instance_id
)
auth_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"Accept": "application/json",
}
text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body = WatsonxOrchestrateTransformation.build_wxo_run_body(
wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id
)
run_response = await client.post(
f"{base_url}/v1/orchestrate/runs",
json=body,
headers=auth_headers,
)
run_response.raise_for_status()
run_data: Dict[str, Any] = run_response.json()
run_data = await WatsonxOrchestrateHandler._get_successful_run_data(
run_data=run_data,
base_url=base_url,
auth_headers=auth_headers,
client=client,
)
response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(
run_data
)
return WatsonxOrchestrateTransformation.build_a2a_message_response(
request_id=request_id, text=response_text
)
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=120.0)
token = await WatsonxOrchestrateHandler._get_bearer_token(
cp4d_host=wxo.cp4d_host,
auth_mode=wxo.auth_mode,
api_key=wxo.api_key,
username=wxo.username,
client=client,
)
base_url = WatsonxOrchestrateTransformation.get_api_base_url(
wxo.cp4d_host, wxo.instance_id
)
auth_headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"Accept": "text/event-stream, application/json",
}
text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body = WatsonxOrchestrateTransformation.build_wxo_run_body(
wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id
)
try:
response = await client.post(
f"{base_url}/v1/orchestrate/runs/stream",
json=body,
headers=auth_headers,
stream=True,
)
response.raise_for_status()
except httpx.TransportError as exc:
verbose_logger.warning(
f"WXO: Streaming request failed before a run was submitted "
f"({exc!r}), falling back to non-streaming + fake streaming",
exc_info=True,
)
result = await WatsonxOrchestrateHandler.handle_non_streaming(
request_id=request_id,
params=params,
litellm_params=litellm_params,
)
response_text = (
WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(
result
)
)
async for (
chunk
) in WatsonxOrchestrateTransformation.fake_streaming_from_text(
text=response_text,
request_id=request_id,
chunk_size=chunk_size,
delay_ms=delay_ms,
):
yield chunk
return
content_type = response.headers.get("content-type", "").lower()
if "text/event-stream" not in content_type:
response_body = await response.aread()
result = json.loads(response_body)
result = await WatsonxOrchestrateHandler._get_successful_run_data(
run_data=result,
base_url=base_url,
auth_headers=auth_headers,
client=client,
)
accumulated_text = (
WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result)
)
else:
accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(
response
)
async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text(
text=accumulated_text,
request_id=request_id,
chunk_size=chunk_size,
delay_ms=delay_ms,
):
yield chunk

View file

@ -0,0 +1,224 @@
"""
Transformation layer for IBM watsonx Orchestrate (WXO) agent provider.
WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model:
POST /v1/orchestrate/runs → submit run, get run_id
GET /v1/orchestrate/runs/{id} → poll until terminal state
POST /v1/orchestrate/runs/stream → native SSE streaming
"""
import asyncio
from typing import Any, AsyncIterator, Dict, Optional
from uuid import uuid4
from litellm._logging import verbose_logger
class WatsonxOrchestrateTransformation:
"""
Handles request/response transformation between A2A and the WXO REST API.
"""
TERMINAL_STATES = frozenset(
{"completed", "succeeded", "failed", "error", "cancelled"}
)
SUCCESS_STATES = frozenset({"completed", "succeeded"})
@staticmethod
def get_api_base_url(cp4d_host: str, instance_id: str) -> str:
"""Build the WXO API base URL from host and instance ID."""
return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}"
@staticmethod
def extract_text_from_a2a_params(params: Dict[str, Any]) -> str:
"""
Extract user message text from A2A MessageSendParams.
A2A format: params.message.parts[*] where part.kind == "text"
"""
message = params.get("message", {})
parts = message.get("parts", [])
texts = []
for part in parts:
if not isinstance(part, dict):
continue
kind = part.get("kind")
if kind in (None, "", "text") and part.get("text"):
texts.append(part["text"])
return " ".join(texts) or ""
@staticmethod
def build_wxo_run_body(
wxo_agent_id: str,
text: str,
thread_id: Optional[str] = None,
) -> Dict[str, Any]:
"""Build the WXO POST /v1/orchestrate/runs request body."""
body: Dict[str, Any] = {
"agent_id": wxo_agent_id,
"message": {
"role": "user",
"content": [
{
"response_type": "text",
"text": text,
}
],
},
}
if thread_id:
body["thread_id"] = thread_id
return body
@staticmethod
def extract_text_from_wxo_result(result: Any) -> str:
"""
Extract response text from a WXO run result.
WXO can return text in several locations; checks in priority order per the API spec.
"""
if not isinstance(result, dict):
return ""
# Primary: last_message.content[0].text
try:
text = result["last_message"]["content"][0]["text"]
if text:
return str(text)
except (KeyError, IndexError, TypeError):
pass
# Secondary: result.data.message.content[0].text
try:
text = result["result"]["data"]["message"]["content"][0]["text"]
if text:
return str(text)
except (KeyError, IndexError, TypeError):
pass
# Tertiary: results as a raw string
results = result.get("results")
if results and isinstance(results, str):
return results
return ""
@staticmethod
def extract_text_from_a2a_message_response(a2a_response: Dict[str, Any]) -> str:
result = a2a_response.get("result")
if not isinstance(result, dict):
verbose_logger.warning("WXO: A2A response missing result object")
return ""
parts = result.get("parts")
if not isinstance(parts, list):
verbose_logger.warning("WXO: A2A result has no parts list")
return ""
for part in parts:
if (
isinstance(part, dict)
and part.get("kind") == "text"
and part.get("text")
):
return str(part["text"])
verbose_logger.warning("WXO: A2A result parts contained no text")
return ""
@staticmethod
def build_a2a_message_response(request_id: str, text: str) -> Dict[str, Any]:
"""
Build a standard A2A non-streaming SendMessageResponse (kind=message).
"""
return {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": text}],
"messageId": str(uuid4()),
},
}
@staticmethod
async def fake_streaming_from_text(
text: str,
request_id: str,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
"""
Emit standard A2A streaming events from a completed text response.
Event sequence:
1. task (kind="task", state="submitted")
2. status-update (kind="status-update", state="working")
3. artifact-update chunks
4. status-update (kind="status-update", state="completed", final=True)
"""
task_id = str(uuid4())
context_id = str(uuid4())
artifact_id = str(uuid4())
# 1. Task submitted
yield {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"id": task_id,
"kind": "task",
"status": {"state": "submitted"},
},
}
# 2. Working
yield {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"final": False,
"kind": "status-update",
"status": {"state": "working"},
"taskId": task_id,
},
}
await asyncio.sleep(delay_ms / 1000.0)
# 3. Artifact chunks (always emit at least one chunk, even for empty text)
text_to_chunk = text or ""
for i in range(0, max(len(text_to_chunk), 1), chunk_size):
chunk_text = text_to_chunk[i : i + chunk_size]
is_last = (i + chunk_size) >= max(len(text_to_chunk), 1)
yield {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"kind": "artifact-update",
"taskId": task_id,
"artifact": {
"artifactId": artifact_id,
"parts": [{"kind": "text", "text": chunk_text}],
},
},
}
if not is_last:
await asyncio.sleep(delay_ms / 1000.0)
# 4. Completed
yield {
"jsonrpc": "2.0",
"id": request_id,
"result": {
"contextId": context_id,
"final": True,
"kind": "status-update",
"status": {"state": "completed"},
"taskId": task_id,
},
}
verbose_logger.debug(
f"WXO: Fake streaming completed for request_id={request_id}"
)

View file

@ -60,6 +60,12 @@ class A2ARequestUtils:
if not isinstance(result, dict):
return ""
# Direct message format (A2A spec): detect by explicit kind tag only.
# The "parts" heuristic is too broad and would match any future result
# type that happens to include a "parts" field.
if result.get("kind") == "message":
return A2ARequestUtils.extract_text_from_message(result)
message = result.get("message", {})
return A2ARequestUtils.extract_text_from_message(message)

View file

@ -75,7 +75,7 @@
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"fine-grained-tool-streaming-2025-05-14": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": null,
"mcp-client-2025-11-20": null,
"mcp-client-2025-04-04": null,
@ -106,7 +106,7 @@
"effort-2025-11-24": "effort-2025-11-24",
"fast-mode-2026-02-01": null,
"files-api-2025-04-14": null,
"fine-grained-tool-streaming-2025-05-14": null,
"fine-grained-tool-streaming-2025-05-14": "fine-grained-tool-streaming-2025-05-14",
"interleaved-thinking-2025-05-14": null,
"mcp-client-2025-11-20": null,
"mcp-client-2025-04-04": null,
@ -129,7 +129,7 @@
"bash_20241022": null,
"bash_20250124": null,
"code-execution-2025-08-25": null,
"compact-2026-01-12": null,
"compact-2026-01-12": "compact-2026-01-12",
"computer-use-2025-01-24": "computer-use-2025-01-24",
"computer-use-2025-11-24": "computer-use-2025-11-24",
"context-1m-2025-08-07": "context-1m-2025-08-07",

View file

@ -10,7 +10,7 @@ This is an __init__.py file to allow the following interface
"""
from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union
from typing import Any, AsyncIterator, Coroutine, Dict, Iterator, List, Optional, Union
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
anthropic_messages as _async_anthropic_messages,
@ -100,8 +100,11 @@ def create(
**kwargs,
) -> Union[
AnthropicMessagesResponse,
Iterator[bytes],
AsyncIterator[Any],
Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any]]],
Coroutine[
Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]
],
]:
"""
Async wrapper for Anthropic's messages API

View file

@ -309,9 +309,13 @@ class Cache:
param_value = kwargs[param]
cache_key += f"{str(param)}: {str(param_value)}"
verbose_logger.debug("\nCreated cache key: %s", cache_key)
hashed_cache_key = Cache._get_hashed_cache_key(cache_key)
hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
verbose_logger.debug(
"\nCreated cache key: %s (source material length: %d)",
hashed_cache_key,
len(cache_key),
)
# Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError
# when kwargs already contains preset_cache_key from upstream callers
kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"}
@ -497,6 +501,34 @@ class Cache:
return cached_response
return cached_result
@staticmethod
def _get_safe_cache_lookup_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
cache_lookup_kwargs: Dict[str, Any] = {}
for prompt_kwarg in ("messages", "input"):
if prompt_kwarg in kwargs:
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
if isinstance(kwargs.get("metadata"), dict):
cache_lookup_kwargs["metadata"] = {}
return cache_lookup_kwargs
@staticmethod
def _update_metadata_from_cache_lookup_kwargs(
original_kwargs: Dict[str, Any], cache_lookup_kwargs: Dict[str, Any]
) -> None:
original_metadata = original_kwargs.get("metadata")
cache_lookup_metadata = cache_lookup_kwargs.get("metadata")
if not isinstance(original_metadata, dict) or not isinstance(
cache_lookup_metadata, dict
):
return
if "semantic-similarity" in cache_lookup_metadata:
original_metadata["semantic-similarity"] = cache_lookup_metadata[
"semantic-similarity"
]
def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
"""
Retrieves the cached result for the given arguments.
@ -511,7 +543,6 @@ class Cache:
try: # never block execution
if self.should_use_cache(**kwargs) is not True:
return
messages = kwargs.get("messages", [])
if "cache_key" in kwargs:
cache_key = kwargs["cache_key"]
else:
@ -523,12 +554,19 @@ class Cache:
or cache_control_args.get("s-max-age")
or float("inf")
)
cache_lookup_kwargs = self._get_safe_cache_lookup_kwargs(kwargs)
if dynamic_cache_object is not None:
cached_result = dynamic_cache_object.get_cache(
cache_key, messages=messages
cache_key, **cache_lookup_kwargs
)
else:
cached_result = self.cache.get_cache(cache_key, messages=messages)
cached_result = self.cache.get_cache(
cache_key, **cache_lookup_kwargs
)
self._update_metadata_from_cache_lookup_kwargs(
original_kwargs=kwargs,
cache_lookup_kwargs=cache_lookup_kwargs,
)
return self._get_cache_logic(
cached_result=cached_result, max_age=max_age
)
@ -549,7 +587,6 @@ class Cache:
if self.should_use_cache(**kwargs) is not True:
return
kwargs.get("messages", [])
if "cache_key" in kwargs:
cache_key = kwargs["cache_key"]
else:
@ -654,6 +691,7 @@ class Cache:
self,
embedding_response: Any,
model: Optional[str],
prompt_tokens: Optional[int] = None,
prompt_tokens_details: Optional[dict] = None,
) -> CachedEmbedding:
"""
@ -666,6 +704,7 @@ class Cache:
"index": embedding_response.get("index"),
"object": embedding_response.get("object"),
"model": model,
"prompt_tokens": prompt_tokens,
"prompt_tokens_details": prompt_tokens_details,
}
elif hasattr(embedding_response, "model_dump"):
@ -675,6 +714,7 @@ class Cache:
"index": data.get("index"),
"object": data.get("object"),
"model": model,
"prompt_tokens": prompt_tokens,
"prompt_tokens_details": prompt_tokens_details,
}
else:
@ -684,6 +724,7 @@ class Cache:
"index": data.get("index"),
"object": data.get("object"),
"model": model,
"prompt_tokens": prompt_tokens,
"prompt_tokens_details": prompt_tokens_details,
}
except KeyError as e:
@ -732,6 +773,29 @@ class Cache:
per_item[key] = value
return per_item if per_item else None
def _get_per_item_prompt_tokens(
self,
result: EmbeddingResponse,
idx_in_result_data: int,
) -> Optional[int]:
"""
Extract the per-item prompt_tokens from a response for caching.
Single-item responses store the full usage.prompt_tokens. Multi-item
responses distribute it evenly (with remainder) so that summing all
per-item values on retrieval reconstructs the original total.
"""
if result.usage is None or result.usage.prompt_tokens is None:
return None
total = result.usage.prompt_tokens
num_items = len(result.data)
if num_items <= 1:
return total
quotient, remainder = divmod(total, num_items)
return quotient + (1 if idx_in_result_data < remainder else 0)
def add_embedding_response_to_cache(
self,
result: EmbeddingResponse,
@ -743,7 +807,11 @@ class Cache:
kwargs["cache_key"] = preset_cache_key
embedding_response = result.data[idx_in_result_data]
# Extract per-item prompt_tokens_details from response usage
# Extract per-item prompt_tokens + details from response usage
prompt_tokens = self._get_per_item_prompt_tokens(
result=result,
idx_in_result_data=idx_in_result_data,
)
prompt_tokens_details = self._get_per_item_prompt_tokens_details(
result=result,
idx_in_result_data=idx_in_result_data,
@ -754,6 +822,7 @@ class Cache:
embedding_dict: CachedEmbedding = self._convert_to_cached_embedding(
embedding_response,
model_name,
prompt_tokens=prompt_tokens,
prompt_tokens_details=prompt_tokens_details,
)

View file

@ -394,7 +394,7 @@ class LLMCachingHandler:
return cr["model"]
return None
def _process_async_embedding_cached_response(
def _process_async_embedding_cached_response( # noqa: PLR0915
self,
final_embedding_cached_response: Optional[EmbeddingResponse],
cached_result: List[Optional[CachedEmbedding]],
@ -456,7 +456,10 @@ class LLMCachingHandler:
index=idx,
object="embedding",
)
if isinstance(kwargs_input_as_list[idx], str):
cached_prompt_tokens = cr.get("prompt_tokens")
if cached_prompt_tokens is not None:
prompt_tokens += cached_prompt_tokens
elif isinstance(kwargs_input_as_list[idx], str):
from litellm.utils import token_counter
prompt_tokens += token_counter(

View file

@ -22,6 +22,7 @@ import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import (
DEFAULT_REDIS_MAJOR_VERSION,
REDIS_CIRCUIT_BREAKER_ENABLED,
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
)
@ -114,15 +115,23 @@ class RedisCircuitBreaker:
OPEN = "open"
HALF_OPEN = "half_open"
def __init__(self, failure_threshold: int, recovery_timeout: int) -> None:
def __init__(
self,
failure_threshold: int,
recovery_timeout: int,
enabled: bool = True,
) -> None:
self.failure_threshold = failure_threshold
self.recovery_timeout = recovery_timeout
self.enabled = enabled
self._failure_count = 0
self._opened_at: Optional[float] = None
self._state = self.CLOSED
def is_open(self) -> bool:
"""Returns True if Redis calls should be skipped."""
if not self.enabled:
return False
if self._state == self.HALF_OPEN:
# Probe already in flight — fast-fail all concurrent requests.
# Only the one call that caused the OPEN→HALF_OPEN transition
@ -136,6 +145,8 @@ class RedisCircuitBreaker:
return False
def record_failure(self) -> None:
if not self.enabled:
return
self._failure_count += 1
self._opened_at = time.time()
if self._failure_count >= self.failure_threshold:
@ -149,6 +160,8 @@ class RedisCircuitBreaker:
self._state = self.OPEN
def record_success(self) -> None:
if not self.enabled:
return
if self._state == self.HALF_OPEN:
verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered")
self._failure_count = 0
@ -243,6 +256,7 @@ class RedisCache(BaseCache):
self._circuit_breaker = RedisCircuitBreaker(
failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
enabled=REDIS_CIRCUIT_BREAKER_ENABLED,
)
self._setup_health_pings()

View file

@ -213,6 +213,78 @@ class RedisSemanticCache(BaseCache):
ttl = int(ttl)
return ttl
@classmethod
def _get_prompt_from_kwargs(cls, **kwargs) -> Optional[str]:
"""
Extract a semantic-cache prompt from chat or Responses API request kwargs.
"""
messages = kwargs.get("messages")
if messages:
return get_str_from_messages(messages)
if "input" not in kwargs:
return None
prompt_parts: List[str] = []
cls._collect_responses_input_text(kwargs.get("input"), prompt_parts)
prompt = "\n".join(prompt_parts).strip()
return prompt or None
@classmethod
def _collect_responses_input_text(cls, value: Any, prompt_parts: List[str]) -> None:
value = cls._coerce_response_input_value(value)
if value is None:
return
if isinstance(value, str):
stripped_value = value.strip()
if stripped_value:
prompt_parts.append(stripped_value)
return
if isinstance(value, (list, tuple)):
for item in value:
cls._collect_responses_input_text(item, prompt_parts)
return
if isinstance(value, dict):
content = value.get("content")
if content is not None:
cls._collect_responses_input_text(content, prompt_parts)
return
for text_key in ("text", "output", "input_text", "output_text"):
text_value = value.get(text_key)
if isinstance(text_value, str):
stripped_text = text_value.strip()
if stripped_text:
prompt_parts.append(stripped_text)
return
return
content = getattr(value, "content", None)
if content is not None:
cls._collect_responses_input_text(content, prompt_parts)
return
for text_key in ("text", "output", "input_text", "output_text"):
text_value = getattr(value, text_key, None)
if isinstance(text_value, str):
stripped_text = text_value.strip()
if stripped_text:
prompt_parts.append(stripped_text)
return
@staticmethod
def _coerce_response_input_value(value: Any) -> Any:
model_dump = getattr(value, "model_dump", None)
if callable(model_dump):
return model_dump()
dict_method = getattr(value, "dict", None)
if callable(dict_method):
return dict_method()
return value
def _get_embedding(self, prompt: str) -> List[float]:
"""
Generate an embedding vector for the given prompt using the configured embedding model.
@ -278,13 +350,11 @@ class RedisSemanticCache(BaseCache):
value_str: Optional[str] = None
try:
# Extract the prompt from messages
messages = kwargs.get("messages", [])
if not messages:
print_verbose("No messages provided for semantic caching")
prompt = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
print_verbose("No prompt provided for semantic caching")
return
prompt = get_str_from_messages(messages)
value_str = str(value)
store_kwargs: Dict[str, Any] = {
@ -315,14 +385,12 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Redis semantic-cache get_cache, kwargs: {kwargs}")
try:
# Extract the prompt from messages
messages = kwargs.get("messages", [])
if not messages:
print_verbose("No messages provided for semantic cache lookup")
prompt = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
print_verbose("No prompt provided for semantic cache lookup")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
prompt = get_str_from_messages(messages)
# Check the cache for semantically similar prompts in this exact
# LiteLLM cache-key scope.
check_kwargs: Dict[str, Any] = {
@ -428,13 +496,11 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Async Redis semantic-cache set_cache, kwargs: {kwargs}")
try:
# Extract the prompt from messages
messages = kwargs.get("messages", [])
if not messages:
print_verbose("No messages provided for semantic caching")
prompt = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
print_verbose("No prompt provided for semantic caching")
return
prompt = get_str_from_messages(messages)
value_str = str(value)
# Generate embedding for the value (response) to cache
@ -471,15 +537,12 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Async Redis semantic-cache get_cache, kwargs: {kwargs}")
try:
# Extract the prompt from messages
messages = kwargs.get("messages", [])
if not messages:
print_verbose("No messages provided for semantic cache lookup")
prompt = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
print_verbose("No prompt provided for semantic cache lookup")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
return None
prompt = get_str_from_messages(messages)
# Generate embedding for the prompt
prompt_embedding = await self._get_async_embedding(prompt, **kwargs)

View file

@ -171,6 +171,8 @@ class ResponsesToCompletionBridgeHandler:
model_response = validated_kwargs["model_response"]
logging_obj = validated_kwargs["logging_obj"]
custom_llm_provider = validated_kwargs["custom_llm_provider"]
if kwargs.get("stream") is True and "stream" not in optional_params:
optional_params = {**optional_params, "stream": True}
request_data = self.transformation_handler.transform_request(
model=model,
@ -182,6 +184,14 @@ class ResponsesToCompletionBridgeHandler:
client=kwargs.get("client"),
)
# Pin the resolved provider so `responses()` doesn't re-run
# `get_llm_provider()` on the model string and strip a second
# provider prefix (see GitHub issue #28505). request_data already
# carries `custom_llm_provider` via the spread of
# `sanitized_litellm_params`; overwriting it on the dict (rather
# than adding an explicit kwarg) avoids the duplicate-keyword
# TypeError that would otherwise fire on the real bridge path.
request_data["custom_llm_provider"] = custom_llm_provider
result = responses(
**request_data,
)
@ -255,6 +265,8 @@ class ResponsesToCompletionBridgeHandler:
model_response = validated_kwargs["model_response"]
logging_obj = validated_kwargs["logging_obj"]
custom_llm_provider = validated_kwargs["custom_llm_provider"]
if kwargs.get("stream") is True and "stream" not in optional_params:
optional_params = {**optional_params, "stream": True}
try:
request_data = self.transformation_handler.transform_request(
@ -268,6 +280,13 @@ class ResponsesToCompletionBridgeHandler:
except Exception as e:
raise e
# Pin the resolved provider so `aresponses()` doesn't re-run
# `get_llm_provider()` on the model string and strip a second
# provider prefix (see GitHub issue #28505). Set on request_data
# rather than passed as a separate kwarg to avoid the duplicate-
# keyword TypeError when `sanitized_litellm_params` already
# carries `custom_llm_provider`.
request_data["custom_llm_provider"] = custom_llm_provider
result = await aresponses(
**request_data,
aresponses=True,

View file

@ -402,6 +402,20 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
instructions,
) = self.convert_chat_completion_messages_to_responses_api(messages)
# OpenAI's Responses API rejects an empty input. For a system-only
# request, carry the system message as a system-role input item instead
# of instructions, mirroring how non-string system content is already
# handled in convert_chat_completion_messages_to_responses_api.
if not input_items and instructions is not None:
input_items = [
{
"type": "message",
"role": "system",
"content": [{"type": "input_text", "text": instructions}],
}
]
instructions = None
optional_params = self._extract_extra_body_params(optional_params)
# Build responses API request using the reverse transformation logic

View file

@ -398,6 +398,9 @@ REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD = int(
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT = int(
os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)
)
REDIS_CIRCUIT_BREAKER_ENABLED = (
os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true"
)
# Default Redis major version to assume when version cannot be determined
# Using 7 as it's the modern version that supports LPOP with count parameter
DEFAULT_REDIS_MAJOR_VERSION = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7))
@ -418,6 +421,7 @@ REPLICATE_POLLING_DELAY_SECONDS = float(
DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS = int(
os.getenv("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096)
)
DEFAULT_OCI_CHAT_MAX_TOKENS = 4096
TOGETHER_AI_4_B = int(os.getenv("TOGETHER_AI_4_B", 4))
TOGETHER_AI_8_B = int(os.getenv("TOGETHER_AI_8_B", 8))
TOGETHER_AI_21_B = int(os.getenv("TOGETHER_AI_21_B", 21))
@ -585,6 +589,7 @@ LITELLM_CHAT_PROVIDERS = [
"volcengine",
"codestral",
"text-completion-codestral",
"text-completion-inception",
"deepseek",
"sambanova",
"maritalk",
@ -620,6 +625,7 @@ LITELLM_CHAT_PROVIDERS = [
"oci",
"morph",
"lambda_ai",
"inception",
"vercel_ai_gateway",
"wandb",
"ovhcloud",
@ -676,6 +682,7 @@ OPENAI_CHAT_COMPLETION_PARAMS = [
"extra_headers",
"thinking",
"web_search_options",
"include_server_side_tool_invocations",
"service_tier",
"prompt_cache_key",
"prompt_cache_retention",
@ -737,6 +744,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES = {
"verbosity": None,
"thinking": None,
"web_search_options": None,
"include_server_side_tool_invocations": None,
"service_tier": None,
"safety_identifier": None,
"prompt_cache_key": None,
@ -771,6 +779,7 @@ openai_compatible_endpoints: List = [
"https://api.moonshot.ai/v1",
"https://api.publicai.co/v1",
"https://api.synthetic.new/openai/v1",
"https://serverless.tensormesh.ai/v1",
"https://api.stima.tech/v1",
"https://nano-gpt.com/api/v1",
"https://api.poe.com/v1",
@ -778,6 +787,7 @@ openai_compatible_endpoints: List = [
"https://api.v0.dev/v1",
"https://api.morphllm.com/v1",
"https://api.lambda.ai/v1",
"https://api.inceptionlabs.ai/v1",
"https://api.hyperbolic.xyz/v1",
"https://ai-gateway.helicone.ai/",
"https://ai-gateway.vercel.sh/v1",
@ -820,10 +830,12 @@ openai_compatible_providers: List = [
"meta_llama",
"publicai", # PublicAI - JSON-configured provider
"synthetic", # Synthetic - JSON-configured provider
"tensormesh", # Tensormesh - JSON-configured provider
"apertis", # Apertis - JSON-configured provider
"nano-gpt", # Nano-GPT - JSON-configured provider
"poe", # Poe - JSON-configured provider
"chutes", # Chutes - JSON-configured provider
"parasail", # Parasail - JSON-configured provider
"featherless_ai",
"nscale",
"nebius",
@ -833,6 +845,7 @@ openai_compatible_providers: List = [
"helicone",
"morph",
"lambda_ai",
"inception",
"hyperbolic",
"vercel_ai_gateway",
"aiml",
@ -855,6 +868,7 @@ openai_text_completion_compatible_providers: List = (
"moonshot",
"publicai",
"synthetic",
"tensormesh",
"apertis",
"nano-gpt",
"poe",
@ -868,6 +882,7 @@ openai_text_completion_compatible_providers: List = (
_openai_like_providers: List = [
"predibase",
"databricks",
"lemonade",
"watsonx",
] # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk
# well supported replicate llms
@ -1147,6 +1162,8 @@ BEDROCK_CONVERSE_MODELS = [
"openai.gpt-oss-120b-1:0",
"anthropic.claude-haiku-4-5-20251001-v1:0",
"anthropic.claude-sonnet-4-5-20250929-v1:0",
"anthropic.claude-fable-5",
"anthropic.claude-opus-4-8",
"anthropic.claude-opus-4-7",
"anthropic.claude-opus-4-6-v1:0",
"anthropic.claude-opus-4-6-v1",
@ -1408,6 +1425,13 @@ LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs"
# Prometheus metrics, audit trails, or any other downstream consumer.
LITELLM_PROXY_MASTER_KEY_ALIAS = "litellm_proxy_master_key"
# Marker placed in ``model_call_details`` on a synthetic ``Logging`` object that
# records a proxy-gate error (auth/rate-limit rejection) for a request that never
# reached an upstream provider. Tracing callbacks key off it to avoid fabricating
# an LLM-call span for a call that did not happen. See
# ``ProxyLogging._handle_logging_proxy_only_error``.
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL = "litellm_no_upstream_llm_call"
# Key Rotation Constants
LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false")
LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int(
@ -1460,6 +1484,7 @@ DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job"
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job"
PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics"
CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data"
MAVVRIK_FOCUS_EXPORT_JOB_NAME = "mavvrik_focus_export_usage_data"
CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)
)
@ -1474,6 +1499,10 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float(
os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
)
SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(
os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)
)
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int(

View file

@ -133,6 +133,8 @@ _VIDEO_CALL_TYPES = frozenset(
{
CallTypes.create_video.value,
CallTypes.acreate_video.value,
CallTypes.video_edit.value,
CallTypes.avideo_edit.value,
CallTypes.video_remix.value,
CallTypes.avideo_remix.value,
}
@ -417,9 +419,36 @@ def cost_per_token( # noqa: PLR0915
prompt_tokens_cost_usd_dollar: float = 0
completion_tokens_cost_usd_dollar: float = 0
model_cost_ref = litellm.model_cost
# Only callers that explicitly pass `custom_llm_provider` get the
# dedup/prefix-join treatment. When provider is omitted, preserve legacy
# behavior: `model_with_provider` stays equal to the raw `model` string
# (provider is detected below for downstream use only).
caller_supplied_provider = custom_llm_provider is not None
# `model` is normally a string, but callers that mock the transport can pass
# non-string objects. Only run the string-based dedup/prefix-join when it is
# actually a string — e.g. a MagicMock's `.startswith()` is always truthy and
# its slices return new mocks, which would spin the dedup loop forever.
model_is_str = isinstance(model, str)
# Router/proxy deployments may repeat the provider segment (e.g. model_name
# "openai/openai/gpt-5.5"). Strip duplicated `{provider}/` chains before joining.
if caller_supplied_provider and model_is_str:
_dup_prefix = f"{custom_llm_provider}/"
while model.startswith(_dup_prefix):
_remainder = model[len(_dup_prefix) :]
if _remainder.startswith(_dup_prefix):
model = _remainder
else:
break
model_with_provider = model
if custom_llm_provider is not None:
model_with_provider = custom_llm_provider + "/" + model
if caller_supplied_provider:
_prov_prefix = f"{custom_llm_provider}/"
if model_is_str and model.startswith(_prov_prefix):
model_with_provider = model
else:
model_with_provider = f"{custom_llm_provider}/{model}"
if region_name is not None:
model_with_provider_and_region = (
f"{custom_llm_provider}/{region_name}/{model}"
@ -430,6 +459,9 @@ def cost_per_token( # noqa: PLR0915
model_with_provider = model_with_provider_and_region
else:
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model)
assert custom_llm_provider is not None # caller-supplied or get_llm_provider
model_without_prefix = model
model_parts = model.split("/", 1)
if len(model_parts) > 1:
@ -2393,12 +2425,11 @@ class BaseTokenUsageProcessor:
if not attr.startswith("_") and not callable(
getattr(usage.completion_tokens_details, attr)
):
current_val = getattr(
combined.completion_tokens_details, attr, 0
current_val = (
getattr(combined.completion_tokens_details, attr, 0) or 0
)
new_val = getattr(usage.completion_tokens_details, attr, 0)
if new_val is not None and current_val is not None:
new_val = getattr(usage.completion_tokens_details, attr, 0) or 0
if isinstance(new_val, (int, float)):
setattr(
combined.completion_tokens_details,
attr,
@ -2457,6 +2488,11 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
)
_TRANSCRIPTION_COMPLETED_EVENT_TYPE = (
"conversation.item.input_audio_transcription.completed"
)
def handle_realtime_stream_cost_calculation(
results: OpenAIRealtimeStreamList,
combined_usage_object: Usage,
@ -2502,4 +2538,99 @@ def handle_realtime_stream_cost_calculation(
break # exit if we find a valid model
total_cost = input_cost_per_token + output_cost_per_token
if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results):
total_cost += handle_realtime_transcription_cost_calculation(
results=results,
custom_llm_provider=custom_llm_provider,
litellm_model_name=litellm_model_name,
)
return total_cost
def handle_realtime_transcription_cost_calculation(
results: OpenAIRealtimeStreamList,
custom_llm_provider: str,
litellm_model_name: str,
) -> float:
"""
Cost for realtime transcription sessions (e.g. gpt-realtime-whisper).
Transcription sessions emit no `response.done` events; instead each
`conversation.item.input_audio_transcription.completed` event carries a
`usage` object billed by the ASR model. The usage is one of:
- {"type": "duration", "seconds": <float>} → priced via input_cost_per_second
- {"type": "tokens", "input_tokens": ...} → priced via input/audio token cost
"""
completed_events = [
cast(dict, result)
for result in results
if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE
]
if not completed_events:
return 0.0
model_name = (
_get_transcription_model_name_from_results(results) or litellm_model_name
)
try:
model_info = litellm.get_model_info(
model=model_name, custom_llm_provider=custom_llm_provider
)
except Exception:
model_info = None
total_cost = 0.0
for event in completed_events:
usage = event.get("usage") or {}
total_cost += _transcription_usage_cost(usage, model_info)
return total_cost
def _get_transcription_model_name_from_results(
results: OpenAIRealtimeStreamList,
) -> Optional[str]:
"""Resolve the ASR model from a transcription_session.* / session.* event."""
for result in results:
if result.get("type") in (
"transcription_session.created",
"transcription_session.updated",
"session.created",
"session.updated",
):
session = cast(dict, result).get("session", {}) or {}
transcription = (
(session.get("audio", {}) or {}).get("input", {}) or {}
).get("transcription", {}) or session.get("input_audio_transcription", {})
model = (transcription or {}).get("model") or session.get("model")
if model:
return model
return None
def _transcription_usage_cost(usage: dict, model_info: Optional[ModelInfo]) -> float:
if model_info is None:
return 0.0
usage_type = usage.get("type")
if usage_type == "duration":
seconds = usage.get("seconds") or 0.0
per_second = model_info.get("input_cost_per_second") or 0.0
return float(seconds) * float(per_second)
if usage_type == "tokens":
input_token_details = usage.get("input_token_details") or {}
audio_tokens = input_token_details.get("audio_tokens") or 0
text_tokens = input_token_details.get("text_tokens") or 0
output_tokens = usage.get("output_tokens") or 0
audio_cost = float(audio_tokens) * float(
model_info.get("input_cost_per_audio_token")
or model_info.get("input_cost_per_token")
or 0.0
)
text_cost = float(text_tokens) * float(
model_info.get("input_cost_per_token") or 0.0
)
output_cost = float(output_tokens) * float(
model_info.get("output_cost_per_token") or 0.0
)
return audio_cost + text_cost + output_cost
return 0.0

View file

@ -9,13 +9,109 @@
## LiteLLM versions of the OpenAI Exception Types
from typing import Any, Dict, Optional
import enum
from typing import Any, Dict, Optional, Union
import httpx
import openai
from litellm.types.utils import LiteLLMCommonStrings
class RateLimitErrorCategory(str, enum.Enum):
"""
Category of a rate limit error, allowing callers to distinguish where the rate
limit originated. Exposed on every :class:`RateLimitError` instance via the
``category`` attribute.
Use these values to switch on the rate limit source, e.g.::
try:
...
except litellm.RateLimitError as e:
if e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT:
... # litellm's own limiter (key/team/user/model RPM/TPM/budget)
elif e.category == RateLimitErrorCategory.VENDOR_RATE_LIMIT:
... # the upstream LLM provider returned 429
"""
VENDOR_RATE_LIMIT = "vendor_rate_limit"
"""The upstream LLM provider returned a rate-limit response (e.g. OpenAI 429)."""
VENDOR_BATCH_RATE_LIMIT = "vendor_batch_rate_limit"
"""The upstream LLM provider returned a rate-limit response on a batch endpoint."""
LITELLM_RATE_LIMIT = "litellm_rate_limit"
"""LiteLLM's own rate limiter (key/team/user/model RPM/TPM, budget, parallel-requests, etc.) blocked the request."""
LITELLM_BATCH_RATE_LIMIT = "litellm_batch_rate_limit"
"""LiteLLM's own batch rate limiter (token/request budget across a batch input file) blocked the request."""
class RateLimitType(str, enum.Enum):
"""
The dimension that was exceeded when a rate-limit error fired.
This is orthogonal to :class:`RateLimitErrorCategory` — *category* tells
callers **who** rate-limited the request (the upstream vendor vs. one of
litellm's own limiters), while *type* tells them **which limit dimension**
was exceeded (an RPM ceiling, a TPM ceiling, a max-parallel-requests
ceiling, a budget cap, or a max-iterations cap).
Surfaced both on every :class:`RateLimitError` instance via the
``rate_limit_type`` attribute and on the structured
``StandardLoggingPayload.error_information.error_rate_limit_type`` field
so custom callbacks / metrics consumers can split rate-limit failures by
cause without parsing free-text error messages.
"""
REQUESTS = "requests"
"""Requests-per-minute (RPM) or requests-per-window ceiling exceeded."""
TOKENS = "tokens"
"""Tokens-per-minute (TPM) or tokens-per-window ceiling exceeded."""
CONCURRENT_REQUESTS = "concurrent_requests"
"""``max_parallel_requests`` — too many in-flight requests at once."""
BUDGET = "budget"
"""Spend budget cap reached (key, team, user, or per-session)."""
MAX_ITERATIONS = "max_iterations"
"""Per-session max-iterations cap reached (agent-style flows)."""
_RATE_LIMIT_CATEGORY_VALUES = frozenset(c.value for c in RateLimitErrorCategory)
_RATE_LIMIT_TYPE_VALUES = frozenset(t.value for t in RateLimitType)
def validate_rate_limit_category(value: Any) -> Optional[str]:
"""Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`.
Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus
labels) to reject `.category` strings set by unrelated third-party exceptions
— otherwise those would leak into custom-callback payloads and Prometheus
label cardinality.
"""
if isinstance(value, RateLimitErrorCategory):
return value.value
if isinstance(value, str) and value in _RATE_LIMIT_CATEGORY_VALUES:
return value
return None
def validate_rate_limit_type(value: Any) -> Optional[str]:
"""Return ``value`` only if it matches a known :class:`RateLimitType`.
See :func:`validate_rate_limit_category` for the rationale.
"""
if isinstance(value, RateLimitType):
return value.value
if isinstance(value, str) and value in _RATE_LIMIT_TYPE_VALUES:
return value
return None
_MINIMAL_ERROR_RESPONSE: Optional[httpx.Response] = None
@ -321,6 +417,18 @@ class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore
class RateLimitError(openai.RateLimitError): # type: ignore
"""
Unified rate-limit error.
Every rate-limit condition surfaced by litellm — whether it originated from
an upstream LLM provider, a vendor batch endpoint, or one of litellm's own
proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,
max-iterations, etc.) — is raised as an instance of this class.
The :attr:`category` attribute lets callers distinguish the source. See
:class:`RateLimitErrorCategory` for the available values.
"""
def __init__(
self,
message,
@ -330,6 +438,12 @@ class RateLimitError(openai.RateLimitError): # type: ignore
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
category: Union[str, RateLimitErrorCategory] = (
RateLimitErrorCategory.VENDOR_RATE_LIMIT
),
rate_limit_type: Optional[Union[str, RateLimitType]] = None,
headers: Optional[Dict[str, str]] = None,
detail: Any = None,
):
self.status_code = 429
self.message = "litellm.RateLimitError: {}".format(message)
@ -338,9 +452,39 @@ class RateLimitError(openai.RateLimitError): # type: ignore
self.litellm_debug_info = litellm_debug_info
self.max_retries = max_retries
self.num_retries = num_retries
self.category = (
category.value if isinstance(category, RateLimitErrorCategory) else category
)
# Which dimension was exceeded — request count, token count, parallel
# requests, budget, max iterations. None when the source didn't
# classify the failure (e.g. legacy vendor 429 with no header hints).
self.rate_limit_type: Optional[str] = (
rate_limit_type.value
if isinstance(rate_limit_type, RateLimitType)
else rate_limit_type
)
# Headers explicitly attached to the error (e.g. retry-after,
# rate_limit_type, reset_at). Preserved across the proxy boundary so
# clients can react appropriately.
#
# IMPORTANT: we deliberately do NOT auto-populate self.headers from
# response.headers when only `response` is provided. A vendor 429 can
# set arbitrary response headers (Set-Cookie, CORS overrides, …); if
# those leaked into e.headers and a downstream proxy serializer
# forwarded them to the client, a malicious upstream could inject
# browser-interpreted headers for the proxy origin. Vendor response
# headers stay reachable on `e.response.headers` for callers that
# explicitly want them; only the proxy-supplied `headers=` kwarg
# makes it onto `self.headers`.
_response_headers = (
getattr(response, "headers", None) if response is not None else None
)
self.headers: Optional[Dict[str, str]] = (
{k: str(v) for k, v in headers.items()} if headers else None
)
# Mirrors FastAPI HTTPException.detail so the same instance can be
# serialized through both the ProxyException and HTTPException paths.
self.detail = detail if detail is not None else self.message
self.response = httpx.Response(
status_code=429,
headers=_response_headers,
@ -843,11 +987,24 @@ LITELLM_EXCEPTION_TYPES = [
class BudgetExceededError(Exception):
def __init__(
self, current_cost: float, max_budget: float, message: Optional[str] = None
self,
current_cost: float,
max_budget: float,
message: Optional[str] = None,
llm_provider: Optional[str] = None,
):
self.current_cost = current_cost
self.max_budget = max_budget
self.status_code = 429
self.llm_provider = llm_provider or ""
# Surface unified rate-limit fields without joining the RateLimitError
# hierarchy so existing `except BudgetExceededError:` handlers keep
# working; custom callbacks reading StandardLoggingPayload pick these
# up via the same `category` / `rate_limit_type` attributes the rest
# of the unified rate-limit error path uses. Stored as plain strings
# to match the normalization RateLimitError.__init__ performs.
self.category: str = RateLimitErrorCategory.LITELLM_RATE_LIMIT.value
self.rate_limit_type: str = RateLimitType.BUDGET.value
message = (
message
or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}"
@ -1062,3 +1219,37 @@ class GuardrailInterventionNormalStringError(
def __repr__(self):
return self.__str__()
class SensitiveDataRouteException(Exception):
"""
Exception raised when a guardrail detects sensitive data and wants to reroute the request.
Instead of blocking the request, this exception signals that the request should be
routed to a different model (typically an on-premise model for data privacy).
The proxy catches this exception and:
1. Reroutes the current request to the specified model
2. When sticky_session_routing is True, stores the routing decision in session
cache so all subsequent requests in the same session are routed to the same model
"""
def __init__(
self,
route_to_model: str,
session_id: str,
guardrail_name: Optional[str] = None,
detection_info: Optional[Dict[str, Any]] = None,
message: Optional[str] = None,
sticky_session_routing: bool = True,
):
self.route_to_model = route_to_model
self.session_id = session_id
self.guardrail_name = guardrail_name
self.detection_info = detection_info or {}
self.sticky_session_routing = sticky_session_routing
self.message = (
message
or f"Sensitive data detected by {guardrail_name}. Routing to model: {route_to_model}"
)
super().__init__(self.message)

View file

@ -4,6 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
import os
from typing import (
Any,
Awaitable,
@ -16,7 +17,6 @@ from typing import (
TypeVar,
Union,
)
import httpx
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
from mcp.client.sse import sse_client
@ -42,9 +42,8 @@ from mcp.types import (
)
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.types.llms.custom_http import VerifyTypes
from litellm.types.mcp import (
@ -61,13 +60,33 @@ def to_basic_auth(auth_value: str) -> str:
return base64.b64encode(auth_value.encode("utf-8")).decode()
def _strip_header_whitespace(headers: Dict[str, str]) -> Dict[str, str]:
return {
(key.strip() if isinstance(key, str) else key): (
value.strip() if isinstance(value, str) else value
)
for key, value in headers.items()
}
def _first_non_cancelled_cause(exc: BaseException) -> Optional[BaseException]:
queue: List[BaseException] = [exc]
while queue:
current = queue.pop(0)
nested = getattr(current, "exceptions", None)
if nested:
queue.extend(nested)
elif not isinstance(current, asyncio.CancelledError):
return current
return None
TSessionResult = TypeVar("TSessionResult")
class MCPSigV4Auth(httpx.Auth):
"""
httpx Auth class that signs each request with AWS SigV4.
This is used for MCP servers that require AWS SigV4 authentication,
such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow()
for every outgoing request, enabling per-request signature computation.
@ -92,10 +111,8 @@ class MCPSigV4Auth(httpx.Auth):
"Missing botocore to use AWS SigV4 authentication. "
"Run 'pip install boto3'."
)
self.service_name = aws_service_name or "bedrock-agentcore"
self.region_name = aws_region_name or "us-east-1"
# Note: os.environ/ prefixed values are already resolved by
# ProxyConfig._check_for_os_environ_vars() at config load time.
# Values arrive here as plain strings.
@ -143,20 +160,17 @@ class MCPSigV4Auth(httpx.Auth):
session_name = (
aws_session_name or f"litellm-mcp-{int(__import__('time').time())}"
)
sts_kwargs: dict = {"region_name": aws_region_name}
if aws_access_key_id and aws_secret_access_key:
sts_kwargs["aws_access_key_id"] = aws_access_key_id
sts_kwargs["aws_secret_access_key"] = aws_secret_access_key
if aws_session_token:
sts_kwargs["aws_session_token"] = aws_session_token
sts_client = boto3.client("sts", **sts_kwargs)
sts_response = sts_client.assume_role(
RoleArn=aws_role_name,
RoleSessionName=session_name,
)
sts_creds = sts_response["Credentials"]
return Credentials(
access_key=sts_creds["AccessKeyId"],
@ -178,17 +192,14 @@ class MCPSigV4Auth(httpx.Auth):
data=request.content,
headers=dict(request.headers),
)
# Sign the request — SigV4Auth.add_auth() adds Authorization,
# X-Amz-Date, and X-Amz-Security-Token (if session token present).
# Host header is derived automatically from the URL.
sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name)
sigv4.add_auth(aws_request)
# Copy SigV4 headers back to the httpx request
for header_name, header_value in aws_request.headers.items():
request.headers[header_name] = header_value
yield request
@ -198,6 +209,8 @@ class MCPClient:
SSE and HTTP transports
Authentication via Bearer token, Basic Auth, or API Key
Tool calling with error handling and result parsing
Sampling callbacks for upstream server LLM requests
Elicitation callbacks for upstream server user-input requests
"""
def __init__(
@ -211,6 +224,9 @@ class MCPClient:
extra_headers: Optional[Dict[str, str]] = None,
ssl_verify: Optional[VerifyTypes] = None,
aws_auth: Optional[httpx.Auth] = None,
sampling_callback: Optional[Callable] = None,
elicitation_callback: Optional[Callable] = None,
logging_callback: Optional[Callable] = None,
):
self.server_url: str = server_url
self.transport_type: MCPTransport = transport_type
@ -222,6 +238,9 @@ class MCPClient:
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
self._aws_auth: Optional[httpx.Auth] = aws_auth
self._last_initialize_instructions: Optional[str] = None
self._sampling_callback: Optional[Callable] = sampling_callback
self._elicitation_callback: Optional[Callable] = elicitation_callback
self._logging_callback: Optional[Callable] = logging_callback
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
@ -231,23 +250,20 @@ class MCPClient:
) -> Tuple[Any, Optional[httpx.AsyncClient]]:
"""
Create the appropriate transport context based on transport type.
Returns:
Tuple of (transport_context, http_client).
http_client is only set for HTTP transport and needs cleanup.
"""
http_client: Optional[httpx.AsyncClient] = None
if self.transport_type == MCPTransport.stdio:
if not self.stdio_config:
raise ValueError("stdio_config is required for stdio transport")
server_params = StdioServerParameters(
command=self.stdio_config.get("command", ""),
args=self.stdio_config.get("args", []),
env=self.stdio_config.get("env", {}),
env=self._get_safe_stdio_env(self.stdio_config.get("env")),
)
return stdio_client(server_params), None
if self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
@ -260,14 +276,12 @@ class MCPClient:
),
None,
)
# HTTP transport (default)
if streamable_http_client is None:
raise ImportError(
"streamable_http_client is not available. "
"Please install mcp with HTTP support."
)
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
@ -281,6 +295,54 @@ class MCPClient:
)
return transport_ctx, http_client
def _get_safe_stdio_env(
self, provided_env: Optional[Dict[str, str]]
) -> Optional[Dict[str, str]]:
"""
Return a safe environment for the stdio subprocess.
If provided_env is set, we use it as-is.
If provided_env is None, we return a minimal allowlist from the parent environment
to avoid leaking sensitive LiteLLM keys (OPENAI_API_KEY, etc.) to sub-processes.
"""
if provided_env is not None:
return provided_env
# Minimal allowlist of safe/standard environment variables
safe_keys = {
"PATH",
"HOME",
"USER",
"LOGNAME",
"TMPDIR",
"TMP",
"TEMP",
"SHELL",
"LANG",
"LC_ALL",
# Node/Package manager caches
"NPM_CONFIG_CACHE",
"PNPM_HOME",
"XDG_CACHE_HOME",
"XDG_CONFIG_HOME",
"XDG_DATA_HOME",
# System info
"SYSTEMROOT",
"COMSPEC",
"PATHEXT",
"WINDIR",
}
safe_env = {}
for key in safe_keys:
if key in os.environ:
safe_env[key] = os.environ[key]
if "NPM_CONFIG_CACHE" not in safe_env:
safe_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
return safe_env
async def _execute_session_operation(
self,
transport_ctx: Any,
@ -288,13 +350,24 @@ class MCPClient:
) -> TSessionResult:
"""
Execute an operation within a transport and session context.
Handles entering/exiting contexts and running the operation.
Passes sampling/elicitation/logging callbacks to the ClientSession
so that upstream MCP servers can request LLM inference (sampling),
user input (elicitation), or send log messages.
"""
transport = await transport_ctx.__aenter__()
in_flight_error: Optional[BaseException] = None
try:
read_stream, write_stream = transport[0], transport[1]
session_ctx = ClientSession(read_stream, write_stream)
# Build session kwargs with optional callbacks
session_kwargs: Dict[str, Any] = {}
if self._sampling_callback is not None:
session_kwargs["sampling_callback"] = self._sampling_callback
if self._elicitation_callback is not None:
session_kwargs["elicitation_callback"] = self._elicitation_callback
if self._logging_callback is not None:
session_kwargs["logging_callback"] = self._logging_callback
session_ctx = ClientSession(read_stream, write_stream, **session_kwargs)
session = await session_ctx.__aenter__()
try:
init_result = await session.initialize()
@ -309,11 +382,21 @@ class MCPClient:
await session_ctx.__aexit__(None, None, None)
except BaseException as e:
verbose_logger.debug(f"Error during session context exit: {e}")
except BaseException as e:
in_flight_error = e
raise
finally:
try:
await transport_ctx.__aexit__(None, None, None)
except BaseException as e:
verbose_logger.debug(f"Error during transport context exit: {e}")
except BaseException as exit_error:
verbose_logger.debug(
f"Error during transport context exit: {exit_error}"
)
root_cause = _first_non_cancelled_cause(exit_error)
if root_cause is not None and isinstance(
in_flight_error, asyncio.CancelledError
):
raise root_cause from in_flight_error
async def run_with_session(
self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]
@ -351,7 +434,6 @@ class MCPClient:
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers = {}
if self._mcp_auth_value:
if isinstance(self._mcp_auth_value, str):
if self.auth_type == MCPAuth.bearer_token:
@ -373,17 +455,14 @@ class MCPClient:
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
# signing (including the body hash), so it uses httpx.Auth flow instead
# of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
# update the headers with the extra headers
if self.extra_headers:
headers.update(self.extra_headers)
return headers
return _strip_header_whitespace(headers)
def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]:
"""
Create a custom httpx client factory that uses LiteLLM's SSL configuration.
This factory follows the same CA bundle path logic as http_handler.py:
1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle)
2. Check SSL_VERIFY environment variable
@ -400,17 +479,14 @@ class MCPClient:
"""Create an httpx.AsyncClient with LiteLLM's SSL configuration."""
# Get unified SSL configuration using the same logic as http_handler.py
ssl_config = get_ssl_configuration(self.ssl_verify)
verbose_logger.debug(
f"MCP client using SSL configuration: {type(ssl_config).__name__}"
)
# Use SigV4 auth if configured and no explicit auth provided.
# The MCP SDK's sse_client and streamable_http_client call this
# factory without passing auth=, so self._aws_auth is used.
# For non-SigV4 clients, self._aws_auth is None — no behavior change.
effective_auth = auth if auth is not None else self._aws_auth
return httpx.AsyncClient(
headers=headers,
timeout=timeout,
@ -421,8 +497,16 @@ class MCPClient:
return factory
async def list_tools(self) -> List[MCPTool]:
"""List available tools from the server."""
async def list_tools(self, raise_on_error: bool = False) -> List[MCPTool]:
"""List available tools from the server.
Args:
raise_on_error: When True, re-raise exceptions instead of returning
an empty list. Used by the proxy's pass-through MCP flow so it
can surface upstream HTTP 401 responses as a proper 401 to the
MCP client (triggering the upstream OAuth flow) rather than
masking them as "connected, no tools".
"""
verbose_logger.debug(
f"MCP client listing tools from {self.server_url or 'stdio'}"
)
@ -450,7 +534,6 @@ class MCPClient:
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
@ -458,6 +541,8 @@ class MCPClient:
"the MCP server may have crashed, disconnected, or timed out"
)
if raise_on_error:
raise
# Return empty list instead of raising to allow graceful degradation
return []
@ -481,7 +566,6 @@ class MCPClient:
f"MCP Tool '{call_tool_request_params.name}' progress: "
f"{progress}/{total} ({percentage:.0f}%) - {message or ''}"
)
# Forward to Host if callback provided
if host_progress_callback:
try:
@ -504,14 +588,15 @@ class MCPClient:
)
return tool_result
except asyncio.CancelledError:
verbose_logger.warning("MCP client tool call was cancelled")
verbose_logger.warning(
f"MCP client tool call timed out after {self.timeout}s for {self.server_url}"
)
raise
except Exception as e:
import traceback
error_trace = traceback.format_exc()
verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}")
# Log detailed error information
error_type = type(e).__name__
verbose_logger.error(
@ -522,14 +607,12 @@ class MCPClient:
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream - "
"the MCP server may have crashed, disconnected, or timed out."
)
# Return a default error result instead of raising
return MCPCallToolResult(
content=[
@ -567,14 +650,12 @@ class MCPClient:
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream during list_tools - "
"the MCP server may have crashed, disconnected, or timed out"
)
# Return empty list instead of raising to allow graceful degradation
return []
@ -607,7 +688,6 @@ class MCPClient:
error_trace = traceback.format_exc()
verbose_logger.debug(f"MCP client get_prompt traceback:\n{error_trace}")
# Log detailed error information
error_type = type(e).__name__
verbose_logger.error(
@ -618,14 +698,12 @@ class MCPClient:
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream during get_prompt - "
"the MCP server may have crashed, disconnected, or timed out."
)
raise
async def list_resources(self) -> list[Resource]:
@ -657,14 +735,12 @@ class MCPClient:
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream during list_resources - "
"the MCP server may have crashed, disconnected, or timed out"
)
# Return empty list instead of raising to allow graceful degradation
return []
@ -699,14 +775,12 @@ class MCPClient:
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream during list_resource_templates - "
"the MCP server may have crashed, disconnected, or timed out"
)
# Return empty list instead of raising to allow graceful degradation
return []
@ -732,7 +806,6 @@ class MCPClient:
error_trace = traceback.format_exc()
verbose_logger.debug(f"MCP client read_resource traceback:\n{error_trace}")
# Log detailed error information
error_type = type(e).__name__
verbose_logger.error(
@ -743,12 +816,10 @@ class MCPClient:
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream during read_resource - "
"the MCP server may have crashed, disconnected, or timed out."
)
raise

View file

@ -18,6 +18,42 @@ else:
GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging()
def _encode_google_genai_sse_event(event_lines: List[str]) -> bytes:
return ("\n".join(event_lines) + "\n\n").encode("utf-8")
def _next_google_genai_sse_chunk(line_iter) -> bytes:
event_lines: List[str] = []
while True:
try:
line = next(line_iter)
except StopIteration:
if event_lines:
return _encode_google_genai_sse_event(event_lines)
raise
if line == "":
if event_lines:
return _encode_google_genai_sse_event(event_lines)
continue
event_lines.append(line)
async def _anext_google_genai_sse_chunk(line_iter) -> bytes:
event_lines: List[str] = []
while True:
try:
line = await line_iter.__anext__()
except StopAsyncIteration:
if event_lines:
return _encode_google_genai_sse_event(event_lines)
raise
if line == "":
if event_lines:
return _encode_google_genai_sse_event(event_lines)
continue
event_lines.append(line)
class BaseGoogleGenAIGenerateContentStreamingIterator:
"""
Base class for Google GenAI Generate Content streaming iterators that provides common logic
@ -91,18 +127,17 @@ class GoogleGenAIGenerateContentStreamingIterator(
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
# Store the iterator once to avoid multiple stream consumption
self.stream_iterator = response.iter_bytes()
# Gemini streamGenerateContent uses SSE line framing; iter_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.iter_lines()
def __iter__(self):
return self
def __next__(self):
try:
# Get the next chunk from the stored iterator
chunk = next(self.stream_iterator)
chunk = _next_google_genai_sse_chunk(self.stream_iterator)
self.collected_chunks.append(chunk)
# Just yield raw bytes
return chunk
except StopIteration:
raise StopIteration
@ -147,18 +182,17 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
# Store the async iterator once to avoid multiple stream consumption
self.stream_iterator = response.aiter_bytes()
# Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.aiter_lines()
def __aiter__(self):
return self
async def __anext__(self):
try:
# Get the next chunk from the stored async iterator
chunk = await self.stream_iterator.__anext__()
chunk = await _anext_google_genai_sse_chunk(self.stream_iterator)
self.collected_chunks.append(chunk)
# Just yield raw bytes
return chunk
except StopAsyncIteration:
await self._handle_async_streaming_logging()

View file

@ -1,7 +1,7 @@
from abc import ABC, abstractmethod
from typing import Literal
from litellm.proxy._types import CallInfo
from litellm.proxy._types import CallInfo, Litellm_EntityType
class BaseBudgetAlertType(ABC):
@ -31,6 +31,8 @@ class SoftBudgetAlert(BaseBudgetAlertType):
return "Soft Budget Crossed: "
def get_id(self, user_info: CallInfo) -> str:
if user_info.event_group == Litellm_EntityType.TEAM:
return user_info.team_id or "default_id"
return user_info.token or "default_id"

View file

@ -8,6 +8,7 @@ Notes:
"""
import asyncio
import time
from typing import TYPE_CHECKING, Any, Optional
import litellm
@ -36,11 +37,15 @@ class AlertingHangingRequestCheck:
slack_alerting_object: SlackAlerting,
):
self.slack_alerting_object = slack_alerting_object
# checks run every alerting_threshold / 2 seconds, so entries must
# stay cached for at least 1.5x the threshold to guarantee a check
# happens after they cross it
self.hanging_request_cache_ttl = int(
self.slack_alerting_object.alerting_threshold * 1.5
+ HANGING_ALERT_BUFFER_TIME_SECONDS
)
self.hanging_request_cache = InMemoryCache(
default_ttl=int(
self.slack_alerting_object.alerting_threshold
+ HANGING_ALERT_BUFFER_TIME_SECONDS
),
default_ttl=self.hanging_request_cache_ttl,
)
async def add_request_to_hanging_request_check(
@ -76,10 +81,7 @@ class AlertingHangingRequestCheck:
await self.hanging_request_cache.async_set_cache(
key=hanging_request_data.request_id,
value=hanging_request_data,
ttl=int(
self.slack_alerting_object.alerting_threshold
+ HANGING_ALERT_BUFFER_TIME_SECONDS
),
ttl=self.hanging_request_cache_ttl,
)
return
@ -111,6 +113,9 @@ class AlertingHangingRequestCheck:
if hanging_request_data is None:
continue
if hanging_request_data.alerted:
continue
request_status = (
await proxy_logging_obj.internal_usage_cache.async_get_cache(
key="request_status:{}".format(hanging_request_data.request_id),
@ -127,12 +132,21 @@ class AlertingHangingRequestCheck:
)
continue
request_age_seconds = time.time() - hanging_request_data.created_at
if request_age_seconds < self.slack_alerting_object.alerting_threshold:
# in-flight but below the alerting threshold; keep it cached
# so a later check can alert if it never completes
continue
################
# Send the Alert on Slack
################
await self.send_hanging_request_alert(
hanging_request_data=hanging_request_data
)
# flag so the entry is skipped on later ticks; one alert per hang,
# with the existing TTL still handling cleanup
hanging_request_data.alerted = True
return

View file

@ -37,6 +37,8 @@ from litellm.proxy._types import (
VirtualKeyEvent,
WebhookEvent,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.types.integrations.slack_alerting import *
from ..email_templates.templates import *
@ -1231,7 +1233,7 @@ Model Info:
and recipient_user_id is not None
and prisma_client is not None
):
user_row = await prisma_client.db.litellm_usertable.find_unique(
user_row = await UserRepository(prisma_client).table.find_unique(
where={"user_id": recipient_user_id}
)
@ -1263,7 +1265,7 @@ Model Info:
team_id = webhook_event.team_id
team_name = "Default Team"
if team_id is not None and prisma_client is not None:
team_row = await prisma_client.db.litellm_teamtable.find_unique(
team_row = await TeamRepository(prisma_client).table.find_unique(
where={"team_id": team_id}
)
if team_row is not None:

View file

@ -8,18 +8,23 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor
BaseLLMObsOTELAttributes,
safe_set_attribute,
)
from litellm.litellm_core_utils.redact_messages import (
should_redact_message_logging,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.types.utils import StandardLoggingPayload
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm.integrations._types.open_inference import (
MessageAttributes,
ImageAttributes,
SpanAttributes,
AudioAttributes,
EmbeddingAttributes,
ImageAttributes,
MessageAttributes,
MessageContentAttributes,
OpenInferenceSpanKindValues,
SpanAttributes,
ToolCallAttributes,
)
@ -53,40 +58,24 @@ class ArizeOTELAttributes(BaseLLMObsOTELAttributes):
msg.get("content", ""),
)
@staticmethod
@override
def set_response_output_messages(span: "Span", response_obj):
"""
Sets output message attributes on the span from the LLM response.
Args:
span: The OpenTelemetry span to set attributes on
response_obj: The response object containing choices with messages
"""
from litellm.integrations._types.open_inference import (
MessageAttributes,
SpanAttributes,
)
# Additive: emit structured tool_calls / multimodal content
# so Arize/Phoenix can render tool-using and image-bearing
# turns. These set NEW attribute keys (MESSAGE_TOOL_CALLS /
# MESSAGE_NAME / MESSAGE_TOOL_CALL_ID / MESSAGE_CONTENTS.*) —
# never replace the MESSAGE_CONTENT write above.
_safe_emit(
f"input message extras (idx={idx})",
_emit_input_message_extras,
span,
prefix,
msg,
)
for idx, choice in enumerate(response_obj.get("choices", [])):
response_message = choice.get("message", {})
safe_set_attribute(
span,
SpanAttributes.OUTPUT_VALUE,
response_message.get("content", ""),
)
# This shows up under `output_messages` tab on the span page.
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}"
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
response_message.get("role"),
)
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
response_message.get("content", ""),
)
# Note: `BaseLLMObsOTELAttributes.set_response_output_messages` is not
# overridden here. The live code path uses `_set_choice_outputs` (called
# via `_set_response_attributes` from `set_attributes`) which handles
# tool_calls, multimodal output, embeddings, audio, images, and structured
# outputs in a single place.
def _set_response_attributes(span: "Span", response_obj):
@ -106,11 +95,17 @@ def _set_response_attributes(span: "Span", response_obj):
def _set_choice_outputs(span: "Span", response_obj, msg_attrs, span_attrs):
for idx, choice in enumerate(response_obj.get("choices", [])):
response_message = choice.get("message", {})
safe_set_attribute(
span,
span_attrs.OUTPUT_VALUE,
response_message.get("content", ""),
)
content = response_message.get("content", "")
# Tool-only assistant responses have empty content; serialize the
# tool_calls into OUTPUT_VALUE so Arize's "Output" pane isn't blank.
output_value = content
if not output_value:
tool_calls = _get_tool_calls(response_message)
if tool_calls:
output_value = _summarize_tool_calls_for_output(tool_calls)
safe_set_attribute(span, span_attrs.OUTPUT_VALUE, output_value)
prefix = f"{span_attrs.LLM_OUTPUT_MESSAGES}.{idx}"
safe_set_attribute(
span,
@ -120,7 +115,18 @@ def _set_choice_outputs(span: "Span", response_obj, msg_attrs, span_attrs):
safe_set_attribute(
span,
f"{prefix}.{msg_attrs.MESSAGE_CONTENT}",
response_message.get("content", ""),
content,
)
# Additive: emit assistant tool_calls so tool-using turns render in
# Arize/Phoenix. Sets new MESSAGE_TOOL_CALLS keys only — does not
# change MESSAGE_CONTENT/MESSAGE_ROLE writes above.
_safe_emit(
f"output tool_calls (idx={idx})",
_emit_message_tool_calls,
span,
prefix,
response_message,
)
@ -278,6 +284,43 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs):
reasoning_tokens,
)
# Additive: cache token breakdown so prompt-caching savings render in
# Arize. Sources covered:
# - OpenAI Chat Completions: `prompt_tokens_details.cached_tokens`
# - Anthropic / Bedrock-Anthropic: `cache_read_input_tokens`,
# `cache_creation_input_tokens`
# All emits are conditional, so when none of these fields exist (the
# situation in the existing test fixtures) no extra attributes are set.
prompt_token_details = _safe_get(usage, "prompt_tokens_details") or _safe_get(
usage, "input_tokens_details"
)
cache_read = _safe_get(prompt_token_details, "cached_tokens") or _safe_get(
usage, "cache_read_input_tokens"
)
if cache_read:
safe_set_attribute(
span,
span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ,
cache_read,
)
# Anthropic / Bedrock-Anthropic only — OpenAI's `prompt_tokens_details`
# does not expose a cache-write count, so we read straight off `usage`.
cache_write = _safe_get(usage, "cache_creation_input_tokens")
if cache_write:
safe_set_attribute(
span,
span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE,
cache_write,
)
audio_prompt_tokens = _safe_get(prompt_token_details, "audio_tokens")
if audio_prompt_tokens:
safe_set_attribute(
span,
span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO,
audio_prompt_tokens,
)
def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
"""
@ -321,6 +364,10 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
"videos",
"realtime",
"pass_through",
# `passthrough` (no underscore) is what real call_types use:
# `allm_passthrough_route`, `llm_passthrough_route`. Without
# this they fell through to UNKNOWN, blanking span.kind.
"passthrough",
"anthropic_messages",
"ocr",
)
@ -396,6 +443,18 @@ def set_attributes(
"""
Populates span with OpenInference-compliant LLM attributes for Arize and Phoenix tracing.
"""
# Coerce non-dict response objects (e.g. httpx.Response from passthrough
# routes) into a dict so downstream `.get()` calls don't crash. Existing
# dict / `.get()`-bearing objects (incl. Pydantic OpenAI Responses API
# models) are returned unchanged, preserving the existing test behavior.
response_obj_for_attrs = _coerce_response_obj_for_attrs(response_obj)
# Set span.kind defensively before anything else. If a downstream step
# throws, the span still has a kind so Arize can render it correctly
# (an LLM call instead of UNKNOWN). This is the single source of truth
# for span.kind — no late re-write happens below.
_safe_emit("early span kind", _set_early_span_kind, span, kwargs)
try:
optional_params = _sanitize_optional_params(kwargs.get("optional_params"))
litellm_params = kwargs.get("litellm_params", {}) or {}
@ -415,25 +474,22 @@ def set_attributes(
metadata_tools = _extract_metadata_tools(metadata)
optional_tools = _extract_optional_tools(optional_params)
call_type = standard_logging_payload.get("call_type")
_set_request_attributes(
span=span,
kwargs=kwargs,
standard_logging_payload=standard_logging_payload,
optional_params=optional_params,
litellm_params=litellm_params,
response_obj=response_obj,
response_obj=response_obj_for_attrs,
span_attrs=SpanAttributes,
)
span_kind = _infer_open_inference_span_kind(call_type=call_type)
# span.kind was already set above by `_set_early_span_kind`. We do
# NOT re-write it here based on tool presence: a chat completion
# that passes `tools=[...]` (or returns `tool_calls`) is still an
# LLM call per the OpenInference spec — TOOL is reserved for actual
# tool execution spans, not LLM calls that request tools.
_set_tool_attributes(span, optional_tools, metadata_tools)
if (
optional_tools or metadata_tools
) and span_kind != OpenInferenceSpanKindValues.TOOL.value:
span_kind = OpenInferenceSpanKindValues.TOOL.value
safe_set_attribute(span, SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind)
attributes.set_messages(span, kwargs)
model_params = (
@ -443,7 +499,7 @@ def set_attributes(
)
_set_model_params(span, model_params, SpanAttributes)
_set_response_attributes(span=span, response_obj=response_obj)
_set_response_attributes(span=span, response_obj=response_obj_for_attrs)
except Exception as e:
verbose_logger.error(
@ -452,6 +508,22 @@ def set_attributes(
if hasattr(span, "record_exception"):
span.record_exception(e)
# Additive emitters. Each is independently guarded so a failure can never
# blank the attributes set by the main try-block above. New attributes are
# written under new keys; existing attributes are not overwritten.
slp = kwargs.get("standard_logging_object")
_safe_emit("session/user attrs", _set_session_and_user_attrs, span, kwargs, slp)
_safe_emit("response cost", _set_response_cost_attr, span, slp)
_safe_emit(
"passthrough normalization",
_maybe_normalize_passthrough,
span,
kwargs,
response_obj,
response_obj_for_attrs,
slp,
)
def _sanitize_optional_params(optional_params: Optional[dict]) -> dict:
if not isinstance(optional_params, dict):
@ -534,3 +606,529 @@ def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) ->
user_id = model_params.get("user")
if user_id is not None:
safe_set_attribute(span, span_attrs.USER_ID, user_id)
# ---------------------------------------------------------------------------
# Additive rendering helpers (introduced to enhance Arize/Phoenix rendering
# without changing any previously-emitted attribute keys or values).
# ---------------------------------------------------------------------------
def _safe_emit(label: str, fn, *args, **kwargs) -> None:
"""Run an additive attribute emitter, swallowing any error so it cannot
blank attributes set elsewhere on the span. Failures are logged at debug.
"""
try:
fn(*args, **kwargs)
except Exception as e:
verbose_logger.debug("[Arize] %s skipped: %s", label, e)
def _set_early_span_kind(span: "Span", kwargs: dict) -> None:
"""Defensively set OPENINFERENCE_SPAN_KIND before any other logic runs."""
slp = kwargs.get("standard_logging_object")
call_type = slp.get("call_type") if isinstance(slp, dict) else None
safe_set_attribute(
span,
SpanAttributes.OPENINFERENCE_SPAN_KIND,
_infer_open_inference_span_kind(call_type=call_type),
)
def _coerce_response_obj_for_attrs(response_obj):
"""Return a `.get`-compatible view of `response_obj` when possible.
- dicts and Pydantic models that already expose `.get` are returned
unchanged (preserves all current behavior, including the Responses API
flow which relies on Pydantic attribute access).
- `httpx.Response` and other text-only responses (passthrough routes)
are JSON-decoded so the standard extraction paths can read fields like
`id`, `model`, and `usage`. On failure the original object is returned
so behavior is no worse than today.
"""
if response_obj is None or hasattr(response_obj, "get"):
return response_obj
text = getattr(response_obj, "text", None)
if isinstance(text, str) and text:
try:
parsed = json.loads(text)
if isinstance(parsed, dict):
return parsed
except Exception:
pass
return response_obj
def _coerce_text(value) -> Optional[str]:
"""Best-effort text extraction from a message-content value.
Returns None when no textual portion can be derived. Handles:
- plain strings
- lists of OpenAI-style content parts (`{"type": "text", "text": ...}`)
- lists of Anthropic-style content parts (`{"type": "text", "text": ...}`
or `{"type": "input_text", "text": ...}`)
"""
if value is None:
return None
if isinstance(value, str):
return value
if isinstance(value, list):
parts = []
for part in value:
if isinstance(part, str):
parts.append(part)
elif isinstance(part, dict):
text = part.get("text") or part.get("input_text")
if isinstance(text, str):
parts.append(text)
if parts:
return "\n".join(parts)
return None
def _to_plain_dict(value):
"""Best-effort: coerce a value (Pydantic model / dict / None) to a dict.
Returns the original value when no safe conversion exists. Used to bridge
OpenAI Pydantic message/tool_call objects into the dict-based helpers.
"""
if value is None or isinstance(value, dict):
return value
model_dump = getattr(value, "model_dump", None)
if callable(model_dump):
try:
return model_dump()
except Exception:
pass
return value
def _get_tool_calls(message) -> Optional[list]:
"""Return ``message.tool_calls`` only when it's a non-empty list.
Works for dicts and Pydantic message objects via ``_safe_get``.
"""
tool_calls = _safe_get(message, "tool_calls")
return tool_calls if isinstance(tool_calls, list) and tool_calls else None
def _normalize_tool_call(raw_tc) -> Optional[Dict[str, Any]]:
"""Normalize a single tool_call (dict or Pydantic) into a stable shape:
{"id": str|None, "type": str, "function": {"name": str|None, "arguments": str|None}}
Arguments are coerced to a JSON string per OpenInference convention.
Returns ``None`` when ``raw_tc`` cannot be coerced to a dict.
"""
tc = _to_plain_dict(raw_tc)
if not isinstance(tc, dict):
return None
function = _to_plain_dict(tc.get("function"))
name = function.get("name") if isinstance(function, dict) else None
args = function.get("arguments") if isinstance(function, dict) else None
if args is not None and not isinstance(args, str):
try:
args = json.dumps(args)
except Exception:
args = str(args)
return {
"id": tc.get("id"),
"type": tc.get("type", "function"),
"function": {"name": name, "arguments": args},
}
def _summarize_tool_calls_for_output(tool_calls) -> str:
"""Render a tool_calls list as a compact JSON string for OUTPUT_VALUE.
Best-effort: returns ``str(tool_calls)`` if anything unexpected happens
so OUTPUT_VALUE is never blanked on a malformed payload.
"""
try:
normalized = [n for n in (_normalize_tool_call(tc) for tc in tool_calls) if n]
return json.dumps({"tool_calls": normalized})
except Exception:
return str(tool_calls)
def _emit_message_tool_calls(span: "Span", prefix: str, message) -> None:
"""Emit ``MESSAGE_TOOL_CALLS.*`` for an assistant message that requested
tool calls. Pure addition: only writes when ``tool_calls`` is non-empty.
Accepts dicts or Pydantic message objects (e.g. ``litellm.Message``); the
same applies to each tool_call entry.
"""
tool_calls = _get_tool_calls(message)
if not tool_calls:
return
for tc_idx, raw_tc in enumerate(tool_calls):
tc = _normalize_tool_call(raw_tc)
if tc is None:
continue
tc_prefix = f"{prefix}.{MessageAttributes.MESSAGE_TOOL_CALLS}.{tc_idx}"
if tc["id"]:
safe_set_attribute(
span, f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_ID}", tc["id"]
)
fn = tc["function"]
if fn["name"]:
safe_set_attribute(
span,
f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}",
fn["name"],
)
if fn["arguments"] is not None:
safe_set_attribute(
span,
f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
fn["arguments"],
)
def _emit_input_message_extras(span: "Span", prefix: str, message: dict) -> None:
"""Emit additive attributes for an input message:
- `MESSAGE_NAME` and `MESSAGE_TOOL_CALL_ID` (commonly set on tool-result
messages so traces show which tool produced which result).
- `MESSAGE_TOOL_CALLS.*` when an assistant message requested tools.
- `MESSAGE_CONTENTS.*` structured content for list-shaped content
(multimodal text + image parts). The plain `MESSAGE_CONTENT` write is
still performed by the caller, so renderers that only read the legacy
key continue to work.
"""
if not isinstance(message, dict):
return
name = message.get("name")
if name:
safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_NAME}", name)
tool_call_id = message.get("tool_call_id")
if tool_call_id:
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_TOOL_CALL_ID}",
tool_call_id,
)
_emit_message_tool_calls(span, prefix, message)
content = message.get("content")
if isinstance(content, list):
contents_prefix = f"{prefix}.{MessageAttributes.MESSAGE_CONTENTS}"
for part_idx, part in enumerate(content):
if not isinstance(part, dict):
continue
part_prefix = f"{contents_prefix}.{part_idx}"
part_type = part.get("type")
if part_type in ("text", "input_text"):
text = part.get("text")
if isinstance(text, str):
safe_set_attribute(
span,
f"{part_prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TYPE}",
"text",
)
safe_set_attribute(
span,
f"{part_prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TEXT}",
text,
)
elif part_type in ("image_url", "image", "input_image"):
url = None
image = part.get("image_url")
if isinstance(image, dict):
url = image.get("url")
elif isinstance(image, str):
url = image
if not url:
# Anthropic-style source.{type=base64,media_type,data}
source = part.get("source")
if isinstance(source, dict) and source.get("data"):
media_type = source.get("media_type", "image/jpeg")
url = f"data:{media_type};base64,{source['data']}"
elif isinstance(part.get("url"), str):
url = part["url"]
if url:
safe_set_attribute(
span,
f"{part_prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TYPE}",
"image",
)
safe_set_attribute(
span,
f"{part_prefix}.message_content.image.image.url",
url,
)
def _set_session_and_user_attrs(
span: "Span", kwargs: dict, standard_logging_payload
) -> None:
"""Emit `SESSION_ID` / `USER_ID` / team metadata when source data exists.
`SESSION_ID` is emitted only when an explicit end-user identifier exists
(`metadata.user_api_key_end_user_id`). We deliberately do NOT fall back
to `trace_id`, because that would create a distinct "session" for every
single request and distort Arize's Session-grouping analytics. The
`trace_id` is still emitted under its own `litellm.trace_id` key so
spans remain filterable by trace.
USER_ID is *only* emitted when no upstream path (model_params.user or
optional_params.user) has already set it, to avoid overwriting an
existing value with a possibly-different one from API-key metadata.
"""
if not isinstance(standard_logging_payload, dict):
return
metadata = standard_logging_payload.get("metadata") or {}
if not isinstance(metadata, dict):
return
session_id = metadata.get("user_api_key_end_user_id")
if session_id:
safe_set_attribute(span, SpanAttributes.SESSION_ID, str(session_id))
trace_id = standard_logging_payload.get("trace_id")
if trace_id:
safe_set_attribute(span, "litellm.trace_id", str(trace_id))
optional_params = kwargs.get("optional_params") or {}
model_params = standard_logging_payload.get("model_parameters") or {}
has_user_already = bool(
(isinstance(optional_params, dict) and optional_params.get("user"))
or (isinstance(model_params, dict) and model_params.get("user"))
)
if not has_user_already:
user_id = metadata.get("user_api_key_user_id")
if user_id:
safe_set_attribute(span, SpanAttributes.USER_ID, str(user_id))
team_id = metadata.get("user_api_key_team_id")
if team_id:
safe_set_attribute(span, "litellm.team_id", str(team_id))
team_alias = metadata.get("user_api_key_team_alias")
if team_alias:
safe_set_attribute(span, "litellm.team_alias", str(team_alias))
key_alias = metadata.get("user_api_key_alias")
if key_alias:
safe_set_attribute(span, "litellm.key_alias", str(key_alias))
def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None:
"""Emit cost attributes from the StandardLoggingPayload when present.
Uses the OpenInference `llm.cost.total` key so Arize / Phoenix can
surface the cost in their "Total Cost" column. LiteLLM only tracks a
single total in `StandardLoggingPayload.response_cost`, so we cannot
split it into prompt/completion. We also keep the legacy
`llm.response.cost` key for back-compat with any consumer querying it.
"""
if not isinstance(standard_logging_payload, dict):
return
cost = standard_logging_payload.get("response_cost")
if cost is None:
return
try:
cost_value = float(cost)
except (TypeError, ValueError):
return
safe_set_attribute(span, "llm.cost.total", cost_value)
safe_set_attribute(span, "llm.response.cost", cost_value)
def _is_passthrough_call_type(call_type: Optional[str]) -> bool:
if not call_type:
return False
lowered = str(call_type).lower()
return "passthrough" in lowered or "pass_through" in lowered
def _maybe_normalize_passthrough(
span: "Span",
kwargs: dict,
raw_response_obj,
coerced_response_obj,
standard_logging_payload,
) -> None:
"""Surface input/output text for passthrough routes (e.g. Bedrock
InvokeModel) so the parent span renders as more than `usage` numbers.
Only runs when `call_type` is a passthrough variant. Reads from:
- `kwargs["additional_args"]["complete_input_dict"]` for input
- the coerced response (or `kwargs["original_response"]`) for output
All emits are best-effort: if the provider shape isn't recognized the
helper exits silently. Existing chat/completion paths never enter this
helper because their call_type doesn't contain "passthrough".
TEMPORARY BRIDGE: passthrough handlers don't populate the
StandardLoggingPayload `messages` field today (they call
`transform_response(messages=[])`), so the input is only available via
`additional_args.complete_input_dict`. The proper fix is upstream in
`base_passthrough_logging_handler._create_response_logging_payload()`:
once that populates SLP `messages`/`response`, every callback gets
passthrough I/O (with central redaction) for free and this helper's
`complete_input_dict` fallback can be deleted. See follow-up issue.
"""
call_type = (
standard_logging_payload.get("call_type")
if isinstance(standard_logging_payload, dict)
else None
)
if not _is_passthrough_call_type(call_type):
return
# Respect LiteLLM's central message-redaction contract. The normal
# chat/completion path is redacted by `perform_redaction` before
# callbacks run, but `complete_input_dict` (read below) is NOT covered by
# that layer — so without this gate, an operator who enabled redaction
# would still see raw passthrough prompts in Arize. Skip entirely when
# redaction is on so neither input nor output leaks through this bridge.
if should_redact_message_logging(kwargs):
return
# --- INPUT --------------------------------------------------------------
additional_args = kwargs.get("additional_args") or {}
complete_input_dict = (
additional_args.get("complete_input_dict")
if isinstance(additional_args, dict)
else None
)
if isinstance(complete_input_dict, dict):
_set_passthrough_input_attributes(span, complete_input_dict.get("messages"))
# --- OUTPUT -------------------------------------------------------------
parsed_response = _parse_passthrough_response(
raw_response_obj, coerced_response_obj, kwargs
)
if not isinstance(parsed_response, dict):
return
_set_passthrough_output_attributes(span, parsed_response)
def _set_passthrough_input_attributes(span: "Span", messages) -> None:
"""Render passthrough request messages into INPUT_VALUE + LLM_INPUT_MESSAGES."""
if not (isinstance(messages, list) and messages):
return
# Set INPUT_VALUE from the last user message text if discoverable.
last_text = None
for msg in reversed(messages):
if isinstance(msg, dict):
last_text = _coerce_text(msg.get("content"))
if last_text:
break
if last_text:
safe_set_attribute(span, SpanAttributes.INPUT_VALUE, last_text)
# Mirror messages into LLM_INPUT_MESSAGES so the input pane renders.
for idx, msg in enumerate(messages):
if not isinstance(msg, dict):
continue
prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}"
role = msg.get("role")
if role:
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
role,
)
text = _coerce_text(msg.get("content"))
if text is not None:
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
text,
)
def _set_passthrough_output_attributes(span: "Span", parsed_response: dict) -> None:
"""Render passthrough response into OUTPUT_VALUE + LLM_OUTPUT_MESSAGES."""
# Anthropic / Bedrock-Anthropic: `content` is a list of typed parts.
content_list = parsed_response.get("content")
if isinstance(content_list, list) and content_list:
texts = []
for part in content_list:
if isinstance(part, dict) and isinstance(part.get("text"), str):
texts.append(part["text"])
joined = "\n\n".join(t for t in texts if t)
if joined:
safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, joined)
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0"
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
parsed_response.get("role", "assistant"),
)
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
joined,
)
# OpenAI-style passthrough: `choices[0].message.content`
choices = parsed_response.get("choices")
if isinstance(choices, list) and choices:
first = choices[0]
if isinstance(first, dict):
msg = first.get("message")
if isinstance(msg, dict):
text = _coerce_text(msg.get("content"))
if text:
safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, text)
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0"
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
msg.get("role", "assistant"),
)
safe_set_attribute(
span,
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
text,
)
def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs):
"""Return a dict view of the provider response for passthrough routes."""
# Prefer the coerced view (already JSON-parsed for httpx.Response).
candidates = []
if isinstance(coerced_response_obj, dict):
candidates.append(coerced_response_obj)
if (
isinstance(raw_response_obj, dict)
and raw_response_obj is not coerced_response_obj
):
candidates.append(raw_response_obj)
for candidate in candidates:
# StandardPassThroughResponseObject wrapper: {"response": "..."}.
if (
"response" in candidate
and "content" not in candidate
and "choices" not in candidate
):
inner = candidate.get("response")
if isinstance(inner, str):
try:
parsed = json.loads(inner)
if isinstance(parsed, dict):
return parsed
except Exception:
continue
if isinstance(inner, dict):
return inner
else:
return candidate
# Fallback: kwargs["original_response"] from the OTel base path.
original = kwargs.get("original_response") if isinstance(kwargs, dict) else None
if isinstance(original, dict):
return original
if isinstance(original, str):
try:
parsed = json.loads(original)
if isinstance(parsed, dict):
return parsed
except Exception:
return None
return None

View file

@ -1,5 +1,7 @@
import os
from typing import TYPE_CHECKING, Any, Optional, Union
import threading
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
from litellm._logging import verbose_logger
from litellm.integrations.arize import _utils
@ -8,8 +10,10 @@ from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig
if TYPE_CHECKING:
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SpanProcessor
from opentelemetry.trace import Span as _Span
from opentelemetry.trace import SpanKind
from opentelemetry.trace import Tracer
from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry
from litellm.integrations.opentelemetry import (
@ -21,20 +25,27 @@ if TYPE_CHECKING:
OpenTelemetryConfig = _OpenTelemetryConfig
Span = Union[_Span, Any]
OpenTelemetry = _OpenTelemetry
LITELLM_TRACER_NAME: str
else:
Protocol = Any
OpenTelemetryConfig = Any
Span = Any
Tracer = Any
TracerProvider = Any
SpanKind = Any
# Import OpenTelemetry at runtime
SpanProcessor = Any
try:
from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.integrations.opentelemetry import (
LITELLM_TRACER_NAME,
OpenTelemetry,
)
except ImportError:
LITELLM_TRACER_NAME = "litellm"
OpenTelemetry = None # type: ignore
ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces"
_MAX_PROJECT_PROVIDERS = 64
class ArizePhoenixLogger(OpenTelemetry): # type: ignore
@ -48,37 +59,142 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
def _init_tracing(self, tracer_provider):
"""
Override to always create a *private* TracerProvider for Arize Phoenix.
Override to create per-project TracerProviders (LRU-cached) for Arize Phoenix.
The base ``OpenTelemetry._init_tracing`` falls back to the global
TracerProvider when one already exists. That causes whichever
integration initialises second to silently reuse the first one's
exporter, so spans only reach one destination.
By creating our own provider we guarantee Arize Phoenix always gets
its own exporter pipeline, regardless of initialisation order.
"""
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import SpanKind
if tracer_provider is not None:
# Explicitly supplied (e.g. in tests) — honour it.
self.tracer = tracer_provider.get_tracer("litellm")
self._use_injected_tracer_provider = True
self._shared_span_processor = None
self.tracer = tracer_provider.get_tracer(LITELLM_TRACER_NAME)
self.span_kind = SpanKind
return
# Always create a dedicated provider — never touch the global one.
provider = TracerProvider(resource=self._get_litellm_resource(self.config))
provider.add_span_processor(self._get_span_processor())
self.tracer = provider.get_tracer("litellm")
self._use_injected_tracer_provider = False
self._project_providers: OrderedDict[str, TracerProvider] = OrderedDict()
self._project_providers_lock = threading.Lock()
self._shared_span_processor = self._get_span_processor()
self.span_kind = SpanKind
default_project = self._resolve_project_name({})
self.tracer = self._get_tracer_for(default_project)
verbose_logger.debug(
"ArizePhoenixLogger: Created dedicated TracerProvider "
"(endpoint=%s, exporter=%s)",
"ArizePhoenixLogger: Initialized per-project TracerProvider cache "
"(default_project=%s, endpoint=%s, exporter=%s)",
default_project,
self.config.endpoint,
self.config.exporter,
)
def flush_tracer_providers(self) -> None:
"""
Flush all cached per-project providers and the shared span processor.
Call on graceful proxy shutdown. Do not call on LRU eviction — in-flight
spans may still reference evicted providers.
"""
if getattr(self, "_use_injected_tracer_provider", False):
return
shared_processor = getattr(self, "_shared_span_processor", None)
if shared_processor is not None:
try:
shared_processor.force_flush()
except Exception as e:
verbose_logger.debug(
"ArizePhoenixLogger: shared span processor force_flush failed: %s",
e,
)
with getattr(self, "_project_providers_lock", threading.Lock()):
providers = list(getattr(self, "_project_providers", {}).values())
for provider in providers:
try:
provider.force_flush()
except Exception as e:
verbose_logger.debug(
"ArizePhoenixLogger: TracerProvider force_flush failed: %s", e
)
def _get_litellm_resource_for_project(self, project_name: str):
"""
Build an OTEL Resource with project routing attrs that win over env detector.
Phoenix uses ``openinference.project.name``; Arize AX uses ``model_id`` and
``service.name``. Project attrs are merged last so OTEL_RESOURCE_ATTRIBUTES
from init does not pin every provider to one project.
"""
from opentelemetry.sdk.resources import OTELResourceDetector, Resource
project_attributes: dict[str, str] = {
"openinference.project.name": project_name,
"model_id": project_name,
"service.name": project_name,
}
deployment_environment = getattr(self.config, "deployment_environment", None)
if deployment_environment is not None:
project_attributes["deployment.environment"] = deployment_environment
env_resource = OTELResourceDetector().detect()
project_resource = Resource.create(project_attributes) # type: ignore[arg-type]
return env_resource.merge(project_resource)
def _build_tracer_provider_for_project(self, project_name: str) -> TracerProvider:
"""Create a TracerProvider for *project_name* (caller holds no cache lock)."""
from opentelemetry.sdk.trace import TracerProvider
provider = TracerProvider(
resource=self._get_litellm_resource_for_project(project_name)
)
provider.add_span_processor(self._shared_span_processor)
return provider
def _get_tracer_for(self, project_name: str) -> Tracer:
"""Return a tracer for *project_name*, creating/caching a provider on miss."""
if getattr(self, "_use_injected_tracer_provider", False):
return self.tracer
with self._project_providers_lock:
if project_name in self._project_providers:
self._project_providers.move_to_end(project_name)
return self._project_providers[project_name].get_tracer(
LITELLM_TRACER_NAME
)
# OTELResourceDetector().detect() is synchronous; build outside the lock so
# concurrent requests for other projects are not blocked on cache misses.
new_provider = self._build_tracer_provider_for_project(project_name)
with self._project_providers_lock:
if project_name in self._project_providers:
self._project_providers.move_to_end(project_name)
return self._project_providers[project_name].get_tracer(
LITELLM_TRACER_NAME
)
if len(self._project_providers) >= _MAX_PROJECT_PROVIDERS:
self._project_providers.popitem(last=False)
self._project_providers[project_name] = new_provider
return new_provider.get_tracer(LITELLM_TRACER_NAME)
def _resolve_tracer_for_kwargs(self, kwargs: dict) -> Tuple[str, Tracer]:
"""Resolve project name once and return the matching tracer."""
project_name = self._resolve_project_name(kwargs)
return project_name, self._get_tracer_for(project_name)
def get_tracer_to_use_for_request(self, kwargs: dict) -> Tracer:
"""Route guardrail/raw-request spans to the same per-project tracer as the request."""
if getattr(self, "_use_injected_tracer_provider", False):
return self.tracer
return self._resolve_tracer_for_kwargs(kwargs)[1]
def _init_otel_logger_on_litellm_proxy(self):
"""
Override: Arize Phoenix should NOT overwrite the proxy's
@ -93,56 +209,109 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
@staticmethod
def set_arize_phoenix_attributes(span: Span, kwargs, response_obj):
from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import (
safe_set_attribute,
)
_utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes)
# Dynamic project name: check metadata first, then fall back to env var config
dynamic_project_name = ArizePhoenixLogger._get_dynamic_project_name(kwargs)
if dynamic_project_name:
safe_set_attribute(span, "openinference.project.name", dynamic_project_name)
else:
# Fall back to static config from env var
config = ArizePhoenixLogger.get_arize_phoenix_config()
if config.project_name:
safe_set_attribute(
span, "openinference.project.name", config.project_name
)
return
@staticmethod
def _get_dynamic_project_name(kwargs) -> Optional[str]:
"""
Retrieve dynamic Phoenix project name from request metadata.
def _normalize_project_name(name: Optional[str]) -> Optional[str]:
if name is None:
return None
normalized = str(name).strip()
return normalized if normalized else None
Users can set `metadata.phoenix_project_name` in their request to route
traces to different Phoenix projects dynamically.
"""
standard_logging_payload = kwargs.get("standard_logging_object")
if isinstance(standard_logging_payload, dict):
metadata = standard_logging_payload.get("metadata")
@staticmethod
def _iter_metadata_dicts_from_kwargs(kwargs: dict):
"""Yield request metadata dicts; standard_logging_object before litellm_params."""
for key in ("standard_logging_object", "litellm_params"):
found_key = kwargs.get(key)
if not isinstance(found_key, dict):
continue
metadata = found_key.get("metadata")
if isinstance(metadata, dict):
project_name = metadata.get("phoenix_project_name")
if project_name:
return str(project_name)
yield metadata
# Also check litellm_params.metadata for SDK usage
@staticmethod
def _is_proxy_request(kwargs: dict) -> bool:
"""True when the call is routed through the LiteLLM proxy.
Proxy mode is determined solely by the server-set ``proxy_server_request``
field in ``litellm_params``. Checking request metadata for
``user_api_key_auth_metadata`` is intentionally avoided: that field is
user-supplied and would let an authenticated caller fake proxy-mode
detection to route their telemetry into arbitrary Arize/Phoenix projects.
"""
litellm_params = kwargs.get("litellm_params")
if isinstance(litellm_params, dict):
metadata = litellm_params.get("metadata") or {}
else:
metadata = {}
if isinstance(metadata, dict):
project_name = metadata.get("phoenix_project_name")
if project_name:
return str(project_name)
return isinstance(litellm_params, dict) and bool(
litellm_params.get("proxy_server_request")
)
@staticmethod
def _project_from_metadata_dict(
metadata: dict, metadata_key: str, *, proxy_mode: bool
) -> Optional[str]:
"""
Read a Phoenix project field from proxy/SDK metadata.
On the proxy, only ``user_api_key_auth_metadata`` (team/key config) may
select the project. SDK callers may still set project fields directly on
``metadata``.
"""
auth_metadata = metadata.get("user_api_key_auth_metadata")
if isinstance(auth_metadata, dict):
project = ArizePhoenixLogger._normalize_project_name(
auth_metadata.get(metadata_key)
)
if project:
return project
if not proxy_mode:
return ArizePhoenixLogger._normalize_project_name(
metadata.get(metadata_key)
)
return None
def _get_phoenix_context(self, kwargs):
@staticmethod
def _metadata_project_from_kwargs(kwargs: dict, metadata_key: str) -> Optional[str]:
proxy_mode = ArizePhoenixLogger._is_proxy_request(kwargs)
for metadata in ArizePhoenixLogger._iter_metadata_dicts_from_kwargs(kwargs):
project = ArizePhoenixLogger._project_from_metadata_dict(
metadata, metadata_key, proxy_mode=proxy_mode
)
if project:
return project
return None
@staticmethod
def _resolve_project_name(kwargs: dict) -> str:
"""
Resolve the target Phoenix/Arize project for this request.
Proxy priority: ``user_api_key_auth_metadata.phoenix_project_name_override``,
``user_api_key_auth_metadata.phoenix_project_name``, env, then ``default``.
SDK priority: request metadata fields, then env, then ``default``.
"""
override = ArizePhoenixLogger._metadata_project_from_kwargs(
kwargs, "phoenix_project_name_override"
)
if override:
return override
phoenix_name = ArizePhoenixLogger._metadata_project_from_kwargs(
kwargs, "phoenix_project_name"
)
if phoenix_name:
return phoenix_name
env_name = ArizePhoenixLogger._normalize_project_name(
os.environ.get("PHOENIX_PROJECT_NAME")
or os.environ.get("ARIZE_PROJECT_NAME")
)
if env_name:
return env_name
return "default"
def _get_phoenix_context(self, kwargs, tracer: Optional[Tracer] = None):
"""
Build a trace context for Phoenix's dedicated TracerProvider.
@ -159,11 +328,13 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
"""
from opentelemetry import trace
if tracer is None:
tracer = self._resolve_tracer_for_kwargs(kwargs)[1]
litellm_params = kwargs.get("litellm_params", {}) or {}
proxy_server_request = litellm_params.get("proxy_server_request", {}) or {}
headers = proxy_server_request.get("headers", {}) or {}
# Propagate distributed trace context if the caller sent a traceparent
traceparent_ctx = (
self.get_traceparent_from_header(headers=headers)
if headers.get("traceparent")
@ -173,10 +344,8 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
is_proxy_mode = bool(proxy_server_request)
if is_proxy_mode:
# Create a parent span on Phoenix's own tracer so both parent
# and child are exported to Phoenix.
start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time"))
parent_span = self.tracer.start_span(
parent_span = tracer.start_span(
name="litellm_proxy_request",
start_time=(
self._to_ns(start_time_val) if start_time_val is not None else None
@ -187,100 +356,77 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
ctx = trace.set_span_in_context(parent_span)
return ctx, parent_span
# SDK mode — no parent span needed
return traceparent_ctx, None
def _handle_success(self, kwargs, response_obj, start_time, end_time):
"""
Override to always create spans on ArizePhoenixLogger's dedicated TracerProvider.
The base class's ``_get_span_context`` would find the parent span created by
the ``otel`` callback on the *global* TracerProvider. That span is invisible
in Phoenix (different exporter pipeline), so we ignore it and build our own
hierarchy via ``_get_phoenix_context``.
"""
from opentelemetry.trace import Status, StatusCode
verbose_logger.debug(
"ArizePhoenixLogger: Logging kwargs: %s, OTEL config settings=%s",
kwargs,
self.config,
self._handle_phoenix_trace(
kwargs, response_obj, start_time, end_time, success=True
)
ctx, parent_span = self._get_phoenix_context(kwargs)
# Create litellm_request span (child of our parent when in proxy mode)
span = self.tracer.start_span(
name=self._get_span_name(kwargs),
start_time=self._to_ns(start_time),
context=ctx,
)
span.set_status(Status(StatusCode.OK))
self.set_attributes(span, kwargs, response_obj)
# Raw-request sub-span (if enabled) — must be created before
# ending the parent span so the hierarchy is valid.
self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span)
span.end(end_time=self._to_ns(end_time))
# Guardrail span
self._create_guardrail_span(kwargs=kwargs, context=ctx)
# Annotate and close our proxy parent span
if parent_span is not None:
parent_span.set_status(Status(StatusCode.OK))
self.set_attributes(parent_span, kwargs, response_obj)
parent_span.end(end_time=self._to_ns(end_time))
# Metrics & cost recording
self._record_metrics(kwargs, response_obj, start_time, end_time)
# Semantic logs
if self.config.enable_events:
self._emit_semantic_logs(kwargs, response_obj, span)
def _handle_failure(self, kwargs, response_obj, start_time, end_time):
"""
Override to always create failure spans on ArizePhoenixLogger's dedicated
TracerProvider. Mirrors ``_handle_success`` but sets ERROR status.
"""
self._handle_phoenix_trace(
kwargs, response_obj, start_time, end_time, success=False
)
def _handle_phoenix_trace(
self,
kwargs,
response_obj,
start_time,
end_time,
*,
success: bool,
):
from opentelemetry.trace import Status, StatusCode
verbose_logger.debug(
"ArizePhoenixLogger: Failure - Logging kwargs: %s, OTEL config settings=%s",
"ArizePhoenixLogger: %s - kwargs: %s, OTEL config settings=%s",
"success" if success else "failure",
kwargs,
self.config,
)
ctx, parent_span = self._get_phoenix_context(kwargs)
_project_name, tracer = self._resolve_tracer_for_kwargs(kwargs)
ctx, parent_span = self._get_phoenix_context(kwargs, tracer=tracer)
# Create litellm_request span (child of our parent when in proxy mode)
span = self.tracer.start_span(
status = Status(StatusCode.OK if success else StatusCode.ERROR)
span = tracer.start_span(
name=self._get_span_name(kwargs),
start_time=self._to_ns(start_time),
context=ctx,
)
span.set_status(Status(StatusCode.ERROR))
span.set_status(status)
self.set_attributes(span, kwargs, response_obj)
self._record_exception_on_span(span=span, kwargs=kwargs)
if not success:
self._record_exception_on_span(span=span, kwargs=kwargs)
if success:
self._maybe_log_raw_request(
kwargs, response_obj, start_time, end_time, span
)
span.end(end_time=self._to_ns(end_time))
# Guardrail span
self._create_guardrail_span(kwargs=kwargs, context=ctx)
# Annotate and close our proxy parent span
if parent_span is not None:
parent_span.set_status(Status(StatusCode.ERROR))
parent_span.set_status(status)
self.set_attributes(parent_span, kwargs, response_obj)
self._record_exception_on_span(span=parent_span, kwargs=kwargs)
if not success:
self._record_exception_on_span(span=parent_span, kwargs=kwargs)
parent_span.end(end_time=self._to_ns(end_time))
if success:
self._record_metrics(kwargs, response_obj, start_time, end_time)
if self.config.enable_events:
self._emit_semantic_logs(kwargs, response_obj, span)
@staticmethod
def get_arize_phoenix_config() -> ArizePhoenixConfig:
"""
Retrieves the Arize Phoenix configuration based on environment variables.
Returns:
ArizePhoenixConfig: A Pydantic model containing Arize Phoenix configuration.
"""
api_key = os.environ.get("PHOENIX_API_KEY", None)
@ -295,18 +441,15 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
protocol: Protocol = "otlp_http"
if collector_endpoint:
# Parse the endpoint to determine protocol
if collector_endpoint.startswith("grpc://") or (
":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint
):
endpoint = collector_endpoint
protocol = "otlp_grpc"
else:
# Phoenix Cloud endpoints (app.phoenix.arize.com) include the space in the URL
if "app.phoenix.arize.com" in collector_endpoint:
endpoint = collector_endpoint
protocol = "otlp_http"
# For other HTTP endpoints, ensure they have the correct path
elif "/v1/traces" not in collector_endpoint:
if collector_endpoint.endswith("/v1"):
endpoint = collector_endpoint + "/traces"
@ -318,7 +461,6 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
endpoint = collector_endpoint
protocol = "otlp_http"
else:
# If no endpoint specified, self hosted phoenix
endpoint = "http://localhost:6006/v1/traces"
protocol = "otlp_http"
verbose_logger.debug(
@ -329,12 +471,11 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
if api_key is not None:
otlp_auth_headers = f"Authorization=Bearer {api_key}"
elif "app.phoenix.arize.com" in endpoint:
# Phoenix Cloud requires an API key
raise ValueError(
"PHOENIX_API_KEY must be set when using Phoenix Cloud (app.phoenix.arize.com)."
)
project_name = os.environ.get("PHOENIX_PROJECT_NAME", "default")
project_name = os.environ.get("PHOENIX_PROJECT_NAME") or "default"
return ArizePhoenixConfig(
otlp_auth_headers=otlp_auth_headers,
@ -343,8 +484,6 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
project_name=project_name,
)
## cannot suppress additional proxy server spans, removed previous methods.
async def async_health_check(self):
config = self.get_arize_phoenix_config()

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