mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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
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:
commit
fa81017e12
145 changed files with 14804 additions and 2136 deletions
|
|
@ -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
|
||||
|
|
|
|||
1
.github/workflows/test-unit-proxy-db.yml
vendored
1
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -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
3
.gitignore
vendored
|
|
@ -100,4 +100,5 @@ STABILIZATION_TODO.md
|
|||
**/playwright-report
|
||||
**/*.storageState.json
|
||||
**/coverage
|
||||
test-config
|
||||
test-config
|
||||
.vscode
|
||||
20
Dockerfile
20
Dockerfile
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 }}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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: {}
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
BIN
docs/images/local-testing/hosted-vllm-custom-tool-local-test.png
Normal file
BIN
docs/images/local-testing/hosted-vllm-custom-tool-local-test.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 82 KiB |
6
litellm-js/spend-logs/package-lock.json
generated
6
litellm-js/spend-logs/package-lock.json
generated
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -5255,7 +5255,6 @@ class BaseLLMHTTPHandler:
|
|||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"OpenAI-Beta": "realtime=v1",
|
||||
}
|
||||
|
||||
if extra_headers:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
0
litellm/llms/nvidia_riva/__init__.py
Normal file
0
litellm/llms/nvidia_riva/__init__.py
Normal file
0
litellm/llms/nvidia_riva/audio_transcription/__init__.py
Normal file
0
litellm/llms/nvidia_riva/audio_transcription/__init__.py
Normal file
232
litellm/llms/nvidia_riva/audio_transcription/audio_utils.py
Normal file
232
litellm/llms/nvidia_riva/audio_transcription/audio_utils.py
Normal 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)
|
||||
444
litellm/llms/nvidia_riva/audio_transcription/handler.py
Normal file
444
litellm/llms/nvidia_riva/audio_transcription/handler.py
Normal 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
|
||||
284
litellm/llms/nvidia_riva/audio_transcription/transformation.py
Normal file
284
litellm/llms/nvidia_riva/audio_transcription/transformation.py
Normal 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
|
||||
92
litellm/llms/nvidia_riva/common_utils.py
Normal file
92
litellm/llms/nvidia_riva/common_utils.py
Normal 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)
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
196
litellm/proxy/_experimental/mcp_server/auth/token_exchange.py
Normal file
196
litellm/proxy/_experimental/mcp_server/auth/token_exchange.py
Normal 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()
|
||||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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] = (
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
121
litellm/proxy/middleware/request_size_limit_middleware.py
Normal file
121
litellm/proxy/middleware/request_size_limit_middleware.py
Normal 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})
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
276
scripts/mock_bedrock_passthrough_target.py
Normal file
276
scripts/mock_bedrock_passthrough_target.py
Normal 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
82
scripts/tpm_headline_test.sh
Executable 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 ~1–2 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
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
29
tests/litellm/test_sambanova_model_metadata.py
Normal file
29
tests/litellm/test_sambanova_model_metadata.py
Normal 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"
|
||||
|
|
@ -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",
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
create→test→delete 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):
|
||||
|
|
|
|||
220
tests/llm_translation/test_vcr_filters.py
Normal file
220
tests/llm_translation/test_vcr_filters.py
Normal 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
|
||||
|
|
@ -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="")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
135
tests/proxy_unit_tests/test_request_size_limit_middleware.py
Normal file
135
tests/proxy_unit_tests/test_request_size_limit_middleware.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
Loading…
Add table
Reference in a new issue