scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

Merge remote-tracking branch 'origin/litellm_mcp_listed_tool_metadata' into litellm_mcp_gateway_sign_in_provider

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	litellm/proxy/guardrails/guardrail_hooks/agent_365/agent_365.py
#	tests/unit/proxy/_experimental/mcp_server/test_caller_sign_in.py
This commit is contained in:
yucheng 2026-10-02 08:14:47 +00:00
commit 6adc2344a0
1482 changed files with 40957 additions and 10162 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

134
.circleci/config.yml
View file

@ -408,7 +408,7 @@ jobs:
- run:
name: Run Windows-specific test
command: |
uv run --no-sync python -m pytest tests/windows_tests/ -v
uv run --no-sync python -m pytest --tb=short tests/windows_tests/ -v
windows_release_wheel:
executor:
@ -486,6 +486,7 @@ jobs:
- install_rust
- run:
name: Build the wheel
no_output_timeout: 30m
environment:
UV_HTTP_TIMEOUT: "300"
command: |
@ -550,7 +551,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="tr ' ' '\\n' | 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 --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
@ -624,7 +625,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="tr ' ' '\\n' | 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 --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
@ -696,7 +697,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -751,7 +752,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/proxy_admin_ui_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -814,7 +815,7 @@ jobs:
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="tr ' ' '\\n' | 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 --tb=short \
-v \
-k 'router' \
-n 4 \
@ -858,7 +859,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/router_unit_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -903,7 +904,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -947,7 +948,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="tr ' ' '\\n' | 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 --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=20 \
@ -985,7 +986,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/llm_translation/realtime/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1030,7 +1031,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/agent_tests/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1074,7 +1075,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/guardrails_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1120,7 +1121,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/unified_google_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1175,7 +1176,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/llm_responses_api_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1209,7 +1210,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/ocr_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1253,7 +1254,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/search_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1297,7 +1298,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/batches_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1341,7 +1342,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/litellm_utils_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1386,7 +1387,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/pass_through_unit_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1431,7 +1432,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/image_gen_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1465,7 +1466,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/logging_callback_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
-n 4 \
@ -1510,7 +1511,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/audio_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
@ -1530,61 +1531,6 @@ jobs:
paths:
- audio_coverage.xml
- audio_coverage
redis_caching_unit_tests:
docker:
- *python312_image
working_directory: ~/project
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
# Run pytest and generate JUnit XML report
- run:
name: Run tests
command: |
mkdir -p test-results
TEST_FILES=$(printf "%s\n" \
tests/local_testing/test_dual_cache.py \
tests/local_testing/test_redis_batch_optimizations.py \
tests/local_testing/test_redis_increment_with_floor.py \
tests/local_testing/test_router_utils.py)
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -s \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 -n 2 \
--reruns 2 --reruns-delay 1"
no_output_timeout: 20m
- run:
name: Rename the coverage files
command: |
mv coverage.xml redis_caching_coverage.xml
mv .coverage redis_caching_coverage
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- redis_caching_coverage.xml
- redis_caching_coverage
installing_litellm_on_python:
docker:
- *python312_image
@ -1604,7 +1550,7 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
uv run --no-sync python -m pytest --tb=short -vv tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
installing_litellm_on_python_3_13:
docker:
@ -1628,7 +1574,7 @@ jobs:
- run:
name: Run tests
command: |
uv run --no-sync python -m pytest -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
uv run --no-sync python -m pytest --tb=short -v tests/local_testing/test_basic_python_version.py -k "not legacy_resolver"
installing_litellm_on_python_v2_migration_resolver:
docker:
@ -1659,7 +1605,7 @@ jobs:
- run:
name: Run both migration resolvers against Postgres
command: |
uv run --no-sync python -m pytest -vv \
uv run --no-sync python -m pytest --tb=short -vv \
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings \
tests/local_testing/test_basic_python_version.py::test_litellm_proxy_server_config_no_general_settings_legacy_resolver
@ -1828,7 +1774,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
@ -1925,7 +1871,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-s -v \
--junitxml=test-results/junit.xml \
-n 4 \
@ -2012,7 +1958,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/openai_endpoints_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-s -vv \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2095,7 +2041,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/otel_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2147,7 +2093,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit-2.xml \
--durations=5"
@ -2228,7 +2174,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/spend_tracking_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2333,7 +2279,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/multi_instance_e2e_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2405,7 +2351,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/store_model_in_db_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2490,7 +2436,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/basic_proxy_startup_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv \
--junitxml=test-results/junit-2.xml \
--durations=5"
@ -2587,7 +2533,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/pass_through_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2658,7 +2604,7 @@ jobs:
TEST_FILES=$(circleci tests glob "tests/proxy_e2e_anthropic_messages_tests/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest --tb=short \
-vv -s \
--junitxml=test-results/junit.xml \
--durations=5"
@ -2688,7 +2634,7 @@ jobs:
- run:
name: Combine Coverage
command: |
uv tool run --from 'coverage[toml]==7.10.6' coverage combine realtime_translation_coverage ocr_coverage search_coverage logging_coverage audio_coverage local_testing_part1_coverage local_testing_part2_coverage pass_through_unit_tests_coverage batches_coverage guardrails_coverage redis_caching_coverage agent_coverage google_generate_content_endpoint_coverage litellm_utils_coverage router_unit_tests_coverage auth_ui_unit_tests_coverage
uv tool run --from 'coverage[toml]==7.10.6' coverage combine realtime_translation_coverage ocr_coverage search_coverage logging_coverage audio_coverage local_testing_part1_coverage local_testing_part2_coverage pass_through_unit_tests_coverage batches_coverage guardrails_coverage agent_coverage google_generate_content_endpoint_coverage litellm_utils_coverage router_unit_tests_coverage auth_ui_unit_tests_coverage
uv tool run --from 'coverage[toml]==7.10.6' coverage xml
- codecov/upload:
file: ./coverage.xml
@ -3188,7 +3134,7 @@ jobs:
name: Test provider capture and replay harness
command: |
mkdir -p test-results/provider-replay-harness
uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
uv run --no-sync pytest --tb=short -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
--junitxml=test-results/provider-replay-harness/junit.xml \
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
@ -3491,7 +3437,6 @@ workflows:
- image_gen_testing
- logging_testing
- audio_testing
- redis_caching_unit_tests
- upload-coverage:
requires:
- realtime_translation_testing
@ -3506,7 +3451,6 @@ workflows:
- image_gen_testing
- logging_testing
- audio_testing
- redis_caching_unit_tests
- langfuse_logging_unit_tests
- local_testing_part1
- local_testing_part2

4
.circleci/scripts/run_integration.sh
View file

@ -168,6 +168,7 @@ start_proxy() {
"${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
@ -190,7 +191,7 @@ if [ "$suite" = management ] || [ "$suite" = mcp ]; then
fi
if [ "$suite" = providers ]; then
INTEGRATION_RUN_ID="$integration_identity" .venv/bin/python -m pytest --noconftest -o addopts= \
INTEGRATION_RUN_ID="$integration_identity" .venv/bin/python -m pytest --tb=short --noconftest -o addopts= \
--strict-markers --strict-config -p no:pytest-retry -p no:rerunfailures --timeout=30 \
tests/e2e/test_provider_edge.py::TestReplayMode::test_content_drift_returns_the_miss_status_naming_both_keys \
tests/e2e/test_provider_edge.py::TestReplayMode::test_exhausted_key_returns_the_miss_status \
@ -228,6 +229,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \

7
.circleci/scripts/unit_selection.sh
View file

@ -77,6 +77,7 @@ legacy_paths() {
echo tests/unit/embeddings
echo tests/unit/endpoints
echo tests/unit/files
echo tests/unit/harness
echo tests/unit/images
echo tests/unit/interactions
echo tests/unit/messages
@ -107,7 +108,7 @@ legacy_paths() {
echo tests/unit/proxy/test_update_spend.py
echo tests/unit/skills/test_skills_db.py ;;
proxy-db-endpoints-and-responses)
echo tests/unit/proxy/engine
echo tests/unit/proxy/lens
echo tests/unit/proxy/auth/test_models_fallback_endpoint.py
echo tests/unit/proxy/common_utils/test_check_batch_cost.py
echo tests/unit/proxy/common_utils/test_check_responses_cost.py
@ -116,7 +117,7 @@ legacy_paths() {
echo tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py
echo tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py
echo tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py
echo tests/unit/proxy/response_polling/test_response_polling_handler.py
echo tests/unit/proxy/response_polling
echo tests/unit/proxy/test_custom_tokenizer_bug.py
echo tests/unit/proxy/test_get_favicon.py
echo tests/unit/proxy/test_get_image.py
@ -145,12 +146,14 @@ legacy_paths() {
echo tests/unit/proxy/test_proxy_token_counter.py
echo tests/unit/proxy/test_server_root_path.py ;;
proxy-db-proxy-server-core)
echo tests/unit/proxy/test__lazy_features.py
echo tests/unit/proxy/test_aproxy_startup.py
echo tests/unit/proxy/test_proxy_server.py ;;
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;
proxy-extras) echo tests/unit/litellm_proxy_extras ;;
proxy-infra)
echo tests/unit/gateway
echo tests/unit/proxy/management
echo tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py
echo tests/unit/proxy/roi_calculator ;;
responses-caching-types)

13
.github/actions/cache-cargo-build/action.yml vendored
View file

@ -25,6 +25,7 @@ runs:
using: composite
steps:
- name: Restore the Cargo registry and target directory
if: github.ref == 'refs/heads/main'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
@ -34,3 +35,15 @@ runs:
key: ${{ runner.os }}-maturin-${{ inputs.profile }}-${{ hashFiles('litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-maturin-${{ inputs.profile }}-
- name: Restore the Cargo registry and target directory
if: github.ref != 'refs/heads/main'
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-maturin-${{ inputs.profile }}-${{ hashFiles('litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-maturin-${{ inputs.profile }}-

10
.github/actions/cache-prisma-binaries/action.yml vendored
View file

@ -30,6 +30,7 @@ runs:
echo "version=${version}" >> "$GITHUB_OUTPUT"
- name: Restore Prisma binaries
if: github.ref == 'refs/heads/main'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
# ~/.cache/prisma-python holds the npm install tree prisma-client-py
@ -38,3 +39,12 @@ runs:
~/.cache/prisma-python
~/.cache/prisma
key: ${{ runner.os }}-prisma-binaries-${{ steps.version.outputs.version }}
- name: Restore Prisma binaries
if: github.ref != 'refs/heads/main'
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/prisma-python
~/.cache/prisma
key: ${{ runner.os }}-prisma-binaries-${{ steps.version.outputs.version }}

25
.github/actions/cache-uv-downloads/action.yml vendored Normal file
View file

@ -0,0 +1,25 @@
name: "Cache uv downloads"
description: >-
Restore the uv download cache on every run and save it only from main, so pull
requests reuse main's cache instead of evicting it with their own copies.
runs:
using: composite
steps:
- name: Restore and save the uv download cache
if: github.ref == 'refs/heads/main'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: ${{ env.UV_CACHE_DIR }}
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-
- name: Restore the uv download cache
if: github.ref != 'refs/heads/main'
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: ${{ env.UV_CACHE_DIR }}
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-

3
.github/actions/setup-uv-with-retries/action.yml vendored
View file

@ -17,6 +17,7 @@ runs:
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}
save-cache: ${{ github.ref == 'refs/heads/main' }}
- name: Wait before attempt 2
if: steps.attempt-1.outcome == 'failure'
@ -30,6 +31,7 @@ runs:
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}
save-cache: ${{ github.ref == 'refs/heads/main' }}
- name: Wait before attempt 3
if: steps.attempt-2.outcome == 'failure'
@ -41,3 +43,4 @@ runs:
uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1
with:
version: ${{ inputs.version }}
save-cache: ${{ github.ref == 'refs/heads/main' }}

8
.github/ci-coverage-allowlist.yml vendored
View file

@ -4,6 +4,14 @@ description: >-
by a job nor listed here, so every entry below is a decision on the record.
test_paths:
- reason: >-
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and
tests/unit/llms/*/harness
paths:
- tests/harness_e2e
- reason: >-
The Rust/Python parity harness is run manually through its local CLI. Recorded replay,
fixture generation, and harness checks are intentionally outside pull request CI

1
.github/e2e-stack/select_tests.py vendored
View file

@ -11,6 +11,7 @@ UNSUPPORTED: Final = re.compile(
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
r"|^tests/e2e/logging/test_otel_v2_langfuse_generation_output_e2e\.py$"
r"|^tests/e2e/logging/test_langsmith_batch_serialization_e2e\.py$"
r"|^tests/e2e/logging/test_s3_log_e2e\.py$"
r"|^tests/e2e/secret_manager/"
)
HARNESS: Final = re.compile(

48
.github/scripts/assert_ci_coverage.py vendored
View file

@ -9,6 +9,7 @@ import sys
import warnings
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
import yaml
@ -35,7 +36,7 @@ GLOB_CHARS = frozenset("*?")
# itself decomposed one level deeper and is checked through its own entry.
SHARDED_ROOTS: tuple[str, ...] = (
"tests/test_litellm",
"tests/test_litellm/proxy",
"tests/unit/proxy",
)
@ -119,11 +120,48 @@ def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
)
def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]:
SELECTION_ARM_RE = re.compile(r"(?ms)^\s*([A-Za-z0-9_|*-]+)\)\s*(.*?);;")
def _unit_selection_arms(repo_root: pathlib.Path = REPO_ROOT) -> Mapping[str, frozenset[str]]:
script: Final = repo_root / ".circleci/scripts/unit_selection.sh"
if not script.is_file():
return frozenset()
return frozenset(match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(_uncommented(script.read_text())))
return MappingProxyType({})
text: Final = _uncommented(script.read_text())
return MappingProxyType(
{
label: frozenset(
match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(body)
)
for label, body in SELECTION_ARM_RE.findall(text)
}
)
def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]:
return frozenset(
token for tokens in _unit_selection_arms(repo_root).values() for token in tokens
)
def _wired_unit_flags(scalars: Iterable[Scalar]) -> frozenset[str]:
return frozenset(
scalar.value
for scalar in scalars
if scalar.key == "unit-flag" and "${{" not in scalar.value
)
def _shard_tokens(
scalars: Iterable[Scalar], arms: Mapping[str, frozenset[str]]
) -> frozenset[str]:
wired: Final = _wired_unit_flags(scalars)
return _invoked_test_tokens(scalars) | frozenset(
token
for label, tokens in arms.items()
if label in wired
for token in tokens
)
def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
@ -480,7 +518,7 @@ def _check_slices() -> int:
def _check_shards() -> int:
findings = _unassigned_shard_children(_invoked_test_tokens(_all_scalars()))
findings = _unassigned_shard_children(_shard_tokens(_all_scalars(), _unit_selection_arms()))
if findings:
_report(
"test directories and files that no shard claims",

11
.github/workflows/_test-unit-base.yml vendored
View file

@ -132,12 +132,7 @@ jobs:
- name: Cache uv dependencies
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: ${{ env.UV_CACHE_DIR }}
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-
uses: ./.github/actions/cache-uv-downloads
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
@ -274,7 +269,7 @@ jobs:
- name: Upload to Codecov
id: codecov-upload
continue-on-error: true
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
with:
use_oidc: true
directory: coverage-reports
@ -285,7 +280,7 @@ jobs:
- name: Upload to Codecov (retry)
if: steps.codecov-upload.outcome == 'failure'
continue-on-error: true
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
with:
use_oidc: true
directory: coverage-reports

10
.github/workflows/lens-worker.yml vendored
View file

@ -5,13 +5,13 @@ on:
branches: [main, litellm_oss_branch, "litellm_**"]
paths:
- deploy/lens/**
- litellm/proxy/engine/**
- litellm/proxy/lens/**
- .github/workflows/lens-worker.yml
push:
branches: [main, litellm_agent_engine]
branches: [main]
paths:
- deploy/lens/**
- litellm/proxy/engine/**
- litellm/proxy/lens/**
- .github/workflows/lens-worker.yml
workflow_dispatch:
@ -41,8 +41,8 @@ jobs:
--security-opt no-new-privileges --entrypoint python \
lens-worker:${{ github.sha }} -c '
import os
import engine.worker
from engine.trace_store import trace_store
import lens.worker
from lens.trace_store import trace_store
assert os.getuid() == 65532
with trace_store() as store:
assert store.count() == 0

12
.github/workflows/mutation-test.yml vendored
View file

@ -44,6 +44,7 @@ jobs:
version: "0.10.9"
- name: Cache uv dependencies
if: github.ref == 'refs/heads/main'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
@ -53,6 +54,17 @@ jobs:
restore-keys: |
${{ runner.os }}-uv-
- name: Cache uv dependencies
if: github.ref != 'refs/heads/main'
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build

12
.github/workflows/test-code-quality.yml vendored
View file

@ -44,6 +44,7 @@ jobs:
version: "0.10.9"
- name: Cache uv dependencies
if: github.ref == 'refs/heads/main'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
@ -53,6 +54,17 @@ jobs:
restore-keys: |
${{ runner.os }}-uv-
- name: Cache uv dependencies
if: github.ref != 'refs/heads/main'
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build

14
.github/workflows/test-postgres.yml vendored
View file

@ -95,7 +95,7 @@ jobs:
version: "0.10.9"
- name: Cache uv dependencies
if: steps.changes.outputs.decision != 'skip'
if: steps.changes.outputs.decision != 'skip' && github.ref == 'refs/heads/main'
timeout-minutes: 5
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
@ -106,6 +106,18 @@ jobs:
restore-keys: |
${{ runner.os }}-uv-postgres-
- name: Cache uv dependencies
if: steps.changes.outputs.decision != 'skip' && github.ref != 'refs/heads/main'
timeout-minutes: 5
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-postgres-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-postgres-
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 12

2
.github/workflows/test-redis-compat.yml vendored
View file

@ -98,7 +98,7 @@ jobs:
- name: Upload Redis coverage
if: matrix.redis-version == '5.3.1'
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5.5.5
with:
use_oidc: true
files: coverage-redis.xml

3
.github/workflows/test-rust.yml vendored
View file

@ -83,6 +83,7 @@ jobs:
with:
workspaces: litellm-rust
cache-on-failure: true
save-if: ${{ github.ref == 'refs/heads/main' }}
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
@ -121,6 +122,7 @@ jobs:
with:
workspaces: litellm-rust
cache-on-failure: true
save-if: ${{ github.ref == 'refs/heads/main' }}
- run: cargo nextest run --workspace --locked
@ -162,6 +164,7 @@ jobs:
with:
workspaces: litellm-rust
cache-on-failure: true
save-if: ${{ github.ref == 'refs/heads/main' }}
- run: uv build --wheel --out-dir dist

12
.github/workflows/test-terraform-provider.yml vendored
View file

@ -77,6 +77,7 @@ jobs:
version: "0.10.9"
- name: Cache uv dependencies
if: github.ref == 'refs/heads/main'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
@ -86,6 +87,17 @@ jobs:
restore-keys: |
${{ runner.os }}-uv-
- name: Cache uv dependencies
if: github.ref != 'refs/heads/main'
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build

13
.github/workflows/test-unit-documentation.yml vendored
View file

@ -54,7 +54,7 @@ jobs:
version: "0.10.9"
- name: Cache uv dependencies
if: steps.changes.outputs.decision != 'skip'
if: steps.changes.outputs.decision != 'skip' && github.ref == 'refs/heads/main'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
@ -64,6 +64,17 @@ jobs:
restore-keys: |
${{ runner.os }}-uv-
- name: Cache uv dependencies
if: steps.changes.outputs.decision != 'skip' && github.ref != 'refs/heads/main'
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-cargo-build

141
.github/workflows/test-unit.yml vendored
View file

@ -140,38 +140,46 @@ jobs:
- shard: proxy-endpoints
artifact-name: proxy-endpoints
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/list_api
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers
tests/test_litellm/proxy/anthropic_endpoints
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/container_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/ocr_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/credential_endpoints
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/rerank_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
tests/test_litellm/proxy/utils
tests/unit/proxy/analytics_endpoints
tests/unit/proxy/management_endpoints
tests/unit/proxy/list_api
tests/unit/proxy/memory
tests/unit/proxy/guardrails
tests/unit/proxy/management_helpers
--ignore=tests/unit/proxy/management_endpoints/test_jwt_key_mapping.py
--ignore=tests/unit/proxy/management_endpoints/test_key_generate_prisma.py
--ignore=tests/unit/proxy/management_endpoints/test_roi_calculator_endpoints.py
--ignore=tests/unit/proxy/management_helpers/test_audit_logs_proxy.py
--ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py
--ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py
--ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py
--ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py
tests/unit/proxy/anthropic_endpoints
tests/unit/proxy/google_endpoints
tests/unit/proxy/openai_files_endpoint
tests/unit/proxy/batches_endpoints
tests/unit/proxy/container_endpoints
tests/unit/proxy/fine_tuning_endpoints
tests/unit/proxy/vector_store_files_endpoints
tests/unit/proxy/video_endpoints
tests/unit/proxy/response_api_endpoints
tests/unit/proxy/image_endpoints
tests/unit/proxy/ocr_endpoints
tests/unit/proxy/vector_store_endpoints
tests/unit/proxy/agent_endpoints
tests/unit/proxy/a2a
tests/unit/proxy/credential_endpoints
tests/unit/proxy/discovery_endpoints
tests/unit/proxy/health_endpoints
tests/unit/proxy/shutdown
tests/unit/proxy/public_endpoints
tests/unit/proxy/prompts
tests/unit/proxy/rag_endpoints
tests/unit/proxy/rerank_endpoints
tests/unit/proxy/realtime_endpoints
tests/unit/proxy/ui_crud_endpoints
tests/unit/proxy/config_resolvers
tests/unit/proxy/utils
workers: 4
reruns: 2
timeout-minutes: 20
@ -179,7 +187,7 @@ jobs:
- shard: proxy-server
artifact-name: proxy-server
test-path: "tests/test_litellm/proxy/proxy_server"
test-path: "tests/unit/proxy/proxy_server"
workers: 4
reruns: 2
timeout-minutes: 60
@ -188,25 +196,66 @@ jobs:
- shard: proxy-infra
artifact-name: proxy-infra
test-path: >-
tests/test_litellm/proxy/db
tests/test_litellm/proxy/middleware
tests/test_litellm/proxy/spend_tracking
tests/test_litellm/proxy/pass_through_endpoints
tests/test_litellm/proxy/_experimental
tests/test_litellm/proxy/experimental
tests/test_litellm/proxy/common_utils
tests/test_litellm/proxy/enterprise_billing
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py
tests/unit/proxy/db
--ignore=tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py
--ignore=tests/unit/proxy/db/test_update_daily_tag_spend.py
tests/unit/proxy/middleware
--ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py
tests/unit/proxy/spend_tracking
--ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py
tests/unit/proxy/pass_through_endpoints
tests/unit/proxy/_experimental
--ignore=tests/unit/proxy/_experimental/mcp_server
tests/unit/proxy/experimental
tests/unit/proxy/common_utils
--ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py
--ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py
--ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py
--ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py
--ignore=tests/unit/proxy/common_utils/test_realtime_cache.py
tests/unit/proxy/enterprise_billing
tests/unit/proxy/types_utils
tests/unit/proxy/logging_endpoints
unit-flag: proxy-infra
workers: 4
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: proxy-infra-root
artifact-name: proxy-infra-root
test-path: >-
tests/unit/proxy/test_*.py
--ignore=tests/unit/proxy/test_aproxy_startup.py
--ignore=tests/unit/proxy/test_credential_slot_registry.py
--ignore=tests/unit/proxy/test_custom_callback_input.py
--ignore=tests/unit/proxy/test_custom_logger_s3_gcs.py
--ignore=tests/unit/proxy/test_custom_tokenizer_bug.py
--ignore=tests/unit/proxy/test_db_schema_changes.py
--ignore=tests/unit/proxy/test_deprecated_key_grace_period.py
--ignore=tests/unit/proxy/test_get_favicon.py
--ignore=tests/unit/proxy/test_get_image.py
--ignore=tests/unit/proxy/test_prisma_client_backoff_retry.py
--ignore=tests/unit/proxy/test_prompt_test_endpoint.py
--ignore=tests/unit/proxy/test_proxy_config_unit_test.py
--ignore=tests/unit/proxy/test_proxy_custom_auth.py
--ignore=tests/unit/proxy/test_proxy_reject_logging.py
--ignore=tests/unit/proxy/test_proxy_server.py
--ignore=tests/unit/proxy/test_proxy_setting_guardrails.py
--ignore=tests/unit/proxy/test_proxy_token_counter.py
--ignore=tests/unit/proxy/test_proxy_utils.py
--ignore=tests/unit/proxy/test_reducto_ocr_route.py
--ignore=tests/unit/proxy/test_response_polling_pre_call_checks.py
--ignore=tests/unit/proxy/test_server_root_path.py
--ignore=tests/unit/proxy/test_ui_path_detection.py
--ignore=tests/unit/proxy/test_unit_test_proxy_hooks.py
--ignore=tests/unit/proxy/test_update_spend.py
--ignore=tests/unit/proxy/test_zero_cost_model_budget_bypass.py
workers: 4
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: caching-local
artifact-name: caching-local
test-path: ""

2
.gitignore vendored
View file

@ -104,7 +104,7 @@ litellm_config.yaml
.cursor
litellm/proxy/to_delete_loadtest_work/*
update_model_cost_map.py
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
scripts/test_vertex_ai_search.py
LAZY_LOADING_IMPROVEMENTS.md
STABILIZATION_TODO.md

2
AGENTS.md
View file

@ -62,7 +62,7 @@ Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
If you get an LIT001 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing

4
CONTRIBUTING.md
View file

@ -98,7 +98,7 @@ Add your tests to the [`tests/unit/` directory](https://github.com/BerriAI/litel
The `tests/unit/` directory follows the same structure as `litellm/`:
- `litellm/proxy/caching_routes.py` → `tests/test_litellm/proxy/test_caching_routes.py`
- `litellm/proxy/caching_routes.py` → `tests/unit/proxy/test_caching_routes.py`
- `litellm/utils.py` → `tests/unit/test_utils.py`
### Example Test
@ -136,7 +136,7 @@ If you're running broader test suites, proxy tests, or anything that touches Pos
make install-test-deps
```
This syncs the locked test environment used across the repo, including `psycopg` v3 plus `psycopg-binary` (used by `pytest-postgresql`), `psycopg2-binary` (used by some proxy E2E tests), and a generated Prisma client for DB-backed proxy tests, so pytest startup matches CI without manual package installs.
This syncs the locked test environment used across the repo, including `psycopg` v3 plus `psycopg-binary`, `psycopg2-binary` (used by some proxy E2E tests), and a generated Prisma client for DB-backed proxy tests, so pytest startup matches CI without manual package installs.
### Running Linting and Formatting Checks

12
Makefile
View file

@ -1,7 +1,7 @@
# LiteLLM Makefile
# Simple Makefile for running tests and basic development tasks
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc test-unit-proxy-root \
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
test-rust-extension rust-sqlx-prepare \
@ -47,6 +47,7 @@ help:
@echo " make test-unit-proxy-guardrails - Run proxy guardrails+mgmt tests (~51 files)"
@echo " make test-unit-proxy-core - Run proxy auth+client+db+hooks tests (~52 files)"
@echo " make test-unit-proxy-misc - Run proxy misc tests (~77 files)"
@echo " make test-unit-proxy-root - Run proxy root-file tests (tests/unit/proxy/test_*.py)"
@echo " make test-unit-integrations - Run integration tests (~60 files)"
@echo " make test-unit-core-utils - Run core utils tests (~32 files)"
@echo " make test-unit-other - Run other tests (caching, responses, etc., ~69 files)"
@ -321,13 +322,16 @@ test-unit-llms: install-test-deps
$(UV_RUN) pytest tests/unit/llms --tb=short -vv -n 4 --durations=20
test-unit-proxy-guardrails: install-test-deps
$(UV_RUN) pytest tests/test_litellm/proxy/guardrails tests/test_litellm/proxy/management_endpoints tests/test_litellm/proxy/management_helpers --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/proxy/guardrails tests/unit/proxy/management_endpoints tests/unit/proxy/management_helpers --tb=short -vv -n 4 --durations=20
test-unit-proxy-core: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/auth tests/unit/proxy/client tests/test_litellm/proxy/db tests/unit/proxy/hooks tests/unit/proxy/policy_engine --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/proxy/auth tests/unit/proxy/client tests/unit/proxy/db tests/unit/proxy/hooks tests/unit/proxy/policy_engine --ignore=tests/unit/proxy/db/db_transaction_queue/test_e2e_pod_lock_manager.py --ignore=tests/unit/proxy/db/test_update_daily_tag_spend.py --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/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 tests/unit/proxy/test_proxy_server_endpoints_and_startup.py tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py --tb=short -vv -n 4 --durations=20
$(UV_RUN) pytest tests/unit/proxy/agent_endpoints tests/unit/proxy/anthropic_endpoints tests/unit/proxy/common_utils --ignore=tests/unit/proxy/common_utils/test_cache_aware_routing.py --ignore=tests/unit/proxy/common_utils/test_check_batch_cost.py --ignore=tests/unit/proxy/common_utils/test_check_responses_cost.py --ignore=tests/unit/proxy/common_utils/test_proxy_encrypt_decrypt.py --ignore=tests/unit/proxy/common_utils/test_realtime_cache.py tests/unit/proxy/discovery_endpoints tests/unit/proxy/experimental tests/unit/proxy/google_endpoints tests/unit/proxy/health_endpoints tests/unit/proxy/image_endpoints tests/unit/proxy/middleware --ignore=tests/unit/proxy/middleware/test_request_size_limit_middleware.py tests/unit/proxy/openai_files_endpoint tests/unit/proxy/pass_through_endpoints tests/unit/proxy/prompts tests/unit/proxy/public_endpoints tests/unit/proxy/response_api_endpoints tests/unit/proxy/shutdown tests/unit/proxy/spend_tracking --ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py tests/unit/proxy/ui_crud_endpoints tests/unit/proxy/vector_store_endpoints tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py --ignore=tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py --ignore=tests/unit/proxy/google_endpoints/test_google_endpoint_routing.py --ignore=tests/unit/proxy/google_endpoints/test_google_gemini_proxy_request.py --ignore=tests/unit/proxy/public_endpoints/test_blog_posts_endpoint.py --tb=short -vv -n 4 --durations=20
test-unit-proxy-root: install-test-deps
$(UV_RUN) pytest tests/unit/proxy/test_*.py --ignore=tests/unit/proxy/test_aproxy_startup.py --ignore=tests/unit/proxy/test_credential_slot_registry.py --ignore=tests/unit/proxy/test_custom_callback_input.py --ignore=tests/unit/proxy/test_custom_logger_s3_gcs.py --ignore=tests/unit/proxy/test_custom_tokenizer_bug.py --ignore=tests/unit/proxy/test_db_schema_changes.py --ignore=tests/unit/proxy/test_deprecated_key_grace_period.py --ignore=tests/unit/proxy/test_get_favicon.py --ignore=tests/unit/proxy/test_get_image.py --ignore=tests/unit/proxy/test_prisma_client_backoff_retry.py --ignore=tests/unit/proxy/test_prompt_test_endpoint.py --ignore=tests/unit/proxy/test_proxy_config_unit_test.py --ignore=tests/unit/proxy/test_proxy_custom_auth.py --ignore=tests/unit/proxy/test_proxy_reject_logging.py --ignore=tests/unit/proxy/test_proxy_server.py --ignore=tests/unit/proxy/test_proxy_setting_guardrails.py --ignore=tests/unit/proxy/test_proxy_token_counter.py --ignore=tests/unit/proxy/test_proxy_utils.py --ignore=tests/unit/proxy/test_reducto_ocr_route.py --ignore=tests/unit/proxy/test_response_polling_pre_call_checks.py --ignore=tests/unit/proxy/test_server_root_path.py --ignore=tests/unit/proxy/test_ui_path_detection.py --ignore=tests/unit/proxy/test_unit_test_proxy_hooks.py --ignore=tests/unit/proxy/test_update_spend.py --ignore=tests/unit/proxy/test_zero_cost_model_budget_bypass.py --tb=short -vv -n 4 --durations=20
test-unit-integrations: install-test-deps
$(UV_RUN) pytest tests/unit/integrations --tb=short -vv -n 4 --durations=20

25
README.md
View file

@ -268,6 +268,31 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
</details>
<details>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
### Python SDK - Agents
```python
import litellm
from litellm import Harness, sandbox
result = litellm.agent(
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
"Find why tests/test_router.py is flaky and fix it.",
sandbox=sandbox.local("./repo"),
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
)
print(result.text, result.cost, [f.path for f in result.files])
```
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)
</details>
### Supported Providers ([Website Supported Models](https://models.litellm.ai/) | [Docs](https://docs.litellm.ai/docs/providers))
| Provider | `/chat/completions` | `/messages` | `/responses` | `/embeddings` | `/image/generations` | `/audio/transcriptions` | `/audio/speech` | `/moderations` | `/batches` | `/rerank` |

4
backend/routes/allowlist.py
View file

@ -81,7 +81,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
# Spend / analytics
"/spend/",
"/analytics/",
"/engine/",
"/lens/",
"/v1/traces",
"/global/",
"/user_agent",
@ -146,7 +146,7 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
{
"/",
"/routes",
"/engine",
"/lens",
"/openapi.json",
"/docs",
"/docs/oauth2-redirect",

4
deploy/lens/Dockerfile
View file

@ -1,6 +1,6 @@
FROM python:3.12-slim
WORKDIR /app
RUN pip install --no-cache-dir httpx==0.28.1 pydantic==2.11.7
COPY litellm/proxy/engine/__init__.py litellm/proxy/engine/models.py litellm/proxy/engine/trace_store.py litellm/proxy/engine/analysis.py litellm/proxy/engine/worker.py /app/engine/
COPY litellm/proxy/lens/__init__.py litellm/proxy/lens/models.py litellm/proxy/lens/trace_store.py litellm/proxy/lens/analysis.py litellm/proxy/lens/worker.py /app/lens/
USER 65532:65532
CMD ["python", "-m", "engine.worker"]
CMD ["python", "-m", "lens.worker"]

10
deploy/lens/Dockerfile.dockerignore
View file

@ -1,8 +1,8 @@
**
!litellm/
!litellm/proxy/
!litellm/proxy/engine/
!litellm/proxy/engine/__init__.py
!litellm/proxy/engine/models.py
!litellm/proxy/engine/analysis.py
!litellm/proxy/engine/worker.py
!litellm/proxy/lens/
!litellm/proxy/lens/__init__.py
!litellm/proxy/lens/models.py
!litellm/proxy/lens/analysis.py
!litellm/proxy/lens/worker.py

18
deploy/lens/README.md
View file

@ -74,7 +74,7 @@ V1 requires ClickHouse for both sources. It does not reconstruct sessions from u
The UI and API use the same scan lifecycle. Authenticate with a proxy administrator credential for writes, or a proxy-admin viewer credential for reads. Worker credentials are only for worker operations
```bash
curl "$LITELLM_URL/engine" -H "Authorization: Bearer $LITELLM_API_KEY" \
curl "$LITELLM_URL/lens" -H "Authorization: Bearer $LITELLM_API_KEY" \
-H 'Content-Type: application/json' -d '{
"name": "Research quality", "model": "your-model-alias",
"context": "Answer the requested question using cited, retrieved evidence.",
@ -83,14 +83,14 @@ curl "$LITELLM_URL/engine" -H "Authorization: Bearer $LITELLM_API_KEY" \
"enabled": true, "interval_minutes": 1440, "monthly_budget": 50
}'
curl "$LITELLM_URL/engine/$LENS_ID/runs" -X POST \
curl "$LITELLM_URL/lens/$LENS_ID/runs" -X POST \
-H "Authorization: Bearer $LITELLM_API_KEY" -H 'Content-Type: application/json' -d '{}'
curl "$LITELLM_URL/engine/$LENS_ID/runs?offset=0" -H "Authorization: Bearer $LITELLM_API_KEY"
curl "$LITELLM_URL/engine/$LENS_ID/runs/$BATCH_ID" -H "Authorization: Bearer $LITELLM_API_KEY"
curl "$LITELLM_URL/lens/$LENS_ID/runs?offset=0" -H "Authorization: Bearer $LITELLM_API_KEY"
curl "$LITELLM_URL/lens/$LENS_ID/runs/$BATCH_ID" -H "Authorization: Bearer $LITELLM_API_KEY"
```
Creation queues the first batch. Posting to `/engine/{id}/runs` queues another, or returns the existing active batch. The run response contains its ID under `jobs[0].id`. Poll the batch URL for status, findings and assessments. List responses omit large result payloads; request a batch to retrieve them. Supply an optional complete `settings` object on the runs POST for a one-off override; the saved lens stays unchanged. Selection accepts `team_id`, exact `filters`, and opaque `execution_ids` returned by `/engine/preview/sample`. Preview accepts `offset` and `as_of` to keep the time window fixed while paging. Feedback uses `PATCH /engine/{id}/findings/{finding_id}` with `status` and `reason`
Creation queues the first batch. Posting to `/lens/{id}/runs` queues another, or returns the existing active batch. The run response contains its ID under `jobs[0].id`. Poll the batch URL for status, findings and assessments. List responses omit large result payloads; request a batch to retrieve them. Supply an optional complete `settings` object on the runs POST for a one-off override; the saved lens stays unchanged. Selection accepts `team_id`, exact `filters`, and opaque `execution_ids` returned by `/lens/preview/sample`. Preview accepts `offset` and `as_of` to keep the time window fixed while paging. Feedback uses `PATCH /lens/{id}/findings/{finding_id}` with `status` and `reason`
## Quality evaluation
@ -107,3 +107,11 @@ Set `LITELLM_API_KEY` privately. This makes paid model calls. Inspect missed and
The worker uses temporary disk space for trace content while reviewing it, and removes those files after each review. The Docker command supplies a writable temporary mount while keeping the application filesystem read-only
To check that accepted behavior stays accepted without hiding new problems, run the evaluator with `--dataset tests/proxy_behavior/lens/feedback_cases.json`. Reports include elapsed time, model call count, reported cost when the proxy provides it, missed checks, unexpected checks, and inconclusive candidates
## Upgrading from the original Lens API
The Lens API now uses `/lens` instead of `/engine`, list responses use `lenses`, and worker claims use `lens_id`. Upgrade the proxy and recreate every worker with the image shown by the upgraded dashboard before starting new scans. Update API clients to the new paths and response fields. Old worker images cannot poll the renamed API
Stop workers and let active scans finish before upgrading. Deploy proxy instances together: older proxies cannot use the renamed database tables. The schema migration renames the three Lens tables and the run-history identifier column in place, preserving saved investigations, findings, history, worker credentials, and billing assignments. Existing migration files retain their original names and checksums
Upgrades using `--use_prisma_db_push` stop before schema changes if any legacy Lens table exists, preventing Prisma from dropping saved data. Apply `litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql` to the configured database schema before retrying. Deployments already using migration history can instead start without `--use_prisma_db_push` to apply the shipped migration normally. Fresh databases and databases already using the renamed tables can continue using database push

2
deploy/lens/compose.yaml
View file

@ -1,6 +1,6 @@
services:
lens-worker:
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:c41e932eaf3e4efbcaf8cc5027c7e93021e5b2823f21cb8785cd107e37b91c9a}
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:a8e8731d954916594eea462969946b9292fb771681ff515a9fd296b53f856c77}
environment:
LITELLM_URL: ${LITELLM_URL:?Set the URL reachable from this container}
LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:?Create a worker credential in the Lens UI}

BIN
deploy/lens/screenshots/after.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 95 KiB

BIN
deploy/lens/screenshots/before.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 6.9 KiB

BIN
deploy/lens/screenshots/finding.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 89 KiB

BIN
deploy/lens/screenshots/progress.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 80 KiB

BIN
deploy/lens/screenshots/setup.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 70 KiB

BIN
deploy/lens/screenshots/trace.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 132 KiB

BIN
deploy/lens/screenshots/worker-billing-after.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 59 KiB

BIN
deploy/lens/screenshots/worker-billing-before.png
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 54 KiB

2
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -616,7 +616,7 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
data["prompt"] = self.redact_text(prompt, source="prompt")
return 1
if isinstance(prompt, list):
data["prompt"] = [ # mutable-ok: data["prompt"] is a list on the wire
data["prompt"] = [
self.redact_text(item, source="prompt")
if isinstance(item, str) and item
else item

18
litellm-proxy-extras/litellm_proxy_extras/migration_lock.py
View file

@ -87,3 +87,21 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N
f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. "
f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
)
@contextmanager
def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
"""A session-level, non-blocking hold of the migration coordinator lock on an autocommit
connection, for DDL that cannot run inside a transaction (`CREATE INDEX CONCURRENTLY`).
Yields whether the lock was acquired; a v2 resolver or another migration job's index build
holding it yields False. Released on exit."""
from psycopg.rows import class_row
with connection.cursor(row_factory=class_row(_LockResult)) as cursor:
row: Final = cursor.execute("SELECT pg_try_advisory_lock(%s) AS acquired", (MIGRATION_LOCK_KEY,)).fetchone()
acquired: Final = row is not None and row.acquired
try:
yield acquired
finally:
if acquired:
connection.execute("SELECT pg_advisory_unlock(%s)", (MIGRATION_LOCK_KEY,))

46
litellm-proxy-extras/litellm_proxy_extras/migration_recovery.py
View file

@ -1,4 +1,5 @@
import hashlib
import re
import subprocess
from collections.abc import Mapping
from dataclasses import dataclass
@ -156,3 +157,48 @@ def baseline_current_schema(
"review any feature-specific backfill requirements.",
len(migrations),
)
_LINE_COMMENT_RE: Final = re.compile(r"--[^\n]*")
_BLOCK_COMMENT_RE: Final = re.compile(r"/\*.*?\*/", re.DOTALL)
_NO_OP_STATEMENT_RE: Final = re.compile(r"^\s*SELECT\s+1\s*$", re.IGNORECASE)
def is_inert_migration(script: str) -> bool:
"""Whether a migration file changes nothing: only comments and `SELECT 1`, so
applying it can neither repeat nor skip a database change."""
stripped: Final = _LINE_COMMENT_RE.sub("", _BLOCK_COMMENT_RE.sub("", script))
return all(not part.strip() or _NO_OP_STATEMENT_RE.match(part) for part in stripped.split(";"))
def roll_back_failed_inert_migration(coordinator: MigrationCoordinator, schema: str, migration: Path) -> bool:
"""Roll back the failed ledger row of a migration whose file in this build is inert,
so `migrate deploy` applies the inert file on its next pass. The row records an
earlier build's attempt at SQL this build no longer ships (an index now built by the
migration job), so no database change can be repeated or skipped by replaying
the empty file. The caller commits this checkpoint before the next Prisma command.
"""
from psycopg import sql
if not is_inert_migration(migration.read_text(encoding="utf-8")):
return False
coordinator.acquire_prisma_lock()
records: Final = _migration_records(coordinator.connection, schema, migration)
unfinished: Final = tuple(record for record in records if not record.finished)
if len(unfinished) != 1:
return False
result: Final = coordinator.connection.execute(
sql.SQL(
"UPDATE {} SET rolled_back_at = current_timestamp "
"WHERE id = %s AND finished_at IS NULL AND rolled_back_at IS NULL"
).format(sql.Identifier(schema, "_prisma_migrations")),
(unfinished[0].id,),
)
if result.rowcount != 1:
raise RuntimeError("Could not roll back the failed inert migration history row; rerun the database setup.")
logger.info(
"Rolled back the failed history row of %s: this build ships it as an inert migration, "
"its index is built by the migration job",
migration.parent.name,
)
return True

8
litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql
View file

@ -1,2 +1,6 @@
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");
-- The (api_key, startTime) index on LiteLLM_SpendLogs is built after migrate deploy,
-- through litellm_proxy_extras/request_log_indexes.py: concurrently on a plain table and
-- per partition on a partitioned one. The migration job builds it; a serving proxy that
-- ran the migrations itself builds it in the background once it serves. A migration
-- cannot do either without blocking spend-log writes or failing on a partitioned table.
SELECT 1;

18
litellm-proxy-extras/litellm_proxy_extras/migrations/20260831120001_spend_logs_litellm_call_id_index/migration.sql
View file

@ -1,12 +1,6 @@
-- CreateIndex (CONCURRENTLY)
--
-- Disclaimer:
-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
-- interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
-- - Do not edit this file after it has been applied to any database: Prisma checksums
-- migrations; add a new migration instead.
-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
-- without IF NOT EXISTS if you must support older versions).
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_SpendLogs_litellm_call_id_idx" ON "LiteLLM_SpendLogs"("litellm_call_id");
-- The litellm_call_id index on LiteLLM_SpendLogs is built after migrate deploy, through
-- litellm_proxy_extras/request_log_indexes.py: concurrently on a plain table and per
-- partition on a partitioned one. The migration job builds it; a serving proxy that ran
-- the migrations itself builds it in the background once it serves. Postgres refuses
-- CREATE INDEX CONCURRENTLY on a partitioned parent, so this migration no longer runs it.
SELECT 1;

18
litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql Normal file
View file

@ -0,0 +1,18 @@
DO $$
BEGIN
ALTER TABLE IF EXISTS "LiteLLM_Engine" RENAME TO "LiteLLM_Lens";
ALTER TABLE IF EXISTS "LiteLLM_EngineRun" RENAME TO "LiteLLM_LensRun";
ALTER TABLE IF EXISTS "LiteLLM_EngineWorker" RENAME TO "LiteLLM_LensWorker";
IF EXISTS (
SELECT 1 FROM pg_attribute
WHERE attrelid = to_regclass('"LiteLLM_LensRun"')
AND attname = 'engine_id' AND NOT attisdropped
) THEN
ALTER TABLE "LiteLLM_LensRun" RENAME COLUMN "engine_id" TO "lens_id";
END IF;
ALTER INDEX IF EXISTS "LiteLLM_Engine_pkey" RENAME TO "LiteLLM_Lens_pkey";
ALTER INDEX IF EXISTS "LiteLLM_EngineRun_pkey" RENAME TO "LiteLLM_LensRun_pkey";
ALTER INDEX IF EXISTS "LiteLLM_EngineWorker_pkey" RENAME TO "LiteLLM_LensWorker_pkey";
ALTER INDEX IF EXISTS "LiteLLM_EngineWorker_token_hash_key" RENAME TO "LiteLLM_LensWorker_token_hash_key";
ALTER INDEX IF EXISTS "LiteLLM_EngineRun_engine_id_created_at_idx" RENAME TO "LiteLLM_LensRun_lens_id_created_at_idx";
END $$;

448
litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py Normal file
View file

@ -0,0 +1,448 @@
"""The request-log indexes built after `prisma migrate deploy` instead of by a migration:
by the migration job, or by a serving proxy that ran the migrations itself (in the
background, once it serves).
A migration cannot build them: a plain `CREATE INDEX` blocks spend-log inserts for the
whole build, and `CREATE INDEX CONCURRENTLY` is refused on a partitioned parent
(db_scripts/partition_spend_logs.sql). `REQUEST_LOG_INDEXES` is the one list to extend;
names match what Prisma derives from the `@@index` declarations in schema.prisma, so an
index a database already has is recognized and never rebuilt.
"""
import hashlib
import random
import re
import time
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.migration_lock import held_migration_lock
if TYPE_CHECKING:
import psycopg
from psycopg import sql
@dataclass(frozen=True, slots=True)
class RequestLogIndex:
"""One index the migration job owns: the table, the exact Prisma index name and the
column list as it would be written after `ON <table>`."""
table: str
name: str
definition: str
@property
def columns(self) -> tuple[str, ...]:
return tuple(re.findall(r'"([^"]+)"', self.definition))
def partition_index_name(self, partition: str) -> str:
"""The child index name for one partition, built the way Postgres names the
children of a partitioned index, and kept within the 63 byte identifier limit."""
name: Final = f"{partition}_{self.name.removeprefix(f'{self.table}_')}"
if len(name.encode()) <= _IDENTIFIER_MAX_BYTES:
return name
digest: Final = hashlib.sha256(name.encode()).hexdigest()[:_DIGEST_LENGTH]
budget: Final = _IDENTIFIER_MAX_BYTES - _DIGEST_LENGTH - 1
kept: Final = next(name[:length] for length in range(len(name), 0, -1) if len(name[:length].encode()) <= budget)
return f"{kept}_{digest}"
REQUEST_LOG_INDEXES: Final = (
RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'),
RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'),
)
_IDENTIFIER_MAX_BYTES: Final = 63
_PARENT_LOCK_TIMEOUT: Final = "2s"
_PARENT_LOCK_ATTEMPTS: Final = 30
_LOCK_HANDOVER_SECONDS: Final = 2.0
_DIGEST_LENGTH: Final = 8
_CREATE_INDEX_STATEMENT: Final = re.compile(
r'^\s*CREATE\s+(?:UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?"(?P<index>[^"]+)"\s+ON\b',
re.IGNORECASE,
)
_TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)"
_CHILDREN_WITHOUT_THE_INDEX_SQL: Final = (
"SELECT child.relname AS name, n.nspname AS schema, child.relkind = 'p' AS partitioned "
"FROM pg_inherits i JOIN pg_class child ON child.oid = i.inhrelid "
"JOIN pg_namespace n ON n.oid = child.relnamespace "
"WHERE i.inhparent = to_regclass(%s) AND NOT EXISTS ("
"SELECT 1 FROM pg_inherits attached JOIN pg_index x ON x.indexrelid = attached.inhrelid "
"WHERE attached.inhparent = to_regclass(%s) AND x.indrelid = child.oid) "
"ORDER BY child.relname"
)
_EQUIVALENT_INDEXES_SQL: Final = (
"SELECT i.relname AS name, x.indisvalid AS valid "
"FROM pg_index x JOIN pg_class i ON i.oid = x.indexrelid JOIN pg_am am ON am.oid = i.relam "
"WHERE x.indrelid = to_regclass(%s) AND i.relname <> %s AND am.amname = 'btree' AND NOT x.indisunique "
"AND x.indexprs IS NULL AND x.indpred IS NULL AND x.indnkeyatts = x.indnatts "
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indoption::int2[]) o WHERE o <> 0) "
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indclass::oid[]) c JOIN pg_opclass oc ON oc.oid = c WHERE NOT oc.opcdefault) "
"AND NOT EXISTS (SELECT 1 FROM unnest(x.indcollation::oid[]) WITH ORDINALITY c(coll, ord) "
"JOIN unnest(x.indkey::int2[]) WITH ORDINALITY k(attnum, ord) ON k.ord = c.ord "
"JOIN pg_attribute a ON a.attrelid = x.indrelid AND a.attnum = k.attnum "
"WHERE c.coll <> 0 AND c.coll <> a.attcollation) "
"AND (SELECT array_agg(a.attname::text ORDER BY k.ord) FROM unnest(x.indkey::int2[]) WITH ORDINALITY k(attnum, ord) "
"JOIN pg_attribute a ON a.attrelid = x.indrelid AND a.attnum = k.attnum) = %s::text[] "
"AND NOT EXISTS (SELECT 1 FROM pg_inherits WHERE inhrelid = x.indexrelid) "
"ORDER BY x.indisvalid DESC, i.relname"
)
_INDEX_STATE_SQL: Final = (
'SELECT x.indisvalid AS valid, t.relname AS "table" '
"FROM pg_index x JOIN pg_class t ON t.oid = x.indrelid WHERE x.indexrelid = to_regclass(%s)"
)
@dataclass(frozen=True, slots=True)
class _Relation:
name: str
schema: str
partitioned: bool
@dataclass(frozen=True, slots=True)
class _IndexState:
valid: bool
table: str
@dataclass(frozen=True, slots=True)
class _EquivalentIndex:
name: str
valid: bool
@dataclass(frozen=True, slots=True)
class _TableKind:
partitioned: bool
def filter_request_log_index_diff(diff_sql: str, indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES) -> str:
"""The `prisma migrate diff` script without the statements that create a migration-job-owned
index, which the schema declares and the migrations deliberately do not build."""
names: Final = frozenset(index.name for index in indexes)
statements: Final = diff_sql.split(";")
kept: Final = tuple(statement for statement in statements if not _creates_one_of(statement, names))
return ";".join(kept) if any(part.strip() for part in kept) else ""
def _creates_one_of(statement: str, names: frozenset[str]) -> bool:
match: Final = _CREATE_INDEX_STATEMENT.match(_without_comments(statement))
return match is not None and match["index"] in names
def _without_comments(statement: str) -> str:
return "\n".join(line for line in statement.splitlines() if not line.lstrip().startswith("--"))
def _connect(database_url: str) -> "psycopg.Connection[tuple[object, ...]]":
import psycopg
return psycopg.connect(database_url, connect_timeout=10, autocommit=True)
def ensure_request_log_indexes(
database_url: str,
schema: str,
indexes: tuple[RequestLogIndex, ...] = REQUEST_LOG_INDEXES,
connect: "Callable[[str], psycopg.Connection[tuple[object, ...]]]" = _connect,
) -> bool:
"""Build every listed index that is missing or invalid. Each build step runs under
the migration coordinator lock, held per statement so a resolver booting on another
replica gets in between partitions rather than waiting for the whole table. Any
failure is logged and left for the next index build; the result says whether
every index ended up valid. Never raises."""
import psycopg
try:
with connect(database_url) as connection:
connection.execute("SET statement_timeout = 0")
results: Final = tuple(_ensure_index(connection, schema, index) for index in indexes)
except psycopg.Error as exc:
logger.warning("Could not build the request-log indexes, leaving them for the next index build: %s", exc)
return False
if not all(results):
logger.warning("Some request-log indexes are not in place yet, leaving them for the next index build")
return False
logger.info("Request-log indexes are all in place")
return True
def _under_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]", step: Callable[[], bool]) -> bool:
with held_migration_lock(connection) as held:
if not held:
logger.info(
"Another process holds the migration lock, leaving the request-log indexes to the next index build"
)
return False
return step()
def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex) -> bool:
from psycopg.rows import class_row
with connection.cursor(row_factory=class_row(_TableKind)) as cursor:
table: Final = cursor.execute(_TABLE_KIND_SQL, (_regclass_name(connection, schema, index.table),)).fetchone()
if table is None:
logger.info("Table %s does not exist yet, skipping index %s", index.table, index.name)
return True
if table.partitioned:
return build_index_on_partitioned_table(connection, schema, index)
return _build_leaf_index(connection, schema, index.table, index.name, index)
def _regclass_name(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, name: str) -> str:
from psycopg import sql
return sql.Identifier(schema, name).as_string(connection)
def _create_index_statement(
connection: "psycopg.Connection[tuple[object, ...]]", prefix: "sql.Composed", definition: str
) -> bytes:
return (prefix.as_string(connection) + definition).encode()
def _index_state(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: str) -> "_IndexState | None":
from psycopg.rows import class_row
with connection.cursor(row_factory=class_row(_IndexState)) as cursor:
return cursor.execute(_INDEX_STATE_SQL, (_regclass_name(connection, schema, index),)).fetchone()
def _equivalent_indexes(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
table: str,
name: str,
index: RequestLogIndex,
) -> tuple[_EquivalentIndex, ...]:
"""The indexes on `table` other than `name` with the same definition: default btree
over the same columns in the same order, no expression, predicate, DESC or custom
opclass or collation, and not attached under a partitioned index. Valid ones first."""
from psycopg.rows import class_row
with connection.cursor(row_factory=class_row(_EquivalentIndex)) as cursor:
return tuple(
cursor.execute(
_EQUIVALENT_INDEXES_SQL, (_regclass_name(connection, schema, table), name, list(index.columns))
).fetchall()
)
def _adopt_equivalent_index(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
table: str,
name: str,
index: RequestLogIndex,
) -> bool:
"""Rename a valid index of the same definition under another name (an operator's
hand-built copy, say) to the name this code expects, instead of building a second
one. RENAME on an index is a catalog change that lets writes through."""
from psycopg import sql
equivalent: Final = next(
(found for found in _equivalent_indexes(connection, schema, table, name, index) if found.valid), None
)
if equivalent is None:
return False
logger.info(
"Renaming the equivalent index %s on %s to %s instead of building a second one", equivalent.name, table, name
)
connection.execute(
sql.SQL("ALTER INDEX {} RENAME TO {}").format(sql.Identifier(schema, equivalent.name), sql.Identifier(name))
)
return True
def _report_second_copies(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
table: str,
name: str,
index: RequestLogIndex,
concurrently: bool,
) -> None:
"""Log every other index of the same definition with the statement that removes it.
Dropping is the operator's call: a second copy costs writes and disk, never results."""
from psycopg import sql
drop: Final = "DROP INDEX CONCURRENTLY" if concurrently else "DROP INDEX"
for copy in _equivalent_indexes(connection, schema, table, name, index):
logger.warning(
"Index %s on %s is a second copy of %s and only costs writes and disk; remove it with: %s %s",
copy.name,
table,
name,
drop,
sql.Identifier(schema, copy.name).as_string(connection),
)
def _children_without_the_index(
connection: "psycopg.Connection[tuple[object, ...]]", schema: str, table: str, index: str
) -> tuple[_Relation, ...]:
from psycopg.rows import class_row
with connection.cursor(row_factory=class_row(_Relation)) as cursor:
return tuple(
cursor.execute(
_CHILDREN_WITHOUT_THE_INDEX_SQL,
(_regclass_name(connection, schema, table), _regclass_name(connection, schema, index)),
).fetchall()
)
def _build_leaf_index(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
table: str,
name: str,
index: RequestLogIndex,
) -> bool:
"""Build one plain table's or partition's index with CONCURRENTLY so writes keep
flowing. The catalog is read under the migration lock, so a replica that saw an
invalid index before the lock finds the valid one another replica just built and
leaves it. An invalid index left by an interrupted build is dropped and rebuilt; a
valid index of the same definition under another name is renamed rather than
duplicated; an index of that name on another table is a collision this code will
not touch."""
from psycopg import sql
def build() -> bool:
existing: Final = _index_state(connection, schema, name)
if existing is not None and existing.table != table:
logger.warning(
"Index %s already exists on %s rather than %s, leaving it alone", name, existing.table, table
)
return False
if existing is not None and existing.valid:
return True
if existing is not None:
logger.info("Dropping the invalid index %s left by an interrupted build on %s", name, table)
connection.execute(sql.SQL("DROP INDEX CONCURRENTLY {}").format(sql.Identifier(schema, name)))
elif _adopt_equivalent_index(connection, schema, table, name, index):
return True
logger.info("Building index %s on %s concurrently", name, table)
prefix: Final = sql.SQL("CREATE INDEX CONCURRENTLY IF NOT EXISTS {} ON {} ").format(
sql.Identifier(name), sql.Identifier(schema, table)
)
connection.execute(_create_index_statement(connection, prefix, index.definition))
built: Final = _index_state(connection, schema, name)
return built is not None and built.valid
current: Final = _index_state(connection, schema, name)
if current is None or not current.valid or current.table != table:
if not _under_migration_lock(connection, build):
return False
time.sleep(_LOCK_HANDOVER_SECONDS)
_report_second_copies(connection, schema, table, name, index, concurrently=True)
return True
def build_index_on_partitioned_table(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
index: RequestLogIndex,
table: "str | None" = None,
name: "str | None" = None,
) -> bool:
"""Build the index the way Postgres allows on a partitioned parent: a metadata-only
parent index ON ONLY the parent, one CONCURRENTLY build per partition, and ATTACH
PARTITION for each child. Partitions that are themselves partitioned get the same
treatment one level down. Every step checks the catalog before acting, so an
interrupted run resumes where it stopped and a second run finds nothing to do; a
parent or child index of the same definition under another name is renamed and
used rather than duplicated. The connection must be in autocommit mode. True when
the parent index ends up valid."""
parent_table: Final = index.table if table is None else table
parent_index: Final = index.name if name is None else name
existing: Final = _index_state(connection, schema, parent_index)
if existing is not None and existing.table != parent_table:
logger.warning(
"Index %s already exists on %s rather than %s, leaving it alone", parent_index, existing.table, parent_table
)
return False
if existing is None and not _under_migration_lock(
connection,
lambda: (
_adopt_equivalent_index(connection, schema, parent_table, parent_index, index)
or _create_parent_index(connection, schema, parent_index, parent_table, index)
),
):
return False
children: Final = _children_without_the_index(connection, schema, parent_table, parent_index)
if not all(_attach_child_index(connection, schema, parent_index, child, index) for child in children):
return False
final: Final = _index_state(connection, schema, parent_index)
if final is None or not final.valid:
return False
_report_second_copies(connection, schema, parent_table, parent_index, index, concurrently=False)
return True
def _create_parent_index(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
name: str,
table: str,
index: RequestLogIndex,
) -> bool:
"""Create the metadata-only parent index. Postgres takes a SHARE lock on the
parent for that statement, so it waits for in-flight writes and queues new ones
behind it; a short lock_timeout with retries keeps every such pause bounded."""
import psycopg
from psycopg import sql
prefix: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format(
sql.Identifier(name), sql.Identifier(schema, table)
)
statement: Final = _create_index_statement(connection, prefix, index.definition)
connection.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(_PARENT_LOCK_TIMEOUT)))
try:
for _ in range(_PARENT_LOCK_ATTEMPTS):
try:
connection.execute(statement)
return True
except psycopg.errors.LockNotAvailable:
logger.info("Waiting for in-flight writes to %s before creating the parent index %s", table, name)
time.sleep(random.uniform(0.1, 0.5))
finally:
connection.execute("SET lock_timeout = 0")
logger.warning("Could not get the parent lock on %s to create %s, leaving it for the next index build", table, name)
return False
def _attach_child_index(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
parent_index: str,
child: _Relation,
index: RequestLogIndex,
) -> bool:
from psycopg import sql
child_index: Final = index.partition_index_name(child.name)
built: Final = (
build_index_on_partitioned_table(connection, child.schema, index, child.name, child_index)
if child.partitioned
else _build_leaf_index(connection, child.schema, child.name, child_index, index)
)
if not built:
return False
def attach() -> bool:
connection.execute(
sql.SQL("ALTER INDEX {} ATTACH PARTITION {}").format(
sql.Identifier(schema, parent_index), sql.Identifier(child.schema, child_index)
)
)
logger.info("Attached index %s on partition %s to %s", child_index, child.name, parent_index)
return True
return _under_migration_lock(connection, attach)

10
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1895,22 +1895,22 @@ model LiteLLM_WorkflowMessage {
@@index([run_id])
}
model LiteLLM_Engine {
model LiteLLM_Lens {
id String @id
version Int @default(0)
data Json
}
model LiteLLM_EngineRun {
model LiteLLM_LensRun {
id String @id
engine_id String
lens_id String
created_at DateTime
data Json
@@index([engine_id, created_at])
@@index([lens_id, created_at])
}
model LiteLLM_EngineWorker {
model LiteLLM_LensWorker {
id String @id
token_hash String @unique
data Json

211
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -5,6 +5,7 @@ import re
import shutil
import subprocess
import tempfile
import threading
import time
from collections.abc import Callable
from dataclasses import dataclass, replace
@ -13,6 +14,7 @@ from typing import TYPE_CHECKING, Final, Optional
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@ -24,6 +26,7 @@ from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.request_log_indexes import ensure_request_log_indexes, filter_request_log_index_diff
if TYPE_CHECKING:
import psycopg
@ -433,6 +436,21 @@ class ProxyExtrasDBManager:
return True
return False
@staticmethod
def _filter_migration_job_owned_drift(diff_sql: str, partitioned: bool | None = None) -> str:
"""The drift script without the indexes the migration job builds (the schema
declares them, the migrations deliberately do not) and, when LiteLLM_SpendLogs
is partitioned, without its primary-key rewrite and partitioning artifacts."""
without_indexes: Final = filter_request_log_index_diff(diff_sql)
is_partitioned: Final = ProxyExtrasDBManager.spend_logs_is_partitioned() if partitioned is None else partitioned
if not is_partitioned:
return without_indexes
logger.info(
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
"rewrite and partitioning artifacts from the drift script"
)
return filter_partitioned_spend_logs_diff(without_indexes)
@staticmethod
def _resolve_all_migrations(
migrations_dir: str, schema_path: str, mark_all_applied: bool = True
@ -513,21 +531,14 @@ class ProxyExtrasDBManager:
return
logger.info(f"Migration diff created at {diff_sql_path}")
if ProxyExtrasDBManager.spend_logs_is_partitioned():
filtered_sql = filter_partitioned_spend_logs_diff(
diff_sql_path.read_text()
)
diff_sql_path.write_text(filtered_sql)
logger.info(
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
"rewrite and partitioning artifacts from the drift script"
)
if not filtered_sql.strip():
logger.info("Drift script is empty after filtering; nothing to apply")
if not mark_all_applied:
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
filtered_sql: Final = ProxyExtrasDBManager._filter_migration_job_owned_drift(diff_sql_path.read_text())
diff_sql_path.write_text(filtered_sql)
if not filtered_sql.strip():
logger.info("Drift script is empty after filtering; nothing to apply")
if not mark_all_applied:
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
return
# 2. Run prisma db execute to apply the migration
applied_ok = False
@ -590,6 +601,36 @@ class ProxyExtrasDBManager:
f"Failed to resolve migration {migration_name}: {e.stderr}"
)
@staticmethod
def raise_if_lens_rename_pending() -> None:
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return
try:
import psycopg
except ImportError as exc:
raise RuntimeError("Install psycopg to verify Lens data safety before prisma db push.") from exc
try:
with psycopg.connect(
ProxyExtrasDBManager._strip_prisma_query_params(database_url), connect_timeout=10, autocommit=True
) as connection:
legacy: Final = connection.execute(
"SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid=c.relnamespace "
"WHERE n.nspname=%s AND c.relname IN ('LiteLLM_Engine', 'LiteLLM_EngineRun', 'LiteLLM_EngineWorker') "
"LIMIT 1",
(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",),
).fetchone()
except psycopg.Error as exc:
raise RuntimeError(
"Cannot verify Lens data safety; refusing prisma db push. Check database connectivity and psycopg installation."
) from exc
if legacy is not None:
raise RuntimeError(
"Legacy Lens tables exist. prisma db push would drop saved Lens data. "
"Apply the shipped 20261001100000_rename_lens migration to this database schema before retrying. "
"Deployments using migration history can upgrade without --use_prisma_db_push instead."
)
@staticmethod
def spend_logs_is_partitioned() -> bool:
"""True when the connected database's LiteLLM_SpendLogs is a
@ -770,7 +811,7 @@ class ProxyExtrasDBManager:
conn.execute(statement)
except psycopg.Error as e:
logger.warning(
"Could not repair invalid index %s.%s, will retry on the next startup. "
"Could not repair invalid index %s.%s, will retry on the next database setup run. "
"If this keeps happening, run `%s` by hand as the index owner. Error: %s",
index.schema,
index.name,
@ -781,16 +822,21 @@ class ProxyExtrasDBManager:
logger.info("%s invalid index %s.%s", action, index.schema, index.name)
@staticmethod
def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
def repair_invalid_indexes(
lock_timeout: str = "30s",
repair: "Callable[[psycopg.Connection[tuple[str, str, str]], _InvalidIndex], None] | None" = None,
) -> bool:
"""Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
INVALID (a migration deadlock between replicas is the usual cause; the
retried migration skips them because of IF NOT EXISTS). Never raises:
returns True when no invalid index remains, False when the repair was
skipped or failed and will be retried on the next startup. Looks in the
skipped or failed and will be retried on the next database setup run. Looks in the
schema DATABASE_URL names, the only URL Prisma migrates through, but
connects over DIRECT_URL when set: the session settings, the advisory
lock and REINDEX CONCURRENTLY all need one server session, which a
transaction pooler does not give."""
transaction pooler does not give. Each rebuild holds the migration
coordinator lock on its own, like the migration job's index build, so a resolver
booting on another replica waits for one index at most."""
prisma_url: Final = os.getenv("DATABASE_URL")
if not prisma_url:
return False
@ -826,20 +872,53 @@ class ProxyExtrasDBManager:
if lock_row is None or not lock_row[0]:
logger.info("Another replica is already rebuilding the invalid indexes, skipping")
return False
for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
ProxyExtrasDBManager._repair_index(conn, index)
repair_one: Final = repair or ProxyExtrasDBManager._repair_index
repaired: Final = all(
ProxyExtrasDBManager._repair_under_migration_lock(conn, schema, index, repair_one)
for index in found
)
if not repaired:
return False
remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
except psycopg.Error as e:
logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
logger.warning(
"Could not check for invalid indexes, will retry on the next database setup run. Error: %s", e
)
return False
return not remaining
@staticmethod
def _repair_under_migration_lock(
conn: "psycopg.Connection[tuple[str, str, str]]",
schema: str,
index: _InvalidIndex,
repair: "Callable[[psycopg.Connection[tuple[str, str, str]], _InvalidIndex], None]",
) -> bool:
"""Rebuild one index under the migration coordinator lock, skipping it when a
migration job finished or dropped it in the meantime. False when another process
holds the lock, so the check waits for the next database setup run."""
with held_migration_lock(conn) as held:
if not held:
logger.info(
"Another process is building indexes under the migration lock, leaving the "
"invalid index check to the next database setup run"
)
return False
still_invalid: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
if any(found.schema == index.schema and found.name == index.name for found in still_invalid):
repair(conn, index)
return True
@staticmethod
def _setup_database_v2(use_migrate: bool) -> bool:
if not use_migrate:
return ProxyExtrasDBManager._run_database_v2(False)
from litellm_proxy_extras.migration_lock import migration_environment, migration_lock
from litellm_proxy_extras.migration_recovery import baseline_current_schema, recover_completed_migration
from litellm_proxy_extras.migration_recovery import (
baseline_current_schema,
recover_completed_migration,
roll_back_failed_inert_migration,
)
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
@ -854,7 +933,9 @@ class ProxyExtrasDBManager:
if not migration.is_file():
return False
with migration_lock(lock_url) as coordinator:
return recover_completed_migration(coordinator, schema, migration)
return recover_completed_migration(coordinator, schema, migration) or roll_back_failed_inert_migration(
coordinator, schema, migration
)
def baseline_existing(migrations_dir: str) -> None:
with migration_lock(lock_url) as coordinator:
@ -895,6 +976,7 @@ class ProxyExtrasDBManager:
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
if not use_migrate:
ProxyExtrasDBManager.raise_if_lens_rename_pending()
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
original_dir = os.getcwd()
@ -1146,13 +1228,16 @@ class ProxyExtrasDBManager:
)
@staticmethod
def setup_database(
use_migrate: bool = False, use_v2_resolver: bool = False
) -> bool:
def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool:
"""
Set up the database using either prisma migrate or prisma db push
Uses migrations from litellm-proxy-extras package
The request-log indexes in `REQUEST_LOG_INDEXES` are not built here: the
migration job builds them through `run_migration_job`, and a serving proxy that
ran the migrations itself starts them through `start_request_log_index_build`
once it is ready to serve.
Args:
use_migrate: Whether to use prisma migrate instead of db push
use_v2_resolver: Opt into the v2 migration resolver (safer during
@ -1169,10 +1254,48 @@ class ProxyExtrasDBManager:
migrated = ProxyExtrasDBManager._run_migrations(
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
if migrated:
ProxyExtrasDBManager.repair_invalid_indexes()
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return migrated
if not migrated:
return False
ProxyExtrasDBManager.repair_invalid_indexes()
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return True
@staticmethod
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
"""Build the indexes in `REQUEST_LOG_INDEXES` on the writer, in the schema the
migrations target. Idempotent and never raises; False when an index is still
missing or invalid, so the migration job reports it and gets rerun instead of
leaving the table unindexed until the next deploy."""
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return True
direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
os.environ.get("DIRECT_URL") or database_url
)
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
return build(direct_url, schema)
@staticmethod
def run_migration_job(
use_migrate: bool = False,
use_v2_resolver: bool = False,
setup: Callable[[bool, bool], bool] = setup_database,
build: Callable[[], bool] = build_request_log_indexes,
) -> bool:
"""The migration job's whole run: `setup_database`, then the request-log indexes,
built synchronously so the job exits only once they are in place. False when the
migrations failed or an index could not be built, so the Job is rerun."""
return setup(use_migrate, use_v2_resolver) and build()
@staticmethod
def start_request_log_index_build(build: Callable[[], bool] = build_request_log_indexes) -> threading.Thread:
"""A serving proxy that ran the migrations itself (schema updates not disabled)
builds the request-log indexes on a daemon thread, so a long build never delays
readiness. A build that could not finish is logged and picked up by the next boot
or the migration job."""
thread: Final = threading.Thread(target=build, name="litellm-request-log-indexes", daemon=True)
thread.start()
return thread
@staticmethod
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
@ -1216,15 +1339,16 @@ class ProxyExtrasDBManager:
logger.info("✅ Post-migration sanity check completed")
return True
except subprocess.CalledProcessError as e:
logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}")
if "P3009" in e.stderr:
stderr: Final = str(e.stderr or "")
logger.info(f"prisma db error: {stderr}, e: {e.stdout}")
if "P3009" in stderr:
# Extract the failed migration name from the error message
migration_match = re.search(
r"`(\d+_.*)` migration", e.stderr
r"`(\d+_.*)` migration", stderr
)
if migration_match:
failed_migration = migration_match.group(1)
if ProxyExtrasDBManager._is_idempotent_error(e.stderr):
if ProxyExtrasDBManager._is_idempotent_error(stderr):
logger.info(
f"Migration {failed_migration} failed due to idempotent error (e.g., column already exists), resolving as applied"
)
@ -1280,8 +1404,8 @@ class ProxyExtrasDBManager:
f"✅ Migration {failed_migration} marked as rolled back... retrying"
)
elif (
"P3005" in e.stderr
and "database schema is not empty" in e.stderr
"P3005" in stderr
and "database schema is not empty" in stderr
):
logger.info(
"Database schema is not empty, creating baseline migration. In read-only file system, please set an environment variable `LITELLM_MIGRATION_DIR` to a writable directory to enable migrations. Learn more - https://docs.litellm.ai/docs/proxy/prod#read-only-file-system"
@ -1295,13 +1419,13 @@ class ProxyExtrasDBManager:
)
logger.info("✅ All migrations resolved.")
return True
elif "P3018" in e.stderr:
elif "P3018" in stderr:
# Check if this is a permission error or idempotent error
if ProxyExtrasDBManager._is_permission_error(e.stderr):
if ProxyExtrasDBManager._is_permission_error(stderr):
# Permission errors should NOT be marked as applied
# Extract migration name for logging
migration_match = re.search(
r"Migration name: (\d+_.*)", e.stderr
r"Migration name: (\d+_.*)", stderr
)
migration_name = (
migration_match.group(1)
@ -1311,7 +1435,7 @@ class ProxyExtrasDBManager:
logger.error(
f"❌ Migration {migration_name} failed due to insufficient permissions. "
f"Please check database user privileges. Error: {e.stderr}"
f"Please check database user privileges. Error: {stderr}"
)
# Mark as rolled back and exit with error
@ -1334,7 +1458,7 @@ class ProxyExtrasDBManager:
f"was NOT applied. Please grant necessary database permissions and retry."
) from e
elif ProxyExtrasDBManager._is_idempotent_error(e.stderr):
elif ProxyExtrasDBManager._is_idempotent_error(stderr):
# Idempotent errors mean the migration has effectively been applied
logger.info(
"Migration failed due to idempotent error (e.g., column already exists), "
@ -1342,7 +1466,7 @@ class ProxyExtrasDBManager:
)
# Extract the migration name from the error message
migration_match = re.search(
r"Migration name: (\d+_.*)", e.stderr
r"Migration name: (\d+_.*)", stderr
)
if migration_match:
migration_name = migration_match.group(1)
@ -1391,13 +1515,14 @@ class ProxyExtrasDBManager:
logger.warning(
f"P3018 error encountered but could not classify "
f"as permission or idempotent error. "
f"Error: {e.stderr}"
f"Error: {stderr}"
)
raise
else:
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
# Use prisma db push with increased timeout
ProxyExtrasDBManager.raise_if_lens_rename_pending()
prisma_toolchain.run_prisma(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
timeout=prisma_command_timeout(),

4
litellm-proxy-extras/pyproject.toml
View file

@ -4,6 +4,10 @@ version = "0.4.103"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
dependencies = [
"psycopg>=3.2,<4.0",
"psycopg-binary>=3.2,<4.0",
]
license = "MIT"
license-files = ["LICENSE"]
authors = [

24
litellm-rust/Cargo.lock generated
View file

@ -4086,9 +4086,11 @@ dependencies = [
"litellm-secrets",
"litellm-secrets-aws",
"litellm-secrets-types",
"litellm-storage-clickhouse",
"litellm-token-counter",
"litellm-traces",
"litellm-tracing",
"prost",
"pyo3",
"pyo3-async-runtimes",
"qdrant-client",
@ -4288,6 +4290,20 @@ dependencies = [
"veil",
]
[[package]]
name = "litellm-storage-clickhouse"
version = "0.1.0"
dependencies = [
"flate2",
"litellm-http",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"url",
]
[[package]]
name = "litellm-testkit"
version = "0.1.0"
@ -4369,19 +4385,22 @@ name = "litellm-traces"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"flate2",
"litellm-http",
"litellm-storage-clickhouse",
"opentelemetry-proto",
"prost",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"strum",
"testcontainers-modules",
"thiserror 2.0.19",
"time",
"tokio",
"url",
"wiremock",
]
[[package]]
@ -4804,6 +4823,7 @@ dependencies = [
"js-sys",
"pin-project-lite",
"thiserror 2.0.19",
"tracing",
]
[[package]]
@ -4818,6 +4838,8 @@ dependencies = [
"opentelemetry_sdk 0.33.0",
"prost",
"serde",
"tonic",
"tonic-prost",
]
[[package]]

4
litellm-rust/Cargo.toml
View file

@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" }
litellm-router = { path = "crates/router" }
litellm-tracing = { path = "crates/tracing" }
litellm-traces = { path = "crates/traces" }
litellm-storage-clickhouse = { path = "crates/storage-clickhouse" }
litellm-core = { path = "crates/core" }
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
litellm-gateway = { path = "crates/gateway" }
@ -81,6 +82,7 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "mul
qdrant-client = { version = "1.19.0", default-features = false }
uuid = { version = "1", features = ["v4"] }
rstest = "0.26.1"
wiremock = "0.6.5"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustify = "=0.7.0"
@ -115,6 +117,8 @@ time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"
veil = "0.3.0"
prost = "0.14.4"
opentelemetry-proto = "0.33"
[profile.release]
opt-level = 3

2
litellm-rust/crates/cache-azure-blob/Cargo.toml
View file

@ -26,4 +26,4 @@ litellm-cache-testing.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
wiremock = "0.6.5"
wiremock.workspace = true

2
litellm-rust/crates/cache-gcs/Cargo.toml
View file

@ -21,4 +21,4 @@ litellm-cache-testing.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true

2
litellm-rust/crates/cache-response/Cargo.toml
View file

@ -21,4 +21,4 @@ redis = "1.7.0"
redis-test = "1.0.4"
rstest.workspace = true
tokio.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true

2
litellm-rust/crates/cache-s3/Cargo.toml
View file

@ -23,6 +23,6 @@ tokio.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-testing.workspace = true
rstest.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

2
litellm-rust/crates/core/Cargo.toml
View file

@ -47,4 +47,4 @@ litellm-host-native.workspace = true
litellm-llms = { workspace = true, features = ["test-support"] }
rstest.workspace = true
rstest_reuse.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true

2
litellm-rust/crates/gateway-inference/Cargo.toml
View file

@ -30,4 +30,4 @@ futures-util.workspace = true
tokio = { workspace = true, features = ["io-util"] }
rstest.workspace = true
tower = { version = "0.5.3", features = ["util"] }
wiremock = "0.6.5"
wiremock.workspace = true

57
litellm-rust/crates/host-python/src/conversion_cache.rs Normal file
View file

@ -0,0 +1,57 @@
use std::collections::{HashMap, hash_map::Entry};
use pyo3::prelude::*;
pub struct ToPythonCache<'a, 'py, T> {
entries: HashMap<usize, (&'a T, Bound<'py, PyAny>)>,
}
impl<T> Default for ToPythonCache<'_, '_, T> {
fn default() -> Self {
Self {
entries: HashMap::new(),
}
}
}
impl<'a, 'py, T> ToPythonCache<'a, 'py, T> {
pub fn get_or_try_insert_with(
&mut self,
value: &'a T,
convert: impl FnOnce(&'a T) -> PyResult<Bound<'py, PyAny>>,
) -> PyResult<&Bound<'py, PyAny>> {
let identity = std::ptr::from_ref(value) as usize;
let entry = match self.entries.entry(identity) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert((value, convert(value)?)),
};
Ok(&entry.1)
}
}
pub struct FromPythonCache<'py, T> {
entries: HashMap<usize, (Bound<'py, PyAny>, T)>,
}
impl<T> Default for FromPythonCache<'_, T> {
fn default() -> Self {
Self {
entries: HashMap::new(),
}
}
}
impl<'py, T> FromPythonCache<'py, T> {
pub fn get_or_try_insert_with(
&mut self,
value: &Bound<'py, PyAny>,
convert: impl FnOnce(&Bound<'py, PyAny>) -> PyResult<T>,
) -> PyResult<&T> {
let identity = value.as_ptr() as usize;
let entry = match self.entries.entry(identity) {
Entry::Occupied(entry) => entry.into_mut(),
Entry::Vacant(entry) => entry.insert((value.clone(), convert(value)?)),
};
Ok(&entry.1)
}
}

2
litellm-rust/crates/host-python/src/lib.rs
View file

@ -5,6 +5,7 @@
mod argument;
mod binding;
mod conversion_cache;
mod driver;
mod error;
mod file_reader;
@ -20,6 +21,7 @@ mod services;
pub use argument::lookup;
pub use binding::PythonBinding;
pub use conversion_cache::{FromPythonCache, ToPythonCache};
pub use driver::{CallOptions, run_call};
pub use error::{InvokeError, missing_state};
pub use file_reader::{FileContent, PythonFileReader, py_bytes};

121
litellm-rust/crates/host-python/tests/conversion_cache.rs Normal file
View file

@ -0,0 +1,121 @@
use std::{cell::Cell, rc::Rc};
use litellm_host_python::{FromPythonCache, Pythonized, ToPythonCache};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use rstest::{fixture, rstest};
#[fixture]
fn python() {
Python::initialize();
}
#[rstest]
fn rust_identity_reuses_python_objects_without_merging_equal_values(#[from(python)] _python: ()) {
Python::attach(|py| {
let original = Rc::new(vec![1, 2]);
let cloned = original.clone();
let equal = Rc::new(vec![1, 2]);
let mut cache = ToPythonCache::default();
let first = cache
.get_or_try_insert_with(original.as_ref(), |value| {
Pythonized(value).into_pyobject(py)
})
.unwrap()
.clone();
let second = cache
.get_or_try_insert_with(cloned.as_ref(), |_| panic!("must reuse conversion"))
.unwrap()
.clone();
let third = cache
.get_or_try_insert_with(equal.as_ref(), |value| Pythonized(value).into_pyobject(py))
.unwrap();
assert!(first.is(&second));
assert!(!first.is(third));
assert!(first.eq(third).unwrap());
});
}
#[rstest]
fn python_identity_reuses_rust_values_without_merging_equal_objects(#[from(python)] _python: ()) {
Python::attach(|py| {
let original = PyDict::new(py);
original.set_item("value", 1).unwrap();
let equal = original.copy().unwrap();
let calls = Cell::new(0);
let mut cache = FromPythonCache::default();
let convert = |value: &Bound<'_, PyAny>| {
calls.set(calls.get() + 1);
value.get_item("value")?.extract::<i32>().map(Rc::new)
};
let first = cache
.get_or_try_insert_with(original.as_any(), convert)
.unwrap()
.clone();
let second = cache
.get_or_try_insert_with(original.as_any(), convert)
.unwrap()
.clone();
let third = cache
.get_or_try_insert_with(equal.as_any(), convert)
.unwrap();
assert!(Rc::ptr_eq(&first, &second));
assert!(!Rc::ptr_eq(&first, third));
assert_eq!(&first, third);
assert_eq!(calls.get(), 2);
});
}
#[rstest]
fn python_sources_stay_alive_until_the_cache_is_dropped(#[from(python)] _python: ()) {
Python::attach(|py| {
let value = py
.eval(pyo3::ffi::c_str!("type('Tracked', (), {})()"), None, None)
.unwrap();
let weak = py
.import("weakref")
.unwrap()
.call_method1("ref", (&value,))
.unwrap();
let mut cache = FromPythonCache::default();
cache.get_or_try_insert_with(&value, |_| Ok(42)).unwrap();
drop(value);
assert!(!weak.call0().unwrap().is_none());
drop(cache);
assert!(weak.call0().unwrap().is_none());
});
}
#[rstest]
#[case::to_python(true)]
#[case::from_python(false)]
fn failed_conversions_preserve_exceptions_and_can_be_retried(
#[from(python)] _python: (),
#[case] to_python: bool,
) {
Python::attach(|py| {
let failure = PyValueError::new_err("conversion failed");
if to_python {
let source = vec![1, 2];
let mut cache = ToPythonCache::default();
let error = cache
.get_or_try_insert_with(&source, |_| Err(failure.clone_ref(py)))
.unwrap_err();
assert!(error.value(py).is(failure.value(py)));
let result = cache
.get_or_try_insert_with(&source, |value| Pythonized(value).into_pyobject(py))
.unwrap();
assert_eq!(result.extract::<Vec<i32>>().unwrap(), source);
} else {
let source = PyDict::new(py).into_any();
let mut cache = FromPythonCache::default();
let error = cache
.get_or_try_insert_with(&source, |_| Err(failure.clone_ref(py)))
.unwrap_err();
assert!(error.value(py).is(failure.value(py)));
assert_eq!(
*cache.get_or_try_insert_with(&source, |_| Ok(42)).unwrap(),
42
);
}
});
}

4
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -22,6 +22,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
fancy-regex.workspace = true
litellm-tracing.workspace = true
litellm-traces.workspace = true
litellm-storage-clickhouse.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
@ -51,6 +52,7 @@ litellm-llms-types.workspace = true
litellm-host-python.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
@ -72,7 +74,7 @@ futures-util.workspace = true
rstest.workspace = true
sha2.workspace = true
tokio-tungstenite.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true
aws-sdk-secretsmanager = "1.117.0"
[[bench]]

3
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -44,7 +44,7 @@ mod _native {
#[pymodule_export]
use crate::routes::token_counter::TokenCounter;
#[pymodule_export]
use crate::routes::traces::{NativeTraceStorage, trace_decode_otlp};
use crate::routes::traces::{NativeTraceStorage, trace_decode_otlp, trace_encode_error};
#[cfg(feature = "huggingface")]
#[pymodule_export]
use crate::tokenizer::HuggingFaceEncoding;
@ -111,6 +111,7 @@ mod tests {
"NativeDiagnosticProcessor",
"NativeTraceStorage",
"trace_decode_otlp",
"trace_encode_error",
"TokenCounter",
"Tokenizer",
"gil_stats",

186
litellm-rust/crates/python-bridge/src/routes/traces.rs
View file

@ -1,12 +1,33 @@
use std::collections::BTreeMap;
use litellm_host_python::{FromPythonCache, ToPythonCache};
use litellm_http::ClientVariant;
use litellm_traces::{Connection, Error, InsertTable, Parameter, ReadQuery};
use litellm_storage_clickhouse::Storage;
use litellm_traces::{Error, InsertTable, Parameter, ReadQuery, Shared};
use prost::Message;
use pyo3::{
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
prelude::*,
types::{PyBytes, PyDict, PyList, PyMapping, PyString},
};
#[derive(Message)]
struct OtlpErrorStatus {
#[prost(int32, tag = "1")]
code: i32,
#[prost(string, tag = "2")]
message: String,
}
#[pyfunction]
pub fn trace_encode_error<'py>(py: Python<'py>, message: &str) -> Bound<'py, PyBytes> {
let status = OtlpErrorStatus {
code: 0,
message: message.to_owned(),
};
PyBytes::new(py, &status.encode_to_vec())
}
fn map_error(error: Error) -> PyErr {
match error {
Error::InvalidRow
@ -27,9 +48,7 @@ fn map_error(error: Error) -> PyErr {
#[pyclass]
pub struct NativeTraceStorage {
database: String,
writer: Connection,
reader: Option<Connection>,
storage: Storage,
}
#[pymethods]
@ -39,12 +58,7 @@ impl NativeTraceStorage {
fn new(database: String, url: &str, reader_url: Option<&str>) -> PyResult<Self> {
litellm_traces::schema_statements(&database, 1, 1).map_err(map_error)?;
Ok(Self {
writer: Connection::writer(url).map_err(map_error)?,
reader: reader_url
.map(|value| Connection::reader(value, &database))
.transpose()
.map_err(map_error)?,
database,
storage: Storage::new(database, url, reader_url).map_err(map_error)?,
})
}
@ -55,8 +69,8 @@ impl NativeTraceStorage {
spend_log_retention_days: u32,
) -> PyResult<Bound<'py, PyAny>> {
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
let connection = self.writer.clone();
let database = self.database.clone();
let connection = self.storage.writer().clone();
let database = self.storage.database().to_owned();
crate::execution::run_async(
py,
async move {
@ -77,18 +91,17 @@ impl NativeTraceStorage {
&self,
py: Python<'py>,
table: &str,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] rows: Vec<
BTreeMap<String, serde_json::Value>,
>,
#[pyo3(from_py_with = insert_rows_from_py)] rows: Vec<litellm_traces::InsertRow>,
) -> PyResult<Bound<'py, PyAny>> {
let table = InsertTable::parse(table).map_err(map_error)?;
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
let connection = self.writer.clone();
let database = self.database.clone();
let connection = self.storage.writer().clone();
let database = self.storage.database().to_owned();
crate::execution::run_async(
py,
async move {
litellm_traces::insert_rows(&client, &connection, &database, table, rows).await
litellm_traces::insert_shared_rows(&client, &connection, &database, table, rows)
.await
},
map_error,
)
@ -104,7 +117,7 @@ impl NativeTraceStorage {
>,
) -> PyResult<Bound<'py, PyAny>> {
let query = litellm_traces::LensQuery::parse(name).map_err(map_error)?;
let connection = self.reader.clone().ok_or_else(|| {
let connection = self.storage.reader().cloned().ok_or_else(|| {
PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
})?;
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
@ -127,7 +140,7 @@ impl NativeTraceStorage {
>,
) -> PyResult<Bound<'py, PyAny>> {
let query = ReadQuery::parse(query).map_err(map_error)?;
let connection = self.reader.clone().ok_or_else(|| {
let connection = self.storage.reader().cloned().ok_or_else(|| {
PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
})?;
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
@ -146,21 +159,132 @@ pub fn trace_decode_otlp<'py>(
py: Python<'py>,
body: &[u8],
content_type: Option<&str>,
content_encoding: Option<&str>,
max_decompressed_bytes: usize,
) -> PyResult<Bound<'py, PyAny>> {
let spans = py
.detach(|| {
litellm_traces::decode_otlp(
body,
content_type,
content_encoding,
max_decompressed_bytes,
)
})
.detach(|| litellm_traces::decode_otlp(body, content_type))
.map_err(|error| match error {
litellm_traces::DecodeError::TooLarge => PyOverflowError::new_err(error.to_string()),
_ => PyValueError::new_err(error.to_string()),
})?;
litellm_host_python::Pythonized(spans).into_pyobject(py)
spans_to_py(py, &spans).map(Bound::into_any)
}
fn insert_rows_from_py(value: &Bound<'_, PyAny>) -> PyResult<Vec<litellm_traces::InsertRow>> {
let mut resources = FromPythonCache::default();
value
.try_iter()?
.map(|row| {
let row = row?;
let mut fields = BTreeMap::new();
for item in row.cast::<PyMapping>()?.items()?.iter() {
let (key, value): (String, Bound<'_, PyAny>) = item.extract()?;
let converted = if matches!(
key.as_str(),
"ResourceAttributes" | "ScopeName" | "ScopeVersion"
) {
resources
.get_or_try_insert_with(&value, |value| {
litellm_host_python::from_py_argument::<serde_json::Value>(value)
.map(Shared::new)
})?
.clone()
} else {
Shared::new(litellm_host_python::from_py_argument(&value)?)
};
fields.insert(key, converted);
}
Ok(fields)
})
.collect()
}
fn spans_to_py<'py>(
py: Python<'py>,
spans: &[litellm_traces::DecodedSpan],
) -> PyResult<Bound<'py, PyList>> {
let mut resources = ToPythonCache::default();
let mut scopes = ToPythonCache::default();
let result = PyList::empty(py);
for span in spans {
let resource = resources
.get_or_try_insert_with(span.resource_attributes.as_ref(), |value| {
litellm_host_python::Pythonized(value).into_pyobject(py)
})?;
let row = PyDict::new(py);
row.set_item("trace_id", &span.trace_id)?;
row.set_item("span_id", &span.span_id)?;
row.set_item("parent_span_id", &span.parent_span_id)?;
row.set_item("trace_state", &span.trace_state)?;
row.set_item("name", &span.name)?;
row.set_item("kind", &span.kind)?;
row.set_item("resource_attributes", resource)?;
for (key, value) in [
("scope_name", &span.scope_name),
("scope_version", &span.scope_version),
] {
let value = scopes.get_or_try_insert_with(value.as_ref(), |value| {
Ok(PyString::new(py, value).into_any())
})?;
row.set_item(key, value)?;
}
row.set_item("attributes", &span.attributes)?;
row.set_item("start_ns", span.start_ns)?;
row.set_item("end_ns", span.end_ns)?;
row.set_item("status_code", &span.status_code)?;
row.set_item("status_message", &span.status_message)?;
row.set_item(
"events",
litellm_host_python::Pythonized(&span.events).into_pyobject(py)?,
)?;
result.append(row)?;
}
Ok(result)
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[rstest]
fn insert_projection_preserves_identity_without_merging_equal_resources() {
Python::initialize();
Python::attach(|py| {
let resource = PyDict::new(py);
resource.set_item("service.name", "shared").unwrap();
let equal_resource = resource.copy().unwrap();
let rows = PyList::empty(py);
for value in [&resource, &resource, &equal_resource] {
let row = PyDict::new(py);
row.set_item("ResourceAttributes", value).unwrap();
rows.append(row).unwrap();
}
let projected = insert_rows_from_py(rows.as_any()).unwrap();
assert!(Shared::shares_storage_with(
&projected[0]["ResourceAttributes"],
&projected[1]["ResourceAttributes"]
));
assert!(!Shared::shares_storage_with(
&projected[0]["ResourceAttributes"],
&projected[2]["ResourceAttributes"]
));
assert_eq!(projected[0], projected[2]);
});
}
#[rstest]
fn shared_conversion_preserves_every_decoded_field() {
Python::initialize();
Python::attach(|py| {
let spans = litellm_traces::decode_otlp(
include_bytes!("../../../../../tests/test_litellm/tracing/fixtures/langsmith_deep_agent_export.json"),
Some("application/json"),
).unwrap();
let expected = litellm_host_python::Pythonized(&spans)
.into_pyobject(py)
.unwrap();
let actual = spans_to_py(py, &spans).unwrap();
assert!(actual.eq(expected).unwrap());
});
}
}

2
litellm-rust/crates/secrets-aws/Cargo.toml
View file

@ -21,5 +21,5 @@ aws-credential-types = "1.3.0"
base64.workspace = true
rstest.workspace = true
tokio.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true
tempfile = "3"

2
litellm-rust/crates/secrets-azure/Cargo.toml
View file

@ -20,7 +20,7 @@ percent-encoding = "2.3"
[dev-dependencies]
litellm-http = { workspace = true, features = ["test-support"] }
wiremock = "0.6.5"
wiremock.workspace = true
rstest.workspace = true
serde_json.workspace = true
sha2.workspace = true

2
litellm-rust/crates/secrets-cyberark/Cargo.toml
View file

@ -25,6 +25,6 @@ rcgen = "0.14.10"
rstest.workspace = true
tempfile = "3.27.0"
tokio.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true
serde.workspace = true
serde_json.workspace = true

2
litellm-rust/crates/secrets-google/Cargo.toml
View file

@ -28,4 +28,4 @@ reqwest.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
google-cloud-auth.workspace = true
rstest.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true

2
litellm-rust/crates/secrets-hashicorp/Cargo.toml
View file

@ -21,4 +21,4 @@ veil.workspace = true
rstest.workspace = true
tempfile = "3"
tokio.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true

2
litellm-rust/crates/secrets/Cargo.toml
View file

@ -36,7 +36,7 @@ tokio = { workspace = true, features = ["fs"] }
[dev-dependencies]
litellm-http = { workspace = true, features = ["test-support"] }
rstest.workspace = true
wiremock = "0.6.5"
wiremock.workspace = true
tempfile = "3"
aws-sdk-kms = "1.120.0"
google-cloud-kms-v1 = "1.14.0"

20
litellm-rust/crates/storage-clickhouse/Cargo.toml Normal file
View file

@ -0,0 +1,20 @@
[package]
name = "litellm-storage-clickhouse"
version = "0.1.0"
description = "Shared ClickHouse connection and HTTP storage for LiteLLM features"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
flate2.workspace = true
litellm-http.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
url.workspace = true
[dev-dependencies]
litellm-http = { workspace = true, features = ["test-support"] }
rstest.workspace = true
tokio.workspace = true

5
litellm-rust/crates/storage-clickhouse/README.md Normal file
View file

@ -0,0 +1,5 @@
# ClickHouse storage
`litellm-storage-clickhouse` exports `Storage`, a shared writer connection and optional reader connection for one ClickHouse database. It also exports bounded HTTP read and insert execution
The crate has no trace tables, OTLP types, or named trace queries. `litellm-traces` supplies those rules and uses this storage for both trace rows and spend rows

29
litellm-rust/crates/storage-clickhouse/src/error.rs Normal file
View file

@ -0,0 +1,29 @@
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("invalid ClickHouse insert row")]
InvalidRow,
#[error("invalid ClickHouse insert table")]
InvalidTable,
#[error("invalid ClickHouse HTTP URL")]
InvalidUrl,
#[error("database must be a nonempty SQL identifier and retention must be positive")]
InvalidSchema,
#[error("SQL query must not be empty")]
EmptySql,
#[error("unknown ClickHouse read query")]
InvalidQuery,
#[error("ClickHouse query failed with HTTP status {0}")]
QueryFailed(u16),
#[error("ClickHouse insert failed with HTTP status {0}")]
InsertFailed(u16),
#[error("ClickHouse insert exceeds the encoded size limit")]
InsertTooLarge,
#[error("ClickHouse schema setup failed with HTTP status {0}")]
SchemaFailed(u16),
#[error("ClickHouse query exceeded the response size limit")]
ResponseTooLarge,
#[error("ClickHouse returned an invalid or failed JSON query response")]
InvalidResponse,
#[error("ClickHouse query transport failed")]
Transport,
}

87
litellm-rust/crates/storage-clickhouse/src/insert.rs Normal file
View file

@ -0,0 +1,87 @@
use std::{io::Write, time::Duration};
use flate2::{Compression, write::GzEncoder};
use litellm_http::Client;
use crate::{Connection, Error, valid_identifier};
const INSERT_TIMEOUT: Duration = Duration::from_secs(30);
pub async fn insert_encoded_rows(
client: &Client,
connection: &Connection,
database: &str,
table: &str,
token: &str,
encoded: &str,
) -> Result<(), Error> {
if !valid_identifier(database) {
return Err(Error::InvalidSchema);
}
if !valid_identifier(table) {
return Err(Error::InvalidTable);
}
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder
.write_all(encoded.as_bytes())
.map_err(|_| Error::InvalidRow)?;
let body = encoder.finish().map_err(|_| Error::InvalidRow)?;
insert_compressed_rows(client, connection, database, table, token, body).await
}
pub async fn insert_compressed_rows(
client: &Client,
connection: &Connection,
database: &str,
table: &str,
token: &str,
body: Vec<u8>,
) -> Result<(), Error> {
if !valid_identifier(database) {
return Err(Error::InvalidSchema);
}
if !valid_identifier(table) {
return Err(Error::InvalidTable);
}
let mut url = connection.url().clone();
let existing_pairs: Vec<(String, String)> = url
.query_pairs()
.filter(|(key, _)| {
!matches!(
key.as_ref(),
"query"
| "async_insert"
| "async_insert_deduplicate"
| "wait_for_async_insert"
| "input_format_skip_unknown_fields"
| "date_time_input_format"
)
})
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
url.query_pairs_mut()
.clear()
.extend_pairs(existing_pairs)
.append_pair(
"query",
&format!("INSERT INTO `{database}`.{} FORMAT JSONEachRow", table),
)
.append_pair("insert_deduplication_token", token)
.append_pair("async_insert", "1")
.append_pair("async_insert_deduplicate", "1")
.append_pair("wait_for_async_insert", "1")
.append_pair("input_format_skip_unknown_fields", "0")
.append_pair("date_time_input_format", "best_effort");
let response = client
.post(url)
.timeout(INSERT_TIMEOUT)
.header("Content-Encoding", "gzip")
.body(body)
.send()
.await
.map_err(|_| Error::Transport)?;
if !response.status().is_success() {
return Err(Error::InsertFailed(response.status().as_u16()));
}
Ok(())
}

127
litellm-rust/crates/storage-clickhouse/src/lib.rs Normal file
View file

@ -0,0 +1,127 @@
mod error;
mod insert;
mod read;
pub use error::Error;
pub use insert::{insert_compressed_rows, insert_encoded_rows};
pub use read::{Parameter, execute_read};
use url::Url;
#[derive(Clone)]
pub struct Connection {
url: Url,
}
impl Connection {
pub fn parse(value: &str) -> Result<Self, Error> {
let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?;
if !matches!(url.scheme(), "http" | "https") || url.host().is_none() {
return Err(Error::InvalidUrl);
}
Ok(Self { url })
}
pub fn configured(
url: &str,
database: &str,
user: &str,
password: &str,
) -> Result<Self, Error> {
let mut connection = Self::parse(url)?;
connection
.url
.set_username(user)
.map_err(|_| Error::InvalidUrl)?;
connection
.url
.set_password(Some(password))
.map_err(|_| Error::InvalidUrl)?;
let pairs: Vec<_> = connection
.url
.query_pairs()
.filter(|(key, _)| !matches!(key.as_ref(), "database" | "user" | "password"))
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
connection
.url
.query_pairs_mut()
.clear()
.extend_pairs(pairs)
.append_pair("database", database);
Ok(connection)
}
pub fn writer(url: &str) -> Result<Self, Error> {
let mut connection = Self::parse(url)?;
let pairs: Vec<_> = connection
.url
.query_pairs()
.filter(|(key, _)| !matches!(key.as_ref(), "database" | "readonly" | "query"))
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
connection.url.query_pairs_mut().clear().extend_pairs(pairs);
Ok(connection)
}
pub fn reader(url: &str, database: &str) -> Result<Self, Error> {
let mut connection = Self::parse(url)?;
let pairs: Vec<_> = connection
.url
.query_pairs()
.filter(|(key, _)| key != "database")
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
connection
.url
.query_pairs_mut()
.clear()
.extend_pairs(pairs)
.append_pair("database", database);
Ok(connection)
}
pub fn url(&self) -> &Url {
&self.url
}
}
#[derive(Clone)]
pub struct Storage {
database: String,
writer: Connection,
reader: Option<Connection>,
}
impl Storage {
pub fn new(database: String, url: &str, reader_url: Option<&str>) -> Result<Self, Error> {
if !valid_identifier(&database) {
return Err(Error::InvalidSchema);
}
Ok(Self {
writer: Connection::writer(url)?,
reader: reader_url
.map(|value| Connection::reader(value, &database))
.transpose()?,
database,
})
}
pub fn database(&self) -> &str {
&self.database
}
pub fn writer(&self) -> &Connection {
&self.writer
}
pub fn reader(&self) -> Option<&Connection> {
self.reader.as_ref()
}
}
pub(crate) fn valid_identifier(value: &str) -> bool {
!value.is_empty()
&& value
.bytes()
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
}

113
litellm-rust/crates/storage-clickhouse/src/read.rs Normal file
View file

@ -0,0 +1,113 @@
use std::{collections::BTreeMap, time::Duration};
use litellm_http::Client;
use serde::Deserialize;
use crate::{Connection, Error};
const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;
#[derive(Debug, Deserialize)]
#[serde(untagged)]
pub enum Parameter {
Text(String),
Integer(i64),
Strings(Vec<String>),
}
impl Parameter {
fn encoded(&self) -> String {
match self {
Self::Text(value) => escaped(value),
Self::Integer(value) => value.to_string(),
Self::Strings(values) => format!(
"[{}]",
values
.iter()
.map(|value| format!("'{}'", escaped(value).replace('\'', "\\'")))
.collect::<Vec<_>>()
.join(",")
),
}
}
}
fn escaped(value: &str) -> String {
value
.replace('\\', "\\\\")
.replace('\t', "\\t")
.replace('\n', "\\n")
.replace('\r', "\\r")
.replace('\0', "\\0")
}
pub async fn execute_read(
client: &Client,
connection: &Connection,
sql: &str,
parameters: &BTreeMap<String, Parameter>,
) -> Result<String, Error> {
if sql.trim().is_empty() {
return Err(Error::EmptySql);
}
let mut url = connection.url().clone();
let existing_pairs: Vec<(String, String)> = url
.query_pairs()
.filter(|(key, _)| {
!key.starts_with("param_")
&& !matches!(
key.as_ref(),
"query"
| "readonly"
| "default_format"
| "max_result_rows"
| "result_overflow_mode"
| "max_execution_time"
| "wait_end_of_query"
)
})
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
url.query_pairs_mut()
.clear()
.extend_pairs(existing_pairs)
.append_pair("readonly", "1")
.append_pair("max_result_rows", "1000")
.append_pair("result_overflow_mode", "throw")
.append_pair("max_execution_time", "10")
.append_pair("wait_end_of_query", "1")
.append_pair("default_format", "JSON");
url.query_pairs_mut().extend_pairs(
parameters
.iter()
.map(|(name, value)| (format!("param_{name}"), value.encoded())),
);
let request = client
.post(url)
.timeout(Duration::from_secs(15))
.body(sql.to_owned());
let mut response = request.send().await.map_err(|_| Error::Transport)?;
if !response.status().is_success() {
return Err(Error::QueryFailed(response.status().as_u16()));
}
let mut body = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? {
if body.len() + chunk.len() > MAX_RESPONSE_BYTES {
return Err(Error::ResponseTooLarge);
}
body.extend_from_slice(&chunk);
}
let json: serde_json::Value =
serde_json::from_slice(&body).map_err(|_| Error::InvalidResponse)?;
if json.get("exception").is_some() || !json.get("data").is_some_and(serde_json::Value::is_array)
{
return Err(Error::InvalidResponse);
}
String::from_utf8(body).map_err(|_| Error::InvalidResponse)
}

34
litellm-rust/crates/storage-clickhouse/tests/connection.rs Normal file
View file

@ -0,0 +1,34 @@
use litellm_storage_clickhouse::{Connection, Storage};
use rstest::rstest;
#[rstest]
#[case::http("http://localhost:8123", true)]
#[case::https("https://localhost:8443", true)]
#[case::tcp("tcp://localhost:9000", false)]
#[case::missing_host("http://", false)]
fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool) {
assert_eq!(Connection::parse(value).is_ok(), expected);
}
#[rstest]
#[case::writer_only(None, false)]
#[case::separate_reader(Some("http://localhost:8124"), true)]
fn storage_exports_writer_and_optional_reader(
#[case] reader_url: Option<&str>,
#[case] has_reader: bool,
) {
let storage = Storage::new("litellm".to_owned(), "http://localhost:8123", reader_url)
.expect("valid ClickHouse URLs");
assert_eq!(storage.database(), "litellm");
assert_eq!(storage.writer().url().host_str(), Some("localhost"));
assert_eq!(storage.writer().url().port(), Some(8123));
assert_eq!(storage.reader().is_some(), has_reader);
}
#[rstest]
#[case::empty("")]
#[case::injection("db; DROP DATABASE default")]
fn storage_rejects_invalid_database(#[case] database: &str) {
assert!(Storage::new(database.to_owned(), "http://localhost:8123", None).is_err());
}

34
litellm-rust/crates/storage-clickhouse/tests/transport.rs Normal file
View file

@ -0,0 +1,34 @@
use std::collections::BTreeMap;
use litellm_http::Client;
use litellm_storage_clickhouse::{Connection, Error, execute_read, insert_encoded_rows};
use rstest::rstest;
#[rstest]
#[case::invalid_database("db; DROP DATABASE default", "spend_logs", true)]
#[case::invalid_table("litellm", "spend_logs; DROP TABLE otel_traces", false)]
#[tokio::test]
async fn insert_rejects_invalid_identifiers(
#[case] database: &str,
#[case] table: &str,
#[case] invalid_database: bool,
) {
let client = Client::no_redirect_for_test();
let connection = Connection::writer("http://localhost:8123").expect("valid URL");
let result = insert_encoded_rows(&client, &connection, database, table, "token", "{}").await;
assert!(matches!(&result, Err(Error::InvalidSchema)) == invalid_database);
assert!(matches!(&result, Err(Error::InvalidTable)) == !invalid_database);
}
#[rstest]
#[tokio::test]
async fn read_rejects_empty_sql() {
let client = Client::no_redirect_for_test();
let connection = Connection::reader("http://localhost:8123", "litellm").expect("valid URL");
assert!(matches!(
execute_read(&client, &connection, " ", &BTreeMap::new()).await,
Err(Error::EmptySql)
));
}

2
litellm-rust/crates/traces/AGENTS.md
View file

@ -1,4 +1,4 @@
- Rust owns OTLP wire decoding, ClickHouse schema, row encoding, named reads, connection validation and transport
- Keep OTLP decoding, trace schema, row encoding and named query selection here. Generic ClickHouse connections and HTTP execution belong in `litellm-storage-clickhouse`
- Keep this crate independent of Python; PyO3 conversion and public Python exceptions belong in `python-bridge`
- Keep the SQL migrations here as the only ClickHouse schema definition
- Use typed query parameters and a dedicated SELECT-only reader with server-side limits

15
litellm-rust/crates/traces/Cargo.toml
View file

@ -8,18 +8,25 @@ repository.workspace = true
[dependencies]
base64.workspace = true
flate2.workspace = true
opentelemetry-proto = { version = "0.33.0", default-features = false, features = ["gen-tonic-messages", "trace", "with-serde"] }
prost = "0.14.4"
opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] }
prost.workspace = true
time = { workspace = true, features = ["formatting"] }
litellm-http.workspace = true
litellm-storage-clickhouse.workspace = true
sha2.workspace = true
serde.workspace = true
serde = { workspace = true, features = ["rc"] }
serde_json.workspace = true
strum.workspace = true
thiserror.workspace = true
url.workspace = true
[dev-dependencies]
criterion.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
rstest.workspace = true
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
tokio.workspace = true
wiremock.workspace = true
[[bench]]
name = "resource-fanout"
harness = false

39
litellm-rust/crates/traces/benches/resource-fanout.rs Normal file
View file

@ -0,0 +1,39 @@
use std::{collections::BTreeMap, hint::black_box, time::Duration};
use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use litellm_traces::Shared;
fn fanout<T: Clone>(resource: &T, spans: usize) -> Vec<T> {
(0..spans).map(|_| resource.clone()).collect()
}
fn resource_fanout(c: &mut Criterion) {
let mut group = c.benchmark_group("resource_fanout");
for (attribute_bytes, spans) in [(256, 1), (256, 64), (8192, 1024), (16384, 1024)] {
let attributes = BTreeMap::from([
("service.name".to_owned(), "benchmark".to_owned()),
("payload".to_owned(), "x".repeat(attribute_bytes)),
]);
let owned = Box::new(attributes.clone());
let shared = Shared::new(attributes);
let case = format!("{attribute_bytes}B_{spans}_spans");
group.throughput(Throughput::Elements(spans as u64));
group.bench_with_input(BenchmarkId::new("owned", &case), &owned, |b, resource| {
b.iter(|| black_box(fanout(black_box(resource), spans)));
});
group.bench_with_input(BenchmarkId::new("shared", &case), &shared, |b, resource| {
b.iter(|| black_box(fanout(black_box(resource), spans)));
});
}
group.finish();
}
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(20)
.warm_up_time(Duration::from_secs(1))
.measurement_time(Duration::from_secs(2));
targets = resource_fanout
}
criterion_main!(benches);

13
litellm-rust/crates/traces/query/span_error.sql Normal file
View file

@ -0,0 +1,13 @@
SELECT SpanId AS span_id,
substringUTF8(StatusMessage, {error_offset:UInt64} + 1, 16384) AS message,
lengthUTF8(StatusMessage) AS total_chars,
hex(SHA256(StatusMessage)) AS version
FROM otel_traces
WHERE TraceId = {trace_id:String} AND SpanId = {span_id:String}
AND (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
AND ({api_key_hash:String} = '' OR ApiKeyHash = {api_key_hash:String})
AND ({trace_ref:String} = '' OR
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
AND ({error_version:String} = '' OR hex(SHA256(StatusMessage)) = {error_version:String})
ORDER BY Timestamp, EngineReceivedMs, StatusMessage
LIMIT 1

5
litellm-rust/crates/traces/query/trace_spans.sql
View file

@ -1,6 +1,7 @@
SELECT o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name,
o.ObservationType AS type, o.AgentName AS agent, o.StatusCode AS status,
o.StatusMessage AS status_message,
substringUTF8(o.StatusMessage, 1, 128) AS status_message,
lengthUTF8(o.StatusMessage) > 128 AS error_truncated,
toUnixTimestamp64Nano(o.Timestamp) AS start_ns, o.Duration AS duration_ns,
o.ServiceName AS service, o.InputPreview AS input_preview, o.Model AS model,
o.InputTokens AS input_tokens, o.OutputTokens AS output_tokens,
@ -12,5 +13,5 @@ WHERE o.TraceId = {trace_id:String}
AND ({api_key_hash:String} = '' OR o.ApiKeyHash = {api_key_hash:String})
AND ({trace_ref:String} = '' OR
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
ORDER BY o.Timestamp
ORDER BY o.Timestamp, o.EngineReceivedMs, o.StatusMessage
LIMIT 1 BY o.SpanId

32
litellm-rust/crates/traces/src/error.rs
View file

@ -1,37 +1,7 @@
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("invalid ClickHouse insert row")]
InvalidRow,
#[error("invalid ClickHouse insert table")]
InvalidTable,
#[error("invalid ClickHouse HTTP URL")]
InvalidUrl,
#[error("database must be a nonempty SQL identifier and retention must be positive")]
InvalidSchema,
#[error("SQL query must not be empty")]
EmptySql,
#[error("unknown ClickHouse read query")]
InvalidQuery,
#[error("ClickHouse query failed with HTTP status {0}")]
QueryFailed(u16),
#[error("ClickHouse insert failed with HTTP status {0}")]
InsertFailed(u16),
#[error("ClickHouse insert exceeds the encoded size limit")]
InsertTooLarge,
#[error("ClickHouse schema setup failed with HTTP status {0}")]
SchemaFailed(u16),
#[error("ClickHouse query exceeded the response size limit")]
ResponseTooLarge,
#[error("ClickHouse returned an invalid or failed JSON query response")]
InvalidResponse,
#[error("ClickHouse query transport failed")]
Transport,
}
#[derive(Debug, thiserror::Error)]
pub enum DecodeError {
#[error("invalid OTLP trace payload")]
InvalidPayload,
#[error("OTLP trace payload exceeds the decompressed size limit")]
#[error("OTLP trace payload exceeds the decoding budget")]
TooLarge,
}

327
litellm-rust/crates/traces/src/insert.rs
View file

@ -1,4 +1,10 @@
use std::{collections::BTreeMap, io::Write, time::Duration};
use std::{
borrow::Cow,
collections::BTreeMap,
io::{BufWriter, Write},
};
use serde::{Serialize, Serializer, ser::SerializeMap};
use flate2::{Compression, write::GzEncoder};
use litellm_http::Client;
@ -6,10 +12,11 @@ use serde_json::Value;
use sha2::{Digest, Sha256};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use crate::{Connection, Error};
use crate::{Connection, Error, Shared};
const MAX_INSERT_BYTES: usize = 64 * 1024 * 1024;
const INSERT_TIMEOUT: Duration = Duration::from_secs(30);
pub type InsertRow = BTreeMap<String, Shared<Value>>;
pub enum InsertTable {
OtelTraces,
@ -39,125 +46,176 @@ pub async fn insert_rows(
database: &str,
table: InsertTable,
rows: Vec<BTreeMap<String, Value>>,
) -> Result<(), Error> {
insert_shared_rows(client, connection, database, table, shared_rows(rows)).await
}
pub async fn insert_shared_rows(
client: &Client,
connection: &Connection,
database: &str,
table: InsertTable,
rows: Vec<InsertRow>,
) -> Result<(), Error> {
if rows.is_empty() {
return Ok(());
}
let token = format!(
"{:x}",
Sha256::digest(encode_rows_with_limit(rows.clone(), MAX_INSERT_BYTES)?.as_bytes())
);
let received_ms = OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000;
let rows = rows
.into_iter()
let received_ms = (OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64;
let (token, body) = prepare_insert(&rows, received_ms, MAX_INSERT_BYTES)?;
litellm_storage_clickhouse::insert_compressed_rows(
client,
connection,
database,
table.name(),
&token,
body,
)
.await
}
fn shared_rows(rows: Vec<BTreeMap<String, Value>>) -> Vec<InsertRow> {
rows.into_iter()
.map(|row| {
row.into_iter()
.filter(|(key, _)| key != "EngineReceivedMs")
.chain(std::iter::once((
"EngineReceivedMs".to_owned(),
Value::from(received_ms as u64),
)))
.map(|(key, value)| (key, Shared::new(value)))
.collect()
})
.collect();
let encoded = encode_rows_with_limit(rows, MAX_INSERT_BYTES)?;
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder
.write_all(encoded.as_bytes())
.map_err(|_| Error::InvalidRow)?;
let body = encoder.finish().map_err(|_| Error::InvalidRow)?;
let mut url = connection.url().clone();
let existing_pairs: Vec<(String, String)> = url
.query_pairs()
.filter(|(key, _)| {
!matches!(
key.as_ref(),
"query"
| "async_insert"
| "async_insert_deduplicate"
| "wait_for_async_insert"
| "input_format_skip_unknown_fields"
| "date_time_input_format"
)
})
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
url.query_pairs_mut()
.clear()
.extend_pairs(existing_pairs)
.append_pair(
"query",
&format!(
"INSERT INTO `{database}`.{} FORMAT JSONEachRow",
table.name()
),
)
.append_pair("insert_deduplication_token", &token)
.append_pair("async_insert", "1")
.append_pair("async_insert_deduplicate", "1")
.append_pair("wait_for_async_insert", "1")
.append_pair("input_format_skip_unknown_fields", "0")
.append_pair("date_time_input_format", "best_effort");
let response = client
.post(url)
.timeout(INSERT_TIMEOUT)
.header("Content-Encoding", "gzip")
.body(body)
.send()
.await
.map_err(|_| Error::Transport)?;
if !response.status().is_success() {
return Err(Error::InsertFailed(response.status().as_u16()));
}
Ok(())
.collect()
}
pub fn encode_rows(rows: Vec<BTreeMap<String, Value>>) -> Result<String, Error> {
encode_rows_with_limit(rows, usize::MAX)
}
fn encode_rows_with_limit(
rows: Vec<BTreeMap<String, Value>>,
limit: usize,
) -> Result<String, Error> {
let mut body = Vec::new();
for row in rows {
let encoded = row
.into_iter()
.map(|(name, value)| insert_value(&name, value).map(|value| (name, value)))
.collect::<Result<BTreeMap<_, _>, _>>()?;
let record = serde_json::to_vec(&encoded).map_err(|_| Error::InvalidRow)?;
let size = body
.len()
.checked_add(record.len())
.and_then(|size| size.checked_add(usize::from(!body.is_empty())))
.ok_or(Error::InsertTooLarge)?;
if size > limit {
return Err(Error::InsertTooLarge);
}
if !body.is_empty() {
body.push(b'\n');
}
body.extend_from_slice(&record);
}
let body = write_rows(&shared_rows(rows), None, Vec::new(), usize::MAX)?;
String::from_utf8(body).map_err(|_| Error::InvalidRow)
}
fn insert_value(name: &str, value: Value) -> Result<Value, Error> {
fn prepare_insert(
rows: &[InsertRow],
received_ms: u64,
limit: usize,
) -> Result<(String, Vec<u8>), Error> {
let hash = write_rows(rows, None, HashWriter(Sha256::new()), limit)?;
let token = format!("{:x}", hash.0.finalize());
let encoder = write_rows(
rows,
Some(received_ms),
BufWriter::new(GzEncoder::new(Vec::new(), Compression::default())),
limit,
)?;
let body = encoder
.into_inner()
.map_err(|_| Error::InvalidRow)?
.finish()
.map_err(|_| Error::InvalidRow)?;
Ok((token, body))
}
struct HashWriter(Sha256);
impl Write for HashWriter {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
self.0.update(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
struct LimitedWriter<W> {
inner: W,
remaining: usize,
exceeded: bool,
}
impl<W: Write> Write for LimitedWriter<W> {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
if bytes.len() > self.remaining {
self.exceeded = true;
return Err(std::io::Error::other(Error::InsertTooLarge));
}
let written = self.inner.write(bytes)?;
self.remaining -= written;
Ok(written)
}
fn flush(&mut self) -> std::io::Result<()> {
self.inner.flush()
}
}
fn write_rows<W: Write>(
rows: &[InsertRow],
received_ms: Option<u64>,
writer: W,
limit: usize,
) -> Result<W, Error> {
let mut writer = LimitedWriter {
inner: writer,
remaining: limit,
exceeded: false,
};
for (index, row) in rows.iter().enumerate() {
let result = (|| {
if index != 0 {
writer.write_all(b"\n").map_err(serde_json::Error::io)?;
}
serde_json::to_writer(&mut writer, &EncodedRow { row, received_ms })
})();
if result.is_err() {
return Err(if writer.exceeded {
Error::InsertTooLarge
} else {
Error::InvalidRow
});
}
}
Ok(writer.inner)
}
struct EncodedRow<'a> {
row: &'a InsertRow,
received_ms: Option<u64>,
}
impl Serialize for EncodedRow<'_> {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
let mut map = serializer.serialize_map(None)?;
let mut received_ms = self.received_ms;
for (name, value) in self.row {
if name.as_str() >= "EngineReceivedMs"
&& let Some(timestamp) = received_ms.take()
{
map.serialize_entry("EngineReceivedMs", &timestamp)?;
}
if name == "EngineReceivedMs" && self.received_ms.is_some() {
continue;
}
let value = insert_value(name, value).map_err(serde::ser::Error::custom)?;
map.serialize_entry(name, &value)?;
}
if let Some(timestamp) = received_ms {
map.serialize_entry("EngineReceivedMs", &timestamp)?;
}
map.end()
}
}
fn insert_value<'a>(name: &str, value: &'a Value) -> Result<Cow<'a, Value>, Error> {
let multiplier = match name {
"Timestamp" => 1,
"start_time" | "end_time" | "completion_start_time" => 1_000_000,
_ => return Ok(value),
_ => return Ok(Cow::Borrowed(value)),
};
if name == "completion_start_time" && value.is_null() {
return Ok(value);
return Ok(Cow::Borrowed(value));
}
let timestamp = value.as_i64().ok_or(Error::InvalidRow)?;
let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * multiplier)
.map_err(|_| Error::InvalidRow)?;
datetime
.format(&Rfc3339)
.map(Value::String)
.map(|value| Cow::Owned(Value::String(value)))
.map_err(|_| Error::InvalidRow)
}
@ -168,20 +226,83 @@ mod tests {
use rstest::rstest;
use serde_json::json;
use super::encode_rows_with_limit;
use super::{shared_rows, write_rows};
use crate::Error;
#[rstest]
fn encoded_limit_counts_utf8_bytes_across_rows() {
let rows = vec![
let rows = shared_rows(vec![
BTreeMap::from([("Input".to_owned(), json!("雪"))]),
BTreeMap::from([("Input".to_owned(), json!("雪"))]),
];
let encoded = encode_rows_with_limit(rows.clone(), usize::MAX).expect("valid rows");
]);
let encoded = write_rows(&rows, None, Vec::new(), usize::MAX).expect("valid rows");
assert!(encode_rows_with_limit(rows.clone(), encoded.len()).is_ok());
assert!(write_rows(&rows, None, Vec::new(), encoded.len()).is_ok());
assert!(matches!(
encode_rows_with_limit(rows, encoded.len() - 1),
write_rows(&rows, None, Vec::new(), encoded.len() - 1),
Err(Error::InsertTooLarge)
));
}
#[rstest]
#[case::absent(None)]
#[case::submitted(Some(123))]
fn streamed_insert_preserves_token_and_stamps_receive_time(#[case] submitted: Option<u64>) {
use flate2::read::GzDecoder;
use sha2::{Digest, Sha256};
use std::io::Read;
let mut row = BTreeMap::from([
("ApiKeyHash".into(), json!("key")),
("ResourceAttributes".into(), json!({"message": "雪\n\""})),
("Timestamp".into(), json!(1_234_567_890)),
]);
if let Some(value) = submitted {
row.insert("EngineReceivedMs".into(), json!(value));
}
let legacy = match submitted {
Some(_) => {
"{\"ApiKeyHash\":\"key\",\"EngineReceivedMs\":123,\"ResourceAttributes\":{\"message\":\"雪\\n\\\"\"},\"Timestamp\":\"1970-01-01T00:00:01.23456789Z\"}"
}
None => {
"{\"ApiKeyHash\":\"key\",\"ResourceAttributes\":{\"message\":\"雪\\n\\\"\"},\"Timestamp\":\"1970-01-01T00:00:01.23456789Z\"}"
}
};
let rows = shared_rows(vec![row.clone(), row]);
let (token, body) = super::prepare_insert(&rows, 456, 4096).unwrap();
assert_eq!(
token,
format!("{:x}", Sha256::digest(format!("{legacy}\n{legacy}")))
);
let mut decoded = String::new();
GzDecoder::new(body.as_slice())
.read_to_string(&mut decoded)
.unwrap();
let expected = json!({
"ApiKeyHash": "key", "EngineReceivedMs": 456,
"ResourceAttributes": {"message": "雪\n\""},
"Timestamp": "1970-01-01T00:00:01.23456789Z",
});
assert_eq!(
decoded
.lines()
.map(|line| serde_json::from_str::<serde_json::Value>(line).unwrap())
.collect::<Vec<_>>(),
vec![expected.clone(), expected]
);
assert_eq!(
rows[0]
.get("EngineReceivedMs")
.map(|value| value.as_u64().unwrap()),
submitted
);
}
#[rstest]
fn stamped_insert_enforces_the_encoded_limit() {
let rows = shared_rows(vec![BTreeMap::new()]);
assert!(super::prepare_insert(&rows, 1, 22).is_ok());
assert!(matches!(
super::prepare_insert(&rows, 1, 21),
Err(Error::InsertTooLarge)
));
}

88
litellm-rust/crates/traces/src/lib.rs
View file

@ -2,89 +2,13 @@ mod error;
mod insert;
mod otlp;
mod schema;
mod shared;
mod sql;
pub use error::{DecodeError, Error};
pub use insert::{InsertTable, encode_rows, insert_rows};
pub use error::DecodeError;
pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
pub use litellm_storage_clickhouse::{Connection, Error, Parameter, execute_read};
pub use otlp::{DecodedSpan, decode_otlp};
pub use schema::{ensure_schema, schema_statements};
pub use sql::{LensQuery, Parameter, ReadQuery, execute_named_read, execute_read};
use url::Url;
#[derive(Clone)]
pub struct Connection {
url: Url,
}
impl Connection {
pub fn parse(value: &str) -> Result<Self, Error> {
let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?;
if !matches!(url.scheme(), "http" | "https") || url.host().is_none() {
return Err(Error::InvalidUrl);
}
Ok(Self { url })
}
pub fn configured(
url: &str,
database: &str,
user: &str,
password: &str,
) -> Result<Self, Error> {
let mut connection = Self::parse(url)?;
connection
.url
.set_username(user)
.map_err(|_| Error::InvalidUrl)?;
connection
.url
.set_password(Some(password))
.map_err(|_| Error::InvalidUrl)?;
let pairs: Vec<_> = connection
.url
.query_pairs()
.filter(|(key, _)| !matches!(key.as_ref(), "database" | "user" | "password"))
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
connection
.url
.query_pairs_mut()
.clear()
.extend_pairs(pairs)
.append_pair("database", database);
Ok(connection)
}
pub fn writer(url: &str) -> Result<Self, Error> {
let mut connection = Self::parse(url)?;
let pairs: Vec<_> = connection
.url
.query_pairs()
.filter(|(key, _)| !matches!(key.as_ref(), "database" | "readonly" | "query"))
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
connection.url.query_pairs_mut().clear().extend_pairs(pairs);
Ok(connection)
}
pub fn reader(url: &str, database: &str) -> Result<Self, Error> {
let mut connection = Self::parse(url)?;
let pairs: Vec<_> = connection
.url
.query_pairs()
.filter(|(key, _)| key != "database")
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
connection
.url
.query_pairs_mut()
.clear()
.extend_pairs(pairs)
.append_pair("database", database);
Ok(connection)
}
pub fn url(&self) -> &Url {
&self.url
}
}
pub use shared::{Shared, SharedIdentity};
pub use sql::{LensQuery, ReadQuery, execute_named_read};

221
litellm-rust/crates/traces/src/otlp.rs
View file

@ -1,221 +0,0 @@
use std::{collections::BTreeMap, io::Read};
use base64::Engine;
use flate2::read::GzDecoder;
use opentelemetry_proto::tonic::{
collector::trace::v1::ExportTraceServiceRequest,
common::v1::{AnyValue, KeyValue, any_value::Value as AttributeValue},
trace::v1::{Span, span::SpanKind, status::StatusCode},
};
use prost::Message;
use serde::Serialize;
use serde_json::Value;
use crate::DecodeError;
#[derive(Serialize)]
pub struct DecodedEvent {
pub name: String,
pub attributes: BTreeMap<String, String>,
}
#[derive(Serialize)]
pub struct DecodedSpan {
pub trace_id: String,
pub span_id: String,
pub parent_span_id: String,
pub trace_state: String,
pub name: String,
pub kind: String,
pub resource_attributes: BTreeMap<String, String>,
pub scope_name: String,
pub scope_version: String,
pub attributes: BTreeMap<String, String>,
pub start_ns: u64,
pub end_ns: u64,
pub status_code: String,
pub status_message: String,
pub events: Vec<DecodedEvent>,
}
pub fn decode_otlp(
body: &[u8],
content_type: Option<&str>,
content_encoding: Option<&str>,
max_decompressed_bytes: usize,
) -> Result<Vec<DecodedSpan>, DecodeError> {
let payload = if content_encoding == Some("gzip") || body.starts_with(&[0x1f, 0x8b]) {
let limit = u64::try_from(max_decompressed_bytes).map_err(|_| DecodeError::TooLarge)?;
let mut decoded = Vec::new();
GzDecoder::new(body)
.take(limit + 1)
.read_to_end(&mut decoded)
.map_err(|_| DecodeError::InvalidPayload)?;
decoded
} else {
body.to_vec()
};
if payload.len() > max_decompressed_bytes {
return Err(DecodeError::TooLarge);
}
let request = if content_type.is_some_and(|value| value.contains("json")) {
let value: Value =
serde_json::from_slice(&payload).map_err(|_| DecodeError::InvalidPayload)?;
serde_json::from_value(normalize_json_ids(value)?)
.map_err(|_| DecodeError::InvalidPayload)?
} else {
ExportTraceServiceRequest::decode(payload.as_slice())
.map_err(|_| DecodeError::InvalidPayload)?
};
Ok(request
.resource_spans
.into_iter()
.flat_map(|resource_spans| {
let resource_attributes = attributes(
resource_spans
.resource
.map(|resource| resource.attributes)
.unwrap_or_default(),
);
resource_spans
.scope_spans
.into_iter()
.flat_map(move |scope_spans| {
let scope = scope_spans.scope.unwrap_or_default();
let resource_attributes = resource_attributes.clone();
scope_spans.spans.into_iter().map(move |span| {
decoded_span(span, &resource_attributes, &scope.name, &scope.version)
})
})
})
.collect())
}
fn normalize_json_ids(value: Value) -> Result<Value, DecodeError> {
match value {
Value::Object(fields) => fields
.into_iter()
.map(|(name, value)| {
let normalized = if matches!(name.as_str(), "traceId" | "spanId" | "parentSpanId") {
let encoded = value.as_str().ok_or(DecodeError::InvalidPayload)?;
let bytes = base64::engine::general_purpose::STANDARD
.decode(encoded)
.map_err(|_| DecodeError::InvalidPayload)?;
Value::String(hex_bytes(&bytes))
} else if name == "kind" && value.is_string() {
let kind = SpanKind::from_str_name(value.as_str().unwrap_or_default())
.ok_or(DecodeError::InvalidPayload)?;
Value::from(kind as i32)
} else if name == "code" && value.is_string() {
let code = StatusCode::from_str_name(value.as_str().unwrap_or_default())
.ok_or(DecodeError::InvalidPayload)?;
Value::from(code as i32)
} else {
normalize_json_ids(value)?
};
Ok((name, normalized))
})
.collect::<Result<serde_json::Map<_, _>, _>>()
.map(Value::Object),
Value::Array(values) => values
.into_iter()
.map(normalize_json_ids)
.collect::<Result<Vec<_>, _>>()
.map(Value::Array),
value => Ok(value),
}
}
fn hex_bytes(bytes: &[u8]) -> String {
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}
fn decoded_span(
span: Span,
resource_attributes: &BTreeMap<String, String>,
scope_name: &str,
scope_version: &str,
) -> DecodedSpan {
let status = span.status.unwrap_or_default();
DecodedSpan {
trace_id: hex_bytes(&span.trace_id),
span_id: hex_bytes(&span.span_id),
parent_span_id: hex_bytes(&span.parent_span_id),
trace_state: span.trace_state,
name: span.name,
kind: SpanKind::try_from(span.kind)
.unwrap_or(SpanKind::Unspecified)
.as_str_name()
.to_owned(),
resource_attributes: resource_attributes.clone(),
scope_name: scope_name.to_owned(),
scope_version: scope_version.to_owned(),
attributes: attributes(span.attributes),
start_ns: span.start_time_unix_nano,
end_ns: span.end_time_unix_nano,
status_code: StatusCode::try_from(status.code)
.unwrap_or(StatusCode::Unset)
.as_str_name()
.to_owned(),
status_message: status.message,
events: span
.events
.into_iter()
.map(|event| DecodedEvent {
name: event.name,
attributes: attributes(event.attributes),
})
.collect(),
}
}
fn attributes(values: Vec<KeyValue>) -> BTreeMap<String, String> {
values
.into_iter()
.map(|entry| {
(
entry.key,
entry.value.as_ref().map(attribute_text).unwrap_or_default(),
)
})
.collect()
}
fn attribute_text(value: &AnyValue) -> String {
match value.value.as_ref() {
Some(AttributeValue::StringValue(value)) => value.clone(),
Some(AttributeValue::BoolValue(value)) => value.to_string(),
Some(AttributeValue::IntValue(value)) => value.to_string(),
Some(AttributeValue::DoubleValue(value)) => {
serde_json::to_string(value).unwrap_or_default()
}
Some(AttributeValue::BytesValue(value)) => String::from_utf8_lossy(value).into_owned(),
Some(AttributeValue::ArrayValue(value)) => format!(
"[{}]",
value
.values
.iter()
.map(|value| serde_json::to_string(&attribute_text(value)).unwrap_or_default())
.collect::<Vec<_>>()
.join(", ")
),
Some(AttributeValue::KvlistValue(value)) => format!(
"{{{}}}",
value
.values
.iter()
.map(|entry| format!(
"{}: {}",
serde_json::to_string(&entry.key).unwrap_or_default(),
serde_json::to_string(
&entry.value.as_ref().map(attribute_text).unwrap_or_default()
)
.unwrap_or_default()
))
.collect::<Vec<_>>()
.join(", ")
),
Some(AttributeValue::StringValueStrindex(value)) => value.to_string(),
None => String::new(),
}
}

101
litellm-rust/crates/traces/src/otlp/attributes.rs Normal file
View file

@ -0,0 +1,101 @@
use std::{collections::BTreeMap, io::Write};
use opentelemetry_proto::tonic::common::v1::{
AnyValue, KeyValue, any_value::Value as AttributeValue,
};
use serde::{
Serialize, Serializer,
ser::{SerializeMap, SerializeSeq},
};
use super::limits::{Budget, MAX_ATTRIBUTES};
use crate::DecodeError;
struct AttributeWriter<'a> {
body: Vec<u8>,
budget: &'a mut Budget,
}
impl Write for AttributeWriter<'_> {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
self.budget
.consume(bytes.len())
.map_err(std::io::Error::other)?;
self.body.extend_from_slice(bytes);
Ok(bytes.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
pub(super) fn attributes(
values: Vec<KeyValue>,
budget: &mut Budget,
) -> Result<BTreeMap<String, String>, DecodeError> {
if values.len() > MAX_ATTRIBUTES {
return Err(DecodeError::TooLarge);
}
values
.into_iter()
.map(|entry| {
budget.consume(entry.key.len() + 96)?;
let text = match entry.value {
Some(AnyValue {
value: Some(AttributeValue::StringValue(value)),
}) => {
budget.consume(value.len())?;
value
}
Some(AnyValue {
value: Some(AttributeValue::BytesValue(value)),
}) => {
budget.consume(value.len().saturating_mul(3))?;
String::from_utf8_lossy(&value).into_owned()
}
value => {
let mut writer = AttributeWriter {
body: Vec::new(),
budget,
};
serde_json::to_writer(&mut writer, &AttributeJson(value.as_ref()))
.map_err(|_| DecodeError::TooLarge)?;
String::from_utf8(writer.body).map_err(|_| DecodeError::InvalidPayload)?
}
};
Ok((entry.key, text))
})
.collect()
}
struct AttributeJson<'a>(Option<&'a AnyValue>);
impl Serialize for AttributeJson<'_> {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
match self.0.and_then(|value| value.value.as_ref()) {
Some(AttributeValue::StringValue(value)) => serializer.serialize_str(value),
Some(AttributeValue::BoolValue(value)) => serializer.serialize_bool(*value),
Some(AttributeValue::IntValue(value)) => serializer.serialize_i64(*value),
Some(AttributeValue::DoubleValue(value)) => serializer.serialize_f64(*value),
Some(AttributeValue::BytesValue(value)) => {
serializer.serialize_str(&String::from_utf8_lossy(value))
}
Some(AttributeValue::ArrayValue(value)) => {
let mut sequence = serializer.serialize_seq(Some(value.values.len()))?;
for entry in &value.values {
sequence.serialize_element(&AttributeJson(Some(entry)))?;
}
sequence.end()
}
Some(AttributeValue::KvlistValue(value)) => {
let mut map = serializer.serialize_map(Some(value.values.len()))?;
for entry in &value.values {
map.serialize_entry(&entry.key, &AttributeJson(entry.value.as_ref()))?;
}
map.end()
}
Some(AttributeValue::StringValueStrindex(value)) => serializer.serialize_i32(*value),
None => serializer.serialize_unit(),
}
}
}

212
litellm-rust/crates/traces/src/otlp/limits.rs Normal file
View file

@ -0,0 +1,212 @@
use std::fmt;
use prost::encoding::{DecodeContext, WireType, decode_key, decode_varint, skip_field};
use serde::de::{DeserializeSeed, MapAccess, SeqAccess, Visitor};
use crate::{DecodeError, Shared};
pub(super) const MAX_DEPTH: usize = 32;
pub(super) const MAX_NODES: usize = 65_536;
pub(super) const MAX_SPANS: usize = 4_096;
pub(super) const MAX_ATTRIBUTES: usize = 256;
pub(super) const MAX_EVENTS: usize = 256;
pub(super) const MAX_DECODED_SPAN_BYTES: usize = 16 * 1024 * 1024;
pub(super) fn json_preflight(payload: &[u8]) -> Result<(), DecodeError> {
let mut nodes = 0;
let mut exceeded = false;
let mut decoder = serde_json::Deserializer::from_slice(payload);
let result = JsonBudget {
nodes: &mut nodes,
exceeded: &mut exceeded,
depth: 0,
}
.deserialize(&mut decoder)
.and_then(|()| decoder.end());
if exceeded {
return Err(DecodeError::TooLarge);
}
result.map_err(|_| DecodeError::InvalidPayload)
}
struct JsonBudget<'a> {
nodes: &'a mut usize,
exceeded: &'a mut bool,
depth: usize,
}
impl<'de> DeserializeSeed<'de> for JsonBudget<'_> {
type Value = ();
fn deserialize<D: serde::Deserializer<'de>>(self, decoder: D) -> Result<(), D::Error> {
*self.nodes += 1;
if *self.nodes > MAX_NODES || self.depth > MAX_DEPTH {
*self.exceeded = true;
return Err(serde::de::Error::custom("OTLP structure exceeds budget"));
}
decoder.deserialize_any(self)
}
}
impl<'de> Visitor<'de> for JsonBudget<'_> {
type Value = ();
fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str("OTLP JSON")
}
fn visit_bool<E: serde::de::Error>(self, _: bool) -> Result<(), E> {
Ok(())
}
fn visit_i64<E: serde::de::Error>(self, _: i64) -> Result<(), E> {
Ok(())
}
fn visit_u64<E: serde::de::Error>(self, _: u64) -> Result<(), E> {
Ok(())
}
fn visit_f64<E: serde::de::Error>(self, _: f64) -> Result<(), E> {
Ok(())
}
fn visit_str<E: serde::de::Error>(self, _: &str) -> Result<(), E> {
Ok(())
}
fn visit_unit<E: serde::de::Error>(self) -> Result<(), E> {
Ok(())
}
fn visit_seq<A: SeqAccess<'de>>(self, mut sequence: A) -> Result<(), A::Error> {
while sequence
.next_element_seed(JsonBudget {
nodes: self.nodes,
exceeded: self.exceeded,
depth: self.depth + 1,
})?
.is_some()
{}
Ok(())
}
fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<(), A::Error> {
while map
.next_key_seed(JsonBudget {
nodes: self.nodes,
exceeded: self.exceeded,
depth: self.depth + 1,
})?
.is_some()
{
map.next_value_seed(JsonBudget {
nodes: self.nodes,
exceeded: self.exceeded,
depth: self.depth + 1,
})?;
}
Ok(())
}
}
#[derive(Clone, Copy)]
enum MessageKind {
Export,
ResourceSpans,
Resource,
ScopeSpans,
Scope,
Span,
Event,
Link,
Status,
KeyValue,
AnyValue,
Array,
KvList,
}
impl MessageKind {
fn child(self, tag: u32) -> Option<Self> {
match (self, tag) {
(Self::Export, 1) => Some(Self::ResourceSpans),
(Self::ResourceSpans, 1) => Some(Self::Resource),
(Self::ResourceSpans, 2) => Some(Self::ScopeSpans),
(Self::Resource, 1)
| (Self::Scope, 3)
| (Self::Span, 9)
| (Self::Event, 3)
| (Self::Link, 4)
| (Self::KvList, 1) => Some(Self::KeyValue),
(Self::ScopeSpans, 1) => Some(Self::Scope),
(Self::ScopeSpans, 2) => Some(Self::Span),
(Self::Span, 11) => Some(Self::Event),
(Self::Span, 13) => Some(Self::Link),
(Self::Span, 15) => Some(Self::Status),
(Self::KeyValue, 2) | (Self::Array, 1) => Some(Self::AnyValue),
(Self::AnyValue, 5) => Some(Self::Array),
(Self::AnyValue, 6) => Some(Self::KvList),
_ => None,
}
}
}
pub(super) fn protobuf_preflight(payload: &[u8]) -> Result<(), DecodeError> {
scan_message(payload, MessageKind::Export, 0, &mut 0)
}
fn scan_message(
mut payload: &[u8],
kind: MessageKind,
depth: usize,
nodes: &mut usize,
) -> Result<(), DecodeError> {
if depth > MAX_DEPTH {
return Err(DecodeError::TooLarge);
}
while !payload.is_empty() {
*nodes += 1;
if *nodes > MAX_NODES {
return Err(DecodeError::TooLarge);
}
let (tag, wire) = decode_key(&mut payload).map_err(|_| DecodeError::InvalidPayload)?;
if let (WireType::LengthDelimited, Some(child)) = (wire, kind.child(tag)) {
let length = decode_varint(&mut payload).map_err(|_| DecodeError::InvalidPayload)?;
let length = usize::try_from(length).map_err(|_| DecodeError::InvalidPayload)?;
let (message, rest) = payload
.split_at_checked(length)
.ok_or(DecodeError::InvalidPayload)?;
scan_message(message, child, depth + 1, nodes)?;
payload = rest;
} else {
skip_field(wire, tag, &mut payload, DecodeContext::default())
.map_err(|_| DecodeError::InvalidPayload)?;
}
}
Ok(())
}
pub(super) struct Budget {
remaining: usize,
}
impl Budget {
pub(super) fn new(remaining: usize) -> Self {
Self { remaining }
}
pub(super) fn clone_shared<T: Clone>(
&mut self,
value: &Shared<T>,
allocated_bytes: impl FnOnce(&T) -> usize,
) -> Result<Shared<T>, DecodeError> {
let cloned = value.clone();
if !value.shares_storage_with(&cloned) {
self.consume(allocated_bytes(value))?;
}
Ok(cloned)
}
pub(super) fn consume(&mut self, bytes: usize) -> Result<(), DecodeError> {
self.remaining = self
.remaining
.checked_sub(bytes)
.ok_or(DecodeError::TooLarge)?;
Ok(())
}
}

42
litellm-rust/crates/traces/src/otlp/mod.rs Normal file
View file

@ -0,0 +1,42 @@
mod attributes;
mod limits;
mod span;
mod wire;
use serde::Serialize;
use std::collections::BTreeMap;
use crate::{DecodeError, Shared};
#[derive(Serialize)]
pub struct DecodedEvent {
pub name: String,
pub attributes: BTreeMap<String, String>,
}
#[derive(Serialize)]
pub struct DecodedSpan {
pub trace_id: String,
pub span_id: String,
pub parent_span_id: String,
pub trace_state: String,
pub name: String,
pub kind: String,
pub resource_attributes: Shared<BTreeMap<String, String>>,
pub scope_name: Shared<String>,
pub scope_version: Shared<String>,
pub attributes: BTreeMap<String, String>,
pub start_ns: u64,
pub end_ns: u64,
pub status_code: String,
pub status_message: String,
pub events: Vec<DecodedEvent>,
}
pub fn decode_otlp(
body: &[u8],
content_type: Option<&str>,
) -> Result<Vec<DecodedSpan>, DecodeError> {
let request = wire::decode(body, content_type)?;
span::flatten(request)
}

166
litellm-rust/crates/traces/src/otlp/span.rs Normal file
View file

@ -0,0 +1,166 @@
use std::collections::BTreeMap;
use opentelemetry_proto::tonic::{
collector::trace::v1::ExportTraceServiceRequest,
trace::v1::{ResourceSpans, ScopeSpans, Span, span::SpanKind, status::StatusCode},
};
use super::{
DecodedEvent, DecodedSpan,
attributes::attributes,
limits::{Budget, MAX_ATTRIBUTES, MAX_DECODED_SPAN_BYTES, MAX_EVENTS, MAX_SPANS},
};
use crate::{DecodeError, Shared};
pub(super) fn flatten(request: ExportTraceServiceRequest) -> Result<Vec<DecodedSpan>, DecodeError> {
let mut budget = Budget::new(MAX_DECODED_SPAN_BYTES);
let mut spans = Vec::new();
for resource in request.resource_spans {
append_resource(resource, &mut budget, &mut spans)?;
}
Ok(spans)
}
fn append_resource(
resource: ResourceSpans,
budget: &mut Budget,
spans: &mut Vec<DecodedSpan>,
) -> Result<(), DecodeError> {
let attributes = Shared::new(attributes(
resource
.resource
.map(|resource| resource.attributes)
.unwrap_or_default(),
budget,
)?);
for scope in resource.scope_spans {
append_scope(scope, &attributes, budget, spans)?;
}
Ok(())
}
fn append_scope(
scope_spans: ScopeSpans,
resource: &Shared<BTreeMap<String, String>>,
budget: &mut Budget,
spans: &mut Vec<DecodedSpan>,
) -> Result<(), DecodeError> {
let scope = scope_spans.scope.unwrap_or_default();
if scope.attributes.len() > MAX_ATTRIBUTES {
return Err(DecodeError::TooLarge);
}
budget.consume(scope.name.len() + scope.version.len())?;
let scope_name: Shared<String> = scope.name.into();
let scope_version: Shared<String> = scope.version.into();
for span in scope_spans.spans {
if spans.len() >= MAX_SPANS {
return Err(DecodeError::TooLarge);
}
validate_span(&span)?;
budget.consume(
span.name.len()
+ span.trace_state.len()
+ span
.status
.as_ref()
.map_or(0, |status| status.message.len())
+ size_of::<DecodedSpan>()
+ 128,
)?;
spans.push(decoded_span(
span,
resource,
&scope_name,
&scope_version,
budget,
)?);
}
Ok(())
}
fn valid_id(value: &[u8], length: usize) -> bool {
value.len() == length && value.iter().any(|byte| *byte != 0)
}
fn validate_span(span: &Span) -> Result<(), DecodeError> {
if !valid_id(&span.trace_id, 16)
|| !valid_id(&span.span_id, 8)
|| (!span.parent_span_id.is_empty() && !valid_id(&span.parent_span_id, 8))
|| span.start_time_unix_nano > i64::MAX as u64
|| span.end_time_unix_nano > i64::MAX as u64
|| span.end_time_unix_nano < span.start_time_unix_nano
|| span
.links
.iter()
.any(|link| !valid_id(&link.trace_id, 16) || !valid_id(&link.span_id, 8))
{
return Err(DecodeError::InvalidPayload);
}
if span.events.len() > MAX_EVENTS
|| span.links.len() > MAX_EVENTS
|| span.attributes.len() > MAX_ATTRIBUTES
|| span
.links
.iter()
.any(|link| link.attributes.len() > MAX_ATTRIBUTES)
|| span
.events
.iter()
.any(|event| event.attributes.len() > MAX_ATTRIBUTES)
{
return Err(DecodeError::TooLarge);
}
Ok(())
}
fn hex_bytes(bytes: &[u8]) -> String {
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}
fn decoded_span(
span: Span,
resource_attributes: &Shared<BTreeMap<String, String>>,
scope_name: &Shared<String>,
scope_version: &Shared<String>,
budget: &mut Budget,
) -> Result<DecodedSpan, DecodeError> {
let status = span.status.unwrap_or_default();
Ok(DecodedSpan {
trace_id: hex_bytes(&span.trace_id),
span_id: hex_bytes(&span.span_id),
parent_span_id: hex_bytes(&span.parent_span_id),
trace_state: span.trace_state,
name: span.name,
kind: SpanKind::try_from(span.kind)
.unwrap_or(SpanKind::Unspecified)
.as_str_name()
.to_owned(),
resource_attributes: budget.clone_shared(resource_attributes, |attributes| {
attributes
.iter()
.map(|(key, value)| key.len() + value.len() + 96)
.sum()
})?,
scope_name: budget.clone_shared(scope_name, String::len)?,
scope_version: budget.clone_shared(scope_version, String::len)?,
attributes: attributes(span.attributes, budget)?,
start_ns: span.start_time_unix_nano,
end_ns: span.end_time_unix_nano,
status_code: StatusCode::try_from(status.code)
.unwrap_or(StatusCode::Unset)
.as_str_name()
.to_owned(),
status_message: status.message,
events: span
.events
.into_iter()
.map(|event| {
budget.consume(event.name.len() + 96)?;
Ok(DecodedEvent {
name: event.name,
attributes: attributes(event.attributes, budget)?,
})
})
.collect::<Result<Vec<_>, DecodeError>>()?,
})
}

43
litellm-rust/crates/traces/src/otlp/wire.rs Normal file
View file

@ -0,0 +1,43 @@
use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
use prost::Message;
use super::limits::{json_preflight, protobuf_preflight};
use crate::DecodeError;
#[derive(strum::EnumString)]
#[strum(ascii_case_insensitive)]
enum OtlpMediaType {
#[strum(serialize = "application/json")]
Json,
#[strum(
serialize = "application/x-protobuf",
serialize = "application/protobuf"
)]
Protobuf,
}
pub(super) fn decode(
body: &[u8],
content_type: Option<&str>,
) -> Result<ExportTraceServiceRequest, DecodeError> {
let media_type = content_type
.unwrap_or("application/x-protobuf")
.split(';')
.next()
.unwrap_or_default()
.trim()
.parse::<OtlpMediaType>()
.map_err(|_| DecodeError::InvalidPayload)?;
let request = match media_type {
OtlpMediaType::Json => {
json_preflight(body)?;
serde_json::from_slice(body).map_err(|_| DecodeError::InvalidPayload)?
}
OtlpMediaType::Protobuf => {
protobuf_preflight(body)?;
ExportTraceServiceRequest::decode(body).map_err(|_| DecodeError::InvalidPayload)?
}
};
Ok(request)
}

46
litellm-rust/crates/traces/src/shared.rs Normal file
View file

@ -0,0 +1,46 @@
use std::ops::Deref;
use serde::Serialize;
type Storage<T> = std::sync::Arc<T>;
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(transparent)]
pub struct Shared<T>(Storage<T>);
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct SharedIdentity(usize);
impl<T> Shared<T> {
pub fn new(value: T) -> Self {
Self(Storage::new(value))
}
pub fn identity(&self) -> SharedIdentity {
SharedIdentity(std::ptr::from_ref(self.as_ref()) as usize)
}
pub fn shares_storage_with(&self, other: &Self) -> bool {
self.identity() == other.identity()
}
}
impl<T> From<T> for Shared<T> {
fn from(value: T) -> Self {
Self::new(value)
}
}
impl<T> AsRef<T> for Shared<T> {
fn as_ref(&self) -> &T {
self.0.as_ref()
}
}
impl<T> Deref for Shared<T> {
type Target = T;
fn deref(&self) -> &T {
self.as_ref()
}
}

116
litellm-rust/crates/traces/src/sql.rs
View file

@ -1,17 +1,14 @@
use std::{collections::BTreeMap, time::Duration};
use serde::Deserialize;
use std::collections::BTreeMap;
use litellm_http::Client;
use crate::{Connection, Error};
const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;
use crate::{Connection, Error, Parameter, execute_read};
pub enum ReadQuery {
ListTraces,
TraceSpans,
SpanDetail,
SpanError,
SpendByResponseIds,
}
@ -21,6 +18,7 @@ impl ReadQuery {
"list_traces" => Ok(Self::ListTraces),
"trace_spans" => Ok(Self::TraceSpans),
"span_detail" => Ok(Self::SpanDetail),
"span_error" => Ok(Self::SpanError),
"spend_by_response_ids" => Ok(Self::SpendByResponseIds),
_ => Err(Error::InvalidQuery),
}
@ -31,116 +29,12 @@ impl ReadQuery {
Self::ListTraces => include_str!("../query/list_traces.sql"),
Self::TraceSpans => include_str!("../query/trace_spans.sql"),
Self::SpanDetail => include_str!("../query/span_detail.sql"),
Self::SpanError => include_str!("../query/span_error.sql"),
Self::SpendByResponseIds => include_str!("../query/spend_by_response_ids.sql"),
}
}
}
#[derive(Debug, Deserialize)]
#[serde(untagged)]
pub enum Parameter {
Text(String),
Integer(i64),
Strings(Vec<String>),
}
impl Parameter {
fn encoded(&self) -> String {
match self {
Self::Text(value) => escaped(value),
Self::Integer(value) => value.to_string(),
Self::Strings(values) => format!(
"[{}]",
values
.iter()
.map(|value| format!("'{}'", escaped(value).replace('\'', "\\'")))
.collect::<Vec<_>>()
.join(",")
),
}
}
}
fn escaped(value: &str) -> String {
value
.replace('\\', "\\\\")
.replace('\t', "\\t")
.replace('\n', "\\n")
.replace('\r', "\\r")
.replace('\0', "\\0")
}
pub async fn execute_read(
client: &Client,
connection: &Connection,
sql: &str,
parameters: &BTreeMap<String, Parameter>,
) -> Result<String, Error> {
if sql.trim().is_empty() {
return Err(Error::EmptySql);
}
let mut url = connection.url().clone();
let existing_pairs: Vec<(String, String)> = url
.query_pairs()
.filter(|(key, _)| {
!key.starts_with("param_")
&& !matches!(
key.as_ref(),
"query"
| "readonly"
| "default_format"
| "max_result_rows"
| "result_overflow_mode"
| "max_execution_time"
| "wait_end_of_query"
)
})
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect();
url.query_pairs_mut()
.clear()
.extend_pairs(existing_pairs)
.append_pair("readonly", "1")
.append_pair("max_result_rows", "1000")
.append_pair("result_overflow_mode", "throw")
.append_pair("max_execution_time", "10")
.append_pair("wait_end_of_query", "1")
.append_pair("default_format", "JSON");
url.query_pairs_mut().extend_pairs(
parameters
.iter()
.map(|(name, value)| (format!("param_{name}"), value.encoded())),
);
let request = client
.post(url)
.timeout(Duration::from_secs(15))
.body(sql.to_owned());
let mut response = request.send().await.map_err(|_| Error::Transport)?;
if !response.status().is_success() {
return Err(Error::QueryFailed(response.status().as_u16()));
}
let mut body = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? {
if body.len() + chunk.len() > MAX_RESPONSE_BYTES {
return Err(Error::ResponseTooLarge);
}
body.extend_from_slice(&chunk);
}
let json: serde_json::Value =
serde_json::from_slice(&body).map_err(|_| Error::InvalidResponse)?;
if json.get("exception").is_some() || !json.get("data").is_some_and(serde_json::Value::is_array)
{
return Err(Error::InvalidResponse);
}
String::from_utf8(body).map_err(|_| Error::InvalidResponse)
}
#[derive(Clone, Copy)]
pub enum LensQuery {
Sample,

105
litellm-rust/crates/traces/tests/insert.rs
View file

@ -1,8 +1,107 @@
use std::collections::BTreeMap;
use std::{
collections::BTreeMap,
io::{BufRead, BufReader},
};
use litellm_traces::encode_rows;
use rstest::rstest;
use flate2::read::GzDecoder;
use litellm_http::Client;
use litellm_traces::{
Connection, Error, InsertRow, InsertTable, Shared, encode_rows, insert_shared_rows,
};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{header, method},
};
#[fixture]
fn shared_rows(#[default(16 * 1024)] attribute_bytes: usize) -> Vec<InsertRow> {
let resource = Shared::new(json!({"shared": "x".repeat(attribute_bytes)}));
(0..1024)
.map(|index| {
BTreeMap::from([
("ResourceAttributes".into(), resource.clone()),
("SpanId".into(), Shared::new(json!(format!("{index:016x}")))),
("Timestamp".into(), Shared::new(json!(1))),
])
})
.collect()
}
#[rstest]
#[case::one_request(1)]
#[case::concurrent_requests(2)]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn shared_fanout_survives_gzip_insert_over_http(
shared_rows: Vec<InsertRow>,
#[case] concurrency: usize,
) {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(header("Content-Encoding", "gzip"))
.respond_with(ResponseTemplate::new(200))
.expect(concurrency as u64)
.mount(&server)
.await;
let client = Client::no_redirect_for_test();
let connection = Connection::parse(&server.uri()).unwrap();
let expected_resource = shared_rows[0]["ResourceAttributes"].clone();
let expected_count = shared_rows.len();
let mut requests = tokio::task::JoinSet::new();
for _ in 0..concurrency {
let client = client.clone();
let connection = connection.clone();
let rows = shared_rows.clone();
requests.spawn(async move {
insert_shared_rows(
&client,
&connection,
"traces",
InsertTable::OtelTraces,
rows,
)
.await
});
}
while let Some(result) = requests.join_next().await {
result.unwrap().unwrap();
}
let received = server.received_requests().await.unwrap();
assert_eq!(received.len(), concurrency);
for request in received {
let decoder = GzDecoder::new(request.body.as_slice());
let mut count = 0;
for (index, line) in BufReader::new(decoder).lines().enumerate() {
let row: Value = serde_json::from_str(&line.unwrap()).unwrap();
assert_eq!(&row["ResourceAttributes"], expected_resource.as_ref());
assert_eq!(row["SpanId"], format!("{index:016x}"));
assert_eq!(row["Timestamp"], "1970-01-01T00:00:00.000000001Z");
assert!(row["EngineReceivedMs"].as_u64().unwrap() > 0);
count += 1;
}
assert_eq!(count, expected_count);
}
}
#[rstest]
#[tokio::test]
async fn shared_fanout_over_insert_limit_never_reaches_http(
#[with(64 * 1024)] shared_rows: Vec<InsertRow>,
) {
let server = MockServer::start().await;
let connection = Connection::parse(&server.uri()).unwrap();
let result = insert_shared_rows(
&Client::no_redirect_for_test(),
&connection,
"traces",
InsertTable::OtelTraces,
shared_rows,
)
.await;
assert!(matches!(result, Err(Error::InsertTooLarge)));
assert!(server.received_requests().await.unwrap().is_empty());
}
#[rstest]
#[case::span("Timestamp", json!(1_234_567_890), json!("1970-01-01T00:00:01.23456789Z"))]

145
litellm-rust/crates/traces/tests/migrations.rs
View file

@ -835,3 +835,148 @@ async fn lens_content_keeps_output_visible_after_long_input(
assert_eq!(recovered, original);
Ok(())
}
#[rstest]
#[case::ascii(10, format!("ParentCommand: {}", "x".repeat(460_000)))]
#[case::multibyte(1_000, "\u{1f9ea}".repeat(1_024))]
#[case::escaped(1_000, "\0\n\"\\".repeat(1_024))]
#[tokio::test]
async fn trace_error_previews_preserve_paginated_diagnostics(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
#[case] span_count: usize,
#[case] message: String,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let rows = (0..span_count)
.map(|index| {
serde_json::from_value(serde_json::json!({
"Timestamp": timestamp + index as i64, "TraceId": "diagnostic-trace",
"SpanId": format!("span-{index}"), "SpanName": "tool",
"StatusCode": "STATUS_CODE_ERROR", "StatusMessage": message,
}))
})
.collect::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
insert_rows(&database, "otel_traces", rows).await?;
let reader = Connection::reader(&database.url, "trace_test")?;
let mut parameters = BTreeMap::from([
(
"trace_id".into(),
Parameter::Text("diagnostic-trace".into()),
),
("team_ids".into(), Parameter::Strings(vec![])),
("api_key_hash".into(), Parameter::Text(String::new())),
("trace_ref".into(), Parameter::Text(String::new())),
]);
let body = execute_named_read(
&database.client,
&reader,
ReadQuery::TraceSpans,
&parameters,
)
.await?;
let response: serde_json::Value = serde_json::from_str(&body)?;
let spans = response["data"].as_array().expect("trace spans");
assert_eq!(spans.len(), span_count);
let prefix: String = message.chars().take(128).collect();
assert!(!prefix.is_empty());
assert!(
spans
.iter()
.all(|span| span["status_message"] == prefix && span["error_truncated"] == 1)
);
parameters.insert("span_id".into(), Parameter::Text("span-0".into()));
parameters.insert("error_version".into(), Parameter::Text(String::new()));
let mut recovered = String::new();
loop {
parameters.insert(
"error_offset".into(),
Parameter::Integer(recovered.chars().count() as i64),
);
let body = execute_named_read(&database.client, &reader, ReadQuery::SpanError, &parameters)
.await?;
assert!(body.len() < 128 * 1024);
let response: serde_json::Value = serde_json::from_str(&body)?;
let chunk = response["data"][0]["message"]
.as_str()
.expect("diagnostic chunk");
assert!(!chunk.is_empty());
recovered.push_str(chunk);
let version = response["data"][0]["version"]
.as_str()
.expect("diagnostic version");
parameters.insert("error_version".into(), Parameter::Text(version.into()));
if recovered.chars().count() >= message.chars().count() {
break;
}
}
assert_eq!(recovered, message);
parameters.insert(
"api_key_hash".into(),
Parameter::Text("unrelated-key".into()),
);
let denied =
execute_named_read(&database.client, &reader, ReadQuery::SpanError, &parameters).await?;
assert_eq!(
serde_json::from_str::<serde_json::Value>(&denied)?["data"],
serde_json::json!([])
);
Ok(())
}
#[rstest]
#[case::different_start(1, 0)]
#[case::different_receive(0, 1)]
#[case::tied_timestamps(0, 0)]
#[tokio::test]
async fn duplicate_span_preview_matches_diagnostic(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
#[case] start_delta: i64,
#[case] receive_delta: i64,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let message = "a".repeat(200);
let rows = [
(start_delta, receive_delta, "z".repeat(200)),
(0, 0, message.clone()),
]
.into_iter()
.map(|(start_delta, receive_delta, message)| {
serde_json::from_value(serde_json::json!({
"Timestamp": timestamp + start_delta, "EngineReceivedMs": 100 + receive_delta,
"TraceId": "duplicate-trace", "SpanId": "duplicate-span", "StatusMessage": message,
}))
})
.collect::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
insert_rows(&database, "otel_traces", rows).await?;
let reader = Connection::reader(&database.url, "trace_test")?;
let parameters = BTreeMap::from([
("trace_id".into(), Parameter::Text("duplicate-trace".into())),
("span_id".into(), Parameter::Text("duplicate-span".into())),
("team_ids".into(), Parameter::Strings(vec![])),
("api_key_hash".into(), Parameter::Text(String::new())),
("trace_ref".into(), Parameter::Text(String::new())),
("error_version".into(), Parameter::Text(String::new())),
("error_offset".into(), Parameter::Integer(0)),
]);
let preview = execute_named_read(
&database.client,
&reader,
ReadQuery::TraceSpans,
&parameters,
)
.await?;
let diagnostic =
execute_named_read(&database.client, &reader, ReadQuery::SpanError, &parameters).await?;
let preview: serde_json::Value = serde_json::from_str(&preview)?;
let diagnostic: serde_json::Value = serde_json::from_str(&diagnostic)?;
assert_eq!(preview["data"].as_array().unwrap().len(), 1);
assert_eq!(preview["data"][0]["status_message"], message[..128]);
assert_eq!(diagnostic["data"][0]["message"], message);
Ok(())
}

350
litellm-rust/crates/traces/tests/otlp.rs
View file

@ -1,33 +1,19 @@
use flate2::{Compression, write::GzEncoder};
use litellm_traces::Shared;
use litellm_traces::decode_otlp;
use rstest::rstest;
use std::io::Write;
const FIXTURE: &[u8] = include_bytes!(
"../../../../tests/test_litellm/tracing/fixtures/langsmith_deep_agent_export.json"
);
#[rstest]
#[case::json(FIXTURE, Some("application/json"), None)]
#[case::gzip_json(FIXTURE, Some("application/json"), Some("gzip"))]
fn decodes_neutral_spans(
#[case] body: &[u8],
#[case] content_type: Option<&str>,
#[case] content_encoding: Option<&str>,
) {
let payload = if content_encoding == Some("gzip") {
let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
encoder.write_all(body).expect("gzip input");
encoder.finish().expect("gzip payload")
} else {
body.to_vec()
};
let spans = decode_otlp(&payload, content_type, content_encoding, 8 * 1024 * 1024)
.expect("valid OTLP export");
#[case::json(FIXTURE, Some("application/json"))]
fn decodes_neutral_spans(#[case] body: &[u8], #[case] content_type: Option<&str>) {
let spans = decode_otlp(body, content_type).expect("valid OTLP export");
assert_eq!(spans.len(), 6);
assert_eq!(spans[0].trace_id, "4bad42b84e9de3ba46fc870185f8f023");
assert_eq!(spans[0].resource_attributes["service.name"], "agent-demo");
assert_eq!(spans[0].scope_name, "langsmith");
assert_eq!(spans[0].scope_name.as_ref(), "langsmith");
assert!(
spans
.iter()
@ -36,12 +22,322 @@ fn decodes_neutral_spans(
}
#[rstest]
#[case::invalid(b"not protobuf", None, 8 * 1024 * 1024)]
#[case::too_large(FIXTURE, Some("application/json"), 1)]
fn rejects_invalid_or_oversized_payload(
#[case] body: &[u8],
#[case] content_type: Option<&str>,
#[case] limit: usize,
) {
assert!(decode_otlp(body, content_type, None, limit).is_err());
fn accepts_trace_larger_than_eight_mib(mut span: opentelemetry_proto::tonic::trace::v1::Span) {
use prost::Message;
span.name = "x".repeat(9 * 1024 * 1024);
let body = request_with(span).encode_to_vec();
let decoded = decode_otlp(&body, None).expect("16 MiB default accepts a 9 MiB trace");
assert_eq!(decoded[0].name.len(), 9 * 1024 * 1024);
}
#[rstest]
fn rejects_invalid_payload() {
assert!(decode_otlp(b"not protobuf", None).is_err());
}
#[rstest]
fn decoder_does_not_enforce_the_http_body_limit() {
let body = format!("{{\"ignored\":\"{}\"}}", "x".repeat(16 * 1024 * 1024 + 1));
assert!(
decode_otlp(body.as_bytes(), Some("application/json"))
.unwrap()
.is_empty()
);
}
fn request_with(
span: opentelemetry_proto::tonic::trace::v1::Span,
) -> opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest {
use opentelemetry_proto::tonic::{
collector::trace::v1::ExportTraceServiceRequest,
trace::v1::{ResourceSpans, ScopeSpans},
};
ExportTraceServiceRequest {
resource_spans: vec![ResourceSpans {
scope_spans: vec![ScopeSpans {
spans: vec![span],
..Default::default()
}],
..Default::default()
}],
}
}
#[rstest::fixture]
fn span() -> opentelemetry_proto::tonic::trace::v1::Span {
opentelemetry_proto::tonic::trace::v1::Span {
trace_id: vec![1; 16],
span_id: vec![2; 8],
start_time_unix_nano: 1,
end_time_unix_nano: 2,
..Default::default()
}
}
#[rstest]
fn standard_json_and_protobuf_preserve_the_same_identifiers(
span: opentelemetry_proto::tonic::trace::v1::Span,
) {
use prost::Message;
let request = request_with(span);
let json = serde_json::to_vec(&request).unwrap();
let binary = request.encode_to_vec();
let json_spans = decode_otlp(&json, Some("application/json; charset=utf-8")).unwrap();
let binary_spans = decode_otlp(&binary, Some("application/x-protobuf")).unwrap();
assert_eq!(
serde_json::to_value(&json_spans).unwrap(),
serde_json::to_value(&binary_spans).unwrap()
);
assert_eq!(json_spans[0].trace_id, "01".repeat(16));
assert_eq!(json_spans[0].span_id, "02".repeat(8));
}
#[rstest]
#[case::json("APPLICATION/JSON; charset=utf-8", b"{}")]
#[case::protobuf("application/x-protobuf; charset=binary", b"")]
#[case::protobuf_alias("APPLICATION/PROTOBUF", b"")]
fn supported_content_types_select_the_decoder(#[case] content_type: &str, #[case] body: &[u8]) {
assert!(decode_otlp(body, Some(content_type)).is_ok());
}
#[rstest]
#[case::missing_content_type(None)]
#[case::unsupported_content_type(Some("text/plain"))]
fn content_type_defaults_to_protobuf_and_rejects_unknown_values(
#[case] content_type: Option<&str>,
) {
let result = decode_otlp(b"", content_type);
assert_eq!(result.is_ok(), content_type.is_none());
}
#[rstest]
#[case::short_trace(vec![1; 15], vec![2;8], 1, 2)]
#[case::zero_trace(vec![0; 16], vec![2;8], 1, 2)]
#[case::short_span(vec![1; 16], vec![2;7], 1, 2)]
#[case::timestamp_overflow(vec![1;16], vec![2;8], i64::MAX as u64 + 1, i64::MAX as u64 + 1)]
#[case::negative_duration(vec![1;16], vec![2;8], 3, 2)]
fn rejects_ids_and_timestamps_that_cannot_be_stored(
#[case] trace_id: Vec<u8>,
#[case] span_id: Vec<u8>,
#[case] start: u64,
#[case] end: u64,
) {
use prost::Message;
let span = opentelemetry_proto::tonic::trace::v1::Span {
trace_id,
span_id,
start_time_unix_nano: start,
end_time_unix_nano: end,
..Default::default()
};
assert!(matches!(
decode_otlp(&request_with(span).encode_to_vec(), None),
Err(litellm_traces::DecodeError::InvalidPayload)
));
}
#[rstest]
fn resource_fanout_shares_one_allocation(span: opentelemetry_proto::tonic::trace::v1::Span) {
use opentelemetry_proto::tonic::{
common::v1::{AnyValue, KeyValue, any_value::Value},
resource::v1::Resource,
};
use prost::Message;
let mut request = request_with(span.clone());
request.resource_spans[0].resource = Some(Resource {
attributes: vec![KeyValue {
key: "shared".into(),
value: Some(AnyValue {
value: Some(Value::StringValue("x".repeat(16 * 1024))),
}),
..Default::default()
}],
..Default::default()
});
request.resource_spans[0].scope_spans[0].spans = vec![span; 1024];
let second_scope = request.resource_spans[0].scope_spans[0].clone();
request.resource_spans[0].scope_spans.push(second_scope);
request
.resource_spans
.push(request.resource_spans[0].clone());
let body = request.encode_to_vec();
let decoded = decode_otlp(&body, None).expect("shared resources do not expand with span count");
assert_eq!(decoded.len(), 4096);
assert!(decoded[..2048].iter().all(|span| {
Shared::shares_storage_with(&span.resource_attributes, &decoded[0].resource_attributes)
}));
assert!(!Shared::shares_storage_with(
&decoded[0].resource_attributes,
&decoded[2048].resource_attributes
));
assert_eq!(
*decoded[0].resource_attributes,
*decoded[2048].resource_attributes
);
}
#[rstest]
fn nested_values_are_serialized_once(span: opentelemetry_proto::tonic::trace::v1::Span) {
use opentelemetry_proto::tonic::common::v1::{
AnyValue, ArrayValue, KeyValue, any_value::Value,
};
use prost::Message;
let nested = (0..8).fold(
AnyValue {
value: Some(Value::StringValue("quoted \"value\"".into())),
},
|child, _| AnyValue {
value: Some(Value::ArrayValue(ArrayValue {
values: vec![child],
})),
},
);
let mut request = request_with(span);
request.resource_spans[0].scope_spans[0].spans[0].attributes = vec![KeyValue {
key: "nested".into(),
value: Some(nested),
..Default::default()
}];
let spans = decode_otlp(&request.encode_to_vec(), None).unwrap();
let expected = (0..8).fold(serde_json::json!("quoted \"value\""), |child, _| {
serde_json::json!([child])
});
assert_eq!(
serde_json::from_str::<serde_json::Value>(&spans[0].attributes["nested"]).unwrap(),
expected
);
assert!(spans[0].attributes["nested"].len() < 64);
}
#[rstest]
#[case::nesting(format!("{}0{}", "[".repeat(40), "]".repeat(40)).into_bytes())]
#[case::nodes(format!("[{}]", vec!["0"; 65537].join(",")).into_bytes())]
fn rejects_json_structure_before_building_a_tree(#[case] body: Vec<u8>) {
assert!(matches!(
decode_otlp(&body, Some("application/json")),
Err(litellm_traces::DecodeError::TooLarge)
));
}
#[rstest]
#[case::depth(40, 1)]
#[case::nodes(0, 65537)]
fn protobuf_preflight_rejects_expansion_before_prost_allocates(
span: opentelemetry_proto::tonic::trace::v1::Span,
#[case] depth: usize,
#[case] count: usize,
) {
use opentelemetry_proto::tonic::common::v1::{
AnyValue, ArrayValue, KeyValue, any_value::Value,
};
use prost::Message;
let value = (0..depth).fold(
AnyValue {
value: Some(Value::BoolValue(true)),
},
|child, _| AnyValue {
value: Some(Value::ArrayValue(ArrayValue {
values: vec![child],
})),
},
);
let mut request = request_with(span);
request.resource_spans[0].scope_spans[0].spans[0].attributes = vec![KeyValue {
key: "deep".into(),
value: Some(value),
..Default::default()
}];
request.resource_spans = vec![request.resource_spans[0].clone(); count];
let body = request.encode_to_vec();
assert!(matches!(
decode_otlp(&body, None),
Err(litellm_traces::DecodeError::TooLarge)
));
}
#[rstest]
fn scope_fanout_shares_name_and_version(span: opentelemetry_proto::tonic::trace::v1::Span) {
use opentelemetry_proto::tonic::common::v1::InstrumentationScope;
use prost::Message;
let mut request = request_with(span.clone());
request.resource_spans[0].scope_spans[0].scope = Some(InstrumentationScope {
name: "n".repeat(16 * 1024),
version: "v".repeat(16 * 1024),
..Default::default()
});
request.resource_spans[0].scope_spans[0].spans = vec![span; 1024];
let decoded = decode_otlp(&request.encode_to_vec(), None).unwrap();
assert!(
decoded
.iter()
.all(|span| Shared::shares_storage_with(&span.scope_name, &decoded[0].scope_name))
);
assert!(
decoded.iter().all(|span| Shared::shares_storage_with(
&span.scope_version,
&decoded[0].scope_version
))
);
assert_eq!(decoded[0].scope_name.len(), 16 * 1024);
assert_eq!(decoded[0].scope_version.len(), 16 * 1024);
}
#[rstest]
fn unique_attribute_expansion_still_respects_decoded_budget(
span: opentelemetry_proto::tonic::trace::v1::Span,
) {
use opentelemetry_proto::tonic::common::v1::{AnyValue, KeyValue, any_value::Value};
use prost::Message;
let mut request = request_with(span.clone());
request.resource_spans[0].scope_spans[0].spans = (0..1024)
.map(|index| {
let mut span = span.clone();
span.attributes = vec![KeyValue {
key: "unique".into(),
value: Some(AnyValue {
value: Some(Value::StringValue(format!(
"{index:04}{}",
"x".repeat(16_300)
))),
}),
..Default::default()
}];
span
})
.collect();
let body = request.encode_to_vec();
assert!(body.len() < 16 * 1024 * 1024);
assert!(matches!(
decode_otlp(&body, None),
Err(litellm_traces::DecodeError::TooLarge)
));
}
#[rstest]
fn escaped_attribute_expansion_is_bounded_below_four_mib(
span: opentelemetry_proto::tonic::trace::v1::Span,
) {
use opentelemetry_proto::tonic::common::v1::{
AnyValue, ArrayValue, KeyValue, any_value::Value,
};
use prost::Message;
let mut request = request_with(span);
request.resource_spans[0].scope_spans[0].spans[0].attributes = vec![KeyValue {
key: "escaped".into(),
value: Some(AnyValue {
value: Some(Value::ArrayValue(ArrayValue {
values: vec![AnyValue {
value: Some(Value::StringValue("\0".repeat(3 * 1024 * 1024))),
}],
})),
}),
..Default::default()
}];
let body = request.encode_to_vec();
assert!(body.len() < 4 * 1024 * 1024);
assert!(matches!(
decode_otlp(&body, None),
Err(litellm_traces::DecodeError::TooLarge)
));
}

11
litellm-rust/crates/traces/tests/queries.rs
View file

@ -1,11 +0,0 @@
use litellm_traces::Connection;
use rstest::rstest;
#[rstest]
#[case::http("http://localhost:8123", true)]
#[case::https("https://localhost:8443", true)]
#[case::tcp("tcp://localhost:9000", false)]
#[case::missing_host("http://", false)]
fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool) {
assert_eq!(Connection::parse(value).is_ok(), expected);
}

23
litellm-rust/crates/traces/tests/shared.rs Normal file
View file

@ -0,0 +1,23 @@
use litellm_traces::Shared;
use rstest::rstest;
#[rstest]
fn clones_preserve_values_and_serialize_transparently() {
let original = Shared::new(vec!["value".to_owned()]);
let cloned = original.clone();
assert_eq!(cloned.as_ref(), original.as_ref());
assert_eq!(
serde_json::to_value(&cloned).unwrap(),
serde_json::json!(["value"])
);
}
#[rstest]
fn clones_share_storage_without_merging_equal_values() {
let original = Shared::new("value".to_owned());
let cloned = original.clone();
let equal = Shared::new("value".to_owned());
assert!(original.shares_storage_with(&cloned));
assert!(!original.shares_storage_with(&equal));
assert_eq!(*original, *equal);
}

29
litellm/__init__.py
View file

@ -663,7 +663,7 @@ azure_anthropic_models: Set = set()
azure_text_models: Set = set()
anyscale_models: Set = set()
cerebras_models: Set = set()
nadir_models: Set = set() # mutable-ok: provider registry, filled from model_cost at import like every sibling provider
nadir_models: Set = set()
galadriel_models: Set = set()
nvidia_nim_models: Set = set()
nvidia_riva_models: Set = set()
@ -697,7 +697,7 @@ recraft_models: Set = set()
cometapi_models: Set = set()
oci_models: Set = set()
vercel_ai_gateway_models: Set = set()
edenai_models: Set = set() # mutable-ok: filled from the price map at import, like the sibling provider sets
edenai_models: Set = set()
volcengine_models: Set = set()
wandb_models: Set = set(WANDB_MODELS)
ovhcloud_models: Set = set()
@ -2282,6 +2282,24 @@ if TYPE_CHECKING:
# Track if async client cleanup has been registered (for lazy loading)
_async_client_cleanup_registered = False
# litellm.agent() entrypoints, resolved lazily from litellm.harness by __getattr__.
_AGENT_EXPORTS: Final = frozenset(
{
"agent",
"aagent",
"agent_session",
"aagent_session",
"agent_resume",
"aagent_resume",
"agent_capabilities",
"Harness",
"ClaudeCodeOptions",
"CodexOptions",
"OpenCodeOptions",
"DeepAgentsOptions",
}
)
# Eager loading for backwards compatibility with VCR and other HTTP recording tools
# When LITELLM_DISABLE_LAZY_LOADING is set, lazy-loaded attributes are loaded at import time
# For now, this only affects encoding (tiktoken) as it was the only reported issue
@ -2315,6 +2333,13 @@ def __getattr__(name: str) -> Any:
handler_func: Final = registry[name]
return handler_func(name)
# litellm.agent() and friends: imported on first access (not needed for completion calls)
if name == "harness" or name in _AGENT_EXPORTS:
import importlib
harness_module = importlib.import_module("litellm.harness")
return harness_module if name == "harness" else getattr(harness_module, name)
# Lazy load encoding from main.py to avoid heavy tiktoken import
if name == "encoding":
from ._lazy_imports import get_litellm_globals

16
litellm/_logging.py
View file

@ -352,13 +352,9 @@ def _replace_string_leaves(value: object, values: Iterator[str]) -> object:
if isinstance(value, str):
return next(values)
if isinstance(value, dict):
return { # mutable-ok: LogRecord extras must keep JSON dict shape for handlers
key: _replace_string_leaves(child, values) for key, child in value.items()
}
return {key: _replace_string_leaves(child, values) for key, child in value.items()}
if isinstance(value, list):
return [ # mutable-ok: LogRecord extras must keep JSON list shape for handlers
_replace_string_leaves(child, values) for child in value
]
return [_replace_string_leaves(child, values) for child in value]
if isinstance(value, tuple):
return tuple(_replace_string_leaves(child, values) for child in value)
return value
@ -368,13 +364,9 @@ def _sort_processed_sets(original: object, processed: object) -> object:
if isinstance(original, set) and isinstance(processed, list):
return sorted(processed)
if isinstance(original, dict) and isinstance(processed, dict):
return { # mutable-ok: sorting nested sets must preserve the surrounding JSON dict
key: _sort_processed_sets(original.get(key), value) for key, value in processed.items()
}
return {key: _sort_processed_sets(original.get(key), value) for key, value in processed.items()}
if isinstance(original, list) and isinstance(processed, list):
return [ # mutable-ok: sorting nested sets must preserve the surrounding JSON list
_sort_processed_sets(before, after) for before, after in zip(original, processed)
]
return [_sort_processed_sets(before, after) for before, after in zip(original, processed)]
if isinstance(original, tuple) and isinstance(processed, tuple):
return tuple(_sort_processed_sets(before, after) for before, after in zip(original, processed))
return processed

4
litellm/_redis.py
View file

@ -233,7 +233,7 @@ def _coerce_redis_kwargs_types(
"socket_keepalive": bool,
}
)
result: Final = dict(redis_kwargs) # mutable-ok: per-key try/except coercion below needs to drop individual keys
result: Final = dict(redis_kwargs)
for key, value in redis_kwargs.items():
if not isinstance(value, str):
continue
@ -803,7 +803,7 @@ def _credential_provider_auth_kwargs(redis_kwargs: dict) -> dict:
superseded: Final = frozenset({"redis_connect_func", "username", "password"})
kept: Final = ((k, v) for k, v in redis_kwargs.items() if k not in superseded)
return dict(kept, credential_provider=credential_provider) # mutable-ok: the branches below mutate these kwargs
return dict(kept, credential_provider=credential_provider)
def get_redis_client(**env_overrides):

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