mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 10:21:52 +00:00
Merge pull request #29861 from BerriAI/litellm_internal_staging
chore(ci): promote internal staging to main
This commit is contained in:
yuneng-jiang
+213
-39
@@ -111,6 +111,28 @@ commands:
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:6379
|
||||
timeout: "60"
|
||||
start_openai_record_replay_proxy:
|
||||
description: "Start the record/replay proxy (tests/_openai_record_replay_proxy.py) on host port 8090 and wait until healthy. Models whose api_base points here replay recorded provider responses, so the E2E run neither pays for nor depends on the live provider. The default upstream is OpenAI; a non-OpenAI model must point its api_base at /__recorder_upstream/<host>/ so the recorder forwards there instead of defaulting to OpenAI. Run after uv deps are synced."
|
||||
steps:
|
||||
- run:
|
||||
name: Start record/replay proxy
|
||||
background: true
|
||||
command: |
|
||||
CASSETTE_REDIS_URL="$CASSETTE_REDIS_URL" \
|
||||
RECORDER_UPSTREAM_BASE_URL="https://api.openai.com" \
|
||||
uv run --no-sync python tests/_openai_record_replay_proxy.py --host 0.0.0.0 --port 8090
|
||||
- run:
|
||||
name: Wait for record/replay proxy
|
||||
command: |
|
||||
for i in $(seq 1 30); do
|
||||
if curl -sf http://localhost:8090/__recorder_health >/dev/null 2>&1; then
|
||||
echo "record/replay proxy is up"
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "record/replay proxy did not become ready" >&2
|
||||
exit 1
|
||||
setup_litellm_enterprise_pip:
|
||||
steps:
|
||||
- run:
|
||||
@@ -182,7 +204,14 @@ jobs:
|
||||
- run:
|
||||
name: Run Windows-specific test
|
||||
command: |
|
||||
uv run --no-sync python -m pytest tests/windows_tests/test_litellm_on_windows.py -v
|
||||
uv run --no-sync python -m pytest tests/windows_tests/ -v
|
||||
- run:
|
||||
name: Guard against MAX_PATH-busting packaged wheel paths
|
||||
environment:
|
||||
UV_HTTP_TIMEOUT: "300"
|
||||
command: |
|
||||
uv build --wheel --out-dir dist
|
||||
uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py
|
||||
|
||||
local_testing_part1:
|
||||
docker:
|
||||
@@ -228,7 +257,7 @@ jobs:
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--split-by=timings \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml \
|
||||
@@ -242,8 +271,15 @@ jobs:
|
||||
- run:
|
||||
name: Rename the coverage files
|
||||
command: |
|
||||
mv coverage.xml local_testing_part1_coverage.xml
|
||||
mv .coverage local_testing_part1_coverage
|
||||
# When CI reruns only the failed tests, a parallel node can receive
|
||||
# zero tests and pytest never writes coverage. Emit empty placeholders
|
||||
# so persist_to_workspace and the downstream coverage combine stay green.
|
||||
if [ -f coverage.xml ]; then
|
||||
mv coverage.xml local_testing_part1_coverage.xml
|
||||
mv .coverage local_testing_part1_coverage
|
||||
else
|
||||
touch local_testing_part1_coverage.xml local_testing_part1_coverage
|
||||
fi
|
||||
|
||||
# Store test results
|
||||
- store_test_results:
|
||||
@@ -293,7 +329,7 @@ jobs:
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--split-by=timings \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml \
|
||||
@@ -307,8 +343,15 @@ jobs:
|
||||
- run:
|
||||
name: Rename the coverage files
|
||||
command: |
|
||||
mv coverage.xml local_testing_part2_coverage.xml
|
||||
mv .coverage local_testing_part2_coverage
|
||||
# When CI reruns only the failed tests, a parallel node can receive
|
||||
# zero tests and pytest never writes coverage. Emit empty placeholders
|
||||
# so persist_to_workspace and the downstream coverage combine stay green.
|
||||
if [ -f coverage.xml ]; then
|
||||
mv coverage.xml local_testing_part2_coverage.xml
|
||||
mv .coverage local_testing_part2_coverage
|
||||
else
|
||||
touch local_testing_part2_coverage.xml local_testing_part2_coverage
|
||||
fi
|
||||
|
||||
# Store test results
|
||||
- store_test_results:
|
||||
@@ -356,7 +399,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
@@ -409,7 +452,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/proxy_admin_ui_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -431,6 +474,120 @@ jobs:
|
||||
- auth_ui_unit_tests_coverage.xml
|
||||
- auth_ui_unit_tests_coverage
|
||||
|
||||
proxy_behavior_tests:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: litellm_test
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Seed DB schema via prisma db push
|
||||
command: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
- run:
|
||||
name: Generate Prisma Client
|
||||
command: uv run --no-sync python -m prisma generate
|
||||
- run:
|
||||
name: Run proxy management behavior tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
uv run --no-sync python -m pytest tests/proxy_behavior \
|
||||
-v --junitxml=test-results/junit.xml --durations=10
|
||||
no_output_timeout: 15m
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
|
||||
proxy_security_tests:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: litellm_test
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Seed DB schema via prisma db push
|
||||
command: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
- run:
|
||||
name: Generate Prisma Client
|
||||
command: uv run --no-sync python -m prisma generate
|
||||
- run:
|
||||
name: Run proxy security tests
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
uv run --no-sync python -m pytest tests/proxy_security_tests \
|
||||
-v --junitxml=test-results/junit.xml --durations=10
|
||||
no_output_timeout: 15m
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
|
||||
schema_migration_check:
|
||||
docker:
|
||||
- *python312_image
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_DB: litellm_test
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
# An empty database; the test applies every committed migration itself.
|
||||
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "60"
|
||||
- run:
|
||||
name: Generate Prisma Client
|
||||
command: uv run --no-sync python -m prisma generate
|
||||
- run:
|
||||
name: Check schema.prisma is in sync with committed migrations
|
||||
command: |
|
||||
mkdir -p test-results
|
||||
uv run --no-sync python -m pytest tests/proxy_migration_tests \
|
||||
-v --junitxml=test-results/junit.xml --durations=10
|
||||
no_output_timeout: 15m
|
||||
- store_test_results:
|
||||
path: test-results
|
||||
|
||||
litellm_router_testing: # Runs all tests with the "router" keyword
|
||||
docker:
|
||||
- *python312_image
|
||||
@@ -457,12 +614,17 @@ jobs:
|
||||
- run:
|
||||
name: Run tests
|
||||
command: |
|
||||
# On a "rerun failed tests" build a parallel node can receive no
|
||||
# tests, so the test command never creates test-results. Pre-create it
|
||||
# so store_test_results doesn't fail the node on a missing path.
|
||||
mkdir -p test-results
|
||||
|
||||
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
|
||||
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--split-by=timings \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v \
|
||||
-k 'router' \
|
||||
-n 4 \
|
||||
@@ -504,7 +666,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/router_unit_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -547,7 +709,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
@@ -589,7 +751,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/llm_translation/**/test_*.py" | grep -v "^tests/llm_translation/realtime/")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=20 \
|
||||
@@ -625,7 +787,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/llm_translation/realtime/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -668,7 +830,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/agent_tests/**/test_*.py" | grep -v "^tests/agent_tests/local_only_agent_tests/")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -710,7 +872,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/guardrails_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -754,7 +916,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/unified_google_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -805,7 +967,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/llm_responses_api_testing/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
@@ -836,7 +998,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -878,7 +1040,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/search_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -922,7 +1084,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/enterprise/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit-enterprise.xml \
|
||||
--durations=10 \
|
||||
@@ -952,7 +1114,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/batches_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -994,7 +1156,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/litellm_utils_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -1037,7 +1199,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/pass_through_unit_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -1080,7 +1242,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/image_gen_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5 \
|
||||
@@ -1112,7 +1274,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/logging_callback_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
-n 4 \
|
||||
@@ -1155,7 +1317,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/audio_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -1206,7 +1368,7 @@ jobs:
|
||||
tests/local_testing/test_router_utils.py)
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--cov=./litellm --cov-report=xml \
|
||||
--junitxml=test-results/junit.xml \
|
||||
@@ -1456,7 +1618,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit-2.xml \
|
||||
--durations=5"
|
||||
@@ -1485,6 +1647,7 @@ jobs:
|
||||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker tag litellm-docker-database:ci my-app:latest
|
||||
- start_openai_record_replay_proxy
|
||||
- run:
|
||||
name: Run Docker container
|
||||
command: |
|
||||
@@ -1515,6 +1678,7 @@ jobs:
|
||||
-e LANGFUSE_PROJECT2_PUBLIC=$LANGFUSE_PROJECT2_PUBLIC \
|
||||
-e LANGFUSE_PROJECT1_SECRET=$LANGFUSE_PROJECT1_SECRET \
|
||||
-e LANGFUSE_PROJECT2_SECRET=$LANGFUSE_PROJECT2_SECRET \
|
||||
-e RECORDER_OPENAI_BASE_URL=http://host.docker.internal:8090/v1 \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
--name my-app \
|
||||
-v $(pwd)/proxy_server_config.yaml:/app/config.yaml \
|
||||
@@ -1539,7 +1703,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-s -v -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
-n 4 \
|
||||
@@ -1622,7 +1786,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/openai_endpoints_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-s -vv \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
@@ -1652,6 +1816,7 @@ jobs:
|
||||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker images | grep litellm-docker-database
|
||||
- start_openai_record_replay_proxy
|
||||
- run:
|
||||
name: Run Docker container
|
||||
# intentionally give bad redis credentials here
|
||||
@@ -1675,6 +1840,7 @@ jobs:
|
||||
-e DD_SITE=$DD_SITE \
|
||||
-e AWS_REGION_NAME=$AWS_REGION_NAME \
|
||||
-e COHERE_API_KEY=$COHERE_API_KEY \
|
||||
-e RECORDER_COHERE_BASE_URL=http://host.docker.internal:8090/__recorder_upstream/api.cohere.com \
|
||||
-e GCS_FLUSH_INTERVAL="1" \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
--name my-app \
|
||||
@@ -1698,7 +1864,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/otel_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
@@ -1748,7 +1914,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit-2.xml \
|
||||
--durations=5"
|
||||
@@ -1824,7 +1990,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/spend_tracking_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
@@ -1922,7 +2088,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/multi_instance_e2e_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
@@ -1985,7 +2151,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/store_model_in_db_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
@@ -2065,7 +2231,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x \
|
||||
--junitxml=test-results/junit-2.xml \
|
||||
--durations=5"
|
||||
@@ -2209,7 +2375,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/pass_through_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-v -x \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
@@ -2240,6 +2406,7 @@ jobs:
|
||||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker images | grep litellm-docker-database
|
||||
- start_openai_record_replay_proxy
|
||||
- run:
|
||||
name: Run Docker container with test config
|
||||
command: |
|
||||
@@ -2248,6 +2415,7 @@ jobs:
|
||||
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
|
||||
-e RECORDER_ANTHROPIC_BASE_URL=http://host.docker.internal:8090/__recorder_upstream/api.anthropic.com \
|
||||
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
|
||||
-e AWS_REGION_NAME="us-east-1" \
|
||||
@@ -2275,7 +2443,7 @@ jobs:
|
||||
TEST_FILES=$(circleci tests glob "tests/proxy_e2e_anthropic_messages_tests/**/test_*.py")
|
||||
echo "$TEST_FILES" | circleci tests run \
|
||||
--verbose \
|
||||
--command="awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
|
||||
-vv -x -s \
|
||||
--junitxml=test-results/junit.xml \
|
||||
--durations=5"
|
||||
@@ -2617,6 +2785,12 @@ workflows:
|
||||
filters: *main_branches
|
||||
- auth_ui_unit_tests:
|
||||
filters: *main_branches
|
||||
- proxy_behavior_tests:
|
||||
filters: *main_branches
|
||||
- proxy_security_tests:
|
||||
filters: *main_branches
|
||||
- schema_migration_check:
|
||||
filters: *main_branches
|
||||
- build_docker_database_image:
|
||||
filters: *main_branches
|
||||
- e2e_ui_testing:
|
||||
|
||||
@@ -8,3 +8,6 @@
|
||||
|
||||
# Update pydantic code to fix warnings (GH-3600)
|
||||
876840e9957bc7e9f7d6a2b58c4d7c53dad16481
|
||||
|
||||
# style(ui): run prettier --write across the dashboard (#29622)
|
||||
7edf3a9cb55548b143df1692f4ed7c4681d7fcf7
|
||||
|
||||
+2
-1
@@ -1 +1,2 @@
|
||||
*.ipynb linguist-vendored
|
||||
*.ipynb linguist-vendored
|
||||
ui/litellm-dashboard/src/lib/http/schema.d.ts linguist-generated
|
||||
@@ -27,6 +27,11 @@ on:
|
||||
required: false
|
||||
type: number
|
||||
default: 10
|
||||
dist:
|
||||
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
|
||||
required: false
|
||||
type: string
|
||||
default: "loadscope"
|
||||
artifact-name:
|
||||
description: "Unique name for the coverage artifact (must be unique per run)"
|
||||
required: true
|
||||
@@ -82,18 +87,31 @@ jobs:
|
||||
MAX_FAILURES: ${{ inputs.max-failures }}
|
||||
WORKERS: ${{ inputs.workers }}
|
||||
RERUNS: ${{ inputs.reruns }}
|
||||
DIST: ${{ inputs.dist }}
|
||||
run: |
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
-n "${WORKERS}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--dist=loadscope \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
else
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
-n "${WORKERS}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--dist="${DIST}" \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
fi
|
||||
|
||||
- name: Save coverage report
|
||||
if: always()
|
||||
|
||||
@@ -1,190 +0,0 @@
|
||||
name: _Unit Test Services Base (Reusable)
|
||||
|
||||
on:
|
||||
workflow_call:
|
||||
inputs:
|
||||
test-path:
|
||||
description: "Pytest path(s) to run"
|
||||
required: true
|
||||
type: string
|
||||
workers:
|
||||
description: "Number of pytest-xdist workers (0 = no parallelism)"
|
||||
required: false
|
||||
type: number
|
||||
default: 2
|
||||
reruns:
|
||||
description: "Number of reruns for flaky tests"
|
||||
required: false
|
||||
type: number
|
||||
default: 2
|
||||
timeout-minutes:
|
||||
description: "Job timeout in minutes"
|
||||
required: false
|
||||
type: number
|
||||
default: 20
|
||||
max-failures:
|
||||
description: "Stop after this many failures"
|
||||
required: false
|
||||
type: number
|
||||
default: 10
|
||||
enable-postgres:
|
||||
description: "Start a local Postgres service container and run Prisma migrations"
|
||||
required: false
|
||||
type: boolean
|
||||
default: false
|
||||
dist:
|
||||
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
|
||||
required: false
|
||||
type: string
|
||||
default: "loadscope"
|
||||
artifact-name:
|
||||
description: "Unique name for the coverage artifact (must be unique per run)"
|
||||
required: false
|
||||
type: string
|
||||
default: "run"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
# The postgres service container below is spawned per-job on localhost and
|
||||
# destroyed with the job. Nothing outside the runner can reach it. The
|
||||
# user/password/database here are not secrets — they're bootstrap values
|
||||
# for a throwaway container — so we hardcode them instead of attaching
|
||||
# every matrix shard to a GHA environment just to read three "secrets"
|
||||
# (which also produces a "temporarily deployed to …" notification on the
|
||||
# PR timeline per shard per push).
|
||||
jobs:
|
||||
run:
|
||||
name: Run tests
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: ${{ inputs.timeout-minutes }}
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres@sha256:705a5d5b5836f3fcba0d02c4d281e6a7dd9ed2dd4078640f08a1e1e9896e097d # postgres:14
|
||||
env:
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: litellm
|
||||
POSTGRES_DB: litellm_test
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready"
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-services-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-services-
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Run Prisma migrations
|
||||
if: ${{ inputs.enable-postgres }}
|
||||
env:
|
||||
DATABASE_URL: "postgresql://litellm:litellm@localhost:5432/litellm_test"
|
||||
run: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
|
||||
- name: Run tests
|
||||
env:
|
||||
TEST_PATH: ${{ inputs.test-path }}
|
||||
MAX_FAILURES: ${{ inputs.max-failures }}
|
||||
WORKERS: ${{ inputs.workers }}
|
||||
RERUNS: ${{ inputs.reruns }}
|
||||
DIST: ${{ inputs.dist }}
|
||||
DATABASE_URL: ${{ inputs.enable-postgres && 'postgresql://litellm:litellm@localhost:5432/litellm_test' || '' }}
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
else
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
-n "${WORKERS}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--dist="${DIST}" \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
fi
|
||||
|
||||
- name: Save coverage report
|
||||
if: always()
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
path: coverage.xml
|
||||
retention-days: 1
|
||||
|
||||
upload-coverage:
|
||||
name: Upload coverage to Codecov
|
||||
needs: run
|
||||
if: always()
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Download coverage report
|
||||
uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e # v4.2.1
|
||||
with:
|
||||
pattern: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
path: coverage-reports
|
||||
merge-multiple: true
|
||||
|
||||
- name: Upload to Codecov
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
with:
|
||||
use_oidc: true
|
||||
directory: coverage-reports
|
||||
root_dir: ${{ github.workspace }}
|
||||
flags: ${{ inputs.artifact-name }}
|
||||
fail_ci_if_error: false
|
||||
@@ -0,0 +1,84 @@
|
||||
name: Check UI API Types Sync
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
paths:
|
||||
- "litellm/proxy/**"
|
||||
- "litellm/types/**"
|
||||
- "ui/litellm-dashboard/src/lib/http/schema.d.ts"
|
||||
- "ui/litellm-dashboard/scripts/gen-api-types.mjs"
|
||||
- "ui/litellm-dashboard/package.json"
|
||||
- "ui/litellm-dashboard/package-lock.json"
|
||||
- ".github/workflows/check-ui-api-types.yml"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
check-sync:
|
||||
name: Verify schema.d.ts matches the proxy OpenAPI spec
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Install backend dependencies
|
||||
run: uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "npm"
|
||||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dashboard dependencies
|
||||
working-directory: ui/litellm-dashboard
|
||||
run: npm ci
|
||||
|
||||
- name: Regenerate types from the live spec
|
||||
working-directory: ui/litellm-dashboard
|
||||
env:
|
||||
LITELLM_PYTHON: "uv run --no-sync python"
|
||||
run: npm run gen:api
|
||||
|
||||
- name: Fail if types are stale
|
||||
run: |
|
||||
if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
|
||||
echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec."
|
||||
echo ""
|
||||
echo "A backend route or model changed without regenerating the dashboard types."
|
||||
echo "To fix, run from ui/litellm-dashboard:"
|
||||
echo " npm run gen:api"
|
||||
echo "then commit the updated src/lib/http/schema.d.ts."
|
||||
exit 1
|
||||
fi
|
||||
echo "schema.d.ts is in sync with the proxy OpenAPI spec."
|
||||
@@ -63,3 +63,28 @@ jobs:
|
||||
sha: commitHash,
|
||||
});
|
||||
core.info(`Created branch ${branchName} at ${commitHash}`);
|
||||
|
||||
- name: Create stable line branch
|
||||
env:
|
||||
TAG: ${{ inputs.tag }}
|
||||
COMMIT_HASH: ${{ inputs.commit_hash }}
|
||||
uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
|
||||
with:
|
||||
script: |
|
||||
const tag = process.env.TAG;
|
||||
const commitHash = process.env.COMMIT_HASH;
|
||||
|
||||
const match = tag.match(/^v?(\d+)\.(\d+)\.0$/);
|
||||
if (!match) {
|
||||
core.info(`Tag ${tag} is not the X.Y.0 stable opener; skipping stable line branch`);
|
||||
return;
|
||||
}
|
||||
const lineBranch = `stable/${match[1]}.${match[2]}.x`;
|
||||
|
||||
await github.rest.git.createRef({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
ref: `refs/heads/${lineBranch}`,
|
||||
sha: commitHash,
|
||||
});
|
||||
core.info(`Created branch ${lineBranch} at ${commitHash}`);
|
||||
|
||||
@@ -36,3 +36,79 @@ jobs:
|
||||
|
||||
- name: Build
|
||||
run: npm run build
|
||||
|
||||
frontend-lint:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 8
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ui/litellm-dashboard
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
|
||||
- name: Collect changed files
|
||||
id: changed
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
: > "$RUNNER_TEMP/prettier_files.txt"
|
||||
: > "$RUNNER_TEMP/eslint_files.txt"
|
||||
while IFS= read -r f; do
|
||||
[ -f "$f" ] || continue
|
||||
case "$f" in
|
||||
*.js | *.jsx | *.ts | *.tsx | *.mjs | *.cjs)
|
||||
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt"
|
||||
printf '%s\n' "$f" >> "$RUNNER_TEMP/eslint_files.txt" ;;
|
||||
*.json | *.css | *.scss | *.md | *.mdx | *.yml | *.yaml | *.html)
|
||||
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt" ;;
|
||||
esac
|
||||
done < <(git diff --name-only --diff-filter=ACMR --relative "$BASE_SHA"...HEAD -- .)
|
||||
if [ -s "$RUNNER_TEMP/prettier_files.txt" ] || [ -s "$RUNNER_TEMP/eslint_files.txt" ]; then
|
||||
echo "has_files=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "has_files=false" >> "$GITHUB_OUTPUT"
|
||||
echo "No lintable UI files changed in this PR; nothing to check."
|
||||
fi
|
||||
|
||||
- name: Setup Node.js
|
||||
if: steps.changed.outputs.has_files == 'true'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "npm"
|
||||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dependencies
|
||||
if: steps.changed.outputs.has_files == 'true'
|
||||
run: npm ci
|
||||
|
||||
- name: Lint changed files (prettier + eslint)
|
||||
if: steps.changed.outputs.has_files == 'true'
|
||||
run: |
|
||||
prettier_files=()
|
||||
eslint_files=()
|
||||
while IFS= read -r f; do prettier_files+=("$f"); done < "$RUNNER_TEMP/prettier_files.txt"
|
||||
while IFS= read -r f; do eslint_files+=("$f"); done < "$RUNNER_TEMP/eslint_files.txt"
|
||||
status=0
|
||||
if [ ${#prettier_files[@]} -gt 0 ]; then
|
||||
echo "::group::Prettier (${#prettier_files[@]} files)"
|
||||
npx prettier --check "${prettier_files[@]}" || { status=1; echo "::error::Unformatted files. Fix with: npm run format"; }
|
||||
echo "::endgroup::"
|
||||
fi
|
||||
if [ ${#eslint_files[@]} -gt 0 ]; then
|
||||
echo "::group::ESLint (${#eslint_files[@]} files)"
|
||||
npx eslint --no-warn-ignored --pass-on-unpruned-suppressions "${eslint_files[@]}" || status=1
|
||||
echo "::endgroup::"
|
||||
fi
|
||||
exit $status
|
||||
|
||||
- name: Check lint budgets
|
||||
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
|
||||
run: |
|
||||
npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true
|
||||
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
name: "Unit Tests: Proxy DB Operations"
|
||||
|
||||
# Uses DATABASE_URL secret — only runs on trusted branches, not PRs.
|
||||
on:
|
||||
push:
|
||||
branches: [main, "litellm_**"]
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
@@ -30,9 +31,6 @@ concurrency:
|
||||
# xdist balances its 188 parametrized cases across workers instead of
|
||||
# pinning the whole file to one worker (the default --dist=loadscope
|
||||
# behavior for single-file targets).
|
||||
# * test_db_schema_migration.py is isolated because one test in it
|
||||
# (test_aaaasschema_migration_check) takes ~170s — by itself it
|
||||
# determines the shard's wall-clock floor.
|
||||
jobs:
|
||||
# Fast guard — fails the workflow if a test_*.py file under
|
||||
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
|
||||
@@ -166,18 +164,6 @@ jobs:
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
|
||||
# ---- db-and-spend: isolate the 170s schema-migration test ----
|
||||
# test_db_schema_migration.py has exactly one test, and that test
|
||||
# is mostly waiting on `prisma migrate deploy` / `prisma migrate
|
||||
# diff` subprocesses (~170s). It does no CPU-bound Python work
|
||||
# inside the test. Running with workers=0 (serial, no xdist)
|
||||
# skips the 4-worker cold-start cost we'd otherwise pay for a
|
||||
# single test, saving ~4 minutes of wall-clock.
|
||||
- test-group: schema-migration
|
||||
test-path: "tests/proxy_unit_tests/test_db_schema_migration.py"
|
||||
workers: 0
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
- test-group: db-and-spend
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_prisma_client_backoff_retry.py
|
||||
@@ -232,12 +218,11 @@ jobs:
|
||||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: ${{ matrix.test-path }}
|
||||
workers: ${{ matrix.workers }}
|
||||
reruns: 2
|
||||
timeout-minutes: ${{ matrix.timeout }}
|
||||
enable-postgres: true
|
||||
dist: ${{ matrix.dist }}
|
||||
artifact-name: proxy-db-${{ matrix.test-group }}
|
||||
|
||||
@@ -36,11 +36,13 @@ jobs:
|
||||
tests/test_litellm/proxy/a2a
|
||||
tests/test_litellm/proxy/discovery_endpoints
|
||||
tests/test_litellm/proxy/health_endpoints
|
||||
tests/test_litellm/proxy/shutdown
|
||||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
tests/test_litellm/proxy/utils
|
||||
workers: 2
|
||||
reruns: 2
|
||||
artifact-name: proxy-endpoints
|
||||
|
||||
@@ -1,34 +0,0 @@
|
||||
name: "Unit Tests: Proxy Management-Endpoint Behavior Pinning"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
proxy-mgmt-behavior:
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
with:
|
||||
test-path: tests/proxy_behavior
|
||||
# workers=0 (no xdist): the world seed is a single shared Postgres
|
||||
# state — two xdist workers both call seed_world() and race on the
|
||||
# ``behavior-pin-budget`` row, producing UniqueViolation + cascading
|
||||
# missing-membership FK failures. The whole suite is ~7s sequentially,
|
||||
# so the cost of disabling parallelism here is negligible.
|
||||
workers: 0
|
||||
reruns: 0
|
||||
enable-postgres: true
|
||||
artifact-name: proxy-mgmt-behavior
|
||||
timeout-minutes: 15
|
||||
@@ -1,28 +0,0 @@
|
||||
name: "Unit Tests: Security"
|
||||
|
||||
# Kept push-only (was previously required by DATABASE_URL secret scoping;
|
||||
# now the postgres credentials are ephemeral localhost values but the
|
||||
# push-trigger stays to match the proxy-db workflow cadence).
|
||||
on:
|
||||
push:
|
||||
branches: [main, "litellm_**"]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
security:
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
with:
|
||||
test-path: "tests/proxy_security_tests/"
|
||||
workers: 1
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
enable-postgres: true
|
||||
artifact-name: security
|
||||
@@ -9,12 +9,15 @@ Don't assume that the existing code is correct or the right way of doing things
|
||||
- readable
|
||||
- easy to maintain/change
|
||||
- modern
|
||||
|
||||
In that order of importance
|
||||
|
||||
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
|
||||
|
||||
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
|
||||
|
||||
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
|
||||
|
||||
When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose
|
||||
|
||||
Always use @.github/pull_request_template.md as a guide for your PR body
|
||||
@@ -26,6 +29,8 @@ If you ever make public-facing PR descriptions, comments, issues, commit message
|
||||
- don't use "—". Instead, reach for ";", ".", etc.
|
||||
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
|
||||
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
|
||||
- don't add a trailing "." at the end of paragraphs (just like this file)
|
||||
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
|
||||
|
||||
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
|
||||
|
||||
@@ -37,7 +42,7 @@ When you must use real LLM models to, for example, write e2e tests, write a QA r
|
||||
|
||||
If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
|
||||
|
||||
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
|
||||
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
|
||||
|
||||
When working on a PR, keep the PR description in sync with new commits being made
|
||||
|
||||
@@ -49,22 +54,22 @@ CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bas
|
||||
|
||||
## Think Before Coding
|
||||
|
||||
**Don't assume. Don't hide confusion. Surface tradeoffs.**
|
||||
**Don't assume. Don't hide confusion. Surface tradeoffs**
|
||||
|
||||
Before implementing:
|
||||
- State your assumptions explicitly. If uncertain, ask.
|
||||
- If multiple interpretations exist, present them. Don't pick silently.
|
||||
- If a simpler approach exists, say so. Push back when warranted.
|
||||
- If something is unclear, stop. Name what's confusing. Ask.
|
||||
- State your assumptions explicitly. If uncertain, ask
|
||||
- If multiple interpretations exist, present them. Don't pick silently
|
||||
- If a simpler approach exists, say so. Push back when warranted
|
||||
- If something is unclear, stop. Name what's confusing. Ask
|
||||
|
||||
## Simplicity First
|
||||
|
||||
**Minimum code that solves the problem. Nothing speculative.**
|
||||
**Minimum code that solves the problem. Nothing speculative**
|
||||
|
||||
- No features beyond what was asked.
|
||||
- No abstractions for single-use code.
|
||||
- No "flexibility" or "configurability" that wasn't requested.
|
||||
- No error handling for impossible scenarios.
|
||||
- If you write 200 lines and it could be 50, rewrite it.
|
||||
- No features beyond what was asked
|
||||
- No abstractions for single-use code
|
||||
- No "flexibility" or "configurability" that wasn't requested
|
||||
- No error handling for impossible scenarios
|
||||
- If you write 200 lines and it could be 50, rewrite it
|
||||
|
||||
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify.
|
||||
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify
|
||||
|
||||
@@ -146,7 +146,7 @@ test-unit-proxy-core: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/proxy/auth tests/test_litellm/proxy/client tests/test_litellm/proxy/db tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-proxy-misc: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20
|
||||
$(UV_RUN) pytest tests/test_litellm/proxy/_experimental tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/anthropic_endpoints tests/test_litellm/proxy/common_utils tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/experimental tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/middleware tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/pass_through_endpoints tests/test_litellm/proxy/prompts tests/test_litellm/proxy/public_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/shutdown tests/test_litellm/proxy/spend_tracking tests/test_litellm/proxy/ui_crud_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/test_*.py --tb=short -vv -n 4 --durations=20
|
||||
|
||||
test-unit-integrations: install-test-deps
|
||||
$(UV_RUN) pytest tests/test_litellm/integrations --tb=short -vv -n 4 --durations=20
|
||||
|
||||
@@ -37,7 +37,7 @@
|
||||
</a>
|
||||
</h4>
|
||||
|
||||
<img width="2688" height="1600" alt="Group 7154 (1)" src="https://github.com/user-attachments/assets/c5ee0412-6fb5-4fb6-ab5b-bafae4209ca6" />
|
||||
<img alt="LiteLLM AI Gateway" src="https://github.com/user-attachments/assets/c5ee0412-6fb5-4fb6-ab5b-bafae4209ca6" />
|
||||
|
||||
---
|
||||
|
||||
@@ -407,7 +407,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature
|
||||
### Run in Developer Mode
|
||||
#### Services
|
||||
1. Setup .env file in root
|
||||
2. Run dependant services `docker-compose up db prometheus`
|
||||
2. Run dependent services `docker-compose up db prometheus`
|
||||
|
||||
#### Backend
|
||||
1. (In root) create virtual environment `python -m venv .venv`
|
||||
|
||||
@@ -285,11 +285,31 @@ db:
|
||||
deployStandalone: true
|
||||
|
||||
# Lifecycle hooks for the LiteLLM container
|
||||
#
|
||||
# Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks
|
||||
# the pod NotReady and blocks only until in-flight requests actually finish
|
||||
# (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always
|
||||
# waiting the worst-case duration. The drain runs once (the preStop hook and
|
||||
# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few
|
||||
# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before
|
||||
# SIGKILL.
|
||||
#
|
||||
# /health/drain is off by default; enable it with
|
||||
# general_settings.enable_drain_endpoint: true. The kubelet calls preStop
|
||||
# hooks without proxy credentials, so when the health port is reachable from
|
||||
# other pods (the common case) also set
|
||||
# general_settings.drain_endpoint_token (or the DRAIN_ENDPOINT_TOKEN env
|
||||
# var) and send the same value on the X-Drain-Token header from the hook.
|
||||
# Calls missing/wrong the token get a 401 and have no side effect.
|
||||
# Example:
|
||||
# lifecycle:
|
||||
# preStop:
|
||||
# exec:
|
||||
# command: ["/bin/sh", "-c", "sleep 10"]
|
||||
# httpGet:
|
||||
# path: /health/drain
|
||||
# port: 4000
|
||||
# httpHeaders:
|
||||
# - name: X-Drain-Token
|
||||
# value: <same value as drain_endpoint_token>
|
||||
lifecycle: {}
|
||||
|
||||
# Settings for Bitnami postgresql chart (if db.deployStandalone is true, ignored
|
||||
|
||||
@@ -106,6 +106,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
||||
# Health & ops
|
||||
"/health",
|
||||
"/metrics",
|
||||
"/watsonx"
|
||||
)
|
||||
|
||||
GATEWAY_EXACT_PATHS: frozenset[str] = frozenset(
|
||||
|
||||
@@ -12,9 +12,14 @@ spec:
|
||||
{{- include "litellm.backend.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
{{- with .Values.backend.podAnnotations }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.podAnnotations }}
|
||||
annotations:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
checksum/config: {{ include (print $.Template.BasePath "/gateway/configmap.yaml") . | sha256sum }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.podAnnotations }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
labels:
|
||||
{{- include "litellm.backend.selectorLabels" . | nindent 8 }}
|
||||
@@ -35,7 +40,17 @@ spec:
|
||||
protocol: TCP
|
||||
env:
|
||||
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: CONFIG_FILE_PATH
|
||||
value: /app/config/config.yaml
|
||||
{{- end }}
|
||||
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
volumeMounts:
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
{{- end }}
|
||||
{{- with .Values.backend.livenessProbe }}
|
||||
livenessProbe:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
@@ -46,6 +61,12 @@ spec:
|
||||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.backend.resources | nindent 12 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
volumes:
|
||||
- name: gateway-config
|
||||
configMap:
|
||||
name: {{ include "litellm.gateway.fullname" . }}-config
|
||||
{{- end }}
|
||||
{{- with .Values.backend.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
-- AlterTable: add admin-configured env_vars to MCP server table
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "env_vars" JSONB DEFAULT '[]';
|
||||
|
||||
-- CreateTable: per-user env var values for MCP servers
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MCPUserEnvVars" (
|
||||
"id" TEXT NOT NULL,
|
||||
"user_id" TEXT NOT NULL,
|
||||
"server_id" TEXT NOT NULL,
|
||||
"values_b64" TEXT NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_MCPUserEnvVars_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_user_id_server_id_key" ON "LiteLLM_MCPUserEnvVars"("user_id", "server_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_user_id_idx" ON "LiteLLM_MCPUserEnvVars"("user_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_server_id_idx" ON "LiteLLM_MCPUserEnvVars"("server_id");
|
||||
+2
@@ -0,0 +1,2 @@
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "oauth_passthrough" BOOLEAN NOT NULL DEFAULT false;
|
||||
+2
@@ -0,0 +1,2 @@
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "oauth2_flow" TEXT;
|
||||
+3
@@ -0,0 +1,3 @@
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "timeout" DOUBLE PRECISION;
|
||||
|
||||
@@ -311,6 +311,11 @@ model LiteLLM_MCPServerTable {
|
||||
tool_name_to_description Json? @default("{}")
|
||||
extra_headers String[] @default([])
|
||||
static_headers Json? @default("{}")
|
||||
// Admin-configured environment variables interpolated into static_headers
|
||||
// via ${NAME} syntax. Stored as an array of
|
||||
// {name, value, scope, description}. scope is "global" (value used as-is)
|
||||
// or "user" (value supplied per-user via LiteLLM_MCPUserEnvVars).
|
||||
env_vars Json? @default("[]")
|
||||
// Health check status
|
||||
status String? @default("unknown")
|
||||
last_health_check DateTime?
|
||||
@@ -322,13 +327,16 @@ model LiteLLM_MCPServerTable {
|
||||
authorization_url String?
|
||||
token_url String?
|
||||
registration_url String?
|
||||
oauth2_flow String?
|
||||
allow_all_keys Boolean @default(false)
|
||||
available_on_public_internet Boolean @default(true)
|
||||
delegate_auth_to_upstream Boolean @default(false)
|
||||
oauth_passthrough Boolean @default(false)
|
||||
is_byok Boolean @default(false)
|
||||
byok_description String[] @default([])
|
||||
byok_api_key_help_url String?
|
||||
source_url String?
|
||||
timeout Float?
|
||||
// BYOM submission lifecycle
|
||||
approval_status String? @default("active")
|
||||
submitted_by String?
|
||||
@@ -363,6 +371,21 @@ model LiteLLM_MCPUserCredentials {
|
||||
@@unique([user_id, server_id])
|
||||
}
|
||||
|
||||
// Per-user environment variable values for MCP servers.
|
||||
// values_b64 is an encrypted JSON object: {VAR_NAME: "value", ...}.
|
||||
model LiteLLM_MCPUserEnvVars {
|
||||
id String @id @default(uuid())
|
||||
user_id String
|
||||
server_id String
|
||||
values_b64 String
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@unique([user_id, server_id])
|
||||
@@index([user_id])
|
||||
@@index([server_id])
|
||||
}
|
||||
|
||||
// Generate Tokens for Proxy
|
||||
model LiteLLM_VerificationToken {
|
||||
token String @id
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.73"
|
||||
version = "0.4.74"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
||||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.73"
|
||||
version = "0.4.74"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
||||
+37
-1
@@ -16,8 +16,17 @@ import os
|
||||
# Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available
|
||||
import dotenv as _dotenv
|
||||
|
||||
|
||||
def _dev_env_hot_reload_enabled() -> bool:
|
||||
"""The proxy exports this flag when started with ``--reload``. A reloaded
|
||||
worker is a fresh process that inherits the reloader's environment, so an
|
||||
edited ``.env`` value stays masked by the stale inherited one unless we
|
||||
let the file win; overriding makes the edit take effect on reload."""
|
||||
return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True"
|
||||
|
||||
|
||||
if os.getenv("LITELLM_MODE", "DEV") == "DEV":
|
||||
_dotenv.load_dotenv()
|
||||
_dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
|
||||
|
||||
from typing import (
|
||||
Callable,
|
||||
@@ -240,6 +249,7 @@ api_key: Optional[str] = None
|
||||
openai_key: Optional[str] = None
|
||||
groq_key: Optional[str] = None
|
||||
gigachat_key: Optional[str] = None
|
||||
xai_key: Optional[str] = None
|
||||
databricks_key: Optional[str] = None
|
||||
openai_like_key: Optional[str] = None
|
||||
azure_key: Optional[str] = None
|
||||
@@ -277,6 +287,7 @@ ovhcloud_key: Optional[str] = None
|
||||
lemonade_key: Optional[str] = None
|
||||
sap_service_key: Optional[str] = None
|
||||
amazon_nova_api_key: Optional[str] = None
|
||||
inception_key: Optional[str] = None
|
||||
common_cloud_provider_auth_params: dict = {
|
||||
"params": ["project", "region_name", "token"],
|
||||
"providers": ["vertex_ai", "bedrock", "watsonx", "azure", "vertex_ai_beta"],
|
||||
@@ -442,6 +453,7 @@ disable_copilot_system_to_assistant: bool = (
|
||||
False # If false (default), converts all 'system' role messages to 'assistant' for GitHub Copilot compatibility. Set to true to disable this behavior.
|
||||
)
|
||||
public_mcp_servers: Optional[List[str]] = None
|
||||
public_mcp_hub_strict_whitelist: bool = True
|
||||
public_model_groups: Optional[List[str]] = None
|
||||
public_agent_groups: Optional[List[str]] = None
|
||||
# Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
|
||||
@@ -550,6 +562,7 @@ cohere_models: Set = set()
|
||||
cohere_chat_models: Set = set()
|
||||
mistral_chat_models: Set = set()
|
||||
text_completion_codestral_models: Set = set()
|
||||
text_completion_inception_models: Set = set()
|
||||
anthropic_models: Set = set()
|
||||
openrouter_models: Set = set()
|
||||
datarobot_models: Set = set()
|
||||
@@ -608,6 +621,7 @@ cerebras_models: Set = set()
|
||||
galadriel_models: Set = set()
|
||||
nvidia_nim_models: Set = set()
|
||||
nvidia_riva_models: Set = set()
|
||||
soniox_models: Set = set()
|
||||
sambanova_models: Set = set()
|
||||
sambanova_embedding_models: Set = set()
|
||||
novita_models: Set = set()
|
||||
@@ -627,6 +641,7 @@ publicai_models: Set = set()
|
||||
v0_models: Set = set()
|
||||
morph_models: Set = set()
|
||||
lambda_ai_models: Set = set()
|
||||
inception_models: Set = set()
|
||||
hyperbolic_models: Set = set()
|
||||
black_forest_labs_models: Set = set()
|
||||
recraft_models: Set = set()
|
||||
@@ -791,6 +806,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
||||
fireworks_ai_embedding_models.add(key)
|
||||
elif value.get("litellm_provider") == "text-completion-codestral":
|
||||
text_completion_codestral_models.add(key)
|
||||
elif value.get("litellm_provider") == "text-completion-inception":
|
||||
text_completion_inception_models.add(key)
|
||||
elif value.get("litellm_provider") == "xai":
|
||||
xai_models.add(key)
|
||||
elif value.get("litellm_provider") == "zai":
|
||||
@@ -837,6 +854,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
||||
nvidia_nim_models.add(key)
|
||||
elif value.get("litellm_provider") == "nvidia_riva":
|
||||
nvidia_riva_models.add(key)
|
||||
elif value.get("litellm_provider") == "soniox":
|
||||
soniox_models.add(key)
|
||||
elif value.get("litellm_provider") == "sambanova":
|
||||
sambanova_models.add(key)
|
||||
elif value.get("litellm_provider") == "sambanova-embedding-models":
|
||||
@@ -877,6 +896,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
||||
morph_models.add(key)
|
||||
elif value.get("litellm_provider") == "lambda_ai":
|
||||
lambda_ai_models.add(key)
|
||||
elif value.get("litellm_provider") == "inception":
|
||||
inception_models.add(key)
|
||||
elif value.get("litellm_provider") == "hyperbolic":
|
||||
hyperbolic_models.add(key)
|
||||
elif value.get("litellm_provider") == "black_forest_labs":
|
||||
@@ -979,6 +1000,7 @@ model_list = list(
|
||||
| watsonx_models
|
||||
| gemini_models
|
||||
| text_completion_codestral_models
|
||||
| text_completion_inception_models
|
||||
| xai_models
|
||||
| zai_models
|
||||
| fal_ai_models
|
||||
@@ -999,6 +1021,7 @@ model_list = list(
|
||||
| galadriel_models
|
||||
| nvidia_nim_models
|
||||
| nvidia_riva_models
|
||||
| soniox_models
|
||||
| sambanova_models
|
||||
| azure_text_models
|
||||
| novita_models
|
||||
@@ -1017,6 +1040,7 @@ model_list = list(
|
||||
| v0_models
|
||||
| morph_models
|
||||
| lambda_ai_models
|
||||
| inception_models
|
||||
| black_forest_labs_models
|
||||
| recraft_models
|
||||
| cometapi_models
|
||||
@@ -1073,6 +1097,7 @@ models_by_provider: dict = {
|
||||
"fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models,
|
||||
"aleph_alpha": aleph_alpha_models,
|
||||
"text-completion-codestral": text_completion_codestral_models,
|
||||
"text-completion-inception": text_completion_inception_models,
|
||||
"xai": xai_models,
|
||||
"zai": zai_models,
|
||||
"fal_ai": fal_ai_models,
|
||||
@@ -1097,6 +1122,7 @@ models_by_provider: dict = {
|
||||
"galadriel": galadriel_models,
|
||||
"nvidia_nim": nvidia_nim_models,
|
||||
"nvidia_riva": nvidia_riva_models,
|
||||
"soniox": soniox_models,
|
||||
"sambanova": sambanova_models | sambanova_embedding_models,
|
||||
"novita": novita_models,
|
||||
"nebius": nebius_models | nebius_embedding_models,
|
||||
@@ -1117,6 +1143,7 @@ models_by_provider: dict = {
|
||||
"v0": v0_models,
|
||||
"morph": morph_models,
|
||||
"lambda_ai": lambda_ai_models,
|
||||
"inception": inception_models,
|
||||
"hyperbolic": hyperbolic_models,
|
||||
"black_forest_labs": black_forest_labs_models,
|
||||
"recraft": recraft_models,
|
||||
@@ -1727,6 +1754,9 @@ if TYPE_CHECKING:
|
||||
from .llms.openrouter.responses.transformation import (
|
||||
OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig,
|
||||
)
|
||||
from .llms.bedrock_mantle.responses.transformation import (
|
||||
BedrockMantleResponsesAPIConfig as BedrockMantleResponsesAPIConfig,
|
||||
)
|
||||
from .llms.gemini.interactions.transformation import (
|
||||
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
|
||||
)
|
||||
@@ -1868,6 +1898,9 @@ if TYPE_CHECKING:
|
||||
from .llms.codestral.completion.transformation import (
|
||||
CodestralTextCompletionConfig as CodestralTextCompletionConfig,
|
||||
)
|
||||
from .llms.inception.completion.transformation import (
|
||||
InceptionTextCompletionConfig as InceptionTextCompletionConfig,
|
||||
)
|
||||
from .llms.azure.azure import (
|
||||
AzureOpenAIAssistantsAPIConfig as AzureOpenAIAssistantsAPIConfig,
|
||||
)
|
||||
@@ -1936,6 +1969,9 @@ if TYPE_CHECKING:
|
||||
from .llms.lambda_ai.chat.transformation import (
|
||||
LambdaAIChatConfig as LambdaAIChatConfig,
|
||||
)
|
||||
from .llms.inception.chat.transformation import (
|
||||
InceptionChatConfig as InceptionChatConfig,
|
||||
)
|
||||
from .llms.hyperbolic.chat.transformation import (
|
||||
HyperbolicChatConfig as HyperbolicChatConfig,
|
||||
)
|
||||
|
||||
@@ -237,6 +237,7 @@ LLM_CONFIG_NAMES = (
|
||||
"PerplexityResponsesConfig",
|
||||
"DatabricksResponsesAPIConfig",
|
||||
"OpenRouterResponsesAPIConfig",
|
||||
"BedrockMantleResponsesAPIConfig",
|
||||
"GoogleAIStudioInteractionsConfig",
|
||||
"OpenAIOSeriesConfig",
|
||||
"AnthropicSkillsConfig",
|
||||
@@ -267,6 +268,7 @@ LLM_CONFIG_NAMES = (
|
||||
"AIMLChatConfig",
|
||||
"VolcEngineChatConfig",
|
||||
"CodestralTextCompletionConfig",
|
||||
"InceptionTextCompletionConfig",
|
||||
"AzureOpenAIAssistantsAPIConfig",
|
||||
"HerokuChatConfig",
|
||||
"CometAPIConfig",
|
||||
@@ -310,6 +312,7 @@ LLM_CONFIG_NAMES = (
|
||||
"MorphChatConfig",
|
||||
"RAGFlowConfig",
|
||||
"LambdaAIChatConfig",
|
||||
"InceptionChatConfig",
|
||||
"HyperbolicChatConfig",
|
||||
"VercelAIGatewayConfig",
|
||||
"OVHCloudChatConfig",
|
||||
@@ -318,6 +321,7 @@ LLM_CONFIG_NAMES = (
|
||||
"LemonadeChatConfig",
|
||||
"SnowflakeEmbeddingConfig",
|
||||
"AmazonNovaChatConfig",
|
||||
"SonioxAudioTranscriptionConfig",
|
||||
)
|
||||
|
||||
# Types that support lazy loading via _lazy_import_types
|
||||
@@ -956,6 +960,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
||||
".llms.openrouter.responses.transformation",
|
||||
"OpenRouterResponsesAPIConfig",
|
||||
),
|
||||
"BedrockMantleResponsesAPIConfig": (
|
||||
".llms.bedrock_mantle.responses.transformation",
|
||||
"BedrockMantleResponsesAPIConfig",
|
||||
),
|
||||
"GoogleAIStudioInteractionsConfig": (
|
||||
".llms.gemini.interactions.transformation",
|
||||
"GoogleAIStudioInteractionsConfig",
|
||||
@@ -1040,6 +1048,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
||||
".llms.codestral.completion.transformation",
|
||||
"CodestralTextCompletionConfig",
|
||||
),
|
||||
"InceptionTextCompletionConfig": (
|
||||
".llms.inception.completion.transformation",
|
||||
"InceptionTextCompletionConfig",
|
||||
),
|
||||
"AzureOpenAIAssistantsAPIConfig": (
|
||||
".llms.azure.azure",
|
||||
"AzureOpenAIAssistantsAPIConfig",
|
||||
@@ -1154,6 +1166,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
||||
"MorphChatConfig": (".llms.morph.chat.transformation", "MorphChatConfig"),
|
||||
"RAGFlowConfig": (".llms.ragflow.chat.transformation", "RAGFlowConfig"),
|
||||
"LambdaAIChatConfig": (".llms.lambda_ai.chat.transformation", "LambdaAIChatConfig"),
|
||||
"InceptionChatConfig": (
|
||||
".llms.inception.chat.transformation",
|
||||
"InceptionChatConfig",
|
||||
),
|
||||
"HyperbolicChatConfig": (
|
||||
".llms.hyperbolic.chat.transformation",
|
||||
"HyperbolicChatConfig",
|
||||
@@ -1180,6 +1196,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
||||
".llms.amazon_nova.chat.transformation",
|
||||
"AmazonNovaChatConfig",
|
||||
),
|
||||
"SonioxAudioTranscriptionConfig": (
|
||||
".llms.soniox.audio_transcription.transformation",
|
||||
"SonioxAudioTranscriptionConfig",
|
||||
),
|
||||
}
|
||||
|
||||
# Import map for utils module lazy imports
|
||||
|
||||
@@ -371,6 +371,8 @@ class ServiceLogging(CustomLogger):
|
||||
service=ServiceTypes.LITELLM,
|
||||
duration=_duration,
|
||||
call_type=kwargs.get("call_type", "unknown"),
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
@@ -20,9 +20,20 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
|
||||
)
|
||||
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
|
||||
|
||||
# litellm_params key carrying the authenticated principal (hashed virtual key) so
|
||||
# A2A provider configs can scope provider-side state (e.g. LangFlow session memory)
|
||||
# per key instead of trusting the client-supplied A2A contextId.
|
||||
A2A_USER_API_KEY_HASH_PARAM = "litellm_a2a_user_api_key_hash"
|
||||
|
||||
# Agent metadata fields stored in litellm_params that are not valid litellm.acompletion() kwargs
|
||||
_AGENT_ONLY_PARAMS = frozenset(
|
||||
{"is_public", "agent_name", "agent_id", "agent_card_params"}
|
||||
{
|
||||
"is_public",
|
||||
"agent_name",
|
||||
"agent_id",
|
||||
"agent_card_params",
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -37,6 +48,8 @@ class A2ACompletionBridgeHandler:
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Handle non-streaming A2A request via litellm.acompletion.
|
||||
@@ -50,25 +63,24 @@ class A2ACompletionBridgeHandler:
|
||||
Returns:
|
||||
A2A SendMessageResponse dict
|
||||
"""
|
||||
# Get provider config for custom_llm_provider
|
||||
custom_llm_provider = litellm_params.get("custom_llm_provider")
|
||||
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model=litellm_params.get("model"),
|
||||
)
|
||||
|
||||
# If provider config exists, use it
|
||||
if a2a_provider_config is not None:
|
||||
verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}")
|
||||
|
||||
response_data = await a2a_provider_config.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
if not _skip_a2a_provider_routing:
|
||||
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model=litellm_params.get("model"),
|
||||
)
|
||||
|
||||
return response_data
|
||||
if a2a_provider_config is not None:
|
||||
verbose_logger.info(
|
||||
f"A2A: Using provider config for {custom_llm_provider}"
|
||||
)
|
||||
|
||||
return await a2a_provider_config.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
# Extract message from params
|
||||
message = params.get("message", {})
|
||||
@@ -137,6 +149,8 @@ class A2ACompletionBridgeHandler:
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Handle streaming A2A request via litellm.acompletion with stream=True.
|
||||
@@ -156,28 +170,27 @@ class A2ACompletionBridgeHandler:
|
||||
Yields:
|
||||
A2A streaming response events
|
||||
"""
|
||||
# Get provider config for custom_llm_provider
|
||||
custom_llm_provider = litellm_params.get("custom_llm_provider")
|
||||
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model=litellm_params.get("model"),
|
||||
)
|
||||
|
||||
# If provider config exists, use it
|
||||
if a2a_provider_config is not None:
|
||||
verbose_logger.info(
|
||||
f"A2A: Using provider config for {custom_llm_provider} (streaming)"
|
||||
if not _skip_a2a_provider_routing:
|
||||
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model=litellm_params.get("model"),
|
||||
)
|
||||
|
||||
async for chunk in a2a_provider_config.handle_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
):
|
||||
yield chunk
|
||||
if a2a_provider_config is not None:
|
||||
verbose_logger.info(
|
||||
f"A2A: Using provider config for {custom_llm_provider} (streaming)"
|
||||
)
|
||||
|
||||
return
|
||||
async for chunk in a2a_provider_config.handle_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
):
|
||||
yield chunk
|
||||
|
||||
return
|
||||
|
||||
# Extract message from params
|
||||
message = params.get("message", {})
|
||||
|
||||
@@ -159,7 +159,9 @@ async def _send_message_via_completion_bridge(
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
return LiteLLMSendMessageResponse.from_dict(response_dict)
|
||||
return LiteLLMSendMessageResponse.from_dict(
|
||||
response_dict, request_id=str(request.id)
|
||||
)
|
||||
|
||||
|
||||
async def _execute_a2a_send_with_retry(
|
||||
@@ -317,15 +319,6 @@ async def asend_message(
|
||||
)
|
||||
card_url = getattr(agent_card, "url", None) if agent_card else None
|
||||
|
||||
context_id = trace_id or str(uuid.uuid4())
|
||||
message = request.params.message
|
||||
if isinstance(message, dict):
|
||||
if message.get("context_id") is None:
|
||||
message["context_id"] = context_id
|
||||
else:
|
||||
if getattr(message, "context_id", None) is None:
|
||||
message.context_id = context_id
|
||||
|
||||
a2a_response = await _execute_a2a_send_with_retry(
|
||||
a2a_client=a2a_client,
|
||||
request=request,
|
||||
@@ -338,7 +331,9 @@ async def asend_message(
|
||||
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
|
||||
|
||||
# Wrap in LiteLLM response type for _hidden_params support
|
||||
response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response)
|
||||
response = LiteLLMSendMessageResponse.from_a2a_response(
|
||||
a2a_response, request_id=str(request.id)
|
||||
)
|
||||
|
||||
# Calculate token usage from request and response
|
||||
response_dict = a2a_response.model_dump(mode="json", exclude_none=True)
|
||||
|
||||
@@ -48,4 +48,16 @@ class A2AProviderConfigManager:
|
||||
|
||||
return BedrockAgentCoreA2AConfig()
|
||||
|
||||
if custom_llm_provider == "langflow":
|
||||
from litellm.a2a_protocol.providers.langflow.config import LangFlowA2AConfig
|
||||
|
||||
return LangFlowA2AConfig()
|
||||
|
||||
if custom_llm_provider == "watsonx_orchestrate":
|
||||
from litellm.a2a_protocol.providers.watsonx_orchestrate.config import (
|
||||
WatsonxOrchestrateA2AConfig,
|
||||
)
|
||||
|
||||
return WatsonxOrchestrateA2AConfig()
|
||||
|
||||
return None
|
||||
|
||||
@@ -0,0 +1,62 @@
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
|
||||
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
|
||||
A2A_USER_API_KEY_HASH_PARAM,
|
||||
A2ACompletionBridgeHandler,
|
||||
)
|
||||
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
|
||||
from litellm.llms.langflow.a2a import merge_a2a_session_into_litellm_params
|
||||
|
||||
|
||||
class LangFlowA2AConfig(BaseA2AProviderConfig):
|
||||
"""A2A bridge for LangFlow: scopes contextId to the authenticated key as the
|
||||
LangFlow session_id, then uses completion."""
|
||||
|
||||
async def handle_non_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Dict[str, Any]:
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if not litellm_params:
|
||||
raise ValueError(
|
||||
"litellm_params is required for LangFlowA2AConfig "
|
||||
"(must contain custom_llm_provider and model)"
|
||||
)
|
||||
litellm_params = merge_a2a_session_into_litellm_params(
|
||||
litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)
|
||||
)
|
||||
return await A2ACompletionBridgeHandler.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
_skip_a2a_provider_routing=True,
|
||||
)
|
||||
|
||||
async def handle_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if not litellm_params:
|
||||
raise ValueError(
|
||||
"litellm_params is required for LangFlowA2AConfig "
|
||||
"(must contain custom_llm_provider and model)"
|
||||
)
|
||||
litellm_params = merge_a2a_session_into_litellm_params(
|
||||
litellm_params, params, litellm_params.get(A2A_USER_API_KEY_HASH_PARAM)
|
||||
)
|
||||
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
_skip_a2a_provider_routing=True,
|
||||
):
|
||||
yield chunk
|
||||
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
IBM watsonx Orchestrate (WXO) A2A provider.
|
||||
"""
|
||||
@@ -0,0 +1,55 @@
|
||||
"""
|
||||
A2A provider configuration for IBM watsonx Orchestrate (WXO).
|
||||
"""
|
||||
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
|
||||
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
|
||||
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
|
||||
WatsonxOrchestrateHandler,
|
||||
)
|
||||
|
||||
|
||||
class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
|
||||
"""A2A bridge for IBM watsonx Orchestrate (REST runs API + poll/SSE)."""
|
||||
|
||||
async def handle_non_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> Dict[str, Any]:
|
||||
"""Handle a non-streaming A2A request via WXO runs API."""
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if not litellm_params:
|
||||
raise ValueError(
|
||||
"litellm_params is required for WatsonxOrchestrateA2AConfig "
|
||||
"(must contain cp4d_host, instance_id, wxo_agent_id, api_key)"
|
||||
)
|
||||
return await WatsonxOrchestrateHandler.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
async def handle_streaming(
|
||||
self,
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""Handle a streaming A2A request via WXO streaming runs API."""
|
||||
litellm_params = kwargs.get("litellm_params")
|
||||
if not litellm_params:
|
||||
raise ValueError(
|
||||
"litellm_params is required for WatsonxOrchestrateA2AConfig "
|
||||
"(must contain cp4d_host, instance_id, wxo_agent_id, api_key)"
|
||||
)
|
||||
async for chunk in WatsonxOrchestrateHandler.handle_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
):
|
||||
yield chunk
|
||||
@@ -0,0 +1,373 @@
|
||||
"""
|
||||
Handler for IBM watsonx Orchestrate (WXO) agent provider.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from typing import Any, AsyncIterator, Dict, NamedTuple, Optional, Tuple, cast
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import (
|
||||
WatsonxOrchestrateTransformation,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
||||
_IBM_CLOUD_IAM_URL = "https://iam.cloud.ibm.com/identity/token"
|
||||
_POLL_INTERVAL_S = 2.0
|
||||
_MAX_POLL_ATTEMPTS = 90
|
||||
_TOKEN_CACHE_TTL_BUFFER_S = 60
|
||||
_token_cache: Dict[str, Tuple[str, float]] = {}
|
||||
|
||||
|
||||
class WXORequestParams(NamedTuple):
|
||||
cp4d_host: str
|
||||
instance_id: str
|
||||
wxo_agent_id: str
|
||||
api_key: str
|
||||
username: Optional[str]
|
||||
auth_mode: str
|
||||
thread_id: Optional[str]
|
||||
|
||||
|
||||
class WatsonxOrchestrateHandler:
|
||||
@staticmethod
|
||||
def _http_client(timeout: float = 90.0) -> AsyncHTTPHandler:
|
||||
return get_async_httpx_client(
|
||||
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
|
||||
params={"timeout": timeout},
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _token_cache_key(
|
||||
auth_mode: str,
|
||||
cp4d_host: str,
|
||||
api_key: str,
|
||||
username: Optional[str],
|
||||
) -> str:
|
||||
material = f"{auth_mode}:{cp4d_host}:{username or ''}:{api_key}"
|
||||
return hashlib.sha256(material.encode()).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def _cp4d_token_ttl_seconds(
|
||||
expiration: Any, now_wall: Optional[float] = None
|
||||
) -> int:
|
||||
# CP4D returns expiration as absolute Unix epoch seconds, not a duration.
|
||||
expires_at = int(expiration)
|
||||
wall = now_wall if now_wall is not None else time.time()
|
||||
return max(expires_at - int(wall), 0)
|
||||
|
||||
@staticmethod
|
||||
async def _get_bearer_token(
|
||||
cp4d_host: str,
|
||||
auth_mode: str,
|
||||
api_key: str,
|
||||
username: Optional[str] = None,
|
||||
client: Optional[AsyncHTTPHandler] = None,
|
||||
) -> str:
|
||||
cache_key = WatsonxOrchestrateHandler._token_cache_key(
|
||||
auth_mode, cp4d_host, api_key, username
|
||||
)
|
||||
now = time.monotonic()
|
||||
cached = _token_cache.get(cache_key)
|
||||
if cached and cached[1] > now:
|
||||
return cached[0]
|
||||
|
||||
if client is None:
|
||||
client = WatsonxOrchestrateHandler._http_client(timeout=30.0)
|
||||
|
||||
if auth_mode == "ibm_cloud":
|
||||
response = await client.post(
|
||||
_IBM_CLOUD_IAM_URL,
|
||||
data={
|
||||
"grant_type": "urn:ibm:params:oauth:grant-type:apikey",
|
||||
"apikey": api_key,
|
||||
},
|
||||
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
token = str(payload["access_token"])
|
||||
ttl_s = int(payload.get("expires_in", 3600))
|
||||
else:
|
||||
if not username:
|
||||
raise ValueError(
|
||||
"'username' is required in litellm_params when auth_mode='cp4d'"
|
||||
)
|
||||
token_url = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize"
|
||||
response = await client.post(
|
||||
token_url,
|
||||
json={"username": username, "api_key": api_key},
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
token = str(payload["token"])
|
||||
expiration = payload.get("expiration")
|
||||
if expiration is None:
|
||||
ttl_s = 3600
|
||||
else:
|
||||
ttl_s = WatsonxOrchestrateHandler._cp4d_token_ttl_seconds(expiration)
|
||||
|
||||
expires_at = now + max(ttl_s - _TOKEN_CACHE_TTL_BUFFER_S, 0)
|
||||
_token_cache[cache_key] = (token, expires_at)
|
||||
for stale_key, (_, stale_expires_at) in list(_token_cache.items()):
|
||||
if stale_expires_at <= now:
|
||||
del _token_cache[stale_key]
|
||||
return token
|
||||
|
||||
@staticmethod
|
||||
async def _poll_run(
|
||||
base_url: str,
|
||||
run_id: str,
|
||||
auth_headers: Dict[str, str],
|
||||
client: AsyncHTTPHandler,
|
||||
max_attempts: int = _MAX_POLL_ATTEMPTS,
|
||||
interval_s: float = _POLL_INTERVAL_S,
|
||||
) -> Dict[str, Any]:
|
||||
url = f"{base_url}/v1/orchestrate/runs/{run_id}"
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
await asyncio.sleep(interval_s)
|
||||
response = await client.get(url, headers=auth_headers)
|
||||
response.raise_for_status()
|
||||
result: Dict[str, Any] = response.json()
|
||||
status = result.get("status", "")
|
||||
verbose_logger.debug(
|
||||
f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'"
|
||||
)
|
||||
if status in WatsonxOrchestrateTransformation.TERMINAL_STATES:
|
||||
return result
|
||||
|
||||
raise asyncio.TimeoutError(
|
||||
f"WXO run '{run_id}' did not reach a terminal state after "
|
||||
f"{max_attempts * interval_s:.0f}s"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _get_successful_run_data(
|
||||
run_data: Dict[str, Any],
|
||||
base_url: str,
|
||||
auth_headers: Dict[str, str],
|
||||
client: AsyncHTTPHandler,
|
||||
) -> Dict[str, Any]:
|
||||
status = run_data.get("status", "")
|
||||
if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
|
||||
run_id = run_data.get("run_id") or run_data.get("id") or ""
|
||||
if not run_id:
|
||||
raise ValueError(f"WXO: No run_id in response: {run_data}")
|
||||
run_data = await WatsonxOrchestrateHandler._poll_run(
|
||||
base_url=base_url,
|
||||
run_id=run_id,
|
||||
auth_headers=auth_headers,
|
||||
client=client,
|
||||
)
|
||||
status = run_data.get("status", "")
|
||||
|
||||
if status not in WatsonxOrchestrateTransformation.SUCCESS_STATES:
|
||||
raise RuntimeError(
|
||||
f"WXO run ended with non-success status '{status}': {run_data}"
|
||||
)
|
||||
|
||||
return run_data
|
||||
|
||||
@staticmethod
|
||||
async def _accumulate_wxo_sse_text(response: Any) -> str:
|
||||
accumulated_text = ""
|
||||
async for line in response.aiter_lines():
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
data_str = line[5:].strip()
|
||||
if not data_str or data_str == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
event = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(
|
||||
event
|
||||
)
|
||||
if chunk_text:
|
||||
accumulated_text += chunk_text
|
||||
return accumulated_text
|
||||
|
||||
@staticmethod
|
||||
def _extract_litellm_params(litellm_params: Dict[str, Any]) -> WXORequestParams:
|
||||
cp4d_host = litellm_params.get("cp4d_host") or ""
|
||||
instance_id = litellm_params.get("instance_id") or ""
|
||||
wxo_agent_id = litellm_params.get("wxo_agent_id") or ""
|
||||
api_key = litellm_params.get("api_key") or ""
|
||||
|
||||
if not cp4d_host:
|
||||
raise ValueError("'cp4d_host' is required in litellm_params for WXO agents")
|
||||
if not instance_id:
|
||||
raise ValueError(
|
||||
"'instance_id' is required in litellm_params for WXO agents"
|
||||
)
|
||||
if not wxo_agent_id:
|
||||
raise ValueError(
|
||||
"'wxo_agent_id' is required in litellm_params for WXO agents"
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError("'api_key' is required in litellm_params for WXO agents")
|
||||
|
||||
return WXORequestParams(
|
||||
cp4d_host=cp4d_host,
|
||||
instance_id=instance_id,
|
||||
wxo_agent_id=wxo_agent_id,
|
||||
api_key=api_key,
|
||||
username=litellm_params.get("username") or None,
|
||||
auth_mode=litellm_params.get("auth_mode") or "cp4d",
|
||||
thread_id=litellm_params.get("thread_id") or None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def handle_non_streaming(
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
|
||||
|
||||
client = WatsonxOrchestrateHandler._http_client(timeout=90.0)
|
||||
token = await WatsonxOrchestrateHandler._get_bearer_token(
|
||||
cp4d_host=wxo.cp4d_host,
|
||||
auth_mode=wxo.auth_mode,
|
||||
api_key=wxo.api_key,
|
||||
username=wxo.username,
|
||||
client=client,
|
||||
)
|
||||
base_url = WatsonxOrchestrateTransformation.get_api_base_url(
|
||||
wxo.cp4d_host, wxo.instance_id
|
||||
)
|
||||
auth_headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
|
||||
body = WatsonxOrchestrateTransformation.build_wxo_run_body(
|
||||
wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id
|
||||
)
|
||||
|
||||
run_response = await client.post(
|
||||
f"{base_url}/v1/orchestrate/runs",
|
||||
json=body,
|
||||
headers=auth_headers,
|
||||
)
|
||||
run_response.raise_for_status()
|
||||
run_data: Dict[str, Any] = run_response.json()
|
||||
|
||||
run_data = await WatsonxOrchestrateHandler._get_successful_run_data(
|
||||
run_data=run_data,
|
||||
base_url=base_url,
|
||||
auth_headers=auth_headers,
|
||||
client=client,
|
||||
)
|
||||
|
||||
response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(
|
||||
run_data
|
||||
)
|
||||
return WatsonxOrchestrateTransformation.build_a2a_message_response(
|
||||
request_id=request_id, text=response_text
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def handle_streaming(
|
||||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
chunk_size: int = 50,
|
||||
delay_ms: int = 10,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
|
||||
|
||||
client = WatsonxOrchestrateHandler._http_client(timeout=120.0)
|
||||
token = await WatsonxOrchestrateHandler._get_bearer_token(
|
||||
cp4d_host=wxo.cp4d_host,
|
||||
auth_mode=wxo.auth_mode,
|
||||
api_key=wxo.api_key,
|
||||
username=wxo.username,
|
||||
client=client,
|
||||
)
|
||||
base_url = WatsonxOrchestrateTransformation.get_api_base_url(
|
||||
wxo.cp4d_host, wxo.instance_id
|
||||
)
|
||||
auth_headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream, application/json",
|
||||
}
|
||||
text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
|
||||
body = WatsonxOrchestrateTransformation.build_wxo_run_body(
|
||||
wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id
|
||||
)
|
||||
|
||||
try:
|
||||
response = await client.post(
|
||||
f"{base_url}/v1/orchestrate/runs/stream",
|
||||
json=body,
|
||||
headers=auth_headers,
|
||||
stream=True,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.TransportError as exc:
|
||||
verbose_logger.warning(
|
||||
f"WXO: Streaming request failed before a run was submitted "
|
||||
f"({exc!r}), falling back to non-streaming + fake streaming",
|
||||
exc_info=True,
|
||||
)
|
||||
result = await WatsonxOrchestrateHandler.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
response_text = (
|
||||
WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(
|
||||
result
|
||||
)
|
||||
)
|
||||
async for (
|
||||
chunk
|
||||
) in WatsonxOrchestrateTransformation.fake_streaming_from_text(
|
||||
text=response_text,
|
||||
request_id=request_id,
|
||||
chunk_size=chunk_size,
|
||||
delay_ms=delay_ms,
|
||||
):
|
||||
yield chunk
|
||||
return
|
||||
|
||||
content_type = response.headers.get("content-type", "").lower()
|
||||
if "text/event-stream" not in content_type:
|
||||
response_body = await response.aread()
|
||||
result = json.loads(response_body)
|
||||
result = await WatsonxOrchestrateHandler._get_successful_run_data(
|
||||
run_data=result,
|
||||
base_url=base_url,
|
||||
auth_headers=auth_headers,
|
||||
client=client,
|
||||
)
|
||||
accumulated_text = (
|
||||
WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result)
|
||||
)
|
||||
else:
|
||||
accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(
|
||||
response
|
||||
)
|
||||
|
||||
async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text(
|
||||
text=accumulated_text,
|
||||
request_id=request_id,
|
||||
chunk_size=chunk_size,
|
||||
delay_ms=delay_ms,
|
||||
):
|
||||
yield chunk
|
||||
@@ -0,0 +1,224 @@
|
||||
"""
|
||||
Transformation layer for IBM watsonx Orchestrate (WXO) agent provider.
|
||||
|
||||
WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model:
|
||||
POST /v1/orchestrate/runs → submit run, get run_id
|
||||
GET /v1/orchestrate/runs/{id} → poll until terminal state
|
||||
POST /v1/orchestrate/runs/stream → native SSE streaming
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
|
||||
class WatsonxOrchestrateTransformation:
|
||||
"""
|
||||
Handles request/response transformation between A2A and the WXO REST API.
|
||||
"""
|
||||
|
||||
TERMINAL_STATES = frozenset(
|
||||
{"completed", "succeeded", "failed", "error", "cancelled"}
|
||||
)
|
||||
SUCCESS_STATES = frozenset({"completed", "succeeded"})
|
||||
|
||||
@staticmethod
|
||||
def get_api_base_url(cp4d_host: str, instance_id: str) -> str:
|
||||
"""Build the WXO API base URL from host and instance ID."""
|
||||
return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}"
|
||||
|
||||
@staticmethod
|
||||
def extract_text_from_a2a_params(params: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Extract user message text from A2A MessageSendParams.
|
||||
|
||||
A2A format: params.message.parts[*] where part.kind == "text"
|
||||
"""
|
||||
message = params.get("message", {})
|
||||
parts = message.get("parts", [])
|
||||
texts = []
|
||||
for part in parts:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
kind = part.get("kind")
|
||||
if kind in (None, "", "text") and part.get("text"):
|
||||
texts.append(part["text"])
|
||||
return " ".join(texts) or ""
|
||||
|
||||
@staticmethod
|
||||
def build_wxo_run_body(
|
||||
wxo_agent_id: str,
|
||||
text: str,
|
||||
thread_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Build the WXO POST /v1/orchestrate/runs request body."""
|
||||
body: Dict[str, Any] = {
|
||||
"agent_id": wxo_agent_id,
|
||||
"message": {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"response_type": "text",
|
||||
"text": text,
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
if thread_id:
|
||||
body["thread_id"] = thread_id
|
||||
return body
|
||||
|
||||
@staticmethod
|
||||
def extract_text_from_wxo_result(result: Any) -> str:
|
||||
"""
|
||||
Extract response text from a WXO run result.
|
||||
|
||||
WXO can return text in several locations; checks in priority order per the API spec.
|
||||
"""
|
||||
if not isinstance(result, dict):
|
||||
return ""
|
||||
|
||||
# Primary: last_message.content[0].text
|
||||
try:
|
||||
text = result["last_message"]["content"][0]["text"]
|
||||
if text:
|
||||
return str(text)
|
||||
except (KeyError, IndexError, TypeError):
|
||||
pass
|
||||
|
||||
# Secondary: result.data.message.content[0].text
|
||||
try:
|
||||
text = result["result"]["data"]["message"]["content"][0]["text"]
|
||||
if text:
|
||||
return str(text)
|
||||
except (KeyError, IndexError, TypeError):
|
||||
pass
|
||||
|
||||
# Tertiary: results as a raw string
|
||||
results = result.get("results")
|
||||
if results and isinstance(results, str):
|
||||
return results
|
||||
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def extract_text_from_a2a_message_response(a2a_response: Dict[str, Any]) -> str:
|
||||
result = a2a_response.get("result")
|
||||
if not isinstance(result, dict):
|
||||
verbose_logger.warning("WXO: A2A response missing result object")
|
||||
return ""
|
||||
parts = result.get("parts")
|
||||
if not isinstance(parts, list):
|
||||
verbose_logger.warning("WXO: A2A result has no parts list")
|
||||
return ""
|
||||
for part in parts:
|
||||
if (
|
||||
isinstance(part, dict)
|
||||
and part.get("kind") == "text"
|
||||
and part.get("text")
|
||||
):
|
||||
return str(part["text"])
|
||||
verbose_logger.warning("WXO: A2A result parts contained no text")
|
||||
return ""
|
||||
|
||||
@staticmethod
|
||||
def build_a2a_message_response(request_id: str, text: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Build a standard A2A non-streaming SendMessageResponse (kind=message).
|
||||
"""
|
||||
return {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"role": "agent",
|
||||
"parts": [{"kind": "text", "text": text}],
|
||||
"messageId": str(uuid4()),
|
||||
},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
async def fake_streaming_from_text(
|
||||
text: str,
|
||||
request_id: str,
|
||||
chunk_size: int = 50,
|
||||
delay_ms: int = 10,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Emit standard A2A streaming events from a completed text response.
|
||||
|
||||
Event sequence:
|
||||
1. task (kind="task", state="submitted")
|
||||
2. status-update (kind="status-update", state="working")
|
||||
3. artifact-update chunks
|
||||
4. status-update (kind="status-update", state="completed", final=True)
|
||||
"""
|
||||
task_id = str(uuid4())
|
||||
context_id = str(uuid4())
|
||||
artifact_id = str(uuid4())
|
||||
|
||||
# 1. Task submitted
|
||||
yield {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"id": task_id,
|
||||
"kind": "task",
|
||||
"status": {"state": "submitted"},
|
||||
},
|
||||
}
|
||||
|
||||
# 2. Working
|
||||
yield {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"final": False,
|
||||
"kind": "status-update",
|
||||
"status": {"state": "working"},
|
||||
"taskId": task_id,
|
||||
},
|
||||
}
|
||||
await asyncio.sleep(delay_ms / 1000.0)
|
||||
|
||||
# 3. Artifact chunks (always emit at least one chunk, even for empty text)
|
||||
text_to_chunk = text or ""
|
||||
for i in range(0, max(len(text_to_chunk), 1), chunk_size):
|
||||
chunk_text = text_to_chunk[i : i + chunk_size]
|
||||
is_last = (i + chunk_size) >= max(len(text_to_chunk), 1)
|
||||
yield {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"kind": "artifact-update",
|
||||
"taskId": task_id,
|
||||
"artifact": {
|
||||
"artifactId": artifact_id,
|
||||
"parts": [{"kind": "text", "text": chunk_text}],
|
||||
},
|
||||
},
|
||||
}
|
||||
if not is_last:
|
||||
await asyncio.sleep(delay_ms / 1000.0)
|
||||
|
||||
# 4. Completed
|
||||
yield {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"final": True,
|
||||
"kind": "status-update",
|
||||
"status": {"state": "completed"},
|
||||
"taskId": task_id,
|
||||
},
|
||||
}
|
||||
|
||||
verbose_logger.debug(
|
||||
f"WXO: Fake streaming completed for request_id={request_id}"
|
||||
)
|
||||
@@ -182,6 +182,14 @@ class ResponsesToCompletionBridgeHandler:
|
||||
client=kwargs.get("client"),
|
||||
)
|
||||
|
||||
# Pin the resolved provider so `responses()` doesn't re-run
|
||||
# `get_llm_provider()` on the model string and strip a second
|
||||
# provider prefix (see GitHub issue #28505). request_data already
|
||||
# carries `custom_llm_provider` via the spread of
|
||||
# `sanitized_litellm_params`; overwriting it on the dict (rather
|
||||
# than adding an explicit kwarg) avoids the duplicate-keyword
|
||||
# TypeError that would otherwise fire on the real bridge path.
|
||||
request_data["custom_llm_provider"] = custom_llm_provider
|
||||
result = responses(
|
||||
**request_data,
|
||||
)
|
||||
@@ -268,6 +276,13 @@ class ResponsesToCompletionBridgeHandler:
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
# Pin the resolved provider so `aresponses()` doesn't re-run
|
||||
# `get_llm_provider()` on the model string and strip a second
|
||||
# provider prefix (see GitHub issue #28505). Set on request_data
|
||||
# rather than passed as a separate kwarg to avoid the duplicate-
|
||||
# keyword TypeError when `sanitized_litellm_params` already
|
||||
# carries `custom_llm_provider`.
|
||||
request_data["custom_llm_provider"] = custom_llm_provider
|
||||
result = await aresponses(
|
||||
**request_data,
|
||||
aresponses=True,
|
||||
|
||||
@@ -585,6 +585,7 @@ LITELLM_CHAT_PROVIDERS = [
|
||||
"volcengine",
|
||||
"codestral",
|
||||
"text-completion-codestral",
|
||||
"text-completion-inception",
|
||||
"deepseek",
|
||||
"sambanova",
|
||||
"maritalk",
|
||||
@@ -620,6 +621,7 @@ LITELLM_CHAT_PROVIDERS = [
|
||||
"oci",
|
||||
"morph",
|
||||
"lambda_ai",
|
||||
"inception",
|
||||
"vercel_ai_gateway",
|
||||
"wandb",
|
||||
"ovhcloud",
|
||||
@@ -676,6 +678,7 @@ OPENAI_CHAT_COMPLETION_PARAMS = [
|
||||
"extra_headers",
|
||||
"thinking",
|
||||
"web_search_options",
|
||||
"include_server_side_tool_invocations",
|
||||
"service_tier",
|
||||
"prompt_cache_key",
|
||||
"prompt_cache_retention",
|
||||
@@ -737,6 +740,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES = {
|
||||
"verbosity": None,
|
||||
"thinking": None,
|
||||
"web_search_options": None,
|
||||
"include_server_side_tool_invocations": None,
|
||||
"service_tier": None,
|
||||
"safety_identifier": None,
|
||||
"prompt_cache_key": None,
|
||||
@@ -771,6 +775,7 @@ openai_compatible_endpoints: List = [
|
||||
"https://api.moonshot.ai/v1",
|
||||
"https://api.publicai.co/v1",
|
||||
"https://api.synthetic.new/openai/v1",
|
||||
"https://serverless.tensormesh.ai/v1",
|
||||
"https://api.stima.tech/v1",
|
||||
"https://nano-gpt.com/api/v1",
|
||||
"https://api.poe.com/v1",
|
||||
@@ -778,6 +783,7 @@ openai_compatible_endpoints: List = [
|
||||
"https://api.v0.dev/v1",
|
||||
"https://api.morphllm.com/v1",
|
||||
"https://api.lambda.ai/v1",
|
||||
"https://api.inceptionlabs.ai/v1",
|
||||
"https://api.hyperbolic.xyz/v1",
|
||||
"https://ai-gateway.helicone.ai/",
|
||||
"https://ai-gateway.vercel.sh/v1",
|
||||
@@ -820,6 +826,7 @@ openai_compatible_providers: List = [
|
||||
"meta_llama",
|
||||
"publicai", # PublicAI - JSON-configured provider
|
||||
"synthetic", # Synthetic - JSON-configured provider
|
||||
"tensormesh", # Tensormesh - JSON-configured provider
|
||||
"apertis", # Apertis - JSON-configured provider
|
||||
"nano-gpt", # Nano-GPT - JSON-configured provider
|
||||
"poe", # Poe - JSON-configured provider
|
||||
@@ -833,6 +840,7 @@ openai_compatible_providers: List = [
|
||||
"helicone",
|
||||
"morph",
|
||||
"lambda_ai",
|
||||
"inception",
|
||||
"hyperbolic",
|
||||
"vercel_ai_gateway",
|
||||
"aiml",
|
||||
@@ -855,6 +863,7 @@ openai_text_completion_compatible_providers: List = (
|
||||
"moonshot",
|
||||
"publicai",
|
||||
"synthetic",
|
||||
"tensormesh",
|
||||
"apertis",
|
||||
"nano-gpt",
|
||||
"poe",
|
||||
@@ -868,6 +877,7 @@ openai_text_completion_compatible_providers: List = (
|
||||
_openai_like_providers: List = [
|
||||
"predibase",
|
||||
"databricks",
|
||||
"lemonade",
|
||||
"watsonx",
|
||||
] # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk
|
||||
# well supported replicate llms
|
||||
|
||||
@@ -2425,12 +2425,11 @@ class BaseTokenUsageProcessor:
|
||||
if not attr.startswith("_") and not callable(
|
||||
getattr(usage.completion_tokens_details, attr)
|
||||
):
|
||||
current_val = getattr(
|
||||
combined.completion_tokens_details, attr, 0
|
||||
current_val = (
|
||||
getattr(combined.completion_tokens_details, attr, 0) or 0
|
||||
)
|
||||
new_val = getattr(usage.completion_tokens_details, attr, 0)
|
||||
|
||||
if new_val is not None and current_val is not None:
|
||||
new_val = getattr(usage.completion_tokens_details, attr, 0) or 0
|
||||
if isinstance(new_val, (int, float)):
|
||||
setattr(
|
||||
combined.completion_tokens_details,
|
||||
attr,
|
||||
|
||||
@@ -1062,3 +1062,37 @@ class GuardrailInterventionNormalStringError(
|
||||
|
||||
def __repr__(self):
|
||||
return self.__str__()
|
||||
|
||||
|
||||
class SensitiveDataRouteException(Exception):
|
||||
"""
|
||||
Exception raised when a guardrail detects sensitive data and wants to reroute the request.
|
||||
|
||||
Instead of blocking the request, this exception signals that the request should be
|
||||
routed to a different model (typically an on-premise model for data privacy).
|
||||
|
||||
The proxy catches this exception and:
|
||||
1. Reroutes the current request to the specified model
|
||||
2. When sticky_session_routing is True, stores the routing decision in session
|
||||
cache so all subsequent requests in the same session are routed to the same model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
route_to_model: str,
|
||||
session_id: str,
|
||||
guardrail_name: Optional[str] = None,
|
||||
detection_info: Optional[Dict[str, Any]] = None,
|
||||
message: Optional[str] = None,
|
||||
sticky_session_routing: bool = True,
|
||||
):
|
||||
self.route_to_model = route_to_model
|
||||
self.session_id = session_id
|
||||
self.guardrail_name = guardrail_name
|
||||
self.detection_info = detection_info or {}
|
||||
self.sticky_session_routing = sticky_session_routing
|
||||
self.message = (
|
||||
message
|
||||
or f"Sensitive data detected by {guardrail_name}. Routing to model: {route_to_model}"
|
||||
)
|
||||
super().__init__(self.message)
|
||||
|
||||
@@ -4,6 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
from typing import (
|
||||
Any,
|
||||
Awaitable,
|
||||
@@ -16,7 +17,6 @@ from typing import (
|
||||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
import httpx
|
||||
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
|
||||
from mcp.client.sse import sse_client
|
||||
@@ -42,9 +42,8 @@ from mcp.types import (
|
||||
)
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
|
||||
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
from litellm.types.llms.custom_http import VerifyTypes
|
||||
from litellm.types.mcp import (
|
||||
@@ -61,13 +60,33 @@ def to_basic_auth(auth_value: str) -> str:
|
||||
return base64.b64encode(auth_value.encode("utf-8")).decode()
|
||||
|
||||
|
||||
def _strip_header_whitespace(headers: Dict[str, str]) -> Dict[str, str]:
|
||||
return {
|
||||
(key.strip() if isinstance(key, str) else key): (
|
||||
value.strip() if isinstance(value, str) else value
|
||||
)
|
||||
for key, value in headers.items()
|
||||
}
|
||||
|
||||
|
||||
def _first_non_cancelled_cause(exc: BaseException) -> Optional[BaseException]:
|
||||
queue: List[BaseException] = [exc]
|
||||
while queue:
|
||||
current = queue.pop(0)
|
||||
nested = getattr(current, "exceptions", None)
|
||||
if nested:
|
||||
queue.extend(nested)
|
||||
elif not isinstance(current, asyncio.CancelledError):
|
||||
return current
|
||||
return None
|
||||
|
||||
|
||||
TSessionResult = TypeVar("TSessionResult")
|
||||
|
||||
|
||||
class MCPSigV4Auth(httpx.Auth):
|
||||
"""
|
||||
httpx Auth class that signs each request with AWS SigV4.
|
||||
|
||||
This is used for MCP servers that require AWS SigV4 authentication,
|
||||
such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow()
|
||||
for every outgoing request, enabling per-request signature computation.
|
||||
@@ -92,10 +111,8 @@ class MCPSigV4Auth(httpx.Auth):
|
||||
"Missing botocore to use AWS SigV4 authentication. "
|
||||
"Run 'pip install boto3'."
|
||||
)
|
||||
|
||||
self.service_name = aws_service_name or "bedrock-agentcore"
|
||||
self.region_name = aws_region_name or "us-east-1"
|
||||
|
||||
# Note: os.environ/ prefixed values are already resolved by
|
||||
# ProxyConfig._check_for_os_environ_vars() at config load time.
|
||||
# Values arrive here as plain strings.
|
||||
@@ -143,20 +160,17 @@ class MCPSigV4Auth(httpx.Auth):
|
||||
session_name = (
|
||||
aws_session_name or f"litellm-mcp-{int(__import__('time').time())}"
|
||||
)
|
||||
|
||||
sts_kwargs: dict = {"region_name": aws_region_name}
|
||||
if aws_access_key_id and aws_secret_access_key:
|
||||
sts_kwargs["aws_access_key_id"] = aws_access_key_id
|
||||
sts_kwargs["aws_secret_access_key"] = aws_secret_access_key
|
||||
if aws_session_token:
|
||||
sts_kwargs["aws_session_token"] = aws_session_token
|
||||
|
||||
sts_client = boto3.client("sts", **sts_kwargs)
|
||||
sts_response = sts_client.assume_role(
|
||||
RoleArn=aws_role_name,
|
||||
RoleSessionName=session_name,
|
||||
)
|
||||
|
||||
sts_creds = sts_response["Credentials"]
|
||||
return Credentials(
|
||||
access_key=sts_creds["AccessKeyId"],
|
||||
@@ -178,17 +192,14 @@ class MCPSigV4Auth(httpx.Auth):
|
||||
data=request.content,
|
||||
headers=dict(request.headers),
|
||||
)
|
||||
|
||||
# Sign the request — SigV4Auth.add_auth() adds Authorization,
|
||||
# X-Amz-Date, and X-Amz-Security-Token (if session token present).
|
||||
# Host header is derived automatically from the URL.
|
||||
sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name)
|
||||
sigv4.add_auth(aws_request)
|
||||
|
||||
# Copy SigV4 headers back to the httpx request
|
||||
for header_name, header_value in aws_request.headers.items():
|
||||
request.headers[header_name] = header_value
|
||||
|
||||
yield request
|
||||
|
||||
|
||||
@@ -198,6 +209,8 @@ class MCPClient:
|
||||
SSE and HTTP transports
|
||||
Authentication via Bearer token, Basic Auth, or API Key
|
||||
Tool calling with error handling and result parsing
|
||||
Sampling callbacks for upstream server LLM requests
|
||||
Elicitation callbacks for upstream server user-input requests
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -211,6 +224,9 @@ class MCPClient:
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
ssl_verify: Optional[VerifyTypes] = None,
|
||||
aws_auth: Optional[httpx.Auth] = None,
|
||||
sampling_callback: Optional[Callable] = None,
|
||||
elicitation_callback: Optional[Callable] = None,
|
||||
logging_callback: Optional[Callable] = None,
|
||||
):
|
||||
self.server_url: str = server_url
|
||||
self.transport_type: MCPTransport = transport_type
|
||||
@@ -222,6 +238,9 @@ class MCPClient:
|
||||
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
|
||||
self._aws_auth: Optional[httpx.Auth] = aws_auth
|
||||
self._last_initialize_instructions: Optional[str] = None
|
||||
self._sampling_callback: Optional[Callable] = sampling_callback
|
||||
self._elicitation_callback: Optional[Callable] = elicitation_callback
|
||||
self._logging_callback: Optional[Callable] = logging_callback
|
||||
# handle the basic auth value if provided
|
||||
if auth_value:
|
||||
self.update_auth_value(auth_value)
|
||||
@@ -231,23 +250,20 @@ class MCPClient:
|
||||
) -> Tuple[Any, Optional[httpx.AsyncClient]]:
|
||||
"""
|
||||
Create the appropriate transport context based on transport type.
|
||||
|
||||
Returns:
|
||||
Tuple of (transport_context, http_client).
|
||||
http_client is only set for HTTP transport and needs cleanup.
|
||||
"""
|
||||
http_client: Optional[httpx.AsyncClient] = None
|
||||
|
||||
if self.transport_type == MCPTransport.stdio:
|
||||
if not self.stdio_config:
|
||||
raise ValueError("stdio_config is required for stdio transport")
|
||||
server_params = StdioServerParameters(
|
||||
command=self.stdio_config.get("command", ""),
|
||||
args=self.stdio_config.get("args", []),
|
||||
env=self.stdio_config.get("env", {}),
|
||||
env=self._get_safe_stdio_env(self.stdio_config.get("env")),
|
||||
)
|
||||
return stdio_client(server_params), None
|
||||
|
||||
if self.transport_type == MCPTransport.sse:
|
||||
headers = self._get_auth_headers()
|
||||
httpx_client_factory = self._create_httpx_client_factory()
|
||||
@@ -260,14 +276,12 @@ class MCPClient:
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
# HTTP transport (default)
|
||||
if streamable_http_client is None:
|
||||
raise ImportError(
|
||||
"streamable_http_client is not available. "
|
||||
"Please install mcp with HTTP support."
|
||||
)
|
||||
|
||||
headers = self._get_auth_headers()
|
||||
httpx_client_factory = self._create_httpx_client_factory()
|
||||
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
|
||||
@@ -281,6 +295,54 @@ class MCPClient:
|
||||
)
|
||||
return transport_ctx, http_client
|
||||
|
||||
def _get_safe_stdio_env(
|
||||
self, provided_env: Optional[Dict[str, str]]
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Return a safe environment for the stdio subprocess.
|
||||
|
||||
If provided_env is set, we use it as-is.
|
||||
If provided_env is None, we return a minimal allowlist from the parent environment
|
||||
to avoid leaking sensitive LiteLLM keys (OPENAI_API_KEY, etc.) to sub-processes.
|
||||
"""
|
||||
if provided_env is not None:
|
||||
return provided_env
|
||||
|
||||
# Minimal allowlist of safe/standard environment variables
|
||||
safe_keys = {
|
||||
"PATH",
|
||||
"HOME",
|
||||
"USER",
|
||||
"LOGNAME",
|
||||
"TMPDIR",
|
||||
"TMP",
|
||||
"TEMP",
|
||||
"SHELL",
|
||||
"LANG",
|
||||
"LC_ALL",
|
||||
# Node/Package manager caches
|
||||
"NPM_CONFIG_CACHE",
|
||||
"PNPM_HOME",
|
||||
"XDG_CACHE_HOME",
|
||||
"XDG_CONFIG_HOME",
|
||||
"XDG_DATA_HOME",
|
||||
# System info
|
||||
"SYSTEMROOT",
|
||||
"COMSPEC",
|
||||
"PATHEXT",
|
||||
"WINDIR",
|
||||
}
|
||||
|
||||
safe_env = {}
|
||||
for key in safe_keys:
|
||||
if key in os.environ:
|
||||
safe_env[key] = os.environ[key]
|
||||
|
||||
if "NPM_CONFIG_CACHE" not in safe_env:
|
||||
safe_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
|
||||
|
||||
return safe_env
|
||||
|
||||
async def _execute_session_operation(
|
||||
self,
|
||||
transport_ctx: Any,
|
||||
@@ -288,13 +350,24 @@ class MCPClient:
|
||||
) -> TSessionResult:
|
||||
"""
|
||||
Execute an operation within a transport and session context.
|
||||
|
||||
Handles entering/exiting contexts and running the operation.
|
||||
Passes sampling/elicitation/logging callbacks to the ClientSession
|
||||
so that upstream MCP servers can request LLM inference (sampling),
|
||||
user input (elicitation), or send log messages.
|
||||
"""
|
||||
transport = await transport_ctx.__aenter__()
|
||||
in_flight_error: Optional[BaseException] = None
|
||||
try:
|
||||
read_stream, write_stream = transport[0], transport[1]
|
||||
session_ctx = ClientSession(read_stream, write_stream)
|
||||
# Build session kwargs with optional callbacks
|
||||
session_kwargs: Dict[str, Any] = {}
|
||||
if self._sampling_callback is not None:
|
||||
session_kwargs["sampling_callback"] = self._sampling_callback
|
||||
if self._elicitation_callback is not None:
|
||||
session_kwargs["elicitation_callback"] = self._elicitation_callback
|
||||
if self._logging_callback is not None:
|
||||
session_kwargs["logging_callback"] = self._logging_callback
|
||||
session_ctx = ClientSession(read_stream, write_stream, **session_kwargs)
|
||||
session = await session_ctx.__aenter__()
|
||||
try:
|
||||
init_result = await session.initialize()
|
||||
@@ -309,11 +382,21 @@ class MCPClient:
|
||||
await session_ctx.__aexit__(None, None, None)
|
||||
except BaseException as e:
|
||||
verbose_logger.debug(f"Error during session context exit: {e}")
|
||||
except BaseException as e:
|
||||
in_flight_error = e
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
await transport_ctx.__aexit__(None, None, None)
|
||||
except BaseException as e:
|
||||
verbose_logger.debug(f"Error during transport context exit: {e}")
|
||||
except BaseException as exit_error:
|
||||
verbose_logger.debug(
|
||||
f"Error during transport context exit: {exit_error}"
|
||||
)
|
||||
root_cause = _first_non_cancelled_cause(exit_error)
|
||||
if root_cause is not None and isinstance(
|
||||
in_flight_error, asyncio.CancelledError
|
||||
):
|
||||
raise root_cause from in_flight_error
|
||||
|
||||
async def run_with_session(
|
||||
self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]
|
||||
@@ -351,7 +434,6 @@ class MCPClient:
|
||||
def _get_auth_headers(self) -> dict:
|
||||
"""Generate authentication headers based on auth type."""
|
||||
headers = {}
|
||||
|
||||
if self._mcp_auth_value:
|
||||
if isinstance(self._mcp_auth_value, str):
|
||||
if self.auth_type == MCPAuth.bearer_token:
|
||||
@@ -373,17 +455,14 @@ class MCPClient:
|
||||
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
|
||||
# signing (including the body hash), so it uses httpx.Auth flow instead
|
||||
# of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
|
||||
|
||||
# update the headers with the extra headers
|
||||
if self.extra_headers:
|
||||
headers.update(self.extra_headers)
|
||||
|
||||
return headers
|
||||
return _strip_header_whitespace(headers)
|
||||
|
||||
def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]:
|
||||
"""
|
||||
Create a custom httpx client factory that uses LiteLLM's SSL configuration.
|
||||
|
||||
This factory follows the same CA bundle path logic as http_handler.py:
|
||||
1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle)
|
||||
2. Check SSL_VERIFY environment variable
|
||||
@@ -400,17 +479,14 @@ class MCPClient:
|
||||
"""Create an httpx.AsyncClient with LiteLLM's SSL configuration."""
|
||||
# Get unified SSL configuration using the same logic as http_handler.py
|
||||
ssl_config = get_ssl_configuration(self.ssl_verify)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"MCP client using SSL configuration: {type(ssl_config).__name__}"
|
||||
)
|
||||
|
||||
# Use SigV4 auth if configured and no explicit auth provided.
|
||||
# The MCP SDK's sse_client and streamable_http_client call this
|
||||
# factory without passing auth=, so self._aws_auth is used.
|
||||
# For non-SigV4 clients, self._aws_auth is None — no behavior change.
|
||||
effective_auth = auth if auth is not None else self._aws_auth
|
||||
|
||||
return httpx.AsyncClient(
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
@@ -421,8 +497,16 @@ class MCPClient:
|
||||
|
||||
return factory
|
||||
|
||||
async def list_tools(self) -> List[MCPTool]:
|
||||
"""List available tools from the server."""
|
||||
async def list_tools(self, raise_on_error: bool = False) -> List[MCPTool]:
|
||||
"""List available tools from the server.
|
||||
|
||||
Args:
|
||||
raise_on_error: When True, re-raise exceptions instead of returning
|
||||
an empty list. Used by the proxy's pass-through MCP flow so it
|
||||
can surface upstream HTTP 401 responses as a proper 401 to the
|
||||
MCP client (triggering the upstream OAuth flow) rather than
|
||||
masking them as "connected, no tools".
|
||||
"""
|
||||
verbose_logger.debug(
|
||||
f"MCP client listing tools from {self.server_url or 'stdio'}"
|
||||
)
|
||||
@@ -450,7 +534,6 @@ class MCPClient:
|
||||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
@@ -458,6 +541,8 @@ class MCPClient:
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
||||
if raise_on_error:
|
||||
raise
|
||||
# Return empty list instead of raising to allow graceful degradation
|
||||
return []
|
||||
|
||||
@@ -481,7 +566,6 @@ class MCPClient:
|
||||
f"MCP Tool '{call_tool_request_params.name}' progress: "
|
||||
f"{progress}/{total} ({percentage:.0f}%) - {message or ''}"
|
||||
)
|
||||
|
||||
# Forward to Host if callback provided
|
||||
if host_progress_callback:
|
||||
try:
|
||||
@@ -504,14 +588,15 @@ class MCPClient:
|
||||
)
|
||||
return tool_result
|
||||
except asyncio.CancelledError:
|
||||
verbose_logger.warning("MCP client tool call was cancelled")
|
||||
verbose_logger.warning(
|
||||
f"MCP client tool call timed out after {self.timeout}s for {self.server_url}"
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_trace = traceback.format_exc()
|
||||
verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}")
|
||||
|
||||
# Log detailed error information
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.error(
|
||||
@@ -522,14 +607,12 @@ class MCPClient:
|
||||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream - "
|
||||
"the MCP server may have crashed, disconnected, or timed out."
|
||||
)
|
||||
|
||||
# Return a default error result instead of raising
|
||||
return MCPCallToolResult(
|
||||
content=[
|
||||
@@ -567,14 +650,12 @@ class MCPClient:
|
||||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during list_tools - "
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
||||
# Return empty list instead of raising to allow graceful degradation
|
||||
return []
|
||||
|
||||
@@ -607,7 +688,6 @@ class MCPClient:
|
||||
|
||||
error_trace = traceback.format_exc()
|
||||
verbose_logger.debug(f"MCP client get_prompt traceback:\n{error_trace}")
|
||||
|
||||
# Log detailed error information
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.error(
|
||||
@@ -618,14 +698,12 @@ class MCPClient:
|
||||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during get_prompt - "
|
||||
"the MCP server may have crashed, disconnected, or timed out."
|
||||
)
|
||||
|
||||
raise
|
||||
|
||||
async def list_resources(self) -> list[Resource]:
|
||||
@@ -657,14 +735,12 @@ class MCPClient:
|
||||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during list_resources - "
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
||||
# Return empty list instead of raising to allow graceful degradation
|
||||
return []
|
||||
|
||||
@@ -699,14 +775,12 @@ class MCPClient:
|
||||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during list_resource_templates - "
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
||||
# Return empty list instead of raising to allow graceful degradation
|
||||
return []
|
||||
|
||||
@@ -732,7 +806,6 @@ class MCPClient:
|
||||
|
||||
error_trace = traceback.format_exc()
|
||||
verbose_logger.debug(f"MCP client read_resource traceback:\n{error_trace}")
|
||||
|
||||
# Log detailed error information
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.error(
|
||||
@@ -743,12 +816,10 @@ class MCPClient:
|
||||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during read_resource - "
|
||||
"the MCP server may have crashed, disconnected, or timed out."
|
||||
)
|
||||
|
||||
raise
|
||||
|
||||
@@ -8,18 +8,23 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor
|
||||
BaseLLMObsOTELAttributes,
|
||||
safe_set_attribute,
|
||||
)
|
||||
from litellm.litellm_core_utils.redact_messages import (
|
||||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
from litellm.integrations._types.open_inference import (
|
||||
MessageAttributes,
|
||||
ImageAttributes,
|
||||
SpanAttributes,
|
||||
AudioAttributes,
|
||||
EmbeddingAttributes,
|
||||
ImageAttributes,
|
||||
MessageAttributes,
|
||||
MessageContentAttributes,
|
||||
OpenInferenceSpanKindValues,
|
||||
SpanAttributes,
|
||||
ToolCallAttributes,
|
||||
)
|
||||
|
||||
|
||||
@@ -53,40 +58,24 @@ class ArizeOTELAttributes(BaseLLMObsOTELAttributes):
|
||||
msg.get("content", ""),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@override
|
||||
def set_response_output_messages(span: "Span", response_obj):
|
||||
"""
|
||||
Sets output message attributes on the span from the LLM response.
|
||||
Args:
|
||||
span: The OpenTelemetry span to set attributes on
|
||||
response_obj: The response object containing choices with messages
|
||||
"""
|
||||
from litellm.integrations._types.open_inference import (
|
||||
MessageAttributes,
|
||||
SpanAttributes,
|
||||
)
|
||||
# Additive: emit structured tool_calls / multimodal content
|
||||
# so Arize/Phoenix can render tool-using and image-bearing
|
||||
# turns. These set NEW attribute keys (MESSAGE_TOOL_CALLS /
|
||||
# MESSAGE_NAME / MESSAGE_TOOL_CALL_ID / MESSAGE_CONTENTS.*) —
|
||||
# never replace the MESSAGE_CONTENT write above.
|
||||
_safe_emit(
|
||||
f"input message extras (idx={idx})",
|
||||
_emit_input_message_extras,
|
||||
span,
|
||||
prefix,
|
||||
msg,
|
||||
)
|
||||
|
||||
for idx, choice in enumerate(response_obj.get("choices", [])):
|
||||
response_message = choice.get("message", {})
|
||||
safe_set_attribute(
|
||||
span,
|
||||
SpanAttributes.OUTPUT_VALUE,
|
||||
response_message.get("content", ""),
|
||||
)
|
||||
|
||||
# This shows up under `output_messages` tab on the span page.
|
||||
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}"
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
|
||||
response_message.get("role"),
|
||||
)
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
|
||||
response_message.get("content", ""),
|
||||
)
|
||||
# Note: `BaseLLMObsOTELAttributes.set_response_output_messages` is not
|
||||
# overridden here. The live code path uses `_set_choice_outputs` (called
|
||||
# via `_set_response_attributes` from `set_attributes`) which handles
|
||||
# tool_calls, multimodal output, embeddings, audio, images, and structured
|
||||
# outputs in a single place.
|
||||
|
||||
|
||||
def _set_response_attributes(span: "Span", response_obj):
|
||||
@@ -106,11 +95,17 @@ def _set_response_attributes(span: "Span", response_obj):
|
||||
def _set_choice_outputs(span: "Span", response_obj, msg_attrs, span_attrs):
|
||||
for idx, choice in enumerate(response_obj.get("choices", [])):
|
||||
response_message = choice.get("message", {})
|
||||
safe_set_attribute(
|
||||
span,
|
||||
span_attrs.OUTPUT_VALUE,
|
||||
response_message.get("content", ""),
|
||||
)
|
||||
content = response_message.get("content", "")
|
||||
|
||||
# Tool-only assistant responses have empty content; serialize the
|
||||
# tool_calls into OUTPUT_VALUE so Arize's "Output" pane isn't blank.
|
||||
output_value = content
|
||||
if not output_value:
|
||||
tool_calls = _get_tool_calls(response_message)
|
||||
if tool_calls:
|
||||
output_value = _summarize_tool_calls_for_output(tool_calls)
|
||||
|
||||
safe_set_attribute(span, span_attrs.OUTPUT_VALUE, output_value)
|
||||
prefix = f"{span_attrs.LLM_OUTPUT_MESSAGES}.{idx}"
|
||||
safe_set_attribute(
|
||||
span,
|
||||
@@ -120,7 +115,18 @@ def _set_choice_outputs(span: "Span", response_obj, msg_attrs, span_attrs):
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{msg_attrs.MESSAGE_CONTENT}",
|
||||
response_message.get("content", ""),
|
||||
content,
|
||||
)
|
||||
|
||||
# Additive: emit assistant tool_calls so tool-using turns render in
|
||||
# Arize/Phoenix. Sets new MESSAGE_TOOL_CALLS keys only — does not
|
||||
# change MESSAGE_CONTENT/MESSAGE_ROLE writes above.
|
||||
_safe_emit(
|
||||
f"output tool_calls (idx={idx})",
|
||||
_emit_message_tool_calls,
|
||||
span,
|
||||
prefix,
|
||||
response_message,
|
||||
)
|
||||
|
||||
|
||||
@@ -278,6 +284,43 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs):
|
||||
reasoning_tokens,
|
||||
)
|
||||
|
||||
# Additive: cache token breakdown so prompt-caching savings render in
|
||||
# Arize. Sources covered:
|
||||
# - OpenAI Chat Completions: `prompt_tokens_details.cached_tokens`
|
||||
# - Anthropic / Bedrock-Anthropic: `cache_read_input_tokens`,
|
||||
# `cache_creation_input_tokens`
|
||||
# All emits are conditional, so when none of these fields exist (the
|
||||
# situation in the existing test fixtures) no extra attributes are set.
|
||||
prompt_token_details = _safe_get(usage, "prompt_tokens_details") or _safe_get(
|
||||
usage, "input_tokens_details"
|
||||
)
|
||||
cache_read = _safe_get(prompt_token_details, "cached_tokens") or _safe_get(
|
||||
usage, "cache_read_input_tokens"
|
||||
)
|
||||
if cache_read:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ,
|
||||
cache_read,
|
||||
)
|
||||
# Anthropic / Bedrock-Anthropic only — OpenAI's `prompt_tokens_details`
|
||||
# does not expose a cache-write count, so we read straight off `usage`.
|
||||
cache_write = _safe_get(usage, "cache_creation_input_tokens")
|
||||
if cache_write:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE,
|
||||
cache_write,
|
||||
)
|
||||
|
||||
audio_prompt_tokens = _safe_get(prompt_token_details, "audio_tokens")
|
||||
if audio_prompt_tokens:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
span_attrs.LLM_TOKEN_COUNT_PROMPT_DETAILS_AUDIO,
|
||||
audio_prompt_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
|
||||
"""
|
||||
@@ -321,6 +364,10 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
|
||||
"videos",
|
||||
"realtime",
|
||||
"pass_through",
|
||||
# `passthrough` (no underscore) is what real call_types use:
|
||||
# `allm_passthrough_route`, `llm_passthrough_route`. Without
|
||||
# this they fell through to UNKNOWN, blanking span.kind.
|
||||
"passthrough",
|
||||
"anthropic_messages",
|
||||
"ocr",
|
||||
)
|
||||
@@ -396,6 +443,18 @@ def set_attributes(
|
||||
"""
|
||||
Populates span with OpenInference-compliant LLM attributes for Arize and Phoenix tracing.
|
||||
"""
|
||||
# Coerce non-dict response objects (e.g. httpx.Response from passthrough
|
||||
# routes) into a dict so downstream `.get()` calls don't crash. Existing
|
||||
# dict / `.get()`-bearing objects (incl. Pydantic OpenAI Responses API
|
||||
# models) are returned unchanged, preserving the existing test behavior.
|
||||
response_obj_for_attrs = _coerce_response_obj_for_attrs(response_obj)
|
||||
|
||||
# Set span.kind defensively before anything else. If a downstream step
|
||||
# throws, the span still has a kind so Arize can render it correctly
|
||||
# (an LLM call instead of UNKNOWN). This is the single source of truth
|
||||
# for span.kind — no late re-write happens below.
|
||||
_safe_emit("early span kind", _set_early_span_kind, span, kwargs)
|
||||
|
||||
try:
|
||||
optional_params = _sanitize_optional_params(kwargs.get("optional_params"))
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
@@ -415,25 +474,22 @@ def set_attributes(
|
||||
metadata_tools = _extract_metadata_tools(metadata)
|
||||
optional_tools = _extract_optional_tools(optional_params)
|
||||
|
||||
call_type = standard_logging_payload.get("call_type")
|
||||
_set_request_attributes(
|
||||
span=span,
|
||||
kwargs=kwargs,
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
response_obj=response_obj,
|
||||
response_obj=response_obj_for_attrs,
|
||||
span_attrs=SpanAttributes,
|
||||
)
|
||||
|
||||
span_kind = _infer_open_inference_span_kind(call_type=call_type)
|
||||
# span.kind was already set above by `_set_early_span_kind`. We do
|
||||
# NOT re-write it here based on tool presence: a chat completion
|
||||
# that passes `tools=[...]` (or returns `tool_calls`) is still an
|
||||
# LLM call per the OpenInference spec — TOOL is reserved for actual
|
||||
# tool execution spans, not LLM calls that request tools.
|
||||
_set_tool_attributes(span, optional_tools, metadata_tools)
|
||||
if (
|
||||
optional_tools or metadata_tools
|
||||
) and span_kind != OpenInferenceSpanKindValues.TOOL.value:
|
||||
span_kind = OpenInferenceSpanKindValues.TOOL.value
|
||||
|
||||
safe_set_attribute(span, SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind)
|
||||
attributes.set_messages(span, kwargs)
|
||||
|
||||
model_params = (
|
||||
@@ -443,7 +499,7 @@ def set_attributes(
|
||||
)
|
||||
_set_model_params(span, model_params, SpanAttributes)
|
||||
|
||||
_set_response_attributes(span=span, response_obj=response_obj)
|
||||
_set_response_attributes(span=span, response_obj=response_obj_for_attrs)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.error(
|
||||
@@ -452,6 +508,22 @@ def set_attributes(
|
||||
if hasattr(span, "record_exception"):
|
||||
span.record_exception(e)
|
||||
|
||||
# Additive emitters. Each is independently guarded so a failure can never
|
||||
# blank the attributes set by the main try-block above. New attributes are
|
||||
# written under new keys; existing attributes are not overwritten.
|
||||
slp = kwargs.get("standard_logging_object")
|
||||
_safe_emit("session/user attrs", _set_session_and_user_attrs, span, kwargs, slp)
|
||||
_safe_emit("response cost", _set_response_cost_attr, span, slp)
|
||||
_safe_emit(
|
||||
"passthrough normalization",
|
||||
_maybe_normalize_passthrough,
|
||||
span,
|
||||
kwargs,
|
||||
response_obj,
|
||||
response_obj_for_attrs,
|
||||
slp,
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_optional_params(optional_params: Optional[dict]) -> dict:
|
||||
if not isinstance(optional_params, dict):
|
||||
@@ -534,3 +606,529 @@ def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) ->
|
||||
user_id = model_params.get("user")
|
||||
if user_id is not None:
|
||||
safe_set_attribute(span, span_attrs.USER_ID, user_id)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Additive rendering helpers (introduced to enhance Arize/Phoenix rendering
|
||||
# without changing any previously-emitted attribute keys or values).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _safe_emit(label: str, fn, *args, **kwargs) -> None:
|
||||
"""Run an additive attribute emitter, swallowing any error so it cannot
|
||||
blank attributes set elsewhere on the span. Failures are logged at debug.
|
||||
"""
|
||||
try:
|
||||
fn(*args, **kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("[Arize] %s skipped: %s", label, e)
|
||||
|
||||
|
||||
def _set_early_span_kind(span: "Span", kwargs: dict) -> None:
|
||||
"""Defensively set OPENINFERENCE_SPAN_KIND before any other logic runs."""
|
||||
slp = kwargs.get("standard_logging_object")
|
||||
call_type = slp.get("call_type") if isinstance(slp, dict) else None
|
||||
safe_set_attribute(
|
||||
span,
|
||||
SpanAttributes.OPENINFERENCE_SPAN_KIND,
|
||||
_infer_open_inference_span_kind(call_type=call_type),
|
||||
)
|
||||
|
||||
|
||||
def _coerce_response_obj_for_attrs(response_obj):
|
||||
"""Return a `.get`-compatible view of `response_obj` when possible.
|
||||
|
||||
- dicts and Pydantic models that already expose `.get` are returned
|
||||
unchanged (preserves all current behavior, including the Responses API
|
||||
flow which relies on Pydantic attribute access).
|
||||
- `httpx.Response` and other text-only responses (passthrough routes)
|
||||
are JSON-decoded so the standard extraction paths can read fields like
|
||||
`id`, `model`, and `usage`. On failure the original object is returned
|
||||
so behavior is no worse than today.
|
||||
"""
|
||||
if response_obj is None or hasattr(response_obj, "get"):
|
||||
return response_obj
|
||||
text = getattr(response_obj, "text", None)
|
||||
if isinstance(text, str) and text:
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
except Exception:
|
||||
pass
|
||||
return response_obj
|
||||
|
||||
|
||||
def _coerce_text(value) -> Optional[str]:
|
||||
"""Best-effort text extraction from a message-content value.
|
||||
|
||||
Returns None when no textual portion can be derived. Handles:
|
||||
- plain strings
|
||||
- lists of OpenAI-style content parts (`{"type": "text", "text": ...}`)
|
||||
- lists of Anthropic-style content parts (`{"type": "text", "text": ...}`
|
||||
or `{"type": "input_text", "text": ...}`)
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
parts = []
|
||||
for part in value:
|
||||
if isinstance(part, str):
|
||||
parts.append(part)
|
||||
elif isinstance(part, dict):
|
||||
text = part.get("text") or part.get("input_text")
|
||||
if isinstance(text, str):
|
||||
parts.append(text)
|
||||
if parts:
|
||||
return "\n".join(parts)
|
||||
return None
|
||||
|
||||
|
||||
def _to_plain_dict(value):
|
||||
"""Best-effort: coerce a value (Pydantic model / dict / None) to a dict.
|
||||
|
||||
Returns the original value when no safe conversion exists. Used to bridge
|
||||
OpenAI Pydantic message/tool_call objects into the dict-based helpers.
|
||||
"""
|
||||
if value is None or isinstance(value, dict):
|
||||
return value
|
||||
model_dump = getattr(value, "model_dump", None)
|
||||
if callable(model_dump):
|
||||
try:
|
||||
return model_dump()
|
||||
except Exception:
|
||||
pass
|
||||
return value
|
||||
|
||||
|
||||
def _get_tool_calls(message) -> Optional[list]:
|
||||
"""Return ``message.tool_calls`` only when it's a non-empty list.
|
||||
|
||||
Works for dicts and Pydantic message objects via ``_safe_get``.
|
||||
"""
|
||||
tool_calls = _safe_get(message, "tool_calls")
|
||||
return tool_calls if isinstance(tool_calls, list) and tool_calls else None
|
||||
|
||||
|
||||
def _normalize_tool_call(raw_tc) -> Optional[Dict[str, Any]]:
|
||||
"""Normalize a single tool_call (dict or Pydantic) into a stable shape:
|
||||
|
||||
{"id": str|None, "type": str, "function": {"name": str|None, "arguments": str|None}}
|
||||
|
||||
Arguments are coerced to a JSON string per OpenInference convention.
|
||||
Returns ``None`` when ``raw_tc`` cannot be coerced to a dict.
|
||||
"""
|
||||
tc = _to_plain_dict(raw_tc)
|
||||
if not isinstance(tc, dict):
|
||||
return None
|
||||
function = _to_plain_dict(tc.get("function"))
|
||||
name = function.get("name") if isinstance(function, dict) else None
|
||||
args = function.get("arguments") if isinstance(function, dict) else None
|
||||
if args is not None and not isinstance(args, str):
|
||||
try:
|
||||
args = json.dumps(args)
|
||||
except Exception:
|
||||
args = str(args)
|
||||
return {
|
||||
"id": tc.get("id"),
|
||||
"type": tc.get("type", "function"),
|
||||
"function": {"name": name, "arguments": args},
|
||||
}
|
||||
|
||||
|
||||
def _summarize_tool_calls_for_output(tool_calls) -> str:
|
||||
"""Render a tool_calls list as a compact JSON string for OUTPUT_VALUE.
|
||||
|
||||
Best-effort: returns ``str(tool_calls)`` if anything unexpected happens
|
||||
so OUTPUT_VALUE is never blanked on a malformed payload.
|
||||
"""
|
||||
try:
|
||||
normalized = [n for n in (_normalize_tool_call(tc) for tc in tool_calls) if n]
|
||||
return json.dumps({"tool_calls": normalized})
|
||||
except Exception:
|
||||
return str(tool_calls)
|
||||
|
||||
|
||||
def _emit_message_tool_calls(span: "Span", prefix: str, message) -> None:
|
||||
"""Emit ``MESSAGE_TOOL_CALLS.*`` for an assistant message that requested
|
||||
tool calls. Pure addition: only writes when ``tool_calls`` is non-empty.
|
||||
|
||||
Accepts dicts or Pydantic message objects (e.g. ``litellm.Message``); the
|
||||
same applies to each tool_call entry.
|
||||
"""
|
||||
tool_calls = _get_tool_calls(message)
|
||||
if not tool_calls:
|
||||
return
|
||||
for tc_idx, raw_tc in enumerate(tool_calls):
|
||||
tc = _normalize_tool_call(raw_tc)
|
||||
if tc is None:
|
||||
continue
|
||||
tc_prefix = f"{prefix}.{MessageAttributes.MESSAGE_TOOL_CALLS}.{tc_idx}"
|
||||
if tc["id"]:
|
||||
safe_set_attribute(
|
||||
span, f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_ID}", tc["id"]
|
||||
)
|
||||
fn = tc["function"]
|
||||
if fn["name"]:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}",
|
||||
fn["name"],
|
||||
)
|
||||
if fn["arguments"] is not None:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{tc_prefix}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
|
||||
fn["arguments"],
|
||||
)
|
||||
|
||||
|
||||
def _emit_input_message_extras(span: "Span", prefix: str, message: dict) -> None:
|
||||
"""Emit additive attributes for an input message:
|
||||
|
||||
- `MESSAGE_NAME` and `MESSAGE_TOOL_CALL_ID` (commonly set on tool-result
|
||||
messages so traces show which tool produced which result).
|
||||
- `MESSAGE_TOOL_CALLS.*` when an assistant message requested tools.
|
||||
- `MESSAGE_CONTENTS.*` structured content for list-shaped content
|
||||
(multimodal text + image parts). The plain `MESSAGE_CONTENT` write is
|
||||
still performed by the caller, so renderers that only read the legacy
|
||||
key continue to work.
|
||||
"""
|
||||
if not isinstance(message, dict):
|
||||
return
|
||||
|
||||
name = message.get("name")
|
||||
if name:
|
||||
safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_NAME}", name)
|
||||
|
||||
tool_call_id = message.get("tool_call_id")
|
||||
if tool_call_id:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_TOOL_CALL_ID}",
|
||||
tool_call_id,
|
||||
)
|
||||
|
||||
_emit_message_tool_calls(span, prefix, message)
|
||||
|
||||
content = message.get("content")
|
||||
if isinstance(content, list):
|
||||
contents_prefix = f"{prefix}.{MessageAttributes.MESSAGE_CONTENTS}"
|
||||
for part_idx, part in enumerate(content):
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
part_prefix = f"{contents_prefix}.{part_idx}"
|
||||
part_type = part.get("type")
|
||||
if part_type in ("text", "input_text"):
|
||||
text = part.get("text")
|
||||
if isinstance(text, str):
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{part_prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TYPE}",
|
||||
"text",
|
||||
)
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{part_prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TEXT}",
|
||||
text,
|
||||
)
|
||||
elif part_type in ("image_url", "image", "input_image"):
|
||||
url = None
|
||||
image = part.get("image_url")
|
||||
if isinstance(image, dict):
|
||||
url = image.get("url")
|
||||
elif isinstance(image, str):
|
||||
url = image
|
||||
if not url:
|
||||
# Anthropic-style source.{type=base64,media_type,data}
|
||||
source = part.get("source")
|
||||
if isinstance(source, dict) and source.get("data"):
|
||||
media_type = source.get("media_type", "image/jpeg")
|
||||
url = f"data:{media_type};base64,{source['data']}"
|
||||
elif isinstance(part.get("url"), str):
|
||||
url = part["url"]
|
||||
if url:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{part_prefix}.{MessageContentAttributes.MESSAGE_CONTENT_TYPE}",
|
||||
"image",
|
||||
)
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{part_prefix}.message_content.image.image.url",
|
||||
url,
|
||||
)
|
||||
|
||||
|
||||
def _set_session_and_user_attrs(
|
||||
span: "Span", kwargs: dict, standard_logging_payload
|
||||
) -> None:
|
||||
"""Emit `SESSION_ID` / `USER_ID` / team metadata when source data exists.
|
||||
|
||||
`SESSION_ID` is emitted only when an explicit end-user identifier exists
|
||||
(`metadata.user_api_key_end_user_id`). We deliberately do NOT fall back
|
||||
to `trace_id`, because that would create a distinct "session" for every
|
||||
single request and distort Arize's Session-grouping analytics. The
|
||||
`trace_id` is still emitted under its own `litellm.trace_id` key so
|
||||
spans remain filterable by trace.
|
||||
|
||||
USER_ID is *only* emitted when no upstream path (model_params.user or
|
||||
optional_params.user) has already set it, to avoid overwriting an
|
||||
existing value with a possibly-different one from API-key metadata.
|
||||
"""
|
||||
if not isinstance(standard_logging_payload, dict):
|
||||
return
|
||||
metadata = standard_logging_payload.get("metadata") or {}
|
||||
if not isinstance(metadata, dict):
|
||||
return
|
||||
|
||||
session_id = metadata.get("user_api_key_end_user_id")
|
||||
if session_id:
|
||||
safe_set_attribute(span, SpanAttributes.SESSION_ID, str(session_id))
|
||||
|
||||
trace_id = standard_logging_payload.get("trace_id")
|
||||
if trace_id:
|
||||
safe_set_attribute(span, "litellm.trace_id", str(trace_id))
|
||||
|
||||
optional_params = kwargs.get("optional_params") or {}
|
||||
model_params = standard_logging_payload.get("model_parameters") or {}
|
||||
has_user_already = bool(
|
||||
(isinstance(optional_params, dict) and optional_params.get("user"))
|
||||
or (isinstance(model_params, dict) and model_params.get("user"))
|
||||
)
|
||||
if not has_user_already:
|
||||
user_id = metadata.get("user_api_key_user_id")
|
||||
if user_id:
|
||||
safe_set_attribute(span, SpanAttributes.USER_ID, str(user_id))
|
||||
|
||||
team_id = metadata.get("user_api_key_team_id")
|
||||
if team_id:
|
||||
safe_set_attribute(span, "litellm.team_id", str(team_id))
|
||||
team_alias = metadata.get("user_api_key_team_alias")
|
||||
if team_alias:
|
||||
safe_set_attribute(span, "litellm.team_alias", str(team_alias))
|
||||
key_alias = metadata.get("user_api_key_alias")
|
||||
if key_alias:
|
||||
safe_set_attribute(span, "litellm.key_alias", str(key_alias))
|
||||
|
||||
|
||||
def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None:
|
||||
"""Emit cost attributes from the StandardLoggingPayload when present.
|
||||
|
||||
Uses the OpenInference `llm.cost.total` key so Arize / Phoenix can
|
||||
surface the cost in their "Total Cost" column. LiteLLM only tracks a
|
||||
single total in `StandardLoggingPayload.response_cost`, so we cannot
|
||||
split it into prompt/completion. We also keep the legacy
|
||||
`llm.response.cost` key for back-compat with any consumer querying it.
|
||||
"""
|
||||
if not isinstance(standard_logging_payload, dict):
|
||||
return
|
||||
cost = standard_logging_payload.get("response_cost")
|
||||
if cost is None:
|
||||
return
|
||||
try:
|
||||
cost_value = float(cost)
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
safe_set_attribute(span, "llm.cost.total", cost_value)
|
||||
safe_set_attribute(span, "llm.response.cost", cost_value)
|
||||
|
||||
|
||||
def _is_passthrough_call_type(call_type: Optional[str]) -> bool:
|
||||
if not call_type:
|
||||
return False
|
||||
lowered = str(call_type).lower()
|
||||
return "passthrough" in lowered or "pass_through" in lowered
|
||||
|
||||
|
||||
def _maybe_normalize_passthrough(
|
||||
span: "Span",
|
||||
kwargs: dict,
|
||||
raw_response_obj,
|
||||
coerced_response_obj,
|
||||
standard_logging_payload,
|
||||
) -> None:
|
||||
"""Surface input/output text for passthrough routes (e.g. Bedrock
|
||||
InvokeModel) so the parent span renders as more than `usage` numbers.
|
||||
|
||||
Only runs when `call_type` is a passthrough variant. Reads from:
|
||||
- `kwargs["additional_args"]["complete_input_dict"]` for input
|
||||
- the coerced response (or `kwargs["original_response"]`) for output
|
||||
|
||||
All emits are best-effort: if the provider shape isn't recognized the
|
||||
helper exits silently. Existing chat/completion paths never enter this
|
||||
helper because their call_type doesn't contain "passthrough".
|
||||
|
||||
TEMPORARY BRIDGE: passthrough handlers don't populate the
|
||||
StandardLoggingPayload `messages` field today (they call
|
||||
`transform_response(messages=[])`), so the input is only available via
|
||||
`additional_args.complete_input_dict`. The proper fix is upstream in
|
||||
`base_passthrough_logging_handler._create_response_logging_payload()`:
|
||||
once that populates SLP `messages`/`response`, every callback gets
|
||||
passthrough I/O (with central redaction) for free and this helper's
|
||||
`complete_input_dict` fallback can be deleted. See follow-up issue.
|
||||
"""
|
||||
call_type = (
|
||||
standard_logging_payload.get("call_type")
|
||||
if isinstance(standard_logging_payload, dict)
|
||||
else None
|
||||
)
|
||||
if not _is_passthrough_call_type(call_type):
|
||||
return
|
||||
|
||||
# Respect LiteLLM's central message-redaction contract. The normal
|
||||
# chat/completion path is redacted by `perform_redaction` before
|
||||
# callbacks run, but `complete_input_dict` (read below) is NOT covered by
|
||||
# that layer — so without this gate, an operator who enabled redaction
|
||||
# would still see raw passthrough prompts in Arize. Skip entirely when
|
||||
# redaction is on so neither input nor output leaks through this bridge.
|
||||
if should_redact_message_logging(kwargs):
|
||||
return
|
||||
|
||||
# --- INPUT --------------------------------------------------------------
|
||||
additional_args = kwargs.get("additional_args") or {}
|
||||
complete_input_dict = (
|
||||
additional_args.get("complete_input_dict")
|
||||
if isinstance(additional_args, dict)
|
||||
else None
|
||||
)
|
||||
if isinstance(complete_input_dict, dict):
|
||||
_set_passthrough_input_attributes(span, complete_input_dict.get("messages"))
|
||||
|
||||
# --- OUTPUT -------------------------------------------------------------
|
||||
parsed_response = _parse_passthrough_response(
|
||||
raw_response_obj, coerced_response_obj, kwargs
|
||||
)
|
||||
if not isinstance(parsed_response, dict):
|
||||
return
|
||||
|
||||
_set_passthrough_output_attributes(span, parsed_response)
|
||||
|
||||
|
||||
def _set_passthrough_input_attributes(span: "Span", messages) -> None:
|
||||
"""Render passthrough request messages into INPUT_VALUE + LLM_INPUT_MESSAGES."""
|
||||
if not (isinstance(messages, list) and messages):
|
||||
return
|
||||
# Set INPUT_VALUE from the last user message text if discoverable.
|
||||
last_text = None
|
||||
for msg in reversed(messages):
|
||||
if isinstance(msg, dict):
|
||||
last_text = _coerce_text(msg.get("content"))
|
||||
if last_text:
|
||||
break
|
||||
if last_text:
|
||||
safe_set_attribute(span, SpanAttributes.INPUT_VALUE, last_text)
|
||||
# Mirror messages into LLM_INPUT_MESSAGES so the input pane renders.
|
||||
for idx, msg in enumerate(messages):
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
prefix = f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}"
|
||||
role = msg.get("role")
|
||||
if role:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
|
||||
role,
|
||||
)
|
||||
text = _coerce_text(msg.get("content"))
|
||||
if text is not None:
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
|
||||
text,
|
||||
)
|
||||
|
||||
|
||||
def _set_passthrough_output_attributes(span: "Span", parsed_response: dict) -> None:
|
||||
"""Render passthrough response into OUTPUT_VALUE + LLM_OUTPUT_MESSAGES."""
|
||||
# Anthropic / Bedrock-Anthropic: `content` is a list of typed parts.
|
||||
content_list = parsed_response.get("content")
|
||||
if isinstance(content_list, list) and content_list:
|
||||
texts = []
|
||||
for part in content_list:
|
||||
if isinstance(part, dict) and isinstance(part.get("text"), str):
|
||||
texts.append(part["text"])
|
||||
joined = "\n\n".join(t for t in texts if t)
|
||||
if joined:
|
||||
safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, joined)
|
||||
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0"
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
|
||||
parsed_response.get("role", "assistant"),
|
||||
)
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
|
||||
joined,
|
||||
)
|
||||
|
||||
# OpenAI-style passthrough: `choices[0].message.content`
|
||||
choices = parsed_response.get("choices")
|
||||
if isinstance(choices, list) and choices:
|
||||
first = choices[0]
|
||||
if isinstance(first, dict):
|
||||
msg = first.get("message")
|
||||
if isinstance(msg, dict):
|
||||
text = _coerce_text(msg.get("content"))
|
||||
if text:
|
||||
safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, text)
|
||||
prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0"
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_ROLE}",
|
||||
msg.get("role", "assistant"),
|
||||
)
|
||||
safe_set_attribute(
|
||||
span,
|
||||
f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}",
|
||||
text,
|
||||
)
|
||||
|
||||
|
||||
def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs):
|
||||
"""Return a dict view of the provider response for passthrough routes."""
|
||||
# Prefer the coerced view (already JSON-parsed for httpx.Response).
|
||||
candidates = []
|
||||
if isinstance(coerced_response_obj, dict):
|
||||
candidates.append(coerced_response_obj)
|
||||
if (
|
||||
isinstance(raw_response_obj, dict)
|
||||
and raw_response_obj is not coerced_response_obj
|
||||
):
|
||||
candidates.append(raw_response_obj)
|
||||
|
||||
for candidate in candidates:
|
||||
# StandardPassThroughResponseObject wrapper: {"response": "..."}.
|
||||
if (
|
||||
"response" in candidate
|
||||
and "content" not in candidate
|
||||
and "choices" not in candidate
|
||||
):
|
||||
inner = candidate.get("response")
|
||||
if isinstance(inner, str):
|
||||
try:
|
||||
parsed = json.loads(inner)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
except Exception:
|
||||
continue
|
||||
if isinstance(inner, dict):
|
||||
return inner
|
||||
else:
|
||||
return candidate
|
||||
|
||||
# Fallback: kwargs["original_response"] from the OTel base path.
|
||||
original = kwargs.get("original_response") if isinstance(kwargs, dict) else None
|
||||
if isinstance(original, dict):
|
||||
return original
|
||||
if isinstance(original, str):
|
||||
try:
|
||||
parsed = json.loads(original)
|
||||
if isinstance(parsed, dict):
|
||||
return parsed
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
@@ -104,6 +104,51 @@
|
||||
},
|
||||
"description": "Datadog Custom Metrics Integration"
|
||||
},
|
||||
{
|
||||
"id": "galileo",
|
||||
"displayName": "Galileo",
|
||||
"logo": "galileo.ico",
|
||||
"supports_key_team_logging": false,
|
||||
"dynamic_params": {
|
||||
"GALILEO_API_KEY": {
|
||||
"type": "password",
|
||||
"ui_name": "API Key",
|
||||
"description": "Galileo Cloud API key (app.galileo.ai). Omit for enterprise username/password auth.",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_PROJECT_ID": {
|
||||
"type": "text",
|
||||
"ui_name": "Project ID",
|
||||
"description": "Galileo project ID to log traces to",
|
||||
"required": true
|
||||
},
|
||||
"GALILEO_LOG_STREAM_ID": {
|
||||
"type": "text",
|
||||
"ui_name": "Log Stream ID",
|
||||
"description": "Galileo log stream ID for v2 spans logging (optional)",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_BASE_URL": {
|
||||
"type": "text",
|
||||
"ui_name": "Base URL",
|
||||
"description": "Galileo API base URL (e.g. https://api.galileo.ai for Cloud, or your enterprise API URL)",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_USERNAME": {
|
||||
"type": "text",
|
||||
"ui_name": "Username",
|
||||
"description": "Galileo enterprise username (legacy Observe auth; use instead of API key)",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_PASSWORD": {
|
||||
"type": "password",
|
||||
"ui_name": "Password",
|
||||
"description": "Galileo enterprise password (legacy Observe auth)",
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "Galileo AI Observability Integration"
|
||||
},
|
||||
{
|
||||
"id": "datadog_cost_management",
|
||||
"displayName": "Datadog Cost Management",
|
||||
|
||||
@@ -47,9 +47,29 @@ from litellm.exceptions import (
|
||||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
ModifyResponseException,
|
||||
SensitiveDataRouteException,
|
||||
)
|
||||
|
||||
|
||||
def get_session_id_from_request_data(request_data: Dict[str, Any]) -> Optional[str]:
|
||||
"""Extract session_id from request data (litellm_session_id or metadata)."""
|
||||
session_id = request_data.get("litellm_session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
metadata = request_data.get("metadata") or {}
|
||||
session_id = metadata.get("session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
litellm_metadata = request_data.get("litellm_metadata") or {}
|
||||
session_id = litellm_metadata.get("session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
@@ -68,6 +88,9 @@ class CustomGuardrail(CustomLogger):
|
||||
end_session_after_n_fails: Optional[int] = None,
|
||||
on_violation: Optional[str] = None,
|
||||
realtime_violation_message: Optional[str] = None,
|
||||
on_sensitive_data: Optional[str] = None,
|
||||
sensitive_data_route_to_model: Optional[str] = None,
|
||||
sticky_session_routing: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
@@ -83,6 +106,9 @@ class CustomGuardrail(CustomLogger):
|
||||
end_session_after_n_fails: For /v1/realtime sessions, end the session after this many violations
|
||||
on_violation: For /v1/realtime sessions, 'warn' or 'end_session'
|
||||
realtime_violation_message: Message the bot speaks aloud when a /v1/realtime guardrail fires
|
||||
on_sensitive_data: Action when sensitive data is detected. 'block' (default) or 'route'
|
||||
sensitive_data_route_to_model: Model to route to when on_sensitive_data='route'
|
||||
sticky_session_routing: When True, all subsequent requests in the session use the same model
|
||||
"""
|
||||
self.guardrail_name = guardrail_name
|
||||
self.supported_event_hooks = supported_event_hooks
|
||||
@@ -96,6 +122,11 @@ class CustomGuardrail(CustomLogger):
|
||||
self.end_session_after_n_fails: Optional[int] = end_session_after_n_fails
|
||||
self.on_violation: Optional[str] = on_violation
|
||||
self.realtime_violation_message: Optional[str] = realtime_violation_message
|
||||
self.on_sensitive_data: Optional[str] = on_sensitive_data
|
||||
self.sensitive_data_route_to_model: Optional[str] = (
|
||||
sensitive_data_route_to_model
|
||||
)
|
||||
self.sticky_session_routing: bool = sticky_session_routing
|
||||
|
||||
if supported_event_hooks:
|
||||
## validate event_hook is in supported_event_hooks
|
||||
@@ -167,6 +198,108 @@ class CustomGuardrail(CustomLogger):
|
||||
detection_info=detection_info,
|
||||
)
|
||||
|
||||
def raise_sensitive_data_route_exception(
|
||||
self,
|
||||
route_to_model: str,
|
||||
request_data: Dict[str, Any],
|
||||
detection_info: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Raise an exception to reroute the request to a different model.
|
||||
|
||||
Use this when sensitive data is detected and the guardrail is configured
|
||||
to route to an on-premise model instead of blocking.
|
||||
|
||||
The exception will reroute this request to the specified model. When
|
||||
sticky_session_routing is enabled (the default), it also stores the
|
||||
routing decision so subsequent requests in this session reuse the model.
|
||||
|
||||
Args:
|
||||
route_to_model: The model to route this request (and session) to
|
||||
request_data: The original request data dictionary
|
||||
detection_info: Optional non-sensitive detection metadata (e.g. matched
|
||||
entity types, rule ids, scores). This is surfaced in request metadata
|
||||
and logs, so it must not contain the raw detected sensitive values.
|
||||
|
||||
Raises:
|
||||
SensitiveDataRouteException: Always raises to trigger rerouting
|
||||
"""
|
||||
session_id = self._get_session_id_from_request_data(request_data)
|
||||
if not session_id:
|
||||
raise ValueError(
|
||||
"Cannot route sensitive data without a session_id. "
|
||||
"Ensure the request includes a session_id in metadata or headers."
|
||||
)
|
||||
|
||||
raise SensitiveDataRouteException(
|
||||
route_to_model=route_to_model,
|
||||
session_id=session_id,
|
||||
guardrail_name=self.guardrail_name,
|
||||
detection_info=detection_info,
|
||||
sticky_session_routing=self.sticky_session_routing,
|
||||
)
|
||||
|
||||
def _get_session_id_from_request_data(
|
||||
self, request_data: Dict[str, Any]
|
||||
) -> Optional[str]:
|
||||
"""Extract session_id from request data."""
|
||||
return get_session_id_from_request_data(request_data)
|
||||
|
||||
def should_route_on_sensitive_data(self) -> bool:
|
||||
"""
|
||||
Returns True if this guardrail is configured to route requests
|
||||
to a different model when sensitive data is detected.
|
||||
"""
|
||||
return (
|
||||
self.on_sensitive_data == "route"
|
||||
and self.sensitive_data_route_to_model is not None
|
||||
)
|
||||
|
||||
def handle_sensitive_data_detection(
|
||||
self,
|
||||
request_data: Dict[str, Any],
|
||||
detection_info: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Handle sensitive data detection based on guardrail configuration.
|
||||
|
||||
If on_sensitive_data='route', raises SensitiveDataRouteException to reroute.
|
||||
Otherwise, raises GuardrailRaisedException to block. When routing is
|
||||
configured but the request carries no session_id, routing is not possible
|
||||
so the request falls back to a graceful block.
|
||||
|
||||
Args:
|
||||
request_data: The request data dictionary
|
||||
detection_info: Optional non-sensitive detection metadata. When routing,
|
||||
this is surfaced in request metadata and logs, so it must not contain
|
||||
the raw detected sensitive values.
|
||||
|
||||
Raises:
|
||||
SensitiveDataRouteException: When configured to route and a session_id is present
|
||||
GuardrailRaisedException: When configured to block, or when routing is
|
||||
configured but no session_id is available
|
||||
"""
|
||||
if self.should_route_on_sensitive_data():
|
||||
try:
|
||||
self.raise_sensitive_data_route_exception(
|
||||
route_to_model=self.sensitive_data_route_to_model, # type: ignore
|
||||
request_data=request_data,
|
||||
detection_info=detection_info,
|
||||
)
|
||||
except ValueError:
|
||||
raise GuardrailRaisedException(
|
||||
message=(
|
||||
f"Sensitive data detected by {self.guardrail_name} "
|
||||
"(routing skipped: request has no session_id)"
|
||||
),
|
||||
guardrail_name=self.guardrail_name,
|
||||
)
|
||||
else:
|
||||
raise GuardrailRaisedException(
|
||||
message=f"Sensitive data detected by {self.guardrail_name}",
|
||||
guardrail_name=self.guardrail_name,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
|
||||
"""
|
||||
@@ -662,6 +795,16 @@ class CustomGuardrail(CustomLogger):
|
||||
request_data["metadata"] = {}
|
||||
_append_guardrail_info(request_data["metadata"])
|
||||
|
||||
# Emit the otel guardrail span here, where every guardrail execution lands,
|
||||
# rather than relying on a post-call hook that does not fire on every path
|
||||
# (e.g. a pass-through request that passes its guardrails).
|
||||
try:
|
||||
from litellm.integrations.otel.logger import emit_guardrail_span
|
||||
|
||||
emit_guardrail_span(slg)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
@@ -743,12 +886,20 @@ class CustomGuardrail(CustomLogger):
|
||||
Guardrails signal intentional blocks by raising:
|
||||
- GuardrailRaisedException (generic guardrail API, tool permission)
|
||||
- BlockedPiiEntityError (Presidio PII detection)
|
||||
- SensitiveDataRouteException (sensitive-data reroute to on-premise model)
|
||||
- HTTPException with status 400 (content policy violation)
|
||||
- ModifyResponseException (passthrough mode violation)
|
||||
"""
|
||||
if isinstance(e, ModifyResponseException):
|
||||
return True
|
||||
if isinstance(e, (GuardrailRaisedException, BlockedPiiEntityError)):
|
||||
if isinstance(
|
||||
e,
|
||||
(
|
||||
GuardrailRaisedException,
|
||||
BlockedPiiEntityError,
|
||||
SensitiveDataRouteException,
|
||||
),
|
||||
):
|
||||
return True
|
||||
if (
|
||||
HTTPException is not None
|
||||
|
||||
@@ -41,6 +41,7 @@ from litellm.integrations.datadog.datadog_handler import (
|
||||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
MaskedHTTPStatusError,
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
@@ -68,6 +69,22 @@ DD_LOGGED_SUCCESS_SERVICE_TYPES = [
|
||||
]
|
||||
|
||||
|
||||
def _resolve_dd_batch_size() -> int:
|
||||
raw = os.getenv("DD_BATCH_SIZE")
|
||||
if raw is None:
|
||||
return DD_MAX_BATCH_SIZE
|
||||
try:
|
||||
value = int(raw)
|
||||
except ValueError:
|
||||
verbose_logger.warning(
|
||||
"Datadog: ignoring invalid DD_BATCH_SIZE=%r, using %s",
|
||||
raw,
|
||||
DD_MAX_BATCH_SIZE,
|
||||
)
|
||||
return DD_MAX_BATCH_SIZE
|
||||
return max(1, min(value, DD_MAX_BATCH_SIZE))
|
||||
|
||||
|
||||
class DataDogLogger(
|
||||
CustomBatchLogger,
|
||||
AdditionalLoggingUtils,
|
||||
@@ -128,7 +145,9 @@ class DataDogLogger(
|
||||
asyncio.create_task(self.periodic_flush())
|
||||
self.flush_lock = asyncio.Lock()
|
||||
super().__init__(
|
||||
**kwargs, flush_lock=self.flush_lock, batch_size=DD_MAX_BATCH_SIZE
|
||||
**kwargs,
|
||||
flush_lock=self.flush_lock,
|
||||
batch_size=_resolve_dd_batch_size(),
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
@@ -339,28 +358,14 @@ class DataDogLogger(
|
||||
"[DATADOG MOCK] Mock mode enabled - API calls will be intercepted"
|
||||
)
|
||||
|
||||
response = await self.async_send_compressed_data(batch_to_send)
|
||||
if response.status_code == 413:
|
||||
verbose_logger.exception(DD_ERRORS.DATADOG_413_ERROR.value)
|
||||
self.log_queue = batch_to_send + self.log_queue
|
||||
return
|
||||
|
||||
response.raise_for_status()
|
||||
if response.status_code != 202:
|
||||
raise Exception(
|
||||
f"Response from datadog API status_code: {response.status_code}, text: {response.text}"
|
||||
)
|
||||
undelivered = await self._send_with_413_split(batch_to_send)
|
||||
if undelivered:
|
||||
self.log_queue = undelivered + self.log_queue
|
||||
|
||||
if self.is_mock_mode:
|
||||
verbose_logger.debug(
|
||||
f"[DATADOG MOCK] Batch of {len(batch_to_send)} events successfully mocked"
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"Datadog: Response from datadog API status_code: %s, text: %s",
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
self.log_queue = batch_to_send + self.log_queue
|
||||
@@ -368,6 +373,62 @@ class DataDogLogger(
|
||||
f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}"
|
||||
)
|
||||
|
||||
async def _send_with_413_split(self, batch: List) -> List:
|
||||
"""
|
||||
Send a batch, halving any sub-batch that 413s (payload too large) and retrying the
|
||||
halves, since Datadog enforces a 5MB uncompressed limit per request.
|
||||
|
||||
A 413 surfaces as a raised MaskedHTTPStatusError (httpx raise_for_status), not a
|
||||
returned response, so both paths are handled. A lone event that still 413s is
|
||||
dropped to avoid wedging the queue on an undeliverable payload. Returns the events
|
||||
that could not be delivered because of a non-413 (transient) error, so the caller
|
||||
re-queues only those and never the events already accepted by Datadog.
|
||||
"""
|
||||
pending: List[List] = [batch]
|
||||
while pending:
|
||||
chunk = pending.pop()
|
||||
if not chunk:
|
||||
continue
|
||||
try:
|
||||
response = await self.async_send_compressed_data(chunk)
|
||||
except Exception as e:
|
||||
if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413:
|
||||
response = e.response
|
||||
else:
|
||||
verbose_logger.exception(
|
||||
f"Datadog Error sending batch API - {str(e)}"
|
||||
)
|
||||
return self._undelivered(chunk, pending)
|
||||
|
||||
if response.status_code == 413:
|
||||
if len(chunk) == 1:
|
||||
verbose_logger.error(DD_ERRORS.DATADOG_413_ERROR.value)
|
||||
continue
|
||||
mid = len(chunk) // 2
|
||||
pending.append(chunk[mid:])
|
||||
pending.append(chunk[:mid])
|
||||
continue
|
||||
|
||||
if response.status_code != 202:
|
||||
verbose_logger.error(
|
||||
"Datadog: unexpected response status_code=%s, text=%s",
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
return self._undelivered(chunk, pending)
|
||||
|
||||
verbose_logger.debug(
|
||||
"Datadog: delivered %s events, status_code=%s, text=%s",
|
||||
len(chunk),
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _undelivered(chunk: List, pending: List[List]) -> List:
|
||||
return chunk + [event for remaining in reversed(pending) for event in remaining]
|
||||
|
||||
async def flush_queue(self):
|
||||
if self.flush_lock is None:
|
||||
return
|
||||
|
||||
@@ -95,7 +95,9 @@ class FocusTransformer:
|
||||
pl.lit("Usage-Based").alias("ChargeFrequency"),
|
||||
fmt(pl.col("ChargePeriodEnd")).alias("ChargePeriodEnd"),
|
||||
fmt(pl.col("ChargePeriodStart")).alias("ChargePeriodStart"),
|
||||
dec(pl.lit(1.0)).alias("ConsumedQuantity"),
|
||||
dec(
|
||||
pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0)
|
||||
).alias("ConsumedQuantity"),
|
||||
pl.lit("Requests").alias("ConsumedUnit"),
|
||||
dec(pl.col("spend").fill_null(0.0)).alias("ContractedCost"),
|
||||
none_str.alias("ContractedUnitPrice"),
|
||||
@@ -107,7 +109,9 @@ class FocusTransformer:
|
||||
none_str.alias("AvailabilityZone"),
|
||||
pl.lit("USD").alias("PricingCurrency"),
|
||||
none_str.alias("PricingCategory"),
|
||||
dec(pl.lit(1.0)).alias("PricingQuantity"),
|
||||
dec(
|
||||
pl.col("api_requests").cast(pl.Int64).cast(pl.Float64).fill_null(0.0)
|
||||
).alias("PricingQuantity"),
|
||||
none_dec.alias("PricingCurrencyContractedUnitPrice"),
|
||||
dec(pl.col("spend").fill_null(0.0)).alias("PricingCurrencyEffectiveCost"),
|
||||
none_dec.alias("PricingCurrencyListUnitPrice"),
|
||||
|
||||
+549
-79
@@ -1,8 +1,13 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Tuple, cast
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import litellm
|
||||
@@ -12,11 +17,15 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
get_content_from_model_response,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
HttpxBinaryResponseContent,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
||||
GALILEO_CLOUD_API_BASE_URL = "https://api.galileo.ai"
|
||||
# Cap the in-memory buffer so persistent flush failures (e.g. Galileo
|
||||
@@ -33,6 +42,11 @@ class LLMResponse(BaseModel):
|
||||
model: str
|
||||
num_input_tokens: int
|
||||
num_output_tokens: int
|
||||
num_total_tokens: int
|
||||
cost: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Total cost of the LLM call in USD as computed by LiteLLM.",
|
||||
)
|
||||
output_logprobs: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="Optional. When available, logprobs are used to compute Uncertainty.",
|
||||
@@ -121,10 +135,14 @@ class GalileoObserve(CustomLogger):
|
||||
|
||||
@staticmethod
|
||||
def _galileo_input_messages(
|
||||
messages: Optional[List[Any]], input_text: str
|
||||
messages: Optional[Any], input_text: str
|
||||
) -> List[Dict[str, str]]:
|
||||
if isinstance(messages, dict):
|
||||
messages = messages.get("messages")
|
||||
if not messages:
|
||||
return [{"role": "user", "content": input_text}]
|
||||
if not isinstance(messages, list):
|
||||
return [{"role": "user", "content": input_text}]
|
||||
|
||||
galileo_messages: List[Dict[str, str]] = []
|
||||
for message in messages:
|
||||
@@ -147,13 +165,59 @@ class GalileoObserve(CustomLogger):
|
||||
return [{"role": "user", "content": input_text}]
|
||||
|
||||
@staticmethod
|
||||
def _record_to_v2_span(record: Dict[str, Any]) -> Dict[str, Any]:
|
||||
created_at = record.get("created_at", "")
|
||||
def _local_timezone():
|
||||
return datetime.now().astimezone().tzinfo or timezone.utc
|
||||
|
||||
@staticmethod
|
||||
def _format_created_at(dt: Union[datetime, Any]) -> str:
|
||||
"""Serialize timestamps as UTC ISO-8601 for Galileo."""
|
||||
if not isinstance(dt, datetime):
|
||||
return str(dt)
|
||||
|
||||
if dt.tzinfo is None:
|
||||
# LiteLLM often passes naive datetimes in local time; convert to UTC
|
||||
# instead of appending Z to local time (which shifts Traces tab sorting).
|
||||
dt = dt.replace(tzinfo=GalileoObserve._local_timezone())
|
||||
|
||||
return dt.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
@staticmethod
|
||||
def _normalize_created_at(created_at: str) -> str:
|
||||
if created_at and not re.search(r"(Z|[+-]\d{2}:?\d{2})$", created_at):
|
||||
created_at = f"{created_at}Z"
|
||||
return f"{created_at}Z"
|
||||
return created_at
|
||||
|
||||
@staticmethod
|
||||
def _token_metrics_from_record(record: Dict[str, Any]) -> Dict[str, Any]:
|
||||
num_input_tokens = int(record.get("num_input_tokens") or 0)
|
||||
num_output_tokens = int(record.get("num_output_tokens") or 0)
|
||||
num_total_tokens = int(record.get("num_total_tokens") or 0)
|
||||
if num_total_tokens == 0 and (num_input_tokens or num_output_tokens):
|
||||
num_total_tokens = num_input_tokens + num_output_tokens
|
||||
metrics: Dict[str, Any] = {
|
||||
"num_input_tokens": num_input_tokens,
|
||||
"num_output_tokens": num_output_tokens,
|
||||
"num_total_tokens": num_total_tokens,
|
||||
}
|
||||
cost = record.get("cost")
|
||||
if cost is not None:
|
||||
metrics["cost"] = float(cost)
|
||||
return metrics
|
||||
|
||||
@staticmethod
|
||||
def _record_to_v2_span(
|
||||
record: Dict[str, Any],
|
||||
*,
|
||||
trace_id: str,
|
||||
span_id: str,
|
||||
) -> Dict[str, Any]:
|
||||
created_at = GalileoObserve._normalize_created_at(record.get("created_at", ""))
|
||||
|
||||
span: Dict[str, Any] = {
|
||||
"type": "llm",
|
||||
"id": span_id,
|
||||
"trace_id": trace_id,
|
||||
"parent_id": trace_id,
|
||||
"name": record.get("node_type", "litellm"),
|
||||
"created_at": created_at,
|
||||
"input": GalileoObserve._galileo_input_messages(
|
||||
@@ -167,14 +231,49 @@ class GalileoObserve(CustomLogger):
|
||||
"model": record.get("model"),
|
||||
"metrics": {
|
||||
"duration_ns": int(record.get("latency_ms", 0)) * 1_000_000,
|
||||
"num_input_tokens": record.get("num_input_tokens"),
|
||||
"num_output_tokens": record.get("num_output_tokens"),
|
||||
**GalileoObserve._token_metrics_from_record(record),
|
||||
},
|
||||
}
|
||||
if record.get("tags"):
|
||||
span["tags"] = record["tags"]
|
||||
return span
|
||||
|
||||
@staticmethod
|
||||
def _record_to_v2_trace(record: Dict[str, Any]) -> Dict[str, Any]:
|
||||
trace_id = str(uuid.uuid4())
|
||||
span_id = str(uuid.uuid4())
|
||||
created_at = GalileoObserve._normalize_created_at(record.get("created_at", ""))
|
||||
|
||||
return {
|
||||
"type": "trace",
|
||||
"id": trace_id,
|
||||
"name": record.get("node_type", "litellm"),
|
||||
"created_at": created_at,
|
||||
"input": record.get("input_text", ""),
|
||||
"output": record.get("output_text", ""),
|
||||
"status_code": record.get("status_code", 200),
|
||||
"metrics": {
|
||||
"duration_ns": int(record.get("latency_ms", 0)) * 1_000_000,
|
||||
**GalileoObserve._token_metrics_from_record(record),
|
||||
},
|
||||
"spans": [
|
||||
GalileoObserve._record_to_v2_span(
|
||||
record, trace_id=trace_id, span_id=span_id
|
||||
)
|
||||
],
|
||||
}
|
||||
|
||||
def _build_traces_payload(self, records: List[dict]) -> Dict[str, Any]:
|
||||
payload: Dict[str, Any] = {
|
||||
"traces": [self._record_to_v2_trace(record) for record in records],
|
||||
"logging_method": "api_direct",
|
||||
"reliable": False,
|
||||
"is_complete": True,
|
||||
}
|
||||
if self.log_stream_id:
|
||||
payload["log_stream_id"] = self.log_stream_id
|
||||
return payload
|
||||
|
||||
def _get_ingest_request(self) -> Optional[Tuple[str, Dict[str, Any]]]:
|
||||
if not self.base_url or not self.project_id:
|
||||
return None
|
||||
@@ -184,105 +283,457 @@ class GalileoObserve(CustomLogger):
|
||||
# flush_in_memory_records) aren't silently dropped when we later clear
|
||||
# the in-memory buffer.
|
||||
records = list(self.in_memory_records)
|
||||
payload = self._build_traces_payload(records)
|
||||
|
||||
if self.use_v2_api:
|
||||
payload: Dict[str, Any] = {
|
||||
"spans": [self._record_to_v2_span(record) for record in records],
|
||||
"reliable": False,
|
||||
}
|
||||
if self.log_stream_id:
|
||||
payload["log_stream_id"] = self.log_stream_id
|
||||
return (
|
||||
f"{self.base_url}/v2/projects/{self.project_id}/spans",
|
||||
f"{self.base_url}/ingest/traces/{self.project_id}",
|
||||
payload,
|
||||
)
|
||||
|
||||
# Username/password auth logs in for a JWT and uses the standard v2 traces API.
|
||||
return (
|
||||
f"{self.base_url}/projects/{self.project_id}/observe/ingest",
|
||||
{"records": records},
|
||||
f"{self.base_url}/v2/projects/{self.project_id}/traces",
|
||||
payload,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _redact_headers(headers: Optional[Dict[str, str]]) -> Dict[str, str]:
|
||||
if not headers:
|
||||
return {}
|
||||
redacted: Dict[str, str] = {}
|
||||
for key, value in headers.items():
|
||||
if key.lower() in {"authorization", "galileo-api-key"} and value:
|
||||
redacted[key] = (
|
||||
f"{value[:8]}...{value[-4:]}" if len(value) > 12 else "***"
|
||||
)
|
||||
else:
|
||||
redacted[key] = value
|
||||
return redacted
|
||||
|
||||
def _log_flush_config(self) -> None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush config: use_v2_api=%s base_url=%s project_id=%s "
|
||||
"log_stream_id=%s api_key_set=%s username_set=%s record_count=%s",
|
||||
self.use_v2_api,
|
||||
self.base_url,
|
||||
self.project_id,
|
||||
self.log_stream_id,
|
||||
bool(self.api_key),
|
||||
bool(self.username),
|
||||
len(self.in_memory_records),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _log_v2_payload_validation(payload: Dict[str, Any]) -> None:
|
||||
missing_fields: List[str] = []
|
||||
traces = payload.get("traces", [])
|
||||
if not traces:
|
||||
missing_fields.append("traces")
|
||||
|
||||
for trace_index, trace in enumerate(traces):
|
||||
if not isinstance(trace, dict):
|
||||
continue
|
||||
for field in ("id", "type", "spans"):
|
||||
if field not in trace:
|
||||
missing_fields.append(f"traces[{trace_index}].{field}")
|
||||
|
||||
trace_id = trace.get("id")
|
||||
for span_index, span in enumerate(trace.get("spans", [])):
|
||||
if not isinstance(span, dict):
|
||||
continue
|
||||
for field in ("id", "trace_id", "parent_id"):
|
||||
if field not in span:
|
||||
missing_fields.append(
|
||||
f"traces[{trace_index}].spans[{span_index}].{field}"
|
||||
)
|
||||
if trace_id and span.get("trace_id") != trace_id:
|
||||
missing_fields.append(
|
||||
f"traces[{trace_index}].spans[{span_index}].trace_id mismatch"
|
||||
)
|
||||
|
||||
if missing_fields:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: ingest /traces payload validation issues: %s",
|
||||
missing_fields,
|
||||
)
|
||||
|
||||
def _log_flush_payload(self, url: str, payload: Dict[str, Any]) -> None:
|
||||
traces = payload.get("traces", [])
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush URL: %s trace_count=%s",
|
||||
url,
|
||||
len(traces) if isinstance(traces, list) else 0,
|
||||
)
|
||||
if self.use_v2_api and "/ingest/traces/" in url:
|
||||
self._log_v2_payload_validation(payload)
|
||||
|
||||
@staticmethod
|
||||
def _log_http_status_error(error: httpx.HTTPStatusError, url: str) -> None:
|
||||
response = error.response
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error: status=%s url=%s",
|
||||
response.status_code,
|
||||
url,
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error response body: %s",
|
||||
response.text,
|
||||
)
|
||||
try:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error response json: %s",
|
||||
response.json(),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _build_prompt(kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
optional_params = kwargs.get("optional_params", {}) or {}
|
||||
prompt: Dict[str, Any] = {"messages": kwargs.get("messages")}
|
||||
if optional_params.get("functions") is not None:
|
||||
prompt["functions"] = optional_params["functions"]
|
||||
if optional_params.get("tools") is not None:
|
||||
prompt["tools"] = optional_params["tools"]
|
||||
return prompt
|
||||
|
||||
@staticmethod
|
||||
def _serialize_galileo_output(value: Any) -> Optional[str]:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
|
||||
def _json_default(obj: Any) -> Any:
|
||||
if hasattr(obj, "model_dump"):
|
||||
return obj.model_dump()
|
||||
return str(obj)
|
||||
|
||||
return json.dumps(value, default=_json_default)
|
||||
|
||||
@staticmethod
|
||||
def _prompt_to_input_text(prompt: Dict[str, Any]) -> str:
|
||||
messages = prompt.get("messages")
|
||||
if messages is not None:
|
||||
text = GalileoObserve._input_text_from_messages(messages)
|
||||
if text:
|
||||
return text
|
||||
return json.dumps(prompt, default=str)
|
||||
|
||||
@staticmethod
|
||||
def _get_chat_content_for_galileo(response_obj: litellm.ModelResponse) -> Any:
|
||||
if response_obj.choices and len(response_obj.choices) > 0:
|
||||
message = response_obj["choices"][0]["message"]
|
||||
if hasattr(message, "json"):
|
||||
message_json = message.json()
|
||||
if isinstance(message_json, str):
|
||||
return json.loads(message_json)
|
||||
return message_json
|
||||
return message
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_text_completion_content_for_galileo(
|
||||
response_obj: litellm.TextCompletionResponse,
|
||||
) -> Optional[str]:
|
||||
if response_obj.choices and len(response_obj.choices) > 0:
|
||||
return response_obj.choices[0].text
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_responses_api_content_for_galileo(
|
||||
response_obj: ResponsesAPIResponse,
|
||||
) -> Any:
|
||||
if hasattr(response_obj, "output") and response_obj.output:
|
||||
return response_obj.output
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _langfuse_style_rerank_prompt(kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Match Langfuse rerank input: prompt = {"messages": kwargs.get("messages")}."""
|
||||
return {"messages": kwargs.get("messages")}
|
||||
|
||||
def _get_galileo_input_output_content(
|
||||
self,
|
||||
kwargs: Dict[str, Any],
|
||||
response_obj: Any,
|
||||
level: str = "DEFAULT",
|
||||
status_message: Optional[str] = None,
|
||||
) -> Tuple[str, Optional[str], Any]:
|
||||
"""
|
||||
Mirror Langfuse _get_langfuse_input_output_content for Galileo ingest.
|
||||
|
||||
Returns (input_text, output_text, messages_for_span). output_text None skips ingest.
|
||||
"""
|
||||
call_type = kwargs.get("call_type")
|
||||
prompt = self._build_prompt(kwargs)
|
||||
|
||||
if (
|
||||
level == "ERROR"
|
||||
and status_message is not None
|
||||
and isinstance(status_message, str)
|
||||
):
|
||||
return self._prompt_to_input_text(prompt), status_message, prompt
|
||||
|
||||
if response_obj is not None and (
|
||||
call_type == "embedding"
|
||||
or isinstance(response_obj, litellm.EmbeddingResponse)
|
||||
):
|
||||
return self._prompt_to_input_text(prompt), None, prompt
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, litellm.ModelResponse):
|
||||
output = self._get_chat_content_for_galileo(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, HttpxBinaryResponseContent
|
||||
):
|
||||
return self._prompt_to_input_text(prompt), "speech-output", prompt
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, litellm.TextCompletionResponse
|
||||
):
|
||||
output = self._get_text_completion_content_for_galileo(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, litellm.ImageResponse):
|
||||
output = response_obj.get("data", None)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, litellm.TranscriptionResponse
|
||||
):
|
||||
output = response_obj.get("text", None)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, litellm.RerankResponse
|
||||
):
|
||||
output = response_obj.results
|
||||
rerank_prompt = self._langfuse_style_rerank_prompt(kwargs)
|
||||
return (
|
||||
json.dumps(rerank_prompt, default=str),
|
||||
self._serialize_galileo_output(output),
|
||||
rerank_prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, ResponsesAPIResponse):
|
||||
output = self._get_responses_api_content_for_galileo(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
if (
|
||||
call_type == "_arealtime"
|
||||
and response_obj is not None
|
||||
and isinstance(response_obj, list)
|
||||
):
|
||||
input_val = kwargs.get("input")
|
||||
return (
|
||||
self._serialize_galileo_output(input_val) or "",
|
||||
self._serialize_galileo_output(response_obj),
|
||||
input_val,
|
||||
)
|
||||
|
||||
if (
|
||||
call_type == "pass_through_endpoint"
|
||||
and response_obj is not None
|
||||
and isinstance(response_obj, dict)
|
||||
):
|
||||
output = response_obj.get("response", "")
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, dict):
|
||||
output = get_content_from_model_response(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
return self._prompt_to_input_text(prompt), None, kwargs.get("messages") or []
|
||||
|
||||
def get_output_str_from_response(
|
||||
self, response_obj: Any, kwargs: Dict[str, Any]
|
||||
) -> Optional[str]:
|
||||
if response_obj is None:
|
||||
return None
|
||||
if kwargs.get("call_type", None) == "embedding" or isinstance(
|
||||
response_obj, litellm.EmbeddingResponse
|
||||
):
|
||||
return None
|
||||
if isinstance(response_obj, litellm.TextCompletionResponse):
|
||||
return response_obj.choices[0].text
|
||||
if isinstance(response_obj, litellm.ImageResponse):
|
||||
return json.dumps(response_obj["data"], default=str)
|
||||
if isinstance(response_obj, (litellm.ModelResponse, dict)):
|
||||
return get_content_from_model_response(response_obj)
|
||||
return None
|
||||
_, output_text, _ = self._get_galileo_input_output_content(
|
||||
kwargs=kwargs, response_obj=response_obj
|
||||
)
|
||||
return output_text
|
||||
|
||||
@staticmethod
|
||||
def _input_text_from_messages(messages: Any) -> str:
|
||||
"""Return a plain-string summary of the input suitable for the trace-level input field."""
|
||||
if isinstance(messages, str):
|
||||
return messages
|
||||
if not isinstance(messages, list):
|
||||
return ""
|
||||
# Use the last user/human message so the trace table shows the actual prompt
|
||||
for msg in reversed(messages):
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
if str(msg.get("role", "")).lower() in ("user", "human"):
|
||||
content = msg.get("content") or ""
|
||||
if isinstance(content, list):
|
||||
content = " ".join(
|
||||
b.get("text", "") if isinstance(b, dict) else str(b)
|
||||
for b in content
|
||||
)
|
||||
if content:
|
||||
return str(content)
|
||||
# Fallback: first non-empty content of any role
|
||||
for msg in messages:
|
||||
if isinstance(msg, dict):
|
||||
content = msg.get("content") or ""
|
||||
if isinstance(content, list):
|
||||
content = " ".join(
|
||||
b.get("text", "") if isinstance(b, dict) else str(b)
|
||||
for b in content
|
||||
)
|
||||
if content:
|
||||
return str(content)
|
||||
return ""
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any
|
||||
):
|
||||
verbose_logger.debug("On Async Success")
|
||||
try:
|
||||
await self._async_log_success_event_impl(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
except Exception:
|
||||
verbose_logger.exception(
|
||||
"Galileo Logger: unexpected error in async_log_success_event"
|
||||
)
|
||||
|
||||
async def _async_log_success_event_impl(
|
||||
self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any
|
||||
):
|
||||
if not self._is_configured():
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: skipping flush — set GALILEO_PROJECT_ID and "
|
||||
"either GALILEO_API_KEY (hosted) or GALILEO_USERNAME/GALILEO_PASSWORD "
|
||||
"(enterprise Observe)."
|
||||
"Galileo Logger: skipping — GALILEO_PROJECT_ID=%s GALILEO_API_KEY=%s GALILEO_BASE_URL=%s",
|
||||
bool(self.project_id),
|
||||
bool(self.api_key),
|
||||
bool(self.base_url),
|
||||
)
|
||||
return
|
||||
|
||||
_latency_ms = int((end_time - start_time).total_seconds() * 1000)
|
||||
_call_type = kwargs.get("call_type", "litellm")
|
||||
input_text = litellm.utils.get_formatted_prompt(
|
||||
data=kwargs, call_type=_call_type
|
||||
slo: Optional[Dict[str, Any]] = kwargs.get("standard_logging_object")
|
||||
if slo is None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: no standard_logging_object in kwargs, skipping"
|
||||
)
|
||||
return
|
||||
|
||||
_call_type: str = str(
|
||||
slo.get("call_type") or kwargs.get("call_type") or "litellm"
|
||||
)
|
||||
|
||||
_usage = response_obj.get("usage", {}) or {}
|
||||
num_input_tokens = _usage.get("prompt_tokens", 0)
|
||||
num_output_tokens = _usage.get("completion_tokens", 0)
|
||||
input_text, output_text, messages = self._get_galileo_input_output_content(
|
||||
kwargs=kwargs, response_obj=response_obj
|
||||
)
|
||||
if output_text is None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: skipping %s — no text output to log", _call_type
|
||||
)
|
||||
return
|
||||
|
||||
output_text = self.get_output_str_from_response(
|
||||
response_obj=response_obj, kwargs=kwargs
|
||||
raw_start = slo.get("startTime")
|
||||
raw_end = slo.get("endTime")
|
||||
if raw_start is None or raw_end is None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: standard_logging_object missing startTime/endTime, "
|
||||
"falling back to start_time/end_time params"
|
||||
)
|
||||
if not isinstance(start_time, datetime) or not isinstance(
|
||||
end_time, datetime
|
||||
):
|
||||
return
|
||||
start_ts = start_time
|
||||
end_ts = end_time
|
||||
if start_ts.tzinfo is None:
|
||||
start_ts = start_ts.replace(tzinfo=GalileoObserve._local_timezone())
|
||||
if end_ts.tzinfo is None:
|
||||
end_ts = end_ts.replace(tzinfo=GalileoObserve._local_timezone())
|
||||
start_ts = start_ts.astimezone(timezone.utc)
|
||||
end_ts = end_ts.astimezone(timezone.utc)
|
||||
else:
|
||||
start_ts = datetime.fromtimestamp(float(raw_start), tz=timezone.utc)
|
||||
end_ts = datetime.fromtimestamp(float(raw_end), tz=timezone.utc)
|
||||
_latency_ms = max(0, int((end_ts - start_ts).total_seconds() * 1000))
|
||||
num_input_tokens = int(slo.get("prompt_tokens") or 0)
|
||||
num_output_tokens = int(slo.get("completion_tokens") or 0)
|
||||
num_total_tokens = int(slo.get("total_tokens") or 0)
|
||||
if num_total_tokens == 0 and (num_input_tokens or num_output_tokens):
|
||||
num_total_tokens = num_input_tokens + num_output_tokens
|
||||
|
||||
request_record = LLMResponse(
|
||||
latency_ms=_latency_ms,
|
||||
status_code=200,
|
||||
input_text=input_text,
|
||||
output_text=output_text,
|
||||
node_type=_call_type,
|
||||
model=str(slo.get("model") or kwargs.get("model") or "-"),
|
||||
num_input_tokens=num_input_tokens,
|
||||
num_output_tokens=num_output_tokens,
|
||||
num_total_tokens=num_total_tokens,
|
||||
cost=slo.get("response_cost"),
|
||||
created_at=GalileoObserve._format_created_at(start_ts),
|
||||
)
|
||||
|
||||
if output_text is not None:
|
||||
request_record = LLMResponse(
|
||||
latency_ms=_latency_ms,
|
||||
status_code=200,
|
||||
input_text=input_text,
|
||||
output_text=output_text,
|
||||
node_type=_call_type,
|
||||
model=kwargs.get("model", "-"),
|
||||
num_input_tokens=num_input_tokens,
|
||||
num_output_tokens=num_output_tokens,
|
||||
created_at=start_time.strftime(
|
||||
"%Y-%m-%dT%H:%M:%S"
|
||||
), # timestamp str constructed in "%Y-%m-%dT%H:%M:%S" format
|
||||
request_dict = request_record.model_dump()
|
||||
if isinstance(messages, dict):
|
||||
messages = messages.get("messages")
|
||||
if isinstance(messages, list) and messages:
|
||||
request_dict["messages"] = messages
|
||||
self.in_memory_records.append(request_dict)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: queued record, in_memory=%d", len(self.in_memory_records)
|
||||
)
|
||||
|
||||
# Bound the buffer so persistent flush failures cannot grow it
|
||||
# without limit. Drop the oldest records once we exceed the cap.
|
||||
if len(self.in_memory_records) > GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
dropped = len(self.in_memory_records) - GALILEO_MAX_IN_MEMORY_RECORDS
|
||||
self.in_memory_records = self.in_memory_records[
|
||||
-GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
]
|
||||
verbose_logger.warning(
|
||||
"Galileo Logger: in-memory buffer exceeded %s records; "
|
||||
"dropped %s oldest record(s). Check Galileo connectivity/credentials.",
|
||||
GALILEO_MAX_IN_MEMORY_RECORDS,
|
||||
dropped,
|
||||
)
|
||||
|
||||
request_dict = request_record.model_dump()
|
||||
messages = kwargs.get("messages")
|
||||
if messages:
|
||||
request_dict["messages"] = messages
|
||||
self.in_memory_records.append(request_dict)
|
||||
|
||||
# Bound the buffer so persistent flush failures cannot grow it
|
||||
# without limit. Drop the oldest records once we exceed the cap.
|
||||
if len(self.in_memory_records) > GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
dropped = len(self.in_memory_records) - GALILEO_MAX_IN_MEMORY_RECORDS
|
||||
self.in_memory_records = self.in_memory_records[
|
||||
-GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
]
|
||||
verbose_logger.warning(
|
||||
"Galileo Logger: in-memory buffer exceeded %s records; "
|
||||
"dropped %s oldest record(s). Check Galileo connectivity/credentials.",
|
||||
GALILEO_MAX_IN_MEMORY_RECORDS,
|
||||
dropped,
|
||||
)
|
||||
|
||||
if len(self.in_memory_records) >= self.batch_size:
|
||||
await self.flush_in_memory_records()
|
||||
if len(self.in_memory_records) >= self.batch_size:
|
||||
await self.flush_in_memory_records()
|
||||
|
||||
async def flush_in_memory_records(self):
|
||||
if not self.in_memory_records:
|
||||
@@ -296,15 +747,23 @@ class GalileoObserve(CustomLogger):
|
||||
ingest_request = self._get_ingest_request()
|
||||
if ingest_request is None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID"
|
||||
"Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID — skipping flush"
|
||||
)
|
||||
return
|
||||
|
||||
if not await self._ensure_headers():
|
||||
verbose_logger.debug("Galileo Logger: could not set request headers")
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: could not set request headers — skipping flush"
|
||||
)
|
||||
return
|
||||
|
||||
url, payload = ingest_request
|
||||
self._log_flush_config()
|
||||
self._log_flush_payload(url=url, payload=payload)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush headers: %s",
|
||||
self._redact_headers(self.headers),
|
||||
)
|
||||
verbose_logger.debug("flushing in memory records to %s", url)
|
||||
|
||||
try:
|
||||
@@ -313,6 +772,12 @@ class GalileoObserve(CustomLogger):
|
||||
headers=self.headers,
|
||||
json=payload,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
self._log_http_status_error(error=e, url=url)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: failed to flush in memory records: %s", e
|
||||
)
|
||||
return
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: failed to flush in memory records: %s", e
|
||||
@@ -323,6 +788,11 @@ class GalileoObserve(CustomLogger):
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: successfully flushed in memory records"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush response: status=%s body=%s",
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
del self.in_memory_records[:records_in_payload]
|
||||
else:
|
||||
verbose_logger.debug("Galileo Logger: failed to flush in memory records")
|
||||
|
||||
@@ -102,6 +102,18 @@ def langfuse_client_init(
|
||||
if Version(langfuse.version.__version__) >= Version("2.6.0"):
|
||||
parameters["sdk_integration"] = "litellm"
|
||||
|
||||
if Version(langfuse.version.__version__) >= Version("2.7.3"):
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
|
||||
from ...llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
|
||||
parameters["httpx_client"] = httpx.Client(
|
||||
verify=get_ssl_configuration(),
|
||||
cert=os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate),
|
||||
)
|
||||
|
||||
client = Langfuse(**parameters)
|
||||
|
||||
return client
|
||||
|
||||
@@ -65,7 +65,15 @@ class OpenMeterLogger(CustomLogger):
|
||||
"total_tokens": response_obj["usage"].get("total_tokens"),
|
||||
}
|
||||
|
||||
user_param = kwargs.get("user", None) # end-user passed in via 'user' param
|
||||
# OPENMETER_TRUST_REQUEST_USER (default "true"): when set to "false",
|
||||
# the request-supplied `user` field is ignored and the subject is
|
||||
# resolved solely from the key-bound user_api_key_user_id. Proxies
|
||||
# serving multi-tenant traffic enable this to prevent clients from
|
||||
# forging attribution by setting `user` in the request body.
|
||||
trust_request_user = (
|
||||
os.getenv("OPENMETER_TRUST_REQUEST_USER", "true").lower() != "false"
|
||||
)
|
||||
user_param = kwargs.get("user", None) if trust_request_user else None
|
||||
|
||||
# If no user provided directly, try to get it from token user_id
|
||||
if user_param is None:
|
||||
|
||||
@@ -83,6 +83,19 @@ _VALID_CAPTURE_MODES = {
|
||||
}
|
||||
|
||||
|
||||
def _normalize_team_metadata_keys(value: Any) -> List[str]:
|
||||
"""Coerce a team-metadata allowlist from a list or comma-separated string.
|
||||
|
||||
config.yaml passes a YAML list; an env var passes a comma-separated string.
|
||||
Both collapse to a list of stripped, non-empty keys.
|
||||
"""
|
||||
if value is None:
|
||||
return []
|
||||
if isinstance(value, str):
|
||||
return [item.strip() for item in value.split(",") if item.strip()]
|
||||
return [str(item).strip() for item in value if str(item).strip()]
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpenTelemetryConfig:
|
||||
exporter: Union[str, SpanExporter] = "console"
|
||||
@@ -100,6 +113,10 @@ class OpenTelemetryConfig:
|
||||
# One of NO_CONTENT, SPAN_ONLY, EVENT_ONLY, SPAN_AND_EVENT (or "true" as legacy alias).
|
||||
capture_message_content: Optional[str] = None
|
||||
semconv_stability_opt_in: Set[OTELSemconvCategory] = field(default_factory=set)
|
||||
# Sub-keys of the team's free-form metadata stamped onto the inference span
|
||||
# under ``litellm.team.metadata``. Empty by default so none of a team's
|
||||
# metadata leaves the process until explicitly allowlisted.
|
||||
baggage_team_metadata_keys: List[str] = field(default_factory=list)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# If endpoint is specified but exporter is still the default "console",
|
||||
@@ -130,6 +147,11 @@ class OpenTelemetryConfig:
|
||||
self.semconv_stability_opt_in |= parse_semconv_opt_in(
|
||||
os.getenv(OTEL_SEMCONV_STABILITY_OPT_IN_ENV)
|
||||
)
|
||||
self.baggage_team_metadata_keys = _normalize_team_metadata_keys(
|
||||
self.baggage_team_metadata_keys
|
||||
) or _normalize_team_metadata_keys(
|
||||
os.getenv("LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS")
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls):
|
||||
@@ -188,8 +210,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
||||
meter_provider: Optional[Any] = None,
|
||||
**kwargs,
|
||||
):
|
||||
team_metadata_keys_override = kwargs.pop("baggage_team_metadata_keys", None)
|
||||
if config is None:
|
||||
config = OpenTelemetryConfig.from_env()
|
||||
if team_metadata_keys_override is not None:
|
||||
config.baggage_team_metadata_keys = _normalize_team_metadata_keys(
|
||||
team_metadata_keys_override
|
||||
)
|
||||
|
||||
self.config = config
|
||||
self.callback_name = callback_name
|
||||
@@ -985,6 +1012,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
_metadata = litellm_params.get("metadata", {}) or {}
|
||||
proxy_span = _metadata.get("litellm_parent_otel_span", None)
|
||||
|
||||
# Fallback: check litellm_metadata (used by /v1/messages and other
|
||||
# LITELLM_METADATA_ROUTES).
|
||||
if proxy_span is None:
|
||||
_litellm_metadata = litellm_params.get("litellm_metadata", {}) or {}
|
||||
proxy_span = _litellm_metadata.get("litellm_parent_otel_span", None)
|
||||
|
||||
if (
|
||||
proxy_span is not None
|
||||
and getattr(proxy_span, "name", None) == LITELLM_PROXY_REQUEST_SPAN_NAME
|
||||
@@ -1245,7 +1279,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
||||
or {}
|
||||
)
|
||||
team_metadata = self._team_metadata_json(
|
||||
raw_metadata.get("user_api_key_team_metadata")
|
||||
raw_metadata.get("user_api_key_team_metadata"),
|
||||
self.config.baggage_team_metadata_keys,
|
||||
)
|
||||
if team_metadata:
|
||||
self.safe_set_attribute(
|
||||
@@ -1268,15 +1303,20 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _team_metadata_json(value: Any) -> Optional[str]:
|
||||
"""JSON-serialize a team's metadata dict for a single span attribute.
|
||||
def _team_metadata_json(value: Any, allowed_keys: List[str]) -> Optional[str]:
|
||||
"""JSON-serialize only the allowlisted sub-keys of a team's metadata.
|
||||
|
||||
Returns ``None`` for a missing, non-dict, or empty mapping so the
|
||||
empty case is dropped rather than stamping a useless ``"{}"``.
|
||||
Returns ``None`` when nothing is allowlisted or no allowlisted key is
|
||||
present, so the empty case is dropped rather than stamping a useless
|
||||
``"{}"`` (and so a team's metadata never leaves the process until an
|
||||
operator opts each sub-key in via ``baggage_team_metadata_keys``).
|
||||
"""
|
||||
if not isinstance(value, dict) or not value:
|
||||
if not isinstance(value, dict) or not value or not allowed_keys:
|
||||
return None
|
||||
return safe_dumps(value)
|
||||
filtered = {key: value[key] for key in allowed_keys if key in value}
|
||||
if not filtered:
|
||||
return None
|
||||
return safe_dumps(filtered)
|
||||
|
||||
def _record_metrics(self, kwargs, response_obj, start_time, end_time):
|
||||
duration_s = (end_time - start_time).total_seconds()
|
||||
@@ -2635,6 +2675,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
||||
)
|
||||
|
||||
def _to_ns(self, dt):
|
||||
if dt is None:
|
||||
return int(datetime.now().timestamp() * 1e9)
|
||||
if isinstance(dt, (int, float)):
|
||||
return int(dt * 1e9)
|
||||
return int(dt.timestamp() * 1e9)
|
||||
|
||||
def _get_span_name(self, kwargs):
|
||||
@@ -2681,6 +2725,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
||||
_metadata = litellm_params.get("metadata", {}) or {}
|
||||
parent_otel_span = _metadata.get("litellm_parent_otel_span", None)
|
||||
|
||||
# Fallback: check litellm_metadata (used by /v1/messages and other
|
||||
# LITELLM_METADATA_ROUTES that store proxy-internal metadata
|
||||
# separately from the provider's native "metadata" field).
|
||||
if parent_otel_span is None:
|
||||
_litellm_metadata = litellm_params.get("litellm_metadata", {}) or {}
|
||||
parent_otel_span = _litellm_metadata.get("litellm_parent_otel_span", None)
|
||||
|
||||
# Priority 1: Explicit parent span from metadata
|
||||
if parent_otel_span is not None:
|
||||
verbose_logger.debug(
|
||||
@@ -3254,6 +3305,32 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
||||
value=int(status_code),
|
||||
)
|
||||
|
||||
def record_error_attributes_on_span(
|
||||
self,
|
||||
span: Optional[Span],
|
||||
exception: Optional[Exception],
|
||||
status_code: int,
|
||||
) -> None:
|
||||
"""Stamp structured ``error.*`` attributes on the SERVER span from the
|
||||
exception returned to the client, with ``error.code`` pinned to the real
|
||||
response status. Idempotent (overwrites); emits no exception event."""
|
||||
if span is None or exception is None:
|
||||
return
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
StandardLoggingPayloadSetup,
|
||||
)
|
||||
|
||||
error_information = StandardLoggingPayloadSetup.get_error_information(
|
||||
original_exception=exception
|
||||
)
|
||||
error_information["error_code"] = str(status_code)
|
||||
self._record_exception_on_span(
|
||||
span=span,
|
||||
kwargs={
|
||||
"standard_logging_object": {"error_information": error_information}
|
||||
},
|
||||
)
|
||||
|
||||
def set_preprocessing_duration_attribute(
|
||||
self, span: Optional[Span], container: Any
|
||||
) -> None:
|
||||
|
||||
@@ -39,20 +39,32 @@ def extract_opik_metadata(
|
||||
standard_logging_metadata: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Extract and merge Opik metadata from request and requester.
|
||||
Merge Opik metadata from three sources in increasing priority order:
|
||||
|
||||
1. user_api_key_auth_metadata– lowest priority (operator-level defaults)
|
||||
2. litellm_metadata (request)– overrides auth-key defaults
|
||||
3. requester_metadata – highest priority (e.g. proxy header overrides)
|
||||
|
||||
Args:
|
||||
litellm_metadata: Metadata from litellm_params
|
||||
standard_logging_metadata: Metadata from standard_logging_object
|
||||
litellm_metadata: Metadata from litellm_params.mak
|
||||
standard_logging_metadata: Metadata from standard_logging_object.
|
||||
|
||||
Returns:
|
||||
Merged Opik metadata dictionary
|
||||
Merged Opik metadata dictionary.
|
||||
"""
|
||||
opik_meta = litellm_metadata.get("opik", {}).copy()
|
||||
# Start with auth-key defaults (lowest priority).
|
||||
auth_meta = standard_logging_metadata.get("user_api_key_auth_metadata") or {}
|
||||
opik_meta = (auth_meta.get("opik") or {}).copy()
|
||||
|
||||
# Request-level values override auth-key defaults.
|
||||
request_opik = litellm_metadata.get("opik") or {}
|
||||
opik_meta.update(request_opik)
|
||||
|
||||
# Requester-level values win over everything else.
|
||||
requester_metadata = standard_logging_metadata.get("requester_metadata", {}) or {}
|
||||
requester_opik = requester_metadata.get("opik", {}) or {}
|
||||
opik_meta.update(requester_opik)
|
||||
if requester_opik:
|
||||
opik_meta.update(requester_opik)
|
||||
|
||||
_logging.verbose_logger.debug(
|
||||
f"litellm_opik_metadata - {json.dumps(opik_meta, default=str)}"
|
||||
|
||||
@@ -172,10 +172,13 @@ nothing here imports outside it:
|
||||
`capture_span_content` gates whether prompt/response bodies may be written as
|
||||
span attributes; it defaults **off** (`no_content`). The Baggage allowlists are
|
||||
configurable, not hard-coded: set `LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS` /
|
||||
`LITELLM_OTEL_BAGGAGE_METADATA_KEYS` (comma-separated) as env vars, or
|
||||
`baggage_promoted_keys` / `baggage_metadata_keys` (YAML lists) under
|
||||
`callback_settings.otel` in `config.yaml` — the latter reach the config through
|
||||
the logger's constructor kwargs.
|
||||
`LITELLM_OTEL_BAGGAGE_METADATA_KEYS` /
|
||||
`LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS` (comma-separated) as env vars, or
|
||||
`baggage_promoted_keys` / `baggage_metadata_keys` /
|
||||
`baggage_team_metadata_keys` (YAML lists) under `callback_settings.otel` in
|
||||
`config.yaml` — the latter reach the config through the logger's constructor
|
||||
kwargs. `baggage_team_metadata_keys` is empty by default, so none of a team's
|
||||
free-form metadata is promoted until each sub-key is explicitly allowlisted.
|
||||
- [`baggage.py`](./model/baggage.py) — the single definition of which request-identity
|
||||
values are promoted into Baggage (so child spans inherit them) and under which
|
||||
attribute keys.
|
||||
|
||||
@@ -32,20 +32,28 @@ from litellm.integrations.otel.model.payloads import (
|
||||
LLMCallSpanData,
|
||||
LLMRequestParams,
|
||||
LLMUsage,
|
||||
MCPToolCallSpanData,
|
||||
ProxyRequestSpanData,
|
||||
ServerInfo,
|
||||
ServiceSpanData,
|
||||
SpanError,
|
||||
is_mcp_tool_call,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
DB,
|
||||
HTTP,
|
||||
MCP,
|
||||
Client,
|
||||
Error,
|
||||
GenAI,
|
||||
GenAIOperation,
|
||||
GenAIProvider,
|
||||
HTTP,
|
||||
JsonRpc,
|
||||
LiteLLM,
|
||||
MCPMethod,
|
||||
Metric,
|
||||
Network,
|
||||
NetworkTransport,
|
||||
Server,
|
||||
resolve_operation,
|
||||
resolve_provider,
|
||||
@@ -69,13 +77,19 @@ __all__ = [
|
||||
"BAGGAGE_PROMOTED_KEYS",
|
||||
"DB",
|
||||
"DEFAULT_BAGGAGE_METADATA_KEYS",
|
||||
"Client",
|
||||
"Error",
|
||||
"GenAI",
|
||||
"GenAIOperation",
|
||||
"GenAIProvider",
|
||||
"HTTP",
|
||||
"JsonRpc",
|
||||
"LiteLLM",
|
||||
"MCP",
|
||||
"MCPMethod",
|
||||
"Metric",
|
||||
"Network",
|
||||
"NetworkTransport",
|
||||
"Server",
|
||||
"resolve_operation",
|
||||
"resolve_provider",
|
||||
@@ -92,11 +106,13 @@ __all__ = [
|
||||
"LLMCallSpanData",
|
||||
"LLMRequestParams",
|
||||
"LLMUsage",
|
||||
"MCPToolCallSpanData",
|
||||
"ProxyRequestSpanData",
|
||||
"RequestContext",
|
||||
"RequestIdentity",
|
||||
"ServerInfo",
|
||||
"ServiceSpanData",
|
||||
"SpanError",
|
||||
"is_mcp_tool_call",
|
||||
"promoted_baggage",
|
||||
]
|
||||
|
||||
@@ -13,6 +13,7 @@ from litellm.integrations.otel.mappers.base import AttributeMapper, SpanData
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
MCPToolCallSpanData,
|
||||
ServiceSpanData,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.providers import to_otel_span_kind
|
||||
@@ -22,6 +23,7 @@ from litellm.integrations.otel.model.spans import (
|
||||
SpanRole,
|
||||
guardrail_span_name,
|
||||
llm_call_span_name,
|
||||
mcp_tool_call_span_name,
|
||||
service_span_name,
|
||||
)
|
||||
|
||||
@@ -30,6 +32,7 @@ from litellm.integrations.otel.model.spans import (
|
||||
# have no builder here.
|
||||
_NAME_BUILDERS: dict[SpanRole, Callable[..., str]] = {
|
||||
SpanRole.LLM_CALL: llm_call_span_name,
|
||||
SpanRole.MCP_TOOL_CALL: mcp_tool_call_span_name,
|
||||
SpanRole.GUARDRAIL: guardrail_span_name,
|
||||
# DB_CALL and SERVICE are both built from ServiceSpanData; they differ only in
|
||||
# span kind (CLIENT vs INTERNAL) and attribute vocabulary, not in naming.
|
||||
@@ -121,10 +124,14 @@ class SpanEmitter:
|
||||
Return the span, or ``None`` if it was deduplicated away. ``tracer``
|
||||
overrides the bound tracer for this span, used for per-request routing.
|
||||
"""
|
||||
# Only LLM-call spans carry a dedup key; LLM-call and service spans
|
||||
# carry an ``error`` field. ``isinstance`` narrows the type for mypy and
|
||||
# keeps the engine free of duck-typed attribute reads.
|
||||
dedup_key = data.identity.call_id if isinstance(data, LLMCallSpanData) else None
|
||||
# LLM-call and MCP tool-call spans carry a dedup key (their request's
|
||||
# call id), so a sync+async double-firing coalesces. ``isinstance`` narrows
|
||||
# the type for mypy and keeps the engine free of duck-typed attribute reads.
|
||||
dedup_key = (
|
||||
data.identity.call_id
|
||||
if isinstance(data, (LLMCallSpanData, MCPToolCallSpanData))
|
||||
else None
|
||||
)
|
||||
if self._seen(dedup_key, role):
|
||||
return None
|
||||
span = self.start_span(
|
||||
@@ -160,7 +167,15 @@ class SpanEmitter:
|
||||
span.set_attribute(key, value)
|
||||
error = (
|
||||
data.error
|
||||
if isinstance(data, (LLMCallSpanData, ServiceSpanData, GuardrailSpanData))
|
||||
if isinstance(
|
||||
data,
|
||||
(
|
||||
LLMCallSpanData,
|
||||
MCPToolCallSpanData,
|
||||
ServiceSpanData,
|
||||
GuardrailSpanData,
|
||||
),
|
||||
)
|
||||
else None
|
||||
)
|
||||
if error and (error.error_type or error.message):
|
||||
|
||||
@@ -25,14 +25,15 @@ from litellm.integrations.otel.mappers import resolve_mappers
|
||||
from litellm.integrations.otel.model.metadata import (
|
||||
LLMCallEvent,
|
||||
RequestIdentity,
|
||||
guardrail_entries_from_request_data,
|
||||
model_from_request_data,
|
||||
)
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
MCPToolCallSpanData,
|
||||
ServiceSpanData,
|
||||
SpanError,
|
||||
is_mcp_tool_call,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.providers import (
|
||||
build_tracer_provider,
|
||||
@@ -43,7 +44,10 @@ from litellm.integrations.otel.model.spans import SpanRole, span_role_for_servic
|
||||
from litellm.integrations.otel.model.utils import to_ns
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.utils import StandardLoggingGuardrailInformation
|
||||
from litellm.types.utils import (
|
||||
StandardLoggingGuardrailInformation,
|
||||
StandardLoggingPayload,
|
||||
)
|
||||
|
||||
LITELLM_TRACER_NAME = "litellm"
|
||||
|
||||
@@ -200,11 +204,53 @@ class OpenTelemetryV2(CustomLogger):
|
||||
self._open_llm_calls.popitem(last=False)
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
if self._emit_mcp_tool_call(kwargs, start_time, end_time):
|
||||
return
|
||||
self._close_llm_call(kwargs, start_time, end_time)
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
if self._emit_mcp_tool_call(kwargs, start_time, end_time):
|
||||
return
|
||||
self._close_llm_call(kwargs, start_time, end_time)
|
||||
|
||||
def _emit_mcp_tool_call(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
start_time: datetime | float | None,
|
||||
end_time: datetime | float | None,
|
||||
) -> bool:
|
||||
"""Emit an MCP tool-call span when the closed request was a tool call.
|
||||
|
||||
MCP tool calls reach the success/failure callbacks like any other request
|
||||
(with ``call_type`` ``call_mcp_tool``), but they are not LLM calls and have
|
||||
no ``pre_call`` carrier — so they get their own CLIENT span here, parented
|
||||
to the request's server span. Returns whether it handled the event, so the
|
||||
caller skips the LLM-call path. The whole span is emitted at once (there is
|
||||
no boundary to open it at), deduped on the call id by the emitter.
|
||||
"""
|
||||
raw_payload = kwargs.get("standard_logging_object")
|
||||
if not raw_payload or not is_mcp_tool_call(
|
||||
cast(Mapping[str, object], raw_payload)
|
||||
):
|
||||
return False
|
||||
payload = cast("StandardLoggingPayload", raw_payload)
|
||||
data = MCPToolCallSpanData.from_standard_logging_payload(
|
||||
payload, capture_content=self.config.capture_span_content
|
||||
)
|
||||
# A stray LLM carrier from a ``pre_call`` that mis-fired for this id would
|
||||
# otherwise linger until evicted; drop it so it's neither leaked nor closed
|
||||
# as a phantom LLM span.
|
||||
if data.identity.call_id:
|
||||
self._open_llm_calls.pop(data.identity.call_id, None)
|
||||
self._emitter.emit(
|
||||
SpanRole.MCP_TOOL_CALL,
|
||||
data,
|
||||
parent_context=resolve_request_span_context(),
|
||||
start_time_ns=to_ns(start_time),
|
||||
end_time_ns=to_ns(end_time),
|
||||
)
|
||||
return True
|
||||
|
||||
def _close_llm_call(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
@@ -254,6 +300,7 @@ class OpenTelemetryV2(CustomLogger):
|
||||
data.request_model,
|
||||
promoted_keys=tuple(self.config.baggage_promoted_keys),
|
||||
metadata_keys=tuple(self.config.baggage_metadata_keys),
|
||||
team_metadata_keys=tuple(self.config.baggage_team_metadata_keys),
|
||||
)
|
||||
if bag:
|
||||
parent_ctx = set_request_baggage(bag, context=parent_ctx)
|
||||
@@ -380,6 +427,7 @@ class OpenTelemetryV2(CustomLogger):
|
||||
model,
|
||||
promoted_keys=tuple(self.config.baggage_promoted_keys),
|
||||
metadata_keys=tuple(self.config.baggage_metadata_keys),
|
||||
team_metadata_keys=tuple(self.config.baggage_team_metadata_keys),
|
||||
)
|
||||
if bag:
|
||||
# Attach (no detach): the contextvar is scoped to this request's
|
||||
@@ -419,46 +467,27 @@ class OpenTelemetryV2(CustomLogger):
|
||||
)
|
||||
return data
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: Mapping[str, Any],
|
||||
user_api_key_dict: Any,
|
||||
response: Any,
|
||||
) -> Any:
|
||||
self._emit_guardrail_spans(data)
|
||||
return response
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: Mapping[str, Any],
|
||||
original_exception: BaseException | None,
|
||||
user_api_key_dict: Any,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
self._emit_guardrail_spans(request_data)
|
||||
|
||||
def _emit_guardrail_spans(self, request_data: Mapping[str, Any]) -> None:
|
||||
def emit_guardrail_span(self, entry: "StandardLoggingGuardrailInformation") -> None:
|
||||
# Emitted by the guardrail-recording code the moment a guardrail finishes,
|
||||
# not from a post-call hook — that hook does not fire on every path (a
|
||||
# pass-through request that passes its guardrails never reaches it), which
|
||||
# left passing guardrails without a span.
|
||||
#
|
||||
# A guardrail is a sibling of the LLM call under the request's root span,
|
||||
# so parent it to the explicit anchor — not the active span, which on the
|
||||
# failure path can be the live ``auth`` phase span (post-call failure hooks
|
||||
# run from inside it on an auth rejection). Emit with the guardrail's actual
|
||||
# execution window so a pre_call guardrail is placed before the LLM call
|
||||
# rather than at post-call emission time.
|
||||
guardrails = guardrail_entries_from_request_data(request_data)
|
||||
if not guardrails:
|
||||
return
|
||||
parent_ctx = resolve_request_span_context()
|
||||
for entry in guardrails:
|
||||
data = GuardrailSpanData.from_logging_entry(
|
||||
cast("StandardLoggingGuardrailInformation", entry)
|
||||
)
|
||||
self._emitter.emit(
|
||||
SpanRole.GUARDRAIL,
|
||||
data,
|
||||
parent_context=parent_ctx,
|
||||
start_time_ns=to_ns(data.start_time),
|
||||
end_time_ns=to_ns(data.end_time),
|
||||
)
|
||||
# so parent it to the explicit anchor — never the active span, which during
|
||||
# a pre_call guardrail can be the live ``auth`` phase span. Emit with the
|
||||
# guardrail's actual execution window so a pre_call guardrail is placed
|
||||
# before the LLM call rather than at emission time. One entry in, one span
|
||||
# out — the module-level entry point routes each entry to this single
|
||||
# registered logger so a guardrail is never emitted more than once.
|
||||
data = GuardrailSpanData.from_logging_entry(entry)
|
||||
self._emitter.emit(
|
||||
SpanRole.GUARDRAIL,
|
||||
data,
|
||||
parent_context=resolve_request_span_context(),
|
||||
start_time_ns=to_ns(data.start_time),
|
||||
end_time_ns=to_ns(data.end_time),
|
||||
)
|
||||
|
||||
def create_litellm_proxy_request_started_span(
|
||||
self, start_time: datetime, headers: Mapping[str, str] | None
|
||||
@@ -479,6 +508,26 @@ def _registered_v2_logger() -> "OpenTelemetryV2 | None":
|
||||
return logger if isinstance(logger, OpenTelemetryV2) else None
|
||||
|
||||
|
||||
def emit_guardrail_span(entry: "StandardLoggingGuardrailInformation") -> None:
|
||||
"""Emit a guardrail span on the registered v2 OTel logger.
|
||||
|
||||
Called by the guardrail-recording code the moment a guardrail finishes, so a
|
||||
span is produced regardless of whether a post-call hook later runs (it does
|
||||
not on the pass-through allow path). Routes through the single canonical
|
||||
logger — the same one every other v2 entry point uses — so a guardrail
|
||||
recorded once yields exactly one span; fanning out across every reachable
|
||||
``OpenTelemetryV2`` instance double-emits the same entry. Best-effort: span
|
||||
emission must never break guardrail evaluation.
|
||||
"""
|
||||
logger = _registered_v2_logger()
|
||||
if logger is None:
|
||||
return
|
||||
try:
|
||||
logger.emit_guardrail_span(entry)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def seed_request_identity(user_api_key_dict: Any, model: Any = None) -> None:
|
||||
logger = _registered_v2_logger()
|
||||
if logger is not None:
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing_extensions import Protocol, runtime_checkable
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
MCPToolCallSpanData,
|
||||
ServiceSpanData,
|
||||
)
|
||||
|
||||
@@ -21,7 +22,7 @@ AttributeMap = dict[str, AttrValue]
|
||||
# The closed set of span-data types the engine routes through the mapper chain.
|
||||
# Server spans (PROXY_REQUEST + management routes) belong to the mounted FastAPI
|
||||
# instrumentor, not the mapper chain.
|
||||
SpanData = LLMCallSpanData | GuardrailSpanData | ServiceSpanData
|
||||
SpanData = LLMCallSpanData | MCPToolCallSpanData | GuardrailSpanData | ServiceSpanData
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
|
||||
@@ -14,10 +14,18 @@ from litellm.integrations.otel.mappers.utils import collect, drop_none
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
MCPToolCallSpanData,
|
||||
ServiceSpanData,
|
||||
ToolDefinition,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import DB, Error, GenAI, LiteLLM, Server
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
DB,
|
||||
MCP,
|
||||
Error,
|
||||
GenAI,
|
||||
LiteLLM,
|
||||
Server,
|
||||
)
|
||||
from litellm.integrations.otel.model.spans import db_system
|
||||
|
||||
|
||||
@@ -64,6 +72,18 @@ class GenAIMapper:
|
||||
"parameters": lambda t: t.parameters_json or None,
|
||||
}
|
||||
|
||||
_MCP_ATTRS: dict[str, Callable[[MCPToolCallSpanData], AttrValue | None]] = {
|
||||
GenAI.OPERATION_NAME: lambda d: d.operation.value,
|
||||
MCP.METHOD_NAME: lambda d: d.method,
|
||||
MCP.SESSION_ID: lambda d: d.session_id,
|
||||
GenAI.TOOL_NAME: lambda d: d.tool_name or None,
|
||||
GenAI.TOOL_CALL_ARGUMENTS: lambda d: d.arguments_json,
|
||||
GenAI.TOOL_CALL_RESULT: lambda d: d.result_json,
|
||||
LiteLLM.MCP_SERVER_NAME: lambda d: d.server_name,
|
||||
LiteLLM.CALL_ID: lambda d: d.identity.call_id or None,
|
||||
f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost,
|
||||
}
|
||||
|
||||
_GUARDRAIL_ATTRS: dict[str, Callable[[GuardrailSpanData], AttrValue | None]] = {
|
||||
LiteLLM.GUARDRAIL_NAME: lambda d: d.guardrail_name,
|
||||
LiteLLM.GUARDRAIL_MODE: lambda d: d.mode,
|
||||
@@ -92,6 +112,8 @@ class GenAIMapper:
|
||||
match data:
|
||||
case LLMCallSpanData():
|
||||
return self._llm_call(data)
|
||||
case MCPToolCallSpanData():
|
||||
return collect(self._MCP_ATTRS, data)
|
||||
case GuardrailSpanData():
|
||||
return self._guardrail(data)
|
||||
case ServiceSpanData():
|
||||
|
||||
@@ -6,26 +6,36 @@ LLM-call span so that child spans (guardrail, service) inherit them.
|
||||
the allowlisted keys onto every span.
|
||||
|
||||
This module is the single place baggage is defined: ``_PROMOTABLE`` maps each
|
||||
promotable attribute key to how its value is read, and the two ``*_KEYS``
|
||||
defaults select what is promoted unless the config overrides them.
|
||||
promotable attribute key to how its value is read, and the ``*_KEYS`` defaults
|
||||
select what is promoted unless the config overrides them. ``TEAM_METADATA``'s
|
||||
extractor filters the team's free-form metadata to the sub-keys an operator
|
||||
allowlists via ``baggage_team_metadata_keys`` (default none), so the blob is
|
||||
never promoted whole.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.integrations.otel.model.metadata import RequestIdentity
|
||||
from litellm.integrations.otel.model.semconv import GenAI, LiteLLM
|
||||
|
||||
# Attribute key -> value extractor over (identity, request_model). The single
|
||||
# definition of what may be promoted and under which key.
|
||||
_PROMOTABLE: Final[dict[str, Callable[[RequestIdentity, str | None], str | None]]] = {
|
||||
LiteLLM.TEAM_ID: lambda identity, model: identity.team_id,
|
||||
LiteLLM.TEAM_ALIAS: lambda identity, model: identity.team_alias,
|
||||
LiteLLM.TEAM_METADATA: lambda identity, model: identity.team_metadata,
|
||||
LiteLLM.KEY_HASH: lambda identity, model: identity.key_hash,
|
||||
LiteLLM.END_USER: lambda identity, model: identity.end_user,
|
||||
GenAI.REQUEST_MODEL: lambda identity, model: model,
|
||||
LiteLLM.PROVIDER_MODEL: lambda identity, model: identity.provider_model,
|
||||
# Attribute key -> value extractor over (identity, request_model,
|
||||
# team_metadata_keys). The single definition of what may be promoted and under
|
||||
# which key. Only the ``TEAM_METADATA`` extractor consults team_metadata_keys
|
||||
# (to filter the team's metadata to an allowlist); the rest ignore it.
|
||||
_PROMOTABLE: Final[
|
||||
dict[str, Callable[[RequestIdentity, str | None, tuple[str, ...]], str | None]]
|
||||
] = {
|
||||
LiteLLM.TEAM_ID: lambda identity, model, team_metadata_keys: identity.team_id,
|
||||
LiteLLM.TEAM_ALIAS: lambda identity, model, team_metadata_keys: identity.team_alias,
|
||||
LiteLLM.TEAM_METADATA: lambda identity, model, team_metadata_keys: _filtered_team_metadata_json(
|
||||
identity.team_metadata, team_metadata_keys
|
||||
),
|
||||
LiteLLM.KEY_HASH: lambda identity, model, team_metadata_keys: identity.key_hash,
|
||||
LiteLLM.END_USER: lambda identity, model, team_metadata_keys: identity.end_user,
|
||||
GenAI.REQUEST_MODEL: lambda identity, model, team_metadata_keys: model,
|
||||
LiteLLM.PROVIDER_MODEL: lambda identity, model, team_metadata_keys: identity.provider_model,
|
||||
}
|
||||
|
||||
# Keys promoted by default (a subset of ``_PROMOTABLE``). ``END_USER`` is
|
||||
@@ -50,23 +60,31 @@ DEFAULT_BAGGAGE_METADATA_KEYS: Final[tuple[str, ...]] = (
|
||||
"requester_ip_address",
|
||||
)
|
||||
|
||||
# Sub-keys of the team's free-form metadata eligible for promotion under
|
||||
# ``litellm.team.metadata``. Empty by default: a team's metadata can hold
|
||||
# arbitrary operator data, so none of it is promoted until each key is
|
||||
# explicitly allowlisted via ``config.baggage_team_metadata_keys``.
|
||||
DEFAULT_BAGGAGE_TEAM_METADATA_KEYS: Final[tuple[str, ...]] = ()
|
||||
|
||||
|
||||
def promoted_baggage(
|
||||
identity: RequestIdentity,
|
||||
request_model: str | None,
|
||||
promoted_keys: tuple[str, ...],
|
||||
metadata_keys: tuple[str, ...] = DEFAULT_BAGGAGE_METADATA_KEYS,
|
||||
team_metadata_keys: tuple[str, ...] = DEFAULT_BAGGAGE_TEAM_METADATA_KEYS,
|
||||
) -> dict[str, str]:
|
||||
"""Identity values to write into Baggage, filtered to ``promoted_keys``.
|
||||
|
||||
``promoted_keys`` selects from ``_PROMOTABLE``; ``metadata_keys`` selects
|
||||
sub-keys of ``identity.metadata`` to promote under ``litellm.metadata.*``.
|
||||
Empty values are dropped.
|
||||
sub-keys of ``identity.metadata`` to promote under ``litellm.metadata.*``;
|
||||
``team_metadata_keys`` selects sub-keys of the team's metadata to promote
|
||||
under ``litellm.team.metadata``. Empty values are dropped.
|
||||
"""
|
||||
out: dict[str, str] = {}
|
||||
for key, extract in _PROMOTABLE.items():
|
||||
if key in promoted_keys:
|
||||
value = extract(identity, request_model)
|
||||
value = extract(identity, request_model, team_metadata_keys)
|
||||
if value:
|
||||
out[key] = value
|
||||
for meta_key in metadata_keys:
|
||||
@@ -74,3 +92,21 @@ def promoted_baggage(
|
||||
if value:
|
||||
out[f"{LiteLLM.METADATA_PREFIX}{meta_key}"] = value
|
||||
return out
|
||||
|
||||
|
||||
def _filtered_team_metadata_json(
|
||||
metadata: Mapping[str, object] | None,
|
||||
allowed_keys: tuple[str, ...],
|
||||
) -> str | None:
|
||||
"""JSON-serialize only the allowlisted sub-keys of a team's metadata.
|
||||
|
||||
Returns ``None`` when nothing is allowlisted or no allowlisted key is
|
||||
present, so the empty case is dropped rather than promoting ``"{}"``. Keys
|
||||
are sorted for a stable, diff-friendly value.
|
||||
"""
|
||||
if not isinstance(metadata, Mapping) or not allowed_keys:
|
||||
return None
|
||||
filtered = {key: metadata[key] for key in allowed_keys if key in metadata}
|
||||
if not filtered:
|
||||
return None
|
||||
return json.dumps(filtered, default=str, sort_keys=True)
|
||||
|
||||
@@ -9,6 +9,7 @@ from typing_extensions import Annotated
|
||||
from litellm.integrations.otel.model.baggage import (
|
||||
BAGGAGE_PROMOTED_KEYS,
|
||||
DEFAULT_BAGGAGE_METADATA_KEYS,
|
||||
DEFAULT_BAGGAGE_TEAM_METADATA_KEYS,
|
||||
)
|
||||
|
||||
#: Master feature-flag env var. The logger is inert until this is truthy.
|
||||
@@ -168,10 +169,25 @@ class OpenTelemetryV2Config(BaseSettings):
|
||||
"``callback_settings.otel.baggage_metadata_keys`` in config.yaml."
|
||||
),
|
||||
)
|
||||
baggage_team_metadata_keys: Annotated[List[str], NoDecode] = Field(
|
||||
default_factory=lambda: list(DEFAULT_BAGGAGE_TEAM_METADATA_KEYS),
|
||||
validation_alias=AliasChoices(
|
||||
"baggage_team_metadata_keys", "LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS"
|
||||
),
|
||||
description=(
|
||||
"Sub-keys of the team's free-form metadata promoted under "
|
||||
"``litellm.team.metadata``. Empty by default so none of a team's "
|
||||
"metadata leaves the process until explicitly allowlisted. Configure "
|
||||
"via the ``LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS`` env var "
|
||||
"(comma-separated) or "
|
||||
"``callback_settings.otel.baggage_team_metadata_keys`` in config.yaml."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator(
|
||||
"baggage_promoted_keys",
|
||||
"baggage_metadata_keys",
|
||||
"baggage_team_metadata_keys",
|
||||
"mapper_names",
|
||||
mode="before",
|
||||
)
|
||||
|
||||
@@ -36,7 +36,6 @@ model. They coincide on the SDK path, which is correct.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Mapping, cast
|
||||
|
||||
@@ -53,8 +52,10 @@ class RequestIdentity:
|
||||
call_id: str | None = None
|
||||
team_id: str | None = None
|
||||
team_alias: str | None = None
|
||||
# The team's free-form metadata dict, JSON-serialized (empty/missing -> None).
|
||||
team_metadata: str | None = None
|
||||
# The team's free-form metadata, carried raw (empty/missing -> None) and
|
||||
# filtered to an operator allowlist only at Baggage-promotion time, so an
|
||||
# unconfigured deployment never promotes any of it.
|
||||
team_metadata: Mapping[str, Any] | None = None
|
||||
key_hash: str | None = None
|
||||
end_user: str | None = None
|
||||
# The model litellm dispatched to the provider. Only known once the call
|
||||
@@ -86,7 +87,7 @@ class RequestIdentity:
|
||||
or as_str(raw_meta.get("team_id")),
|
||||
team_alias=as_str(raw_meta.get("user_api_key_team_alias"))
|
||||
or as_str(raw_meta.get("team_alias")),
|
||||
team_metadata=_team_metadata_json(
|
||||
team_metadata=_team_metadata_dict(
|
||||
raw_meta.get("user_api_key_team_metadata")
|
||||
),
|
||||
key_hash=as_str(raw_meta.get("user_api_key_hash")),
|
||||
@@ -121,7 +122,7 @@ class RequestIdentity:
|
||||
return cls(
|
||||
team_id=as_str(get("team_id")),
|
||||
team_alias=as_str(get("team_alias")),
|
||||
team_metadata=_team_metadata_json(get("team_metadata")),
|
||||
team_metadata=_team_metadata_dict(get("team_metadata")),
|
||||
key_hash=as_str(get("api_key")),
|
||||
end_user=as_str(get("end_user_id")),
|
||||
# ``provider_model`` is unknown at the auth boundary — routing hasn't
|
||||
@@ -254,26 +255,6 @@ def model_from_request_data(data: object) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def guardrail_entries_from_request_data(
|
||||
request_data: Mapping[str, Any],
|
||||
) -> list[dict]:
|
||||
"""The guardrail-information dicts buried in ``metadata`` of a post-call dict.
|
||||
|
||||
``standard_logging_guardrail_information`` is stored as either a single dict
|
||||
or a list of them; normalize to a list of dicts (dropping non-dict noise) so
|
||||
the caller just iterates. Empty list when none are present.
|
||||
"""
|
||||
metadata = request_data.get("metadata")
|
||||
if not isinstance(metadata, Mapping):
|
||||
return []
|
||||
info = metadata.get("standard_logging_guardrail_information")
|
||||
if isinstance(info, Mapping):
|
||||
return [cast(dict, info)]
|
||||
if isinstance(info, list):
|
||||
return [entry for entry in info if isinstance(entry, dict)]
|
||||
return []
|
||||
|
||||
|
||||
def resolve_provider_model(payload: "StandardLoggingPayload") -> str | None:
|
||||
"""The model litellm dispatched to the provider, from the payload.
|
||||
|
||||
@@ -300,16 +281,14 @@ def _model_info_id(model_info: object) -> str | None:
|
||||
return None
|
||||
|
||||
|
||||
def _team_metadata_json(value: object) -> str | None:
|
||||
"""JSON-serialize a team's metadata dict for a single Baggage value.
|
||||
def _team_metadata_dict(value: object) -> Mapping[str, Any] | None:
|
||||
"""The team's free-form metadata as a raw mapping, or ``None`` when missing
|
||||
or empty.
|
||||
|
||||
Returns ``None`` for a missing, non-dict, or empty mapping so the empty case
|
||||
is dropped rather than promoting a useless ``"{}"``. Keys are sorted for a
|
||||
stable, diff-friendly serialization.
|
||||
Carried raw on the identity and filtered to an operator allowlist only at
|
||||
Baggage-promotion time (see ``baggage.promoted_baggage``), so an empty case
|
||||
is dropped rather than carrying a useless ``{}``.
|
||||
"""
|
||||
if not isinstance(value, Mapping) or not value:
|
||||
return None
|
||||
try:
|
||||
return json.dumps(value, default=str, sort_keys=True)
|
||||
except Exception:
|
||||
return None
|
||||
if isinstance(value, Mapping) and value:
|
||||
return dict(value)
|
||||
return None
|
||||
|
||||
@@ -14,6 +14,7 @@ from litellm.integrations.otel.model.metadata import (
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
GenAIOperation,
|
||||
MCPMethod,
|
||||
resolve_operation,
|
||||
resolve_provider,
|
||||
)
|
||||
@@ -35,11 +36,13 @@ __all__ = [
|
||||
"LLMCallSpanData",
|
||||
"LLMRequestParams",
|
||||
"LLMUsage",
|
||||
"MCPToolCallSpanData",
|
||||
"ProxyRequestSpanData",
|
||||
"ServerInfo",
|
||||
"ServiceSpanData",
|
||||
"SpanError",
|
||||
"ToolDefinition",
|
||||
"is_mcp_tool_call",
|
||||
]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -309,6 +312,77 @@ class LLMCallSpanData:
|
||||
)
|
||||
|
||||
|
||||
# --- the MCP tool-call model ------------------------------------------------- #
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MCPToolCallSpanData:
|
||||
"""One MCP ``tools/call`` execution, parsed from a closed request's payload.
|
||||
|
||||
The proxy is an MCP *client* to the upstream server it forwards the call to,
|
||||
so this is a CLIENT span. ``arguments_json``/``result_json`` are the tool's
|
||||
input/output — sensitive content, so they're only retained when content
|
||||
capture is enabled, mirroring ``LLMCallSpanData``'s message bodies.
|
||||
"""
|
||||
|
||||
operation: GenAIOperation
|
||||
method: str
|
||||
tool_name: str
|
||||
server_name: str | None
|
||||
session_id: str | None
|
||||
arguments_json: str | None
|
||||
result_json: str | None
|
||||
error: SpanError | None
|
||||
response_cost: float | None
|
||||
identity: RequestIdentity
|
||||
|
||||
@classmethod
|
||||
def from_standard_logging_payload(
|
||||
cls, payload: "StandardLoggingPayload", capture_content: bool = False
|
||||
) -> "MCPToolCallSpanData":
|
||||
meta = _mcp_tool_call_metadata(cast(Mapping[str, object], payload))
|
||||
return cls(
|
||||
operation=resolve_operation(as_str(payload.get("call_type"))),
|
||||
method=MCPMethod.TOOLS_CALL.value,
|
||||
tool_name=as_str(meta.get("name")) or "",
|
||||
server_name=as_str(meta.get("mcp_server_name")),
|
||||
session_id=as_str(meta.get("mcp_session_id")),
|
||||
arguments_json=(
|
||||
_json_or_none(meta.get("arguments"))
|
||||
if capture_content and meta.get("arguments") is not None
|
||||
else None
|
||||
),
|
||||
result_json=(
|
||||
_json_or_none(meta.get("result"))
|
||||
if capture_content and meta.get("result") is not None
|
||||
else None
|
||||
),
|
||||
error=_parse_error(payload),
|
||||
response_cost=as_float(payload.get("response_cost")),
|
||||
identity=RequestContext.from_standard_logging_payload(payload).identity,
|
||||
)
|
||||
|
||||
|
||||
def _mcp_tool_call_metadata(payload: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""The MCP gateway's tool-call metadata, which lives under
|
||||
``StandardLoggingPayload.metadata`` (a ``StandardLoggingMetadata`` key), not
|
||||
at the payload's top level."""
|
||||
metadata = payload.get("metadata")
|
||||
if not isinstance(metadata, Mapping):
|
||||
return {}
|
||||
meta = metadata.get("mcp_tool_call_metadata")
|
||||
return meta if isinstance(meta, Mapping) else {}
|
||||
|
||||
|
||||
def is_mcp_tool_call(payload: Mapping[str, object]) -> bool:
|
||||
"""Whether a closed request's payload is an MCP tool call rather than an LLM
|
||||
call — true when the MCP gateway stamped its tool-call metadata, or the call
|
||||
type says so on a path that hasn't populated the metadata yet."""
|
||||
return bool(_mcp_tool_call_metadata(payload)) or (
|
||||
payload.get("call_type") == "call_mcp_tool"
|
||||
)
|
||||
|
||||
|
||||
# --- service event_metadata sanitization ------------------------------------ #
|
||||
|
||||
# Substrings (case-insensitive) of keys that must never reach a span: secrets,
|
||||
|
||||
@@ -16,7 +16,7 @@ class GenAIOperation(str, Enum):
|
||||
GENERATE_CONTENT = "generate_content"
|
||||
CREATE_AGENT = "create_agent" # reserved for future agent spans
|
||||
INVOKE_AGENT = "invoke_agent" # reserved for future agent spans
|
||||
EXECUTE_TOOL = "execute_tool" # reserved for future tool spans
|
||||
EXECUTE_TOOL = "execute_tool" # MCP tool-call spans
|
||||
|
||||
|
||||
class GenAIProvider(str, Enum):
|
||||
@@ -38,6 +38,16 @@ class GenAIProvider(str, Enum):
|
||||
IBM_WATSONX_AI = "ibm.watsonx.ai"
|
||||
|
||||
|
||||
class MCPMethod(str, Enum):
|
||||
"""Well-known values for ``mcp.method.name`` that litellm's MCP gateway
|
||||
serves. The value is the JSON-RPC method exactly as it travels on the wire."""
|
||||
|
||||
TOOLS_CALL = "tools/call"
|
||||
TOOLS_LIST = "tools/list"
|
||||
PROMPTS_GET = "prompts/get"
|
||||
PROMPTS_LIST = "prompts/list"
|
||||
|
||||
|
||||
class GenAI:
|
||||
"""Canonical OTel GenAI span-attribute keys."""
|
||||
|
||||
@@ -68,11 +78,68 @@ class GenAI:
|
||||
SYSTEM_INSTRUCTIONS: Final = "gen_ai.system_instructions"
|
||||
OUTPUT_TYPE: Final = "gen_ai.output.type"
|
||||
CONVERSATION_ID: Final = "gen_ai.conversation.id"
|
||||
# agent / tool (reserved)
|
||||
# agent (reserved)
|
||||
AGENT_ID: Final = "gen_ai.agent.id"
|
||||
AGENT_NAME: Final = "gen_ai.agent.name"
|
||||
# tool / tool-call (stamped on MCP tool-call spans). Arguments and result are
|
||||
# the tool's input/output payloads — sensitive, so they're opt-in and gated by
|
||||
# the same content-capture mode as prompt/response content.
|
||||
TOOL_NAME: Final = "gen_ai.tool.name"
|
||||
TOOL_CALL_ID: Final = "gen_ai.tool.call.id"
|
||||
TOOL_CALL_ARGUMENTS: Final = "gen_ai.tool.call.arguments"
|
||||
TOOL_CALL_RESULT: Final = "gen_ai.tool.call.result"
|
||||
# prompt (MCP ``prompts/get`` etc.)
|
||||
PROMPT_NAME: Final = "gen_ai.prompt.name"
|
||||
|
||||
|
||||
class MCP:
|
||||
"""OTel GenAI MCP (Model Context Protocol) span-attribute keys.
|
||||
|
||||
``METHOD_NAME`` is the only key litellm populates from a closed request today;
|
||||
the rest are part of the convention's vocabulary and are stamped when the
|
||||
corresponding signal (session, protocol version, resource) is available.
|
||||
"""
|
||||
|
||||
METHOD_NAME: Final = "mcp.method.name"
|
||||
SESSION_ID: Final = "mcp.session.id"
|
||||
PROTOCOL_VERSION: Final = "mcp.protocol.version"
|
||||
RESOURCE_URI: Final = "mcp.resource.uri"
|
||||
|
||||
|
||||
class JsonRpc:
|
||||
"""JSON-RPC keys carried on MCP spans. The error/status code lives in the
|
||||
``rpc.*`` namespace per semconv, not ``jsonrpc.*``."""
|
||||
|
||||
REQUEST_ID: Final = "jsonrpc.request.id"
|
||||
PROTOCOL_VERSION: Final = "jsonrpc.protocol.version"
|
||||
RESPONSE_STATUS_CODE: Final = "rpc.response.status_code"
|
||||
|
||||
|
||||
class NetworkTransport(str, Enum):
|
||||
"""Well-known values for ``network.transport``."""
|
||||
|
||||
TCP = "tcp"
|
||||
UDP = "udp"
|
||||
QUIC = "quic"
|
||||
UNIX = "unix"
|
||||
PIPE = "pipe"
|
||||
|
||||
|
||||
class Network:
|
||||
"""OTel network keys, recommended on MCP spans to describe the transport
|
||||
carrying the JSON-RPC messages (stdio pipe, HTTP, websocket, …)."""
|
||||
|
||||
PROTOCOL_NAME: Final = "network.protocol.name"
|
||||
PROTOCOL_VERSION: Final = "network.protocol.version"
|
||||
TRANSPORT: Final = "network.transport"
|
||||
|
||||
|
||||
class Client:
|
||||
"""Peer (client) network keys, stamped on MCP *server* spans the same way
|
||||
``server.*`` is stamped on client spans."""
|
||||
|
||||
ADDRESS: Final = "client.address"
|
||||
PORT: Final = "client.port"
|
||||
|
||||
|
||||
class Error:
|
||||
@@ -137,6 +204,10 @@ class LiteLLM:
|
||||
SERVICE_NAME: Final = "litellm.service.name"
|
||||
SERVICE_CALL_TYPE: Final = "litellm.service.call_type"
|
||||
PREPROCESSING_MS: Final = "litellm.preprocessing.duration_ms"
|
||||
# The logical name of the MCP server a tool call was routed to. There is no
|
||||
# semconv key for an MCP server's *name* (the convention uses ``server.address``
|
||||
# for its network location), so it lives under the vendor namespace.
|
||||
MCP_SERVER_NAME: Final = "litellm.mcp.server.name"
|
||||
|
||||
|
||||
class Metric:
|
||||
@@ -179,6 +250,7 @@ _OPERATION_BY_CALL_TYPE: dict[str, GenAIOperation] = {
|
||||
"aembedding": GenAIOperation.EMBEDDINGS,
|
||||
"responses": GenAIOperation.CHAT,
|
||||
"aresponses": GenAIOperation.CHAT,
|
||||
"call_mcp_tool": GenAIOperation.EXECUTE_TOOL,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -46,6 +46,7 @@ if TYPE_CHECKING:
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
MCPToolCallSpanData,
|
||||
ProxyRequestSpanData,
|
||||
ServiceSpanData,
|
||||
)
|
||||
@@ -54,6 +55,7 @@ if TYPE_CHECKING:
|
||||
class SpanRole(str, Enum):
|
||||
PROXY_REQUEST = "proxy_request"
|
||||
LLM_CALL = "llm_call"
|
||||
MCP_TOOL_CALL = "mcp_tool_call"
|
||||
GUARDRAIL = "guardrail"
|
||||
DB_CALL = "db_call"
|
||||
SERVICE = "service"
|
||||
@@ -81,6 +83,11 @@ SPAN_REGISTRY: dict[SpanRole, SpanSpec] = {
|
||||
SpanRole.LLM_CALL: SpanSpec(
|
||||
SpanRole.LLM_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST
|
||||
),
|
||||
# The proxy is an MCP client to the upstream server it dispatches the tool
|
||||
# call to, so this is a CLIENT span, sibling of the LLM call under the request.
|
||||
SpanRole.MCP_TOOL_CALL: SpanSpec(
|
||||
SpanRole.MCP_TOOL_CALL, LiteLLMSpanKind.CLIENT, parent=SpanRole.PROXY_REQUEST
|
||||
),
|
||||
SpanRole.GUARDRAIL: SpanSpec(
|
||||
SpanRole.GUARDRAIL, LiteLLMSpanKind.INTERNAL, parent=SpanRole.PROXY_REQUEST
|
||||
),
|
||||
@@ -165,6 +172,11 @@ def llm_call_span_name(data: "LLMCallSpanData") -> str:
|
||||
return f"{data.operation.value} {model}".strip()
|
||||
|
||||
|
||||
def mcp_tool_call_span_name(data: "MCPToolCallSpanData") -> str:
|
||||
"""``"{mcp.method.name} {tool}"`` e.g. ``"tools/call get-weather"`` (MCP semconv)."""
|
||||
return f"{data.method} {data.tool_name}".strip()
|
||||
|
||||
|
||||
def proxy_request_span_name(data: "ProxyRequestSpanData") -> str:
|
||||
"""``"{method} {route}"`` (HTTP semconv)."""
|
||||
return f"{data.http_method} {data.route}".strip()
|
||||
|
||||
@@ -511,6 +511,23 @@ class PrometheusLogger(CustomLogger):
|
||||
labelnames=self.get_labels_for_metric("litellm_cached_tokens_metric"),
|
||||
)
|
||||
|
||||
# Provider prompt-caching metrics
|
||||
self.litellm_provider_cache_read_input_tokens_metric = self._counter_factory(
|
||||
name="litellm_provider_cache_read_input_tokens_metric",
|
||||
documentation="Total prompt/input tokens read from provider prompt cache (e.g. OpenAI/Anthropic/Gemini/Bedrock)",
|
||||
labelnames=self.get_labels_for_metric(
|
||||
"litellm_provider_cache_read_input_tokens_metric"
|
||||
),
|
||||
)
|
||||
|
||||
self.litellm_provider_cache_creation_input_tokens_metric = self._counter_factory(
|
||||
name="litellm_provider_cache_creation_input_tokens_metric",
|
||||
documentation="Total prompt/input tokens written to provider prompt cache (e.g. Anthropic/Bedrock)",
|
||||
labelnames=self.get_labels_for_metric(
|
||||
"litellm_provider_cache_creation_input_tokens_metric"
|
||||
),
|
||||
)
|
||||
|
||||
# User and Team count metrics
|
||||
self.litellm_total_users_metric = self._gauge_factory(
|
||||
"litellm_total_users",
|
||||
@@ -1458,11 +1475,11 @@ class PrometheusLogger(CustomLogger):
|
||||
"""
|
||||
cache_hit = standard_logging_payload.get("cache_hit")
|
||||
|
||||
# Only track if cache_hit has a definite value (True or False)
|
||||
if cache_hit is None:
|
||||
return
|
||||
|
||||
if cache_hit is True:
|
||||
# Historically these metrics only tracked LiteLLM caching.
|
||||
# Provider prompt-caching metrics are still emitted below.
|
||||
pass
|
||||
elif cache_hit is True:
|
||||
# Increment cache hits counter
|
||||
PrometheusLogger._inc_labeled_counter(
|
||||
self,
|
||||
@@ -1493,6 +1510,51 @@ class PrometheusLogger(CustomLogger):
|
||||
label_context=label_context,
|
||||
)
|
||||
|
||||
# Provider prompt caching metrics are independent of LiteLLM cache_hit.
|
||||
provider_cache_read_tokens = 0
|
||||
provider_cache_creation_tokens = 0
|
||||
usage_obj = (standard_logging_payload.get("metadata", {}) or {}).get(
|
||||
"usage_object"
|
||||
)
|
||||
if isinstance(usage_obj, dict):
|
||||
# Prefer explicit provider cache fields when available.
|
||||
_read = usage_obj.get("cache_read_input_tokens")
|
||||
_write = usage_obj.get("cache_creation_input_tokens")
|
||||
|
||||
if isinstance(_read, int):
|
||||
provider_cache_read_tokens = _read
|
||||
if isinstance(_write, int):
|
||||
provider_cache_creation_tokens = _write
|
||||
|
||||
# Fallback to prompt_tokens_details.cached_tokens (common normalization point).
|
||||
# Only fallback when the explicit field is genuinely absent (None).
|
||||
if _read is None:
|
||||
prompt_details = usage_obj.get("prompt_tokens_details")
|
||||
if isinstance(prompt_details, dict):
|
||||
cached_tokens = prompt_details.get("cached_tokens")
|
||||
if isinstance(cached_tokens, int):
|
||||
provider_cache_read_tokens = cached_tokens
|
||||
|
||||
if provider_cache_read_tokens > 0:
|
||||
PrometheusLogger._inc_labeled_counter(
|
||||
self,
|
||||
self.litellm_provider_cache_read_input_tokens_metric,
|
||||
"litellm_provider_cache_read_input_tokens_metric",
|
||||
enum_values,
|
||||
label_context=label_context,
|
||||
amount=float(provider_cache_read_tokens),
|
||||
)
|
||||
|
||||
if provider_cache_creation_tokens > 0:
|
||||
PrometheusLogger._inc_labeled_counter(
|
||||
self,
|
||||
self.litellm_provider_cache_creation_input_tokens_metric,
|
||||
"litellm_provider_cache_creation_input_tokens_metric",
|
||||
enum_values,
|
||||
label_context=label_context,
|
||||
amount=float(provider_cache_creation_tokens),
|
||||
)
|
||||
|
||||
async def _increment_remaining_budget_metrics(
|
||||
self,
|
||||
user_api_team: Optional[str],
|
||||
@@ -2628,7 +2690,7 @@ class PrometheusLogger(CustomLogger):
|
||||
Args:
|
||||
guardrail_name: Name of the guardrail
|
||||
latency_seconds: Execution latency in seconds
|
||||
status: "success" or "error"
|
||||
status: "success", "error", or "intervened"
|
||||
error_type: Type of error if any, None otherwise
|
||||
hook_type: "pre_call", "during_call", or "post_call"
|
||||
"""
|
||||
|
||||
@@ -244,6 +244,9 @@ search_tools:
|
||||
- search_tool_name: "my-tavily-tool"
|
||||
litellm_params:
|
||||
search_provider: "tavily"
|
||||
- search_tool_name: "my-you-com-tool"
|
||||
litellm_params:
|
||||
search_provider: "you_com"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
@@ -87,6 +87,7 @@ class ExceptionCheckers:
|
||||
"is longer than the model's context length",
|
||||
"input tokens exceed the configured limit",
|
||||
"`inputs` tokens + `max_new_tokens` must be",
|
||||
"exceeds the available context size", # llama.cpp/Lemonade
|
||||
"exceeds the maximum number of tokens allowed", # Gemini
|
||||
]
|
||||
for substring in known_exception_substrings:
|
||||
@@ -891,12 +892,14 @@ def exception_type( # type: ignore # noqa: PLR0915
|
||||
response=getattr(original_exception, "response", None),
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "model's maximum context limit" in error_str:
|
||||
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
|
||||
exception_mapping_worked = True
|
||||
raise ContextWindowExceededError(
|
||||
message=f"{custom_llm_provider.capitalize()}Exception: Context Window Error - {error_str}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "token_quota_reached" in error_str:
|
||||
exception_mapping_worked = True
|
||||
|
||||
@@ -47,8 +47,9 @@ async def async_completion_with_fallbacks(**kwargs):
|
||||
completion_kwargs = safe_deep_copy(base_kwargs)
|
||||
# Handle dictionary fallback configurations
|
||||
if isinstance(fallback, dict):
|
||||
model = fallback.pop("model", original_model)
|
||||
completion_kwargs.update(fallback)
|
||||
fallback_config = safe_deep_copy(dict(fallback))
|
||||
model = fallback_config.pop("model", original_model)
|
||||
completion_kwargs.update(fallback_config)
|
||||
else:
|
||||
model = fallback
|
||||
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import re
|
||||
from typing import Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
@@ -71,6 +72,25 @@ def _is_azure_claude_model(model: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
_CLAUDE_PATTERN = re.compile(r"^claude-[a-z]+-\d+-\d+(?:-\d{8})?$", re.IGNORECASE)
|
||||
|
||||
|
||||
def _matches_claude_model_pattern(model: str) -> bool:
|
||||
"""
|
||||
Check if a model string matches the Claude model naming pattern.
|
||||
|
||||
Matches patterns like:
|
||||
- claude-opus-4-7
|
||||
- claude-sonnet-4-6
|
||||
- claude-haiku-4-5
|
||||
- claude-opus-5-1-20270101 (with optional date suffix)
|
||||
|
||||
This allows future Claude models to be routed to the Anthropic provider
|
||||
without requiring updates to model_prices_and_context_window.json.
|
||||
"""
|
||||
return _CLAUDE_PATTERN.match(model) is not None
|
||||
|
||||
|
||||
def handle_cohere_chat_model_custom_llm_provider(
|
||||
model: str, custom_llm_provider: Optional[str] = None
|
||||
) -> Tuple[str, Optional[str]]:
|
||||
@@ -353,6 +373,9 @@ def get_llm_provider( # noqa: PLR0915
|
||||
elif endpoint == "https://api.lambda.ai/v1":
|
||||
custom_llm_provider = "lambda_ai"
|
||||
dynamic_api_key = get_secret_str("LAMBDA_API_KEY")
|
||||
elif endpoint == "https://api.inceptionlabs.ai/v1":
|
||||
custom_llm_provider = "inception"
|
||||
dynamic_api_key = get_secret_str("INCEPTION_API_KEY")
|
||||
elif endpoint == "https://api.hyperbolic.xyz/v1":
|
||||
custom_llm_provider = "hyperbolic"
|
||||
dynamic_api_key = get_secret_str("HYPERBOLIC_API_KEY")
|
||||
@@ -398,6 +421,9 @@ def get_llm_provider( # noqa: PLR0915
|
||||
custom_llm_provider = "anthropic_text"
|
||||
else:
|
||||
custom_llm_provider = "anthropic"
|
||||
## anthropic - pattern-based matching for future Claude models
|
||||
elif _matches_claude_model_pattern(model):
|
||||
custom_llm_provider = "anthropic"
|
||||
## cohere
|
||||
elif model in litellm.cohere_models or model in litellm.cohere_embedding_models:
|
||||
custom_llm_provider = "cohere"
|
||||
@@ -633,6 +659,11 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
|
||||
or get_secret_str("NVIDIA_RIVA_API_KEY")
|
||||
or get_secret_str("NVIDIA_NIM_API_KEY")
|
||||
)
|
||||
elif custom_llm_provider == "soniox":
|
||||
api_base = (
|
||||
api_base or get_secret_str("SONIOX_API_BASE") or "https://api.soniox.com"
|
||||
)
|
||||
dynamic_api_key = api_key or get_secret_str("SONIOX_API_KEY")
|
||||
elif custom_llm_provider == "cerebras":
|
||||
api_base = (
|
||||
api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1"
|
||||
@@ -931,6 +962,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
|
||||
) = litellm.LambdaAIChatConfig()._get_openai_compatible_provider_info(
|
||||
api_base, api_key
|
||||
)
|
||||
elif custom_llm_provider == "inception":
|
||||
(
|
||||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.InceptionChatConfig()._get_openai_compatible_provider_info(
|
||||
api_base, api_key
|
||||
)
|
||||
elif custom_llm_provider == "hyperbolic":
|
||||
(
|
||||
api_base,
|
||||
|
||||
@@ -22,9 +22,11 @@ def get_supported_openai_params( # noqa: PLR0915
|
||||
```
|
||||
|
||||
Args:
|
||||
base_model: For Azure, the true underlying model (e.g. ``"azure/gpt-5.2"``)
|
||||
when the deployment name differs. Used for model-type detection so that
|
||||
non-standard deployment names route to the correct config.
|
||||
base_model: An optional capability hint for deployments whose ``model``
|
||||
label isn't recognized on its own (e.g. an Azure deployment name, or a
|
||||
friendly Bedrock alias). It is additive: the result is the union of the
|
||||
params supported by ``model`` and by ``base_model``, so a hint can only
|
||||
add capabilities, never strip ones the real model already supports.
|
||||
|
||||
Returns:
|
||||
- List if custom_llm_provider is mapped
|
||||
@@ -52,7 +54,15 @@ def get_supported_openai_params( # noqa: PLR0915
|
||||
provider_config = None
|
||||
|
||||
if provider_config and request_type == "chat_completion":
|
||||
return provider_config.get_supported_openai_params(model=base_model or model)
|
||||
supported_params = provider_config.get_supported_openai_params(model=model)
|
||||
if base_model and base_model != model:
|
||||
base_model_params = provider_config.get_supported_openai_params(
|
||||
model=base_model
|
||||
)
|
||||
supported_params = list(
|
||||
dict.fromkeys([*supported_params, *base_model_params])
|
||||
)
|
||||
return supported_params
|
||||
|
||||
if custom_llm_provider == "bedrock":
|
||||
return litellm.AmazonConverseConfig().get_supported_openai_params(model=model)
|
||||
@@ -331,6 +341,11 @@ def get_supported_openai_params( # noqa: PLR0915
|
||||
return ElevenLabsAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
elif custom_llm_provider == "soniox":
|
||||
if request_type == "transcription":
|
||||
return litellm.SonioxAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
elif custom_llm_provider in litellm._custom_providers:
|
||||
if request_type == "chat_completion":
|
||||
provider_config = litellm.ProviderConfigManager.get_provider_chat_config(
|
||||
|
||||
@@ -3503,7 +3503,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
else:
|
||||
return None
|
||||
|
||||
def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse:
|
||||
def _handle_anthropic_messages_response_logging(
|
||||
self, result: Any
|
||||
) -> Union[ModelResponse, ResponsesAPIResponse]:
|
||||
"""
|
||||
Handles logging for Anthropic messages responses.
|
||||
|
||||
@@ -3522,6 +3524,15 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
return result
|
||||
elif isinstance(result, ModelResponse):
|
||||
return result
|
||||
elif isinstance(
|
||||
result,
|
||||
(ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent),
|
||||
):
|
||||
# anthropic_messages() can route to OpenAI Responses API; in that path
|
||||
# the assembled streaming result is one of these terminal events rather than
|
||||
# a ModelResponse. Return the inner response so downstream handlers
|
||||
# (_transform_usage_objects, normalize_logging_result) can process it.
|
||||
return result.response
|
||||
|
||||
httpx_response = self.model_call_details.get("httpx_response", None)
|
||||
if httpx_response and isinstance(httpx_response, httpx.Response):
|
||||
@@ -5300,8 +5311,12 @@ class StandardLoggingPayloadSetup:
|
||||
tb_lines[:MAXIMUM_TRACEBACK_LINES_TO_LOG]
|
||||
) # Limit to first 100 lines
|
||||
|
||||
# Get additional error details
|
||||
error_message = str(original_exception)
|
||||
explicit_message = getattr(original_exception, "message", None)
|
||||
error_message = (
|
||||
explicit_message
|
||||
if isinstance(explicit_message, str) and explicit_message
|
||||
else str(original_exception)
|
||||
)
|
||||
|
||||
return StandardLoggingPayloadErrorInformation(
|
||||
error_code=error_status,
|
||||
|
||||
@@ -34,6 +34,14 @@ _IMAGE_RESPONSE_CALL_TYPES = frozenset(
|
||||
_VALID_DATA_RESIDENCIES = frozenset(r.value for r in DataResidency)
|
||||
|
||||
|
||||
def _get_token_detail_value(details: object, key: str) -> Optional[int]:
|
||||
if isinstance(details, dict):
|
||||
value = details.get(key)
|
||||
else:
|
||||
value = getattr(details, key, None)
|
||||
return value if isinstance(value, int) else None
|
||||
|
||||
|
||||
def _is_above_128k(tokens: float) -> bool:
|
||||
if tokens > 128000:
|
||||
return True
|
||||
@@ -870,17 +878,47 @@ def calculate_image_response_cost_from_usage(
|
||||
cached_tokens=0,
|
||||
)
|
||||
|
||||
output_tokens_details = getattr(usage, "completion_tokens_details", None)
|
||||
if output_tokens_details is None:
|
||||
output_tokens_details = getattr(usage, "output_tokens_details", None)
|
||||
|
||||
if output_tokens_details is None:
|
||||
completion_tokens_details = CompletionTokensDetailsWrapper(
|
||||
text_tokens=0,
|
||||
image_tokens=completion_tokens,
|
||||
reasoning_tokens=0,
|
||||
audio_tokens=0,
|
||||
)
|
||||
else:
|
||||
text_tokens = _get_token_detail_value(output_tokens_details, "text_tokens") or 0
|
||||
image_tokens = (
|
||||
_get_token_detail_value(output_tokens_details, "image_tokens") or 0
|
||||
)
|
||||
audio_tokens = (
|
||||
_get_token_detail_value(output_tokens_details, "audio_tokens") or 0
|
||||
)
|
||||
reasoning_tokens = (
|
||||
_get_token_detail_value(output_tokens_details, "reasoning_tokens") or 0
|
||||
)
|
||||
known_output_tokens = (
|
||||
text_tokens + image_tokens + audio_tokens + reasoning_tokens
|
||||
)
|
||||
if completion_tokens > known_output_tokens:
|
||||
text_tokens += completion_tokens - known_output_tokens
|
||||
|
||||
completion_tokens_details = CompletionTokensDetailsWrapper(
|
||||
text_tokens=text_tokens,
|
||||
image_tokens=image_tokens,
|
||||
reasoning_tokens=reasoning_tokens,
|
||||
audio_tokens=audio_tokens,
|
||||
)
|
||||
|
||||
normalized_usage = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(
|
||||
text_tokens=0,
|
||||
image_tokens=completion_tokens,
|
||||
reasoning_tokens=0,
|
||||
audio_tokens=0,
|
||||
),
|
||||
completion_tokens_details=completion_tokens_details,
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
|
||||
@@ -144,6 +144,19 @@ async def convert_to_streaming_response_async(response_object: Optional[dict] =
|
||||
|
||||
choice_list: List[StreamingChoices] = []
|
||||
|
||||
if not response_object.get("choices"):
|
||||
from litellm.exceptions import APIError
|
||||
|
||||
raise APIError(
|
||||
status_code=500,
|
||||
message=(
|
||||
"LiteLLM: provider returned a response with no 'choices'. "
|
||||
f"Raw keys: {list(response_object.keys())}"
|
||||
),
|
||||
llm_provider="",
|
||||
model="",
|
||||
)
|
||||
|
||||
for idx, choice in enumerate(response_object["choices"]):
|
||||
if (
|
||||
choice["message"].get("tool_calls", None) is not None
|
||||
@@ -213,6 +226,20 @@ def convert_to_streaming_response(response_object: Optional[dict] = None):
|
||||
|
||||
model_response_object = ModelResponseStream()
|
||||
choice_list: List[StreamingChoices] = []
|
||||
|
||||
if not response_object.get("choices"):
|
||||
from litellm.exceptions import APIError
|
||||
|
||||
raise APIError(
|
||||
status_code=500,
|
||||
message=(
|
||||
"LiteLLM: provider returned a response with no 'choices'. "
|
||||
f"Raw keys: {list(response_object.keys())}"
|
||||
),
|
||||
llm_provider="",
|
||||
model="",
|
||||
)
|
||||
|
||||
for idx, choice in enumerate(response_object["choices"]):
|
||||
delta = Delta(**choice["message"])
|
||||
finish_reason = choice.get("finish_reason", None)
|
||||
@@ -536,9 +563,20 @@ def convert_to_model_response_object( # noqa: PLR0915
|
||||
return convert_to_streaming_response(response_object=response_object)
|
||||
choice_list: List[Choices] = []
|
||||
|
||||
assert response_object["choices"] is not None and isinstance(
|
||||
if not response_object.get("choices") or not isinstance(
|
||||
response_object["choices"], Iterable
|
||||
)
|
||||
):
|
||||
from litellm.exceptions import APIError
|
||||
|
||||
raise APIError(
|
||||
status_code=500,
|
||||
message=(
|
||||
"LiteLLM: provider returned a response with no 'choices'. "
|
||||
f"Raw keys: {list(response_object.keys())}"
|
||||
),
|
||||
llm_provider="",
|
||||
model="",
|
||||
)
|
||||
|
||||
for idx, choice in enumerate(response_object["choices"]):
|
||||
## HANDLE JSON MODE - anthropic returns single function call]
|
||||
@@ -816,7 +854,12 @@ def convert_to_model_response_object( # noqa: PLR0915
|
||||
model_response_object.results = response_object["results"]
|
||||
|
||||
return model_response_object
|
||||
except Exception:
|
||||
except Exception as e:
|
||||
from litellm.exceptions import APIError
|
||||
|
||||
if isinstance(e, APIError):
|
||||
raise
|
||||
|
||||
received_args = dict(
|
||||
response_object=response_object,
|
||||
model_response_object=model_response_object,
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
from typing import Coroutine, Optional
|
||||
import atexit
|
||||
from typing_extensions import TypedDict
|
||||
@@ -494,31 +495,43 @@ class LoggingWorker:
|
||||
processed = 0
|
||||
start_time = loop.time()
|
||||
|
||||
while not self._queue.empty() and processed < MAX_ITERATIONS_TO_CLEAR_QUEUE:
|
||||
if loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE:
|
||||
self._safe_log(
|
||||
"warning",
|
||||
f"[LoggingWorker] atexit: Reached time limit ({MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush",
|
||||
)
|
||||
break
|
||||
# logging.raiseExceptions is a process-wide global; scope the
|
||||
# suppression to just the drain loop, where shutdown callbacks may
|
||||
# log to already-closed handler streams, so other threads keep their
|
||||
# logging error reporting for as little of the window as possible.
|
||||
previous_raise_exceptions = logging.raiseExceptions
|
||||
logging.raiseExceptions = False
|
||||
try:
|
||||
while (
|
||||
not self._queue.empty()
|
||||
and processed < MAX_ITERATIONS_TO_CLEAR_QUEUE
|
||||
):
|
||||
if loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE:
|
||||
self._safe_log(
|
||||
"warning",
|
||||
f"[LoggingWorker] atexit: Reached time limit ({MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush",
|
||||
)
|
||||
break
|
||||
|
||||
try:
|
||||
task = self._queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
try:
|
||||
task = self._queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
# Run the coroutine synchronously in new loop
|
||||
# Note: We run the coroutine directly, not via create_task,
|
||||
# since we're in a new event loop context
|
||||
try:
|
||||
loop.run_until_complete(task["coroutine"])
|
||||
processed += 1
|
||||
except Exception:
|
||||
# Silent failure to not break user's program
|
||||
pass
|
||||
finally:
|
||||
# Clear reference to prevent memory leaks
|
||||
task = None
|
||||
# Run the coroutine synchronously in new loop
|
||||
# Note: We run the coroutine directly, not via create_task,
|
||||
# since we're in a new event loop context
|
||||
try:
|
||||
loop.run_until_complete(task["coroutine"])
|
||||
processed += 1
|
||||
except Exception:
|
||||
# Silent failure to not break user's program
|
||||
pass
|
||||
finally:
|
||||
# Clear reference to prevent memory leaks
|
||||
task = None
|
||||
finally:
|
||||
logging.raiseExceptions = previous_raise_exceptions
|
||||
|
||||
self._safe_log(
|
||||
"info",
|
||||
|
||||
@@ -850,7 +850,47 @@ def extract_file_data(file_data: FileTypes) -> ExtractedFileData:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def unpack_defs(schema: dict, defs: dict) -> None:
|
||||
def _estimate_json_bytes(obj: Any) -> int:
|
||||
"""Estimate the JSON-serialised byte size of ``obj`` without materialising
|
||||
JSON. Walks iteratively (no recursion stack risk).
|
||||
|
||||
String length is read via ``len()`` (O(1) on Python ``str``) so a target
|
||||
containing a 100MB description costs ~one walk step, not a 100MB
|
||||
serialisation. Escape sequences are not counted exactly, so this is an
|
||||
approximation -- but always within a small constant factor of the real
|
||||
serialised size, which is what a schema-bomb budget needs.
|
||||
"""
|
||||
total = 0
|
||||
stack: list = [obj]
|
||||
while stack:
|
||||
x = stack.pop()
|
||||
if isinstance(x, dict):
|
||||
total += 2 # `{}`
|
||||
for k, v in x.items():
|
||||
total += len(str(k)) + 4 # `"k":,`
|
||||
stack.append(v)
|
||||
elif isinstance(x, list):
|
||||
total += 2 # `[]`
|
||||
total += max(0, len(x) - 1) # commas between items
|
||||
stack.extend(x)
|
||||
elif isinstance(x, str):
|
||||
total += len(x) + 2
|
||||
elif isinstance(x, bool): # bool subclasses int -- check first
|
||||
total += 4 if x else 5
|
||||
elif x is None:
|
||||
total += 4
|
||||
elif isinstance(x, (int, float)):
|
||||
total += 24 # generous upper bound for stringified numbers
|
||||
else:
|
||||
total += 24
|
||||
return total
|
||||
|
||||
|
||||
def unpack_defs(
|
||||
schema: dict,
|
||||
defs: dict,
|
||||
max_inlined_bytes: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Expand *all* ``$ref`` entries pointing into ``$defs`` / ``definitions``.
|
||||
|
||||
This utility walks the entire schema tree (dicts and lists) so it naturally
|
||||
@@ -860,6 +900,15 @@ def unpack_defs(schema: dict, defs: dict) -> None:
|
||||
It mutates *schema* in-place and does **not** return anything. The helper
|
||||
keeps memory overhead low by resolving nodes as it encounters them rather
|
||||
than materialising a fully dereferenced copy first.
|
||||
|
||||
``max_inlined_bytes`` caps the cumulative JSON-byte size of every target
|
||||
that has been inlined and is checked *before* each ``copy.deepcopy``, so
|
||||
an oversized expansion is rejected without first materialising it. A byte
|
||||
bound is the universal measure of expansion -- it simultaneously caps
|
||||
ref-count fan-out, node-count amplification, and scalar-byte amplification
|
||||
(a target containing a large string, ``const``, or ``enum`` entry).
|
||||
Defaults to ``None`` (unbounded) so existing callers are unaffected;
|
||||
raises ``ValueError`` on overflow.
|
||||
"""
|
||||
|
||||
import copy
|
||||
@@ -879,6 +928,7 @@ def unpack_defs(schema: dict, defs: dict) -> None:
|
||||
queue: deque[
|
||||
tuple[Any, Union[dict, list, None], Union[str, int, None], dict, set]
|
||||
] = deque([(schema, None, None, root_defs, set())])
|
||||
inlined_bytes = 0
|
||||
|
||||
while queue:
|
||||
node, parent, key, active_defs, ref_chain = queue.popleft()
|
||||
@@ -899,6 +949,16 @@ def unpack_defs(schema: dict, defs: dict) -> None:
|
||||
if target_schema is None:
|
||||
continue
|
||||
|
||||
if max_inlined_bytes is not None:
|
||||
inlined_bytes += _estimate_json_bytes(target_schema)
|
||||
if inlined_bytes > max_inlined_bytes:
|
||||
raise ValueError(
|
||||
f"unpack_defs: inlined schema exceeded the "
|
||||
f"{max_inlined_bytes:,}-byte budget. Refusing to "
|
||||
f"deep-copy further to prevent schema-bomb "
|
||||
f"resource exhaustion."
|
||||
)
|
||||
|
||||
# Merge defs from the target to capture nested definitions
|
||||
child_defs = {
|
||||
**active_defs,
|
||||
@@ -946,6 +1006,61 @@ def unpack_defs(schema: dict, defs: dict) -> None:
|
||||
queue.append((item, node, idx, active_defs, ref_chain))
|
||||
|
||||
|
||||
def _has_legacy_defs(schema: object) -> bool:
|
||||
if not isinstance(schema, dict):
|
||||
return False
|
||||
components = schema.get("components")
|
||||
return "definitions" in schema or (
|
||||
isinstance(components, dict) and isinstance(components.get("schemas"), dict)
|
||||
)
|
||||
|
||||
|
||||
# Schema-bomb budget for ``unpack_legacy_defs``: cap the cumulative JSON-byte
|
||||
# size of every inlined target. A byte cap is the universal measure of
|
||||
# expansion -- it simultaneously bounds ref-count fan-out, node-count
|
||||
# amplification, and scalar-byte amplification (large ``description`` /
|
||||
# ``const`` / ``enum`` values). Real-world MCP / OpenAPI-derived tool schemas
|
||||
# inline well under 1MB; 10MB sits two orders of magnitude above that, well
|
||||
# below memory-pressure territory, and rejects request-supplied bombs before
|
||||
# the proxy materialises them.
|
||||
_LEGACY_DEFS_MAX_INLINED_BYTES = 10_000_000
|
||||
|
||||
|
||||
def unpack_legacy_defs(
|
||||
schema: dict,
|
||||
*,
|
||||
copy: bool = False,
|
||||
max_inlined_bytes: int = _LEGACY_DEFS_MAX_INLINED_BYTES,
|
||||
) -> dict:
|
||||
"""Inline ``$ref``s backed by draft-04 ``definitions`` / OpenAPI
|
||||
``components.schemas``. ``$defs`` is left untouched.
|
||||
|
||||
Anthropic and Fireworks tool-schema resolvers only recognise ``$defs``;
|
||||
legacy / OpenAPI def blocks are otherwise silently dropped and leave
|
||||
dangling pointers. See https://github.com/BerriAI/litellm/issues/26692.
|
||||
|
||||
Mutates ``schema`` in place and returns it. Pass ``copy=True`` to deep-copy
|
||||
first (only when there is actually work to do). ``max_inlined_bytes``
|
||||
bounds the cumulative JSON-byte size of inlined targets so request-supplied
|
||||
schemas cannot expand into a schema-bomb before reaching the upstream
|
||||
provider -- raises ``ValueError`` on overflow.
|
||||
"""
|
||||
if not _has_legacy_defs(schema):
|
||||
return schema
|
||||
if copy:
|
||||
import copy as _copy
|
||||
|
||||
schema = _copy.deepcopy(schema)
|
||||
# On key collision, ``definitions`` wins over ``components.schemas`` --
|
||||
# ``unpack_defs`` keys refs by last path segment so a single name can only
|
||||
# resolve to one body, and ``definitions`` is the JSON-Schema-native
|
||||
# namespace.
|
||||
defs = schema.pop("components", {}).get("schemas") or {}
|
||||
defs.update(schema.pop("definitions", None) or {})
|
||||
unpack_defs(schema, defs, max_inlined_bytes=max_inlined_bytes)
|
||||
return schema
|
||||
|
||||
|
||||
def _get_image_mime_type_from_url(url: str) -> Optional[str]:
|
||||
"""
|
||||
Get mime type for common image URLs
|
||||
|
||||
@@ -1670,15 +1670,15 @@ def convert_to_gemini_tool_call_result( # noqa: PLR0915
|
||||
if gemini_call_id:
|
||||
_function_response["id"] = gemini_call_id
|
||||
|
||||
# Create part with function_response, and optionally inline_data for images (Computer Use)
|
||||
_part: VertexPartType = {"function_response": _function_response}
|
||||
|
||||
# For Computer Use, if we have images/files, we need separate parts:
|
||||
# - One part with function_response
|
||||
# - One part per inline_data item
|
||||
# Gemini's PartType is a oneof, so we can't have both in the same part
|
||||
# For multimodal function responses, Gemini expects media parts nested
|
||||
# inside functionResponse.parts instead of sibling content parts.
|
||||
if inline_data_list:
|
||||
return [_part] + [{"inline_data": d} for d in inline_data_list]
|
||||
_function_response["parts"] = [
|
||||
{"inline_data": inline_data} for inline_data in inline_data_list
|
||||
]
|
||||
return [_part]
|
||||
|
||||
return _part
|
||||
|
||||
|
||||
@@ -92,8 +92,27 @@ class RealTimeStreaming:
|
||||
# Track whether we have already sent the guardrail turn-detection update
|
||||
# that disables provider auto-response for transcription guardrails.
|
||||
self._guardrail_turn_detection_update_sent: bool = False
|
||||
# Deferred Gemini Live setup: Pipecat may stream audio before session.update.
|
||||
# Buffer client audio until the backend acknowledges setup (setupComplete).
|
||||
self._backend_setup_complete: bool = (
|
||||
provider_config is None or provider_config.requires_session_configuration()
|
||||
)
|
||||
self._flushing_pending_messages_until_setup: bool = False
|
||||
self._pending_messages_until_setup: List[str] = []
|
||||
self._pending_messages_byte_total: int = 0
|
||||
|
||||
# Per-connection caps for pre-setup audio frames (message count + total bytes).
|
||||
_MAX_BUFFERED_MESSAGES: int = 200
|
||||
_MAX_BUFFERED_BYTES: int = 10 * 1024 * 1024 # 10 MB
|
||||
|
||||
_SESSION_EVENT_TYPES = frozenset(["session.created", "session.updated"])
|
||||
_CLIENT_AUDIO_BUFFER_TYPES = frozenset(
|
||||
[
|
||||
"input_audio_buffer.append",
|
||||
"input_audio_buffer.commit",
|
||||
"input_audio_buffer.clear",
|
||||
]
|
||||
)
|
||||
_AUDIO_FORMAT_MAP: Dict[str, Dict[str, Any]] = {
|
||||
"pcm16": {"type": "audio/pcm", "rate": 24000},
|
||||
"g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
|
||||
@@ -285,6 +304,86 @@ class RealTimeStreaming:
|
||||
await self.backend_ws.send(message) # type: ignore[union-attr, attr-defined]
|
||||
return True
|
||||
|
||||
def _uses_deferred_backend_setup(self) -> bool:
|
||||
"""True when setup is deferred until the client's first session.update."""
|
||||
if self.provider_config is None:
|
||||
return False
|
||||
return not self.provider_config.requires_session_configuration()
|
||||
|
||||
def _should_buffer_client_message_until_setup(self, message: str) -> bool:
|
||||
if not self._uses_deferred_backend_setup():
|
||||
return False
|
||||
if (
|
||||
self._backend_setup_complete
|
||||
and not self._flushing_pending_messages_until_setup
|
||||
):
|
||||
return False
|
||||
try:
|
||||
msg_obj = json.loads(message)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return False
|
||||
return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
|
||||
|
||||
def _buffer_pending_message_until_setup(self, message: str) -> None:
|
||||
msg_bytes = len(message.encode("utf-8"))
|
||||
if (
|
||||
len(self._pending_messages_until_setup)
|
||||
< RealTimeStreaming._MAX_BUFFERED_MESSAGES
|
||||
and self._pending_messages_byte_total + msg_bytes
|
||||
<= RealTimeStreaming._MAX_BUFFERED_BYTES
|
||||
):
|
||||
self._pending_messages_until_setup.append(message)
|
||||
self._pending_messages_byte_total += msg_bytes
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"Pre-setup buffer full (%d messages / %d bytes); dropping frame",
|
||||
len(self._pending_messages_until_setup),
|
||||
self._pending_messages_byte_total,
|
||||
)
|
||||
|
||||
async def _flush_pending_messages_until_setup(self) -> bool:
|
||||
pending = self._pending_messages_until_setup
|
||||
self._pending_messages_until_setup = []
|
||||
self._pending_messages_byte_total = 0
|
||||
for idx, message in enumerate(pending):
|
||||
try:
|
||||
await self._send_to_backend(message)
|
||||
except Exception as e:
|
||||
unsent = pending[idx:]
|
||||
self._pending_messages_until_setup = (
|
||||
unsent + self._pending_messages_until_setup
|
||||
)
|
||||
self._pending_messages_byte_total = sum(
|
||||
len(msg.encode("utf-8"))
|
||||
for msg in self._pending_messages_until_setup
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Failed to flush buffered client message after setup: %s "
|
||||
"(%d buffered message(s) retained)",
|
||||
e,
|
||||
len(unsent),
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _send_event_to_client(self, event: Any, event_str: str) -> bool:
|
||||
if self._client_wants_beta and isinstance(event, dict):
|
||||
try:
|
||||
translated = self._translate_event_to_beta(event)
|
||||
if translated is None:
|
||||
return False
|
||||
await self.websocket.send_text(json.dumps(translated))
|
||||
return True
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"Failed to translate %s to beta protocol, forwarding "
|
||||
"untranslated event to client: %s",
|
||||
event.get("type"),
|
||||
e,
|
||||
)
|
||||
await self.websocket.send_text(event_str)
|
||||
return True
|
||||
|
||||
def _cache_session_configuration_request(self, transformed_message: str) -> None:
|
||||
"""Store setup payload once sent to backend.
|
||||
|
||||
@@ -547,6 +646,19 @@ class RealTimeStreaming:
|
||||
isinstance(event, dict) and event.get("type") == "session.created"
|
||||
)
|
||||
if is_session_created_event:
|
||||
if (
|
||||
self._uses_deferred_backend_setup()
|
||||
and not self._backend_setup_complete
|
||||
):
|
||||
self._backend_setup_complete = True
|
||||
self._flushing_pending_messages_until_setup = True
|
||||
try:
|
||||
while self._pending_messages_until_setup:
|
||||
flushed = await self._flush_pending_messages_until_setup()
|
||||
if not flushed:
|
||||
break
|
||||
finally:
|
||||
self._flushing_pending_messages_until_setup = False
|
||||
if self._session_created_sent_to_client:
|
||||
# A synthetic session.created (with placeholder defaults) was
|
||||
# already forwarded to the client when we connected. The
|
||||
@@ -569,7 +681,7 @@ class RealTimeStreaming:
|
||||
## update if a prior attempt was dropped by the provider transform.
|
||||
if is_session_created_event and self._has_audio_transcription_guardrails():
|
||||
self.store_message(event_str)
|
||||
await self.websocket.send_text(event_str)
|
||||
await self._send_event_to_client(event, event_str)
|
||||
await self._maybe_send_guardrail_turn_detection_update()
|
||||
continue
|
||||
## GUARDRAIL: run on transcription events in provider_config path too
|
||||
@@ -581,7 +693,7 @@ class RealTimeStreaming:
|
||||
transcript = event.get("transcript", "")
|
||||
self._collect_user_input_from_backend_event(cast(dict, event))
|
||||
self.store_message(event_str)
|
||||
await self.websocket.send_text(event_str)
|
||||
await self._send_event_to_client(event, event_str)
|
||||
blocked = await self.run_realtime_guardrails(
|
||||
cast(str, transcript),
|
||||
item_id=cast(Optional[str], event.get("item_id")),
|
||||
@@ -591,7 +703,7 @@ class RealTimeStreaming:
|
||||
continue
|
||||
## LOGGING
|
||||
self.store_message(event_str)
|
||||
await self.websocket.send_text(event_str)
|
||||
await self._send_event_to_client(event, event_str)
|
||||
|
||||
async def _handle_raw_backend_message(self, raw_response) -> bool:
|
||||
"""Process a backend message without provider_config (raw path).
|
||||
@@ -880,6 +992,7 @@ class RealTimeStreaming:
|
||||
|
||||
## GUARDRAIL: intercept conversation.item.create for text-based injection.
|
||||
guardrail_turn_detection_injected = False
|
||||
msg_type: Optional[str] = None
|
||||
try:
|
||||
msg_obj = json.loads(message)
|
||||
msg_type = msg_obj.get("type")
|
||||
@@ -1081,6 +1194,29 @@ class RealTimeStreaming:
|
||||
# actually forward to the backend.
|
||||
self.store_input(message=message)
|
||||
|
||||
if self._should_buffer_client_message_until_setup(message):
|
||||
self._buffer_pending_message_until_setup(message)
|
||||
continue
|
||||
|
||||
if self._pending_messages_until_setup:
|
||||
should_send_setup_before_buffered_messages = (
|
||||
not self._backend_setup_complete
|
||||
and not self._flushing_pending_messages_until_setup
|
||||
and msg_type == "session.update"
|
||||
)
|
||||
if not should_send_setup_before_buffered_messages:
|
||||
self._buffer_pending_message_until_setup(message)
|
||||
if (
|
||||
self._backend_setup_complete
|
||||
and not self._flushing_pending_messages_until_setup
|
||||
):
|
||||
await self._flush_pending_messages_until_setup()
|
||||
continue
|
||||
|
||||
if self._flushing_pending_messages_until_setup:
|
||||
self._buffer_pending_message_until_setup(message)
|
||||
continue
|
||||
|
||||
## FORWARD TO BACKEND
|
||||
# Only mark the guardrail turn_detection update as sent after the
|
||||
# backend actually accepted the message. Setting the flag earlier
|
||||
|
||||
@@ -6,10 +6,16 @@ from pydantic import BaseModel
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
|
||||
|
||||
def strip_null_bytes(value: str) -> str:
|
||||
"""Strip NUL bytes, which PostgreSQL text/jsonb columns reject (error 22P05)."""
|
||||
return value.replace("\x00", "")
|
||||
|
||||
|
||||
def safe_dumps(data: Any, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> str:
|
||||
"""
|
||||
Recursively serialize data while detecting circular references.
|
||||
If a circular reference is detected then a marker string is returned.
|
||||
NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors.
|
||||
"""
|
||||
|
||||
def _serialize(obj: Any, seen: set, depth: int) -> Any:
|
||||
@@ -17,7 +23,9 @@ def safe_dumps(data: Any, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> str:
|
||||
if depth > max_depth:
|
||||
return "MaxDepthExceeded"
|
||||
# Base-case: if it is a primitive, simply return it.
|
||||
if isinstance(obj, (str, int, float, bool, type(None))):
|
||||
if isinstance(obj, str):
|
||||
return strip_null_bytes(obj)
|
||||
if isinstance(obj, (int, float, bool, type(None))):
|
||||
return obj
|
||||
# Check for circular reference.
|
||||
if id(obj) in seen:
|
||||
@@ -28,7 +36,7 @@ def safe_dumps(data: Any, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> str:
|
||||
result = {}
|
||||
for k, v in obj.items():
|
||||
if isinstance(k, (str)):
|
||||
result[k] = _serialize(v, seen, depth + 1)
|
||||
result[strip_null_bytes(k)] = _serialize(v, seen, depth + 1)
|
||||
seen.remove(id(obj))
|
||||
return result
|
||||
elif isinstance(obj, list):
|
||||
@@ -51,7 +59,7 @@ def safe_dumps(data: Any, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH) -> str:
|
||||
else:
|
||||
# Fall back to string conversion for non-serializable objects.
|
||||
try:
|
||||
return str(obj)
|
||||
return strip_null_bytes(str(obj))
|
||||
except Exception:
|
||||
return "Unserializable Object"
|
||||
|
||||
|
||||
@@ -50,13 +50,15 @@ def _build_secret_patterns() -> "re.Pattern[str]":
|
||||
r"(?<=://)[^\s'\"]*:[^\s'\"@]+(?=@)",
|
||||
# Databricks personal access tokens
|
||||
r"dapi[0-9a-f]{32}",
|
||||
# Module-level provider keys logged as litellm.<provider>_key=<value>
|
||||
r"litellm\.[A-Za-z0-9_]*_key['\"]?\s*[:=]\s*['\"]?[^\s,'\"})\]{}>]+",
|
||||
# ── Key-name-based redaction ──
|
||||
# Catches secrets inside dicts/config dumps by matching on the KEY name
|
||||
# regardless of what the value looks like.
|
||||
# e.g. 'master_key': 'any-value-here', "database_url": "postgres://..."
|
||||
# private_key with PEM-aware value capture
|
||||
r"""private_key['\"]?\s*[:=]\s*['\"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'\"})\]{}>]+)""",
|
||||
r"(?:master_key|database_url|db_url|connection_string|"
|
||||
r"(?:master_key|xai_key|database_url|db_url|connection_string|"
|
||||
r"signing_key|encryption_key|"
|
||||
r"auth_token|access_token|refresh_token|"
|
||||
r"slack_webhook_url|webhook_url|"
|
||||
|
||||
@@ -1149,6 +1149,32 @@ class CustomStreamWrapper:
|
||||
completion_obj: Dict[str, Any] = {"content": ""}
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, ModelResponseStream)
|
||||
and self.custom_llm_provider is not None
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
_has_content = bool(
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta is not None
|
||||
and (
|
||||
chunk.choices[0].delta.content
|
||||
or chunk.choices[0].delta.tool_calls
|
||||
)
|
||||
)
|
||||
if self.received_finish_reason is not None:
|
||||
if not _has_content:
|
||||
raise StopIteration
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
self.received_finish_reason = chunk.choices[0].finish_reason
|
||||
if not _has_content:
|
||||
return None
|
||||
# Strip finish_reason from the content chunk so it appears
|
||||
# only on the trailing empty-delta chunk (OpenAI spec).
|
||||
# finish_reason_handler() will emit the proper terminal chunk.
|
||||
chunk.choices[0].finish_reason = None # type: ignore[assignment]
|
||||
return chunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, dict)
|
||||
and generic_chunk_has_all_required_fields(
|
||||
|
||||
@@ -29,6 +29,7 @@ from litellm.constants import (
|
||||
RESPONSE_FORMAT_TOOL_NAME,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_legacy_defs
|
||||
from litellm.llms.base_llm.base_utils import type_to_response_format_param
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.types.llms.anthropic import (
|
||||
@@ -80,7 +81,6 @@ from litellm.types.utils import (
|
||||
from litellm.utils import (
|
||||
ModelResponse,
|
||||
Usage,
|
||||
_supports_factory,
|
||||
add_dummy_tool,
|
||||
any_assistant_message_has_thinking_blocks,
|
||||
get_max_tokens,
|
||||
@@ -336,50 +336,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
||||
v in model_lower for v in ("opus-4-7", "opus_4_7", "opus-4.7", "opus_4.7")
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _supports_model_capability(model: str, key: str) -> bool:
|
||||
"""Check a boolean capability ``key`` in the model map.
|
||||
|
||||
Strips bedrock/vertex prefixes so a provider-routed Claude still
|
||||
resolves to the Anthropic model-map entry.
|
||||
"""
|
||||
try:
|
||||
if _supports_factory(
|
||||
model=model,
|
||||
custom_llm_provider="anthropic",
|
||||
key=key,
|
||||
):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
candidates = [model]
|
||||
for prefix in (
|
||||
"bedrock/converse/",
|
||||
"bedrock/invoke/",
|
||||
"bedrock/",
|
||||
"vertex_ai/",
|
||||
):
|
||||
if model.startswith(prefix):
|
||||
candidates.append(model[len(prefix) :])
|
||||
try:
|
||||
from litellm.llms.bedrock.common_utils import BedrockModelInfo
|
||||
|
||||
base = BedrockModelInfo.get_base_model(model)
|
||||
if base:
|
||||
candidates.append(base)
|
||||
candidates.append(f"bedrock/{base}")
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
for cand in candidates:
|
||||
if cand in litellm.model_cost and (
|
||||
litellm.model_cost[cand].get(key) is True
|
||||
):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _supports_effort_level(model: str, level: str) -> bool:
|
||||
"""Check ``supports_{level}_reasoning_effort`` in the model map."""
|
||||
@@ -680,6 +636,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
||||
if "properties" not in _input_schema:
|
||||
_input_schema["properties"] = {}
|
||||
|
||||
# Inline legacy / OpenAPI $refs before the allow-list filter strips
|
||||
# their backing def blocks (https://github.com/BerriAI/litellm/issues/26692).
|
||||
_input_schema = unpack_legacy_defs(_input_schema, copy=True)
|
||||
|
||||
_allowed_properties = set(AnthropicInputSchema.__annotations__.keys())
|
||||
input_schema_filtered = {
|
||||
k: v for k, v in _input_schema.items() if k in _allowed_properties
|
||||
@@ -913,7 +873,39 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
||||
anthropic_tools = []
|
||||
mcp_servers = []
|
||||
for tool in tools:
|
||||
if "input_schema" in tool: # assume in anthropic format
|
||||
if tool.get("type") == "namespace":
|
||||
# Namespace is a grouping container (e.g. codex's multi_agent_v1).
|
||||
# Extract its nested tools and map them individually.
|
||||
for nested in tool.get("tools") or []:
|
||||
if "input_schema" in nested:
|
||||
# Already in Anthropic format.
|
||||
anthropic_tools.append(nested)
|
||||
elif "function" not in nested and "name" in nested:
|
||||
# Flat format: {type, name, description, parameters, ...}.
|
||||
# Normalize to OpenAI-wrapped format before mapping.
|
||||
wrapped = cast(
|
||||
ChatCompletionToolParam,
|
||||
{
|
||||
"type": nested.get("type", "function"),
|
||||
"function": {
|
||||
k: v for k, v in nested.items() if k != "type"
|
||||
},
|
||||
},
|
||||
)
|
||||
nested_tool, nested_mcp = self._map_tool_helper(wrapped)
|
||||
if nested_tool is not None:
|
||||
anthropic_tools.append(nested_tool)
|
||||
if nested_mcp is not None:
|
||||
mcp_servers.append(nested_mcp)
|
||||
elif "function" in nested:
|
||||
nested_tool, nested_mcp = self._map_tool_helper(
|
||||
cast(ChatCompletionToolParam, nested)
|
||||
)
|
||||
if nested_tool is not None:
|
||||
anthropic_tools.append(nested_tool)
|
||||
if nested_mcp is not None:
|
||||
mcp_servers.append(nested_mcp)
|
||||
elif "input_schema" in tool: # assume in anthropic format
|
||||
anthropic_tools.append(tool)
|
||||
else: # assume openai tool call
|
||||
new_tool, mcp_server_tool = self._map_tool_helper(tool)
|
||||
@@ -1973,6 +1965,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
||||
|
||||
# Remove internal LiteLLM parameters that should not be sent to Anthropic API
|
||||
optional_params.pop("is_vertex_request", None)
|
||||
optional_params.pop("client_metadata", None)
|
||||
|
||||
data = {
|
||||
"model": model,
|
||||
|
||||
@@ -272,19 +272,63 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_adaptive_thinking_model(model: str) -> bool:
|
||||
"""Claude 4.6+ models use adaptive thinking with ``output_config.effort``."""
|
||||
def _supports_model_capability(model: str, key: str) -> bool:
|
||||
"""Check a boolean capability ``key`` in the model map.
|
||||
|
||||
Strips bedrock/vertex prefixes so a provider-routed Claude still
|
||||
resolves to the Anthropic model-map entry.
|
||||
"""
|
||||
from litellm.utils import _supports_factory
|
||||
|
||||
try:
|
||||
if _supports_factory(
|
||||
model=model,
|
||||
custom_llm_provider=None,
|
||||
key="supports_adaptive_thinking",
|
||||
custom_llm_provider="anthropic",
|
||||
key=key,
|
||||
):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
candidates = [model]
|
||||
for prefix in (
|
||||
"bedrock/converse/",
|
||||
"bedrock/invoke/",
|
||||
"bedrock/",
|
||||
"vertex_ai/",
|
||||
):
|
||||
if model.startswith(prefix):
|
||||
candidates.append(model[len(prefix) :])
|
||||
try:
|
||||
from litellm.llms.bedrock.common_utils import BedrockModelInfo
|
||||
|
||||
base = BedrockModelInfo.get_base_model(model)
|
||||
if base:
|
||||
candidates.append(base)
|
||||
candidates.append(f"bedrock/{base}")
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
for cand in candidates:
|
||||
if cand in litellm.model_cost and (
|
||||
litellm.model_cost[cand].get(key) is True
|
||||
):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _is_adaptive_thinking_model(model: str) -> bool:
|
||||
"""Claude 4.6+ models use adaptive thinking with ``output_config.effort``.
|
||||
|
||||
Driven by the ``supports_adaptive_thinking`` flag in the model map; the
|
||||
4.6/4.7 name checks remain only as a fallback for provider-routed ids
|
||||
whose map entries predate the flag.
|
||||
"""
|
||||
if AnthropicModelInfo._supports_model_capability(
|
||||
model, "supports_adaptive_thinking"
|
||||
):
|
||||
return True
|
||||
return AnthropicModelInfo._is_claude_4_6_model(
|
||||
model
|
||||
) or AnthropicModelInfo._is_claude_4_7_model(model)
|
||||
|
||||
@@ -1510,6 +1510,17 @@ class LiteLLMAnthropicMessagesAdapter:
|
||||
return "thinking", ChatCompletionThinkingBlock(
|
||||
type="thinking", thinking=thinking, signature=signature
|
||||
)
|
||||
# OpenAI-compatible reasoning backends (e.g. vLLM/SGLang reasoning
|
||||
# parsers) populate ``reasoning_content`` without ``thinking_blocks``.
|
||||
# ``Delta`` deletes the ``thinking_blocks`` attribute when unset, so the
|
||||
# branch above is skipped entirely; open a ``thinking`` block here so the
|
||||
# matching ``thinking_delta`` stream is not emitted into a text block.
|
||||
elif isinstance(choice, StreamingChoices) and getattr(
|
||||
choice.delta, "reasoning_content", None
|
||||
):
|
||||
return "thinking", ChatCompletionThinkingBlock(
|
||||
type="thinking", thinking="", signature=""
|
||||
)
|
||||
|
||||
return "text", TextBlock(type="text", text="")
|
||||
|
||||
|
||||
@@ -229,9 +229,20 @@ async def anthropic_messages(
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Extract modified parameters
|
||||
# Extract modified parameters. Pop every named param of `anthropic_messages`
|
||||
# that we may forward explicitly downstream, so we (a) honor pre-request hook
|
||||
# overrides and (b) avoid duplicate-keyword conflicts when splatting `kwargs`
|
||||
# into call sites that already pass these as named arguments.
|
||||
tools = request_kwargs.pop("tools", tools)
|
||||
stream = request_kwargs.pop("stream", stream)
|
||||
metadata = request_kwargs.pop("metadata", metadata)
|
||||
stop_sequences = request_kwargs.pop("stop_sequences", stop_sequences)
|
||||
system = request_kwargs.pop("system", system)
|
||||
temperature = request_kwargs.pop("temperature", temperature)
|
||||
thinking = request_kwargs.pop("thinking", thinking)
|
||||
tool_choice = request_kwargs.pop("tool_choice", tool_choice)
|
||||
top_k = request_kwargs.pop("top_k", top_k)
|
||||
top_p = request_kwargs.pop("top_p", top_p)
|
||||
# Propagate the provider derived inside pre-request hooks, if not already set.
|
||||
# The litellm_params dict may have been overwritten by **kwargs in
|
||||
# _execute_pre_request_hooks, so fall back to get_llm_provider() if needed.
|
||||
@@ -265,8 +276,8 @@ async def anthropic_messages(
|
||||
return short_circuit_response
|
||||
|
||||
# Run registered MessagesInterceptors (e.g. advisor orchestration loop).
|
||||
# api_key and api_base are explicit params (not in **kwargs) so pass them
|
||||
# explicitly so interceptor sub-calls can route to the same backend.
|
||||
# Named params on `anthropic_messages` are bound to locals, not `**kwargs`,
|
||||
# so forward them explicitly — otherwise interceptor sub-calls drop them.
|
||||
for interceptor in get_messages_interceptors():
|
||||
if interceptor.can_handle(tools, custom_llm_provider):
|
||||
return await interceptor.handle(
|
||||
@@ -278,6 +289,14 @@ async def anthropic_messages(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
metadata=metadata,
|
||||
stop_sequences=stop_sequences,
|
||||
system=system,
|
||||
temperature=temperature,
|
||||
thinking=thinking,
|
||||
tool_choice=tool_choice,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""APISerpent integration for LiteLLM."""
|
||||
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
APISerpent Search API module.
|
||||
"""
|
||||
|
||||
from litellm.llms.apiserpent.search.defaults import APISerpentSearchParams
|
||||
from litellm.llms.apiserpent.search.transformation import APISerpentSearchConfig
|
||||
|
||||
__all__ = ["APISerpentSearchConfig", "APISerpentSearchParams"]
|
||||
@@ -0,0 +1,70 @@
|
||||
"""
|
||||
Default parameter values and shared constants for APISerpent search.
|
||||
|
||||
Single source of truth for the supported request parameters and their
|
||||
package-level defaults. See https://apiserpent.com/docs.
|
||||
"""
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Dict, Literal, Optional
|
||||
|
||||
SearchEngine = Literal["google", "bing", "yahoo", "ddg"]
|
||||
SafeSearch = Literal["off", "moderate", "strict"]
|
||||
Freshness = Literal["h", "1h", "d", "1d", "7d", "w", "m", "1m", "y", "1y"]
|
||||
ResponseFormat = Literal["full", "simple"]
|
||||
|
||||
NUM_MIN = 1
|
||||
NUM_MIN_DEEP = 10
|
||||
NUM_MAX = 100
|
||||
PAGES_MIN = 1
|
||||
PAGES_MAX = 10
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class APISerpentSearchParams:
|
||||
"""
|
||||
Supported APISerpent search parameters with package defaults.
|
||||
|
||||
Fields defaulting to ``None`` are only sent when the caller provides them;
|
||||
the rest are always sent so behavior is deterministic regardless of any
|
||||
server-side defaults.
|
||||
"""
|
||||
|
||||
engine: SearchEngine = "google"
|
||||
country: str = "us"
|
||||
num: int = 10
|
||||
format: ResponseFormat = "full"
|
||||
pages: Optional[int] = None
|
||||
freshness: Optional[Freshness] = None
|
||||
safe: Optional[SafeSearch] = None
|
||||
language: Optional[str] = None
|
||||
pixel_position: Optional[bool] = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# num's deep-search floor (NUM_MIN_DEEP) is endpoint-specific and enforced
|
||||
# in the transform layer; here we only bound the absolute range.
|
||||
if not NUM_MIN <= self.num <= NUM_MAX:
|
||||
raise ValueError(
|
||||
f"num must be between {NUM_MIN} and {NUM_MAX}, got {self.num}"
|
||||
)
|
||||
if self.pages is not None and not PAGES_MIN <= self.pages <= PAGES_MAX:
|
||||
raise ValueError(
|
||||
f"pages must be between {PAGES_MIN} and {PAGES_MAX}, got {self.pages}"
|
||||
)
|
||||
|
||||
def to_request_params(self) -> Dict:
|
||||
"""Return non-None fields as request params, booleans lowercased."""
|
||||
params: Dict = {}
|
||||
for key, value in asdict(self).items():
|
||||
if value is None:
|
||||
continue
|
||||
params[key] = str(value).lower() if isinstance(value, bool) else value
|
||||
return params
|
||||
|
||||
@classmethod
|
||||
def field_names(cls) -> set:
|
||||
return set(cls.__dataclass_fields__.keys())
|
||||
|
||||
|
||||
QUICK_SEARCH_PATH = "/api/search/quick"
|
||||
DEEP_SEARCH_PATH = "/api/search"
|
||||
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
Calls APISerpent's search endpoints to search Google, Bing, Yahoo, or DuckDuckGo.
|
||||
|
||||
Two endpoints under one provider, selected via the ``deep`` boolean param:
|
||||
- ``deep=False`` (default) -> quick search (/api/search/quick)
|
||||
- ``deep=True`` -> deep search (/api/search)
|
||||
|
||||
APISerpent API Reference: https://apiserpent.com/docs
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Literal, Optional, Union, cast
|
||||
from urllib.parse import urlencode
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.apiserpent.search.defaults import (
|
||||
DEEP_SEARCH_PATH,
|
||||
NUM_MAX,
|
||||
NUM_MIN,
|
||||
NUM_MIN_DEEP,
|
||||
QUICK_SEARCH_PATH,
|
||||
APISerpentSearchParams,
|
||||
)
|
||||
from litellm.llms.base_llm.search.transformation import (
|
||||
BaseSearchConfig,
|
||||
SearchResponse,
|
||||
SearchResult,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
DEEP_SEARCH_PARAM = "deep"
|
||||
APISERPENT_BASE = "https://apiserpent.com"
|
||||
APISERPENT_PARAMS_KEY = "_apiserpent_params"
|
||||
|
||||
|
||||
class APISerpentSearchConfig(BaseSearchConfig):
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "APISerpent"
|
||||
|
||||
def get_http_method(self) -> Literal["GET", "POST"]:
|
||||
return "GET"
|
||||
|
||||
@staticmethod
|
||||
def _is_deep_search(optional_params: dict) -> bool:
|
||||
return bool(optional_params.get(DEEP_SEARCH_PARAM))
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Dict:
|
||||
api_key = api_key or get_secret_str("APISERPENT_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"APISERPENT_API_KEY is not set. Set `APISERPENT_API_KEY` environment variable."
|
||||
)
|
||||
headers["X-API-Key"] = api_key
|
||||
headers["Content-Type"] = "application/json"
|
||||
return headers
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
optional_params: dict,
|
||||
data: Optional[Union[Dict, List[Dict]]] = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
"""
|
||||
Build the search URL. APISerpent uses GET, so the transformed request is
|
||||
serialized into the query string. The endpoint path (quick vs deep) is
|
||||
always applied; an ``api_base`` / ``APISERPENT_API_BASE`` override only
|
||||
changes the host. The ``endswith`` guard keeps this idempotent, since the
|
||||
handler re-invokes this method with the already-resolved URL as api_base.
|
||||
"""
|
||||
base = (
|
||||
api_base or get_secret_str("APISERPENT_API_BASE") or APISERPENT_BASE
|
||||
).rstrip("/")
|
||||
path = (
|
||||
DEEP_SEARCH_PATH
|
||||
if self._is_deep_search(optional_params)
|
||||
else QUICK_SEARCH_PATH
|
||||
)
|
||||
if not base.endswith(path):
|
||||
base = f"{base}{path}"
|
||||
|
||||
if data and isinstance(data, dict) and APISERPENT_PARAMS_KEY in data:
|
||||
query_string = urlencode(data[APISERPENT_PARAMS_KEY], doseq=True)
|
||||
return f"{base}?{query_string}"
|
||||
|
||||
return base
|
||||
|
||||
def transform_search_request(
|
||||
self,
|
||||
query: Union[str, List[str]],
|
||||
optional_params: dict,
|
||||
**kwargs,
|
||||
) -> Dict:
|
||||
"""
|
||||
Transform a unified search request into APISerpent query params.
|
||||
|
||||
Unified spec mappings:
|
||||
- query -> q
|
||||
- max_results -> num (clamped to the endpoint's valid range)
|
||||
- country -> country (lowercased)
|
||||
- search_domain_filter -> site: clauses appended to q
|
||||
|
||||
All other APISerpent params (engine, language, freshness, safe, pages,
|
||||
format, pixel_position) pass through, defaulting via APISerpentSearchParams.
|
||||
"""
|
||||
if isinstance(query, list):
|
||||
query = " ".join(query)
|
||||
|
||||
is_deep = self._is_deep_search(optional_params)
|
||||
|
||||
overrides: Dict = {}
|
||||
if "max_results" in optional_params:
|
||||
num_min = NUM_MIN_DEEP if is_deep else NUM_MIN
|
||||
overrides["num"] = max(
|
||||
num_min, min(optional_params["max_results"], NUM_MAX)
|
||||
)
|
||||
if "country" in optional_params:
|
||||
overrides["country"] = cast(str, optional_params["country"]).lower()
|
||||
|
||||
for param, value in optional_params.items():
|
||||
if param in APISerpentSearchParams.field_names() and param not in overrides:
|
||||
overrides[param] = value
|
||||
|
||||
params = {**APISerpentSearchParams(**overrides).to_request_params(), "q": query}
|
||||
|
||||
if "search_domain_filter" in optional_params:
|
||||
domains = optional_params["search_domain_filter"]
|
||||
if isinstance(domains, list) and len(domains) > 0:
|
||||
params["q"] = self._append_domain_filters(str(params["q"]), domains)
|
||||
|
||||
return {APISERPENT_PARAMS_KEY: params}
|
||||
|
||||
@staticmethod
|
||||
def _append_domain_filters(query: str, domains: List[str]) -> str:
|
||||
domain_clauses = " OR ".join(f"site:{domain}" for domain in domains)
|
||||
return f"({query}) ({domain_clauses})"
|
||||
|
||||
def transform_search_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: Optional[LiteLLMLoggingObj],
|
||||
**kwargs,
|
||||
) -> SearchResponse:
|
||||
"""
|
||||
Transform APISerpent response to the unified SearchResponse format.
|
||||
|
||||
Full format nests results under ``results.organic[]``; simple format
|
||||
returns a flat ``results[]`` array. Both expose title/url/snippet.
|
||||
"""
|
||||
response_json = raw_response.json()
|
||||
|
||||
raw_results = response_json.get("results") or {}
|
||||
organic = (
|
||||
raw_results.get("organic", [])
|
||||
if isinstance(raw_results, dict)
|
||||
else raw_results
|
||||
)
|
||||
|
||||
results: List[SearchResult] = []
|
||||
for result in organic:
|
||||
results.append(
|
||||
SearchResult(
|
||||
title=result.get("title", ""),
|
||||
url=result.get("url", ""),
|
||||
snippet=result.get("snippet", ""),
|
||||
date=result.get("date"),
|
||||
last_updated=None,
|
||||
)
|
||||
)
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
object="search",
|
||||
)
|
||||
@@ -97,8 +97,15 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
|
||||
)
|
||||
original_url = httpx.URL(api_base)
|
||||
|
||||
# Extract api_version or use default
|
||||
api_version = cast(Optional[str], litellm_params.get("api_version"))
|
||||
# Resolve api_version: litellm_params > litellm.api_version > AZURE_API_VERSION env > default.
|
||||
# Mirrors the fallback chain used by the Azure chat path in common_utils.py,
|
||||
# so callers that set a global / env api_version don't get an unversioned URL.
|
||||
api_version = (
|
||||
cast(Optional[str], litellm_params.get("api_version"))
|
||||
or litellm.api_version
|
||||
or get_secret_str("AZURE_API_VERSION")
|
||||
or litellm.AZURE_DEFAULT_API_VERSION
|
||||
)
|
||||
|
||||
# Create a new dictionary with existing params
|
||||
query_params = dict(original_url.params)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import enum
|
||||
import re
|
||||
from typing import Any, List, Optional, Tuple, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
@@ -275,21 +276,25 @@ class AzureAIStudioConfig(OpenAIConfig):
|
||||
should_drop_params = litellm_params.get("drop_params") or litellm.drop_params
|
||||
error_text = e.response.text
|
||||
|
||||
if should_drop_params and "Extra inputs are not permitted" in error_text:
|
||||
if "Extra inputs are not permitted" in error_text:
|
||||
if should_drop_params or self._error_has_tool_level_extra_fields(
|
||||
error_text
|
||||
):
|
||||
return True
|
||||
if "unknown field: parameter index is not a valid field" in error_text:
|
||||
return True
|
||||
elif (
|
||||
"unknown field: parameter index is not a valid field" in error_text
|
||||
): # remove index from tool calls
|
||||
return True
|
||||
elif (
|
||||
if (
|
||||
AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value
|
||||
in error_text
|
||||
): # remove extra-parameters from tool calls
|
||||
):
|
||||
return True
|
||||
return super().should_retry_llm_api_inside_llm_translation_on_http_error(
|
||||
e=e, litellm_params=litellm_params
|
||||
)
|
||||
|
||||
def _error_has_tool_level_extra_fields(self, error_text: str) -> bool:
|
||||
return bool(re.search(r"tools\[\d+\]\.", error_text))
|
||||
|
||||
@property
|
||||
def max_retry_on_unprocessable_entity_error(self) -> int:
|
||||
return 2
|
||||
@@ -297,9 +302,10 @@ class AzureAIStudioConfig(OpenAIConfig):
|
||||
def transform_request_on_unprocessable_entity_error(
|
||||
self, e: httpx.HTTPStatusError, request_data: dict
|
||||
) -> dict:
|
||||
error_text = e.response.text
|
||||
_messages = cast(Optional[List[AllMessageValues]], request_data.get("messages"))
|
||||
if (
|
||||
"unknown field: parameter index is not a valid field" in e.response.text
|
||||
"unknown field: parameter index is not a valid field" in error_text
|
||||
and _messages is not None
|
||||
):
|
||||
litellm.remove_index_from_tool_calls(
|
||||
@@ -307,14 +313,31 @@ class AzureAIStudioConfig(OpenAIConfig):
|
||||
)
|
||||
elif (
|
||||
AzureFoundryErrorStrings.SET_EXTRA_PARAMETERS_TO_PASS_THROUGH.value
|
||||
in e.response.text
|
||||
in error_text
|
||||
):
|
||||
request_data = self._drop_extra_params_from_request_data(
|
||||
request_data, e.response.text
|
||||
request_data, error_text
|
||||
)
|
||||
if (
|
||||
"Extra inputs are not permitted" in error_text
|
||||
and self._error_has_tool_level_extra_fields(error_text)
|
||||
):
|
||||
request_data = self._drop_tool_level_extra_fields(request_data, error_text)
|
||||
data = drop_params_from_unprocessable_entity_error(e=e, data=request_data)
|
||||
return data
|
||||
|
||||
def _drop_tool_level_extra_fields(
|
||||
self, request_data: dict, error_text: str
|
||||
) -> dict:
|
||||
fields_to_drop = set(re.findall(r"tools\[\d+\]\.([\w-]+)", error_text))
|
||||
tools = request_data.get("tools")
|
||||
if fields_to_drop and isinstance(tools, list):
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict):
|
||||
for field in fields_to_drop:
|
||||
tool.pop(field, None)
|
||||
return request_data
|
||||
|
||||
def _drop_extra_params_from_request_data(
|
||||
self, request_data: dict, error_text: str
|
||||
) -> dict:
|
||||
@@ -332,9 +355,6 @@ class AzureAIStudioConfig(OpenAIConfig):
|
||||
Error text looks like this"
|
||||
"Extra parameters ['stream_options', 'extra-parameters'] are not allowed when extra-parameters is not set or set to be 'error'.
|
||||
"""
|
||||
import re
|
||||
|
||||
# Extract parameters within square brackets
|
||||
match = re.search(r"\[(.*?)\]", error_text)
|
||||
if not match:
|
||||
return []
|
||||
|
||||
@@ -1534,10 +1534,9 @@ class BaseAWSLLM:
|
||||
)
|
||||
|
||||
sigv4 = SigV4Auth(credentials, service_name, aws_region_name)
|
||||
if headers is not None:
|
||||
headers = headers or {}
|
||||
if not any(header_name.lower() == "content-type" for header_name in headers):
|
||||
headers = {"Content-Type": "application/json", **headers}
|
||||
else:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
aws_signature_headers = self._filter_headers_for_aws_signature(headers)
|
||||
request = AWSRequest(
|
||||
|
||||
@@ -233,6 +233,259 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
||||
# example; add others here as they adopt the same schema.
|
||||
CONVERSE_INVOKE_PROVIDERS = ("nova",)
|
||||
|
||||
# OpenAI batch URL that signals an embedding request. Per OpenAI Batch API
|
||||
# spec, every JSONL record carries a `url` field; we use it as the
|
||||
# authoritative signal to route the line to the embedding code path
|
||||
# instead of inferring from the presence of `input` vs `messages`.
|
||||
OPENAI_EMBEDDINGS_URL = "/v1/embeddings"
|
||||
|
||||
@staticmethod
|
||||
def _is_embedding_record(openai_jsonl_record: Dict[str, Any]) -> bool:
|
||||
"""
|
||||
Decide whether an OpenAI batch JSONL line is an embedding request.
|
||||
|
||||
Precedence (strict - any explicit `url` short-circuits):
|
||||
1. `url == "/v1/embeddings"` -> embedding. Authoritative per the
|
||||
OpenAI Batch API spec.
|
||||
2. Any other non-empty `url` (e.g. `/v1/chat/completions`) -> NOT
|
||||
embedding. We trust the caller's explicit signal even if the
|
||||
body would otherwise suggest embedding; misrouting a chat
|
||||
record into the embedding transformer would corrupt the
|
||||
modelInput, while a chat-shaped body sent to the chat path
|
||||
either succeeds or fails cleanly inside that transformer.
|
||||
3. `url` missing/empty -> fall back to body shape. Requires
|
||||
`input` present AND `messages` absent so a malformed record
|
||||
carrying both keys routes to the chat path (safer default:
|
||||
Anthropic transforms ignore unknown top-level keys, whereas
|
||||
the embedding transformer would silently drop the messages).
|
||||
"""
|
||||
url = openai_jsonl_record.get("url")
|
||||
if url == BedrockFilesConfig.OPENAI_EMBEDDINGS_URL:
|
||||
return True
|
||||
if url:
|
||||
return False
|
||||
body = openai_jsonl_record.get("body", {})
|
||||
if not isinstance(body, dict):
|
||||
return False
|
||||
return "input" in body and "messages" not in body
|
||||
|
||||
# Identifier for the Bedrock Titan v2 InvokeModel body schema as stored
|
||||
# in `model_prices_and_context_window.json`. Centralized so future
|
||||
# embedding-schema variants can add their own value
|
||||
# (e.g. `cohere_v3`, `titan_g1`, `titan_multimodal`) without touching
|
||||
# the detection logic.
|
||||
_TITAN_V2_INVOCATION_SCHEMA = "titan_v2"
|
||||
|
||||
# Substring marker used as a fallback when the registry can't resolve
|
||||
# the model id - notably cross-region inference profile prefixes
|
||||
# (`us.amazon.titan-embed-text-v2:0`) and Bedrock ARN forms, which
|
||||
# `get_model_info` doesn't normalize today.
|
||||
_TITAN_V2_EMBED_MODEL_MARKER = "titan-embed-text-v2"
|
||||
|
||||
# Nested field name under `provider_specific_entry` that identifies the
|
||||
# Bedrock InvokeModel body schema for batch inference.
|
||||
# `provider_specific_entry` is the registry's escape hatch for fields
|
||||
# `get_model_info` doesn't promote to top-level - exactly what we need
|
||||
# here. Documented in the `sample_spec` entry of
|
||||
# `model_prices_and_context_window.json` and surfaced by
|
||||
# `get_model_info` (see `ModelInfo.provider_specific_entry`).
|
||||
_BEDROCK_INVOCATION_SCHEMA_FIELD = "bedrock_invocation_schema"
|
||||
|
||||
@staticmethod
|
||||
def _is_titan_v2_embed_model(model: str) -> bool:
|
||||
"""
|
||||
True iff `model` refers to Amazon Titan Text Embeddings V2.
|
||||
|
||||
Resolution order:
|
||||
1. `model_prices_and_context_window.json` via `get_model_info`.
|
||||
The Titan v2 registry entry carries an explicit
|
||||
`provider_specific_entry.bedrock_invocation_schema` discriminator
|
||||
(`"titan_v2"`). When the registry resolves the id we trust that
|
||||
field as the source of truth - no hardcoded model-id comparison
|
||||
needed.
|
||||
2. Substring fallback (`titan-embed-text-v2` followed by `:`, `/`,
|
||||
or end-of-string) for ids the registry can't normalize. This
|
||||
catches cross-region inference profile prefixes
|
||||
(`us.amazon.titan-embed-text-v2:0`) and Bedrock ARN forms; the
|
||||
marker boundary check rejects lookalikes like
|
||||
`titan-embed-text-v20` or `titan-embed-text-v2-experimental`.
|
||||
|
||||
Tolerant of common id shapes:
|
||||
- "amazon.titan-embed-text-v2:0"
|
||||
- "bedrock/amazon.titan-embed-text-v2:0"
|
||||
- "us.amazon.titan-embed-text-v2:0" (cross-region inference profile)
|
||||
- ARN forms ending in ".../amazon.titan-embed-text-v2:0"
|
||||
"""
|
||||
# Registry-driven path: when get_model_info resolves the id we trust
|
||||
# the registry's discriminator. A resolved id with a different (or
|
||||
# absent) schema value here is intentionally not given a substring
|
||||
# second-chance - the registry is authoritative for ids it knows.
|
||||
registry_schema = BedrockFilesConfig._lookup_provider_specific_field(
|
||||
model, BedrockFilesConfig._BEDROCK_INVOCATION_SCHEMA_FIELD
|
||||
)
|
||||
if registry_schema is not None:
|
||||
return registry_schema == BedrockFilesConfig._TITAN_V2_INVOCATION_SCHEMA
|
||||
|
||||
# Registry silence -> substring fallback for unmapped ids only.
|
||||
normalized = model.lower()
|
||||
if normalized.startswith("bedrock/"):
|
||||
normalized = normalized[len("bedrock/") :]
|
||||
marker = BedrockFilesConfig._TITAN_V2_EMBED_MODEL_MARKER
|
||||
idx = normalized.find(marker)
|
||||
if idx < 0:
|
||||
return False
|
||||
end = idx + len(marker)
|
||||
return end == len(normalized) or normalized[end] in (":", "/")
|
||||
|
||||
@staticmethod
|
||||
def _lookup_provider_specific_field(model_id: str, field: str) -> Optional[str]:
|
||||
"""
|
||||
Read a nested string field from the registry entry's
|
||||
`provider_specific_entry` dict via `litellm.get_model_info`.
|
||||
|
||||
Returns the field's string value when:
|
||||
- the registry resolves `model_id`,
|
||||
- the entry exposes `provider_specific_entry` as a dict, and
|
||||
- that dict has `field` mapped to a non-empty string.
|
||||
Otherwise returns `None`.
|
||||
|
||||
Isolating this means feature detectors (Titan v2 today, future
|
||||
Cohere Embed / Nova Multimodal branches) share one defensive
|
||||
try/except shape instead of duplicating it. The `None` return
|
||||
covers every realistic failure mode: `get_model_info` raises
|
||||
(cross-region profile prefixes, Bedrock ARN forms, unreleased
|
||||
models), returns a non-dict, has no `provider_specific_entry`, or
|
||||
the requested field is missing / non-string / empty.
|
||||
"""
|
||||
try:
|
||||
from litellm import get_model_info
|
||||
|
||||
info = get_model_info(model_id)
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(info, dict):
|
||||
return None
|
||||
provider_specific = info.get("provider_specific_entry")
|
||||
if not isinstance(provider_specific, dict):
|
||||
return None
|
||||
value = provider_specific.get(field)
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
@staticmethod
|
||||
def _coerce_embedding_input_to_string(raw_input: Any, model: str = "") -> str:
|
||||
"""
|
||||
Normalize an OpenAI /v1/embeddings `input` field into the single
|
||||
string that Bedrock Titan v2 InvokeModel expects in `inputText`.
|
||||
|
||||
Accepts: a string, or a single-element list containing one string.
|
||||
Rejects (with actionable messages):
|
||||
- None / missing -> ValueError
|
||||
- Multi-element string lists -> ValueError, prompts caller to
|
||||
emit one JSONL line per input
|
||||
- Pre-tokenized inputs (List[int], List[List[int]]) -> NotImplementedError
|
||||
- Any other type -> ValueError
|
||||
|
||||
Extracted so the validation can be exercised in isolation and so
|
||||
future embedding-provider branches (Titan G1, Cohere) can reuse it
|
||||
without duplicating the type-shaping logic.
|
||||
"""
|
||||
if raw_input is None:
|
||||
raise ValueError(
|
||||
"Embedding batch record is missing required `input` field: "
|
||||
f"model={model}"
|
||||
)
|
||||
|
||||
# Bedrock InvokeModel for Titan v2 takes exactly one string `inputText`
|
||||
# per call. Pre-tokenized inputs and multi-element string lists are
|
||||
# explicitly unsupported so callers emit one JSONL line per embedding
|
||||
# instead of relying on us to silently fan out or concatenate.
|
||||
if isinstance(raw_input, list):
|
||||
if len(raw_input) == 1:
|
||||
candidate = raw_input[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
"Bedrock batch embedding requires one input per JSONL "
|
||||
"record. Got a list with "
|
||||
f"{len(raw_input)} items for model={model}; emit one "
|
||||
"JSONL line per input string instead."
|
||||
)
|
||||
else:
|
||||
candidate = raw_input
|
||||
|
||||
# Catches pre-tokenized inputs (List[int] from OpenAI spec, or a
|
||||
# single int slipping past the list-unwrap above).
|
||||
# NOTE: bool is a subclass of int but treating True/False as a token
|
||||
# is meaningless either way, so the broad check is fine.
|
||||
if isinstance(candidate, (list, int)):
|
||||
raise NotImplementedError(
|
||||
"Bedrock Titan v2 batch embedding does not support "
|
||||
"pre-tokenized integer inputs. Pass `input` as a string "
|
||||
f"(model={model})."
|
||||
)
|
||||
if not isinstance(candidate, str):
|
||||
raise ValueError(
|
||||
"Bedrock batch embedding `input` must be a string (or a "
|
||||
"single-element list of strings). Got type "
|
||||
f"{type(candidate).__name__} for model={model}."
|
||||
)
|
||||
return candidate
|
||||
|
||||
def _map_openai_embedding_to_bedrock_params(
|
||||
self,
|
||||
openai_request_body: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Transform an OpenAI /v1/embeddings request body into the
|
||||
Bedrock InvokeModel `modelInput` for embedding models that AWS
|
||||
supports via batch inference (CreateModelInvocationJob).
|
||||
|
||||
Currently routes Amazon Titan Text Embeddings V2 only; other
|
||||
embedding providers (Titan G1, Titan Multimodal, Cohere Embed,
|
||||
Nova Multimodal Embeddings) raise NotImplementedError until they
|
||||
get a dedicated branch. Splitting them keeps PR scope tight and
|
||||
lets each model's request schema be exercised by its own tests.
|
||||
|
||||
AWS docs (Titan v2 InvokeModel body):
|
||||
https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html
|
||||
"""
|
||||
from litellm.llms.bedrock.embed.amazon_titan_v2_transformation import (
|
||||
AmazonTitanV2Config,
|
||||
)
|
||||
|
||||
_model = openai_request_body.get("model", "")
|
||||
if not self._is_titan_v2_embed_model(_model):
|
||||
# Refuse early instead of silently shaping the body for the wrong
|
||||
# provider. The synchronous /v1/embeddings path supports more
|
||||
# models, but each has a different InvokeModel schema; mapping
|
||||
# them here without dedicated tests would risk corrupt batches.
|
||||
raise NotImplementedError(
|
||||
"Bedrock batch embedding currently supports only Amazon "
|
||||
"Titan Text Embeddings V2 (model id contains "
|
||||
f"'titan-embed-text-v2'). Got model={_model!r}. Track other "
|
||||
"embedding models in https://github.com/BerriAI/litellm/issues."
|
||||
)
|
||||
|
||||
input_text = self._coerce_embedding_input_to_string(
|
||||
openai_request_body.get("input"), model=_model
|
||||
)
|
||||
|
||||
# Map OpenAI-style params (dimensions, encoding_format) onto the
|
||||
# Titan v2 schema (dimensions, embeddingTypes) via the embed config
|
||||
# so this stays in sync with the synchronous /v1/embeddings path.
|
||||
non_default_params = {
|
||||
k: v for k, v in openai_request_body.items() if k not in ("model", "input")
|
||||
}
|
||||
titan_config = AmazonTitanV2Config()
|
||||
inference_params = titan_config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
)
|
||||
return dict(
|
||||
titan_config._transform_request(
|
||||
input=input_text, inference_params=inference_params
|
||||
)
|
||||
)
|
||||
|
||||
def _map_openai_to_bedrock_params(
|
||||
self,
|
||||
openai_request_body: Dict[str, Any],
|
||||
@@ -349,10 +602,19 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
||||
# Determine provider from model name
|
||||
provider = self.get_bedrock_invoke_provider(model)
|
||||
|
||||
# Transform to Bedrock modelInput format
|
||||
model_input = self._map_openai_to_bedrock_params(
|
||||
openai_request_body=openai_body, provider=provider
|
||||
)
|
||||
# Route to the embedding transformer when the OpenAI batch line
|
||||
# targets /v1/embeddings; otherwise fall back to the existing
|
||||
# chat-completion path. We branch here (rather than inside
|
||||
# `_map_openai_to_bedrock_params`) so the chat helper keeps its
|
||||
# narrow contract and the embedding helper can evolve independently.
|
||||
if self._is_embedding_record(_openai_jsonl_content):
|
||||
model_input = self._map_openai_embedding_to_bedrock_params(
|
||||
openai_request_body=openai_body
|
||||
)
|
||||
else:
|
||||
model_input = self._map_openai_to_bedrock_params(
|
||||
openai_request_body=openai_body, provider=provider
|
||||
)
|
||||
|
||||
# Create Bedrock batch record
|
||||
record_id = _openai_jsonl_content.get(
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
"""
|
||||
Amazon Bedrock Mantle - Responses API backend.
|
||||
|
||||
gpt-5.5 / gpt-5.4 on Mantle are exposed ONLY on the `/openai/v1/responses`
|
||||
path (not the standard `/v1/responses`). Payloads and SSE follow the OpenAI
|
||||
Responses spec, so this config inherits OpenAIResponsesAPIConfig and overrides
|
||||
only the endpoint URL and Bearer authentication.
|
||||
|
||||
Auth: AWS Bedrock API key as Bearer token (BEDROCK_MANTLE_API_KEY or the
|
||||
standard AWS_BEARER_TOKEN_BEDROCK), NOT SigV4.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1"
|
||||
|
||||
# Checked longest/most-specific first so a full endpoint URL collapses to host
|
||||
# in one pass and the appended path never doubles.
|
||||
_BASE_SUFFIXES_TO_STRIP = (
|
||||
"/openai/v1/responses",
|
||||
"/v1/responses",
|
||||
"/responses",
|
||||
"/openai/v1",
|
||||
"/v1",
|
||||
)
|
||||
|
||||
|
||||
class BedrockMantleResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
||||
@property
|
||||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.BEDROCK_MANTLE
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
litellm_params: dict,
|
||||
) -> str:
|
||||
region = (
|
||||
get_secret_str("BEDROCK_MANTLE_REGION")
|
||||
or get_secret_str("AWS_REGION")
|
||||
or BEDROCK_MANTLE_DEFAULT_REGION
|
||||
)
|
||||
base = (
|
||||
api_base
|
||||
or get_secret_str("BEDROCK_MANTLE_API_BASE")
|
||||
or f"https://bedrock-mantle.{region}.api.aws"
|
||||
)
|
||||
base = base.rstrip("/")
|
||||
for suffix in _BASE_SUFFIXES_TO_STRIP:
|
||||
if base.endswith(suffix):
|
||||
base = base[: -len(suffix)]
|
||||
break
|
||||
return f"{base}/openai/v1/responses"
|
||||
|
||||
def validate_environment(
|
||||
self, headers: dict, model: str, litellm_params: Optional[GenericLiteLLMParams]
|
||||
) -> dict:
|
||||
litellm_params = litellm_params or GenericLiteLLMParams()
|
||||
api_key = (
|
||||
litellm_params.api_key
|
||||
or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
||||
or get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"Bedrock Mantle API key is required. Set BEDROCK_MANTLE_API_KEY "
|
||||
"(or AWS_BEARER_TOKEN_BEDROCK) or pass api_key."
|
||||
)
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
|
||||
def supports_native_file_search(self) -> bool:
|
||||
return False
|
||||
|
||||
def supports_native_websocket(self) -> bool:
|
||||
return False
|
||||
@@ -5,6 +5,7 @@ Common utilities, constants, and error handling for Black Forest Labs API.
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
@@ -18,6 +19,42 @@ class BlackForestLabsError(BaseLLMException):
|
||||
# API Constants
|
||||
DEFAULT_API_BASE = "https://api.bfl.ai"
|
||||
|
||||
# BFL uses regional subdomains (e.g. gateway.bfl.ai) for polling URLs that
|
||||
# differ from the submission host (api.bfl.ai). We validate against the
|
||||
# registered domain rather than doing a strict same-origin check.
|
||||
_BFL_REGISTERED_DOMAIN = "bfl.ai"
|
||||
|
||||
|
||||
def assert_bfl_polling_url(polling_url: str) -> None:
|
||||
"""Validate that a polling URL points to a BFL-controlled host.
|
||||
|
||||
BFL returns polling URLs on subdomains like ``gateway.bfl.ai`` that differ
|
||||
from the submission host ``api.bfl.ai``. A strict same-origin check would
|
||||
reject these legitimate URLs. Instead we verify the host is ``bfl.ai`` or
|
||||
any subdomain of it, which keeps the SSRF guarantee (credentials only go
|
||||
to BFL-controlled infrastructure) without false-positives on regional hosts.
|
||||
|
||||
Raises:
|
||||
BlackForestLabsError: If the polling URL scheme or host is not trusted.
|
||||
"""
|
||||
parsed = urlparse(polling_url)
|
||||
host = (parsed.hostname or "").lower()
|
||||
|
||||
if parsed.scheme != "https":
|
||||
raise BlackForestLabsError(
|
||||
status_code=502,
|
||||
message="Rejected polling URL: scheme must be https",
|
||||
)
|
||||
|
||||
if host != _BFL_REGISTERED_DOMAIN and not host.endswith(
|
||||
"." + _BFL_REGISTERED_DOMAIN
|
||||
):
|
||||
raise BlackForestLabsError(
|
||||
status_code=502,
|
||||
message="Rejected polling URL: host is not within the bfl.ai domain",
|
||||
)
|
||||
|
||||
|
||||
# Polling configuration
|
||||
DEFAULT_POLLING_INTERVAL = 1.5 # seconds
|
||||
DEFAULT_MAX_POLLING_TIME = 300 # 5 minutes
|
||||
|
||||
@@ -15,7 +15,6 @@ import httpx
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
@@ -29,6 +28,7 @@ from ..common_utils import (
|
||||
DEFAULT_MAX_POLLING_TIME,
|
||||
DEFAULT_POLLING_INTERVAL,
|
||||
BlackForestLabsError,
|
||||
assert_bfl_polling_url,
|
||||
)
|
||||
from .transformation import BlackForestLabsImageEditConfig
|
||||
|
||||
@@ -332,16 +332,11 @@ class BlackForestLabsImageEdit:
|
||||
message="No polling_url in BFL response",
|
||||
)
|
||||
|
||||
# Reject cross-origin polling URLs — the ``x-key`` auth header
|
||||
# would otherwise leak to whatever URL the upstream returns.
|
||||
# VERIA-51.
|
||||
try:
|
||||
assert_same_origin(polling_url, str(initial_response.request.url))
|
||||
except SSRFError as ssrf_err:
|
||||
raise BlackForestLabsError(
|
||||
status_code=502,
|
||||
message=f"Rejected polling URL: {ssrf_err}",
|
||||
)
|
||||
# Reject polling URLs that don't belong to BFL-controlled infrastructure.
|
||||
# BFL uses regional subdomains (e.g. gateway.bfl.ai) that differ from the
|
||||
# submission host (api.bfl.ai), so we validate against the registered
|
||||
# domain rather than doing a strict same-origin check. VERIA-51.
|
||||
assert_bfl_polling_url(polling_url)
|
||||
|
||||
# Get just the auth header for polling
|
||||
polling_headers = {"x-key": headers.get("x-key", "")}
|
||||
@@ -428,16 +423,11 @@ class BlackForestLabsImageEdit:
|
||||
message="No polling_url in BFL response",
|
||||
)
|
||||
|
||||
# Reject cross-origin polling URLs — the ``x-key`` auth header
|
||||
# would otherwise leak to whatever URL the upstream returns.
|
||||
# VERIA-51.
|
||||
try:
|
||||
assert_same_origin(polling_url, str(initial_response.request.url))
|
||||
except SSRFError as ssrf_err:
|
||||
raise BlackForestLabsError(
|
||||
status_code=502,
|
||||
message=f"Rejected polling URL: {ssrf_err}",
|
||||
)
|
||||
# Reject polling URLs that don't belong to BFL-controlled infrastructure.
|
||||
# BFL uses regional subdomains (e.g. gateway.bfl.ai) that differ from the
|
||||
# submission host (api.bfl.ai), so we validate against the registered
|
||||
# domain rather than doing a strict same-origin check. VERIA-51.
|
||||
assert_bfl_polling_url(polling_url)
|
||||
|
||||
# Get just the auth header for polling
|
||||
polling_headers = {"x-key": headers.get("x-key", "")}
|
||||
|
||||
@@ -15,7 +15,6 @@ import httpx
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
@@ -29,6 +28,7 @@ from ..common_utils import (
|
||||
DEFAULT_MAX_POLLING_TIME,
|
||||
DEFAULT_POLLING_INTERVAL,
|
||||
BlackForestLabsError,
|
||||
assert_bfl_polling_url,
|
||||
)
|
||||
from .transformation import BlackForestLabsImageGenerationConfig
|
||||
|
||||
@@ -172,6 +172,10 @@ class BlackForestLabsImageGeneration:
|
||||
raw_response=final_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data=data,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
async def async_image_generation(
|
||||
@@ -274,6 +278,10 @@ class BlackForestLabsImageGeneration:
|
||||
raw_response=final_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data=data,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params_dict,
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
def _poll_for_result_sync(
|
||||
@@ -318,16 +326,11 @@ class BlackForestLabsImageGeneration:
|
||||
message="No polling_url in BFL response",
|
||||
)
|
||||
|
||||
# Reject cross-origin polling URLs — the ``x-key`` auth header
|
||||
# would otherwise leak to whatever URL the upstream returns.
|
||||
# VERIA-51.
|
||||
try:
|
||||
assert_same_origin(polling_url, str(initial_response.request.url))
|
||||
except SSRFError as ssrf_err:
|
||||
raise BlackForestLabsError(
|
||||
status_code=502,
|
||||
message=f"Rejected polling URL: {ssrf_err}",
|
||||
)
|
||||
# Reject polling URLs that don't belong to BFL-controlled infrastructure.
|
||||
# BFL uses regional subdomains (e.g. gateway.bfl.ai) that differ from the
|
||||
# submission host (api.bfl.ai), so we validate against the registered
|
||||
# domain rather than doing a strict same-origin check. VERIA-51.
|
||||
assert_bfl_polling_url(polling_url)
|
||||
|
||||
# Get just the auth header for polling
|
||||
polling_headers = {"x-key": headers.get("x-key", "")}
|
||||
@@ -414,16 +417,11 @@ class BlackForestLabsImageGeneration:
|
||||
message="No polling_url in BFL response",
|
||||
)
|
||||
|
||||
# Reject cross-origin polling URLs — the ``x-key`` auth header
|
||||
# would otherwise leak to whatever URL the upstream returns.
|
||||
# VERIA-51.
|
||||
try:
|
||||
assert_same_origin(polling_url, str(initial_response.request.url))
|
||||
except SSRFError as ssrf_err:
|
||||
raise BlackForestLabsError(
|
||||
status_code=502,
|
||||
message=f"Rejected polling URL: {ssrf_err}",
|
||||
)
|
||||
# Reject polling URLs that don't belong to BFL-controlled infrastructure.
|
||||
# BFL uses regional subdomains (e.g. gateway.bfl.ai) that differ from the
|
||||
# submission host (api.bfl.ai), so we validate against the registered
|
||||
# domain rather than doing a strict same-origin check. VERIA-51.
|
||||
assert_bfl_polling_url(polling_url)
|
||||
|
||||
# Get just the auth header for polling
|
||||
polling_headers = {"x-key": headers.get("x-key", "")}
|
||||
|
||||
@@ -120,6 +120,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig):
|
||||
"stream",
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
@@ -143,7 +144,12 @@ class CohereV2ChatConfig(OpenAIGPTConfig):
|
||||
optional_params["stream"] = value
|
||||
if param == "temperature":
|
||||
optional_params["temperature"] = value
|
||||
if param == "max_tokens":
|
||||
if (
|
||||
param == "max_tokens"
|
||||
and "max_completion_tokens" not in non_default_params
|
||||
):
|
||||
optional_params["max_tokens"] = value
|
||||
if param == "max_completion_tokens":
|
||||
optional_params["max_tokens"] = value
|
||||
if param == "n":
|
||||
optional_params["num_generations"] = value
|
||||
|
||||
@@ -256,7 +256,10 @@ class LiteLLMAiohttpTransport(AiohttpTransport):
|
||||
from yarl import URL as YarlURL
|
||||
|
||||
try:
|
||||
data = request.content
|
||||
# Coerce an empty body to None so aiohttp does not attach a
|
||||
# `Content-Type: application/octet-stream` header for bodyless
|
||||
# requests (e.g. DELETE /responses/{id}), which upstream APIs reject.
|
||||
data = request.content or None
|
||||
except httpx.RequestNotRead:
|
||||
data = request.stream # type: ignore
|
||||
request.headers.pop("transfer-encoding", None) # handled by aiohttp
|
||||
|
||||
@@ -589,6 +589,7 @@ class AsyncHTTPHandler:
|
||||
params: Optional[dict] = None,
|
||||
headers: Optional[dict] = None,
|
||||
follow_redirects: Optional[bool] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
):
|
||||
# Set follow_redirects to UseClientDefault if None
|
||||
_follow_redirects = (
|
||||
@@ -599,7 +600,11 @@ class AsyncHTTPHandler:
|
||||
params.update(HTTPHandler.extract_query_params(url))
|
||||
|
||||
response = await self.client.get(
|
||||
url, params=params, headers=headers, follow_redirects=_follow_redirects # type: ignore
|
||||
url,
|
||||
params=params,
|
||||
headers=headers, # type: ignore
|
||||
follow_redirects=_follow_redirects, # type: ignore
|
||||
timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT,
|
||||
)
|
||||
return response
|
||||
|
||||
@@ -1115,6 +1120,7 @@ class HTTPHandler:
|
||||
params: Optional[dict] = None,
|
||||
headers: Optional[dict] = None,
|
||||
follow_redirects: Optional[bool] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
):
|
||||
# Set follow_redirects to UseClientDefault if None
|
||||
_follow_redirects = (
|
||||
@@ -1128,6 +1134,7 @@ class HTTPHandler:
|
||||
params=params,
|
||||
headers=headers,
|
||||
follow_redirects=_follow_redirects,
|
||||
timeout=timeout if timeout is not None else USE_CLIENT_DEFAULT,
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
@@ -1751,6 +1751,7 @@ class BaseLLMHTTPHandler:
|
||||
api_base=api_base,
|
||||
optional_params=optional_params,
|
||||
data=data,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
@@ -1833,6 +1834,7 @@ class BaseLLMHTTPHandler:
|
||||
api_base=api_base,
|
||||
optional_params=optional_params,
|
||||
data=data,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
@@ -2026,6 +2028,7 @@ class BaseLLMHTTPHandler:
|
||||
litellm_params={
|
||||
"preset_cache_key": None,
|
||||
"stream_response": {},
|
||||
"model_info": kwargs.get("model_info"),
|
||||
**anthropic_messages_optional_request_params,
|
||||
},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
@@ -2585,6 +2588,8 @@ class BaseLLMHTTPHandler:
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
headers.setdefault("Content-Type", "application/json")
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
@@ -2675,6 +2680,8 @@ class BaseLLMHTTPHandler:
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
headers.setdefault("Content-Type", "application/json")
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=input,
|
||||
@@ -5527,6 +5534,7 @@ class BaseLLMHTTPHandler:
|
||||
user_api_key_dict: Optional[Any] = None,
|
||||
litellm_metadata: Optional[Dict[str, Any]] = None,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
first_message: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
@@ -5558,6 +5566,7 @@ class BaseLLMHTTPHandler:
|
||||
api_base=api_base,
|
||||
timeout=timeout,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
first_message=first_message,
|
||||
**kwargs,
|
||||
)
|
||||
await handler.run()
|
||||
@@ -5623,6 +5632,7 @@ class BaseLLMHTTPHandler:
|
||||
logging_obj=logging_obj,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
request_data=_request_data,
|
||||
first_message=first_message,
|
||||
)
|
||||
await streaming.bidirectional_forward()
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user