Merge pull request #27436 from BerriAI/litellm_internal_staging
Some checks failed
Read Version from pyproject.toml / read-version (push) Has been cancelled
CodeQL / Analyze (actions) (push) Has been cancelled
CodeQL / Analyze (javascript-typescript) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CodSpeed Benchmarks / benchmarks (push) Has been cancelled
Helm unit test / unit-test (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
Unit Tests: Caching (Redis) / caching-redis (push) Has been cancelled
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Has been cancelled
Unit Tests: Security / security (push) Has been cancelled
GitHub Actions Security Analysis / zizmor (push) Has been cancelled
Unit Tests: Proxy DB Operations / logging-misc (push) Has been cancelled
Unit Tests: Proxy DB Operations / auth-checks (push) Has been cancelled
Unit Tests: Proxy DB Operations / budgets (push) Has been cancelled
Unit Tests: Proxy DB Operations / custom-logging (push) Has been cancelled
Unit Tests: Proxy DB Operations / db-and-spend (push) Has been cancelled
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Has been cancelled
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Has been cancelled
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Has been cancelled
Unit Tests: Proxy DB Operations / key-generation (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-runtime (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-server-core (push) Has been cancelled
Unit Tests: Proxy DB Operations / schema-migration (push) Has been cancelled
Unit Tests: Proxy DB Operations / proxy-utils (push) Has been cancelled

[Infra] Promote Internal Staging to main
This commit is contained in:
yuneng-jiang 2026-05-07 18:05:34 -07:00 committed by GitHub
commit fa81017e12
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
145 changed files with 14804 additions and 2136 deletions

View file

@ -350,7 +350,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/local_testing -x --junitxml=test-results/junit.xml --durations=5 -k "langfuse"
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"langfuse\""
no_output_timeout: 15m
# Store test results
- store_test_results:
@ -395,7 +403,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/proxy_admin_ui_tests -x --junitxml=test-results/junit.xml --durations=5 -n 2
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2"
no_output_timeout: 15m
# Store test results
@ -471,7 +487,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/router_unit_tests -x --junitxml=test-results/junit.xml --durations=5 -n 4
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
no_output_timeout: 15m
# Store test results
- store_test_results:
@ -495,7 +519,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest tests/local_testing/ -v -k "assistants" -x --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"assistants\""
no_output_timeout: 15m
# Store test results
- store_test_results:
@ -528,14 +560,19 @@ jobs:
# Add --timeout to kill hanging tests after 120s (2 min)
# Add --durations=20 to show 20 slowest tests for debugging
# Subdirectories with dedicated jobs (maintain this list as new jobs are added)
IGNORE_DIRS=(
"tests/llm_translation/realtime"
)
IGNORE_ARGS=""
for dir in "${IGNORE_DIRS[@]}"; do
IGNORE_ARGS="$IGNORE_ARGS --ignore=$dir"
done
uv run --no-sync python -m pytest -v tests/llm_translation $IGNORE_ARGS --junitxml=test-results/junit.xml --durations=20 -n 4 --timeout=120 --timeout_method=thread --retries 2 --retry-delay 5 --max-worker-restart=5
mkdir -p test-results
# Glob excludes the realtime/ subdirectory since it has its own job
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 \
-v \
--junitxml=test-results/junit.xml \
--durations=20 \
-n 4 \
--timeout=120 --timeout_method=thread \
--retries 2 --retry-delay 5 \
--max-worker-restart=5"
no_output_timeout: 15m
# Store test results
@ -560,7 +597,17 @@ jobs:
command: |
# Add --timeout to kill hanging tests after 120s (2 min)
# Add --durations=20 to show 20 slowest tests for debugging
uv run --no-sync python -m pytest -vv tests/llm_translation/realtime --cov=litellm --cov-report=xml -v --junitxml=test-results/junit.xml --durations=20 -n 4 --timeout=120 --timeout_method=thread
mkdir -p test-results
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 \
-vv \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
-n 4 \
--timeout=120 --timeout_method=thread"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -593,7 +640,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/agent_tests --ignore=tests/agent_tests/local_only_agent_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-vv -x -s \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -626,7 +681,18 @@ jobs:
- run:
name: Run tests
command: |
LITELLM_LOG=WARNING uv run --no-sync python -m pytest tests/guardrails_tests -vv --cov=litellm --cov-report=xml --junitxml=test-results/junit.xml --durations=5 -n 2 --timeout=120 --timeout_method=thread
mkdir -p test-results
export LITELLM_LOG=WARNING
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 \
-vv \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2 \
--timeout=120 --timeout_method=thread"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -660,7 +726,16 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/unified_google_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 --retries 3 --retry-delay 5
mkdir -p test-results
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 \
-vv -x -s \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
--retries 3 --retry-delay 5"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -702,7 +777,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/llm_responses_api_testing -x --junitxml=test-results/junit.xml --durations=5 -n 8
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 8"
no_output_timeout: 15m
# Store test results
@ -725,7 +808,16 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/ocr_tests --cov=litellm --cov-report=xml -x -v --junitxml=test-results/junit.xml --durations=5 -n 4
mkdir -p test-results
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 \
-vv -x \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -758,7 +850,16 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/search_tests --cov=litellm --cov-report=xml -x -v --junitxml=test-results/junit.xml --durations=5 -n 4
mkdir -p test-results
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 \
-vv -x \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -793,7 +894,15 @@ jobs:
name: Run enterprise tests
command: |
uv run --no-sync python -m prisma generate
uv run --no-sync python -m pytest -v tests/enterprise -x --junitxml=test-results/junit-enterprise.xml --durations=10 -n 4
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit-enterprise.xml \
--durations=10 \
-n 4"
no_output_timeout: 15m
# Store test results
- store_test_results:
@ -815,7 +924,16 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/batches_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 -n 2
mkdir -p test-results
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 \
-vv -x -s \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -848,7 +966,16 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/litellm_utils_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 -n 2
mkdir -p test-results
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 \
-vv -x -s \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -882,7 +1009,16 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/pass_through_unit_tests --cov=litellm --cov-report=xml -x -v --junitxml=test-results/junit.xml --durations=5 -n 4
mkdir -p test-results
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 \
-vv -x \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -916,7 +1052,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/image_gen_tests -n 4 -x --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
no_output_timeout: 15m
# Store test results
- store_test_results:
@ -939,7 +1083,18 @@ jobs:
- run:
name: Run tests
command: |
LITELLM_LOG=WARNING uv run --no-sync python -m pytest tests/logging_callback_tests -vv --cov=litellm --cov-report=xml -n 4 --junitxml=test-results/junit.xml --durations=5 --timeout=120 --timeout_method=thread
mkdir -p test-results
export LITELLM_LOG=WARNING
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 \
-vv \
--cov=litellm --cov-report=xml \
-n 4 \
--junitxml=test-results/junit.xml \
--durations=5 \
--timeout=120 --timeout_method=thread"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -972,7 +1127,15 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/audio_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-vv -x -s \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
- run:
name: Rename the coverage files
@ -1012,14 +1175,19 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv \
mkdir -p test-results
TEST_FILES=$(printf "%s\n" \
tests/local_testing/test_dual_cache.py \
tests/local_testing/test_redis_batch_optimizations.py \
tests/local_testing/test_router_utils.py \
--cov=litellm --cov-report=xml \
-x -s -v --junitxml=test-results/junit.xml \
--durations=5 -n 2 \
--reruns 2 --reruns-delay 1
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 \
-vv -x -s \
--cov=litellm --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 -n 2 \
--reruns 2 --reruns-delay 1"
no_output_timeout: 20m
- run:
name: Rename the coverage files
@ -1260,8 +1428,17 @@ jobs:
- run:
name: Run Basic Proxy Startup Tests (Health Readiness and Chat Completion)
command: |
uv run --no-sync python -m pytest -v tests/basic_proxy_startup_tests -x --junitxml=test-results/junit-2.xml --durations=5
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
- store_test_results:
path: test-results
build_and_test:
machine:
@ -1331,7 +1508,18 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -s -v tests/*.py -x --junitxml=test-results/junit.xml -n 4 --durations=5 --ignore=tests/otel_tests --ignore=tests/spend_tracking_tests --ignore=tests/pass_through_tests --ignore=tests/proxy_admin_ui_tests --ignore=tests/load_tests --ignore=tests/llm_translation --ignore=tests/llm_responses_api_testing --ignore=tests/mcp_tests --ignore=tests/guardrails_tests --ignore=tests/image_gen_tests --ignore=tests/pass_through_unit_tests
mkdir -p test-results
# Original used `tests/*.py` (top-level only); the `--ignore=...`
# flags were vestigial since shell globbing did not descend into
# subdirectories. Replicate by globbing only top-level test files.
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 \
-s -v -x \
--junitxml=test-results/junit.xml \
-n 4 \
--durations=5"
no_output_timeout: 15m
# Store test results
@ -1406,7 +1594,14 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -s -vv tests/openai_endpoints_tests --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-s -vv \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
# Store test results
@ -1475,7 +1670,14 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/otel_tests --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-v \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
# Clean up first container
- run:
@ -1518,7 +1720,14 @@ jobs:
- run:
name: Run second round of tests
command: |
uv run --no-sync python -m pytest -v tests/basic_proxy_startup_tests -x --junitxml=test-results/junit-2.xml --durations=5
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
# Store test results
@ -1587,8 +1796,17 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/spend_tracking_tests -x --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-vv -x \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
- store_test_results:
path: test-results
- run:
name: Stop and remove first container
when: always
@ -1676,7 +1894,14 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/multi_instance_e2e_tests -x --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-vv -x \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
# Clean up first container
# Store test results
@ -1732,7 +1957,14 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/store_model_in_db_tests -x --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-vv -x \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
- run:
name: Stop and remove containers
@ -1805,9 +2037,18 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/basic_proxy_startup_tests -x --junitxml=test-results/junit-2.xml --durations=5
mkdir -p test-results
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 \
-vv -x \
--junitxml=test-results/junit-2.xml \
--durations=5"
no_output_timeout: 15m
# Clean up first container
- store_test_results:
path: test-results
- run:
name: Stop and remove first container
command: |
@ -1940,7 +2181,14 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/pass_through_tests/ -x --junitxml=test-results/junit.xml --durations=5
mkdir -p test-results
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 \
-v -x \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
# Store test results
@ -1997,9 +2245,16 @@ jobs:
- run:
name: Run Claude Agent SDK E2E Tests
command: |
mkdir -p test-results
export LITELLM_PROXY_URL="http://localhost:4000"
export LITELLM_API_KEY="sk-1234"
uv run --no-sync python -m pytest -vv tests/proxy_e2e_anthropic_messages_tests/ -x -s --junitxml=test-results/junit.xml --durations=5
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 \
-vv -x -s \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
# Store test results

View file

@ -141,6 +141,7 @@ jobs:
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
tests/proxy_unit_tests/test_request_size_limit_middleware.py
workers: 4
dist: loadscope
timeout: 15

3
.gitignore vendored
View file

@ -100,4 +100,5 @@ STABILIZATION_TODO.md
**/playwright-report
**/*.storageState.json
**/coverage
test-config
test-config
.vscode

View file

@ -1,8 +1,8 @@
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
@ -68,8 +68,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervisor && \
npm install -g npm@11.12.1 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 && \
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##*/}"; \
@ -85,17 +85,17 @@ ENV PATH="/app/.venv/bin:${PATH}"
COPY --from=builder /app /app
# Prisma binaries live in $HOME/.cache (default prisma-python location),
# which is /root/.cache here. Copy them from the builder so they survive
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
COPY --from=builder /root/.cache /root/.cache
# 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
# setuptools wheel that surfaces as a CVE finding even though it's not
# on the runtime sys.path).
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
EXPOSE 4000/tcp
COPY docker/supervisord.conf /etc/supervisord.conf
ENTRYPOINT ["docker/prod_entrypoint.sh"]
CMD ["--port", "4000"]

View file

@ -1,3 +1,8 @@
codecov:
require_ci_to_pass: false # post coverage status even if CI has unrelated failures
notify:
wait_for_ci: false # post as soon as expected uploads arrive, don't wait on CI
component_management:
individual_components:
- component_id: "Router"
@ -28,7 +33,7 @@ coverage:
project:
default:
target: auto
threshold: 1% # at maximum allow project coverage to drop by 1%
threshold: 0% # do not allow project coverage to drop
patch:
default:
target: auto

View file

@ -116,21 +116,32 @@ spec:
name: {{ include "redis.secretName" .Subcharts.redis }}
key: {{include "redis.secretPasswordKey" .Subcharts.redis }}
{{- end }}
{{- /*
Inject LITELLM_LOG only when envVars does not already define it.
*/}}
{{- if and .Values.logLevel (not (hasKey (default dict .Values.envVars) "LITELLM_LOG")) }}
- name: LITELLM_LOG
value: {{ .Values.logLevel | quote }}
{{- end }}
{{- if .Values.envVars }}
{{- range $key, $val := .Values.envVars }}
- name: {{ $key }}
value: {{ $val | quote }}
{{- end }}
{{- end }}
{{- if .Values.separateHealthApp }}
- name: SEPARATE_HEALTH_APP
value: "1"
- name: SEPARATE_HEALTH_PORT
value: {{ .Values.separateHealthPort | default "8081" | quote }}
{{- end }}
{{- with .Values.extraEnvVars }}
{{- toYaml . | nindent 12 }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
# one DB on every rollout. Placed last (after envVars and
# extraEnvVars) so this override can't be silently shadowed by a
# user-supplied DISABLE_SCHEMA_UPDATE under last-wins duplicate-env
# semantics — same pattern the migrations Job uses.
- name: DISABLE_SCHEMA_UPDATE
value: "true"
{{- end }}
envFrom:
{{- range .Values.environmentSecrets }}
- secretRef:
@ -158,15 +169,10 @@ spec:
- name: http
containerPort: {{ .Values.service.port }}
protocol: TCP
{{- if .Values.separateHealthApp }}
- name: health
containerPort: {{ .Values.separateHealthPort | default 8081 }}
protocol: TCP
{{- end }}
livenessProbe:
httpGet:
path: {{ .Values.livenessProbe.path | quote }}
port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }}
port: "http"
initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds }}
periodSeconds: {{ .Values.livenessProbe.periodSeconds }}
timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds }}
@ -175,7 +181,7 @@ spec:
readinessProbe:
httpGet:
path: {{ .Values.readinessProbe.path | quote }}
port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }}
port: "http"
initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }}
periodSeconds: {{ .Values.readinessProbe.periodSeconds }}
timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }}
@ -184,7 +190,7 @@ spec:
startupProbe:
httpGet:
path: {{ .Values.startupProbe.path | quote }}
port: {{ if .Values.separateHealthApp }}"health"{{ else }}"http"{{ end }}
port: "http"
initialDelaySeconds: {{ .Values.startupProbe.initialDelaySeconds }}
periodSeconds: {{ .Values.startupProbe.periodSeconds }}
timeoutSeconds: {{ .Values.startupProbe.timeoutSeconds }}

View file

@ -257,16 +257,16 @@ tests:
value: 0
- equal:
path: spec.template.spec.containers[0].livenessProbe.periodSeconds
value: 10
value: 15
- equal:
path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds
value: 1
value: 5
- equal:
path: spec.template.spec.containers[0].livenessProbe.successThreshold
value: 1
- equal:
path: spec.template.spec.containers[0].livenessProbe.failureThreshold
value: 3
value: 5
- equal:
path: spec.template.spec.containers[0].readinessProbe.httpGet.path
value: /health/readiness
@ -278,7 +278,7 @@ tests:
value: 10
- equal:
path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds
value: 1
value: 5
- equal:
path: spec.template.spec.containers[0].readinessProbe.successThreshold
value: 1
@ -296,7 +296,7 @@ tests:
value: 10
- equal:
path: spec.template.spec.containers[0].startupProbe.timeoutSeconds
value: 1
value: 5
- equal:
path: spec.template.spec.containers[0].startupProbe.successThreshold
value: 1

View file

@ -88,26 +88,20 @@ service:
# optionally specify loadBalancerClass
# loadBalancerClass: tailscale
# Separate health app configuration
# When enabled, health checks will use a separate port and the application
# will receive SEPARATE_HEALTH_APP=1 and SEPARATE_HEALTH_PORT from environment variables
separateHealthApp: false
separateHealthPort: 8081
# Probe tuning for proxy container
# Probes for LiteLLM gateway container
livenessProbe:
path: /health/liveliness
initialDelaySeconds: 0
periodSeconds: 10
timeoutSeconds: 1
periodSeconds: 15
timeoutSeconds: 5
successThreshold: 1
failureThreshold: 3
failureThreshold: 5
readinessProbe:
path: /health/readiness
initialDelaySeconds: 0
periodSeconds: 10
timeoutSeconds: 1
timeoutSeconds: 5
successThreshold: 1
failureThreshold: 3
@ -115,7 +109,7 @@ startupProbe:
path: /health/readiness
initialDelaySeconds: 0
periodSeconds: 10
timeoutSeconds: 1
timeoutSeconds: 5
successThreshold: 1
failureThreshold: 30
@ -329,6 +323,17 @@ migrationJob:
helm:
enabled: false
# Log level for the litellm proxy (sets LITELLM_LOG in the deployment env).
# Rendered as a direct `env:` entry, which in Kubernetes takes precedence over
# any `envFrom:` source. If you currently source LITELLM_LOG from an
# environmentSecret or environmentConfigMap, set `logLevel: ""` here to
# disable injection — otherwise this value silently overrides your secret /
# configmap entry.
#
# Setting LITELLM_LOG inside `envVars:` below also wins: the template skips
# this injection entirely when envVars already defines LITELLM_LOG.
logLevel: INFO
# Additional environment variables to be added to the deployment as a map of key-value pairs
envVars: {}

View file

@ -66,7 +66,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervisor && \
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 \
@ -102,7 +102,5 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
EXPOSE 4000/tcp
COPY docker/supervisord.conf /etc/supervisord.conf
ENTRYPOINT ["docker/prod_entrypoint.sh"]
CMD ["--port", "4000"]

View file

@ -103,13 +103,12 @@ RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python3 bash openssl tzdata supervisor libsndfile nodejs && break || sleep 5; \
apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
done
COPY --from=builder /app /app
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets
COPY --from=builder /app/docker/supervisord.conf /etc/supervisord.conf
ENV PATH="/app/.venv/bin:${PATH}" \
PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \

View file

@ -1,14 +1,8 @@
#!/bin/sh
if [ "$SEPARATE_HEALTH_APP" = "1" ]; then
export LITELLM_ARGS="$@"
export SUPERVISORD_STOPWAITSECS="${SUPERVISORD_STOPWAITSECS:-3600}"
exec supervisord -c /etc/supervisord.conf
fi
if [ "$USE_DDTRACE" = "true" ]; then
export DD_TRACE_OPENAI_ENABLED="False"
exec ddtrace-run litellm "$@"
else
exec litellm "$@"
fi
fi

View file

@ -1,46 +0,0 @@
[supervisord]
nodaemon=true
loglevel=info
logfile=/tmp/supervisord.log
pidfile=/tmp/supervisord.pid
[group:litellm]
programs=main,health
[program:main]
command=sh -c 'if [ "$USE_DDTRACE" = "true" ]; then export DD_TRACE_OPENAI_ENABLED="False"; exec ddtrace-run python -m litellm.proxy.proxy_cli --host 0.0.0.0 --port=4000 $LITELLM_ARGS; else exec python -m litellm.proxy.proxy_cli --host 0.0.0.0 --port=4000 $LITELLM_ARGS; fi'
autostart=true
autorestart=true
startretries=3
priority=1
exitcodes=0
stopasgroup=true
killasgroup=true
stopwaitsecs=%(ENV_SUPERVISORD_STOPWAITSECS)s
stdout_logfile=/dev/stdout
stderr_logfile=/dev/stderr
stdout_logfile_maxbytes = 0
stderr_logfile_maxbytes = 0
environment=PYTHONUNBUFFERED=true
[program:health]
command=sh -c '[ "$SEPARATE_HEALTH_APP" = "1" ] && exec uvicorn litellm.proxy.health_endpoints.health_app_factory:build_health_app --factory --host 0.0.0.0 --port=${SEPARATE_HEALTH_PORT:-4001} || exit 0'
autostart=true
autorestart=true
startretries=3
priority=2
exitcodes=0
stopasgroup=true
killasgroup=true
stopwaitsecs=%(ENV_SUPERVISORD_STOPWAITSECS)s
stdout_logfile=/dev/stdout
stderr_logfile=/dev/stderr
stdout_logfile_maxbytes = 0
stderr_logfile_maxbytes = 0
environment=PYTHONUNBUFFERED=true
[eventlistener:process_monitor]
command=python -c "from supervisor import childutils; import os, signal; [os.kill(os.getppid(), signal.SIGTERM) for h,p in iter(lambda: childutils.listener.wait(), None) if h['eventname'] in ['PROCESS_STATE_FATAL', 'PROCESS_STATE_EXITED'] and dict([x.split(':') for x in p.split(' ')])['processname'] in ['main', 'health'] or childutils.listener.ok()]"
events=PROCESS_STATE_EXITED,PROCESS_STATE_FATAL
autostart=true
autorestart=true

Binary file not shown.

After

Width:  |  Height:  |  Size: 82 KiB

View file

@ -535,9 +535,9 @@
}
},
"node_modules/get-tsconfig": {
"version": "4.13.0",
"resolved": "https://registry.npmjs.org/get-tsconfig/-/get-tsconfig-4.13.0.tgz",
"integrity": "sha512-1VKTZJCwBrvbd+Wn3AOgQP/2Av+TfTCOlE4AcRJE72W1ksZXbAx8PPBR9RzgTeSPzlPMHrbANMH3LbltH73wxQ==",
"version": "4.14.0",
"resolved": "https://registry.npmjs.org/get-tsconfig/-/get-tsconfig-4.14.0.tgz",
"integrity": "sha512-yTb+8DXzDREzgvYmh6s9vHsSVCHeC0G3PI5bEXNBHtmshPnO+S5O7qgLEOn0I5QvMy6kpZN8K1NKGyilLb93wA==",
"dev": true,
"license": "MIT",
"dependencies": {

View file

@ -388,6 +388,7 @@ anthropic_beta_headers_url: str = os.getenv(
suppress_debug_info = False
dynamodb_table_name: Optional[str] = None
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
aws_sqs_callback_params: Optional[Dict] = None
@ -414,6 +415,9 @@ custom_prometheus_metadata_labels: List[str] = []
custom_prometheus_tags: List[str] = []
prometheus_metrics_config: Optional[List] = None
prometheus_emit_stream_label: bool = False
prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000
prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0
prometheus_end_user_metrics_cleanup_interval_seconds: Optional[float] = 60.0
disable_add_prefix_to_prompt: bool = (
False # used by anthropic, to disable adding prefix to prompt
)
@ -586,6 +590,7 @@ anyscale_models: Set = set()
cerebras_models: Set = set()
galadriel_models: Set = set()
nvidia_nim_models: Set = set()
nvidia_riva_models: Set = set()
sambanova_models: Set = set()
sambanova_embedding_models: Set = set()
novita_models: Set = set()
@ -812,6 +817,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
galadriel_models.add(key)
elif value.get("litellm_provider") == "nvidia_nim":
nvidia_nim_models.add(key)
elif value.get("litellm_provider") == "nvidia_riva":
nvidia_riva_models.add(key)
elif value.get("litellm_provider") == "sambanova":
sambanova_models.add(key)
elif value.get("litellm_provider") == "sambanova-embedding-models":
@ -971,6 +978,7 @@ model_list = list(
| cerebras_models
| galadriel_models
| nvidia_nim_models
| nvidia_riva_models
| sambanova_models
| azure_text_models
| novita_models
@ -1067,6 +1075,7 @@ models_by_provider: dict = {
"cerebras": cerebras_models,
"galadriel": galadriel_models,
"nvidia_nim": nvidia_nim_models,
"nvidia_riva": nvidia_riva_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"novita": novita_models,
"nebius": nebius_models | nebius_embedding_models,
@ -1618,6 +1627,9 @@ if TYPE_CHECKING:
from .llms.deepgram.audio_transcription.transformation import (
DeepgramAudioTranscriptionConfig as DeepgramAudioTranscriptionConfig,
)
from .llms.nvidia_riva.audio_transcription.transformation import (
NvidiaRivaAudioTranscriptionConfig as NvidiaRivaAudioTranscriptionConfig,
)
from .llms.topaz.image_variations.transformation import (
TopazImageVariationConfig as TopazImageVariationConfig,
)

View file

@ -161,6 +161,11 @@ MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset(
| (set(_MCP_STDIO_EXTRA_COMMANDS.split(",")) - {""})
)
# MCP OAuth2 Token Exchange (OBO) Defaults
MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE = int(
os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500")
)
LITELLM_UI_ALLOW_HEADERS = [
"x-litellm-semantic-filter",
"x-litellm-semantic-filter-tools",
@ -1457,6 +1462,12 @@ KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job"
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job"
SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(
os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
)
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float(
os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
)
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

@ -366,6 +366,8 @@ class MCPClient:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.token:
headers["Authorization"] = f"token {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.oauth2_token_exchange:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request

View file

@ -14,16 +14,18 @@ For batching specific details see CustomBatchLogger class
import asyncio
import os
import time
import traceback
from typing import List, Optional
from typing import List, Optional, Union
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
class AzureSentinelLogger(CustomBatchLogger):
@ -39,6 +41,7 @@ class AzureSentinelLogger(CustomBatchLogger):
tenant_id: Optional[str] = None,
client_id: Optional[str] = None,
client_secret: Optional[str] = None,
audit_stream_name: Optional[str] = None,
**kwargs,
):
"""
@ -57,57 +60,77 @@ class AzureSentinelLogger(CustomBatchLogger):
If not provided, will use AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID env var.
client_secret (str, optional): Azure Client Secret for OAuth2 authentication.
If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var.
audit_stream_name (str, optional): Stream name from DCR for audit logs.
If not provided, audit logs use the standard stream name.
"""
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.dcr_immutable_id = dcr_immutable_id or os.getenv(
resolved_dcr_immutable_id = dcr_immutable_id or os.getenv(
"AZURE_SENTINEL_DCR_IMMUTABLE_ID"
)
self.stream_name = stream_name or os.getenv(
"AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM"
resolved_stream_name = (
stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM"
)
self.endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
self.tenant_id = (
resolved_audit_stream_name = audit_stream_name or resolved_stream_name
resolved_endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
resolved_tenant_id = (
tenant_id
or os.getenv("AZURE_SENTINEL_TENANT_ID")
or os.getenv("AZURE_TENANT_ID")
)
self.client_id = (
resolved_client_id = (
client_id
or os.getenv("AZURE_SENTINEL_CLIENT_ID")
or os.getenv("AZURE_CLIENT_ID")
)
self.client_secret = (
resolved_client_secret = (
client_secret
or os.getenv("AZURE_SENTINEL_CLIENT_SECRET")
or os.getenv("AZURE_CLIENT_SECRET")
)
if not self.dcr_immutable_id:
if not resolved_dcr_immutable_id:
raise ValueError(
"AZURE_SENTINEL_DCR_IMMUTABLE_ID is required. Set it as an environment variable or pass dcr_immutable_id parameter."
)
if not self.endpoint:
if not resolved_endpoint:
raise ValueError(
"AZURE_SENTINEL_ENDPOINT is required. Set it as an environment variable or pass endpoint parameter."
)
if not self.tenant_id:
if not resolved_tenant_id:
raise ValueError(
"AZURE_SENTINEL_TENANT_ID or AZURE_TENANT_ID is required. Set it as an environment variable or pass tenant_id parameter."
)
if not self.client_id:
if not resolved_client_id:
raise ValueError(
"AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID is required. Set it as an environment variable or pass client_id parameter."
)
if not self.client_secret:
if not resolved_client_secret:
raise ValueError(
"AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET is required. Set it as an environment variable or pass client_secret parameter."
)
self.dcr_immutable_id = resolved_dcr_immutable_id
self.stream_name = resolved_stream_name
self.audit_stream_name = resolved_audit_stream_name
self.endpoint = resolved_endpoint
self.tenant_id = resolved_tenant_id
self.client_id = resolved_client_id
self.client_secret = resolved_client_secret
# Build API endpoint: {Endpoint}/dataCollectionRules/{DCR Immutable ID}/streams/{Stream Name}?api-version=2023-01-01
self.api_endpoint = f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01"
self.api_endpoint = self._build_api_endpoint(
endpoint=resolved_endpoint,
dcr_immutable_id=resolved_dcr_immutable_id,
stream_name=resolved_stream_name,
)
self.audit_api_endpoint = self._build_api_endpoint(
endpoint=resolved_endpoint,
dcr_immutable_id=resolved_dcr_immutable_id,
stream_name=resolved_audit_stream_name,
)
# OAuth2 scope for Azure Monitor
self.oauth_scope = "https://monitor.azure.com/.default"
@ -118,6 +141,13 @@ class AzureSentinelLogger(CustomBatchLogger):
super().__init__(**kwargs, flush_lock=self.flush_lock)
asyncio.create_task(self.periodic_flush())
self.log_queue: List[StandardLoggingPayload] = []
self.audit_log_queue: List[StandardAuditLogPayload] = []
@staticmethod
def _build_api_endpoint(
endpoint: str, dcr_immutable_id: str, stream_name: str
) -> str:
return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01"
async def _get_oauth_token(self) -> str:
"""
@ -126,9 +156,6 @@ class AzureSentinelLogger(CustomBatchLogger):
Returns:
Bearer token string
"""
# Check if we have a valid cached token
import time
if (
self.oauth_token
and self.oauth_token_expires_at
@ -170,9 +197,6 @@ class AzureSentinelLogger(CustomBatchLogger):
if not self.oauth_token:
raise Exception("OAuth2 token response did not contain access_token")
# Cache token expiry time
import time
self.oauth_token_expires_at = time.time() + expires_in
return self.oauth_token
@ -246,6 +270,34 @@ class AzureSentinelLogger(CustomBatchLogger):
)
pass
async def async_log_audit_log_event(
self, audit_log: StandardAuditLogPayload
) -> None:
"""
Async log LiteLLM audit log events to Azure Sentinel.
Audit logs are queued separately from standard LLM logs so mixed callback
usage never sends schema-mismatched records in the same ingestion batch.
"""
try:
verbose_logger.debug(
"Azure Sentinel: Logging audit event id=%s action=%s table=%s",
audit_log.get("id"),
audit_log.get("action"),
audit_log.get("table_name"),
)
self.audit_log_queue.append(audit_log)
if len(self.audit_log_queue) >= self.batch_size:
await self.async_send_audit_batch()
except Exception as e:
verbose_logger.exception(
f"Azure Sentinel Audit Log Layer Error - {str(e)}\n{traceback.format_exc()}"
)
pass
async def async_send_batch(self):
"""
Sends the batch of logs to Azure Monitor Logs Ingestion API
@ -253,22 +305,42 @@ class AzureSentinelLogger(CustomBatchLogger):
Raises:
Raises a NON Blocking verbose_logger.exception if an error occurs
"""
await self._async_send_batch_to_api(
log_queue=self.log_queue,
api_endpoint=self.api_endpoint,
log_type="logs",
)
async def async_send_audit_batch(self):
"""
Sends the batch of audit logs to Azure Monitor Logs Ingestion API
"""
await self._async_send_batch_to_api(
log_queue=self.audit_log_queue,
api_endpoint=self.audit_api_endpoint,
log_type="audit logs",
)
async def _async_send_batch_to_api(
self,
log_queue: List[Union[StandardLoggingPayload, StandardAuditLogPayload]],
api_endpoint: str,
log_type: str,
) -> None:
try:
if not self.log_queue:
if not log_queue:
return
verbose_logger.debug(
"Azure Sentinel - about to flush %s events", len(self.log_queue)
"Azure Sentinel - about to flush %s %s", len(log_queue), log_type
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
# Get OAuth2 token
bearer_token = await self._get_oauth_token()
# Convert log queue to JSON array format expected by Logs Ingestion API
# Each log entry should be a JSON object in the array
body = safe_dumps(self.log_queue)
body = safe_dumps(log_queue)
# Set headers for Logs Ingestion API
headers = {
@ -278,7 +350,7 @@ class AzureSentinelLogger(CustomBatchLogger):
# Send the request
response = await self.async_httpx_client.post(
url=self.api_endpoint, data=body.encode("utf-8"), headers=headers
url=api_endpoint, data=body.encode("utf-8"), headers=headers
)
if response.status_code not in [200, 204]:
@ -301,4 +373,15 @@ class AzureSentinelLogger(CustomBatchLogger):
f"Azure Sentinel Error sending batch API - {str(e)}\n{traceback.format_exc()}"
)
finally:
self.log_queue.clear()
log_queue.clear()
async def flush_queue(self):
if self.flush_lock is None:
return
async with self.flush_lock:
if self.log_queue:
await self.async_send_batch()
if self.audit_log_queue:
await self.async_send_audit_batch()
self.last_flush_time = time.time()

View file

@ -25,6 +25,9 @@ from typing import (
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import (
BoundedPrometheusSeriesTracker,
)
from litellm.integrations.prometheus_helpers import (
PrometheusLabelFactoryContext,
_get_cached_end_user_id_for_cost_tracking,
@ -81,6 +84,7 @@ class PrometheusLogger(CustomLogger):
if _custom_buckets is not None
else LATENCY_BUCKETS
)
self._bounded_prometheus_series_tracker = BoundedPrometheusSeriesTracker()
# Create metric factory functions
self._counter_factory = self._create_metric_factory(Counter)
@ -984,6 +988,40 @@ class PrometheusLogger(CustomLogger):
return filtered_labels
def _track_end_user_metric_series(
self,
metric: Any,
metric_name: DEFINED_PROMETHEUS_METRICS,
labels: Dict[str, Optional[str]],
) -> None:
"""
Cap the cardinality of metrics that include the ``end_user`` label.
Called *after* ``metric.labels(...).inc()/observe()`` so the emission is
recorded in prometheus-client's child map before any eviction runs.
Series that get evicted before the next scrape lose updates accrued
since the last scrape this is inherent to any cardinality cap.
"""
labelnames = self.get_labels_for_metric(metric_name)
if UserAPIKeyLabelNames.END_USER.value not in labelnames:
return
if labels.get(UserAPIKeyLabelNames.END_USER.value) is None:
return
max_series = litellm.prometheus_end_user_metrics_max_series_per_metric
ttl_seconds = litellm.prometheus_end_user_metrics_ttl_seconds
if max_series is None and ttl_seconds is None:
return
self._bounded_prometheus_series_tracker.track_series(
metric=metric,
metric_name=metric_name,
label_values=tuple(labels.get(label) for label in labelnames),
max_series=max_series,
ttl_seconds=ttl_seconds,
cleanup_interval_seconds=litellm.prometheus_end_user_metrics_cleanup_interval_seconds,
)
def _inc_labeled_counter(
self,
counter: Any,
@ -998,6 +1036,7 @@ class PrometheusLogger(CustomLogger):
label_context=label_context,
)
counter.labels(**_labels).inc(amount)
self._track_end_user_metric_series(counter, metric_name, _labels)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
# Define prometheus client
@ -1047,21 +1086,9 @@ class PrometheusLogger(CustomLogger):
output_tokens = standard_logging_payload["completion_tokens"]
tokens_used = standard_logging_payload["total_tokens"]
response_cost = standard_logging_payload["response_cost"]
_requester_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
"requester_metadata"
combined_metadata = _get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=standard_logging_payload
)
user_api_key_auth_metadata: Optional[dict] = standard_logging_payload[
"metadata"
].get("user_api_key_auth_metadata")
spend_logs_metadata: Optional[dict] = standard_logging_payload["metadata"].get(
"spend_logs_metadata"
)
combined_metadata: Dict[str, Any] = {
**(_requester_metadata if _requester_metadata else {}),
**(user_api_key_auth_metadata if user_api_key_auth_metadata else {}),
**(spend_logs_metadata if spend_logs_metadata else {}),
}
if standard_logging_payload is not None and isinstance(
standard_logging_payload, dict
):
@ -1416,26 +1443,46 @@ class PrometheusLogger(CustomLogger):
)
remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
remaining_requests = (
metadata.get(remaining_requests_variable_name, sys.maxsize) or sys.maxsize
remaining_requests = metadata.get(remaining_requests_variable_name)
if remaining_requests is None:
remaining_requests = sys.maxsize
remaining_tokens = metadata.get(remaining_tokens_variable_name)
if remaining_tokens is None:
remaining_tokens = sys.maxsize
enum_values = UserAPIKeyLabelValues(
hashed_api_key=user_api_key,
api_key_alias=user_api_key_alias,
model=model_group,
model_id=model_id,
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=_get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=kwargs.get("standard_logging_object")
)
),
)
remaining_tokens = (
metadata.get(remaining_tokens_variable_name, sys.maxsize) or sys.maxsize
label_context = PrometheusLabelFactoryContext(enum_values)
requests_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_requests_for_model"
),
enum_values=enum_values,
label_context=label_context,
)
self.litellm_remaining_api_key_requests_for_model.labels(**requests_labels).set(
remaining_requests
)
self.litellm_remaining_api_key_requests_for_model.labels(
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model_group),
_sanitize_prometheus_label_value(model_id),
).set(remaining_requests)
self.litellm_remaining_api_key_tokens_for_model.labels(
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model_group),
_sanitize_prometheus_label_value(model_id),
).set(remaining_tokens)
tokens_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
"litellm_remaining_api_key_tokens_for_model"
),
enum_values=enum_values,
label_context=label_context,
)
self.litellm_remaining_api_key_tokens_for_model.labels(**tokens_labels).set(
remaining_tokens
)
def _set_latency_metrics(
self,
@ -1471,6 +1518,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_llm_api_time_to_first_token_metric.labels(
**_ttft_labels
).observe(time_to_first_token_seconds)
self._track_end_user_metric_series(
self.litellm_llm_api_time_to_first_token_metric,
"litellm_llm_api_time_to_first_token_metric",
_ttft_labels,
)
else:
verbose_logger.debug(
"Time to first token metric not emitted, stream option in model_parameters is not True"
@ -1491,6 +1543,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_llm_api_latency_metric.labels(**_labels).observe(
api_call_total_time_seconds
)
self._track_end_user_metric_series(
self.litellm_llm_api_latency_metric,
"litellm_llm_api_latency_metric",
_labels,
)
# total request latency
total_time_seconds = self._safe_duration_seconds(
@ -1508,6 +1565,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_request_total_latency_metric.labels(**_labels).observe(
total_time_seconds
)
self._track_end_user_metric_series(
self.litellm_request_total_latency_metric,
"litellm_request_total_latency_metric",
_labels,
)
# request queue time (time from arrival to processing start)
_litellm_params = kwargs.get("litellm_params", {}) or {}
@ -1525,6 +1587,11 @@ class PrometheusLogger(CustomLogger):
self.litellm_request_queue_time_metric.labels(**_labels).observe(
queue_time_seconds
)
self._track_end_user_metric_series(
self.litellm_request_queue_time_metric,
"litellm_request_queue_time_seconds",
_labels,
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
verbose_logger.debug(
@ -1561,18 +1628,27 @@ class PrometheusLogger(CustomLogger):
)
try:
self.litellm_llm_api_failed_requests_metric.labels(
_sanitize_prometheus_label_value(end_user_id),
_sanitize_prometheus_label_value(user_api_key),
_sanitize_prometheus_label_value(user_api_key_alias),
_sanitize_prometheus_label_value(model),
_sanitize_prometheus_label_value(user_api_team),
_sanitize_prometheus_label_value(user_api_team_alias),
_sanitize_prometheus_label_value(user_id),
_sanitize_prometheus_label_value(
standard_logging_payload.get("model_id", "")
enum_values = UserAPIKeyLabelValues(
end_user=end_user_id,
hashed_api_key=user_api_key,
api_key_alias=user_api_key_alias,
model=model,
team=user_api_team,
team_alias=user_api_team_alias,
user=user_id,
model_id=standard_logging_payload.get("model_id", ""),
custom_metadata_labels=get_custom_labels_from_metadata(
metadata=_get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload=standard_logging_payload
)
),
).inc()
)
PrometheusLogger._inc_labeled_counter(
self,
self.litellm_llm_api_failed_requests_metric,
"litellm_llm_api_failed_requests_metric",
enum_values,
)
self.set_llm_deployment_failure_metrics(kwargs)
await self._set_org_budget_metrics_after_api_request(
org_id=user_api_key_org_id,
@ -3622,6 +3698,36 @@ def get_custom_labels_from_metadata(metadata: dict) -> Dict[str, str]:
return result
def _get_combined_custom_metadata_from_standard_logging_payload(
standard_logging_payload: Optional[dict],
) -> Dict[str, Any]:
"""
Combine the metadata sources that can supply custom Prometheus labels.
"""
if not isinstance(standard_logging_payload, dict):
return {}
standard_logging_metadata = standard_logging_payload.get("metadata") or {}
if not isinstance(standard_logging_metadata, dict):
return {}
requester_metadata = standard_logging_metadata.get("requester_metadata")
user_api_key_auth_metadata = standard_logging_metadata.get(
"user_api_key_auth_metadata"
)
spend_logs_metadata = standard_logging_metadata.get("spend_logs_metadata")
return {
**(requester_metadata if isinstance(requester_metadata, dict) else {}),
**(
user_api_key_auth_metadata
if isinstance(user_api_key_auth_metadata, dict)
else {}
),
**(spend_logs_metadata if isinstance(spend_logs_metadata, dict) else {}),
}
def _tag_matches_wildcard_configured_pattern(
tags: Sequence[str], configured_tag: str
) -> bool:

View file

@ -0,0 +1,107 @@
from __future__ import annotations
import time
from collections import OrderedDict
from threading import RLock
from typing import Any, Dict, Optional
class BoundedPrometheusSeriesTracker:
"""
Tracks Prometheus child series and removes stale/excess labelsets.
The tracker is label-agnostic: callers decide which series should be tracked
and pass the full label tuple used by the Prometheus metric.
"""
def __init__(self) -> None:
self._series: Dict[str, OrderedDict[tuple[Optional[str], ...], float]] = {}
self._last_ttl_cleanup: Dict[str, float] = {}
self.lock = RLock()
def track_series(
self,
metric: Any,
metric_name: str,
label_values: tuple[Optional[str], ...],
max_series: Optional[int],
ttl_seconds: Optional[float],
cleanup_interval_seconds: Optional[float],
) -> None:
if max_series is None and ttl_seconds is None:
return
now = time.monotonic()
with self.lock:
series = self._series.setdefault(metric_name, OrderedDict())
series[label_values] = now
series.move_to_end(label_values)
if ttl_seconds is not None and self._should_run_ttl_cleanup(
metric_name=metric_name,
now=now,
cleanup_interval_seconds=cleanup_interval_seconds,
):
expired_label_values = [
tracked_label_values
for tracked_label_values, last_seen in series.items()
if now - last_seen > ttl_seconds
]
for tracked_label_values in expired_label_values:
self._remove_metric_series(metric, series, tracked_label_values)
# max_series <= 0 is treated as "unlimited" so a misconfigured zero
# value cannot silently drop every emission for this metric.
if max_series is not None and max_series > 0:
while len(series) > max_series:
tracked_label_values = next(iter(series))
if not self._remove_metric_child(metric, tracked_label_values):
break
del series[tracked_label_values]
def _should_run_ttl_cleanup(
self,
metric_name: str,
now: float,
cleanup_interval_seconds: Optional[float],
) -> bool:
if cleanup_interval_seconds is None or cleanup_interval_seconds <= 0:
self._last_ttl_cleanup[metric_name] = now
return True
last_cleanup = self._last_ttl_cleanup.get(metric_name)
if last_cleanup is None or now - last_cleanup >= cleanup_interval_seconds:
self._last_ttl_cleanup[metric_name] = now
return True
return False
def _remove_metric_series(
self,
metric: Any,
series: OrderedDict[tuple[Optional[str], ...], float],
label_values: tuple[Optional[str], ...],
) -> None:
if self._remove_metric_child(metric, label_values):
series.pop(label_values, None)
@staticmethod
def _remove_metric_child(
metric: Any, label_values: tuple[Optional[str], ...]
) -> bool:
"""
Remove the Prometheus child for ``label_values`` and report whether the
tracker should commit the matching state change.
Returns ``True`` when the child is no longer present in Prometheus
(either it was just removed or it was already gone), and ``False`` when
``metric.remove()`` raised an unexpected error and the child likely
still exists.
"""
try:
metric.remove(*label_values)
return True
except KeyError:
return True
except (AttributeError, ValueError):
return False

View file

@ -16,6 +16,7 @@ from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
from litellm.integrations.s3 import get_s3_object_key
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
@ -53,15 +54,25 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_strip_base64_files: bool = False,
s3_use_key_prefix: bool = False,
s3_use_virtual_hosted_style: bool = False,
s3_callback_params_override: Optional[dict] = None,
**kwargs,
):
try:
verbose_logger.debug(
f"in init s3 logger - s3_callback_params {litellm.s3_callback_params}"
)
_masker = SensitiveDataMasker()
if s3_callback_params_override is not None:
verbose_logger.debug(
f"in init s3 logger (audit override) - "
f"{_masker.mask_dict(dict(s3_callback_params_override))}"
)
else:
verbose_logger.debug(
f"in init s3 logger - s3_callback_params "
f"{_masker.mask_dict(dict(litellm.s3_callback_params or {}))}"
)
# Initialize S3 params first to get the correct s3_verify value
self._init_s3_params(
params_source=s3_callback_params_override,
s3_bucket_name=s3_bucket_name,
s3_region_name=s3_region_name,
s3_api_version=s3_api_version,
@ -139,94 +150,85 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
s3_strip_base64_files: bool = False,
s3_use_key_prefix: bool = False,
s3_use_virtual_hosted_style: bool = False,
params_source: Optional[dict] = None,
):
"""
Initialize the s3 params for this logging callback
Initialize the s3 params for this logging callback. Reads from
`params_source` if given (e.g. `s3_audit_callback_params` for the
audit-log instance), otherwise falls back to `litellm.s3_callback_params`.
Resolves `os.environ/X` markers into a local dict; never mutates the source.
"""
litellm.s3_callback_params = litellm.s3_callback_params or {}
# read in .env variables - example os.environ/AWS_BUCKET_NAME
for key, value in litellm.s3_callback_params.items():
if isinstance(value, str) and value.startswith("os.environ/"):
litellm.s3_callback_params[key] = litellm.get_secret(value)
if params_source is None:
params_source = litellm.s3_callback_params or {}
params: dict = {
key: (
litellm.get_secret(value)
if isinstance(value, str) and value.startswith("os.environ/")
else value
)
for key, value in params_source.items()
}
self.s3_bucket_name = (
litellm.s3_callback_params.get("s3_bucket_name") or s3_bucket_name
)
self.s3_region_name = (
litellm.s3_callback_params.get("s3_region_name") or s3_region_name
)
self.s3_api_version = (
litellm.s3_callback_params.get("s3_api_version") or s3_api_version
)
self.s3_bucket_name = params.get("s3_bucket_name") or s3_bucket_name
self.s3_region_name = params.get("s3_region_name") or s3_region_name
self.s3_api_version = params.get("s3_api_version") or s3_api_version
self.s3_use_ssl = (
litellm.s3_callback_params.get("s3_use_ssl", True)
if litellm.s3_callback_params.get("s3_use_ssl") is not None
params.get("s3_use_ssl", True)
if params.get("s3_use_ssl") is not None
else s3_use_ssl
)
self.s3_verify = (
litellm.s3_callback_params.get("s3_verify")
if litellm.s3_callback_params.get("s3_verify") is not None
params.get("s3_verify")
if params.get("s3_verify") is not None
else s3_verify
)
self.s3_endpoint_url = (
litellm.s3_callback_params.get("s3_endpoint_url") or s3_endpoint_url
)
self.s3_endpoint_url = params.get("s3_endpoint_url") or s3_endpoint_url
self.s3_aws_access_key_id = (
litellm.s3_callback_params.get("s3_aws_access_key_id")
or s3_aws_access_key_id
params.get("s3_aws_access_key_id") or s3_aws_access_key_id
)
self.s3_aws_secret_access_key = (
litellm.s3_callback_params.get("s3_aws_secret_access_key")
or s3_aws_secret_access_key
params.get("s3_aws_secret_access_key") or s3_aws_secret_access_key
)
self.s3_aws_session_token = (
litellm.s3_callback_params.get("s3_aws_session_token")
or s3_aws_session_token
params.get("s3_aws_session_token") or s3_aws_session_token
)
self.s3_aws_session_name = (
litellm.s3_callback_params.get("s3_aws_session_name") or s3_aws_session_name
params.get("s3_aws_session_name") or s3_aws_session_name
)
self.s3_aws_profile_name = (
litellm.s3_callback_params.get("s3_aws_profile_name") or s3_aws_profile_name
params.get("s3_aws_profile_name") or s3_aws_profile_name
)
self.s3_aws_role_name = (
litellm.s3_callback_params.get("s3_aws_role_name") or s3_aws_role_name
)
self.s3_aws_role_name = params.get("s3_aws_role_name") or s3_aws_role_name
self.s3_aws_web_identity_token = (
litellm.s3_callback_params.get("s3_aws_web_identity_token")
or s3_aws_web_identity_token
params.get("s3_aws_web_identity_token") or s3_aws_web_identity_token
)
self.s3_aws_sts_endpoint = (
litellm.s3_callback_params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint
params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint
)
self.s3_config = litellm.s3_callback_params.get("s3_config") or s3_config
self.s3_path = litellm.s3_callback_params.get("s3_path") or s3_path
# done reading litellm.s3_callback_params
self.s3_config = params.get("s3_config") or s3_config
self.s3_path = params.get("s3_path") or s3_path
self.s3_use_team_prefix = (
bool(litellm.s3_callback_params.get("s3_use_team_prefix", False))
or s3_use_team_prefix
bool(params.get("s3_use_team_prefix", False)) or s3_use_team_prefix
)
self.s3_use_key_prefix = (
bool(litellm.s3_callback_params.get("s3_use_key_prefix", False))
or s3_use_key_prefix
bool(params.get("s3_use_key_prefix", False)) or s3_use_key_prefix
)
self.s3_strip_base64_files = (
bool(litellm.s3_callback_params.get("s3_strip_base64_files", False))
or s3_strip_base64_files
bool(params.get("s3_strip_base64_files", False)) or s3_strip_base64_files
)
self.s3_use_virtual_hosted_style = (
bool(litellm.s3_callback_params.get("s3_use_virtual_hosted_style", False))
bool(params.get("s3_use_virtual_hosted_style", False))
or s3_use_virtual_hosted_style
)

View file

@ -14,6 +14,7 @@ from litellm import _custom_logger_compatible_callbacks_literal
from litellm.integrations.agentops import AgentOps
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.braintrust_logging import BraintrustLogger
@ -73,6 +74,7 @@ class CustomLoggerRegistry:
"opik": OpikLogger,
"argilla": ArgillaLogger,
"opentelemetry": OpenTelemetry,
"azure_sentinel": AzureSentinelLogger,
"azure_storage": AzureBlobStorageLogger,
"humanloop": HumanloopLogger,
# OTEL compatible loggers

View file

@ -621,6 +621,18 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
or "https://integrate.api.nvidia.com/v1"
) # type: ignore
dynamic_api_key = api_key or get_secret_str("NVIDIA_NIM_API_KEY")
elif custom_llm_provider == "nvidia_riva":
# NVIDIA Riva is gRPC-based; api_base must be a host:port like
# `grpc.nvcf.nvidia.com:443` or `localhost:50051`. There is no
# public-default endpoint, so we do not fill one in here.
api_base = api_base or get_secret_str("NVIDIA_RIVA_API_BASE") # type: ignore
# Fall back to NVIDIA_NIM_API_KEY because users running both NVCF
# services typically reuse the same nvapi-* key.
dynamic_api_key = (
api_key
or get_secret_str("NVIDIA_RIVA_API_KEY")
or get_secret_str("NVIDIA_NIM_API_KEY")
)
elif custom_llm_provider == "cerebras":
api_base = (
api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1"

View file

@ -436,12 +436,21 @@ def update_messages_with_model_file_ids(
"""
Updates messages with model file ids.
For managed files (unified file IDs), uses model_file_id_mapping if it
resolves the id, otherwise decodes the base64-encoded unified file ID
and extracts the llm_output_file_id directly. Mirrors the Responses-API
sibling `update_responses_input_with_model_file_ids`.
model_file_id_mapping: Dict[str, Dict[str, str]] = {
"litellm_proxy/file_id": {
"model_id": "provider_file_id"
}
}
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
convert_b64_uid_to_unified_uid,
)
for message in messages:
if message.get("role") == "user":
@ -450,7 +459,13 @@ def update_messages_with_model_file_ids(
if isinstance(content, str):
continue
for c in content:
if c["type"] == "file":
if not isinstance(c, dict):
# Content list items aren't always dicts. e.g.
# text_completion forwards a token-ids list/list-of-
# lists through this path. Skip non-dict items
# instead of indexing into them.
continue
if c.get("type") == "file":
file_object = cast(ChatCompletionFileObject, c)
file_object_file_field = file_object.get("file")
if not isinstance(file_object_file_field, dict):
@ -468,9 +483,23 @@ def update_messages_with_model_file_ids(
if file_id:
provider_file_id = (
model_file_id_mapping.get(file_id, {}).get(model_id)
or file_id
if model_file_id_mapping
else None
)
if (
not provider_file_id
and _is_base64_encoded_unified_file_id(file_id)
):
unified_file_id = convert_b64_uid_to_unified_uid(
file_id
)
if "llm_output_file_id," in unified_file_id:
provider_file_id = unified_file_id.split(
"llm_output_file_id,"
)[1].split(";")[0]
file_object_file_field["file_id"] = (
provider_file_id or file_id
)
file_object_file_field["file_id"] = provider_file_id
if format:
file_object_file_field["format"] = format
return messages

View file

@ -2124,27 +2124,62 @@ def anthropic_process_openai_file_message(
)
_EMPTY_TEXT_PLACEHOLDER = (
"[System: Empty message content sanitised to satisfy protocol]"
)
def _sanitize_empty_text_content(
message: AllMessageValues,
) -> AllMessageValues:
"""
Case C: Sanitize empty text content
- Replace empty or whitespace-only text content with a placeholder message.
- Handles both string content and list-of-blocks content (rewriting only
the empty text blocks in place; non-text blocks like images are left
untouched).
Returns:
The message with sanitized content if needed, otherwise the original message
"""
if message.get("role") in ["user", "assistant"]:
content = message.get("content")
if isinstance(content, str):
if not content or not content.strip():
message = cast(AllMessageValues, dict(message)) # Make a copy
message["content"] = (
"[System: Empty message content sanitised to satisfy protocol]"
)
verbose_logger.debug(
f"_sanitize_empty_text_content: Replaced empty text content in {message.get('role')} message"
)
if message.get("role") not in ["user", "assistant"]:
return message
content = message.get("content")
if isinstance(content, str):
if not content or not content.strip():
message = cast(AllMessageValues, dict(message)) # Make a copy
message["content"] = _EMPTY_TEXT_PLACEHOLDER
verbose_logger.debug(
f"_sanitize_empty_text_content: Replaced empty text content in {message.get('role')} message"
)
return message
if isinstance(content, list):
# Walk the blocks and rewrite any empty text blocks. We rewrite (rather
# than drop) so callers don't end up with an entirely empty content
# list, which Anthropic also rejects.
new_blocks: List[Any] = []
rewrote_any = False
for block in content:
if isinstance(block, dict) and block.get("type") == "text":
text = block.get("text")
if not isinstance(text, str) or not text or not text.strip():
new_block = dict(block)
new_block["text"] = _EMPTY_TEXT_PLACEHOLDER
new_blocks.append(new_block)
rewrote_any = True
continue
new_blocks.append(block)
if rewrote_any:
message = cast(AllMessageValues, dict(message)) # Make a copy
message["content"] = new_blocks # type: ignore
verbose_logger.debug(
f"_sanitize_empty_text_content: Replaced empty text block(s) in {message.get('role')} message"
)
return message
@ -2427,6 +2462,18 @@ def anthropic_messages_pt( # noqa: PLR0915
# Sanitize messages for tool calling issues when modify_params=True
messages = sanitize_messages_for_tool_calling(messages)
# Anthropic rejects empty text content blocks with:
# "messages: text content blocks must be non-empty"
# OpenAI/other providers silently tolerate `{"role": "user", "content": ""}`,
# so callers (and upstream agent frameworks like pydantic-ai) routinely
# send empty user/assistant turns. We always rewrite these to a placeholder
# for Anthropic-shaped requests, independent of `litellm.modify_params`,
# because there is no way to "pass through" an empty text block — the
# request will always 400 otherwise. The richer tool-call sanitization
# (Cases A/B/D in `sanitize_messages_for_tool_calling`) remains gated on
# `modify_params` because it actually mutates conversation structure.
messages = [_sanitize_empty_text_content(m) for m in messages]
# add role=tool support to allow function call result/error submission
user_message_types = {"user", "tool", "function"}
# reformat messages to ensure user/assistant are alternating, if there's either 2 consecutive 'user' messages or 2 consecutive 'assistant' message, merge them.

View file

@ -31,6 +31,8 @@ DefaultLoggedRealTimeEventTypes = [
"session.created",
"response.create",
"response.done",
"conversation.item.added", # GA
"conversation.item.done", # GA
]
@ -44,6 +46,7 @@ class RealTimeStreaming:
model: str = "",
user_api_key_dict: Optional[Any] = None,
request_data: Optional[Dict] = None,
backend_uses_beta_protocol: Optional[bool] = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
@ -54,6 +57,14 @@ class RealTimeStreaming:
self.session_tools: List[Dict] = []
self.tool_calls: List[Dict] = []
# Detect whether the client is explicitly opting into the beta protocol.
self._client_wants_beta = self._detect_beta_header(websocket)
self._backend_uses_beta_protocol = (
self._client_wants_beta
if backend_uses_beta_protocol is None
else backend_uses_beta_protocol
)
_logged_real_time_event_types = litellm.logged_real_time_event_types
if _logged_real_time_event_types is None:
@ -76,6 +87,27 @@ class RealTimeStreaming:
# response.create can be rewritten to include the failure context.
self._pending_guardrail_message: Optional[str] = None
_SESSION_EVENT_TYPES = frozenset(["session.created", "session.updated"])
_AUDIO_FORMAT_MAP: Dict[str, Dict[str, Any]] = {
"pcm16": {"type": "audio/pcm", "rate": 24000},
"g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
"g711_alaw": {"type": "audio/G711-alaw", "rate": 8000},
}
# GA name → beta name (when client WebSocket includes OpenAI-Beta: realtime=v1)
_GA_TO_BETA_EVENT_TYPES: Dict[str, str] = {
"conversation.item.added": "conversation.item.created",
"response.output_text.delta": "response.text.delta",
"response.output_audio.delta": "response.audio.delta",
"response.output_audio_transcript.delta": "response.audio_transcript.delta",
"response.output_text.done": "response.text.done",
"response.output_audio.done": "response.audio.done",
"response.output_audio_transcript.done": "response.audio_transcript.done",
}
_GA_TO_BETA_CONTENT_TYPES: Dict[str, str] = {
"output_text": "text",
"output_audio": "audio",
}
def _should_store_message(
self,
message_obj: Union[dict, OpenAIRealtimeEvents],
@ -92,24 +124,27 @@ class RealTimeStreaming:
if isinstance(message, bytes):
message = message.decode("utf-8")
if isinstance(message, dict):
message_obj = message
# TypedDict union members do not narrow to plain dict for mypy.
message_obj: Dict[str, Any] = cast(Dict[str, Any], message)
else:
message_obj = json.loads(message)
message_obj = cast(Dict[str, Any], json.loads(cast(str, message)))
self._collect_tool_calls_from_response_done(cast(dict, message_obj))
try:
if (
not isinstance(message, dict)
or message_obj.get("type") == "session.created"
or message_obj.get("type") == "session.updated"
):
message_obj = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore
elif not isinstance(message, dict):
message_obj = OpenAIRealtimeStreamResponseBaseObject(**message_obj) # type: ignore
event_type = message_obj.get("type", "")
if event_type in self._SESSION_EVENT_TYPES:
typed_obj = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore
else:
# Use the base object as a safe catch-all for all other event types
# (both beta and GA), so unknown/new event names never raise here.
typed_obj = OpenAIRealtimeStreamResponseBaseObject(**message_obj) # type: ignore
except Exception as e:
verbose_logger.debug(f"Error parsing message for logging: {e}")
raise e
if self._should_store_message(message_obj):
self.messages.append(message_obj)
# Don't re-raise — a parse failure must not drop or delay the message
if self._should_store_message(message_obj):
self.messages.append(message_obj) # type: ignore[arg-type]
return
if self._should_store_message(typed_obj):
self.messages.append(typed_obj)
def _collect_user_input_from_client_event(self, message: Union[str, dict]) -> None:
"""Extract user text content from client WebSocket events for spend logging."""
@ -147,6 +182,8 @@ class RealTimeStreaming:
tools = session.get("tools")
if tools and isinstance(tools, list):
self.session_tools = tools
# GA: session.type is required; log it for traceability but no action needed
verbose_logger.debug(f"Realtime session.type: {session.get('type')}")
except (json.JSONDecodeError, AttributeError, TypeError):
pass
@ -228,6 +265,23 @@ class RealTimeStreaming:
else:
await self.backend_ws.send(message) # type: ignore[union-attr, attr-defined]
def _make_disable_auto_response_message(self) -> str:
"""Return a session.update that disables VAD auto-response."""
if self._backend_uses_beta_protocol:
session: Dict[str, Any] = {
"turn_detection": {"create_response": False},
}
else:
session = {
"type": "realtime",
"audio": {
"input": {
"turn_detection": {"create_response": False},
}
},
}
return json.dumps({"type": "session.update", "session": session})
def _has_realtime_guardrails(self) -> bool:
"""Return True if any callback is registered for realtime guardrail event types."""
from litellm.integrations.custom_guardrail import CustomGuardrail
@ -435,14 +489,7 @@ class RealTimeStreaming:
):
self.store_message(event_str)
await self.websocket.send_text(event_str)
await self._send_to_backend(
json.dumps(
{
"type": "session.update",
"session": {"turn_detection": {"create_response": False}},
}
)
)
await self._send_to_backend(self._make_disable_auto_response_message())
continue
## GUARDRAIL: run on transcription events in provider_config path too
if (
@ -484,14 +531,7 @@ class RealTimeStreaming:
):
self.store_message(raw_response)
await self.websocket.send_text(raw_response)
await self._send_to_backend(
json.dumps(
{
"type": "session.update",
"session": {"turn_detection": {"create_response": False}},
}
)
)
await self._send_to_backend(self._make_disable_auto_response_message())
return True
if (
@ -542,7 +582,20 @@ class RealTimeStreaming:
continue
## LOGGING
self.store_message(raw_response)
await self.websocket.send_text(raw_response)
# If the client opted into beta protocol, translate GA event
# names/shapes back to the beta equivalents before forwarding.
if self._client_wants_beta:
try:
event_dict = json.loads(raw_response)
translated = self._translate_event_to_beta(event_dict)
if translated is None:
continue # drop GA-only events (e.g. conversation.item.done)
await self.websocket.send_text(json.dumps(translated))
except Exception:
await self.websocket.send_text(raw_response)
else:
await self.websocket.send_text(raw_response)
except websockets.exceptions.ConnectionClosed as e: # type: ignore
verbose_logger.exception(
@ -553,6 +606,183 @@ class RealTimeStreaming:
finally:
await self.log_messages()
@staticmethod
def _detect_beta_header(websocket: Any) -> bool:
"""Return True if the client sent 'OpenAI-Beta: realtime=v1'.
Checks the raw ASGI scope headers so it works for both FastAPI WebSocket
objects and any test doubles that expose a .scope dict.
"""
try:
headers = websocket.scope.get("headers", [])
for name, value in headers:
if isinstance(name, bytes):
name = name.decode("latin-1")
if isinstance(value, bytes):
value = value.decode("latin-1")
if name.lower() == "openai-beta" and "realtime=v1" in value.lower():
return True
except Exception:
pass
return False
@staticmethod
def _remap_beta_session_to_ga(session: dict) -> dict:
"""
Convert a beta-style session.update payload to the GA nested schema.
Beta GA field mappings
session.type (inject "realtime" if absent)
session.modalities session.output_modalities
session.voice session.audio.output.voice
session.input_audio_format session.audio.input.format (with type/rate)
session.output_audio_format session.audio.output.format (with type/rate)
session.turn_detection session.audio.input.turn_detection
session.input_audio_transcription session.audio.input.transcription
Fields not in the mapping (instructions, tools, etc.) are passed through.
GA clients that already use the nested shape are unaffected.
"""
# Work on a shallow copy so we don't mutate the caller's dict
session = dict(session)
# 1. Ensure session.type is present
if "type" not in session:
session["type"] = "realtime"
# 2. Rename modalities → output_modalities and normalise combinations.
# Beta allowed ["audio", "text"] together; GA only supports ["audio"] or
# ["text"] as single-element lists. When both are present we prefer
# ["audio"] because audio mode already delivers transcripts via events.
if "modalities" in session:
mods = session.pop("modalities")
if "output_modalities" not in session:
mods_set = {m.lower() for m in (mods or [])}
if "audio" in mods_set:
session["output_modalities"] = ["audio"]
elif "text" in mods_set:
session["output_modalities"] = ["text"]
# 3-7. Lift flat audio fields into the nested audio object
audio: Dict[str, Any] = {}
inp: Dict[str, Any] = {}
out: Dict[str, Any] = {}
# voice → audio.output.voice
if "voice" in session:
out["voice"] = session.pop("voice")
# input_audio_format → audio.input.format
if "input_audio_format" in session:
raw = session.pop("input_audio_format")
inp["format"] = (
RealTimeStreaming._AUDIO_FORMAT_MAP.get(raw, raw)
if isinstance(raw, str)
else raw
)
# output_audio_format → audio.output.format
if "output_audio_format" in session:
raw = session.pop("output_audio_format")
out["format"] = (
RealTimeStreaming._AUDIO_FORMAT_MAP.get(raw, raw)
if isinstance(raw, str)
else raw
)
# turn_detection → audio.input.turn_detection
if "turn_detection" in session:
inp["turn_detection"] = session.pop("turn_detection")
# input_audio_transcription → audio.input.transcription
if "input_audio_transcription" in session:
inp["transcription"] = session.pop("input_audio_transcription")
if inp:
audio["input"] = inp
if out:
audio["output"] = out
if audio:
# Merge with any existing GA-style `audio` block the client already set,
# letting the remapped values take precedence within each sub-key.
existing = session.get("audio") or {}
for sub_key, sub_val in audio.items():
if (
sub_key in existing
and isinstance(existing[sub_key], dict)
and isinstance(sub_val, dict)
):
existing[sub_key] = {**existing[sub_key], **sub_val}
else:
existing[sub_key] = sub_val
session["audio"] = existing
return session
@staticmethod
def _translate_event_to_beta(event: dict) -> Optional[dict]:
"""Translate a single GA event dict to its beta equivalent.
Returns None if the event should be dropped entirely (e.g. the GA-only
conversation.item.done has no beta counterpart).
Returns the (possibly mutated copy of the) event otherwise.
"""
event_type = event.get("type", "")
# conversation.item.done has no beta equivalent — the client already
# received conversation.item.created (translated from .added).
if event_type == "conversation.item.done":
return None
# Shallow-copy so we don't mutate the stored message
translated = dict(event)
# Rename the type field
if event_type in RealTimeStreaming._GA_TO_BETA_EVENT_TYPES:
translated["type"] = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES[event_type]
# Fix content block types inside items (response.done output list,
# conversation.item.created item content, etc.)
if "item" in translated and isinstance(translated["item"], dict):
translated["item"] = RealTimeStreaming._translate_item_content_types(
dict(translated["item"])
)
if "response" in translated and isinstance(translated["response"], dict):
resp = dict(translated["response"])
if "output" in resp and isinstance(resp["output"], list):
resp["output"] = [
(
RealTimeStreaming._translate_item_content_types(dict(o))
if isinstance(o, dict)
else o
)
for o in resp["output"]
]
translated["response"] = resp
return translated
@staticmethod
def _translate_item_content_types(item: dict) -> dict:
"""Replace GA content type names with beta names inside a single item."""
if "content" not in item or not isinstance(item["content"], list):
return item
new_content = []
for block in item["content"]:
if (
isinstance(block, dict)
and block.get("type") in RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES
):
block = dict(block)
block["type"] = RealTimeStreaming._GA_TO_BETA_CONTENT_TYPES[
block["type"]
]
new_content.append(block)
item["content"] = new_content
return item
async def client_ack_messages(self):
try:
while True:
@ -594,6 +824,19 @@ class RealTimeStreaming:
self._pending_guardrail_message = None
continue
# GA compatibility: remap beta-style session fields only when
# the upstream is in GA mode. Beta upstreams expect the flat
# session shape unchanged.
if (
msg_type == "session.update"
and not self._backend_uses_beta_protocol
):
session = msg_obj.get("session", {})
if isinstance(session, dict):
session = self._remap_beta_session_to_ga(session)
msg_obj["session"] = session
message = json.dumps(msg_obj)
except (json.JSONDecodeError, AttributeError):
pass
@ -627,3 +870,8 @@ class RealTimeStreaming:
await forward_task
except asyncio.CancelledError:
pass
def client_sent_openai_beta_realtime_header(websocket: Any) -> bool:
"""True when the client WebSocket includes ``OpenAI-Beta: realtime=v1``."""
return RealTimeStreaming._detect_beta_header(websocket)

View file

@ -65,7 +65,7 @@ from litellm.types.utils import (
from ...base import BaseLLM
from ..common_utils import AnthropicError, process_anthropic_headers
from .transformation import AnthropicConfig
from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
@ -83,6 +83,7 @@ async def make_call(
timeout: Optional[Union[float, httpx.Timeout]],
json_mode: bool,
speed: Optional[str] = None,
tool_name_reverse_map: Optional[Dict[str, str]] = None,
) -> Tuple[Any, httpx.Headers]:
if client is None:
client = litellm.module_level_aclient
@ -117,6 +118,7 @@ async def make_call(
sync_stream=False,
json_mode=json_mode,
speed=speed,
tool_name_reverse_map=tool_name_reverse_map,
)
# LOGGING
@ -141,6 +143,7 @@ def make_sync_call(
timeout: Optional[Union[float, httpx.Timeout]],
json_mode: bool,
speed: Optional[str] = None,
tool_name_reverse_map: Optional[Dict[str, str]] = None,
) -> Tuple[Any, httpx.Headers]:
if client is None:
client = litellm.module_level_client # re-use a module level client
@ -183,6 +186,7 @@ def make_sync_call(
sync_stream=True,
json_mode=json_mode,
speed=speed,
tool_name_reverse_map=tool_name_reverse_map,
)
# LOGGING
@ -237,6 +241,11 @@ class AnthropicChatCompletion(BaseLLM):
timeout=timeout,
json_mode=json_mode,
speed=optional_params.get("speed") if optional_params else None,
tool_name_reverse_map=(
litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY)
if isinstance(litellm_params, dict)
else None
),
)
streamwrapper = CustomStreamWrapper(
completion_stream=completion_stream,
@ -462,6 +471,11 @@ class AnthropicChatCompletion(BaseLLM):
timeout=timeout,
json_mode=json_mode,
speed=optional_params.get("speed") if optional_params else None,
tool_name_reverse_map=(
litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY)
if isinstance(litellm_params, dict)
else None
),
)
return CustomStreamWrapper(
completion_stream=completion_stream,
@ -526,6 +540,7 @@ class ModelResponseIterator:
sync_stream: bool,
json_mode: Optional[bool] = False,
speed: Optional[str] = None,
tool_name_reverse_map: Optional[Dict[str, str]] = None,
):
self.streaming_response = streaming_response
self.response_iterator = self.streaming_response
@ -533,6 +548,13 @@ class ModelResponseIterator:
self.tool_index = -1
self.json_mode = json_mode
self.speed = speed
# rewritten-name -> caller's original. Built per-request from the
# forward map in AnthropicConfig._build_request_tool_name_maps; only
# contains entries we actually rewrote, so a tool legitimately named
# `foo_bar` is *not* reverse-mapped just because some other tool was
# rewritten to `foo_bar` in a different request. Empty/None is the
# common case (no '/' or other invalid chars in any tool name).
self.tool_name_reverse_map: Dict[str, str] = tool_name_reverse_map or {}
# Generate response ID once per stream to match OpenAI-compatible behavior
self.response_id = _generate_id()
@ -557,6 +579,10 @@ class ModelResponseIterator:
# Accumulate compaction blocks for multi-turn reconstruction
self.compaction_blocks: List[Dict[str, Any]] = []
# Accumulate streamed thinking text so final usage can split reasoning
# tokens from regular output tokens.
self.reasoning_content_chunks: List[str] = []
# Track server tool use inputs and results for code_interpreter_results
self._server_tool_inputs: Dict[str, Any] = {}
self.tool_results: List[Dict[str, Any]] = []
@ -587,9 +613,14 @@ class ModelResponseIterator:
return False
def _handle_usage(self, anthropic_usage_chunk: Union[dict, UsageDelta]) -> Usage:
reasoning_content = (
"".join(self.reasoning_content_chunks)
if self.reasoning_content_chunks
else None
)
return AnthropicConfig().calculate_usage(
usage_object=cast(dict, anthropic_usage_chunk),
reasoning_content=None,
reasoning_content=reasoning_content,
speed=self.speed,
)
@ -636,10 +667,13 @@ class ModelResponseIterator:
"thinking" in content_block["delta"]
or "signature" in content_block["delta"]
):
thinking_content = content_block["delta"].get("thinking")
if isinstance(thinking_content, str) and thinking_content:
self.reasoning_content_chunks.append(thinking_content)
thinking_blocks = [
ChatCompletionThinkingBlock(
type="thinking",
thinking=content_block["delta"].get("thinking") or "",
thinking=thinking_content or "",
signature=str(content_block["delta"].get("signature") or ""),
)
]
@ -792,6 +826,16 @@ class ModelResponseIterator:
or content_block_start["content_block"]["type"] == "server_tool_use"
):
self.tool_index += 1
# Reverse-map the (sanitized) tool name back to the
# caller's original. No-op when the map is empty.
_stream_tool_name = content_block_start["content_block"]["name"]
if (
self.tool_name_reverse_map
and _stream_tool_name in self.tool_name_reverse_map
):
_stream_tool_name = self.tool_name_reverse_map[
_stream_tool_name
]
# Use empty string for arguments in content_block_start - actual arguments
# come in subsequent content_block_delta chunks and get accumulated.
# Using str(input) here would prepend '{}' causing invalid JSON accumulation.
@ -799,7 +843,7 @@ class ModelResponseIterator:
id=content_block_start["content_block"]["id"],
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=content_block_start["content_block"]["name"],
name=_stream_tool_name,
arguments="",
),
index=self.tool_index,

View file

@ -105,6 +105,115 @@ else:
LoggingClass = Any
# Anthropic requires tool names to match ^[a-zA-Z0-9_-]{1,128}$. Any other
# character (commonly '/' or '.' from OpenAPI-derived MCP tools, e.g.
# "actions/download-job-logs-for-workflow-run") must be replaced before
# the request is sent.
#
# A naive "replace [^a-zA-Z0-9_-] with _" is unsafe because it's lossy:
# `foo/bar` and `foo_bar` both collapse to `foo_bar`. Two tools with the
# same sanitized name would either 400 at Anthropic (duplicate) or, worse,
# cause the response side to mis-translate `foo_bar` (a name the caller
# really did register) back to `foo/bar`.
#
# Instead we build a *per-request* forward map (original -> sanitized)
# whose codomain is unique within the request: when two originals collapse
# to the same candidate, or when a sanitized name collides with an already-
# valid name elsewhere in the request, we append numeric suffixes
# (`_2`, `_3`, ...) until the result is free.
#
# The reverse map (sanitized -> original) only contains entries where the
# original was actually rewritten. So a tool whose name is already valid
# round-trips identically and is *never* mistakenly re-mapped on the
# response side.
_ANTHROPIC_TOOL_NAME_INVALID_CHARS = re.compile(r"[^a-zA-Z0-9_-]")
_ANTHROPIC_TOOL_NAME_MAX_LEN = 128
# Single, internal-only key on ``litellm_params`` used to thread the per-
# request reverse map (sanitized -> original) from request build to response
# parsing. ``litellm_params`` is never serialized to a provider; ``optional_
# params`` IS (it becomes the JSON body via ``data = {**optional_params}``).
# Keep these two channels strictly separate -- never stash internal
# coordination state in ``optional_params``.
ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY = "_anthropic_tool_name_map"
def _basic_sanitize_anthropic_tool_name(name: str) -> str:
"""Lossy: replace [^a-zA-Z0-9_-] with '_' and truncate to 128.
Used as a candidate generator for the per-request forward map.
Callers should NOT use this directly for translation -- always go
through the forward map so collisions are resolved.
"""
if not isinstance(name, str) or not name:
return name
return _ANTHROPIC_TOOL_NAME_INVALID_CHARS.sub("_", name)[
:_ANTHROPIC_TOOL_NAME_MAX_LEN
]
def _build_anthropic_tool_name_maps(
original_names: List[str],
) -> Tuple[Dict[str, str], Dict[str, str]]:
"""Build (forward, reverse) tool-name maps for a single request.
forward[original] = sanitized -- only present when name was rewritten
reverse[sanitized] = original -- inverse of `forward`
Properties:
- All sanitized names satisfy ^[a-zA-Z0-9_-]{1,128}$.
- Sanitized names are unique within the request (no two originals
collide on the wire).
- A name that's already valid AND doesn't collide with another tool's
sanitized form passes through untouched and is absent from the maps.
That's the key correctness property: response-side translation only
runs on entries we actually rewrote, so a tool legitimately named
`foo_bar` is never incorrectly retyped to `foo/bar` just because
some *other* request had that pair.
- Order-dependent: when two originals would clash, the *second* one
seen gets the disambiguating suffix. Callers should preserve the
caller's tool order (we do).
"""
forward: Dict[str, str] = {}
used: set = set()
# First pass: reserve slots for names that are already valid so they
# always have priority regardless of input order.
for original in original_names:
if not isinstance(original, str) or not original:
continue
candidate = _basic_sanitize_anthropic_tool_name(original)
if candidate == original:
used.add(candidate)
# Second pass: sanitize/disambiguate names that need rewriting.
for original in original_names:
if not isinstance(original, str) or not original:
continue
candidate = _basic_sanitize_anthropic_tool_name(original)
if candidate == original:
continue
# Skip duplicates of the same original name. Without this guard the
# second pass would assign a fresh suffix and overwrite the forward
# map entry, causing every reference to map to the suffixed name and
# leaving the original sanitized slot orphaned in `used` with no
# reverse mapping.
if original in forward:
continue
# Disambiguate against names already chosen this request.
unique = candidate
n = 1
while unique in used:
n += 1
suffix = f"_{n}"
# Keep within the 128-char cap.
head = candidate[: _ANTHROPIC_TOOL_NAME_MAX_LEN - len(suffix)]
unique = f"{head}{suffix}"
forward[original] = unique
used.add(unique)
reverse = {v: k for k, v in forward.items()}
return forward, reverse
REASONING_EFFORT_TO_OUTPUT_CONFIG_EFFORT: Dict[str, str] = {
"low": "low",
"minimal": "low",
@ -486,7 +595,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
}
def _map_tool_choice(
self, tool_choice: Optional[str], parallel_tool_use: Optional[bool]
self,
tool_choice: Optional[str],
parallel_tool_use: Optional[bool],
) -> Optional[AnthropicMessagesToolChoice]:
_tool_choice: Optional[AnthropicMessagesToolChoice] = None
if tool_choice == "auto":
@ -527,7 +638,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
return _tool_choice
def _map_tool_helper( # noqa: PLR0915
self, tool: ChatCompletionToolParam
self,
tool: ChatCompletionToolParam,
) -> Tuple[Optional[AllAnthropicToolsValues], Optional[AnthropicMcpServerTool]]:
returned_tool: Optional[AllAnthropicToolsValues] = None
mcp_server: Optional[AnthropicMcpServerTool] = None
@ -783,7 +895,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
return initial_tool
def _map_tools(
self, tools: List
self,
tools: List,
) -> Tuple[List[AllAnthropicToolsValues], List[AnthropicMcpServerTool]]:
anthropic_tools = []
mcp_servers = []
@ -799,6 +912,174 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
mcp_servers.append(mcp_server_tool)
return anthropic_tools, mcp_servers
@staticmethod
def _rewrite_tool_names_in_messages(
messages: List[AllMessageValues],
name_forward_map: Dict[str, str],
) -> List[AllMessageValues]:
"""Return a copy of `messages` with tool_call/function_call names
rewritten using the per-request forward map.
Only mutates messages whose tool_call/function_call name is *in* the
forward map. Names absent from the map (already valid, no collision)
round-trip untouched. We only deep-copy the entries we actually
change to keep this O(turns-with-rewritten-tools), not O(history).
"""
if not name_forward_map:
return messages
new_messages: List[AllMessageValues] = []
for msg in messages:
if not isinstance(msg, dict):
new_messages.append(msg)
continue
tool_calls = msg.get("tool_calls")
function_call = msg.get("function_call")
if not tool_calls and not function_call:
new_messages.append(msg)
continue
new_msg = dict(msg)
if isinstance(tool_calls, list):
new_calls = []
for tc in tool_calls:
if not isinstance(tc, dict):
new_calls.append(tc)
continue
fn = tc.get("function")
fn_name = fn.get("name") if isinstance(fn, dict) else None
if (
isinstance(fn, dict)
and isinstance(fn_name, str)
and fn_name in name_forward_map
):
new_fn = dict(fn)
new_fn["name"] = name_forward_map[fn_name]
new_tc = dict(tc)
new_tc["function"] = new_fn
new_calls.append(new_tc)
else:
new_calls.append(tc)
new_msg["tool_calls"] = new_calls
fc_name = (
function_call.get("name") if isinstance(function_call, dict) else None
)
if (
isinstance(function_call, dict)
and isinstance(fc_name, str)
and fc_name in name_forward_map
):
new_fc = dict(function_call)
new_fc["name"] = name_forward_map[fc_name]
new_msg["function_call"] = new_fc
new_messages.append(cast(AllMessageValues, new_msg))
return new_messages
@staticmethod
def _build_request_tool_name_maps(
tools: List,
) -> Tuple[Dict[str, str], Dict[str, str]]:
"""Build the (forward, reverse) tool-name maps for an OpenAI tools list.
Operates on **OpenAI-format** tool dicts (pre-``_map_tools``). The
production sanitization path uses ``_sanitize_tool_names_in_request``
instead, which operates on **Anthropic-format** tools (post-
``_map_tools``, where ``type == "custom"``). This helper exists for
callers that need to compute the maps from the raw OpenAI shape --
e.g. test setup or future pre-mapping consumers.
See _build_anthropic_tool_name_maps for the collision rules. Pulls
the original name out of either ``{"function": {"name": ...}}``
(legacy OpenAI shape) or ``{"name": ...}`` (rare top-level shape).
"""
original_names: List[str] = []
for tool in tools or []:
if not isinstance(tool, dict):
continue
original = (
tool.get("function", {}).get("name")
if isinstance(tool.get("function"), dict)
else None
)
if original is None:
original = tool.get("name")
if isinstance(original, str) and original:
original_names.append(original)
return _build_anthropic_tool_name_maps(original_names)
@staticmethod
def _sanitize_tool_names_in_request(
optional_params: Dict[str, Any],
) -> Tuple[Dict[str, str], Dict[str, str]]:
"""Sanitize ``optional_params['tools']`` and ``optional_params['tool_choice']``
in place so every name matches Anthropic's ``^[a-zA-Z0-9_-]{1,128}$``.
Returns ``(forward, reverse)`` for use by message-history rewriting
and response translation. ``forward[original] = sanitized`` is only
populated for names that were actually rewritten -- i.e. either
contained an invalid character or collided with another tool's
sanitized form. Names already valid AND unique pass through and are
absent from both maps.
Only ``type == "custom"`` tools (the OpenAI function-tool shape) are
considered. Hosted tools (``web_search``, ``bash``, ``code_execution``,
``computer_*``, ``mcp``, ...) own reserved names defined by Anthropic
and must not be touched.
"""
tools = optional_params.get("tools")
if not isinstance(tools, list) or not tools:
return {}, {}
# 1. Collect originals from the Anthropic-shaped custom-tool entries.
# Order matters: the first occurrence wins the canonical slot;
# later collisions get numeric suffixes (see
# ``_build_anthropic_tool_name_maps``).
original_names: List[str] = []
for t in tools:
if not isinstance(t, dict):
continue
if t.get("type") != "custom":
continue
name = t.get("name")
if isinstance(name, str) and name:
original_names.append(name)
if not original_names:
return {}, {}
forward, reverse = _build_anthropic_tool_name_maps(original_names)
if not forward:
# Every name was already valid -- nothing to do.
return forward, reverse
# 2. Apply forward map. Build a new list with copy-on-change entries
# so a caller reusing the same tool list/dicts across requests
# doesn't see its inputs permanently rewritten (which would also
# drop the original key from `forward` on the next request).
new_tools: List[Any] = []
for t in tools:
if (
isinstance(t, dict)
and t.get("type") == "custom"
and isinstance(t.get("name"), str)
and t["name"] in forward
):
new_tools.append({**t, "name": forward[t["name"]]})
else:
new_tools.append(t)
optional_params["tools"] = new_tools
# 3. Same for ``tool_choice`` when it targets a named tool. Copy
# rather than mutate for the same reason as above.
tool_choice = optional_params.get("tool_choice")
if isinstance(tool_choice, dict) and tool_choice.get("type") == "tool":
tc_name = tool_choice.get("name")
if isinstance(tc_name, str) and tc_name in forward:
optional_params["tool_choice"] = {
**tool_choice,
"name": forward[tc_name],
}
return forward, reverse
def _detect_tool_search_tools(self, tools: Optional[List]) -> bool:
"""Check if tool search tools are present in the tools list."""
if not tools:
@ -1125,6 +1406,17 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
non_default_params=non_default_params
)
# NB: ``map_openai_params`` deliberately does NOT sanitize tool names
# here. Names are the *original* OpenAI names at this stage, and must
# remain so until ``transform_request`` -- which is the single
# chokepoint where Anthropic, Bedrock-Anthropic, and Vertex-Anthropic
# all pass through. Doing it there guarantees:
# 1. one source of truth for the per-request forward/reverse maps,
# 2. the maps land on ``litellm_params`` (internal), never on
# ``optional_params`` (which is serialized into the request body
# via ``data = {**optional_params}`` and would 400 with
# ``Extra inputs are not permitted``).
for param, value in non_default_params.items():
if param == "max_tokens":
optional_params["max_tokens"] = (
@ -1135,7 +1427,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
value if isinstance(value, int) else max(1, int(round(value)))
)
elif param == "tools":
# check if optional params already has tools
anthropic_tools, mcp_servers = self._map_tools(value)
optional_params = self._add_tools_to_optional_params(
optional_params=optional_params, tools=anthropic_tools
@ -1565,6 +1856,34 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
headers=headers, optional_params=optional_params
)
# === Tool-name sanitization (single chokepoint) ===
# Anthropic enforces ^[a-zA-Z0-9_-]{1,128}$ on every tool name. We
# sanitize *here* -- not in map_openai_params -- because:
#
# - This function is the single boundary shared by AnthropicConfig,
# AmazonAnthropicConfig (Bedrock invoke), VertexAIAnthropicConfig,
# and AzureAnthropicConfig (all call ``super().transform_request``
# or ``AnthropicConfig.transform_request(self, ...)``). Sanitizing
# once here covers every Anthropic-shaped request.
# - The forward/reverse maps are coordination state; they belong on
# ``litellm_params`` (internal-only), never on ``optional_params``
# (which becomes the JSON body via ``{**optional_params}``).
# - It keeps ``map_openai_params`` a pure param translator with no
# side-channel state.
#
# The reverse map only contains entries for names that were actually
# rewritten -- so a tool legitimately named ``foo_bar`` is never
# incorrectly retyped to ``foo/bar`` on the response side.
# See _build_anthropic_tool_name_maps for the collision-handling
# rules and rationale.
_name_forward_map, _name_reverse_map = self._sanitize_tool_names_in_request(
optional_params=optional_params,
)
if _name_forward_map:
messages = self._rewrite_tool_names_in_messages(messages, _name_forward_map)
if _name_reverse_map and isinstance(litellm_params, dict):
litellm_params[ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY] = _name_reverse_map
# Separate system prompt from rest of message
anthropic_system_message_list = self.translate_system_message(messages=messages)
# Handling anthropic API Prompt Caching
@ -1837,8 +2156,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
speed: Optional[str] = None,
) -> Usage:
# NOTE: Sometimes the usage object has None set explicitly for token counts, meaning .get() & key access returns None, and we need to account for this
prompt_tokens = usage_object.get("input_tokens", 0) or 0
completion_tokens = usage_object.get("output_tokens", 0) or 0
raw_prompt_tokens = usage_object.get("input_tokens", 0) or 0
prompt_tokens: int = (
int(raw_prompt_tokens) if isinstance(raw_prompt_tokens, (int, float)) else 0
)
raw_completion_tokens = usage_object.get("output_tokens", 0) or 0
completion_tokens: int = (
int(raw_completion_tokens)
if isinstance(raw_completion_tokens, (int, float))
else 0
)
_usage = usage_object
cache_creation_input_tokens: int = 0
cache_read_input_tokens: int = 0
@ -1907,11 +2234,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
text_tokens=raw_input_tokens,
)
# Always populate completion_token_details, not just when there's reasoning_content
reasoning_tokens = (
estimated_reasoning_tokens = (
token_counter(text=reasoning_content, count_response_tokens=True)
if reasoning_content
else 0
)
reasoning_tokens = min(estimated_reasoning_tokens, completion_tokens)
completion_token_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else 0,
text_tokens=(
@ -2041,6 +2369,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
json_mode: Optional[bool] = None,
prefix_prompt: Optional[str] = None,
speed: Optional[str] = None,
tool_name_reverse_map: Optional[Dict[str, str]] = None,
):
_hidden_params: Dict = {}
_hidden_params["additional_headers"] = process_anthropic_headers(
@ -2065,6 +2394,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
compaction_blocks,
) = self.extract_response_content(completion_response=completion_response)
# Reverse-map rewritten tool names back to caller's originals so a
# downstream OpenAI-style dispatcher can match on the registered name.
# See _build_anthropic_tool_name_maps for why this is keyed on the
# per-request reverse map (so a tool legitimately named `foo_bar` is
# never incorrectly retyped to `foo/bar`). No-op when the map is
# empty (the common case).
if tool_name_reverse_map and tool_calls:
for tc in tool_calls:
fn = tc.get("function") if isinstance(tc, dict) else None
if fn is None:
continue
_name = fn.get("name")
if isinstance(_name, str) and _name in tool_name_reverse_map:
fn["name"] = tool_name_reverse_map[_name]
if (
prefix_prompt is not None
and not text_content.startswith(prefix_prompt)
@ -2191,6 +2535,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
prefix_prompt = self.get_prefix_prompt(messages=messages)
speed = optional_params.get("speed")
tool_name_reverse_map: Optional[Dict[str, str]] = None
if isinstance(litellm_params, dict):
_candidate = litellm_params.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY)
if isinstance(_candidate, dict):
tool_name_reverse_map = _candidate
model_response = self.transform_parsed_response(
completion_response=completion_response,
@ -2199,6 +2548,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
json_mode=json_mode,
prefix_prompt=prefix_prompt,
speed=speed,
tool_name_reverse_map=tool_name_reverse_map,
)
return model_response

View file

@ -89,9 +89,10 @@ class AzureOpenAIRealtime(AzureChatCompletion):
if api_base is None:
raise ValueError("api_base is required for Azure OpenAI calls")
if api_version is None and (
backend_uses_beta_protocol = (
realtime_protocol is None or realtime_protocol.upper() not in ("GA", "V1")
):
)
if api_version is None and backend_uses_beta_protocol:
raise ValueError("api_version is required for Azure OpenAI calls")
url = self._construct_url(
@ -114,6 +115,7 @@ class AzureOpenAIRealtime(AzureChatCompletion):
logging_obj,
user_api_key_dict=user_api_key_dict,
request_data={"litellm_metadata": litellm_metadata or {}},
backend_uses_beta_protocol=backend_uses_beta_protocol,
)
await realtime_streaming.bidirectional_forward()

View file

@ -7,6 +7,8 @@ from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
Callable,
ClassVar,
Dict,
List,
Literal,
@ -63,8 +65,16 @@ class AwsAuthError(Exception):
class BaseAWSLLM:
# Process-wide IAM credential cache (shared across instances — Bedrock passthrough is per-request).
# Storage is in-process memory only: default ``DualCache()`` has no Redis backend unless attached
# elsewhere. Entry TTL: static access-key + secret + region use ``_get_default_ttl_for_boto3_credentials``
# (~59 minutes); ambient env (``_auth_with_env_vars`` returns ``ttl=None``) uses ``InMemoryCache``'s
# ``default_ttl`` (600 seconds / 10 minutes). AssumeRole, web identity, profiles, and explicit
# session-token tuples are not cached — see ``get_credentials`` and ``_get_or_set_cached_credentials``.
_shared_iam_cache: ClassVar[DualCache] = DualCache()
def __init__(self) -> None:
self.iam_cache = DualCache()
self.iam_cache = BaseAWSLLM._shared_iam_cache
super().__init__()
self.aws_authentication_params = [
"aws_access_key_id",
@ -103,6 +113,79 @@ class BaseAWSLLM:
credential_str = json.dumps(credential_args, sort_keys=True)
return hashlib.sha256(credential_str.encode()).hexdigest()
def _get_or_set_cached_credentials(
self,
credential_args: Dict[str, Optional[str]],
credential_fetcher: Callable[[], Tuple[Any, Optional[int]]],
) -> Any:
"""
Read-through IAM cache on the process-wide ``DualCache``.
Only the in-memory layer is used by default (no Redis on ``_shared_iam_cache`` unless
configured globally). TTL on write: static access-key fetches pass
``_get_default_ttl_for_boto3_credentials()`` (~59 minutes); ambient env passes ``ttl=None``,
which ``InMemoryCache.set_cache`` resolves to ``default_ttl`` (600 seconds / 10 minutes by
default).
Used only for static access-key credentials and ambient credentials from
``_auth_with_env_vars`` (including when skipping AssumeRole because the runtime identity
already matches ``aws_role_name``).
AssumeRole, web identity exchange, profiles, and explicit session-token tuples are not
cached here shared ``Credentials`` / refresh state must not span logical sessions.
"""
cache_key = self.get_cache_key(credential_args)
_cached = self.iam_cache.get_cache(cache_key)
if _cached:
return _cached
credentials, ttl = credential_fetcher()
self.iam_cache.set_cache(cache_key, credentials, ttl=ttl)
return credentials
@staticmethod
def _is_auth_with_web_identity_token(
aws_web_identity_token: Optional[str],
aws_role_name: Optional[str],
aws_session_name: Optional[str],
) -> bool:
return (
aws_web_identity_token is not None
and aws_role_name is not None
and aws_session_name is not None
)
@staticmethod
def _is_auth_with_aws_role(aws_role_name: Optional[str]) -> bool:
return aws_role_name is not None
@staticmethod
def _is_auth_with_aws_profile(aws_profile_name: Optional[str]) -> bool:
return aws_profile_name is not None
@staticmethod
def _is_auth_with_aws_session_token_tuple(
aws_access_key_id: Optional[str],
aws_secret_access_key: Optional[str],
aws_session_token: Optional[str],
) -> bool:
return (
aws_access_key_id is not None
and aws_secret_access_key is not None
and aws_session_token is not None
)
@staticmethod
def _is_auth_with_access_key_and_secret_key(
aws_access_key_id: Optional[str],
aws_secret_access_key: Optional[str],
aws_region_name: Optional[str],
) -> bool:
return (
aws_access_key_id is not None
and aws_secret_access_key is not None
and aws_region_name is not None
)
@tracer.wrap()
def get_credentials(
self,
@ -184,95 +267,97 @@ class BaseAWSLLM:
aws_external_id,
)
# create cache key for non-expiring auth flows
args = {
k: v
for k, v in locals().items()
if k.startswith("aws_") or k == "ssl_verify"
}
cache_key = self.get_cache_key(args)
_cached_credentials = self.iam_cache.get_cache(cache_key)
if _cached_credentials:
return _cached_credentials
#########################################################
# Handle diff boto3 auth flows
# for each helper
# Return:
# Credentials - boto3.Credentials
# cache ttl - Optional[int]. If None, the credentials are not cached. Some auth flows have no expiry time.
#
# iam_cache: static keys and ambient env only (including skip-AssumeRole path).
# Do not cache AssumeRole / web identity / profile / explicit session-token paths here.
#########################################################
if (
aws_web_identity_token is not None
and aws_role_name is not None
and aws_session_name is not None
if self._is_auth_with_web_identity_token(
aws_web_identity_token,
aws_role_name,
aws_session_name,
):
credentials, _cache_ttl = self._auth_with_web_identity_token(
aws_web_identity_token=aws_web_identity_token,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
aws_web_identity_token=cast(str, aws_web_identity_token),
aws_role_name=cast(str, aws_role_name),
aws_session_name=cast(str, aws_session_name),
aws_region_name=aws_region_name,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
elif aws_role_name is not None:
# Check if we're already running as the target role and can skip assumption
# This handles IRSA (EKS), ECS task roles, and EC2 instance profiles
if self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify):
return credentials
elif self._is_auth_with_aws_role(aws_role_name):
# Same role (IRSA/ECS/EC2): ambient creds via _get_or_set_cached_credentials like the
# default env branch; never pre-read cache (must run _is_already_running_as_role first).
if self._is_already_running_as_role(
cast(str, aws_role_name), ssl_verify=ssl_verify
):
verbose_logger.debug(
"Already running as target role %s, using ambient credentials",
aws_role_name,
)
credentials, _cache_ttl = self._auth_with_env_vars()
else:
verbose_logger.debug(
"Using role assumption: calling _auth_with_aws_role"
return self._get_or_set_cached_credentials(
args, self._auth_with_env_vars
)
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = (
f"litellm-session-{int(datetime.now().timestamp())}"
)
credentials, _cache_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
aws_region_name=aws_region_name,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
ssl_verify=ssl_verify,
)
elif aws_profile_name is not None: ### CHECK SESSION ###
credentials, _cache_ttl = self._auth_with_aws_profile(aws_profile_name)
elif (
aws_access_key_id is not None
and aws_secret_access_key is not None
and aws_session_token is not None
):
credentials, _cache_ttl = self._auth_with_aws_session_token(
verbose_logger.debug("Using role assumption: calling _auth_with_aws_role")
# If aws_session_name is not provided, generate a default one
if aws_session_name is None:
aws_session_name = f"litellm-session-{int(datetime.now().timestamp())}"
credentials, _assume_ttl = self._auth_with_aws_role(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
)
elif (
aws_access_key_id is not None
and aws_secret_access_key is not None
and aws_region_name is not None
):
credentials, _cache_ttl = self._auth_with_access_key_and_secret_key(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_role_name=cast(str, aws_role_name),
aws_session_name=aws_session_name,
aws_region_name=aws_region_name,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
ssl_verify=ssl_verify,
)
return credentials
elif self._is_auth_with_aws_profile(aws_profile_name):
credentials, _cache_ttl = self._auth_with_aws_profile(
cast(str, aws_profile_name)
)
return credentials
elif self._is_auth_with_aws_session_token_tuple(
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
):
credentials, _cache_ttl = self._auth_with_aws_session_token(
aws_access_key_id=cast(str, aws_access_key_id),
aws_secret_access_key=cast(str, aws_secret_access_key),
aws_session_token=cast(str, aws_session_token),
)
return credentials
elif self._is_auth_with_access_key_and_secret_key(
aws_access_key_id,
aws_secret_access_key,
aws_region_name,
):
return self._get_or_set_cached_credentials(
args,
lambda: self._auth_with_access_key_and_secret_key(
aws_access_key_id=cast(str, aws_access_key_id),
aws_secret_access_key=cast(str, aws_secret_access_key),
aws_region_name=cast(str, aws_region_name),
),
)
else:
credentials, _cache_ttl = self._auth_with_env_vars()
self.iam_cache.set_cache(cache_key, credentials, ttl=_cache_ttl)
return credentials
return self._get_or_set_cached_credentials(args, self._auth_with_env_vars)
def _get_aws_region_from_model_arn(self, model: Optional[str]) -> Optional[str]:
try:

View file

@ -299,29 +299,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM):
)
def _get_response_stream_shape(self):
"""Get the response stream shape for parsing, reusing existing logic."""
try:
# Try to reuse the cached shape from the existing decoder
from litellm.llms.bedrock.chat.invoke_handler import (
get_response_stream_shape,
)
from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE
return get_response_stream_shape()
except ImportError:
# Fallback: create our own shape
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model(
"bedrock-runtime", "service-2"
)
bedrock_service_model = ServiceModel(bedrock_service_dict)
return bedrock_service_model.shape_for("ResponseStream")
except Exception as e:
verbose_logger.warning(f"Could not load response stream shape: {e}")
return None
return BEDROCK_RESPONSE_STREAM_SHAPE
def _extract_response_content(self, events: InvokeAgentEventList) -> str:
"""Extract the final response content from parsed events."""

View file

@ -67,9 +67,13 @@ from litellm.types.utils import (
from litellm.utils import CustomStreamWrapper, get_secret
from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockError, ModelResponseIterator, get_bedrock_tool_name
from ..common_utils import (
BEDROCK_RESPONSE_STREAM_SHAPE,
BedrockError,
ModelResponseIterator,
get_bedrock_tool_name,
)
_response_stream_shape_cache = None
bedrock_tool_name_mappings: InMemoryCache = InMemoryCache(
max_size_in_memory=50, default_ttl=600
)
@ -1391,20 +1395,6 @@ class BedrockLLM(BaseAWSLLM):
return None
def get_response_stream_shape():
global _response_stream_shape_cache
if _response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model("bedrock-runtime", "service-2")
bedrock_service_model = ServiceModel(bedrock_service_dict)
_response_stream_shape_cache = bedrock_service_model.shape_for("ResponseStream")
return _response_stream_shape_cache
class AWSEventStreamDecoder:
def __init__(self, model: str, json_mode: Optional[bool] = False) -> None:
from botocore.parsers import EventStreamJSONParser
@ -1838,8 +1828,18 @@ class AWSEventStreamDecoder:
yield self._chunk_parser(chunk_data=_data)
def _parse_message_from_event(self, event) -> Optional[str]:
if BEDROCK_RESPONSE_STREAM_SHAPE is None:
raise BedrockError(
status_code=500,
message=(
"Bedrock event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
parsed_response = self.parser.parse(
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:
decoded_body = response_dict["body"].decode()

View file

@ -14,6 +14,7 @@ if TYPE_CHECKING:
import httpx
import litellm
from litellm import verbose_logger
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
@ -917,38 +918,57 @@ def get_bedrock_chat_config(model: str):
return litellm.AmazonInvokeConfig()
def _load_bedrock_response_stream_shape():
"""
Load the ResponseStream shape from botocore's bundled bedrock-runtime schema.
Called once at module import time; the result is stored in
``BEDROCK_RESPONSE_STREAM_SHAPE`` and reused for the process lifetime.
Returns ``None`` if botocore is unavailable or the service model cannot be
loaded, so the module still imports cleanly.
"""
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
service_dict = loader.load_service_model("bedrock-runtime", "service-2")
return ServiceModel(service_dict).shape_for("ResponseStream")
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-load bedrock-runtime response stream shape "
"— Bedrock event-stream decoding will be unavailable. Error: %s",
e,
)
return None
# Eagerly resolved once per process — avoids per-instance or per-request disk I/O.
BEDROCK_RESPONSE_STREAM_SHAPE = _load_bedrock_response_stream_shape()
class BedrockEventStreamDecoderBase:
"""
Base class for event stream decoding for Bedrock
"""
_response_stream_shape_cache = None
def __init__(self):
from botocore.parsers import EventStreamJSONParser
self.parser = EventStreamJSONParser()
def get_response_stream_shape(self):
if self._response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
bedrock_service_dict = loader.load_service_model(
"bedrock-runtime", "service-2"
)
bedrock_service_model = ServiceModel(bedrock_service_dict)
self._response_stream_shape_cache = bedrock_service_model.shape_for(
"ResponseStream"
)
return self._response_stream_shape_cache
def _parse_message_from_event(self, event) -> Optional[str]:
if BEDROCK_RESPONSE_STREAM_SHAPE is None:
raise BedrockError(
status_code=500,
message=(
"Bedrock event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(
response_dict, self.get_response_stream_shape()
response_dict, BEDROCK_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:

View file

@ -10,14 +10,12 @@ from ..base_aws_llm import BaseAWSLLM
from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo
if TYPE_CHECKING:
from httpx import URL
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.utils import CostResponseTypes
if TYPE_CHECKING:
from httpx import URL
class BedrockPassthroughConfig(
BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig
):

View file

@ -1,4 +1,5 @@
import asyncio
import concurrent.futures
import inspect
import os
import socket
@ -133,6 +134,11 @@ _DEFAULT_TIMEOUT = httpx.Timeout(
timeout=COMPLETION_HTTP_FALLBACK_SECONDS,
connect=HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS,
)
_STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS = 5.0
_STREAMING_ERROR_BODY_READ_EXECUTOR = concurrent.futures.ThreadPoolExecutor(
max_workers=50,
thread_name_prefix="litellm-streaming-error-body-read",
)
def _prepare_request_data_and_content(
@ -386,17 +392,30 @@ def _safe_get_response_text(response: httpx.Response) -> str:
return ""
async def _safe_aread_response(response: httpx.Response) -> bytes:
async def _safe_aread_response(
response: httpx.Response, timeout: Optional[float] = None
) -> bytes:
"""Safely read async response body, falling back to empty bytes on errors."""
try:
if timeout is not None:
return await asyncio.wait_for(response.aread(), timeout=timeout)
return await response.aread()
except Exception:
return b""
def _safe_read_response(response: httpx.Response) -> bytes:
def _safe_read_response(
response: httpx.Response, timeout: Optional[float] = None
) -> bytes:
"""Safely read sync response body, falling back to empty bytes on errors."""
try:
if timeout is not None:
future = _STREAMING_ERROR_BODY_READ_EXECUTOR.submit(response.read)
try:
return future.result(timeout=timeout)
except Exception:
response.close()
return b""
return response.read()
except Exception:
return b""
@ -405,8 +424,19 @@ def _safe_read_response(response: httpx.Response) -> bytes:
def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None:
"""Raise a MaskedHTTPStatusError for sync HTTP handlers."""
if stream:
_body = mask_sensitive_info(_safe_read_response(e.response))
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
try:
_body = mask_sensitive_info(
_safe_read_response(
e.response,
timeout=_STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS,
)
)
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
finally:
try:
e.response.close()
except Exception:
pass
_text = mask_sensitive_info(_safe_get_response_text(e.response))
raise MaskedHTTPStatusError(e, message=_text, text=_text) from None
@ -414,8 +444,19 @@ def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None:
async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> None:
"""Raise a MaskedHTTPStatusError for async HTTP handlers."""
if stream:
_body = mask_sensitive_info(await _safe_aread_response(e.response))
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
try:
_body = mask_sensitive_info(
await _safe_aread_response(
e.response,
timeout=_STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS,
)
)
raise MaskedHTTPStatusError(e, message=_body, text=_body) from None
finally:
try:
await e.response.aclose()
except Exception:
pass
_text = mask_sensitive_info(_safe_get_response_text(e.response))
raise MaskedHTTPStatusError(e, message=_text, text=_text) from None

View file

@ -5255,7 +5255,6 @@ class BaseLLMHTTPHandler:
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json",
"OpenAI-Beta": "realtime=v1",
}
if extra_headers:

View file

@ -2,7 +2,18 @@
Translate from OpenAI's `/v1/chat/completions` to VLLM's `/v1/chat/completions`
"""
from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, overload
from typing import (
Any,
Coroutine,
Dict,
List,
Literal,
Optional,
Tuple,
Union,
cast,
overload,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
_get_image_mime_type_from_url,
@ -21,6 +32,61 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig
class HostedVLLMChatConfig(OpenAIGPTConfig):
def _convert_custom_tools_to_function_tools(
self, tools: List[Dict[str, Any]]
) -> List[Dict[str, Any]]:
"""
vLLM chat completions currently accepts only OpenAI function tools.
Convert custom tools into function tools so request validation does not fail.
"""
converted_tools: List[Dict[str, Any]] = []
for idx, tool in enumerate(tools):
if not isinstance(tool, dict):
converted_tools.append(tool)
continue
if tool.get("type") != "custom":
converted_tools.append(tool)
continue
custom_tool = tool.get("custom", {})
if not isinstance(custom_tool, dict):
custom_tool = {}
tool_name = (
custom_tool.get("name") or tool.get("name") or f"custom_tool_{idx}"
)
tool_description = custom_tool.get("description") or tool.get("description")
tool_parameters = custom_tool.get("input_schema") or tool.get(
"input_schema"
)
if not isinstance(tool_parameters, dict):
tool_parameters = {
"type": "object",
"properties": {
"input": {
"type": "string",
"description": "Raw tool input payload.",
}
},
"required": ["input"],
}
function_tool: Dict[str, Any] = {
"type": "function",
"function": {
"name": str(tool_name),
"parameters": tool_parameters,
},
}
if isinstance(tool_description, str):
function_tool["function"]["description"] = tool_description
converted_tools.append(function_tool)
return converted_tools
def get_supported_openai_params(self, model: str) -> List[str]:
params = super().get_supported_openai_params(model)
params.extend(["reasoning_effort", "thinking"])
@ -39,6 +105,8 @@ class HostedVLLMChatConfig(OpenAIGPTConfig):
_tools = _remove_additional_properties(_tools)
# remove 'strict' from tools
_tools = _remove_strict_from_schema(_tools)
if isinstance(_tools, list):
_tools = self._convert_custom_tools_to_function_tools(_tools)
if _tools is not None:
non_default_params["tools"] = _tools

View file

View file

@ -0,0 +1,232 @@
"""
Audio resampling utilities for the NVIDIA Riva STT provider.
We intentionally avoid a hard dependency on ``ffmpeg`` so this works in
slim Python environments. Format coverage:
- ``soundfile`` handles wav / flac / ogg out of the box (libsndfile).
- ``audioread`` is tried for everything ``soundfile`` cannot decode (mp3,
m4a, mp4, webm, ...). This is a soft optional dependency.
If neither library can decode the input we raise a clear error instructing
the caller to convert the audio upstream.
"""
import io
import os
import tempfile
from dataclasses import dataclass
from typing import Any, Tuple, cast
from litellm.llms.nvidia_riva.audio_transcription.transformation import (
RIVA_TARGET_NUM_CHANNELS,
RIVA_TARGET_SAMPLE_RATE_HZ,
)
from litellm.llms.nvidia_riva.common_utils import NvidiaRivaException
# Keep this as Any: the module intentionally avoids importing numpy at module
# import time (optional dependency), and project-wide mypy config evaluates this
# file in contexts where conditional type aliases can degrade to "FloatArray?".
FloatArray = Any
_INSTALL_HINT = (
"Install Riva STT extras to enable automatic audio resampling: "
"`pip install 'litellm[stt-nvidia-riva]'`"
)
@dataclass
class ResampledAudio:
pcm_bytes: bytes
duration_seconds: float
sample_rate_hz: int
num_channels: int
def resample_to_riva_pcm(file_bytes: bytes) -> ResampledAudio:
"""
Decode ``file_bytes`` and produce 16 kHz mono LINEAR_PCM (int16 little
endian) suitable for streaming to Riva, plus the audio duration in
seconds (used for cost calculation when Riva does not return usage).
"""
try:
import numpy as np # type: ignore
except ImportError as e:
raise NvidiaRivaException(
status_code=500,
message=f"numpy is required for Riva audio resampling. {_INSTALL_HINT}",
) from e
samples_float, source_rate = _decode_to_float32(file_bytes)
# Downmix to mono by averaging channels.
if samples_float.ndim == 2 and samples_float.shape[1] > 1:
samples_float = samples_float.mean(axis=1)
elif samples_float.ndim == 2:
samples_float = samples_float[:, 0]
samples_float = np.asarray(samples_float, dtype=np.float32).ravel()
if source_rate != RIVA_TARGET_SAMPLE_RATE_HZ:
samples_float = _resample(
samples_float, source_rate, RIVA_TARGET_SAMPLE_RATE_HZ
)
# Clip + convert float [-1, 1] to int16 little-endian PCM.
np.clip(samples_float, -1.0, 1.0, out=samples_float)
pcm_int16 = (samples_float * 32767.0).astype("<i2")
pcm_bytes = pcm_int16.tobytes()
duration_seconds = float(pcm_int16.size) / float(RIVA_TARGET_SAMPLE_RATE_HZ)
return ResampledAudio(
pcm_bytes=pcm_bytes,
duration_seconds=duration_seconds,
sample_rate_hz=RIVA_TARGET_SAMPLE_RATE_HZ,
num_channels=RIVA_TARGET_NUM_CHANNELS,
)
def _decode_to_float32(file_bytes: bytes) -> Tuple["FloatArray", int]:
"""
Decode arbitrary audio bytes into a float32 array shaped either
``(n_samples,)`` (mono) or ``(n_samples, n_channels)`` plus the source
sample rate.
Tries ``soundfile`` first (wav/flac/ogg), then falls back to
``audioread`` for compressed formats. Raises a clear error if neither
works.
"""
import numpy as np # type: ignore
sf_error: Exception | None = None
try:
import soundfile as sf # type: ignore
with io.BytesIO(file_bytes) as buf:
data, source_rate = sf.read(buf, dtype="float32", always_2d=False)
return cast("FloatArray", data), int(source_rate)
except ImportError as e:
sf_error = e
except Exception as e:
# soundfile raises RuntimeError / LibsndfileError for formats it
# cannot decode (mp3 on older libsndfile, m4a, webm, ...).
sf_error = e
try:
import audioread # type: ignore
except ImportError as e:
raise NvidiaRivaException(
status_code=400,
message=(
"Could not decode audio for Riva STT. Install audio extras "
f"(`pip install 'litellm[stt-nvidia-riva]'`) or convert your "
f"audio to wav/flac/ogg before calling the API. "
f"Underlying error: {sf_error}"
),
) from e
# audioread backends (FFmpeg subprocess, GStreamer, Core Audio) require a
# filesystem path, so spill the bytes to a temp file. mkstemp is portable
# to Windows where re-opening a NamedTemporaryFile is not allowed.
fd, tmp_path = tempfile.mkstemp(suffix=".audio")
try:
with os.fdopen(fd, "wb") as tmp_file:
tmp_file.write(file_bytes)
try:
with audioread.audio_open(tmp_path) as src:
source_rate = int(src.samplerate)
channels = int(src.channels)
chunks = []
for buf in src:
chunks.append(np.frombuffer(buf, dtype=np.int16))
if not chunks:
raise NvidiaRivaException(
status_code=400,
message="Audio decode produced no samples.",
)
interleaved = np.concatenate(chunks).astype(np.float32) / 32768.0
if channels > 1:
interleaved = interleaved.reshape(-1, channels)
return cast("FloatArray", interleaved), source_rate
except NvidiaRivaException:
raise
except Exception as e:
raise NvidiaRivaException(
status_code=400,
message=(
"Could not decode audio for Riva STT. Convert your audio to "
f"wav/flac/ogg before calling the API. Underlying error: {e}"
),
) from e
finally:
try:
os.unlink(tmp_path)
except OSError:
pass
def _resample(
samples: "FloatArray", source_rate: int, target_rate: int
) -> "FloatArray":
"""
Resample mono float32 ``samples`` from ``source_rate`` to ``target_rate``.
Prefers high-quality polyphase resampling when ``soxr`` or ``scipy`` is
available (anti-aliased, important for downsampling 44.1/48 kHz -> 16 kHz
where naive interpolation folds high frequencies back into the speech
band). Falls back to linear interpolation if neither is installed
acceptable for speech-only mono input but lossy for wideband content.
"""
import numpy as np # type: ignore
if source_rate == target_rate or samples.size == 0:
return samples
try:
import soxr # type: ignore
return cast(
"FloatArray",
np.asarray(
soxr.resample(samples, source_rate, target_rate), dtype=np.float32
),
)
except ImportError:
pass
try:
from math import gcd
from scipy.signal import resample_poly # type: ignore
g = gcd(int(source_rate), int(target_rate))
up = int(target_rate) // g
down = int(source_rate) // g
return cast(
"FloatArray", np.asarray(resample_poly(samples, up, down), dtype=np.float32)
)
except ImportError:
pass
return _linear_resample(samples, source_rate, target_rate)
def _linear_resample(
samples: "FloatArray", source_rate: int, target_rate: int
) -> "FloatArray":
"""Linear-interpolation fallback. See :func:`_resample` for caveats."""
import numpy as np # type: ignore
duration = samples.size / float(source_rate)
target_length = int(round(duration * target_rate))
if target_length <= 1:
return samples.astype(np.float32)
src_indices = np.linspace(0, samples.size - 1, num=target_length, dtype=np.float64)
left = np.floor(src_indices).astype(np.int64)
right = np.minimum(left + 1, samples.size - 1)
frac = (src_indices - left).astype(np.float32)
return ((1.0 - frac) * samples[left] + frac * samples[right]).astype(np.float32)

View file

@ -0,0 +1,444 @@
"""
NVIDIA Riva STT handler.
This module bridges litellm's transcription dispatch to NVIDIA Riva's gRPC
streaming ASR API. We do *not* go through ``base_llm_http_handler`` because
Riva is gRPC-only: HTTP-shaped abstractions (``httpx.Response``,
``api_base/v1/...`` URLs, multipart bodies) do not apply.
The handler is intentionally a thin orchestration layer:
1. Resample the inbound audio to 16 kHz mono LINEAR_PCM (Riva's required
wire format).
2. Build ``RecognitionConfig`` / ``StreamingRecognitionConfig`` protobufs
from the structured dict produced by
:class:`NvidiaRivaAudioTranscriptionConfig`.
3. Construct ``riva.client.Auth`` honoring NVCF (function-id metadata + TLS)
vs self-hosted (any host:port, optional TLS) modes.
4. Stream the audio through Riva's ``streaming_response_generator`` and
aggregate ``is_final`` results into a single transcript.
5. Return a normalized ``TranscriptionResponse`` with ``duration`` exposed
on ``_hidden_params`` so cost calculation works.
``riva-client`` is imported lazily so ``litellm`` core remains usable
without the optional STT extras installed.
"""
import asyncio
import inspect
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from litellm.litellm_core_utils.audio_utils.utils import (
get_audio_file_name,
process_audio_file,
)
from litellm.llms.nvidia_riva.audio_transcription.audio_utils import (
resample_to_riva_pcm,
)
from litellm.llms.nvidia_riva.audio_transcription.transformation import (
NvidiaRivaAudioTranscriptionConfig,
RIVA_TARGET_NUM_CHANNELS,
RIVA_TARGET_SAMPLE_RATE_HZ,
)
from litellm.llms.nvidia_riva.common_utils import (
NvidiaRivaException,
grpc_error_to_litellm_exception,
)
from litellm.types.utils import FileTypes, TranscriptionResponse
from litellm.utils import convert_to_model_response_object
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
# Stream audio to Riva in ~50 ms slices (1600 samples at 16 kHz). Matches
# NVIDIA's recommended chunk size for streaming ASR — small enough for
# responsive endpointing, large enough to keep per-RPC overhead low.
_DEFAULT_CHUNK_SAMPLES = 1600
_DEFAULT_CHUNK_BYTES = _DEFAULT_CHUNK_SAMPLES * 2 # int16 = 2 bytes/sample
_RIVA_INSTALL_HINT = (
"NVIDIA Riva client is not installed. "
"Install with `pip install 'litellm[stt-nvidia-riva]'`."
)
class NvidiaRivaAudioTranscription:
"""Sync + async entry point for Riva ASR."""
def audio_transcriptions(
self,
model: str,
audio_file: FileTypes,
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
logging_obj: "LiteLLMLoggingObj",
api_key: Optional[str],
api_base: Optional[str],
atranscription: bool = False,
provider_config: Optional[NvidiaRivaAudioTranscriptionConfig] = None,
):
if provider_config is None:
provider_config = NvidiaRivaAudioTranscriptionConfig()
if atranscription:
return self.async_audio_transcriptions(
model=model,
audio_file=audio_file,
optional_params=optional_params,
litellm_params=litellm_params,
model_response=model_response,
timeout=timeout,
logging_obj=logging_obj,
api_key=api_key,
api_base=api_base,
provider_config=provider_config,
)
return self._run_sync(
model=model,
audio_file=audio_file,
optional_params=optional_params,
litellm_params=litellm_params,
model_response=model_response,
timeout=timeout,
logging_obj=logging_obj,
api_key=api_key,
api_base=api_base,
provider_config=provider_config,
atranscription=atranscription,
)
async def async_audio_transcriptions(
self,
model: str,
audio_file: FileTypes,
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
logging_obj: "LiteLLMLoggingObj",
api_key: Optional[str],
api_base: Optional[str],
provider_config: Optional[NvidiaRivaAudioTranscriptionConfig] = None,
) -> TranscriptionResponse:
# ``riva-client`` exposes a sync streaming generator, so we offload
# the blocking call to a worker thread to keep the event loop free.
return await asyncio.to_thread(
self._run_sync,
model=model,
audio_file=audio_file,
optional_params=optional_params,
litellm_params=litellm_params,
model_response=model_response,
timeout=timeout,
logging_obj=logging_obj,
api_key=api_key,
api_base=api_base,
provider_config=provider_config or NvidiaRivaAudioTranscriptionConfig(),
atranscription=True,
)
def _run_sync(
self,
model: str,
audio_file: FileTypes,
optional_params: dict,
litellm_params: dict,
model_response: TranscriptionResponse,
timeout: float,
logging_obj: "LiteLLMLoggingObj",
api_key: Optional[str],
api_base: Optional[str],
provider_config: NvidiaRivaAudioTranscriptionConfig,
atranscription: bool = False,
) -> TranscriptionResponse:
if not api_base:
raise NvidiaRivaException(
status_code=400,
message=(
"NVIDIA Riva requires `api_base` (host:port for the gRPC "
"endpoint, e.g. `grpc.nvcf.nvidia.com:443` or "
"`localhost:50051`). Set it in litellm_params or via "
"NVIDIA_RIVA_API_BASE."
),
)
processed = process_audio_file(audio_file)
resampled = resample_to_riva_pcm(processed.file_content)
request_payload = provider_config.transform_audio_transcription_request(
model=model,
audio_file=audio_file,
optional_params=optional_params,
litellm_params={
**litellm_params,
"api_base": api_base,
"api_key": api_key,
},
).data
if not isinstance(request_payload, dict):
raise NvidiaRivaException(
status_code=500,
message="NvidiaRivaAudioTranscriptionConfig produced an unexpected request payload type.",
)
recognition_config_dict: Dict[str, Any] = request_payload["recognition_config"]
# The wire format is fixed by our resampler; override anything stale
# the caller passed in so the gRPC config matches the bytes we send.
recognition_config_dict["sample_rate_hertz"] = RIVA_TARGET_SAMPLE_RATE_HZ
recognition_config_dict["audio_channel_count"] = RIVA_TARGET_NUM_CHANNELS
recognition_config_dict["encoding"] = "LINEAR_PCM"
response_format = request_payload.get("response_format") or "json"
timestamp_granularities = request_payload.get("timestamp_granularities")
riva_module, riva_asr_module = _import_riva()
auth_obj = self._construct_auth(
riva_module=riva_module,
api_base=api_base,
api_key=api_key,
optional_params=optional_params,
)
recognition_config = self._build_recognition_config_proto(
riva_asr_module=riva_asr_module,
recognition_config_dict=recognition_config_dict,
)
streaming_config = riva_asr_module.StreamingRecognitionConfig(
config=recognition_config, interim_results=False
)
logging_obj.pre_call(
input=None,
api_key=api_key,
additional_args={
"api_base": api_base,
"atranscription": atranscription,
"complete_input_dict": {
"recognition_config": recognition_config_dict,
"nvcf_function_id_set": bool(
optional_params.get("nvcf_function_id")
),
"use_ssl": optional_params.get("use_ssl"),
},
},
)
try:
asr_service = riva_module.ASRService(auth_obj)
audio_chunks = self._iter_audio_chunks(resampled.pcm_bytes)
stream_kwargs: Dict[str, Any] = {
"audio_chunks": audio_chunks,
"streaming_config": streaming_config,
}
# Forward the deadline so the stream cannot block forever if the
# server stalls. Older riva-client versions do not accept a
# ``timeout`` kwarg, so pass it only when supported.
if timeout is not None and self._supports_timeout_kwarg(
asr_service.streaming_response_generator
):
stream_kwargs["timeout"] = float(timeout)
stream = asr_service.streaming_response_generator(**stream_kwargs)
final_results = self._collect_final_results(stream)
except NvidiaRivaException:
raise
except Exception as e:
raise grpc_error_to_litellm_exception(e) from e
transcription = NvidiaRivaAudioTranscriptionConfig.build_transcription_response(
final_results=final_results,
response_format=response_format,
duration_seconds=resampled.duration_seconds,
timestamp_granularities=timestamp_granularities,
)
stringified_response = dict(transcription)
logging_obj.post_call(
input=get_audio_file_name(audio_file),
api_key=api_key,
additional_args={"complete_input_dict": recognition_config_dict},
original_response=stringified_response,
)
hidden_params = {
"model": model,
"custom_llm_provider": "nvidia_riva",
"audio_transcription_duration": resampled.duration_seconds,
}
final_response: TranscriptionResponse = convert_to_model_response_object( # type: ignore
response_object=stringified_response,
model_response_object=model_response,
hidden_params=hidden_params,
response_type="audio_transcription",
)
return final_response
def _construct_auth(
self,
riva_module: Any,
api_base: str,
api_key: Optional[str],
optional_params: dict,
) -> Any:
"""
Build a ``riva.client.Auth`` object.
- When ``nvcf_function_id`` is provided we attach the NVCF
``function-id`` and bearer ``authorization`` metadata, and default
``use_ssl`` to True (NVCF endpoints are TLS-only).
- Otherwise (self-hosted) we default ``use_ssl`` to False but still
honor an explicit override self-hosted Riva behind an ingress
with TLS termination is a real deployment topology.
"""
nvcf_function_id = optional_params.get("nvcf_function_id")
use_ssl_override = optional_params.get("use_ssl")
use_ssl = (
bool(use_ssl_override)
if use_ssl_override is not None
else bool(nvcf_function_id)
)
metadata: List[Tuple[str, str]] = []
if nvcf_function_id:
metadata.append(("function-id", str(nvcf_function_id)))
if api_key:
metadata.append(("authorization", f"Bearer {api_key}"))
try:
return riva_module.Auth(
uri=api_base, use_ssl=use_ssl, metadata_args=metadata
)
except TypeError:
# Older riva-client signatures used positional-only args.
return riva_module.Auth(None, use_ssl, api_base, metadata)
def _build_recognition_config_proto(
self, riva_asr_module: Any, recognition_config_dict: Dict[str, Any]
):
encoding_name = (
recognition_config_dict.get("encoding") or "LINEAR_PCM"
).upper()
encoding_enum = getattr(
riva_asr_module.AudioEncoding,
encoding_name,
riva_asr_module.AudioEncoding.LINEAR_PCM,
)
config = riva_asr_module.RecognitionConfig(
encoding=encoding_enum,
sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]),
language_code=recognition_config_dict["language_code"],
audio_channel_count=int(recognition_config_dict["audio_channel_count"]),
enable_automatic_punctuation=bool(
recognition_config_dict.get("enable_automatic_punctuation", True)
),
enable_word_time_offsets=bool(
recognition_config_dict.get("enable_word_time_offsets", False)
),
max_alternatives=int(recognition_config_dict.get("max_alternatives", 1)),
model=recognition_config_dict.get("model", "") or "",
verbatim_transcripts=bool(
recognition_config_dict.get("verbatim_transcripts", False)
),
profanity_filter=bool(
recognition_config_dict.get("profanity_filter", False)
),
)
endpointing = recognition_config_dict.get("endpointing_config")
if isinstance(endpointing, dict) and endpointing:
try:
ep = riva_asr_module.EndpointingConfig(**endpointing)
config.endpointing_config.CopyFrom(ep)
except Exception:
# If the user supplied an unknown EndpointingConfig field
# (older Riva server), fall back to Riva's defaults rather
# than failing the whole request.
pass
return config
@staticmethod
def _supports_timeout_kwarg(callable_obj: Any) -> bool:
try:
sig = inspect.signature(callable_obj)
except (TypeError, ValueError):
return False
params = sig.parameters
if "timeout" in params:
return True
return any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
@staticmethod
def _iter_audio_chunks(pcm_bytes: bytes):
for offset in range(0, len(pcm_bytes), _DEFAULT_CHUNK_BYTES):
chunk = pcm_bytes[offset : offset + _DEFAULT_CHUNK_BYTES]
if not chunk:
continue
yield chunk
@staticmethod
def _collect_final_results(stream) -> List[Dict[str, Any]]:
"""
Walk the gRPC stream, ignore empty / non-final chunks, and return a
list of normalized final-result dicts. Matching the user's note: the
``id`` blocks with no ``results`` are streaming heartbeats and must
be skipped.
"""
final_results: List[Dict[str, Any]] = []
for response in stream:
results = getattr(response, "results", None) or []
for result in results:
if not getattr(result, "is_final", False):
continue
alternatives = getattr(result, "alternatives", None) or []
if not alternatives:
continue
top = alternatives[0]
transcript = getattr(top, "transcript", "") or ""
words_proto = getattr(top, "words", None) or []
words = []
for word in words_proto:
words.append(
{
"word": getattr(word, "word", ""),
"start_time_ms": int(getattr(word, "start_time", 0) or 0),
"end_time_ms": int(getattr(word, "end_time", 0) or 0),
}
)
final_results.append({"transcript": transcript, "words": words})
return final_results
def _import_riva():
"""
Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``.
We try the SDK first (preferred) and fall back to importing the proto
module separately when the SDK packaging changes between versions.
"""
try:
import riva.client as riva_client # type: ignore
except ImportError as e:
raise NvidiaRivaException(status_code=500, message=_RIVA_INSTALL_HINT) from e
riva_asr_module = riva_client
if not hasattr(riva_asr_module, "RecognitionConfig"):
try:
import riva.client.proto.riva_asr_pb2 as riva_asr_pb2 # type: ignore
riva_asr_module = riva_asr_pb2
except ImportError as e:
raise NvidiaRivaException(
status_code=500, message=_RIVA_INSTALL_HINT
) from e
return riva_client, riva_asr_module

View file

@ -0,0 +1,284 @@
"""
Translates from OpenAI's `/v1/audio/transcriptions` to NVIDIA Riva's gRPC
streaming recognize API.
Riva is gRPC-only, so unlike most providers in this directory the request
"transformation" produced here is a structured dict consumed directly by the
gRPC handler (rather than HTTP form-data). The handler builds Riva
``RecognitionConfig`` / ``StreamingRecognitionConfig`` protobufs from this
dict at call time.
Reference: https://docs.nvidia.com/deeplearning/riva/user-guide/docs/asr/asr-overview.html
"""
from typing import Any, Dict, List, Optional, Union
from httpx import Headers, Response
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIAudioTranscriptionOptionalParams,
)
from litellm.types.utils import FileTypes, TranscriptionResponse
from ...base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
)
from ..common_utils import NvidiaRivaException
# Riva expects a fixed wire format for the audio chunks we stream in.
RIVA_TARGET_SAMPLE_RATE_HZ = 16000
RIVA_TARGET_NUM_CHANNELS = 1
RIVA_TARGET_ENCODING = "LINEAR_PCM"
class NvidiaRivaAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
"""
Config for NVIDIA Riva ASR (gRPC).
Supports both NVCF-hosted (``api_base=grpc.nvcf.nvidia.com:443`` +
``nvcf_function_id``) and self-hosted deployments (any ``host:port``,
optional TLS via ``use_ssl``).
"""
def get_supported_openai_params(
self, model: str
) -> List[OpenAIAudioTranscriptionOptionalParams]:
# Riva natively understands language + word timestamps.
# `response_format` is honored at response-shaping time in the handler.
return ["language", "response_format", "timestamp_granularities"]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
for key, value in non_default_params.items():
if value is None:
continue
if key == "language":
optional_params["language_code"] = self._normalize_language_code(value)
elif key == "timestamp_granularities":
# OpenAI accepts ["word"], ["segment"], or both. Riva only
# natively exposes word timing, so we toggle it on whenever
# "word" is requested. Segment timing is reconstructed in the
# response transformer.
if isinstance(value, list) and "word" in value:
optional_params["enable_word_time_offsets"] = True
optional_params["timestamp_granularities"] = value
elif key == "response_format":
# Stored verbatim; consumed by transform_audio_transcription_response.
optional_params["response_format"] = value
else:
optional_params[key] = value
return optional_params
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, Headers]
) -> BaseLLMException:
return NvidiaRivaException(
message=error_message, status_code=status_code, headers=headers
)
def transform_audio_transcription_request(
self,
model: str,
audio_file: FileTypes,
optional_params: dict,
litellm_params: dict,
) -> AudioTranscriptionRequestData:
"""
Build a structured dict that the gRPC handler consumes. We do *not*
construct protobufs here, so this module remains importable without
``nvidia-riva-client`` being installed (matching how other providers
defer SDK imports to handler-call time).
"""
recognition_config = self._build_recognition_config_dict(
model=model,
optional_params=optional_params,
)
endpointing_config = self._build_endpointing_config_dict(optional_params)
if endpointing_config is not None:
recognition_config["endpointing_config"] = endpointing_config
request_payload: Dict[str, Any] = {
"recognition_config": recognition_config,
"response_format": optional_params.get("response_format") or "json",
"timestamp_granularities": optional_params.get("timestamp_granularities"),
}
return AudioTranscriptionRequestData(data=request_payload, files=None)
def transform_audio_transcription_response(
self,
raw_response: Response,
) -> TranscriptionResponse:
# Not used: Riva responses come from a gRPC stream, not an httpx
# response. The handler calls _build_transcription_response directly.
raise NotImplementedError(
"NvidiaRivaAudioTranscriptionConfig.transform_audio_transcription_response "
"is not used. The handler builds the TranscriptionResponse directly "
"from Riva's gRPC streaming results."
)
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
# gRPC auth is constructed in the handler, not via HTTP headers.
return headers
def _build_recognition_config_dict(
self, model: str, optional_params: dict
) -> Dict[str, Any]:
"""
Build the Riva ``RecognitionConfig`` shape as a plain dict.
``model`` is intentionally left empty when the user has not supplied
``riva_model_name``. Riva auto-selects the right deployment from
``language_code`` + ``sample_rate_hertz``. NVIDIA's internal
deployment names (e.g. ``parakeet-1.1b-en-US-asr-streaming-...``)
change across NIM versions, regions, and self-hosted builds, so
hardcoding any name here would break unpredictably.
"""
return {
"language_code": optional_params.get("language_code", "en-US"),
"sample_rate_hertz": optional_params.get(
"sample_rate_hertz", RIVA_TARGET_SAMPLE_RATE_HZ
),
"encoding": optional_params.get("encoding", RIVA_TARGET_ENCODING),
"audio_channel_count": optional_params.get(
"audio_channel_count", RIVA_TARGET_NUM_CHANNELS
),
"enable_automatic_punctuation": optional_params.get(
"enable_automatic_punctuation", True
),
"enable_word_time_offsets": bool(
optional_params.get("enable_word_time_offsets", False)
),
"max_alternatives": optional_params.get("max_alternatives", 1),
"model": optional_params.get("riva_model_name", ""),
"verbatim_transcripts": optional_params.get("verbatim_transcripts", False),
"profanity_filter": optional_params.get("profanity_filter", False),
}
def _build_endpointing_config_dict(
self, optional_params: dict
) -> Optional[Dict[str, Any]]:
"""
Translate an OpenAI-style ``chunking_strategy`` into Riva's
``EndpointingConfig`` shape, or pass through an explicit
``endpointing_config`` dict.
Returns ``None`` when neither is provided so Riva uses its built-in
VAD defaults.
"""
explicit = optional_params.get("endpointing_config")
if isinstance(explicit, dict):
return dict(explicit)
chunking = optional_params.get("chunking_strategy")
if chunking in (None, "auto"):
return None
if isinstance(chunking, dict) and chunking.get("type") == "server_vad":
config: Dict[str, Any] = {}
if "threshold" in chunking:
threshold = float(chunking["threshold"])
config["start_threshold"] = threshold
config["stop_threshold"] = threshold
if "silence_duration_ms" in chunking:
config["stop_history"] = int(chunking["silence_duration_ms"])
if "prefix_padding_ms" in chunking:
config["stop_history_eou"] = int(chunking["prefix_padding_ms"])
return config or None
return None
@staticmethod
def _normalize_language_code(language: str) -> str:
"""
OpenAI accepts bare ISO-639 codes like ``en``; Riva requires BCP-47
like ``en-US``. Normalize the most common bare codes; pass through
anything that already looks like BCP-47.
"""
if not isinstance(language, str) or not language:
return "en-US"
if "-" in language:
return language
bare_to_bcp47 = {
"en": "en-US",
"es": "es-ES",
"de": "de-DE",
"fr": "fr-FR",
"it": "it-IT",
"pt": "pt-BR",
"ja": "ja-JP",
"ko": "ko-KR",
"zh": "zh-CN",
"ru": "ru-RU",
"hi": "hi-IN",
"ar": "ar-SA",
}
return bare_to_bcp47.get(language.lower(), language)
@staticmethod
def build_transcription_response(
final_results: List[Dict[str, Any]],
response_format: str,
duration_seconds: Optional[float],
timestamp_granularities: Optional[List[str]],
) -> TranscriptionResponse:
"""
Aggregate a list of normalized "final result" dicts into a
``TranscriptionResponse`` shaped for the requested ``response_format``.
Each entry in ``final_results`` is expected to look like::
{
"transcript": str,
"words": [{"word": str, "start_time_ms": int, "end_time_ms": int}, ...],
}
which the handler produces by walking the gRPC stream and keeping
only ``result.is_final`` entries (empty/non-final chunks are
ignored).
"""
full_transcript = "".join(
(item.get("transcript") or "") for item in final_results
).strip()
response = TranscriptionResponse(text=full_transcript)
response["task"] = "transcribe"
if response_format == "verbose_json":
words: List[Dict[str, Any]] = []
if timestamp_granularities and "word" in timestamp_granularities:
for item in final_results:
for word in item.get("words", []) or []:
words.append(
{
"word": word.get("word", ""),
"start": (float(word.get("start_time_ms", 0)) / 1000.0),
"end": float(word.get("end_time_ms", 0)) / 1000.0,
}
)
if words:
response["words"] = words
if duration_seconds is not None:
response["duration"] = duration_seconds
return response

View file

@ -0,0 +1,92 @@
"""
Common utilities and exceptions for the NVIDIA Riva STT provider
"""
from typing import Any, Optional
from litellm.llms.base_llm.chat.transformation import BaseLLMException
class NvidiaRivaException(BaseLLMException):
"""
Exception raised for NVIDIA Riva (gRPC) errors.
``status_code`` is an HTTP-equivalent code derived from the underlying
gRPC ``StatusCode`` (when available) so that litellm's existing error
classifiers (RateLimitError, AuthenticationError, etc.) keep working.
"""
pass
# Mapping from grpc.StatusCode.name -> equivalent HTTP status code.
# Kept as a plain dict (rather than importing grpc enums) so this module is
# importable without grpc installed.
_GRPC_STATUS_CODE_TO_HTTP: dict = {
"OK": 200,
"CANCELLED": 499,
"UNKNOWN": 500,
"INVALID_ARGUMENT": 400,
"DEADLINE_EXCEEDED": 504,
"NOT_FOUND": 404,
"ALREADY_EXISTS": 409,
"PERMISSION_DENIED": 403,
"RESOURCE_EXHAUSTED": 429,
"FAILED_PRECONDITION": 400,
"ABORTED": 409,
"OUT_OF_RANGE": 400,
"UNIMPLEMENTED": 501,
"INTERNAL": 500,
"UNAVAILABLE": 503,
"DATA_LOSS": 500,
"UNAUTHENTICATED": 401,
}
def _extract_grpc_status_name(error: Any) -> Optional[str]:
"""
Best-effort extraction of a gRPC StatusCode name from an arbitrary error.
Works for ``grpc.RpcError`` instances (which expose ``.code()``) as well
as plain exceptions whose string representation contains a status name.
"""
code_fn = getattr(error, "code", None)
if callable(code_fn):
try:
code = code_fn()
except Exception:
code = None
name = getattr(code, "name", None)
if isinstance(name, str):
return name
return None
def _extract_grpc_details(error: Any) -> Optional[str]:
"""Best-effort extraction of a human-readable detail string from a gRPC error."""
details_fn = getattr(error, "details", None)
if callable(details_fn):
try:
details = details_fn()
except Exception:
details = None
if isinstance(details, str) and details:
return details
return None
def grpc_error_to_litellm_exception(error: Exception) -> NvidiaRivaException:
"""
Convert a gRPC error (or any exception raised from the Riva client) into
a ``NvidiaRivaException`` with an appropriate HTTP-equivalent status code.
"""
status_name = _extract_grpc_status_name(error)
http_status = _GRPC_STATUS_CODE_TO_HTTP.get(status_name or "", 500)
detail = _extract_grpc_details(error) or str(error)
message = (
f"NVIDIA Riva gRPC error ({status_name}): {detail}"
if status_name
else f"NVIDIA Riva error: {detail}"
)
return NvidiaRivaException(status_code=http_status, message=message)

View file

@ -6,12 +6,15 @@ This requires websockets, and is currently only supported on LiteLLM Proxy.
from typing import Any, Optional, cast
from litellm._logging import _redact_string
from litellm._logging import _redact_string, verbose_logger
from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
from litellm.types.realtime import RealtimeQueryParams
from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from ....litellm_core_utils.realtime_streaming import RealTimeStreaming
from ....litellm_core_utils.realtime_streaming import (
RealTimeStreaming,
client_sent_openai_beta_realtime_header,
)
from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context
from ..openai import OpenAIChatCompletion
@ -33,21 +36,24 @@ class OpenAIRealtime(OpenAIChatCompletion):
"""
return "https://api.openai.com/"
def _get_additional_headers(self, api_key: str) -> dict:
def _get_additional_headers(
self,
api_key: str,
*,
openai_beta_realtime: bool = False,
) -> dict:
"""
Get additional headers beyond Authorization.
Override this in subclasses to customize headers (e.g., remove OpenAI-Beta).
Headers for the upstream OpenAI Realtime WebSocket.
Args:
api_key: API key for authentication
Returns:
Dictionary of additional headers
When the client sent ``OpenAI-Beta: realtime=v1`` on the proxy WebSocket,
``openai_beta_realtime`` is True and the same header is forwarded upstream
so the legacy beta API is used. GA clients omit that header on the client
connection and must send GA-shaped ``session.update`` payloads.
"""
return {
"Authorization": f"Bearer {api_key}",
"OpenAI-Beta": "realtime=v1",
}
headers: dict = {"Authorization": f"Bearer {api_key}"}
if openai_beta_realtime:
headers["OpenAI-Beta"] = "realtime=v1"
return headers
def _get_ssl_config(self, url: str) -> Any:
"""
@ -120,8 +126,16 @@ class OpenAIRealtime(OpenAIChatCompletion):
# Get provider-specific SSL configuration
ssl_config = self._get_ssl_config(url)
# Get provider-specific headers
headers = self._get_additional_headers(api_key)
openai_beta_realtime = client_sent_openai_beta_realtime_header(websocket)
if not openai_beta_realtime:
verbose_logger.debug(
"OpenAI Realtime: connecting with GA protocol (no OpenAI-Beta header). "
"If your client expects beta event names, add 'OpenAI-Beta: realtime=v1' "
"to the WebSocket headers sent to the LiteLLM proxy."
)
headers = self._get_additional_headers(
api_key, openai_beta_realtime=openai_beta_realtime
)
# Log a masked request preview consistent with other endpoints.
logging_obj.pre_call(

View file

@ -9,7 +9,27 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import GenericStreamingChunk as GChunk
from litellm.types.utils import StreamingChatCompletionChunk
_response_stream_shape_cache = None
def _load_sagemaker_response_stream_shape():
try:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
service_dict = loader.load_service_model("sagemaker-runtime", "service-2")
return ServiceModel(service_dict).shape_for(
"InvokeEndpointWithResponseStreamOutput"
)
except Exception as e:
verbose_logger.warning(
"litellm: could not pre-load sagemaker-runtime response stream shape "
"— SageMaker event-stream decoding will be unavailable. Error: %s",
e,
)
return None
SAGEMAKER_RESPONSE_STREAM_SHAPE = _load_sagemaker_response_stream_shape()
class SagemakerError(BaseLLMException):
@ -187,8 +207,18 @@ class AWSEventStreamDecoder:
verbose_logger.error(f"Final error parsing accumulated JSON: {e}")
def _parse_message_from_event(self, event) -> Optional[str]:
if SAGEMAKER_RESPONSE_STREAM_SHAPE is None:
raise SagemakerError(
status_code=500,
message=(
"SageMaker event-stream shape could not be loaded from botocore. "
"Ensure botocore is correctly installed."
),
)
response_dict = event.to_response_dict()
parsed_response = self.parser.parse(response_dict, get_response_stream_shape())
parsed_response = self.parser.parse(
response_dict, SAGEMAKER_RESPONSE_STREAM_SHAPE
)
if response_dict["status_code"] != 200:
raise ValueError(f"Bad response code, expected 200: {response_dict}")
@ -204,20 +234,3 @@ class AWSEventStreamDecoder:
return None
return chunk.decode() # type: ignore[no-any-return]
def get_response_stream_shape():
global _response_stream_shape_cache
if _response_stream_shape_cache is None:
from botocore.loaders import Loader
from botocore.model import ServiceModel
loader = Loader()
sagemaker_service_dict = loader.load_service_model(
"sagemaker-runtime", "service-2"
)
sagemaker_service_model = ServiceModel(sagemaker_service_dict)
_response_stream_shape_cache = sagemaker_service_model.shape_for(
"InvokeEndpointWithResponseStreamOutput"
)
return _response_stream_shape_cache

View file

@ -28,7 +28,12 @@ class XAIRealtime(OpenAIRealtime):
"""xAI uses a different API base URL."""
return XAI_API_BASE
def _get_additional_headers(self, api_key: str) -> dict:
def _get_additional_headers(
self,
api_key: str,
*,
openai_beta_realtime: bool = False,
) -> dict:
"""
xAI does NOT require the OpenAI-Beta header.
Only send Authorization header.

View file

@ -211,6 +211,12 @@ from .llms.oobabooga.chat import oobabooga
from .llms.openai.completion.handler import OpenAITextCompletion
from .llms.openai.image_variations.handler import OpenAIImageVariationsHandler
from .llms.openai.openai import OpenAIChatCompletion
from .llms.nvidia_riva.audio_transcription.handler import (
NvidiaRivaAudioTranscription,
)
from .llms.nvidia_riva.audio_transcription.transformation import (
NvidiaRivaAudioTranscriptionConfig,
)
from .llms.openai.transcriptions.handler import OpenAIAudioTranscription
from .llms.openai_like.chat.handler import OpenAILikeChatHandler
from .llms.openai_like.embedding.handler import OpenAILikeEmbeddingHandler
@ -266,6 +272,7 @@ from .types.utils import (
openai_chat_completions = OpenAIChatCompletion()
openai_text_completions = OpenAITextCompletion()
openai_audio_transcriptions = OpenAIAudioTranscription()
nvidia_riva_audio_transcriptions = NvidiaRivaAudioTranscription()
openai_image_variations = OpenAIImageVariationsHandler()
groq_chat_completions = GroqChatCompletion()
sap_gen_ai_hub_chat_completions = GenAIHubOrchestration()
@ -1452,14 +1459,14 @@ def completion( # type: ignore # noqa: PLR0915
if eos_token:
custom_prompt_dict[model]["eos_token"] = eos_token
if kwargs.get("model_file_id_mapping"):
messages = update_messages_with_model_file_ids(
messages=messages,
model_id=kwargs.get("model_info", {}).get("id", None),
model_file_id_mapping=cast(
Dict[str, Dict[str, str]], kwargs.get("model_file_id_mapping")
),
)
messages = update_messages_with_model_file_ids(
messages=messages,
model_id=kwargs.get("model_info", {}).get("id", None),
model_file_id_mapping=cast(
Dict[str, Dict[str, str]],
kwargs.get("model_file_id_mapping") or {},
),
)
provider_config: Optional[BaseConfig] = None
if custom_llm_provider is not None and custom_llm_provider in [
@ -6605,6 +6612,26 @@ def transcription(
litellm_params=litellm_params_dict,
shared_session=shared_session,
)
elif custom_llm_provider == "nvidia_riva":
# NVIDIA Riva is gRPC-based, not HTTP. It has its own dedicated handler
# rather than going through base_llm_http_handler.
response = nvidia_riva_audio_transcriptions.audio_transcriptions(
model=model,
audio_file=file,
optional_params=optional_params,
litellm_params=litellm_params_dict,
model_response=model_response,
atranscription=atranscription,
timeout=timeout,
logging_obj=litellm_logging_obj,
api_base=api_base,
api_key=api_key,
provider_config=(
provider_config
if isinstance(provider_config, NvidiaRivaAudioTranscriptionConfig)
else None
),
)
elif provider_config is not None:
response = base_llm_http_handler.audio_transcriptions(
model=model,

View file

@ -28874,6 +28874,19 @@
"mode": "chat",
"output_cost_per_token": 0.0
},
"sambanova/MiniMax-M2.7": {
"input_cost_per_token": 3e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 204800,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.sambanova.ai/plans/pricing",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"sambanova/DeepSeek-R1": {
"input_cost_per_token": 5e-06,
"litellm_provider": "sambanova",

View file

@ -0,0 +1,196 @@
"""
OAuth 2.0 Token Exchange (RFC 8693) handler for MCP servers.
Exchanges a user's incoming JWT (subject_token) for a scoped access token
at an IDP's token exchange endpoint. The exchanged token is then used to
authenticate requests to the upstream MCP server.
See: https://datatracker.ietf.org/doc/html/rfc8693
"""
import asyncio
import hashlib
import weakref
from typing import TYPE_CHECKING, Dict, Tuple
import httpx
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
from litellm.types.mcp_server.mcp_server_manager import MCPServer
# RFC 8693 grant type constant
TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
DEFAULT_SUBJECT_TOKEN_TYPE = "urn:ietf:params:oauth:token-type:access_token"
class TokenExchangeHandler:
"""Handles OAuth 2.0 Token Exchange (RFC 8693) for MCP servers.
Caches exchanged tokens keyed by ``hash(subject_token + server_id)`` so
repeated calls with the same user token skip the IDP round-trip.
"""
def __init__(self) -> None:
self._cache = InMemoryCache(
max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE,
default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL,
)
# WeakValueDictionary so locks are GC'd once no coroutine holds a reference,
# preventing unbounded growth with many rotating user tokens.
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
weakref.WeakValueDictionary()
)
def _get_lock(self, cache_key: str) -> asyncio.Lock:
lock = self._locks.get(cache_key)
if lock is None:
lock = asyncio.Lock()
self._locks[cache_key] = lock
return lock
@staticmethod
def _cache_key(subject_token: str, server_id: str) -> str:
raw = f"{subject_token}:{server_id}"
return hashlib.sha256(raw.encode()).hexdigest()
async def exchange_token(
self,
subject_token: str,
server: "MCPServer",
) -> str:
"""Exchange *subject_token* for a scoped access token.
Returns the exchanged ``access_token`` string (suitable for a
``Bearer`` header).
Raises ``ValueError`` on configuration or IDP errors.
"""
cache_key = self._cache_key(subject_token, server.server_id)
# Fast path
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
# Slow path — one exchange at a time per (user, server) pair
async with self._get_lock(cache_key):
cached = self._cache.get_cache(cache_key)
if cached is not None:
return cached
token, ttl = await self._do_exchange(subject_token, server)
self._cache.set_cache(cache_key, token, ttl=ttl)
return token
async def _do_exchange(
self,
subject_token: str,
server: "MCPServer",
) -> Tuple[str, int]:
"""POST to the token exchange endpoint with RFC 8693 parameters.
Returns ``(access_token, ttl_seconds)``.
"""
endpoint = server.token_exchange_endpoint or server.token_url
if not endpoint:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but no token_exchange_endpoint or token_url configured"
)
if not server.client_id or not server.client_secret:
raise ValueError(
f"MCP server '{server.server_id}' has auth_type=oauth2_token_exchange "
f"but missing client_id or client_secret"
)
data: Dict[str, str] = {
"grant_type": TOKEN_EXCHANGE_GRANT_TYPE,
"subject_token": subject_token,
"subject_token_type": server.subject_token_type
or DEFAULT_SUBJECT_TOKEN_TYPE,
"client_id": server.client_id,
"client_secret": server.client_secret,
}
if server.audience:
data["audience"] = server.audience
if server.scopes:
data["scope"] = " ".join(server.scopes)
verbose_logger.debug(
"Exchanging token for MCP server %s at %s (audience=%s)",
server.server_id,
endpoint,
server.audience,
)
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
try:
response = await client.post(endpoint, data=data)
response.raise_for_status()
except httpx.HTTPStatusError as exc:
verbose_logger.debug(
"Token exchange IDP error for MCP server %s (status %d)",
server.server_id,
exc.response.status_code,
)
raise ValueError(
f"Token exchange for MCP server '{server.server_id}' "
f"failed with status {exc.response.status_code}"
) from exc
body = response.json()
if not isinstance(body, dict):
raise ValueError(
f"Token exchange response for MCP server '{server.server_id}' "
f"returned non-object JSON (got {type(body).__name__})"
)
access_token = body.get("access_token")
if not access_token:
raise ValueError(
f"Token exchange response for MCP server '{server.server_id}' "
f"missing 'access_token'"
)
raw_expires_in = body.get("expires_in")
try:
expires_in = (
int(raw_expires_in)
if raw_expires_in is not None
else MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
)
except (TypeError, ValueError):
expires_in = MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL
ttl = max(
expires_in - MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS,
MCP_OAUTH2_TOKEN_CACHE_MIN_TTL,
)
verbose_logger.info(
"Token exchange succeeded for MCP server %s (expires in %ds)",
server.server_id,
expires_in,
)
return access_token, ttl
def invalidate(self, subject_token: str, server_id: str) -> None:
"""Remove a cached exchanged token (e.g. after a 401)."""
cache_key = self._cache_key(subject_token, server_id)
self._cache.delete_cache(cache_key)
# Module-level singleton
mcp_token_exchange_handler = TokenExchangeHandler()

View file

@ -411,6 +411,15 @@ class MCPServerManager:
aws_role_name=server_config.get("aws_role_name", None),
aws_session_name=server_config.get("aws_session_name", None),
instructions=server_config.get("instructions", None),
# Token Exchange (OBO) fields
token_exchange_endpoint=server_config.get(
"token_exchange_endpoint", None
),
audience=server_config.get("audience", None),
subject_token_type=server_config.get(
"subject_token_type",
"urn:ietf:params:oauth:token-type:access_token",
),
)
self._assign_unique_short_prefix(new_server)
self.config_mcp_servers[server_id] = new_server
@ -765,10 +774,23 @@ class MCPServerManager:
aws_role_name=aws_creds.get("aws_role_name"),
aws_session_name=aws_creds.get("aws_session_name"),
instructions=mcp_server.instructions,
# Token Exchange (OBO) fields — read from credentials JSON blob
token_exchange_endpoint=(
credentials_dict.get("token_exchange_endpoint")
if credentials_dict
else None
),
audience=(credentials_dict.get("audience") if credentials_dict else None),
subject_token_type=(
credentials_dict.get("subject_token_type") if credentials_dict else None
)
or "urn:ietf:params:oauth:token-type:access_token",
)
return new_server
async def _maybe_register_openapi_tools(self, server: MCPServer):
async def _maybe_register_openapi_tools(
self, server: MCPServer, *, initialize_mapping: bool = True
):
"""Register OpenAPI tools if the server has a spec_path configured."""
if server.spec_path:
verbose_logger.info(
@ -779,7 +801,8 @@ class MCPServerManager:
server=server,
base_url=server.url or "",
)
self.initialize_tool_name_to_mcp_server_name_mapping()
if initialize_mapping:
self.initialize_tool_name_to_mcp_server_name_mapping()
async def add_server(self, mcp_server: LiteLLM_MCPServerTable):
try:
@ -1136,6 +1159,29 @@ class MCPServerManager:
#########################################################
# Methods that call the upstream MCP servers
#########################################################
@staticmethod
def _extract_bearer_token(
oauth2_headers: Optional[Dict[str, str]],
raw_headers: Optional[Dict[str, str]],
) -> Optional[str]:
"""Extract the bare Bearer token from oauth2_headers or raw_headers.
Returns the token string without the ``Bearer `` prefix, or ``None``
if no Authorization header is found.
"""
auth_value: Optional[str] = None
if oauth2_headers and "Authorization" in oauth2_headers:
auth_value = oauth2_headers["Authorization"]
elif raw_headers:
# raw_headers may have lowercase keys depending on the ASGI server
normalized = {k.lower(): v for k, v in raw_headers.items()}
auth_value = normalized.get("authorization")
if auth_value:
if auth_value.startswith("Bearer "):
return auth_value[len("Bearer ") :]
return auth_value
return None
def _build_stdio_env(
self,
server: MCPServer,
@ -1169,25 +1215,30 @@ class MCPServerManager:
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
stdio_env: Optional[Dict[str, str]] = None,
subject_token: Optional[str] = None,
) -> MCPClient:
"""
Create an MCPClient instance for the given server.
Auth resolution (single place for all auth logic):
1. ``mcp_auth_header`` per-request/per-user override
2. OAuth2 client_credentials token auto-fetched and cached
3. ``server.authentication_token`` static token from config/DB
2. OAuth2 Token Exchange (OBO) exchange user token for scoped token
3. OAuth2 client_credentials token auto-fetched and cached
4. ``server.authentication_token`` static token from config/DB
Args:
server: The server configuration.
mcp_auth_header: Optional per-request auth override.
extra_headers: Additional headers to forward.
stdio_env: Environment variables for stdio transport.
subject_token: Optional user JWT for token exchange (OBO) flow.
Returns:
Configured MCP client instance.
"""
auth_value = await resolve_mcp_auth(server, mcp_auth_header)
auth_value = await resolve_mcp_auth(
server, mcp_auth_header, subject_token=subject_token
)
transport = server.transport or MCPTransport.sse
@ -1978,7 +2029,11 @@ class MCPServerManager:
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
def _assign_unique_short_prefix(self, server: MCPServer) -> None:
def _assign_unique_short_prefix(
self,
server: MCPServer,
registry: Optional[Dict[str, MCPServer]] = None,
) -> None:
"""Resolve and cache a collision-free short tool prefix on ``server``.
Called at registration time for every MCP server entering the
@ -2002,7 +2057,8 @@ class MCPServerManager:
return
used: Dict[str, str] = {}
for other in self.get_registry().values():
registry_for_collision_check = registry or self.get_registry()
for other in registry_for_collision_check.values():
if other.server_id == server.server_id:
continue
if other.short_prefix:
@ -2534,9 +2590,12 @@ class MCPServerManager:
if server_auth_header is None:
server_auth_header = mcp_auth_header
# oauth2 headers
# Extract subject token for OAuth2 Token Exchange (OBO) flow
subject_token: Optional[str] = None
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.auth_type == MCPAuth.oauth2_token_exchange:
subject_token = self._extract_bearer_token(oauth2_headers, raw_headers)
elif mcp_server.auth_type == MCPAuth.oauth2:
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
extra_headers = None
@ -2604,6 +2663,7 @@ class MCPServerManager:
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
)
call_tool_params = MCPCallToolRequestParams(
@ -2916,46 +2976,72 @@ class MCPServerManager:
# against the *full* set so dedup is deterministic regardless of
# iteration order.
for server in db_mcp_servers:
existing_server = previous_registry.get(server.server_id)
try:
existing_server = previous_registry.get(server.server_id)
if (
existing_server is not None
and existing_server.updated_at is not None
and server.updated_at is not None
and existing_server.updated_at == server.updated_at
):
# Re-use existing server instance to avoid re-running build_mcp_server_from_table()
# which can perform network discovery for OAuth2 servers.
new_registry[server.server_id] = existing_server
continue
if (
existing_server is not None
and existing_server.updated_at is not None
and server.updated_at is not None
and existing_server.updated_at == server.updated_at
):
# Re-use existing server instance to avoid re-running build_mcp_server_from_table()
# which can perform network discovery for OAuth2 servers.
new_registry[server.server_id] = existing_server
continue
_warn_on_server_name_fields(
server_id=server.server_id,
alias=getattr(server, "alias", None),
server_name=getattr(server, "server_name", None),
)
verbose_logger.debug(
f"Building server from DB: {server.server_id} ({server.server_name})"
)
new_server = await self.build_mcp_server_from_table(server)
# Carry the cached short_prefix from the previous registry entry
# (if any) so the prefix is stable across reloads.
if existing_server is not None and existing_server.short_prefix:
new_server.short_prefix = existing_server.short_prefix
new_registry[server.server_id] = new_server
_warn_on_server_name_fields(
server_id=server.server_id,
alias=getattr(server, "alias", None),
server_name=getattr(server, "server_name", None),
)
verbose_logger.debug(
f"Building server from DB: {server.server_id} ({server.server_name})"
)
new_server = await self.build_mcp_server_from_table(server)
# Carry the cached short_prefix from the previous registry entry
# (if any) so the prefix is stable across reloads.
if existing_server is not None and existing_server.short_prefix:
new_server.short_prefix = existing_server.short_prefix
new_registry[server.server_id] = new_server
except Exception as e:
verbose_logger.exception(
"Skipping MCP server %s (%s) during DB reload: %s",
server.server_id,
getattr(server, "alias", None),
e,
)
# Swap in the new registry first so _assign_unique_short_prefix
# sees the complete set when checking for collisions.
self.registry = new_registry
for new_server in new_registry.values():
self._assign_unique_short_prefix(new_server)
# Register OpenAPI tools *after* the final short prefix is assigned
# so the tools are stored in the global registry under the same
# prefix that lookups will use.
await self._maybe_register_openapi_tools(new_server)
# Assign short prefixes against the full candidate set without
# publishing the staged registry to concurrent callers.
registered_registry: Dict[str, MCPServer] = {}
registered_openapi_tools = False
for server_id, new_server in new_registry.items():
try:
self._assign_unique_short_prefix(new_server, registry=new_registry)
# Register OpenAPI tools *after* the final short prefix is assigned
# so the tools are stored in the global registry under the same
# prefix that lookups will use.
await self._maybe_register_openapi_tools(
new_server, initialize_mapping=False
)
registered_registry[server_id] = new_server
if new_server.spec_path:
registered_openapi_tools = True
except Exception as e:
verbose_logger.exception(
"Skipping MCP server %s (%s) during DB reload: %s",
new_server.server_id,
getattr(new_server, "alias", None),
e,
)
self.registry = registered_registry
if registered_openapi_tools:
self.initialize_tool_name_to_mcp_server_name_mapping()
verbose_logger.debug(
"MCP registry refreshed (%s servers in registry)", len(new_registry)
"MCP registry refreshed (%s servers in registry)", len(registered_registry)
)
def get_mcp_servers_from_ids(self, server_ids: List[str]) -> List[MCPServer]:

View file

@ -26,6 +26,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
from litellm.proxy._experimental.mcp_server.auth import token_exchange
from litellm.types.llms.custom_http import httpxSpecialProvider
if TYPE_CHECKING:
@ -50,12 +51,23 @@ class MCPOAuth2TokenCache(InMemoryCache):
def _get_lock(self, server_id: str) -> asyncio.Lock:
return self._locks.setdefault(server_id, asyncio.Lock())
async def async_get_token(self, server: "MCPServer") -> Optional[str]:
@staticmethod
def _has_client_credentials_config(server: "MCPServer") -> bool:
return bool(server.client_id and server.client_secret and server.token_url)
async def async_get_token(
self,
server: "MCPServer",
*,
require_client_credentials_flow: bool = True,
) -> Optional[str]:
"""Return a valid access token, fetching or refreshing as needed.
Returns ``None`` when the server lacks client credentials config.
"""
if not server.has_client_credentials:
if require_client_credentials_flow and not server.has_client_credentials:
return None
if not self._has_client_credentials_config(server):
return None
server_id = server.server_id
@ -263,16 +275,38 @@ mcp_per_user_token_cache = MCPPerUserTokenCache()
async def resolve_mcp_auth(
server: "MCPServer",
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
subject_token: Optional[str] = None,
) -> Optional[Union[str, Dict[str, str]]]:
"""Resolve the auth value for an MCP server.
Priority:
1. ``mcp_auth_header`` per-request/per-user override
2. OAuth2 client_credentials token auto-fetched and cached
3. ``server.authentication_token`` static token from config/DB
2. OAuth2 Token Exchange (OBO / RFC 8693) exchange user token for scoped token
3. OAuth2 client_credentials token auto-fetched and cached
4. ``server.authentication_token`` static token from config/DB
"""
if mcp_auth_header:
return mcp_auth_header
if server.has_token_exchange_config:
if subject_token:
return await token_exchange.mcp_token_exchange_handler.exchange_token(
subject_token, server
)
# No subject_token — fall back to client_credentials using the same client
# credentials and token_url so M2M scenarios still work.
if server.client_id and server.client_secret and server.token_url:
return await mcp_oauth2_token_cache.async_get_token(
server,
require_client_credentials_flow=False,
)
# OBO configured but no subject_token and missing client credentials — warn
# rather than silently proceeding unauthenticated.
verbose_logger.warning(
"MCP server '%s' is configured for token exchange (OBO) but no subject_token "
"was provided and client credentials (client_id/client_secret/token_url) are "
"incomplete. The request will proceed without authentication.",
server.server_id,
)
if server.has_client_credentials:
return await mcp_oauth2_token_cache.async_get_token(server)
return server.authentication_token

View file

@ -6,10 +6,34 @@ import asyncio
import contextvars
import json
import os
import re
from pathlib import PurePosixPath
from typing import Any, Dict, List, Optional
from urllib.parse import quote
# Tool names emitted from OpenAPI specs must work across all major LLM providers.
# OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to
# ^[a-zA-Z0-9_-]+$ on tool names. Many specs (notably GitHub's REST API) use
# tag-namespaced operationIds like "actions/download-job-logs-for-workflow-run"
# which include '/'. Sanitize here so the same regex passes everywhere downstream.
_OPENAPI_TOOL_NAME_INVALID_CHARS = re.compile(r"[^a-zA-Z0-9_-]")
_OPENAPI_TOOL_NAME_MAX_LEN = 128
def sanitize_openapi_tool_name(raw_name: str) -> str:
"""Map an OpenAPI operationId / fallback to a provider-safe tool name.
Replaces any character outside ``[a-zA-Z0-9_-]`` with ``_`` and caps the
result at 128 chars (the most restrictive of the major providers).
Lowercased to match the existing convention in
``register_tools_from_openapi``.
"""
if not raw_name:
return raw_name
sanitized = _OPENAPI_TOOL_NAME_INVALID_CHARS.sub("_", raw_name).lower()
return sanitized[:_OPENAPI_TOOL_NAME_MAX_LEN]
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -399,17 +423,36 @@ def create_tool_function(
def register_tools_from_openapi(spec: Dict[str, Any], base_url: str):
"""Register MCP tools from OpenAPI specification."""
paths = spec.get("paths", {})
used_names: set = set()
for path, path_item in paths.items():
for method in ["get", "post", "put", "delete", "patch"]:
if method in path_item:
operation = path_item[method]
# Generate tool name
operation_id = operation.get(
"operationId", f"{method}_{path.replace('/', '_')}"
)
tool_name = operation_id.replace(" ", "_").lower()
# Generate tool name. Sanitize to ^[a-zA-Z0-9_-]+$ (lowercase)
# so the resulting name is valid across OpenAI/Anthropic/Bedrock.
# Many specs (e.g. GitHub REST) use tag-namespaced operationIds
# like "actions/download-job-logs-for-workflow-run" which
# contain '/' and would 400 at the LLM provider boundary.
operation_id = operation.get("operationId", f"{method}_{path}")
tool_name = sanitize_openapi_tool_name(operation_id)
# Disambiguate collisions: two operationIds that differ only
# by sanitized characters (e.g. "foo/list" and "foo.list")
# would both become "foo_list". Append _2, _3, … to keep
# every tool reachable, mirroring the Anthropic-side logic
# in _build_anthropic_tool_name_maps.
unique = tool_name
n = 1
while unique in used_names:
n += 1
suffix = f"_{n}"
unique = (
tool_name[: _OPENAPI_TOOL_NAME_MAX_LEN - len(suffix)] + suffix
)
tool_name = unique
used_names.add(tool_name)
# Get description
description = operation.get(

View file

@ -857,6 +857,7 @@ if MCP_AVAILABLE:
########################################################
from litellm.proxy.management_endpoints.mcp_management_endpoints import (
NewMCPServerRequest,
_inherit_credentials_from_existing_server,
)
def _extract_credentials(
@ -975,9 +976,11 @@ if MCP_AVAILABLE:
async def _preview_openapi_tools(spec_path: str) -> dict:
"""Generate tool previews from an OpenAPI spec without creating a server."""
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
_OPENAPI_TOOL_NAME_MAX_LEN,
build_input_schema,
load_openapi_spec_async,
resolve_operation_params,
sanitize_openapi_tool_name,
)
try:
@ -985,8 +988,9 @@ if MCP_AVAILABLE:
paths = spec.get("paths", {})
components = spec.get("components", {})
tools: List[dict] = []
used_names: set = set()
for path, path_item in paths.items():
for method in ("get", "post", "put", "patch", "delete"):
for method in ("get", "post", "put", "delete", "patch"):
operation = path_item.get(method)
if operation is None:
continue
@ -995,7 +999,23 @@ if MCP_AVAILABLE:
operation, path_item, components
)
op_id = operation.get("operationId", f"{method}_{path}")
raw_op_id = operation.get("operationId", f"{method}_{path}")
# Match what register_tools_from_openapi does so the preview
# the user sees in the dashboard equals the names that get
# registered (and shipped to LLM providers, which enforce
# ^[a-zA-Z0-9_-]+$). See sanitize_openapi_tool_name docstring.
op_id = sanitize_openapi_tool_name(raw_op_id)
unique = op_id
n = 1
while unique in used_names:
n += 1
suffix = f"_{n}"
unique = (
op_id[: _OPENAPI_TOOL_NAME_MAX_LEN - len(suffix)] + suffix
)
op_id = unique
used_names.add(op_id)
summary = operation.get("summary", "")
description = operation.get("description", summary)
input_schema = build_input_schema(resolved_op)
@ -1068,6 +1088,10 @@ if MCP_AVAILABLE:
},
)
new_mcp_server_request = _inherit_credentials_from_existing_server(
new_mcp_server_request
)
# For OpenAPI spec servers, generate tools from the spec directly
if new_mcp_server_request.spec_path:
return await _preview_openapi_tools(new_mcp_server_request.spec_path)

View file

@ -656,6 +656,13 @@ class LiteLLMRoutes(enum.Enum):
"/health/services",
] + info_routes
# Stateless validators on caller-supplied log data; source logs are
# already accessible via spend_tracking_routes, so no scope expansion.
compliance_check_routes = [
"/compliance/eu-ai-act",
"/compliance/gdpr",
]
# Routes in `global_spend_tracking_routes` return proxy-wide spend across
# every team, customer, and api_key. They are intentionally NOT included
# here — non-admin roles must not see other tenants' spend. Admin roles go
@ -675,6 +682,7 @@ class LiteLLMRoutes(enum.Enum):
]
+ spend_tracking_routes
+ key_management_routes
+ compliance_check_routes
)
internal_user_view_only_routes = spend_tracking_routes
@ -3348,6 +3356,19 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
],
)
azure_sentinel: CallbackOnUI = CallbackOnUI(
litellm_callback_name="azure_sentinel",
ui_callback_name="Azure Sentinel",
litellm_callback_params=[
"AZURE_SENTINEL_DCR_IMMUTABLE_ID",
"AZURE_SENTINEL_ENDPOINT",
"AZURE_SENTINEL_TENANT_ID",
"AZURE_SENTINEL_CLIENT_ID",
"AZURE_SENTINEL_CLIENT_SECRET",
"AZURE_SENTINEL_STREAM_NAME",
],
)
openmeter: CallbackOnUI = CallbackOnUI(
litellm_callback_name="openmeter",
ui_callback_name="OpenMeter",

View file

@ -3516,7 +3516,6 @@ async def _check_team_member_budget(
if (
team_object is not None
and team_object.team_id is not None
and user_object is not None
and valid_token is not None
and valid_token.user_id is not None
):
@ -3619,6 +3618,7 @@ async def _check_team_member_model_access(
llm_router=llm_router,
models=member_allowed_models,
object_type="team",
team_id=team_object.team_id,
)
except ProxyException:
raise ProxyException(

View file

@ -1512,7 +1512,7 @@ class ProxyBaseLLMRequestProcessing:
status_code=result.status_code,
headers=HttpPassThroughEndpointHelpers.get_response_headers(
headers=result.headers,
custom_headers=None,
custom_headers=dict(fastapi_response.headers),
),
)

View file

@ -5,8 +5,10 @@ from typing import Optional
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS,
SPEND_LOG_CLEANUP_BATCH_SIZE,
SPEND_LOG_CLEANUP_JOB_NAME,
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES,
SPEND_LOG_RUN_LOOPS,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
@ -74,6 +76,7 @@ class SpendLogCleanup:
"""
total_deleted = 0
run_count = 0
consecutive_failures = 0
while True:
if run_count > SPEND_LOG_RUN_LOOPS:
verbose_proxy_logger.info(
@ -82,18 +85,50 @@ class SpendLogCleanup:
break
# Step 1: Find logs and delete them in one go without fetching to application
# Delete in batches, limited by self.batch_size
deleted_result = await prisma_client.db.execute_raw(
"""
DELETE FROM "LiteLLM_SpendLogs"
WHERE "request_id" IN (
SELECT "request_id" FROM "LiteLLM_SpendLogs"
WHERE "startTime" < $1::timestamptz
LIMIT $2
try:
deleted_result = await prisma_client.db.execute_raw(
"""
DELETE FROM "LiteLLM_SpendLogs"
WHERE "request_id" IN (
SELECT "request_id" FROM "LiteLLM_SpendLogs"
WHERE "startTime" < $1::timestamptz
LIMIT $2
)
""",
cutoff_date,
self.batch_size,
)
""",
cutoff_date,
self.batch_size,
)
except Exception as batch_exc:
# A single batch failure (e.g. Prisma/DB timeout) must not abort
# the whole run — subsequent batches may still succeed.
consecutive_failures += 1
verbose_proxy_logger.exception(
"Spend log cleanup batch failed "
"(run_count=%d, consecutive_failures=%d, batch_size=%d, "
"cutoff=%s, total_deleted_so_far=%d): %s: %s",
run_count,
consecutive_failures,
self.batch_size,
cutoff_date.isoformat(),
total_deleted,
type(batch_exc).__name__,
batch_exc,
)
if (
consecutive_failures
>= SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES
):
verbose_proxy_logger.error(
"Aborting spend log cleanup after %d consecutive batch "
"failures; total deleted before abort: %d",
consecutive_failures,
total_deleted,
)
break
await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
continue
consecutive_failures = 0
deleted_count = 0
if isinstance(deleted_result, int):
@ -168,7 +203,13 @@ class SpendLogCleanup:
verbose_proxy_logger.info(f"Deleted {total_deleted} logs")
except Exception as e:
verbose_proxy_logger.error(f"Error during cleanup: {str(e)}")
# .exception() captures the traceback; str(e) alone on a Prisma/DB
# timeout is often empty and gives operators no signal to diagnose.
verbose_proxy_logger.exception(
"Error during spend log cleanup: %s: %s",
type(e).__name__,
e,
)
return # Return after error handling
finally:
# Only release the lock if it was actually acquired

View file

@ -1,8 +0,0 @@
from fastapi import FastAPI
from litellm.proxy.health_endpoints._health_endpoints import router as health_router
def build_health_app():
health_app = FastAPI(title="LiteLLM Health Endpoints")
health_app.include_router(health_router)
return health_app

View file

@ -498,6 +498,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"error": f"Model capacity reached for {model}. "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Model TPM: {model_group_info.tpm if model_group_info.tpm is not None else 'not configured'}, "
f"Model RPM: {model_group_info.rpm if model_group_info.rpm is not None else 'not configured'}, "
f"Remaining: {status['limit_remaining']}"
},
headers={
@ -515,8 +517,11 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Model: {model}, "
f"Priority: {priority}, "
f"Rate limit type: {status['rate_limit_type']}, "
f"Model TPM: {model_group_info.tpm if model_group_info.tpm is not None else 'not configured'}, "
f"Model RPM: {model_group_info.rpm if model_group_info.rpm is not None else 'not configured'}, "
f"Remaining: {status['limit_remaining']}, "
f"Model saturation: {saturation:.1%}"
},

File diff suppressed because it is too large Load diff

View file

@ -1667,7 +1667,10 @@ async def add_litellm_data_to_request( # noqa: PLR0915
)
if tags is not None and _admin_allow_client_tags:
data[_metadata_variable_name]["tags"] = tags
data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags(
request_tags=data[_metadata_variable_name].get("tags"),
tags_to_add=tags,
)
elif tags is not None:
verbose_proxy_logger.warning(
"Ignored caller-supplied tags from header/root body: this "

View file

@ -1,3 +1,4 @@
import asyncio
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
@ -543,6 +544,13 @@ def _build_aggregated_sql_query(
where_clause = " AND ".join(sql_conditions)
# Postgres computes every rollup level the response needs — per-date
# totals, per-(date, model), per-(date, model, api_key), per-provider,
# etc. — in a single pass via GROUPING SETS. The GROUPING() bitmask
# encodes which level a row belongs to so Python can dispatch rows
# straight into their buckets without re-summing. The leaf grouping
# is omitted on purpose: nothing in the response shape needs it once
# all the rollups are present.
sql_query = f"""
SELECT
date,
@ -552,6 +560,9 @@ def _build_aggregated_sql_query(
custom_llm_provider,
mcp_namespaced_tool_name,
endpoint,
GROUPING(date, api_key, model, model_group,
custom_llm_provider, mcp_namespaced_tool_name,
endpoint) AS group_level,
SUM(spend)::float AS spend,
SUM(prompt_tokens)::bigint AS prompt_tokens,
SUM(completion_tokens)::bigint AS completion_tokens,
@ -562,32 +573,35 @@ def _build_aggregated_sql_query(
SUM(failed_requests)::bigint AS failed_requests
FROM "{pg_table}"
WHERE {where_clause}
GROUP BY date, api_key, model, model_group, custom_llm_provider,
mcp_namespaced_tool_name, endpoint
ORDER BY date DESC
GROUP BY GROUPING SETS (
(date),
(date, api_key),
(date, model),
(date, model, api_key),
(date, model_group),
(date, model_group, api_key),
(date, custom_llm_provider),
(date, custom_llm_provider, api_key),
(date, mcp_namespaced_tool_name),
(date, mcp_namespaced_tool_name, api_key),
(date, endpoint),
(date, endpoint, api_key),
()
)
"""
return sql_query, sql_params
async def _aggregate_spend_records(
def _aggregate_spend_records_sync(
*,
prisma_client: PrismaClient,
records: List[Any],
api_key_metadata: Dict[str, Dict[str, Any]],
entity_id_field: Optional[str],
entity_metadata_field: Optional[Dict[str, dict]],
) -> Dict[str, Any]:
"""Aggregate rows into DailySpendData list and total metrics."""
api_keys: Set[str] = set()
for record in records:
if record.api_key:
api_keys.add(record.api_key)
api_key_metadata: Dict[str, Dict[str, Any]] = {}
model_metadata: Dict[str, Dict[str, Any]] = {}
provider_metadata: Dict[str, Dict[str, Any]] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
results: List[DailySpendData] = []
total_metrics = SpendMetrics()
@ -631,6 +645,228 @@ async def _aggregate_spend_records(
return {"results": results, "totals": total_metrics}
async def _aggregate_spend_records(
*,
prisma_client: PrismaClient,
records: List[Any],
entity_id_field: Optional[str],
entity_metadata_field: Optional[Dict[str, dict]],
) -> Dict[str, Any]:
"""Aggregate rows into DailySpendData list and total metrics.
The per-row loop is offloaded to a worker thread via asyncio.to_thread so
a large result set doesn't peg the event loop.
"""
api_keys: Set[str] = {record.api_key for record in records if record.api_key}
api_key_metadata: Dict[str, Dict[str, Any]] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
return await asyncio.to_thread(
_aggregate_spend_records_sync,
records=records,
api_key_metadata=api_key_metadata,
entity_id_field=entity_id_field,
entity_metadata_field=entity_metadata_field,
)
# GROUPING() bitmask values for each grouping set emitted by
# _build_aggregated_sql_query. Per Postgres semantics, the rightmost argument
# is the least-significant bit. Argument order:
# date, api_key, model, model_group, custom_llm_provider,
# mcp_namespaced_tool_name, endpoint
# A bit is 1 when the corresponding column is rolled up (i.e. NOT in the
# current grouping set's key), 0 when the column is part of the key.
_GROUP_GRAND_TOTAL = 127 # 0b1111111 — all rolled up
_GROUP_DATE = 63 # 0b0111111 — only date kept
_GROUP_DATE_API_KEY = 31 # 0b0011111
_GROUP_DATE_MODEL = 47 # 0b0101111
_GROUP_DATE_MODEL_API_KEY = 15 # 0b0001111
_GROUP_DATE_MODEL_GROUP = 55 # 0b0110111
_GROUP_DATE_MODEL_GROUP_API_KEY = 23 # 0b0010111
_GROUP_DATE_PROVIDER = 59 # 0b0111011
_GROUP_DATE_PROVIDER_API_KEY = 27 # 0b0011011
_GROUP_DATE_MCP = 61 # 0b0111101
_GROUP_DATE_MCP_API_KEY = 29 # 0b0011101
_GROUP_DATE_ENDPOINT = 62 # 0b0111110
_GROUP_DATE_ENDPOINT_API_KEY = 30 # 0b0011110
def _record_to_spend_metrics(record: Any) -> SpendMetrics:
"""Build a SpendMetrics directly from one already-aggregated rollup row."""
return SpendMetrics(
spend=record.spend,
prompt_tokens=record.prompt_tokens,
completion_tokens=record.completion_tokens,
total_tokens=record.prompt_tokens + record.completion_tokens,
cache_read_input_tokens=record.cache_read_input_tokens,
cache_creation_input_tokens=record.cache_creation_input_tokens,
api_requests=record.api_requests,
successful_requests=record.successful_requests,
failed_requests=record.failed_requests,
)
def _key_metadata(
api_key_metadata: Dict[str, Dict[str, Any]], api_key: str
) -> KeyMetadata:
meta = api_key_metadata.get(api_key, {})
return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id"))
def _aggregate_grouping_sets_records_sync( # noqa: PLR0915
*,
records: List[Any],
api_key_metadata: Dict[str, Dict[str, Any]],
) -> Dict[str, Any]:
"""Build the response from rollup rows produced by the GROUPING SETS query.
Each row carries a `group_level` bitmask (from Postgres GROUPING()) that
identifies which rollup level it belongs to. We dispatch the row's
pre-aggregated metrics straight into the matching bucket no per-row
summing in Python and no nested update_metrics calls.
"""
total_metrics = SpendMetrics()
grouped_data: Dict[str, Dict[str, Any]] = {}
def ensure_date(date_str: str) -> Dict[str, Any]:
bucket = grouped_data.get(date_str)
if bucket is None:
bucket = {"metrics": SpendMetrics(), "breakdown": BreakdownMetrics()}
grouped_data[date_str] = bucket
return bucket
def assign_metric_with_metadata(
target: Dict[str, MetricWithMetadata], key: str, metrics: SpendMetrics
) -> None:
existing = target.get(key)
if existing is None:
target[key] = MetricWithMetadata(metrics=metrics, metadata={})
else:
existing.metrics = metrics
def assign_api_key_breakdown(
target: Dict[str, MetricWithMetadata],
parent_key: str,
api_key: str,
metrics: SpendMetrics,
) -> None:
parent = target.get(parent_key)
if parent is None:
parent = MetricWithMetadata(metrics=SpendMetrics(), metadata={})
target[parent_key] = parent
parent.api_key_breakdown[api_key] = KeyMetricWithMetadata(
metrics=metrics, metadata=_key_metadata(api_key_metadata, api_key)
)
for record in records:
level = record.group_level
metrics = _record_to_spend_metrics(record)
if level == _GROUP_GRAND_TOTAL:
total_metrics = metrics
continue
if level == _GROUP_DATE:
ensure_date(record.date)["metrics"] = metrics
continue
breakdown = ensure_date(record.date)["breakdown"]
if level == _GROUP_DATE_API_KEY:
if record.api_key:
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
metrics=metrics,
metadata=_key_metadata(api_key_metadata, record.api_key),
)
elif level == _GROUP_DATE_MODEL:
if record.model:
assign_metric_with_metadata(breakdown.models, record.model, metrics)
elif level == _GROUP_DATE_MODEL_API_KEY:
if record.model and record.api_key:
assign_api_key_breakdown(
breakdown.models, record.model, record.api_key, metrics
)
elif level == _GROUP_DATE_MODEL_GROUP:
if record.model_group:
assign_metric_with_metadata(
breakdown.model_groups, record.model_group, metrics
)
elif level == _GROUP_DATE_MODEL_GROUP_API_KEY:
if record.model_group and record.api_key:
assign_api_key_breakdown(
breakdown.model_groups,
record.model_group,
record.api_key,
metrics,
)
elif level == _GROUP_DATE_PROVIDER:
provider = record.custom_llm_provider or "unknown"
assign_metric_with_metadata(breakdown.providers, provider, metrics)
elif level == _GROUP_DATE_PROVIDER_API_KEY:
if record.api_key:
provider = record.custom_llm_provider or "unknown"
assign_api_key_breakdown(
breakdown.providers, provider, record.api_key, metrics
)
elif level == _GROUP_DATE_MCP:
if record.mcp_namespaced_tool_name:
assign_metric_with_metadata(
breakdown.mcp_servers, record.mcp_namespaced_tool_name, metrics
)
elif level == _GROUP_DATE_MCP_API_KEY:
if record.mcp_namespaced_tool_name and record.api_key:
assign_api_key_breakdown(
breakdown.mcp_servers,
record.mcp_namespaced_tool_name,
record.api_key,
metrics,
)
elif level == _GROUP_DATE_ENDPOINT:
if record.endpoint:
assign_metric_with_metadata(
breakdown.endpoints, record.endpoint, metrics
)
elif level == _GROUP_DATE_ENDPOINT_API_KEY:
if record.endpoint and record.api_key:
assign_api_key_breakdown(
breakdown.endpoints, record.endpoint, record.api_key, metrics
)
results = [
DailySpendData(
date=datetime.strptime(date_str, "%Y-%m-%d").date(),
metrics=data["metrics"],
breakdown=data["breakdown"],
)
for date_str, data in grouped_data.items()
]
results.sort(key=lambda x: x.date, reverse=True)
return {"results": results, "totals": total_metrics}
async def _aggregate_grouping_sets_records(
*,
prisma_client: PrismaClient,
records: List[Any],
) -> Dict[str, Any]:
"""Async wrapper: fetch api_key_metadata, then dispatch on a worker thread."""
api_keys: Set[str] = {r.api_key for r in records if r.api_key}
api_key_metadata: Dict[str, Dict[str, Any]] = {}
if api_keys:
api_key_metadata = await get_api_key_metadata(prisma_client, api_keys)
return await asyncio.to_thread(
_aggregate_grouping_sets_records_sync,
records=records,
api_key_metadata=api_key_metadata,
)
async def get_daily_activity(
prisma_client: Optional[PrismaClient],
table_name: str,
@ -771,21 +1007,18 @@ async def get_daily_activity_aggregated(
timezone_offset_minutes=timezone_offset_minutes,
)
# Execute GROUP BY query — returns pre-aggregated dicts
# Execute GROUPING SETS query — returns one row per rollup level.
rows = await prisma_client.db.query_raw(sql_query, *sql_params)
if rows is None:
rows = []
# Convert dicts to objects for compatibility with _aggregate_spend_records
records = [SimpleNamespace(**row) for row in rows]
# entity_id_field=None skips entity breakdown (entity dimension was
# collapsed by the GROUP BY, so per-entity data is not available)
aggregated = await _aggregate_spend_records(
# The grouping-sets dispatcher places each row directly in its bucket
# using the row's GROUPING() bitmask. No Python-side summing needed.
aggregated = await _aggregate_grouping_sets_records(
prisma_client=prisma_client,
records=records,
entity_id_field=None,
entity_metadata_field=None,
)
return SpendAnalyticsPaginatedResponse(

View file

@ -152,8 +152,14 @@ if MCP_AVAILABLE:
UserAPIKeyAuth,
UserMCPManagementMode,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.auth.user_api_key_auth import (
_user_api_key_auth_builder,
user_api_key_auth,
)
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
populate_request_with_path_params,
)
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.types.mcp import MCPCredentials
@ -1492,6 +1498,55 @@ if MCP_AVAILABLE:
return _redact_mcp_credentials(temp_record)
async def _mcp_oauth_user_api_key_auth(request: Request) -> UserAPIKeyAuth:
"""
Auth dependency for MCP OAuth browser-navigation endpoints (/authorize, /token).
Tries the Authorization header first. Falls back to decoding the UI
'token' session cookie (set by SSO login) to extract the API key, which
allows browser-based OAuth redirects to work without an explicit
Authorization header.
"""
import jwt as _jwt
from litellm.proxy.proxy_server import master_key
auth_header = request.headers.get("Authorization", "")
api_key = auth_header # _get_bearer_token will strip "Bearer " prefix
if not api_key:
token_cookie = request.cookies.get("token")
if token_cookie and master_key:
try:
decoded = _jwt.decode(
token_cookie,
master_key,
algorithms=["HS256"],
# UI session cookies may omit exp; don't require it.
options={"verify_exp": False},
)
if decoded.get("login_method") in ("sso", "username_password"):
cookie_key = decoded.get("key", "")
if cookie_key:
api_key = f"Bearer {cookie_key}"
except _jwt.InvalidTokenError:
pass
request_data = await _read_request_body(request=request)
request_data = populate_request_with_path_params(
request_data=request_data, request=request
)
return await _user_api_key_auth_builder(
request=request,
api_key=api_key,
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data=request_data,
)
async def _get_cached_temporary_mcp_server_or_404(
server_id: str,
user_api_key_dict: UserAPIKeyAuth,
@ -1542,12 +1597,12 @@ if MCP_AVAILABLE:
@router.get(
"/server/oauth/{server_id}/authorize",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
dependencies=[Depends(_mcp_oauth_user_api_key_auth)],
)
async def mcp_authorize(
request: Request,
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth),
client_id: Optional[str] = None,
redirect_uri: str = Query(...),
state: str = "",
@ -1587,12 +1642,12 @@ if MCP_AVAILABLE:
@router.post(
"/server/oauth/{server_id}/token",
include_in_schema=False,
dependencies=[Depends(user_api_key_auth)],
dependencies=[Depends(_mcp_oauth_user_api_key_auth)],
)
async def mcp_token(
request: Request,
server_id: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth),
grant_type: str = Form(...),
code: Optional[str] = Form(None),
redirect_uri: Optional[str] = Form(None),

View file

@ -4,6 +4,7 @@
This is an enterprise feature and requires a premium license.
"""
import re
from typing import Any, Dict, List, Optional, Set, Tuple
from fastapi import (
@ -843,6 +844,18 @@ async def get_service_provider_config(request: Request):
return SCIMServiceProviderConfig(meta=meta)
def _parse_scim_eq_filter(scim_filter: str) -> Optional[Tuple[str, str]]:
"""Parse the SCIM equality filters Okta uses before user lifecycle changes."""
match = re.match(
r"""\s*([\w.]+)\s+eq\s+(['"]?)(.*?)\2\s*$""",
scim_filter,
flags=re.IGNORECASE,
)
if not match:
return None
return match.group(1).lower(), match.group(3)
# User Endpoints
@scim_router.get(
"/Users",
@ -867,15 +880,21 @@ async def get_users(
try:
prisma_client = await _get_prisma_client_or_raise_exception()
# Parse filter if provided (basic support)
where_conditions = {}
where_conditions: Dict[str, Any] = {}
if filter:
# Very basic filter support - only handling userName eq and emails.value eq
if "userName eq" in filter:
user_id = filter.split("userName eq ")[1].strip("\"'")
where_conditions["user_id"] = user_id
elif "emails.value eq" in filter:
email = filter.split("emails.value eq ")[1].strip("\"'")
where_conditions["user_email"] = email
# Okta locates users by userName before deprovisioning. LiteLLM
# exposes SCIM userName from user_email, while older SCIM-created
# users may still have user_id == userName, so support both.
parsed_filter = _parse_scim_eq_filter(filter)
if parsed_filter:
filter_attribute, filter_value = parsed_filter
if filter_attribute == "username":
where_conditions["OR"] = [
{"user_email": filter_value},
{"user_id": filter_value},
]
elif filter_attribute == "emails.value":
where_conditions["user_email"] = filter_value
# Get users from database
users: List[LiteLLM_UserTable] = (

View file

@ -46,25 +46,46 @@ def get_audit_log_changed_by(
def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]:
"""Resolve a string callback name to a CustomLogger instance, with caching."""
"""Resolve a string callback name to a CustomLogger instance, with caching.
For "s3_v2" with `litellm.s3_audit_callback_params` set, constructs a
dedicated `S3Logger` so audit logs can target a different bucket than the
normal-log singleton served by `_init_custom_logger_compatible_class`.
"""
if name in _audit_log_callback_cache:
return _audit_log_callback_cache[name]
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
instance: Optional[CustomLogger]
if (
name == "s3_v2"
and getattr(litellm, "s3_audit_callback_params", None) is not None
):
from litellm.integrations.s3_v2 import S3Logger as S3V2Logger
instance = _init_custom_logger_compatible_class(
logging_integration=name, # type: ignore
internal_usage_cache=None,
llm_router=None,
)
instance = S3V2Logger(
s3_callback_params_override=litellm.s3_audit_callback_params
)
else:
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
)
instance = _init_custom_logger_compatible_class(
logging_integration=name, # type: ignore
internal_usage_cache=None,
llm_router=None,
)
if instance is not None:
_audit_log_callback_cache[name] = instance
return instance
def reset_audit_log_callback_cache() -> None:
"""Clear cached audit-log callback instances. Call on config reload."""
_audit_log_callback_cache.clear()
def _build_audit_log_payload(
request_data: LiteLLM_AuditLogs,
) -> StandardAuditLogPayload:

View file

@ -0,0 +1,121 @@
import json
from typing import Callable, Optional, Union
from starlette.types import ASGIApp, Message, Receive, Scope, Send
MaxRequestSizeGetter = Callable[[], Optional[Union[int, float]]]
RequestSizeLimitEnabledGetter = Callable[[], bool]
class RequestEntityTooLarge(Exception):
pass
class RequestSizeLimitMiddleware:
"""
Reject oversized requests before downstream auth/routes parse the body.
Content-Length can be rejected without reading any body bytes. Requests
without Content-Length are counted as the ASGI stream is consumed, limiting
memory exposure to the configured threshold plus the current chunk.
"""
def __init__(
self,
app: ASGIApp,
get_max_request_size_mb: MaxRequestSizeGetter,
is_request_size_limit_enabled: RequestSizeLimitEnabledGetter,
) -> None:
self.app = app
self.get_max_request_size_mb = get_max_request_size_mb
self.is_request_size_limit_enabled = is_request_size_limit_enabled
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
max_request_size_mb = self.get_max_request_size_mb()
max_request_size_bytes = _mb_to_bytes(max_request_size_mb)
if max_request_size_bytes is None or not self.is_request_size_limit_enabled():
await self.app(scope, receive, send)
return
content_length = _get_content_length(scope=scope)
if content_length is not None and content_length > max_request_size_bytes:
await _send_request_too_large(
send=send, max_request_size_mb=max_request_size_mb
)
return
received_body_bytes = 0
response_started = False
async def limited_receive() -> Message:
nonlocal received_body_bytes
message = await receive()
if message["type"] != "http.request":
return message
received_body_bytes += len(message.get("body", b""))
if received_body_bytes > max_request_size_bytes:
raise RequestEntityTooLarge
return message
async def tracking_send(message: Message) -> None:
nonlocal response_started
if message["type"] == "http.response.start":
response_started = True
await send(message)
try:
await self.app(scope, limited_receive, tracking_send)
except RequestEntityTooLarge:
if response_started:
raise
await _send_request_too_large(
send=send, max_request_size_mb=max_request_size_mb
)
def _mb_to_bytes(max_request_size_mb: Optional[Union[int, float]]) -> Optional[int]:
if max_request_size_mb is None:
return None
if max_request_size_mb <= 0:
return None
return int(max_request_size_mb * 1024 * 1024)
def _get_content_length(scope: Scope) -> Optional[int]:
headers = dict(scope.get("headers") or [])
raw_content_length = headers.get(b"content-length")
if raw_content_length is None:
return None
try:
return int(raw_content_length)
except ValueError:
return None
async def _send_request_too_large(
send: Send,
max_request_size_mb: Optional[Union[int, float]],
) -> None:
body = json.dumps(
{"error": f"Request size is too large. Max size is {max_request_size_mb} MB"},
separators=(",", ":"),
).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 413,
"headers": [
(b"content-type", b"application/json"),
(b"content-length", str(len(body)).encode("latin-1")),
],
}
)
await send({"type": "http.response.body", "body": body, "more_body": False})

View file

@ -169,6 +169,66 @@ class ProxyInitializationHelpers:
)
return uvicorn_args
@staticmethod
def _get_reload_options(config_path: Optional[str]) -> dict:
"""Build uvicorn reload kwargs so --reload also reacts to YAML edits."""
options: dict = {"reload": True}
if not config_path:
return options
config_abs = os.path.abspath(config_path)
config_dir = os.path.dirname(config_abs)
cwd = os.path.abspath(os.getcwd())
reload_dirs = [cwd]
if config_dir and config_dir != cwd:
reload_dirs.append(config_dir)
options["reload_dirs"] = reload_dirs
# Must be a basename, not an absolute path: uvicorn's
# resolve_reload_patterns() calls pathlib.Path.glob(), which raises
# NotImplementedError on absolute patterns (uvicorn discussion #2156).
options["reload_includes"] = ["*.py", os.path.basename(config_abs)]
return options
@staticmethod
def _patch_statreload_for_config(config_path: str) -> bool:
"""Make uvicorn's StatReload reloader notice YAML config changes.
Uvicorn uses WatchFilesReload when the optional `watchfiles` package
is installed, otherwise StatReload. StatReload hard-codes `*.py` in
`iter_py_files()` and silently ignores `reload_includes`, so the
kwargs from `_get_reload_options` alone don't trigger reloads on YAML
edits. We monkey-patch `iter_py_files` to also yield the config path.
Idempotent across calls and a no-op for the WatchFilesReload path.
"""
try:
from uvicorn.supervisors.statreload import StatReload
except ImportError: # pragma: no cover - uvicorn is a hard dep
return False
if not config_path:
return False
from pathlib import Path
config_abs = Path(config_path).resolve()
patched_paths = getattr(StatReload, "_litellm_patched_config_paths", None)
if patched_paths is None:
original_iter = StatReload.iter_py_files
patched_paths = set()
def _iter_with_config(self): # type: ignore[no-untyped-def]
yield from original_iter(self)
for path in StatReload._litellm_patched_config_paths:
if path.exists():
yield path
StatReload.iter_py_files = _iter_with_config # type: ignore[assignment]
StatReload._litellm_patched_config_paths = patched_paths # type: ignore[attr-defined]
patched_paths.add(config_abs)
return True
@staticmethod
def _init_hypercorn_server(
app: FastAPI,
@ -619,7 +679,7 @@ class ProxyInitializationHelpers:
"--reload",
is_flag=True,
default=False,
help="Enable uvicorn hot reload (dev only). Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
)
def run_server( # noqa: PLR0915
host,
@ -990,11 +1050,6 @@ def run_server( # noqa: PLR0915
litellm_settings=litellm_settings if config else None, # type: ignore[possibly-unbound]
)
# --- SEPARATE HEALTH APP LOGIC ---
# To run the health app separately, use:
# uvicorn litellm.proxy.health_app_factory:build_health_app --factory --host 0.0.0.0 --port=4001
# This is compatible with the SEPARATE_HEALTH_APP Docker/supervisord pattern.
# --- END SEPARATE HEALTH APP LOGIC ---
# Skip server startup if requested (after all setup is done)
if skip_server_startup:
print( # noqa
@ -1028,7 +1083,11 @@ def run_server( # noqa: PLR0915
uvicorn_args["loop"] = loop_type
if reload:
uvicorn_args["reload"] = True
uvicorn_args.update(
ProxyInitializationHelpers._get_reload_options(config)
)
if config:
ProxyInitializationHelpers._patch_statreload_for_config(config)
uvicorn.run(
**uvicorn_args,

View file

@ -403,6 +403,9 @@ from litellm.proxy.middleware.in_flight_requests_middleware import (
InFlightRequestsMiddleware,
)
from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware
from litellm.proxy.middleware.request_size_limit_middleware import (
RequestSizeLimitMiddleware,
)
from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
@ -3820,6 +3823,11 @@ class ProxyConfig:
f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}"
) # noqa
elif key == "audit_log_callbacks":
from litellm.proxy.management_helpers.audit_logs import (
reset_audit_log_callback_cache,
)
reset_audit_log_callback_cache()
litellm.audit_log_callbacks = []
for callback in value:
@ -3901,6 +3909,21 @@ class ProxyConfig:
f"{blue_color_code} setting litellm.{key}={value}{reset_color_code}"
)
setattr(litellm, key, value)
if key in {"s3_audit_callback_params", "s3_callback_params"}:
from litellm.proxy.management_helpers.audit_logs import (
reset_audit_log_callback_cache,
)
from litellm.litellm_core_utils.litellm_logging import (
_in_memory_loggers,
)
from litellm.integrations.s3_v2 import S3Logger as S3V2Logger
reset_audit_log_callback_cache()
_in_memory_loggers[:] = [
cb
for cb in _in_memory_loggers
if not isinstance(cb, S3V2Logger)
]
## GENERAL SERVER SETTINGS (e.g. master key,..) # do this after initializing litellm, to ensure sentry logging works for proxylogging
general_settings = config.get("general_settings", {})
@ -14881,6 +14904,11 @@ app.include_router(ui_discovery_endpoints_router)
app.include_router(google_router)
attach_lazy_features(app)
app.add_middleware(
RequestSizeLimitMiddleware,
get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"),
is_request_size_limit_enabled=lambda: premium_user is True,
)
async def _stream_mcp_asgi_response(

View file

@ -2274,6 +2274,13 @@ class ProxyLogging:
Covers:
1. /chat/completions
"""
from litellm.proxy.proxy_server import llm_router
# Merge model-level guardrails before checking which guardrails to run
request_data = _check_and_merge_model_level_guardrails(
data=request_data, llm_router=llm_router
)
current_response = response
for callback in litellm.callbacks:

View file

@ -1878,14 +1878,24 @@ class OpenAIRealtimeStreamSessionEvents(TypedDict):
class OpenAIRealtimeStreamResponseOutputItemContent(TypedDict, total=False):
audio: str
"""Base64-encoded audio bytes, used for 'input_audio' content types"""
"""Base64-encoded audio bytes, used for 'input_audio' / 'audio' / 'output_audio' content types"""
id: str
"""The ID of the previous conversation item for reference"""
text: str
"""The text content, used for 'input_text' and 'text' content types"""
"""The text content, used for 'input_text' / 'text' / 'output_text' content types"""
transcript: str
"""The transcript content, used for 'input_audio' content types"""
type: Literal["input_audio", "input_text", "text", "item_reference", "audio"]
"""The transcript content, used for 'input_audio' / 'audio' content types"""
type: Literal[
"input_audio",
"input_text",
# Beta assistant content types
"text",
"audio",
"item_reference",
# GA assistant content types (aligns with Responses API)
"output_text",
"output_audio",
]
"""The type of content"""
@ -1945,23 +1955,46 @@ class OpenAIRealtimeConversationCreated(TypedDict, total=False):
class OpenAIRealtimeConversationItemCreated(TypedDict, total=False):
"""Beta: single event emitted when a conversation item is created."""
type: Required[Literal["conversation.item.created"]]
item: OpenAIRealtimeStreamResponseOutputItem
event_id: str
previous_item_id: str
previous_item_id: Optional[str] # None when this is the first item
class OpenAIRealtimeConversationItemAdded(TypedDict, total=False):
"""GA: emitted immediately when a conversation item is added (replaces .created)."""
type: Required[Literal["conversation.item.added"]]
item: OpenAIRealtimeStreamResponseOutputItem
event_id: str
previous_item_id: Optional[str] # None when this is the first item
class OpenAIRealtimeConversationItemDone(TypedDict, total=False):
"""GA: emitted when a conversation item is fully complete (e.g. transcription done)."""
type: Required[Literal["conversation.item.done"]]
item: OpenAIRealtimeStreamResponseOutputItem
event_id: str
previous_item_id: Optional[str] # None when this is the first item
class OpenAIRealtimeResponseContentPart(TypedDict, total=False):
audio: str
"""Base64-encoded audio bytes, if type is 'audio'"""
"""Base64-encoded audio bytes, if type is 'audio' or 'output_audio'"""
text: str
"""The text content, if type is 'text'"""
"""The text content, if type is 'text' or 'output_text'"""
transcript: str
"""The transcript content, if type is 'audio'"""
"""The transcript content, if type is 'audio' or 'output_audio'"""
type: Literal["audio", "text"]
type: Union[
Literal["audio", "text"], # beta
Literal["output_audio", "output_text"], # GA
]
"""The type of content"""
@ -1982,7 +2015,14 @@ class OpenAIRealtimeResponseDelta(TypedDict):
item_id: str
output_index: int
response_id: str
type: Union[Literal["response.text.delta"], Literal["response.audio.delta"]]
type: Union[
Literal["response.text.delta"],
Literal["response.audio.delta"],
# GA renamed events
Literal["response.output_text.delta"],
Literal["response.output_audio.delta"],
Literal["response.output_audio_transcript.delta"],
]
class OpenAIRealtimeResponseTextDone(TypedDict):
@ -1992,7 +2032,10 @@ class OpenAIRealtimeResponseTextDone(TypedDict):
output_index: int
response_id: str
text: str
type: Literal["response.text.done"]
type: Union[
Literal["response.text.done"],
Literal["response.output_text.done"], # GA rename
]
class OpenAIRealtimeResponseAudioDone(TypedDict):
@ -2001,7 +2044,11 @@ class OpenAIRealtimeResponseAudioDone(TypedDict):
item_id: str
output_index: int
response_id: str
type: Literal["response.audio.done"]
type: Union[
Literal["response.audio.done"],
Literal["response.output_audio.done"], # GA rename
Literal["response.output_audio_transcript.done"], # GA rename
]
class OpenAIRealtimeContentPartDone(TypedDict):
@ -2046,10 +2093,18 @@ class OpenAIRealtimeDoneEvent(TypedDict):
class OpenAIRealtimeEventTypes(Enum):
SESSION_CREATED = "session.created"
# Beta delta event names
RESPONSE_TEXT_DELTA = "response.text.delta"
RESPONSE_AUDIO_DELTA = "response.audio.delta"
RESPONSE_TEXT_DONE = "response.text.done"
RESPONSE_AUDIO_DONE = "response.audio.done"
# GA renamed delta event names
RESPONSE_OUTPUT_TEXT_DELTA = "response.output_text.delta"
RESPONSE_OUTPUT_AUDIO_DELTA = "response.output_audio.delta"
RESPONSE_OUTPUT_AUDIO_TRANSCRIPT_DELTA = "response.output_audio_transcript.delta"
RESPONSE_OUTPUT_TEXT_DONE = "response.output_text.done"
RESPONSE_OUTPUT_AUDIO_DONE = "response.output_audio.done"
RESPONSE_OUTPUT_AUDIO_TRANSCRIPT_DONE = "response.output_audio_transcript.done"
RESPONSE_DONE = "response.done"
RESPONSE_OUTPUT_ITEM_ADDED = "response.output_item.added"
RESPONSE_CONTENT_PART_ADDED = "response.content_part.added"
@ -2060,7 +2115,11 @@ OpenAIRealtimeEvents = Union[
OpenAIRealtimeStreamSessionEvents,
OpenAIRealtimeStreamResponseOutputItemAdded,
OpenAIRealtimeResponseContentPartAdded,
# Beta conversation item event
OpenAIRealtimeConversationItemCreated,
# GA conversation item events
OpenAIRealtimeConversationItemAdded,
OpenAIRealtimeConversationItemDone,
OpenAIRealtimeConversationCreated,
OpenAIRealtimeResponseDelta,
OpenAIRealtimeResponseTextDone,

View file

@ -37,6 +37,7 @@ class MCPAuth(str, enum.Enum):
oauth2 = "oauth2"
aws_sigv4 = "aws_sigv4"
token = "token"
oauth2_token_exchange = "oauth2_token_exchange"
# MCP Literals
@ -54,6 +55,7 @@ MCPAuthType = Optional[
MCPAuth.oauth2,
MCPAuth.aws_sigv4,
MCPAuth.token,
MCPAuth.oauth2_token_exchange,
]
]
@ -117,6 +119,22 @@ class MCPCredentials(TypedDict, total=False):
aws_session_name: Optional[str]
"""Session name for STS AssumeRole (used in CloudTrail). Not a secret — stored unencrypted."""
audience: Optional[str]
"""
Target audience for OAuth 2.0 Token Exchange (RFC 8693)
"""
token_exchange_endpoint: Optional[str]
"""
IDP token endpoint for OAuth 2.0 Token Exchange (RFC 8693)
"""
subject_token_type: Optional[str]
"""
Subject token type for OAuth 2.0 Token Exchange (RFC 8693).
Default: urn:ietf:params:oauth:token-type:access_token
"""
class MCPServerCostInfo(TypedDict, total=False):
default_cost_per_query: Optional[float]

View file

@ -57,6 +57,10 @@ class MCPServer(BaseModel):
aws_service_name: Optional[str] = None # defaults to "bedrock-agentcore"
aws_role_name: Optional[str] = None # IAM role ARN for STS AssumeRole
aws_session_name: Optional[str] = None # session name for CloudTrail auditing
# Token Exchange (OBO) fields — RFC 8693
token_exchange_endpoint: Optional[str] = None
audience: Optional[str] = None
subject_token_type: str = "urn:ietf:params:oauth:token-type:access_token"
# Stdio-specific fields
command: Optional[str] = None
args: Optional[List[str]] = None
@ -127,3 +131,12 @@ class MCPServer(BaseModel):
return any(h.lower() in auth_header_names for h in self.extra_headers)
return False
@property
def has_token_exchange_config(self) -> bool:
"""True if this server is configured for OAuth2 token exchange (OBO / RFC 8693)."""
return (
self.auth_type == MCPAuth.oauth2_token_exchange
and bool(self.client_id and self.client_secret)
and bool(self.token_exchange_endpoint or self.token_url)
)

View file

@ -3247,6 +3247,7 @@ class LlmProviders(str, Enum):
A2A = "a2a"
GIGACHAT = "gigachat"
NVIDIA_NIM = "nvidia_nim"
NVIDIA_RIVA = "nvidia_riva"
CEREBRAS = "cerebras"
AI21_CHAT = "ai21_chat"
VOLCENGINE = "volcengine"

View file

@ -8545,6 +8545,12 @@ class ProviderConfigManager:
)
return MistralAudioTranscriptionConfig()
elif litellm.LlmProviders.NVIDIA_RIVA == provider:
from litellm.llms.nvidia_riva.audio_transcription.transformation import (
NvidiaRivaAudioTranscriptionConfig,
)
return NvidiaRivaAudioTranscriptionConfig()
return None
@staticmethod

View file

@ -28879,6 +28879,19 @@
"mode": "chat",
"output_cost_per_token": 0.0
},
"sambanova/MiniMax-M2.7": {
"input_cost_per_token": 3e-07,
"litellm_provider": "sambanova",
"max_input_tokens": 204800,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://cloud.sambanova.ai/plans/pricing",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"sambanova/DeepSeek-R1": {
"input_cost_per_token": 5e-06,
"litellm_provider": "sambanova",

View file

@ -1610,6 +1610,22 @@
"interactions": true
}
},
"nvidia_riva": {
"display_name": "Nvidia Riva (`nvidia_riva`)",
"url": "https://docs.litellm.ai/docs/providers/nvidia_riva",
"endpoints": {
"chat_completions": false,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": true,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false
}
},
"oci": {
"display_name": "OCI (`oci`)",
"url": "https://docs.litellm.ai/docs/providers/oci",

View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.84.0"
version = "1.85.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.14"
@ -10,18 +10,22 @@ authors = [
{ name = "BerriAI" },
]
dependencies = [
"fastuuid==0.14.0",
"httpx==0.28.1",
"openai==2.33.0",
"python-dotenv==1.2.2",
"tiktoken==0.12.0",
"importlib-metadata==8.5.0",
"tokenizers==0.23.1",
"click==8.1.8",
"jinja2==3.1.6",
"aiohttp==3.13.4",
"pydantic==2.12.5",
"jsonschema==4.23.0",
# Ranges (not exact pins) so SDK consumers can coexist with their other
# deps. Reproducibility for our Docker/CI comes from `uv.lock`.
# When changing a floor, verify it installs + imports on every supported
# Python with: `uv pip install --resolution=lowest-direct .`
"fastuuid>=0.14.0,<1.0",
"httpx>=0.28.0,<1.0",
"openai>=2.20.0,<3.0.0",
"python-dotenv>=1.0.0,<2.0",
"tiktoken>=0.8.0,<1.0",
"importlib-metadata>=8.0.0,<9.0",
"tokenizers>=0.21.0,<1.0",
"click>=8.0.0,<9.0",
"jinja2>=3.1.0,<4.0",
"aiohttp>=3.10,<4.0",
"pydantic>=2.10.0,<3.0.0",
"jsonschema>=4.0.0,<5.0",
]
[project.urls]
@ -86,6 +90,14 @@ grpc = [
# Newest non-yanked release older than the 30-day cutoff.
"grpcio==1.78.0",
]
stt-nvidia-riva = [
# NVIDIA Riva STT provider (gRPC). These are imported lazily inside the
# provider handler so litellm core remains usable without them.
"nvidia-riva-client>=2.15.0",
"soundfile>=0.12.1",
"audioread>=3.0.1",
"numpy>=1.26.0",
]
google = ["google-cloud-aiplatform==1.133.0"]
proxy-runtime = [
# Historically bundled in the proxy Docker images via requirements.txt.
@ -238,7 +250,7 @@ source-exclude = [
profile = "black"
[tool.commitizen]
version = "1.84.0"
version = "1.85.0"
version_files = [
"pyproject.toml:^version",
]

View file

@ -0,0 +1,276 @@
#!/usr/bin/env python3
"""
Minimal HTTP target for testing LiteLLM **Bedrock pass-through** (`/bedrock/...` on the proxy).
What it does
- Serves a tiny Converse-shaped JSON (and optional invoke-shaped) response so the proxy can
complete a round trip without calling AWS.
- Does **not** verify SigV4 (Bedrock does); any Authorization header is accepted.
How to run
uv run python scripts/mock_bedrock_passthrough_target.py --host 127.0.0.1 --port 9999
Wire LiteLLM to this host (use **one** of these patterns):
1) model_list (recommended) set the Bedrock runtime base to the mock:
model_list:
- model_name: mock-bedrock-claude
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
custom_llm_provider: bedrock
aws_region_name: us-west-2
api_base: "http://127.0.0.1:9999"
2) Environment (see litellm BaseAWSLLM.get_runtime_endpoint)::
export AWS_BEDROCK_RUNTIME_ENDPOINT="http://127.0.0.1:9999"
Then call the proxy, e.g. (model_name must match config)::
curl -sS -X POST "http://127.0.0.1:4000/bedrock/model/mock-bedrock-claude/converse" \
-H "Authorization: Bearer $LITELLM_KEY" -H "Content-Type: application/json" \
-d '{"messages":[{"role":"user","content":[{"text":"hi"}]}]}'
The proxy will forward to: {api_base}/model/<resolved model id>/converse (SigV4-signed).
This mock implements POST .../converse and returns a minimal valid Converse response.
Notes
- `invoke-with-response-stream` returns a real **binary** AWS event stream
(`application/vnd.amazon.eventstream`) with Anthropic-style JSON payloads inside each
`PayloadPart`, matching Bedrock's InvokeModelWithResponseStream wire format. See
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModelWithResponseStream.html
and https://docs.aws.amazon.com/awstreams/latest/devguide/message-formats.html
- `converse-stream` is still JSON-only placeholder (different inner event shapes).
- Use real (or any non-empty) AWS creds in the environment of the **proxy**; signing still runs.
"""
from __future__ import annotations
import argparse
import base64
import json
from binascii import crc32
from struct import pack
from typing import Any, Dict, Iterator, List
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from starlette.responses import StreamingResponse
app = FastAPI(title="Mock Bedrock runtime (pass-through test target)")
# Minimal structure compatible with Converse: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Converse.html
def _converse_response_body() -> Dict[str, Any]:
return {
"output": {
"message": {
"role": "assistant",
"content": [
{"text": "mock: ok from mock_bedrock_passthrough_target.py"}
],
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": 1,
"outputTokens": 2,
"totalTokens": 3,
},
}
# Minimal invoke (Anthropic messages on bedrock) style — adjust if you test /invoke
def _invoke_response_body() -> Dict[str, Any]:
return {
"id": "msg_mock",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "mock invoke response"}],
"model": "mock",
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 2},
}
def _encode_event_stream_message(headers: Dict[str, str], payload: bytes) -> bytes:
"""Single AWS binary event-stream frame (same layout botocore's ``EventStreamBuffer`` parses)."""
header_blob = b""
for name, value in headers.items():
nb = name.encode("utf-8")
vb = value.encode("utf-8")
header_blob += bytes([len(nb)]) + nb + bytes([7]) + pack("!H", len(vb)) + vb
headers_length = len(header_blob)
payload_length = len(payload)
total_length = 12 + headers_length + payload_length + 4
prelude_wo_crc = pack("!II", total_length, headers_length)
prelude_crc_val = crc32(prelude_wo_crc) & 0xFFFFFFFF
prelude = prelude_wo_crc + pack("!I", prelude_crc_val)
wo_msg_crc = prelude + header_blob + payload
msg_crc_val = crc32(wo_msg_crc[8:], prelude_crc_val) & 0xFFFFFFFF
return wo_msg_crc + pack("!I", msg_crc_val)
def _bedrock_payload_part(inner_event: Dict[str, Any]) -> bytes:
"""Outer JSON expected by bedrock-runtime ``ResponseStream`` / ``PayloadPart``."""
inner_bytes = json.dumps(inner_event, separators=(",", ":")).encode("utf-8")
outer = {
"chunk": {
"bytes": base64.b64encode(inner_bytes).decode("ascii"),
}
}
return json.dumps(outer, separators=(",", ":")).encode("utf-8")
def _anthropic_invoke_stream_events(
model_id: str, assistant_text: str
) -> List[Dict[str, Any]]:
"""
Minimal Anthropic Messages stream events as returned inside Bedrock stream chunks.
Mirrors the sequence Amazon emits for Claude on ``invoke-with-response-stream``.
"""
msg_id = "msg_mock_bedrock_stream"
input_tokens = 3
output_tokens = max(1, len(assistant_text) // 4)
events: List[Dict[str, Any]] = [
{
"type": "message_start",
"message": {
"model": model_id,
"id": msg_id,
"type": "message",
"role": "assistant",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {
"input_tokens": input_tokens,
"output_tokens": 1,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"cache_creation": {
"ephemeral_5m_input_tokens": 0,
"ephemeral_1h_input_tokens": 0,
},
},
},
},
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
]
# Split text into small deltas so downstream streaming behavior is visible.
step = 24
for i in range(0, len(assistant_text), step):
events.append(
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "text_delta",
"text": assistant_text[i : i + step],
},
}
)
events.append({"type": "content_block_stop", "index": 0})
events.append(
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
},
}
)
events.append(
{
"type": "message_stop",
"amazon-bedrock-invocationMetrics": {
"inputTokenCount": input_tokens,
"outputTokenCount": output_tokens,
"invocationLatency": 42,
"firstByteLatency": 10,
},
}
)
return events
def _iter_invoke_with_response_stream(model_id: str) -> Iterator[bytes]:
text = (
"mock streaming: ok from scripts/mock_bedrock_passthrough_target.py "
"(invoke-with-response-stream)."
)
headers = {
":event-type": "chunk",
":content-type": "application/json",
":message-type": "event",
}
for ev in _anthropic_invoke_stream_events(model_id, text):
yield _encode_event_stream_message(headers, _bedrock_payload_part(ev))
@app.get("/health")
def health() -> Dict[str, str]:
return {"status": "ok"}
@app.post("/model/{model_path:path}/converse")
async def converse(model_path: str, request: Request) -> JSONResponse:
# Optional: log body for debugging
_ = await request.body()
return JSONResponse(content=_converse_response_body())
@app.post("/model/{model_path:path}/converse-stream")
async def converse_stream(model_path: str, request: Request) -> JSONResponse:
"""
Not a real AWS event stream returns JSON for quick smoke tests only.
"""
_ = await request.body()
return JSONResponse(
content={
"note": "This mock does not implement application/vnd.amazon.eventstream; use /converse for basic tests."
}
)
@app.post("/model/{model_path:path}/invoke")
async def invoke(model_path: str, request: Request) -> JSONResponse:
_ = await request.body()
return JSONResponse(content=_invoke_response_body())
@app.post("/model/{model_path:path}/invoke-with-response-stream")
async def invoke_with_response_stream(
model_path: str, request: Request
) -> StreamingResponse:
"""
Binary ``application/vnd.amazon.eventstream`` body compatible with boto3/botocore
``InvokeModelWithResponseStream`` / LiteLLM's Bedrock invoke streaming path.
"""
_ = await request.body()
return StreamingResponse(
_iter_invoke_with_response_stream(model_id=model_path),
media_type="application/vnd.amazon.eventstream",
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=9999)
args = parser.parse_args()
import uvicorn
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
if __name__ == "__main__":
main()

82
scripts/tpm_headline_test.sh Executable file
View file

@ -0,0 +1,82 @@
#!/usr/bin/env bash
# Concurrent TPM bypass test — mints a virtual key with tpm_limit=100
# (api_key scope in the v3 rate-limiter), races 10 concurrent calls,
# prints a verdict, then deletes the key.
#
# Note: the `tpm: 100` on a model_list deployment is the *router's*
# load-balancing TPM, not a v3 rate-limit descriptor. The v3 limiter
# enforces against limits set on the key/team/user — so we set
# tpm_limit=100 on the key itself.
#
# Pre-PR: ~all 10 return 200 (race lets concurrent requests bypass the limit).
# Post-PR: only ~12 fit under tpm_limit=100, rest return 429.
#
# Setup (separate terminal):
# kubectl port-forward -n litellm svc/yassin-veks-litellm-helm 4000:4000
#
# Run:
# bash scripts/tpm_headline_test.sh
set -u
PROXY="${PROXY:-http://localhost:4000}"
MASTER_KEY="${MASTER_KEY:-sk-perf-test-fixed-do-not-rotate}"
MODEL="${MODEL:-opus-4.6}"
echo "=== Concurrent TPM bypass test ==="
echo "proxy=$PROXY model=$MODEL key tpm_limit=100 concurrency=10 max_tokens=50"
echo
gen_resp=$(curl -s -X POST "$PROXY/key/generate" \
-H "Authorization: Bearer $MASTER_KEY" \
-H "Content-Type: application/json" \
-d "{\"models\":[\"$MODEL\"],\"tpm_limit\":100,\"duration\":\"10m\",\"key_alias\":\"tpm-headline-$$-$(date +%s)\"}")
KEY=$(printf '%s' "$gen_resp" | sed -n 's/.*"key"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p')
if [ -z "$KEY" ]; then
echo "FAIL — could not mint virtual key. Response: $gen_resp"
exit 1
fi
echo "Minted virtual key: ${KEY:0:12}"
echo
cleanup() {
curl -s -X POST "$PROXY/key/delete" \
-H "Authorization: Bearer $MASTER_KEY" \
-H "Content-Type: application/json" \
-d "{\"keys\":[\"$KEY\"]}" > /dev/null 2>&1 || true
[ -n "${tmp:-}" ] && rm -rf "$tmp"
}
trap cleanup EXIT
tmp=$(mktemp -d)
for i in $(seq 1 10); do
( curl -s -o "$tmp/body.$i" -w "%{http_code}" \
"$PROXY/v1/chat/completions" \
-H "Authorization: Bearer $KEY" \
-H "Content-Type: application/json" \
-d "{\"model\":\"$MODEL\",\"messages\":[{\"role\":\"user\",\"content\":\"concurrent tpm test $i\"}],\"max_tokens\":50}" \
> "$tmp/code.$i" ) &
done
wait
ok=0; limited=0; other=0
for i in $(seq 1 10); do
code=$(cat "$tmp/code.$i")
case "$code" in
200) ok=$((ok+1)) ;;
429) limited=$((limited+1)) ;;
*) other=$((other+1)); echo "req $i -> $code: $(cat "$tmp/body.$i" | head -c 200)" ;;
esac
done
echo
echo "Results: 200=$ok 429=$limited other=$other"
if [ "$limited" -ge 1 ] && [ "$ok" -ge 1 ]; then
echo "PASS — reservation enforced under concurrency."
exit 0
elif [ "$ok" -eq 10 ]; then
echo "FAIL — all 10 succeeded; concurrent bypass still possible."
exit 1
else
echo "INCONCLUSIVE — investigate non-200/429 above."
exit 2
fi

View file

@ -7,7 +7,9 @@ from __future__ import annotations
import atexit
import hashlib
import json
import os
import re
import sys
from typing import Iterable
@ -74,6 +76,17 @@ FILTERED_RESPONSE_HEADERS = (
"date",
)
# Tiny placeholder used to replace base64 image payloads in cassettes.
# Decodes to b"test" — short, valid base64 so test code that decodes
# the field still succeeds.
VCR_IMAGE_B64_PLACEHOLDER = "dGVzdA=="
# Fixed boundary substituted into multipart request bodies so the
# ``safe_body`` matcher sees the same bytes across record and replay.
# httpx generates a fresh random boundary per request via os.urandom,
# which otherwise turns every multipart cassette into a permanent miss.
VCR_FIXED_MULTIPART_BOUNDARY = "vcr-static-boundary"
def _scrub_response(response):
if not isinstance(response, dict):
@ -86,8 +99,88 @@ def _scrub_response(response):
return response
def _replace_b64_json_in_place(obj) -> bool:
"""Recursively replace ``b64_json`` string values in a JSON tree.
Returns ``True`` if any value was rewritten. The check on the
existing value's length keeps the function idempotent — once a
value has been swapped to the placeholder, subsequent invocations
are no-ops.
"""
changed = False
if isinstance(obj, dict):
for key, value in obj.items():
if (
key == "b64_json"
and isinstance(value, str)
and len(value) > len(VCR_IMAGE_B64_PLACEHOLDER)
):
obj[key] = VCR_IMAGE_B64_PLACEHOLDER
changed = True
elif _replace_b64_json_in_place(value):
changed = True
elif isinstance(obj, list):
for item in obj:
if _replace_b64_json_in_place(item):
changed = True
return changed
def _strip_image_b64_payloads(response):
"""Replace ``b64_json`` payloads in image-gen responses before save.
Image-edit and image-generation responses carry the full base64
PNG/JPEG (1-10+ MB) in ``data[*].b64_json``. The image_gen tests
only assert response shape the field decodes, schema validates
they never inspect pixel content. Swapping to a 4-byte placeholder
preserves all those checks while shrinking cassettes by ~99%.
"""
if not isinstance(response, dict):
return response
body = response.get("body")
if not isinstance(body, dict):
return response
raw = body.get("string")
if raw is None:
return response
if isinstance(raw, (bytes, bytearray)):
try:
text = bytes(raw).decode("utf-8")
except UnicodeDecodeError:
return response
was_bytes = True
elif isinstance(raw, str):
text = raw
was_bytes = False
else:
return response
try:
payload = json.loads(text)
except (ValueError, TypeError):
return response
if not _replace_b64_json_in_place(payload):
return response
new_text = json.dumps(payload, separators=(",", ":"))
body["string"] = new_text.encode("utf-8") if was_bytes else new_text
headers = response.get("headers")
if isinstance(headers, dict):
new_len_value = str(len(new_text.encode("utf-8")))
for key in list(headers):
if str(key).lower() == "content-length":
value = headers[key]
headers[key] = (
[new_len_value] if isinstance(value, list) else new_len_value
)
return response
def _before_record_response(response):
return filter_non_2xx_response(_scrub_response(response))
return filter_non_2xx_response(_scrub_response(_strip_image_b64_payloads(response)))
def _safe_body_matcher(r1, r2) -> None:
@ -172,8 +265,84 @@ def _strip_headers(headers, names: Iterable[str]) -> None:
pass
def _normalize_multipart_boundary(request) -> None:
"""Rewrite random multipart boundaries to a fixed string in-place.
httpx generates a fresh ``boundary=<random hex>`` for every
multipart request via ``os.urandom``. Without normalization, the
request body bytes differ across runs even when everything else is
identical, the ``safe_body`` matcher misses, and the persister
keeps appending new episodes until ``MAX_EPISODES_PER_CASSETTE``
refuses the save leaving audio-transcription tests effectively
unmocked. Replacing the boundary in both the Content-Type header
and the body bytes makes the request deterministic.
Idempotent vcrpy invokes this hook multiple times per request,
so the second invocation sees ``boundary=vcr-static-boundary``
already and short-circuits.
"""
headers = getattr(request, "headers", None)
if headers is None:
return
content_type_key = None
content_type_value = None
try:
for key in list(headers.keys()):
if str(key).lower() == "content-type":
content_type_key = key
value = headers[key]
content_type_value = value if isinstance(value, str) else str(value)
break
except AttributeError:
return
if not content_type_value or "multipart/" not in content_type_value.lower():
return
fixed_param = f"boundary={VCR_FIXED_MULTIPART_BOUNDARY}"
if fixed_param in content_type_value:
return
match = re.search(r"boundary=([^\s;]+)", content_type_value)
if not match:
return
current_boundary = match.group(1).strip('"')
if current_boundary == VCR_FIXED_MULTIPART_BOUNDARY:
return
try:
headers[content_type_key] = content_type_value.replace(
match.group(0), fixed_param
)
except (TypeError, AttributeError):
return
body = getattr(request, "body", None)
if body is None:
return
if isinstance(body, (bytes, bytearray)):
try:
new_body = bytes(body).replace(
current_boundary.encode("utf-8"),
VCR_FIXED_MULTIPART_BOUNDARY.encode("utf-8"),
)
except (TypeError, ValueError):
return
elif isinstance(body, str):
new_body = body.replace(current_boundary, VCR_FIXED_MULTIPART_BOUNDARY)
else:
return
try:
request.body = new_body
except (AttributeError, TypeError):
pass
def _before_record_request(request):
"""Fingerprint API keys, then scrub them.
"""Fingerprint API keys, scrub them, and normalize multipart boundaries.
Order matters in two ways:
@ -187,7 +356,8 @@ def _before_record_request(request):
auth headers we already stripped, so re-hashing would yield
``"no-key"`` and the stored vs. incoming fingerprints would
diverge. Skip the recompute when the header is already set so
this hook is idempotent.
this hook is idempotent. The boundary normalizer is also
idempotent for the same reason.
"""
headers = getattr(request, "headers", None)
if headers is None:
@ -199,6 +369,7 @@ def _before_record_request(request):
except (TypeError, AttributeError):
pass
_strip_headers(headers, FILTERED_REQUEST_HEADERS)
_normalize_multipart_boundary(request)
return request

View file

@ -26,24 +26,7 @@ import pytest
import litellm
@pytest.mark.parametrize(
"sync_mode",
[True, False],
)
@pytest.mark.parametrize(
"model, api_key, api_base",
[
(
"azure/tts",
os.getenv("AZURE_TTS_API_KEY"),
os.getenv("AZURE_TTS_API_BASE"),
),
("openai/tts-1", os.getenv("OPENAI_API_KEY"), None),
],
) # ,
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_audio_speech_litellm(sync_mode, model, api_base, api_key):
async def _run_audio_speech_litellm(sync_mode, model, api_base, api_key):
litellm._turn_on_debug()
speech_file_path = Path(__file__).parent / "speech.mp3"
@ -85,6 +68,30 @@ async def test_audio_speech_litellm(sync_mode, model, api_base, api_key):
assert isinstance(response, HttpxBinaryResponseContent)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_audio_speech_litellm_azure(sync_mode):
await _run_audio_speech_litellm(
sync_mode=sync_mode,
model="azure/tts",
api_base=os.getenv("AZURE_TTS_API_BASE"),
api_key=os.getenv("AZURE_TTS_API_KEY"),
)
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_audio_speech_litellm_openai(sync_mode):
await _run_audio_speech_litellm(
sync_mode=sync_mode,
model="openai/tts-1",
api_base=None,
api_key=os.getenv("OPENAI_API_KEY"),
)
@pytest.mark.parametrize(
"sync_mode",
[False, True],

View file

@ -39,24 +39,7 @@ import litellm
from litellm import Router
@pytest.mark.parametrize(
"model, api_key, api_base",
[
("whisper-1", None, None),
(
"azure/whisper",
os.getenv("AZURE_WHISPER_API_KEY"),
os.getenv("AZURE_WHISPER_API_BASE"),
),
],
)
@pytest.mark.parametrize(
"response_format, timestamp_granularities",
[("json", None), ("vtt", None), ("verbose_json", ["word"])],
)
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_transcription(
async def _run_transcription(
model, api_key, api_base, response_format, timestamp_granularities
):
transcript = await litellm.atranscription(
@ -74,6 +57,38 @@ async def test_transcription(
assert transcript.text is not None
@pytest.mark.parametrize(
"response_format, timestamp_granularities",
[("json", None), ("vtt", None), ("verbose_json", ["word"])],
)
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_transcription_openai_whisper(response_format, timestamp_granularities):
await _run_transcription(
model="whisper-1",
api_key=None,
api_base=None,
response_format=response_format,
timestamp_granularities=timestamp_granularities,
)
@pytest.mark.parametrize(
"response_format, timestamp_granularities",
[("json", None), ("vtt", None), ("verbose_json", ["word"])],
)
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_transcription_azure_whisper(response_format, timestamp_granularities):
await _run_transcription(
model="azure/whisper",
api_key=os.getenv("AZURE_WHISPER_API_KEY"),
api_base=os.getenv("AZURE_WHISPER_API_BASE"),
response_format=response_format,
timestamp_granularities=timestamp_granularities,
)
@pytest.mark.asyncio()
async def test_transcription_caching():
import litellm

View file

@ -1,7 +1,6 @@
# conftest.py
import asyncio
import importlib
import os
import sys
@ -12,16 +11,6 @@ sys.path.insert(
) # Adds the parent directory to the system path
import litellm # noqa: E402,F401
from tests._vcr_conftest_common import ( # noqa: E402
VerboseReporterState,
apply_vcr_auto_marker_to_items,
record_vcr_outcome,
register_persister_if_enabled,
vcr_config_dict,
)
_verbose_state = VerboseReporterState()
@pytest.fixture(scope="session")
def event_loop():
@ -31,37 +20,3 @@ def event_loop():
loop = asyncio.new_event_loop()
yield loop
loop.close()
@pytest.fixture(scope="module")
def vcr_config():
return vcr_config_dict()
def pytest_recording_configure(config, vcr):
register_persister_if_enabled(vcr)
@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_makereport(item, call):
outcome = yield
rep = outcome.get_result()
setattr(item, f"rep_{rep.when}", rep)
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
yield
record_vcr_outcome(request, vcr)
def pytest_configure(config):
_verbose_state.remember_pluginmanager(config)
def pytest_runtest_logreport(report):
_verbose_state.maybe_emit_verdict(report)
def pytest_collection_modifyitems(config, items):
apply_vcr_auto_marker_to_items(items)

View file

@ -47,6 +47,37 @@ class TestCustomLogger(CustomLogger):
self.standard_logging_object = kwargs["standard_logging_object"]
async def _acreate_fine_tuning_job_with_propagation_retry(
*, max_attempts: int = 12, initial_delay: float = 1.0, **kwargs
):
"""
Wrap litellm.acreate_fine_tuning_job and retry on the eventual-consistency
400 OpenAI returns when a freshly-uploaded training file isn't yet visible
to the fine-tuning endpoint (`'file-... does not exist'`).
Polling the files-retrieve endpoint or `FileObject.status` doesn't help —
OpenAI's `status` field is deprecated, and the retrieve and fine-tuning
endpoints don't share a consistency model. Retrying the operation itself
is the only reliable signal that propagation has finished.
Total budget with defaults: ~70s across 12 attempts (exp backoff capped at
8s).
"""
delay = initial_delay
last_error: Optional[openai.BadRequestError] = None
for _ in range(max_attempts):
try:
return await litellm.acreate_fine_tuning_job(**kwargs)
except openai.BadRequestError as e:
if "does not exist" not in str(e):
raise
last_error = e
await asyncio.sleep(delay)
delay = min(delay * 1.5, 8.0)
assert last_error is not None
raise last_error
@pytest.mark.asyncio
async def test_create_fine_tune_jobs_async():
try:
@ -64,9 +95,11 @@ async def test_create_fine_tune_jobs_async():
)
print("Response from creating file=", file_obj)
create_fine_tuning_response = await litellm.acreate_fine_tuning_job(
model="gpt-3.5-turbo-0125",
training_file=file_obj.id,
create_fine_tuning_response = (
await _acreate_fine_tuning_job_with_propagation_retry(
model="gpt-4o-mini-2024-07-18",
training_file=file_obj.id,
)
)
print(
@ -74,7 +107,7 @@ async def test_create_fine_tune_jobs_async():
)
assert create_fine_tuning_response.id is not None
assert create_fine_tuning_response.model == "gpt-3.5-turbo-0125"
assert create_fine_tuning_response.model == "gpt-4o-mini-2024-07-18"
await asyncio.sleep(2)
_logged_standard_logging_object = custom_logger.standard_logging_object
@ -83,7 +116,7 @@ async def test_create_fine_tune_jobs_async():
"custom_logger.standard_logging_object=",
json.dumps(_logged_standard_logging_object, indent=4),
)
assert _logged_standard_logging_object["model"] == "gpt-3.5-turbo-0125"
assert _logged_standard_logging_object["model"] == "gpt-4o-mini-2024-07-18"
assert _logged_standard_logging_object["id"] == create_fine_tuning_response.id
# list fine tuning jobs
@ -427,10 +460,10 @@ async def test_mock_openai_create_fine_tune_job():
with patch.object(client.fine_tuning.jobs, "create") as mock_create:
mock_create.return_value = FineTuningJob(
id="ft-123",
model="gpt-3.5-turbo-0125",
model="gpt-4o-mini-2024-07-18",
created_at=1677610602,
status="validating_files",
fine_tuned_model="ft:gpt-3.5-turbo-0125:org:custom_suffix:id",
fine_tuned_model="ft:gpt-4o-mini-2024-07-18:org:custom_suffix:id",
object="fine_tuning.job",
hyperparameters=Hyperparameters(
n_epochs=3,
@ -442,7 +475,7 @@ async def test_mock_openai_create_fine_tune_job():
)
response = await litellm.acreate_fine_tuning_job(
model="gpt-3.5-turbo-0125",
model="gpt-4o-mini-2024-07-18",
training_file="file-123",
hyperparameters={"n_epochs": 3},
suffix="custom_suffix",
@ -453,16 +486,19 @@ async def test_mock_openai_create_fine_tune_job():
mock_create.assert_called_once()
request_params = mock_create.call_args.kwargs
assert request_params["model"] == "gpt-3.5-turbo-0125"
assert request_params["model"] == "gpt-4o-mini-2024-07-18"
assert request_params["training_file"] == "file-123"
assert request_params["hyperparameters"] == {"n_epochs": 3}
assert request_params["suffix"] == "custom_suffix"
# Verify the response
assert response.id == "ft-123"
assert response.model == "gpt-3.5-turbo-0125"
assert response.model == "gpt-4o-mini-2024-07-18"
assert response.status == "validating_files"
assert response.fine_tuned_model == "ft:gpt-3.5-turbo-0125:org:custom_suffix:id"
assert (
response.fine_tuned_model
== "ft:gpt-4o-mini-2024-07-18:org:custom_suffix:id"
)
@pytest.mark.asyncio

View file

@ -126,6 +126,7 @@ sentry_sdk: >=2.21.0 # Unknown license
cryptography: >=43.0.1 # Unknown license
tzdata: >=2025.1 # Unknown license
urllib3: >=2.0.0 # MIT license - https://github.com/urllib3/urllib3
audioread: >=3.0.1 # MIT license manually verified - https://github.com/beetbox/audioread
python-dotenv: >=1.0.0 # Unknown license
tiktoken: >=0.8.0 # Unknown license
click: >=8.1.7 # Unknown license

View file

@ -661,14 +661,14 @@ async def test_async_log_failure_event(prometheus_logger):
# litellm_llm_api_failed_requests_metric incremented
# Labels: end_user, hashed_api_key, api_key_alias, model, team, team_alias, user, model_id
prometheus_logger.litellm_llm_api_failed_requests_metric.labels.assert_called_once_with(
None, # end_user_id
"test_hash",
"test_alias",
"gpt-3.5-turbo",
"test_team",
"test_team_alias",
"test_user",
"model-123", # model_id from standard_logging_payload
end_user=None,
hashed_api_key="test_hash",
api_key_alias="test_alias",
model="gpt-3.5-turbo",
team="test_team",
team_alias="test_team_alias",
user="test_user",
model_id="model-123",
)
prometheus_logger.litellm_llm_api_failed_requests_metric.labels().inc.assert_called_once()

View file

@ -0,0 +1,29 @@
import json
from pathlib import Path
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
def test_sambanova_minimax_m27_model_info():
model = "sambanova/MiniMax-M2.7"
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
with open(json_path) as f:
model_cost = json.load(f)
info = model_cost.get(model)
assert (
info is not None
), f"{model} not found in model_prices_and_context_window.json"
assert info["litellm_provider"] == "sambanova"
assert info["mode"] == "chat"
assert info["input_cost_per_token"] > 0
assert info["output_cost_per_token"] > 0
assert info["max_input_tokens"] == 204800
assert info["max_output_tokens"] == 131072
assert info["supports_function_calling"] is True
assert info["supports_reasoning"] is True
assert info["supports_tool_choice"] is True
routed_model, provider, _, _ = get_llm_provider(model=model)
assert routed_model == "MiniMax-M2.7"
assert provider == "sambanova"

View file

@ -853,7 +853,11 @@ class BaseLLMChatTest(ABC):
@pytest.mark.parametrize(
"image_url",
[
"http://img1.etsystatic.com/260/0/7813604/il_fullxfull.4226713999_q86e.jpg",
# In-repo logo served via jsdelivr (sha-pinned, immutable).
# Bedrock fetches the URL and base64-embeds it in the
# Converse request body; using a multi-MB hosted product
# photo here previously bloated cassettes to ~60 MB each.
"https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/ui/litellm-dashboard/public/assets/logos/litellm_logo.jpg",
"https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png",
],
)

View file

@ -101,7 +101,7 @@ async def test_openai_realtime_direct_call_no_intent():
try:
await litellm._arealtime(
model="openai/gpt-4o-realtime-preview-2024-10-01",
model="openai/gpt-4o-realtime-preview",
websocket=websocket_client,
api_key=os.environ.get("OPENAI_API_KEY"),
timeout=60,
@ -250,13 +250,13 @@ async def test_openai_realtime_direct_call_with_intent():
caught_exception = None
query_params: RealtimeQueryParams = {
"model": "openai/gpt-4o-realtime-preview-2024-10-01",
"model": "openai/gpt-4o-realtime-preview",
"intent": "chat",
}
try:
await litellm._arealtime(
model="openai/gpt-4o-realtime-preview-2024-10-01",
model="openai/gpt-4o-realtime-preview",
websocket=websocket_client,
api_key=os.environ.get("OPENAI_API_KEY"),
query_params=query_params,
@ -331,7 +331,7 @@ def test_realtime_query_params_construction():
from litellm.types.realtime import RealtimeQueryParams
# Test case 1: intent is None (should not be included)
model = "gpt-4o-realtime-preview-2024-10-01"
model = "gpt-4o-realtime-preview"
intent = None
query_params: RealtimeQueryParams = {"model": model}
@ -369,17 +369,17 @@ async def test_realtime_query_params_use_normalized_model_name(monkeypatch):
)
def fake_get_llm_provider(model, api_base=None, api_key=None):
return ("gpt-4o-realtime-preview-2024-10-01", "openai", None, None)
return ("gpt-4o-realtime-preview", "openai", None, None)
monkeypatch.setattr(realtime_main, "get_llm_provider", fake_get_llm_provider)
query_params: RealtimeQueryParams = {
"model": "openai/gpt-4o-realtime-preview-2024-10-01",
"model": "openai/gpt-4o-realtime-preview",
"intent": "chat",
}
await realtime_main._arealtime(
model="openai/gpt-4o-realtime-preview-2024-10-01",
model="openai/gpt-4o-realtime-preview",
websocket=MagicMock(),
api_key="sk-test",
query_params=query_params,
@ -387,7 +387,5 @@ async def test_realtime_query_params_use_normalized_model_name(monkeypatch):
)
called_kwargs = mock_async_realtime.call_args.kwargs
assert (
called_kwargs["query_params"]["model"] == "gpt-4o-realtime-preview-2024-10-01"
)
assert called_kwargs["query_params"]["model"] == "gpt-4o-realtime-preview"
assert called_kwargs["query_params"]["intent"] == "chat"

View file

@ -2,6 +2,7 @@
Tests for Evals API operations across providers
"""
import hashlib
import os
import sys
from abc import ABC, abstractmethod
@ -20,6 +21,46 @@ from litellm.types.llms.openai_evals import (
)
def _stable_eval_name(test_node_name: str, suffix: str = "") -> str:
"""Deterministic eval name keyed off the test's node name.
The previous ``f"Test Eval {int(time.time())}"`` pattern embedded a
fresh value into the request body every run, defeating VCR's
``safe_body`` matcher and forcing a real OpenAI ``create`` call on
every CI run. With a stable per-test name the cassette matches on
replay, and provider-side resources stay bounded because each test
deletes the eval it owns on teardown.
"""
nonce = hashlib.sha1(test_node_name.encode()).hexdigest()[:12]
return f"vcr-managed-{nonce}{suffix}"
_TESTING_CRITERIA = [
{
"type": "label_model",
"model": "gpt-4o",
"input": [
{
"role": "developer",
"content": "Classify the sentiment as 'positive' or 'negative'",
},
{"role": "user", "content": "Statement: {{item.input}}"},
],
"passing_labels": ["positive"],
"labels": ["positive", "negative"],
"name": "Sentiment grader",
}
]
_PROVIDER_FLAKINESS = (
litellm.InternalServerError,
litellm.APIConnectionError,
litellm.Timeout,
litellm.ServiceUnavailableError,
)
class BaseEvalsAPITest(ABC):
"""
Base test class for Evals API operations.
@ -41,13 +82,64 @@ class BaseEvalsAPITest(ABC):
"""Return the API base URL for the provider"""
pass
@pytest.fixture
def managed_eval(self, request):
"""Create a stable-named eval for this test; delete on teardown.
Function-scoped so each cassette captures the full
createtestdelete cycle. A class-scoped fixture would push
the create into whichever test ran first and the delete into
whichever ran last, which is fragile under reordering.
Replaces the prior ``list_evals().data[0].id`` pattern, which
made the URL of ``get_eval`` / ``update_eval`` vary across
runs (the "first" eval depends on what other runs left
behind).
"""
custom_llm_provider = self.get_custom_llm_provider()
api_key = self.get_api_key()
api_base = self.get_api_base()
if not api_key:
pytest.skip(f"No API key provided for {custom_llm_provider}")
try:
created = litellm.create_eval(
name=_stable_eval_name(request.node.name),
data_source_config={
"type": "stored_completions",
"metadata": {"usecase": "chatbot", "vcr": "managed"},
},
testing_criteria=_TESTING_CRITERIA,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
except _PROVIDER_FLAKINESS:
pytest.skip("Provider service unavailable")
except litellm.RateLimitError:
pytest.skip("Rate limit exceeded")
yield created
# Best-effort cleanup. OpenAI eval names are not unique-keyed
# (only IDs are), so a failed delete doesn't block the next
# run's create.
try:
litellm.delete_eval(
eval_id=created.id,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
except Exception:
pass
@pytest.mark.flaky(retries=3, delay=2)
def test_create_eval(self):
def test_create_eval(self, request):
"""
Test creating an evaluation.
"""
import time
custom_llm_provider = self.get_custom_llm_provider()
api_key = self.get_api_key()
api_base = self.get_api_base()
@ -56,53 +148,45 @@ class BaseEvalsAPITest(ABC):
pytest.skip(f"No API key provided for {custom_llm_provider}")
litellm.set_verbose = True
unique_name = _stable_eval_name(request.node.name)
# Create eval with stored_completions data source
unique_name = f"Test Eval {int(time.time())}"
created_id = None
try:
response = litellm.create_eval(
name=unique_name,
data_source_config={
"type": "stored_completions",
"metadata": {"usecase": "chatbot"},
},
testing_criteria=[
{
"type": "label_model",
"model": "gpt-4o",
"input": [
{
"role": "developer",
"content": "Classify the sentiment as 'positive' or 'negative'",
},
{"role": "user", "content": "Statement: {{item.input}}"},
],
"passing_labels": ["positive"],
"labels": ["positive", "negative"],
"name": "Sentiment grader",
}
],
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
except (
litellm.InternalServerError,
litellm.APIConnectionError,
litellm.Timeout,
litellm.ServiceUnavailableError,
):
pytest.skip("Provider service unavailable")
except litellm.RateLimitError:
pytest.skip("Rate limit exceeded")
try:
response = litellm.create_eval(
name=unique_name,
data_source_config={
"type": "stored_completions",
"metadata": {"usecase": "chatbot"},
},
testing_criteria=_TESTING_CRITERIA,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
except _PROVIDER_FLAKINESS:
pytest.skip("Provider service unavailable")
except litellm.RateLimitError:
pytest.skip("Rate limit exceeded")
assert response is not None
assert isinstance(response, Eval)
assert response.id is not None
assert response.name == unique_name
print(f"Created eval: {response}")
print(f"Eval ID: {response.id}")
assert response is not None
assert isinstance(response, Eval)
assert response.id is not None
assert response.name == unique_name
created_id = response.id
print(f"Created eval: {response}")
print(f"Eval ID: {response.id}")
finally:
if created_id is not None:
try:
litellm.delete_eval(
eval_id=created_id,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
except Exception:
pass
def test_list_evals(self):
"""
@ -130,7 +214,7 @@ class BaseEvalsAPITest(ABC):
assert hasattr(response, "has_more")
print(f"Listed evals: {len(response.data)} evaluations")
def test_get_eval(self):
def test_get_eval(self, managed_eval):
"""
Test getting a specific evaluation by ID.
"""
@ -138,89 +222,54 @@ class BaseEvalsAPITest(ABC):
api_key = self.get_api_key()
api_base = self.get_api_base()
if not api_key:
pytest.skip(f"No API key provided for {custom_llm_provider}")
litellm.set_verbose = True
# First list existing evals to get an ID
list_response = litellm.list_evals(
limit=1,
response = litellm.get_eval(
eval_id=managed_eval.id,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
assert isinstance(list_response, ListEvalsResponse)
assert response is not None
assert isinstance(response, Eval)
assert response.id == managed_eval.id
print(f"Retrieved eval: {response}")
if list_response.data and len(list_response.data) > 0:
eval_id = list_response.data[0].id
print(f"Testing with eval ID: {eval_id}")
# Get the eval
response = litellm.get_eval(
eval_id=eval_id,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
assert response is not None
assert isinstance(response, Eval)
assert response.id == eval_id
print(f"Retrieved eval: {response}")
else:
pytest.skip("No existing evals to test with")
def test_update_eval(self):
@pytest.mark.flaky(retries=3, delay=2)
def test_update_eval(self, request, managed_eval):
"""
Test updating an evaluation.
"""
import time
custom_llm_provider = self.get_custom_llm_provider()
api_key = self.get_api_key()
api_base = self.get_api_base()
if not api_key:
pytest.skip(f"No API key provided for {custom_llm_provider}")
litellm.set_verbose = True
updated_name = _stable_eval_name(request.node.name, suffix="-updated")
# First list existing evals
list_response = litellm.list_evals(
limit=1,
response = litellm.update_eval(
eval_id=managed_eval.id,
name=updated_name,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
assert isinstance(list_response, ListEvalsResponse)
if list_response.data and len(list_response.data) > 0:
eval_id = list_response.data[0].id
updated_name = f"Updated Eval {int(time.time())}"
# Update the eval
response = litellm.update_eval(
eval_id=eval_id,
name=updated_name,
custom_llm_provider=custom_llm_provider,
api_key=api_key,
api_base=api_base,
)
assert response is not None
assert isinstance(response, Eval)
assert response.id == eval_id
assert response.name == updated_name
print(f"Updated eval: {response}")
else:
pytest.skip("No existing evals to test with")
assert response is not None
assert isinstance(response, Eval)
assert response.id == managed_eval.id
assert response.name == updated_name
print(f"Updated eval: {response}")
def test_delete_eval(self):
"""
Test deleting an evaluation.
Real delete coverage now lives in the ``managed_eval`` fixture
teardown and in ``test_create_eval``'s ``finally`` block, so
this stays a no-op skip rather than creating a fresh resource
just to delete it.
"""
custom_llm_provider = self.get_custom_llm_provider()
api_key = self.get_api_key()
@ -229,8 +278,7 @@ class BaseEvalsAPITest(ABC):
if not api_key:
pytest.skip(f"No API key provided for {custom_llm_provider}")
# Skip this test to avoid deleting production evals
pytest.skip("Skipping delete test to preserve existing evals")
pytest.skip("Delete is exercised via managed_eval fixture teardown.")
class TestOpenAIEvalsAPI(BaseEvalsAPITest):

View file

@ -0,0 +1,220 @@
"""Unit tests for the VCR record-time filters that keep cassettes small.
Covers:
- ``_strip_image_b64_payloads`` replaces base64 image bodies in
image-gen responses so cassettes don't carry MB-class PNG payloads.
- ``_normalize_multipart_boundary`` rewrites random multipart
boundaries to a fixed string so audio-transcription request bodies
match across record and replay.
"""
from __future__ import annotations
import json
import os
import sys
from vcr.request import Request
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
from tests._vcr_conftest_common import ( # noqa: E402
VCR_FIXED_MULTIPART_BOUNDARY,
VCR_IMAGE_B64_PLACEHOLDER,
_normalize_multipart_boundary,
_strip_image_b64_payloads,
)
# ---------------------------------------------------------------------------
# Image b64 stripper
# ---------------------------------------------------------------------------
def _image_response(b64_payload: str, body_type: str = "bytes") -> dict:
body_text = json.dumps({"data": [{"b64_json": b64_payload}]})
body_string = body_text.encode("utf-8") if body_type == "bytes" else body_text
return {
"status": {"code": 200, "message": "OK"},
"headers": {
"content-type": ["application/json"],
"content-length": [str(len(body_text.encode("utf-8")))],
},
"body": {"string": body_string},
}
def test_strip_image_b64_replaces_payload_when_body_is_bytes():
response = _image_response("A" * 5000, body_type="bytes")
out = _strip_image_b64_payloads(response)
payload = json.loads(out["body"]["string"].decode("utf-8"))
assert payload["data"][0]["b64_json"] == VCR_IMAGE_B64_PLACEHOLDER
def test_strip_image_b64_replaces_payload_when_body_is_str():
response = _image_response("A" * 5000, body_type="str")
out = _strip_image_b64_payloads(response)
payload = json.loads(out["body"]["string"])
assert payload["data"][0]["b64_json"] == VCR_IMAGE_B64_PLACEHOLDER
def test_strip_image_b64_updates_content_length():
response = _image_response("A" * 5000)
out = _strip_image_b64_payloads(response)
expected_len = len(out["body"]["string"])
assert out["headers"]["content-length"] == [str(expected_len)]
def test_strip_image_b64_is_idempotent():
response = _image_response("A" * 5000)
once = _strip_image_b64_payloads(response)
twice = _strip_image_b64_payloads(once)
assert once["body"]["string"] == twice["body"]["string"]
def test_strip_image_b64_handles_nested_data():
body_text = json.dumps(
{
"outer": {
"data": [
{"b64_json": "X" * 4000, "label": "first"},
{"b64_json": "Y" * 4000, "label": "second"},
]
}
}
)
response = {
"status": {"code": 200, "message": "OK"},
"headers": {"content-type": ["application/json"]},
"body": {"string": body_text.encode("utf-8")},
}
out = _strip_image_b64_payloads(response)
payload = json.loads(out["body"]["string"].decode("utf-8"))
assert payload["outer"]["data"][0]["b64_json"] == VCR_IMAGE_B64_PLACEHOLDER
assert payload["outer"]["data"][1]["b64_json"] == VCR_IMAGE_B64_PLACEHOLDER
assert payload["outer"]["data"][0]["label"] == "first"
def test_strip_image_b64_leaves_non_image_response_unchanged():
body_text = json.dumps({"choices": [{"message": {"content": "hello"}}]})
response = {
"status": {"code": 200, "message": "OK"},
"headers": {"content-type": ["application/json"]},
"body": {"string": body_text.encode("utf-8")},
}
out = _strip_image_b64_payloads(response)
assert json.loads(out["body"]["string"].decode("utf-8")) == json.loads(body_text)
def test_strip_image_b64_leaves_invalid_json_unchanged():
response = {
"status": {"code": 200, "message": "OK"},
"headers": {"content-type": ["application/octet-stream"]},
"body": {"string": b"\x89PNG\r\n\x1a\n binary stuff not json"},
}
out = _strip_image_b64_payloads(response)
assert out["body"]["string"] == b"\x89PNG\r\n\x1a\n binary stuff not json"
def test_strip_image_b64_skips_short_values():
"""Already-placeholder values aren't re-replaced (idempotency guard)."""
body_text = json.dumps({"data": [{"b64_json": VCR_IMAGE_B64_PLACEHOLDER}]})
response = {
"status": {"code": 200, "message": "OK"},
"headers": {"content-type": ["application/json"]},
"body": {"string": body_text.encode("utf-8")},
}
out = _strip_image_b64_payloads(response)
payload = json.loads(out["body"]["string"].decode("utf-8"))
assert payload["data"][0]["b64_json"] == VCR_IMAGE_B64_PLACEHOLDER
# ---------------------------------------------------------------------------
# Multipart boundary normalizer
# ---------------------------------------------------------------------------
def _multipart_request(boundary: str):
body_text = (
f"--{boundary}\r\n"
'Content-Disposition: form-data; name="file"; filename="audio.wav"\r\n'
"Content-Type: audio/wav\r\n"
"\r\n"
"fake-audio-bytes\r\n"
f"--{boundary}--\r\n"
)
return Request(
method="POST",
uri="https://api.openai.com/v1/audio/transcriptions",
body=body_text.encode("utf-8"),
headers={
"content-type": f"multipart/form-data; boundary={boundary}",
},
)
def test_normalize_multipart_rewrites_header_and_body():
req = _multipart_request("abc123random")
_normalize_multipart_boundary(req)
assert (
req.headers["content-type"]
== f"multipart/form-data; boundary={VCR_FIXED_MULTIPART_BOUNDARY}"
)
assert b"abc123random" not in req.body
assert VCR_FIXED_MULTIPART_BOUNDARY.encode("utf-8") in req.body
def test_normalize_multipart_is_idempotent():
req = _multipart_request("abc123random")
_normalize_multipart_boundary(req)
body_first = req.body
header_first = req.headers["content-type"]
_normalize_multipart_boundary(req)
assert req.body == body_first
assert req.headers["content-type"] == header_first
def test_normalize_multipart_two_distinct_boundaries_match_after_normalize():
"""Whisper-style: two requests with different random boundaries should
end up with byte-identical bodies after normalization."""
req1 = _multipart_request("boundaryAAA")
req2 = _multipart_request("boundaryBBB")
_normalize_multipart_boundary(req1)
_normalize_multipart_boundary(req2)
assert req1.body == req2.body
assert req1.headers["content-type"] == req2.headers["content-type"]
def test_normalize_multipart_skips_non_multipart_requests():
req = Request(
method="POST",
uri="https://api.openai.com/v1/chat/completions",
body=b'{"model":"gpt-4o"}',
headers={"content-type": "application/json"},
)
_normalize_multipart_boundary(req)
assert req.headers["content-type"] == "application/json"
assert req.body == b'{"model":"gpt-4o"}'
def test_normalize_multipart_skips_request_without_content_type():
req = Request(
method="POST",
uri="https://api.openai.com/v1/chat/completions",
body=b"unknown body",
headers={},
)
_normalize_multipart_boundary(req)
assert req.body == b"unknown body"
def test_normalize_multipart_handles_quoted_boundary():
req = Request(
method="POST",
uri="https://api.openai.com/v1/audio/transcriptions",
body=b"--quoted-boundary--body content--quoted-boundary--",
headers={"content-type": 'multipart/form-data; boundary="quoted-boundary"'},
)
_normalize_multipart_boundary(req)
assert b"quoted-boundary" not in req.body
assert VCR_FIXED_MULTIPART_BOUNDARY.encode("utf-8") in req.body

View file

@ -193,19 +193,12 @@ def _azure_ai_image_mock_response(*args, **kwargs):
return new_response
@pytest.mark.parametrize(
"model, api_base, api_key",
[
(
"azure_ai/Cohere-embed-v3-multilingual-2",
os.getenv("AZURE_AI_API_BASE"),
os.getenv("AZURE_AI_API_KEY"),
)
],
)
@pytest.mark.parametrize("sync_mode", [True]) # , False
@pytest.mark.asyncio
async def test_azure_ai_embedding_image(model, api_base, api_key, sync_mode):
async def test_azure_ai_embedding_image(sync_mode):
model = "azure_ai/Cohere-embed-v3-multilingual-2"
api_base = os.getenv("AZURE_AI_API_BASE")
api_key = os.getenv("AZURE_AI_API_KEY")
try:
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")

View file

@ -101,6 +101,7 @@ _SCALAR_ATTRS = (
"redact_messages_in_exceptions",
"redact_user_api_key_info",
"s3_callback_params",
"s3_audit_callback_params",
"datadog_params",
"vector_store_registry",
)
@ -128,6 +129,7 @@ def isolate_litellm_state():
leaking across tests within the same xdist worker.
"""
from litellm.litellm_core_utils import litellm_logging as ll_logging
from litellm.proxy.management_helpers import audit_logs as ll_audit_logs
# Flush cache and clear internal logger instances before test
if hasattr(litellm, "in_memory_llm_clients_cache"):
@ -135,6 +137,7 @@ def isolate_litellm_state():
# Clear cached logger instances (LangsmithLogger, SlackAlerting, etc.)
ll_logging._in_memory_loggers.clear()
ll_audit_logs._audit_log_callback_cache.clear()
# Reset ALL attrs to their true defaults before the test runs.
# This undoes any module-level mutations from test file imports.
@ -156,6 +159,7 @@ def isolate_litellm_state():
litellm.in_memory_llm_clients_cache.flush_cache()
ll_logging._in_memory_loggers.clear()
ll_audit_logs._audit_log_callback_cache.clear()
for attr in _LIST_ATTRS:
if attr in _DEFAULTS:

View file

@ -12,7 +12,15 @@ from abc import ABC, abstractmethod
# Test resources
TEST_IMAGE_PATH = "test_image_edit.png"
TEST_PDF_URL = "https://arxiv.org/pdf/2201.04234"
# Tiny in-repo PDF served via jsdelivr (sha-pinned, immutable). The arxiv
# PDF previously used here was several MB — once base64-encoded into the
# Vertex OCR request it ballooned cassettes past 100 MB per test. Keep
# the URL stable across runs so cassettes don't churn.
TEST_PDF_URL = (
"https://cdn.jsdelivr.net/gh/BerriAI/litellm"
"@d769e81c90d453240c61fc572cdb27fae06a89d0"
"/tests/llm_translation/fixtures/dummy.pdf"
)
class BaseOCRTest(ABC):

View file

@ -501,21 +501,30 @@ def test_is_request_body_safe_model_enabled(
assert expect_error == error_raised
@pytest.mark.parametrize(
"api_key_value, expect_complete",
[
("sk-real-key", True),
("", False),
(None, False),
(" ", False),
],
)
def test_check_complete_credentials_api_key_values(api_key_value, expect_complete):
def _assert_check_complete_credentials(api_key_value, expect_complete):
request_body = {"model": "gpt-3.5-turbo", "api_key": api_key_value}
result = check_complete_credentials(request_body=request_body)
assert result == expect_complete
def test_check_complete_credentials_with_real_key():
_assert_check_complete_credentials(
api_key_value="sk-" + "x" * 8, expect_complete=True
)
def test_check_complete_credentials_with_empty_string():
_assert_check_complete_credentials(api_key_value="", expect_complete=False)
def test_check_complete_credentials_with_none():
_assert_check_complete_credentials(api_key_value=None, expect_complete=False)
def test_check_complete_credentials_with_whitespace():
_assert_check_complete_credentials(api_key_value=" ", expect_complete=False)
def test_reading_openai_org_id_from_headers():
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup

View file

@ -0,0 +1,135 @@
import pytest
from starlette.responses import JSONResponse
from starlette.testclient import TestClient
from starlette.types import Message
from litellm.proxy.middleware.request_size_limit_middleware import (
RequestSizeLimitMiddleware,
)
def test_request_size_limit_middleware_rejects_content_length_before_body_read():
downstream_called = False
async def app(scope, receive, send):
nonlocal downstream_called
downstream_called = True
response = JSONResponse({"ok": True})
await response(scope, receive, send)
client = TestClient(
RequestSizeLimitMiddleware(
app,
get_max_request_size_mb=lambda: 1,
is_request_size_limit_enabled=lambda: True,
)
)
response = client.post(
"/chat/completions",
content=b"x" * (1024 * 1024 + 1),
headers={"content-type": "application/json"},
)
assert response.status_code == 413
assert response.json() == {"error": "Request size is too large. Max size is 1 MB"}
assert response.headers["content-length"] == str(len(response.content))
assert downstream_called is False
def test_request_size_limit_middleware_zero_limit_disables_guard():
downstream_called = False
async def app(scope, receive, send):
nonlocal downstream_called
downstream_called = True
response = JSONResponse({"ok": True})
await response(scope, receive, send)
client = TestClient(
RequestSizeLimitMiddleware(
app,
get_max_request_size_mb=lambda: 0,
is_request_size_limit_enabled=lambda: True,
)
)
response = client.post(
"/chat/completions",
content=b"x",
headers={"content-type": "application/json"},
)
assert response.status_code == 200
assert response.json() == {"ok": True}
assert downstream_called is True
@pytest.mark.asyncio
async def test_request_size_limit_middleware_rejects_streamed_body_without_content_length():
received_body_bytes = 0
async def app(scope, receive, send):
nonlocal received_body_bytes
while True:
message = await receive()
if message["type"] == "http.disconnect":
break
received_body_bytes += len(message.get("body", b""))
if not message.get("more_body", False):
break
response = JSONResponse({"ok": True})
await response(scope, receive, send)
middleware = RequestSizeLimitMiddleware(
app,
get_max_request_size_mb=lambda: 1,
is_request_size_limit_enabled=lambda: True,
)
sent_messages: list[Message] = []
receive_messages: list[Message] = [
{
"type": "http.request",
"body": b"x" * (1024 * 1024),
"more_body": True,
},
{
"type": "http.request",
"body": b"y",
"more_body": False,
},
]
async def receive():
return receive_messages.pop(0)
async def send(message):
sent_messages.append(message)
await middleware(
{
"type": "http",
"method": "POST",
"path": "/chat/completions",
"headers": [(b"content-type", b"application/json")],
},
receive,
send,
)
expected_body = b'{"error":"Request size is too large. Max size is 1 MB"}'
assert sent_messages[0] == {
"type": "http.response.start",
"status": 413,
"headers": [
(b"content-type", b"application/json"),
(b"content-length", str(len(expected_body)).encode("latin-1")),
],
}
assert sent_messages[1] == {
"type": "http.response.body",
"body": expected_body,
"more_body": False,
}
assert received_body_bytes == 1024 * 1024

View file

@ -269,9 +269,7 @@ async def test_aaauser_personal_budgets(key_ownership):
test_user_cache = getattr(litellm.proxy.proxy_server, "user_api_key_cache")
assert (
test_user_cache.get_cache(
key=hash_token(user_key), model_type=UserAPIKeyAuth
)
test_user_cache.get_cache(key=hash_token(user_key), model_type=UserAPIKeyAuth)
== valid_token
)
@ -514,36 +512,59 @@ async def test_auth_not_connected_to_db():
assert valid_token.token == "failed-to-connect-to-db"
@pytest.mark.parametrize(
"headers, custom_header_name, expected_api_key",
[
# Test with valid Bearer token
({"x-custom-api-key": "Bearer sk-12345678"}, "x-custom-api-key", "sk-12345678"),
# Test with raw token (no Bearer prefix)
({"x-custom-api-key": "Bearer sk-12345678"}, "x-custom-api-key", "sk-12345678"),
# Test with empty header value
({"x-custom-api-key": ""}, "x-custom-api-key", ""),
# Test with missing header
({}, "X-Custom-API-Key", ""),
# Test with different header casing
({"X-CUSTOM-API-KEY": "Bearer sk-12345678"}, "X-Custom-API-Key", "sk-12345678"),
],
)
def test_get_api_key_from_custom_header(headers, custom_header_name, expected_api_key):
def _assert_api_key_from_custom_header(headers, custom_header_name, expected_api_key):
verbose_proxy_logger.setLevel(logging.DEBUG)
# Mock the Request object
request = MagicMock(spec=Request)
request.headers = headers
# Call the function and verify it doesn't raise an exception
api_key = get_api_key_from_custom_header(
request=request, custom_litellm_key_header_name=custom_header_name
)
assert api_key == expected_api_key
def test_get_api_key_from_custom_header_bearer_token():
token = "sk-" + "1" * 8
_assert_api_key_from_custom_header(
headers={"x-custom-api-key": f"Bearer {token}"},
custom_header_name="x-custom-api-key",
expected_api_key=token,
)
def test_get_api_key_from_custom_header_raw_token():
token = "sk-" + "1" * 8
_assert_api_key_from_custom_header(
headers={"x-custom-api-key": f"Bearer {token}"},
custom_header_name="x-custom-api-key",
expected_api_key=token,
)
def test_get_api_key_from_custom_header_empty_value():
_assert_api_key_from_custom_header(
headers={"x-custom-api-key": ""},
custom_header_name="x-custom-api-key",
expected_api_key="",
)
def test_get_api_key_from_custom_header_missing_header():
_assert_api_key_from_custom_header(
headers={},
custom_header_name="X-Custom-API-Key",
expected_api_key="",
)
def test_get_api_key_from_custom_header_different_casing():
token = "sk-" + "1" * 8
_assert_api_key_from_custom_header(
headers={"X-CUSTOM-API-KEY": f"Bearer {token}"},
custom_header_name="X-Custom-API-Key",
expected_api_key=token,
)
from litellm.proxy._types import LitellmUserRoles

View file

@ -2,13 +2,18 @@
Test Azure Sentinel logging integration
"""
import datetime
from unittest.mock import AsyncMock, patch
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
def _close_periodic_flush_task(coro):
coro.close()
return None
@pytest.mark.asyncio
@ -20,7 +25,7 @@ async def test_azure_sentinel_oauth_and_send_batch():
test_client_id = "test-client-id"
test_client_secret = "test-client-secret"
with patch("asyncio.create_task"):
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id=test_dcr_id,
endpoint=test_endpoint,
@ -42,9 +47,6 @@ async def test_azure_sentinel_oauth_and_send_batch():
# Add to queue
logger.log_queue.append(standard_payload)
# Mock OAuth token response
from unittest.mock import MagicMock
mock_token_response = MagicMock()
mock_token_response.status_code = 200
mock_token_response.json = MagicMock(
@ -91,3 +93,173 @@ async def test_azure_sentinel_oauth_and_send_batch():
# Verify queue is cleared
assert len(logger.log_queue) == 0
@pytest.mark.asyncio
async def test_azure_sentinel_queues_audit_log_event():
"""Test that Azure Sentinel supports direct audit log callbacks"""
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id="dcr-test123456789",
endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com",
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
)
logger.batch_size = 2
logger.async_send_audit_batch = AsyncMock()
audit_log = StandardAuditLogPayload(
id="audit-123",
updated_at="2026-05-06T04:39:00+00:00",
changed_by="user-1",
changed_by_api_key="sk-test",
action="created",
table_name="LiteLLM_TeamTable",
object_id="team-1",
before_value=None,
updated_values='{"team_alias": "sentinel-demo"}',
)
await logger.async_log_audit_log_event(audit_log)
assert logger.audit_log_queue == [audit_log]
logger.async_send_audit_batch.assert_not_called()
await logger.async_log_audit_log_event(audit_log)
assert logger.audit_log_queue == [audit_log, audit_log]
logger.async_send_audit_batch.assert_awaited_once()
@pytest.mark.asyncio
async def test_azure_sentinel_sends_audit_log_payload_to_ingestion_api():
"""Test that queued audit logs are sent to Azure Monitor Logs Ingestion"""
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id="dcr-test123456789",
endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com",
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
)
audit_log = StandardAuditLogPayload(
id="audit-123",
updated_at="2026-05-06T04:39:00+00:00",
changed_by="user-1",
changed_by_api_key="sk-test",
action="created",
table_name="LiteLLM_TeamTable",
object_id="team-1",
before_value=None,
updated_values='{"team_alias": "sentinel-demo"}',
)
await logger.async_log_audit_log_event(audit_log)
mock_token_response = MagicMock()
mock_token_response.status_code = 200
mock_token_response.json = MagicMock(
return_value={
"access_token": "test-bearer-token",
"expires_in": 3600,
}
)
mock_token_response.text = "Success"
mock_api_response = MagicMock()
mock_api_response.status_code = 204
mock_api_response.text = "Success"
async def mock_post(*args, **kwargs):
if "oauth2/v2.0/token" in kwargs.get("url", ""):
return mock_token_response
return mock_api_response
logger.async_httpx_client.post = AsyncMock(side_effect=mock_post)
await logger.flush_queue()
api_call_args = logger.async_httpx_client.post.call_args_list[-1]
body = json.loads(api_call_args.kwargs["data"].decode("utf-8"))
assert body == [audit_log]
assert "dcr-test123456789" in api_call_args.kwargs["url"]
assert "Custom-LiteLLM" in api_call_args.kwargs["url"]
assert len(logger.audit_log_queue) == 0
@pytest.mark.asyncio
async def test_azure_sentinel_flushes_standard_and_audit_logs_separately():
"""Test mixed callback roles do not send schema-mismatched batches."""
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id="dcr-test123456789",
stream_name="Custom-LiteLLM-Standard",
audit_stream_name="Custom-LiteLLM-Audit",
endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com",
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
)
standard_payload = StandardLoggingPayload(
id="standard-123",
call_type="completion",
model="gpt-3.5-turbo",
status="success",
messages=[{"role": "user", "content": "Hello"}],
response={"choices": [{"message": {"content": "Hi"}}]},
)
audit_log = StandardAuditLogPayload(
id="audit-123",
updated_at="2026-05-06T04:39:00+00:00",
changed_by="user-1",
changed_by_api_key="sk-test",
action="created",
table_name="LiteLLM_TeamTable",
object_id="team-1",
before_value=None,
updated_values='{"team_alias": "sentinel-demo"}',
)
logger.log_queue.append(standard_payload)
await logger.async_log_audit_log_event(audit_log)
mock_token_response = MagicMock()
mock_token_response.status_code = 200
mock_token_response.json = MagicMock(
return_value={
"access_token": "test-bearer-token",
"expires_in": 3600,
}
)
mock_token_response.text = "Success"
mock_api_response = MagicMock()
mock_api_response.status_code = 204
mock_api_response.text = "Success"
async def mock_post(*args, **kwargs):
if "oauth2/v2.0/token" in kwargs.get("url", ""):
return mock_token_response
return mock_api_response
logger.async_httpx_client.post = AsyncMock(side_effect=mock_post)
await logger.flush_queue()
ingestion_calls = [
call
for call in logger.async_httpx_client.post.call_args_list
if "dataCollectionRules" in call.kwargs["url"]
]
assert len(ingestion_calls) == 2
standard_call, audit_call = ingestion_calls
assert "Custom-LiteLLM-Standard" in standard_call.kwargs["url"]
assert json.loads(standard_call.kwargs["data"].decode("utf-8")) == [
standard_payload
]
assert "Custom-LiteLLM-Audit" in audit_call.kwargs["url"]
assert json.loads(audit_call.kwargs["data"].decode("utf-8")) == [audit_log]

View file

@ -0,0 +1,159 @@
import logging
import sys
import pytest
from prometheus_client import REGISTRY
import litellm
from litellm.integrations.prometheus import PrometheusLogger
def _clear_prometheus_registry() -> None:
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
REGISTRY.unregister(collector)
def _create_prometheus_logger_with_custom_labels(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
litellm,
"custom_prometheus_metadata_labels",
["metadata.department", "metadata.environment"],
)
_clear_prometheus_registry()
return PrometheusLogger()
def _standard_logging_payload_with_requester_metadata() -> dict:
return {
"model_id": "model-123",
"model_group": "gpt-4o-mini",
"api_base": "https://api.openai.com",
"custom_llm_provider": "openai",
"metadata": {
"user_api_key_hash": "test-hash",
"user_api_key_alias": "test-alias",
"user_api_key_team_id": "test-team",
"user_api_key_team_alias": "test-team-alias",
"user_api_key_user_id": "test-user",
"user_api_key_user_email": "test@example.com",
"user_api_key_org_id": None,
"requester_metadata": {
"department": "engineering",
"environment": "production",
},
"user_api_key_auth_metadata": None,
"spend_logs_metadata": None,
},
"request_tags": [],
"completion_tokens": 0,
"total_tokens": 0,
"response_cost": 0,
}
def _metric_samples(metric_name: str):
return [
sample
for metric in REGISTRY.collect()
for sample in metric.samples
if sample.name == metric_name
]
@pytest.mark.asyncio
async def test_async_log_failure_event_accepts_custom_metadata_labels(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
):
prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch)
kwargs = {
"model": "gpt-4o-mini",
"litellm_params": {
"metadata": {
"user_api_key_end_user_id": "test-end-user",
}
},
"standard_logging_object": _standard_logging_payload_with_requester_metadata(),
}
with caplog.at_level(logging.ERROR):
await prometheus_logger.async_log_failure_event(
kwargs=kwargs,
response_obj=None,
start_time=None,
end_time=None,
)
assert "Incorrect label count" not in caplog.text
samples = _metric_samples("litellm_llm_api_failed_requests_metric_total")
assert any(
sample.labels.get("metadata_department") == "engineering"
and sample.labels.get("metadata_environment") == "production"
for sample in samples
)
def test_virtual_key_rate_limit_metrics_accept_custom_metadata_labels(
monkeypatch: pytest.MonkeyPatch,
):
prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch)
metadata = {
"model_group": "gpt-4o-mini",
"litellm-key-remaining-requests-gpt-4o-mini": 3,
"litellm-key-remaining-tokens-gpt-4o-mini": 200,
}
kwargs = {
"litellm_params": {
"metadata": metadata,
},
"standard_logging_object": _standard_logging_payload_with_requester_metadata(),
}
prometheus_logger._set_virtual_key_rate_limit_metrics(
user_api_key="test-hash",
user_api_key_alias="test-alias",
kwargs=kwargs,
metadata=metadata,
model_id="model-123",
)
samples = _metric_samples("litellm_remaining_api_key_requests_for_model")
assert any(
sample.labels.get("metadata_department") == "engineering"
and sample.labels.get("metadata_environment") == "production"
and sample.value == 3
for sample in samples
)
def test_virtual_key_rate_limit_metrics_preserve_zero_remaining_values(
monkeypatch: pytest.MonkeyPatch,
):
prometheus_logger = _create_prometheus_logger_with_custom_labels(monkeypatch)
metadata = {
"model_group": "gpt-4o-mini",
"litellm-key-remaining-requests-gpt-4o-mini": 0,
"litellm-key-remaining-tokens-gpt-4o-mini": 0,
}
kwargs = {
"litellm_params": {
"metadata": metadata,
},
"standard_logging_object": _standard_logging_payload_with_requester_metadata(),
}
prometheus_logger._set_virtual_key_rate_limit_metrics(
user_api_key="test-hash",
user_api_key_alias="test-alias",
kwargs=kwargs,
metadata=metadata,
model_id="model-123",
)
request_samples = _metric_samples("litellm_remaining_api_key_requests_for_model")
token_samples = _metric_samples("litellm_remaining_api_key_tokens_for_model")
assert any(sample.value == 0 for sample in request_samples)
assert any(sample.value == 0 for sample in token_samples)
assert not any(sample.value == sys.maxsize for sample in request_samples)
assert not any(sample.value == sys.maxsize for sample in token_samples)

View file

@ -0,0 +1,181 @@
from time import monotonic
import pytest
from prometheus_client import REGISTRY
import litellm
from litellm.integrations.prometheus import PrometheusLogger
from litellm.integrations.prometheus_helpers import bounded_prometheus_series_tracker
from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import (
BoundedPrometheusSeriesTracker,
)
from litellm.types.integrations.prometheus import UserAPIKeyLabelValues
@pytest.fixture(autouse=True)
def cleanup_prometheus_registry():
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
try:
REGISTRY.unregister(collector)
except Exception:
pass
old_enable_end_user = litellm.enable_end_user_cost_tracking_prometheus_only
old_metrics_config = litellm.prometheus_metrics_config
old_max_series = litellm.prometheus_end_user_metrics_max_series_per_metric
old_ttl_seconds = litellm.prometheus_end_user_metrics_ttl_seconds
old_cleanup_interval_seconds = (
litellm.prometheus_end_user_metrics_cleanup_interval_seconds
)
yield
litellm.enable_end_user_cost_tracking_prometheus_only = old_enable_end_user
litellm.prometheus_metrics_config = old_metrics_config
litellm.prometheus_end_user_metrics_max_series_per_metric = old_max_series
litellm.prometheus_end_user_metrics_ttl_seconds = old_ttl_seconds
litellm.prometheus_end_user_metrics_cleanup_interval_seconds = (
old_cleanup_interval_seconds
)
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
try:
REGISTRY.unregister(collector)
except Exception:
pass
def test_prometheus_end_user_series_are_capped_per_metric():
litellm.enable_end_user_cost_tracking_prometheus_only = True
litellm.prometheus_metrics_config = [
{
"group": "end-user-spend",
"metrics": ["litellm_spend_metric"],
"include_labels": ["end_user"],
}
]
litellm.prometheus_end_user_metrics_max_series_per_metric = 3
litellm.prometheus_end_user_metrics_ttl_seconds = None
logger = PrometheusLogger()
for index in range(6):
PrometheusLogger._inc_labeled_counter(
logger,
logger.litellm_spend_metric,
"litellm_spend_metric",
UserAPIKeyLabelValues(end_user=f"end-user-{index}"),
amount=0.01,
)
assert len(logger.litellm_spend_metric._metrics) == 3
assert set(logger.litellm_spend_metric._metrics) == {
("end-user-3",),
("end-user-4",),
("end-user-5",),
}
def test_bounded_prometheus_series_tracker_is_label_agnostic():
class FakeMetric:
def __init__(self):
self.removed_label_values = []
def remove(self, *label_values):
self.removed_label_values.append(label_values)
metric = FakeMetric()
tracker = BoundedPrometheusSeriesTracker()
for index in range(4):
tracker.track_series(
metric=metric,
metric_name="generic_metric",
label_values=(f"route-{index}", "200"),
max_series=2,
ttl_seconds=None,
cleanup_interval_seconds=60.0,
)
assert metric.removed_label_values == [
("route-0", "200"),
("route-1", "200"),
]
def test_bounded_prometheus_series_tracker_treats_zero_max_as_unlimited():
# A misconfigured ``max_series=0`` must not silently evict every emission.
class FakeMetric:
def __init__(self):
self.removed_label_values = []
def remove(self, *label_values):
self.removed_label_values.append(label_values)
metric = FakeMetric()
tracker = BoundedPrometheusSeriesTracker()
for index in range(3):
tracker.track_series(
metric=metric,
metric_name="generic_metric",
label_values=(f"end-user-{index}",),
max_series=0,
ttl_seconds=None,
cleanup_interval_seconds=60.0,
)
assert metric.removed_label_values == []
def test_prometheus_end_user_series_expire_by_ttl(monkeypatch):
litellm.enable_end_user_cost_tracking_prometheus_only = True
litellm.prometheus_metrics_config = [
{
"group": "end-user-spend",
"metrics": ["litellm_spend_metric"],
"include_labels": ["end_user"],
}
]
litellm.prometheus_end_user_metrics_max_series_per_metric = None
litellm.prometheus_end_user_metrics_ttl_seconds = 10.0
litellm.prometheus_end_user_metrics_cleanup_interval_seconds = 0.0
logger = PrometheusLogger()
current_time = [monotonic()]
monkeypatch.setattr(
bounded_prometheus_series_tracker.time,
"monotonic",
lambda: current_time[0],
)
PrometheusLogger._inc_labeled_counter(
logger,
logger.litellm_spend_metric,
"litellm_spend_metric",
UserAPIKeyLabelValues(end_user="stale-end-user"),
amount=0.01,
)
current_time[0] += 11.0
PrometheusLogger._inc_labeled_counter(
logger,
logger.litellm_spend_metric,
"litellm_spend_metric",
UserAPIKeyLabelValues(end_user="fresh-end-user"),
amount=0.01,
)
assert set(logger.litellm_spend_metric._metrics) == {("fresh-end-user",)}
def test_prometheus_end_user_not_tracked_by_default():
litellm.enable_end_user_cost_tracking_prometheus_only = None
labels = PrometheusLogger().get_labels_for_metric("litellm_spend_metric")
assert "end_user" in labels
label_values = UserAPIKeyLabelValues(end_user="not-exported")
from litellm.integrations.prometheus import prometheus_label_factory
prometheus_labels = prometheus_label_factory(labels, label_values)
assert prometheus_labels["end_user"] is None

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