mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
Merge upstream main into fix/bedrock-scheduled-batch-status
This commit is contained in:
lindicaphxag-tech
commit
703904352e
527 changed files with 9163 additions and 5246 deletions
|
|
@ -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:
|
||||
|
|
@ -551,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 \
|
||||
|
|
@ -625,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 \
|
||||
|
|
@ -697,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 \
|
||||
|
|
@ -752,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 \
|
||||
|
|
@ -815,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 \
|
||||
|
|
@ -859,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 \
|
||||
|
|
@ -904,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 \
|
||||
|
|
@ -948,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 \
|
||||
|
|
@ -986,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 \
|
||||
|
|
@ -1031,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 \
|
||||
|
|
@ -1075,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 \
|
||||
|
|
@ -1121,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 \
|
||||
|
|
@ -1176,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 \
|
||||
|
|
@ -1210,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 \
|
||||
|
|
@ -1254,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 \
|
||||
|
|
@ -1298,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 \
|
||||
|
|
@ -1342,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 \
|
||||
|
|
@ -1387,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 \
|
||||
|
|
@ -1432,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 \
|
||||
|
|
@ -1466,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 \
|
||||
|
|
@ -1511,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 \
|
||||
|
|
@ -1531,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
|
||||
|
|
@ -1605,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:
|
||||
|
|
@ -1629,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:
|
||||
|
|
@ -1660,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
|
||||
|
||||
|
|
@ -1829,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"
|
||||
|
|
@ -1926,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 \
|
||||
|
|
@ -2013,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"
|
||||
|
|
@ -2096,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"
|
||||
|
|
@ -2148,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"
|
||||
|
|
@ -2229,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"
|
||||
|
|
@ -2334,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"
|
||||
|
|
@ -2406,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"
|
||||
|
|
@ -2491,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"
|
||||
|
|
@ -2588,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"
|
||||
|
|
@ -2659,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"
|
||||
|
|
@ -2689,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
|
||||
|
|
@ -3189,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 \
|
||||
|
|
@ -3492,7 +3437,6 @@ workflows:
|
|||
- image_gen_testing
|
||||
- logging_testing
|
||||
- audio_testing
|
||||
- redis_caching_unit_tests
|
||||
- upload-coverage:
|
||||
requires:
|
||||
- realtime_translation_testing
|
||||
|
|
@ -3507,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
|
||||
|
|
|
|||
|
|
@ -191,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 \
|
||||
|
|
|
|||
|
|
@ -116,7 +116,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
|
||||
|
|
@ -151,6 +151,7 @@ legacy_paths() {
|
|||
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)
|
||||
|
|
|
|||
48
.github/scripts/assert_ci_coverage.py
vendored
48
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -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",
|
||||
|
|
|
|||
116
.github/workflows/test-unit.yml
vendored
116
.github/workflows/test-unit.yml
vendored
|
|
@ -204,7 +204,6 @@ jobs:
|
|||
tests/unit/proxy/spend_tracking
|
||||
--ignore=tests/unit/proxy/spend_tracking/test_search_api_logging.py
|
||||
tests/unit/proxy/pass_through_endpoints
|
||||
tests/test_litellm/proxy/pass_through_endpoints
|
||||
tests/unit/proxy/_experimental
|
||||
--ignore=tests/unit/proxy/_experimental/mcp_server
|
||||
tests/unit/proxy/experimental
|
||||
|
|
@ -217,93 +216,46 @@ jobs:
|
|||
tests/unit/proxy/enterprise_billing
|
||||
tests/unit/proxy/types_utils
|
||||
tests/unit/proxy/logging_endpoints
|
||||
tests/unit/proxy/test__types.py
|
||||
tests/unit/proxy/test_aiohttp_cleanup_closed.py
|
||||
tests/unit/proxy/test_aiohttp_session_recovery.py
|
||||
tests/unit/proxy/test_api_key_masking_in_errors.py
|
||||
tests/unit/proxy/test_audio_speech_prometheus_hooks.py
|
||||
tests/unit/proxy/test_batch_expiry.py
|
||||
tests/unit/proxy/test_batch_metadata_none_fix.py
|
||||
tests/unit/proxy/test_batch_retrieve_bedrock.py
|
||||
tests/unit/proxy/test_batch_x_litellm_model_encoding.py
|
||||
tests/unit/proxy/test_blocked_response_usage.py
|
||||
tests/unit/proxy/test_body_snapshot_callback_params.py
|
||||
tests/unit/proxy/test_budget_reservation.py
|
||||
tests/unit/proxy/test_bug_report_config.py
|
||||
tests/unit/proxy/test_caching_routes.py
|
||||
tests/unit/proxy/test_chat_completion_metadata.py
|
||||
tests/unit/proxy/test_claude_code_marketplace.py
|
||||
tests/unit/proxy/test_collector.py
|
||||
tests/unit/proxy/test_common_request_processing.py
|
||||
tests/unit/proxy/test_component_allowlists.py
|
||||
tests/unit/proxy/test_conftest.py
|
||||
tests/unit/proxy/test_cors_config.py
|
||||
tests/unit/proxy/test_custom_proxy.py
|
||||
tests/unit/proxy/test_dynamic_mcp_route.py
|
||||
tests/unit/proxy/test_empty_model_list.py
|
||||
tests/unit/proxy/test_enforce_user_param.py
|
||||
tests/unit/proxy/test_fallback_management_endpoints.py
|
||||
tests/unit/proxy/test_fastapi_offline_routes.py
|
||||
tests/unit/proxy/test_filter_models_by_team_access_group.py
|
||||
tests/unit/proxy/test_health_check_functions.py
|
||||
tests/unit/proxy/test_health_check_max_tokens.py
|
||||
tests/unit/proxy/test_init_litellm_callbacks.py
|
||||
tests/unit/proxy/test_langfuse_passthrough_security.py
|
||||
tests/unit/proxy/test_lazy_openapi_snapshot.py
|
||||
tests/unit/proxy/test_litellm_pre_call_utils.py
|
||||
tests/unit/proxy/test_max_budget_env_var.py
|
||||
tests/unit/proxy/test_mcp_asgi_response.py
|
||||
tests/unit/proxy/test_model_based_routing_files_batches.py
|
||||
tests/unit/proxy/test_model_deprecations_endpoint.py
|
||||
tests/unit/proxy/test_model_dump_with_preserved_fields.py
|
||||
tests/unit/proxy/test_model_id_header_propagation.py
|
||||
tests/unit/proxy/test_model_info_default_limits.py
|
||||
tests/unit/proxy/test_model_level_guardrails.py
|
||||
tests/unit/proxy/test_model_list_aliases.py
|
||||
tests/unit/proxy/test_model_list_callback_filter.py
|
||||
tests/unit/proxy/test_model_list_discoverable.py
|
||||
tests/unit/proxy/test_model_list_healthy_only.py
|
||||
tests/unit/proxy/test_modify_response_streaming_passthrough.py
|
||||
tests/unit/proxy/test_native_compaction.py
|
||||
tests/unit/proxy/test_openai_ws_passthrough_routes.py
|
||||
tests/unit/proxy/test_openapi_schema_validation.py
|
||||
tests/unit/proxy/test_plugin_routes.py
|
||||
tests/unit/proxy/test_pointfive_dashboard_config.py
|
||||
tests/unit/proxy/test_pointfive_ui_callback.py
|
||||
tests/unit/proxy/test_pricing_field_strip.py
|
||||
tests/unit/proxy/test_prisma_engine_watchdog.py
|
||||
tests/unit/proxy/test_prisma_migration.py
|
||||
tests/unit/proxy/test_prometheus_cleanup.py
|
||||
tests/unit/proxy/test_prometheus_metrics_server.py
|
||||
tests/unit/proxy/test_provider_url_destination_guard.py
|
||||
tests/unit/proxy/test_proxy_cli.py
|
||||
tests/unit/proxy/test_proxy_logging_hook_detection.py
|
||||
tests/unit/proxy/test_proxy_types.py
|
||||
tests/unit/proxy/test_pyroscope.py
|
||||
tests/unit/proxy/test_read_model_list.py
|
||||
tests/unit/proxy/test_redis_auth_cache_flag.py
|
||||
tests/unit/proxy/test_response_model_sanitization.py
|
||||
tests/unit/proxy/test_route_a2a_models.py
|
||||
tests/unit/proxy/test_route_llm_request.py
|
||||
tests/unit/proxy/test_route_priority.py
|
||||
tests/unit/proxy/test_sensitive_route_auth.py
|
||||
tests/unit/proxy/test_shared_health_check.py
|
||||
tests/unit/proxy/test_spend_log_cleanup.py
|
||||
tests/unit/proxy/test_swagger_chat_completions.py
|
||||
tests/unit/proxy/test_team_member_update.py
|
||||
tests/unit/proxy/test_team_org_move.py
|
||||
tests/unit/proxy/test_tools_allowlist_enforcement.py
|
||||
tests/unit/proxy/test_tracing_endpoints.py
|
||||
tests/unit/proxy/test_update_llm_router_resilience.py
|
||||
tests/unit/proxy/test_zerobus_dashboard_config.py
|
||||
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
|
||||
tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py
|
||||
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: ""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
89
Makefile
89
Makefile
|
|
@ -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)"
|
||||
|
|
@ -326,89 +327,11 @@ test-unit-proxy-guardrails: install-test-deps
|
|||
test-unit-proxy-core: install-test-deps
|
||||
$(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
|
||||
|
||||
PROXY_INFRA_ROOT_TESTS := \
|
||||
tests/unit/proxy/test__types.py \
|
||||
tests/unit/proxy/test_aiohttp_cleanup_closed.py \
|
||||
tests/unit/proxy/test_aiohttp_session_recovery.py \
|
||||
tests/unit/proxy/test_api_key_masking_in_errors.py \
|
||||
tests/unit/proxy/test_audio_speech_prometheus_hooks.py \
|
||||
tests/unit/proxy/test_batch_expiry.py \
|
||||
tests/unit/proxy/test_batch_metadata_none_fix.py \
|
||||
tests/unit/proxy/test_batch_retrieve_bedrock.py \
|
||||
tests/unit/proxy/test_batch_x_litellm_model_encoding.py \
|
||||
tests/unit/proxy/test_blocked_response_usage.py \
|
||||
tests/unit/proxy/test_body_snapshot_callback_params.py \
|
||||
tests/unit/proxy/test_budget_reservation.py \
|
||||
tests/unit/proxy/test_bug_report_config.py \
|
||||
tests/unit/proxy/test_caching_routes.py \
|
||||
tests/unit/proxy/test_chat_completion_metadata.py \
|
||||
tests/unit/proxy/test_claude_code_marketplace.py \
|
||||
tests/unit/proxy/test_collector.py \
|
||||
tests/unit/proxy/test_common_request_processing.py \
|
||||
tests/unit/proxy/test_component_allowlists.py \
|
||||
tests/unit/proxy/test_conftest.py \
|
||||
tests/unit/proxy/test_cors_config.py \
|
||||
tests/unit/proxy/test_custom_proxy.py \
|
||||
tests/unit/proxy/test_dynamic_mcp_route.py \
|
||||
tests/unit/proxy/test_empty_model_list.py \
|
||||
tests/unit/proxy/test_enforce_user_param.py \
|
||||
tests/unit/proxy/test_fallback_management_endpoints.py \
|
||||
tests/unit/proxy/test_fastapi_offline_routes.py \
|
||||
tests/unit/proxy/test_filter_models_by_team_access_group.py \
|
||||
tests/unit/proxy/test_health_check_functions.py \
|
||||
tests/unit/proxy/test_health_check_max_tokens.py \
|
||||
tests/unit/proxy/test_init_litellm_callbacks.py \
|
||||
tests/unit/proxy/test_langfuse_passthrough_security.py \
|
||||
tests/unit/proxy/test_lazy_openapi_snapshot.py \
|
||||
tests/unit/proxy/test_litellm_pre_call_utils.py \
|
||||
tests/unit/proxy/test_max_budget_env_var.py \
|
||||
tests/unit/proxy/test_mcp_asgi_response.py \
|
||||
tests/unit/proxy/test_model_based_routing_files_batches.py \
|
||||
tests/unit/proxy/test_model_deprecations_endpoint.py \
|
||||
tests/unit/proxy/test_model_dump_with_preserved_fields.py \
|
||||
tests/unit/proxy/test_model_id_header_propagation.py \
|
||||
tests/unit/proxy/test_model_info_default_limits.py \
|
||||
tests/unit/proxy/test_model_level_guardrails.py \
|
||||
tests/unit/proxy/test_model_list_aliases.py \
|
||||
tests/unit/proxy/test_model_list_callback_filter.py \
|
||||
tests/unit/proxy/test_model_list_discoverable.py \
|
||||
tests/unit/proxy/test_model_list_healthy_only.py \
|
||||
tests/unit/proxy/test_modify_response_streaming_passthrough.py \
|
||||
tests/unit/proxy/test_native_compaction.py \
|
||||
tests/unit/proxy/test_openai_ws_passthrough_routes.py \
|
||||
tests/unit/proxy/test_openapi_schema_validation.py \
|
||||
tests/unit/proxy/test_plugin_routes.py \
|
||||
tests/unit/proxy/test_pointfive_dashboard_config.py \
|
||||
tests/unit/proxy/test_pointfive_ui_callback.py \
|
||||
tests/unit/proxy/test_pricing_field_strip.py \
|
||||
tests/unit/proxy/test_prisma_engine_watchdog.py \
|
||||
tests/unit/proxy/test_prisma_migration.py \
|
||||
tests/unit/proxy/test_prometheus_cleanup.py \
|
||||
tests/unit/proxy/test_prometheus_metrics_server.py \
|
||||
tests/unit/proxy/test_provider_url_destination_guard.py \
|
||||
tests/unit/proxy/test_proxy_cli.py \
|
||||
tests/unit/proxy/test_proxy_logging_hook_detection.py \
|
||||
tests/unit/proxy/test_proxy_types.py \
|
||||
tests/unit/proxy/test_pyroscope.py \
|
||||
tests/unit/proxy/test_read_model_list.py \
|
||||
tests/unit/proxy/test_redis_auth_cache_flag.py \
|
||||
tests/unit/proxy/test_response_model_sanitization.py \
|
||||
tests/unit/proxy/test_route_a2a_models.py \
|
||||
tests/unit/proxy/test_route_llm_request.py \
|
||||
tests/unit/proxy/test_route_priority.py \
|
||||
tests/unit/proxy/test_sensitive_route_auth.py \
|
||||
tests/unit/proxy/test_shared_health_check.py \
|
||||
tests/unit/proxy/test_spend_log_cleanup.py \
|
||||
tests/unit/proxy/test_swagger_chat_completions.py \
|
||||
tests/unit/proxy/test_team_member_update.py \
|
||||
tests/unit/proxy/test_team_org_move.py \
|
||||
tests/unit/proxy/test_tools_allowlist_enforcement.py \
|
||||
tests/unit/proxy/test_tracing_endpoints.py \
|
||||
tests/unit/proxy/test_update_llm_router_resilience.py \
|
||||
tests/unit/proxy/test_zerobus_dashboard_config.py
|
||||
|
||||
test-unit-proxy-misc: install-test-deps
|
||||
$(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/test_litellm/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 $(PROXY_INFRA_ROOT_TESTS) 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 --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
|
||||
$(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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
24
litellm-rust/Cargo.lock
generated
24
litellm-rust/Cargo.lock
generated
|
|
@ -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]]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
57
litellm-rust/crates/host-python/src/conversion_cache.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
121
litellm-rust/crates/host-python/tests/conversion_cache.rs
Normal 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
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
|
@ -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]]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -21,4 +21,4 @@ veil.workspace = true
|
|||
rstest.workspace = true
|
||||
tempfile = "3"
|
||||
tokio.workspace = true
|
||||
wiremock = "0.6.5"
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -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
20
litellm-rust/crates/storage-clickhouse/Cargo.toml
Normal 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
5
litellm-rust/crates/storage-clickhouse/README.md
Normal 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
29
litellm-rust/crates/storage-clickhouse/src/error.rs
Normal 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
87
litellm-rust/crates/storage-clickhouse/src/insert.rs
Normal 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
127
litellm-rust/crates/storage-clickhouse/src/lib.rs
Normal 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
113
litellm-rust/crates/storage-clickhouse/src/read.rs
Normal 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
34
litellm-rust/crates/storage-clickhouse/tests/connection.rs
Normal 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
34
litellm-rust/crates/storage-clickhouse/tests/transport.rs
Normal 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)
|
||||
));
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
39
litellm-rust/crates/traces/benches/resource-fanout.rs
Normal 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
13
litellm-rust/crates/traces/query/span_error.sql
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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", ×tamp)?;
|
||||
}
|
||||
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", ×tamp)?;
|
||||
}
|
||||
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)
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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
101
litellm-rust/crates/traces/src/otlp/attributes.rs
Normal 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
212
litellm-rust/crates/traces/src/otlp/limits.rs
Normal 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
42
litellm-rust/crates/traces/src/otlp/mod.rs
Normal 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
166
litellm-rust/crates/traces/src/otlp/span.rs
Normal 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
43
litellm-rust/crates/traces/src/otlp/wire.rs
Normal 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
46
litellm-rust/crates/traces/src/shared.rs
Normal 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()
|
||||
}
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"))]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
¶meters,
|
||||
)
|
||||
.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, ¶meters)
|
||||
.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, ¶meters).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,
|
||||
¶meters,
|
||||
)
|
||||
.await?;
|
||||
let diagnostic =
|
||||
execute_named_read(&database.client, &reader, ReadQuery::SpanError, ¶meters).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(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
23
litellm-rust/crates/traces/tests/shared.rs
Normal 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);
|
||||
}
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -145,7 +145,7 @@ def _a2a_cost_params(litellm_params: Mapping[str, object] | None) -> Mapping[str
|
|||
|
||||
|
||||
def _card_http_kwargs(extra_headers: dict[str, str] | None) -> dict[str, object] | None:
|
||||
return {"headers": extra_headers} if extra_headers else None # mutable-ok: a2a-sdk's get_agent_card takes a dict
|
||||
return {"headers": extra_headers} if extra_headers else None
|
||||
|
||||
|
||||
def _agent_card_path(litellm_params: Mapping[str, object]) -> str | None:
|
||||
|
|
@ -612,7 +612,7 @@ def _build_streaming_logging_obj(
|
|||
logging_obj.model_call_details["agent_id"] = agent_id
|
||||
|
||||
_request_context: Final = (("metadata", metadata), ("proxy_server_request", proxy_server_request))
|
||||
_litellm_params: Final = dict( # mutable-ok: Logging.litellm_params is declared as a dict
|
||||
_litellm_params: Final = dict(
|
||||
(*_a2a_cost_params(litellm_params).items(), *((key, value) for key, value in _request_context if value))
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -127,7 +127,7 @@ async def _handle_completed_batch(
|
|||
return BatchCostUsageResult(
|
||||
cost=0.0,
|
||||
usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
|
||||
models=[], # mutable-ok: no output file means no model was ever priced; BatchCostUsageResult.models requires list[str]
|
||||
models=[],
|
||||
successful_requests=0,
|
||||
failed_requests=await count_error_file_failed_requests(
|
||||
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
|
||||
|
|
|
|||
|
|
@ -103,10 +103,10 @@ async def claim_affinity_pin(
|
|||
try:
|
||||
claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT)
|
||||
args: Final = (
|
||||
json.dumps(dict(pin_value)), # mutable-ok: JSON serialization requires dict, not a generic Mapping
|
||||
json.dumps(dict(pin_value)),
|
||||
int(ttl_seconds),
|
||||
*(
|
||||
(json.dumps(tuple(dict(value) for value in eligible_values)),) # mutable-ok: JSON requires dict
|
||||
(json.dumps(tuple(dict(value) for value in eligible_values)),)
|
||||
if eligible_values is not None
|
||||
else ()
|
||||
),
|
||||
|
|
|
|||
|
|
@ -702,14 +702,14 @@ class LLMCachingHandler:
|
|||
)
|
||||
merged: Final = EmbeddingResponse(
|
||||
model=cached.model,
|
||||
data=[ # mutable-ok: EmbeddingResponse.data is a pydantic list field
|
||||
data=[
|
||||
item
|
||||
if item is not None
|
||||
else Embedding(embedding=next(fresh_items)["embedding"], index=position, object="embedding")
|
||||
for position, item in enumerate(cached.data)
|
||||
],
|
||||
usage=merged_usage,
|
||||
hidden_params={ # mutable-ok: EmbeddingResponse._hidden_params is a mutable dict field
|
||||
hidden_params={
|
||||
**cached._hidden_params,
|
||||
"cache_hit": True,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -252,9 +252,7 @@ class DualCache(BaseCache):
|
|||
if value is not None:
|
||||
self.in_memory_cache.set_cache(key, value, **self._backfill_kwargs(kwargs))
|
||||
|
||||
return list( # mutable-ok: public list contract
|
||||
redis_result.get(key) if value is None else value for key, value in zip(keys, result)
|
||||
)
|
||||
return list(redis_result.get(key) if value is None else value for key, value in zip(keys, result))
|
||||
except Exception as e:
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Cache: exception in batch_get_cache", e, with_traceback=True
|
||||
|
|
@ -329,8 +327,8 @@ class DualCache(BaseCache):
|
|||
def reserve_redis_batch_reads(self, keys: Sequence[str]) -> tuple[list[str], dict[str, float | None]]:
|
||||
"""Reserve the memory-missed keys whose throttled Redis reads are due, as a batch read would."""
|
||||
if self.redis_cache is None:
|
||||
return [], {} # mutable-ok: API contract returns an empty list and dictionary
|
||||
key_list: Final = list(keys) # mutable-ok: batch_get_cache takes a list
|
||||
return [], {}
|
||||
key_list: Final = list(keys)
|
||||
memory: Final = self.in_memory_cache
|
||||
in_memory_result: Final = (
|
||||
None
|
||||
|
|
@ -386,7 +384,7 @@ class DualCache(BaseCache):
|
|||
|
||||
async def declare_batch_get(self, keys: Sequence[str], batch: RedisBatch) -> DeclaredBatchRead:
|
||||
pending: Final = await self._prepare_batch_get(
|
||||
list(keys), # mutable-ok: the shared batch read takes a list
|
||||
list(keys),
|
||||
local_only=False,
|
||||
throttle_redis=False,
|
||||
)
|
||||
|
|
@ -627,7 +625,7 @@ class DualCache(BaseCache):
|
|||
parent_otel_span: Span | None = None,
|
||||
) -> None:
|
||||
batch: Final = None if self.redis_cache is None else active_post_call_redis_batch(self.redis_cache)
|
||||
operations: Final = list(increment_list) # mutable-ok: both increment pipelines take a list
|
||||
operations: Final = list(increment_list)
|
||||
if batch is None:
|
||||
await self.async_increment_cache_pipeline(operations, parent_otel_span=parent_otel_span)
|
||||
return
|
||||
|
|
|
|||
|
|
@ -238,7 +238,7 @@ class EvictedClientCloser:
|
|||
the front rather than having to be searched for.
|
||||
"""
|
||||
with self._queue_lock:
|
||||
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque()) # mutable-ok: FIFO by design
|
||||
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque())
|
||||
while bucket and bucket[0].client_ref() is None:
|
||||
bucket.popleft()
|
||||
self._pending_count -= 1
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ from litellm.types.services import ServiceTypes
|
|||
|
||||
_T = TypeVar("_T")
|
||||
_ScriptArg = str | bytes | int | float
|
||||
SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None] # mutable-ok: Callable params
|
||||
SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None]
|
||||
POST_CALL_FLUSH_DEADLINE_SECONDS: Final = 1.0
|
||||
|
||||
|
||||
|
|
@ -139,7 +139,7 @@ class _MGet(_Op[Mapping[str, object]]):
|
|||
)
|
||||
|
||||
async def run_alone(self) -> Mapping[str, object]:
|
||||
found: Mapping[str, object] = await self._redis_cache.async_batch_get_cache(key_list=list(self._keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API # mutable-ok: the cache API takes a list
|
||||
found: Mapping[str, object] = await self._redis_cache.async_batch_get_cache(key_list=list(self._keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API
|
||||
if any(key not in found for key in self._keys):
|
||||
raise ConnectionError("batch get did not return every key")
|
||||
return found
|
||||
|
|
|
|||
|
|
@ -123,13 +123,13 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]:
|
|||
blocks are the fallback for turns that arrived over another API surface.
|
||||
"""
|
||||
items: Final = _get_reasoning_items(msg)
|
||||
stored: Final = [_reasoning_item_to_response_input(item) for item in items] # mutable-ok: API message payload
|
||||
stored: Final = [_reasoning_item_to_response_input(item) for item in items]
|
||||
if stored:
|
||||
return stored
|
||||
raw_blocks: Final = msg.get("thinking_blocks") or ()
|
||||
blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json
|
||||
replayed: Final = responses_reasoning_items_from_thinking_blocks(blocks)
|
||||
return [dict(item) for item in replayed] # mutable-ok: API message payload
|
||||
return [dict(item) for item in replayed]
|
||||
|
||||
|
||||
def _build_reasoning_item(
|
||||
|
|
@ -441,7 +441,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
input_items.extend(_reasoning_input_items(msg))
|
||||
if content:
|
||||
input_items.append(
|
||||
{ # mutable-ok: API message payload
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": self._convert_content_to_responses_format(content, "assistant"),
|
||||
|
|
@ -475,7 +475,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
if role == "assistant":
|
||||
input_items.extend(_reasoning_input_items(msg))
|
||||
input_items.append(
|
||||
{ # mutable-ok: API message payload
|
||||
{
|
||||
"type": "message",
|
||||
"role": role,
|
||||
"content": self._convert_content_to_responses_format(content, cast(str, role)),
|
||||
|
|
@ -531,11 +531,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
) -> "ResponseText":
|
||||
existing: Final = cast( # cast-ok: text field is a ResponseText | dict[str, Any] | None union
|
||||
"dict[str, object]",
|
||||
dict(responses_api_request).get("text") or {}, # mutable-ok: one-shot merge seed
|
||||
dict(responses_api_request).get("text") or {},
|
||||
)
|
||||
return cast( # cast-ok: merged mapping is a valid ResponseText shape
|
||||
"ResponseText",
|
||||
{**existing, **update}, # mutable-ok: one-shot merged payload
|
||||
{**existing, **update},
|
||||
)
|
||||
|
||||
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]:
|
||||
|
|
@ -1506,7 +1506,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
# tool call; per-stream callers already received it via
|
||||
# output_item.added and the argument delta events
|
||||
return ModelResponseStream(
|
||||
choices=[ # mutable-ok: ModelResponseStream coerces only list choices
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(
|
||||
|
|
@ -1612,7 +1612,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
)
|
||||
],
|
||||
usage=usage,
|
||||
provider_specific_fields=dict(provider_metadata) or None, # mutable-ok: field is typed dict
|
||||
provider_specific_fields=dict(provider_metadata) or None,
|
||||
**(
|
||||
MappingProxyType({"service_tier": served_service_tier})
|
||||
if isinstance(served_service_tier, str)
|
||||
|
|
|
|||
|
|
@ -52,10 +52,10 @@ CLICKHOUSE_MAX_BUFFERED_ROWS: Final = get_env_int("CLICKHOUSE_MAX_BUFFERED_ROWS"
|
|||
CLICKHOUSE_MAX_RETRIES: Final = get_env_int("CLICKHOUSE_MAX_RETRIES", 3)
|
||||
AGENT_TRACING_RETENTION_DAYS: Final = get_env_int("AGENT_TRACING_RETENTION_DAYS", 30)
|
||||
AGENT_TRACING_SPEND_LOG_RETENTION_DAYS: Final = get_env_int("AGENT_TRACING_SPEND_LOG_RETENTION_DAYS", 90)
|
||||
OTLP_MAX_BODY_BYTES: Final = get_env_int("OTLP_MAX_BODY_BYTES", 8 * 1024 * 1024)
|
||||
OTLP_MAX_BODY_BYTES: Final = get_env_int("OTLP_MAX_BODY_BYTES", 16 * 1024 * 1024)
|
||||
OTLP_MAX_ATTRIBUTE_VALUE_BYTES: Final = get_env_int("OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 64 * 1024)
|
||||
OTLP_RETRY_AFTER_SECONDS: Final = get_env_int("OTLP_RETRY_AFTER_SECONDS", 2)
|
||||
OTLP_OFFLOAD_DECODE_BYTES: Final = get_env_int("OTLP_OFFLOAD_DECODE_BYTES", 256 * 1024)
|
||||
OTLP_MAX_CONCURRENT_INGESTS: Final = get_env_int("OTLP_MAX_CONCURRENT_INGESTS", 2)
|
||||
AGENT_TRACING_INPUT_PREVIEW_CHARS: Final = get_env_int("AGENT_TRACING_INPUT_PREVIEW_CHARS", 240)
|
||||
AGENT_TRACING_LIST_PAGE_SIZE: Final = get_env_int("AGENT_TRACING_LIST_PAGE_SIZE", 50)
|
||||
DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
|
||||
|
|
|
|||
|
|
@ -2874,9 +2874,7 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
|
|||
collected_usage_objects: Final = ResponsesWebSocketTokenUsageProcessor.collect_usage_from_responses_ws_results(
|
||||
results
|
||||
)
|
||||
return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(
|
||||
list(collected_usage_objects) # mutable-ok: combine_usage_objects requires a list parameter
|
||||
)
|
||||
return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(list(collected_usage_objects))
|
||||
|
||||
|
||||
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
|
||||
|
|
|
|||
|
|
@ -1136,7 +1136,7 @@ class MCPClient:
|
|||
async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult:
|
||||
capabilities: Final = session.server_capabilities
|
||||
if capabilities is not None and capabilities.resources is None:
|
||||
return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload
|
||||
return ListResourceTemplatesResult(resource_templates=[])
|
||||
try:
|
||||
return ListResourceTemplatesResult(
|
||||
resource_templates=await self._list_optional_pages(
|
||||
|
|
@ -1150,7 +1150,7 @@ class MCPClient:
|
|||
verbose_logger.debug(
|
||||
"MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error
|
||||
)
|
||||
return ListResourceTemplatesResult(resource_templates=[]) # mutable-ok: MCP result payload
|
||||
return ListResourceTemplatesResult(resource_templates=[])
|
||||
|
||||
try:
|
||||
result: Final = await self.run_with_session(_list_resource_templates_operation)
|
||||
|
|
|
|||
|
|
@ -171,9 +171,7 @@ async def load_mcp_tools(
|
|||
"""
|
||||
tools: Final = await list_tools_with_pagination(session)
|
||||
if format == "openai":
|
||||
return [ # mutable-ok: public API returns a list
|
||||
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools
|
||||
]
|
||||
return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools]
|
||||
return tools
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1131,7 +1131,7 @@ Model Info:
|
|||
message=message,
|
||||
level=level,
|
||||
alert_type=AlertType.model_deprecation_warnings,
|
||||
alerting_metadata={ # mutable-ok: send_alert takes a dict payload
|
||||
alerting_metadata={
|
||||
"deprecated_count": len(snapshot.deprecated),
|
||||
"imminent_count": len(snapshot.imminent),
|
||||
"upcoming_count": len(snapshot.upcoming),
|
||||
|
|
@ -1245,8 +1245,8 @@ Model Info:
|
|||
try:
|
||||
existing_invitations: Final = TypeAdapter(list[InvitationModel]).validate_python(
|
||||
await InvitationLinkRepository(prisma_client).table.find_many( # pyright: ignore[reportAny] # untyped prisma boundary (any-ok), result validated by TypeAdapter
|
||||
where={"user_id": recipient_user_id}, # mutable-ok: prisma find_many requires a dict where filter
|
||||
order={"created_at": "desc"}, # mutable-ok: prisma find_many requires a dict order arg
|
||||
where={"user_id": recipient_user_id},
|
||||
order={"created_at": "desc"},
|
||||
),
|
||||
from_attributes=True,
|
||||
)
|
||||
|
|
@ -2011,7 +2011,7 @@ Model Info:
|
|||
message="\n\n".join(event.message for event in typed_events),
|
||||
level="High",
|
||||
alert_type=alert_type,
|
||||
alerting_metadata={}, # mutable-ok: send_alert takes a dict payload
|
||||
alerting_metadata={},
|
||||
)
|
||||
for event in typed_events:
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
|
|
|
|||
|
|
@ -337,7 +337,7 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
Raises a NON Blocking verbose_logger.exception if an error occurs
|
||||
"""
|
||||
batch_to_send: Final = tuple(self.log_queue)
|
||||
self.log_queue = [] # mutable-ok: queue ownership is detached before the async send
|
||||
self.log_queue = []
|
||||
try:
|
||||
undelivered: Final = await self._async_send_batch_to_api(
|
||||
log_queue=batch_to_send,
|
||||
|
|
@ -360,7 +360,7 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
Sends the batch of audit logs to Azure Monitor Logs Ingestion API
|
||||
"""
|
||||
batch_to_send: Final = tuple(self.audit_log_queue)
|
||||
self.audit_log_queue = [] # mutable-ok: queue ownership is detached before the async send
|
||||
self.audit_log_queue = []
|
||||
try:
|
||||
undelivered: Final = await self._async_send_batch_to_api(
|
||||
log_queue=batch_to_send,
|
||||
|
|
@ -384,7 +384,7 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
queue: list[_QueuedPayload],
|
||||
log_type: str,
|
||||
) -> list[_QueuedPayload]:
|
||||
merged: Final = [*undelivered, *queue] # mutable-ok: queue trimming returns a mutable logger queue
|
||||
merged: Final = [*undelivered, *queue]
|
||||
overflow: Final = len(merged) - self.max_queue_size
|
||||
if overflow <= 0:
|
||||
return merged
|
||||
|
|
|
|||
|
|
@ -11,7 +11,8 @@ gzip JSONEachRow insert, either every `CLICKHOUSE_FLUSH_INTERVAL_SECONDS` or as
|
|||
import asyncio
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, ClassVar
|
||||
from contextlib import suppress
|
||||
from typing import Any, ClassVar, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
|
|
@ -21,11 +22,11 @@ from litellm.constants import (
|
|||
CLICKHOUSE_MAX_RETRIES,
|
||||
)
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.rust_bridge.traces import TraceStorage
|
||||
from litellm.rust_bridge.traces import ClickHouseStorage
|
||||
|
||||
|
||||
def clickhouse_storage_from_env() -> TraceStorage:
|
||||
return TraceStorage(
|
||||
def clickhouse_storage_from_env() -> ClickHouseStorage:
|
||||
return ClickHouseStorage(
|
||||
database=os.getenv("CLICKHOUSE_DATABASE", "litellm"),
|
||||
url=os.getenv("CLICKHOUSE_URL", ""),
|
||||
)
|
||||
|
|
@ -34,7 +35,7 @@ def clickhouse_storage_from_env() -> TraceStorage:
|
|||
class ClickHouseBatchLogger(CustomBatchLogger):
|
||||
table: ClassVar[str]
|
||||
|
||||
def __init__(self, storage: TraceStorage | None = None) -> None:
|
||||
def __init__(self, storage: ClickHouseStorage | None = None) -> None:
|
||||
self.storage = storage or clickhouse_storage_from_env()
|
||||
self.rows_written = 0
|
||||
self.rows_dropped = 0
|
||||
|
|
@ -45,11 +46,27 @@ class ClickHouseBatchLogger(CustomBatchLogger):
|
|||
flush_interval=CLICKHOUSE_FLUSH_INTERVAL_SECONDS,
|
||||
)
|
||||
self._flush_task: asyncio.Task[None] | None = None
|
||||
self._stop: Final = asyncio.Event()
|
||||
|
||||
def start(self) -> None:
|
||||
if self._flush_task is None or self._flush_task.done():
|
||||
self._flush_task = asyncio.get_running_loop().create_task(self.periodic_flush())
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._stop.set()
|
||||
if self._flush_task is not None:
|
||||
await self._flush_task
|
||||
while self.log_queue:
|
||||
await self.flush_queue()
|
||||
|
||||
async def periodic_flush(self) -> None:
|
||||
while True:
|
||||
with suppress(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(self._stop.wait(), timeout=self.flush_interval)
|
||||
if self._stop.is_set():
|
||||
return
|
||||
await self.flush_queue()
|
||||
|
||||
def is_full(self) -> bool:
|
||||
"""Backpressure signal: producers should reject (429) instead of enqueueing."""
|
||||
return len(self.log_queue) >= CLICKHOUSE_MAX_BUFFERED_ROWS
|
||||
|
|
|
|||
|
|
@ -58,7 +58,7 @@ def _json(value: object) -> str:
|
|||
|
||||
|
||||
def _json_mapping(value: Mapping[str, Any]) -> str:
|
||||
return _json(dict(value)) # mutable-ok: [LIT002] JSON serialization requires a dict
|
||||
return _json(dict(value))
|
||||
|
||||
|
||||
def _find_traceparent(metadata: Mapping[str, Any], kwargs: Mapping[str, Any]) -> tuple[str, str]:
|
||||
|
|
@ -88,8 +88,8 @@ def _cache_tokens(usage: Mapping[str, Any]) -> tuple[int, int]:
|
|||
|
||||
def _request_tags(value: object) -> list[str]:
|
||||
if not isinstance(value, list):
|
||||
return [] # mutable-ok: [LIT002] empty spend-log tag payload
|
||||
return [str(tag) for tag in value] # mutable-ok: [LIT002] SpendLogRecord schema
|
||||
return []
|
||||
return [str(tag) for tag in value]
|
||||
|
||||
|
||||
def _session_id(payload: StandardLoggingPayload, kwargs: Mapping[str, Any]) -> str:
|
||||
|
|
@ -165,6 +165,6 @@ class ClickHouseSpendLogger(ClickHouseBatchLogger):
|
|||
if payload is None or _is_trace_ingest(payload):
|
||||
return
|
||||
row: Final = spend_log_row_from_payload(payload, kwargs)
|
||||
self.enqueue([dict(row)]) # mutable-ok: [LIT002] batch logger API
|
||||
self.enqueue([dict(row)])
|
||||
except Exception as e:
|
||||
verbose_logger.exception("ClickHouseSpendLogger: failed to log request: %s", e)
|
||||
|
|
|
|||
|
|
@ -1,11 +1,11 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.rust_bridge.traces import TraceStorage
|
||||
from litellm.rust_bridge.traces import ClickHouseStorage
|
||||
|
||||
OTEL_TRACES_TABLE: Final = "otel_traces"
|
||||
AGENT_TRACES_BY_KEY_TABLE: Final = "agent_traces_by_key"
|
||||
SPEND_LOGS_TABLE: Final = "spend_logs"
|
||||
|
||||
|
||||
async def ensure_schema(storage: TraceStorage, trace_retention_days: int, spend_log_retention_days: int) -> None:
|
||||
async def ensure_schema(storage: ClickHouseStorage, trace_retention_days: int, spend_log_retention_days: int) -> None:
|
||||
await storage.ensure_schema(trace_retention_days, spend_log_retention_days)
|
||||
|
|
|
|||
|
|
@ -372,7 +372,7 @@ class CustomGuardrail(CustomLogger):
|
|||
land and degrade to blocking instead of silently letting the
|
||||
flagged request through unmodified.
|
||||
"""
|
||||
advisory_message: Final = {"role": "system", "content": message} # mutable-ok: plain dict for live request
|
||||
advisory_message: Final = {"role": "system", "content": message}
|
||||
existing_messages: Final = data.get("messages")
|
||||
existing_input: Final = data.get("input")
|
||||
existing_instructions: Final = data.get("instructions")
|
||||
|
|
@ -383,7 +383,7 @@ class CustomGuardrail(CustomLogger):
|
|||
# model to disregard a trailing warning. Prefer it over "input"
|
||||
# whenever present.
|
||||
if isinstance(existing_messages, list):
|
||||
messages_with_instructions_note: Final = [ # mutable-ok: fresh list
|
||||
messages_with_instructions_note: Final = [
|
||||
*existing_messages,
|
||||
advisory_message,
|
||||
]
|
||||
|
|
@ -395,7 +395,7 @@ class CustomGuardrail(CustomLogger):
|
|||
# real, read field (e.g. a chat-completions call carrying a stray
|
||||
# "input"), so write to both when both are present.
|
||||
if isinstance(existing_messages, list):
|
||||
messages_with_input_note: Final = [*existing_messages, advisory_message] # mutable-ok: fresh list
|
||||
messages_with_input_note: Final = [*existing_messages, advisory_message]
|
||||
data["messages"] = messages_with_input_note # rebind-ok: mutates caller's dict by design
|
||||
# The Responses API reads "input", not "messages" -- appending only to
|
||||
# "messages" would leave the advisory unreachable for that endpoint.
|
||||
|
|
@ -409,10 +409,10 @@ class CustomGuardrail(CustomLogger):
|
|||
# non-delivery so the caller degrades to blocking.
|
||||
return False
|
||||
if isinstance(existing_messages, list):
|
||||
messages_without_input_note: Final = [*existing_messages, advisory_message] # mutable-ok: fresh list
|
||||
messages_without_input_note: Final = [*existing_messages, advisory_message]
|
||||
data["messages"] = messages_without_input_note # rebind-ok: mutates caller's dict by design
|
||||
return True
|
||||
sole_message: Final = [advisory_message] # mutable-ok: plain list for the live JSON request
|
||||
sole_message: Final = [advisory_message]
|
||||
data["messages"] = sole_message # rebind-ok: mutates caller's dict by design
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -397,7 +397,7 @@ class DataDogLogger(
|
|||
verbose_logger.debug("[DATADOG MOCK] Batch of %s events successfully mocked", len(batch_to_send))
|
||||
|
||||
except BatchSendCancelled as cancelled:
|
||||
self.log_queue = list(cancelled.undelivered) + self.log_queue # mutable-ok: logger queue remains appendable
|
||||
self.log_queue = list(cancelled.undelivered) + self.log_queue
|
||||
raise asyncio.CancelledError() from cancelled
|
||||
except Exception as e:
|
||||
self.log_queue = batch_to_send + self.log_queue
|
||||
|
|
@ -425,7 +425,7 @@ class DataDogLogger(
|
|||
drop_error_message=DD_ERRORS.DATADOG_413_ERROR.value,
|
||||
non_success_handler=requeue_after_http_error,
|
||||
)
|
||||
return list(undelivered) # mutable-ok: caller prepends records to the logger queue
|
||||
return list(undelivered)
|
||||
|
||||
@staticmethod
|
||||
def _exceeds_intake_limits(chunk: Sequence[DatadogPayload]) -> bool:
|
||||
|
|
|
|||
|
|
@ -131,7 +131,7 @@ def _guardrail_entry_without_prompt_carriers(entry: Mapping[str, object]) -> Map
|
|||
Built as an allow-list rather than a deny-list: a key neither set classifies is dropped, so a
|
||||
guardrail that records its own extra detail cannot put the caller's prompt on a redacted span.
|
||||
"""
|
||||
return { # mutable-ok: a fresh record built per entry, handed straight to the span serializer
|
||||
return {
|
||||
field: REDACTED_BY_LITELM_STRING if field in PROMPT_CARRYING_GUARDRAIL_FIELDS else value
|
||||
for field, value in entry.items()
|
||||
if field in _CLASSIFIED_GUARDRAIL_FIELDS
|
||||
|
|
|
|||
|
|
@ -868,10 +868,8 @@ class LangFuseLogger:
|
|||
"id": clean_metadata.pop("generation_id", generation_id),
|
||||
"input": masked_input if not mask_input else "redacted-by-litellm",
|
||||
"output": masked_output if not mask_output else "redacted-by-litellm",
|
||||
"cost_details": {"total": cost} # mutable-ok: langfuse serializes this payload
|
||||
if usage is not None and isinstance(cost, (int, float))
|
||||
else None,
|
||||
"metadata": { # mutable-ok: langfuse serializes this payload, a proxy is not json-encodable
|
||||
"cost_details": {"total": cost} if usage is not None and isinstance(cost, (int, float)) else None,
|
||||
"metadata": {
|
||||
**log_requester_metadata(redact_user_api_key_info(metadata=allowlisted_metadata)), # pyright: ignore[reportArgumentType] # TypedDict in, plain metadata dict out
|
||||
**enrichments,
|
||||
**_lookup_ids(litellm_call_id, response_obj),
|
||||
|
|
|
|||
|
|
@ -109,7 +109,7 @@ def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload)
|
|||
|
||||
def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
|
||||
first: Final = bucket_records[0]
|
||||
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
|
||||
attributes: Final[Mapping[str, str]] = {
|
||||
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
|
||||
for key, value in (
|
||||
("team_id", first.team_id),
|
||||
|
|
@ -150,7 +150,7 @@ def _team_budget_gauges(record: NewRelicMetricRecord) -> tuple[NewRelicMetric, .
|
|||
team_max_budget: Final = record.team_max_budget
|
||||
if team_max_budget is None:
|
||||
return ()
|
||||
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
|
||||
attributes: Final[Mapping[str, str]] = {
|
||||
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
|
||||
for key, value in (("team_id", record.team_id), ("team_alias", record.team_alias))
|
||||
if value
|
||||
|
|
@ -265,7 +265,7 @@ class NewRelicMetricsLogger(CustomBatchLogger):
|
|||
dropped,
|
||||
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
|
||||
)
|
||||
self.log_queue[:] = list(survivors) # mutable-ok: leave late arrivals for the next serialized drain
|
||||
self.log_queue[:] = list(survivors)
|
||||
|
||||
async def _drain_flush_once(self) -> None:
|
||||
"""Attempt every queued record once, in ``batch_size`` chunks, without
|
||||
|
|
|
|||
|
|
@ -384,7 +384,7 @@ def metadata_from_request_data(data: object) -> Mapping[str, object] | None:
|
|||
|
||||
def flatten_metadata(raw: Mapping[str, object]) -> Iterator[tuple[str, str]]:
|
||||
"""Scalar leaves of a nested metadata mapping, keyed by their dotted path."""
|
||||
stack: Final = list(tuple(raw.items())[::-1]) # mutable-ok: iterative worklist keeps the walk off the call stack
|
||||
stack: Final = list(tuple(raw.items())[::-1])
|
||||
while stack:
|
||||
key, value = stack.pop()
|
||||
if (nested := as_str_mapping(value)) is not None:
|
||||
|
|
|
|||
|
|
@ -803,7 +803,7 @@ def _joined_choice(parts: tuple[str, ...]) -> tuple[_Choice, ...]:
|
|||
def _text_completion_choice(choice: Mapping[str, object], text: str) -> Mapping[str, object]:
|
||||
synthesized: Final = _text_choice(text, as_str(choice.get("finish_reason")))
|
||||
merged: Final = (*choice.items(), *synthesized.items())
|
||||
return {k: v for k, v in merged if k != "text"} # mutable-ok: mappers json.dumps and isinstance(dict) it
|
||||
return {k: v for k, v in merged if k != "text"}
|
||||
|
||||
|
||||
def _completion_choices(response: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
|
|
|
|||
|
|
@ -79,7 +79,7 @@ def stream_output(chunks: Sequence[object], data: Mapping[str, object]) -> str |
|
|||
def _assembled_chat_stream(chunks: Sequence[object], data: Mapping[str, object]) -> object:
|
||||
try:
|
||||
return litellm.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType] # upstream types chunks as a bare list
|
||||
chunks=list(chunks), # mutable-ok: stream_chunk_builder takes a list
|
||||
chunks=list(chunks),
|
||||
messages=_MESSAGES.validate_python(data.get("messages")),
|
||||
)
|
||||
except (litellm.APIError, ValidationError):
|
||||
|
|
|
|||
|
|
@ -380,10 +380,8 @@ def inject_trace_context(headers: Mapping[str, str], parent_span: object = None)
|
|||
"""
|
||||
context: Final = _outgoing_trace_context(parent_span)
|
||||
if context is None:
|
||||
return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier
|
||||
carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier
|
||||
key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
|
||||
}
|
||||
return dict(headers)
|
||||
carrier: Final = {key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS}
|
||||
_PROPAGATOR.inject(carrier, context=_propagated_context(headers, context))
|
||||
return carrier
|
||||
|
||||
|
|
|
|||
|
|
@ -636,9 +636,7 @@ class TenantFanOutSpanProcessor(SpanProcessor):
|
|||
live: Final = tuple((id(p), p) for p in (*self._processors.values(), *self._retired.values()))
|
||||
closing: Final = tuple(p for ident, p in live if ident not in self._exporting)
|
||||
self._processors.clear()
|
||||
self._retired = OrderedDict( # mutable-ok: the same bounded map, keeping only what is still exporting
|
||||
(ident, p) for ident, p in live if ident in self._exporting
|
||||
)
|
||||
self._retired = OrderedDict((ident, p) for ident, p in live if ident in self._exporting)
|
||||
for processor in closing:
|
||||
self._drain.submit(processor)
|
||||
self._drain.close(timeout=max(0.0, deadline - time.monotonic()))
|
||||
|
|
|
|||
|
|
@ -171,9 +171,7 @@ class TenantTracerCache:
|
|||
# thread-pool workers concurrently with the event loop, so cache
|
||||
# updates, span counts, and retirement must be atomic.
|
||||
self._lock: Final = threading.Lock()
|
||||
self._providers: OrderedDict[_RouteKey, TracerProvider] = (
|
||||
OrderedDict() # mutable-ok: bounded LRU; eviction needs in-place ordered mutation
|
||||
)
|
||||
self._providers: OrderedDict[_RouteKey, TracerProvider] = OrderedDict()
|
||||
self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state
|
||||
# Oldest-first so an overflow of draining providers sheds the stalest.
|
||||
self._retired: OrderedDict[TracerProvider, None] = OrderedDict() # mutable-ok: draining evicted providers
|
||||
|
|
@ -393,7 +391,7 @@ class TenantTracerCache:
|
|||
if project_headers and kind not in _GRPC_KINDS
|
||||
else base
|
||||
)
|
||||
update: Final = { # mutable-ok: model_copy(update=...) requires a plain dict
|
||||
update: Final = {
|
||||
field: value
|
||||
for field, value in (("headers", routed), ("endpoint", endpoint))
|
||||
if (field == "headers" and routed != spec.headers)
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ def langfuse_preset(
|
|||
if not allow_missing_credentials:
|
||||
raise
|
||||
return base.model_copy(
|
||||
update={ # mutable-ok: pydantic model_copy takes a plain update mapping
|
||||
update={
|
||||
"exporters": credential_gated_exporters(base.exporters, ExporterOwner.LANGFUSE_OTEL),
|
||||
"mapper_names": mappers,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -91,5 +91,5 @@ def signoz_dynamic_headers(
|
|||
) -> dict[str, str]: # mutable-ok: DYNAMIC_HEADERS_BY_CALLBACK returns a dict
|
||||
key: Final = params.get("signoz_ingestion_key")
|
||||
if _tenant_endpoint_is_unusable(params) or not key:
|
||||
return {} # mutable-ok: same registry contract
|
||||
return {"signoz-ingestion-key": key} # mutable-ok: same registry contract
|
||||
return {}
|
||||
return {"signoz-ingestion-key": key}
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ def weave_preset(
|
|||
if not allow_missing_credentials:
|
||||
raise
|
||||
return base.model_copy(
|
||||
update={ # mutable-ok: pydantic model_copy takes a plain update mapping
|
||||
update={
|
||||
"exporters": credential_gated_exporters(base.exporters, ExporterOwner.WEAVE_OTEL),
|
||||
"mapper_names": mappers,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -207,9 +207,7 @@ class PointFiveLogger(CustomBatchLogger):
|
|||
the excluded-field list and this callback's own setting are applied here, then the
|
||||
global, per-request and header settings that only the framework's predicate knows.
|
||||
"""
|
||||
details: Final = self.redact_standard_logging_payload_from_model_call_details(
|
||||
dict(kwargs) # mutable-ok: both framework helpers take the call details as a dict
|
||||
)
|
||||
details: Final = self.redact_standard_logging_payload_from_model_call_details(dict(kwargs))
|
||||
payload: Final = details.get("standard_logging_object")
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -147,7 +147,7 @@ class PointFiveUploadClient:
|
|||
response: Final = await self.http_client.post(
|
||||
self.api_url + path,
|
||||
json=request.model_dump(by_alias=True),
|
||||
headers={ # mutable-ok: AsyncHTTPHandler.post types headers as dict
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
|
|
@ -172,7 +172,7 @@ class PointFiveUploadClient:
|
|||
if isinstance(destination, PointFiveUploadFailure):
|
||||
return destination
|
||||
url, host = destination
|
||||
headers: Final = dict(PUT_HEADERS, Host=host) if host else dict(PUT_HEADERS) # mutable-ok: put wants dict
|
||||
headers: Final = dict(PUT_HEADERS, Host=host) if host else dict(PUT_HEADERS)
|
||||
try:
|
||||
await self.http_client.put(url, data=body, headers=headers, follow_redirects=False)
|
||||
except httpx.HTTPStatusError as e:
|
||||
|
|
|
|||
|
|
@ -1195,7 +1195,7 @@ class PrometheusLogger(CustomLogger):
|
|||
return metric_class(*args, **kwargs)
|
||||
|
||||
kept: Final = tuple(name for name in original_labelnames if name not in self.exclude_labels)
|
||||
kept_kwargs: Final = {**kwargs, "labelnames": kept} # mutable-ok: ** needs a mapping to override labelnames
|
||||
kept_kwargs: Final = {**kwargs, "labelnames": kept}
|
||||
real_metric: Final = metric_class(*args, **kept_kwargs)
|
||||
return _ExcludedLabelMetric(real_metric, original_labelnames, self.exclude_labels)
|
||||
|
||||
|
|
|
|||
|
|
@ -642,7 +642,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
#########################################################
|
||||
uploads: Final = self._batch_file_elements(batch) if self._batch_file_mode_active() else batch
|
||||
self._flush_retries = 0
|
||||
self._flush_dropped = {} # mutable-ok: per-flush drop marks read back by _upload_bounded
|
||||
self._flush_dropped = {}
|
||||
stale: Final = min(self._requeued_count, len(uploads)) if len(uploads) == len(batch) else 0
|
||||
order: Final = (*range(stale, len(uploads)), *range(stale))
|
||||
ordered: Final = await asyncio.gather(*(self._upload_outcome(uploads[i]) for i in order))
|
||||
|
|
@ -694,7 +694,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
self.max_queue_size,
|
||||
overflow,
|
||||
)
|
||||
self.log_queue = [ # mutable-ok: log_queue is the flush buffer shared with custom_batch_logger
|
||||
self.log_queue = [
|
||||
*requeued,
|
||||
*arrivals,
|
||||
][overflow:]
|
||||
|
|
|
|||
|
|
@ -357,11 +357,7 @@ class GuardrailRequestSnapshot:
|
|||
if fingerprint is None:
|
||||
return None
|
||||
return GuardrailRequestSnapshot(
|
||||
body=MappingProxyType(
|
||||
_CHAT_REQUEST_ADAPTER.validate_python(
|
||||
independent_snapshot(dict(body)) # mutable-ok: snapshot helper requires a plain dictionary
|
||||
)
|
||||
),
|
||||
body=MappingProxyType(_CHAT_REQUEST_ADAPTER.validate_python(independent_snapshot(dict(body)))),
|
||||
fingerprint=fingerprint,
|
||||
)
|
||||
|
||||
|
|
@ -813,7 +809,7 @@ def _as_active_job(record: object, attempts: int, spend: float) -> ActiveShadowE
|
|||
except ValidationError as e:
|
||||
verbose_logger.debug("shadow_eval: skipping unsamplable job row: %s", e)
|
||||
return None
|
||||
return job.model_copy(update={"attempts": attempts, "spend": spend}) # mutable-ok: pydantic update payload
|
||||
return job.model_copy(update={"attempts": attempts, "spend": spend})
|
||||
|
||||
|
||||
_jobs_cache: Final = InMemoryCache(max_size_in_memory=4, default_ttl=_JOBS_CACHE_TTL_SECONDS)
|
||||
|
|
@ -864,9 +860,9 @@ class ShadowEvalLogger(CustomLogger):
|
|||
return _EMPTY_JOBS
|
||||
try:
|
||||
records: Final = await prisma.db.litellm_shadowevaljob.find_many(
|
||||
where={ # mutable-ok: Prisma filter
|
||||
where={
|
||||
"stopped_at": None,
|
||||
"ends_at": {"gt": datetime.now(timezone.utc)}, # mutable-ok: Prisma filter
|
||||
"ends_at": {"gt": datetime.now(timezone.utc)},
|
||||
},
|
||||
)
|
||||
grouped: Final = (
|
||||
|
|
@ -874,12 +870,12 @@ class ShadowEvalLogger(CustomLogger):
|
|||
by=["job_id"],
|
||||
count=True,
|
||||
sum={"judge_cost": True, "shadow_cost": True, "shadow_classifier_cost": True},
|
||||
where={"job_id": {"in": [str(record.id) for record in records]}}, # mutable-ok: Prisma filter
|
||||
where={"job_id": {"in": [str(record.id) for record in records]}},
|
||||
)
|
||||
if records
|
||||
else ()
|
||||
)
|
||||
attempt_stats: Final = { # mutable-ok: frozen snapshot of the grouped read
|
||||
attempt_stats: Final = {
|
||||
str(row["job_id"]): (
|
||||
int(row["_count"]["_all"]),
|
||||
_leg_eval_spend(row["_sum"] or _EMPTY_METADATA),
|
||||
|
|
@ -952,13 +948,13 @@ class ShadowEvalLogger(CustomLogger):
|
|||
payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # untyped callback kwargs
|
||||
if payload is None:
|
||||
return
|
||||
raw_meta: Final = get_litellm_metadata_from_kwargs(dict(kwargs)) # mutable-ok: helper needs dict
|
||||
raw_meta: Final = get_litellm_metadata_from_kwargs(dict(kwargs))
|
||||
request_metadata: Final = raw_meta if isinstance(raw_meta, Mapping) else _EMPTY_METADATA
|
||||
if request_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
|
||||
return # internal sub-call (our own shadow/judge, a classifier), not user traffic
|
||||
# redaction rewrites logged content before callbacks run, so this hook
|
||||
# only ever sees placeholders for a redacted request
|
||||
if should_redact_message_logging(dict(kwargs)): # mutable-ok: predicate takes a plain dict
|
||||
if should_redact_message_logging(dict(kwargs)):
|
||||
return
|
||||
metadata: Final = payload.get("metadata") or _EMPTY_METADATA
|
||||
# Each identity the request resolved to is a candidate target; JWT-auth
|
||||
|
|
@ -999,7 +995,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
sample: Final = _judgeable_sample(
|
||||
ops,
|
||||
sample_kwargs,
|
||||
MappingProxyType(dict(payload.get("model_parameters") or {})), # mutable-ok: frozen snapshot
|
||||
MappingProxyType(dict(payload.get("model_parameters") or {})),
|
||||
response_obj,
|
||||
)
|
||||
if sample is None:
|
||||
|
|
@ -1246,7 +1242,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
return
|
||||
try:
|
||||
await prisma.db.litellm_shadowevalattempt.create(
|
||||
data={ # mutable-ok: Prisma payload
|
||||
data={
|
||||
"job_id": job.id,
|
||||
"request_id": request_id,
|
||||
"router_name": router_name,
|
||||
|
|
@ -1287,12 +1283,10 @@ class ShadowEvalLogger(CustomLogger):
|
|||
try:
|
||||
response: Final = await router.acompletion(
|
||||
model=target_model,
|
||||
messages=[ # mutable-ok: provider transforms rewrite messages in place, so the router gets its own copy
|
||||
dict(m) for m in messages
|
||||
], # pyright: ignore[reportArgumentType] # snapshot of the SDK's own message dicts
|
||||
messages=[dict(m) for m in messages], # pyright: ignore[reportArgumentType] # snapshot of the SDK's own message dicts
|
||||
metadata=shadow_metadata,
|
||||
num_retries=0,
|
||||
fallbacks=[], # mutable-ok: SDK kwarg; a failed shadow is a recorded error, never a spend multiplier
|
||||
fallbacks=[],
|
||||
**shadow_params,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # provider errors become error rows, not crashes
|
||||
|
|
@ -1341,8 +1335,8 @@ class ShadowEvalLogger(CustomLogger):
|
|||
if m.get("content") is not None
|
||||
)
|
||||
judge_metadata: Final = sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_JUDGE_CALL_ORIGIN)
|
||||
judge_messages: Final = [ # mutable-ok: SDK takes a list
|
||||
{"role": "system", "content": PAIRWISE_JUDGE_SYSTEM_PROMPT}, # mutable-ok: SDK message
|
||||
judge_messages: Final = [
|
||||
{"role": "system", "content": PAIRWISE_JUDGE_SYSTEM_PROMPT},
|
||||
{
|
||||
"role": "user",
|
||||
"content": _judge_user_prompt(conversation, response_a, response_b, _tool_definitions_text(tools)),
|
||||
|
|
|
|||
|
|
@ -1683,7 +1683,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = (
|
||||
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth)
|
||||
)
|
||||
return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches
|
||||
return {
|
||||
**user_api_key_metadata,
|
||||
**parent_correlation.as_search_metadata(),
|
||||
"model_group": search_tool_name,
|
||||
|
|
|
|||
|
|
@ -188,9 +188,7 @@ class ZerobusLogger(CustomBatchLogger):
|
|||
|
||||
def _payload_for(self, kwargs: Mapping[str, object]) -> Mapping[str, object] | None:
|
||||
"""The payload to buffer, redacted the way the framework redacts the success path."""
|
||||
details: Final = self.redact_standard_logging_payload_from_model_call_details(
|
||||
dict(kwargs) # mutable-ok: both framework helpers take the call details as a dict
|
||||
)
|
||||
details: Final = self.redact_standard_logging_payload_from_model_call_details(dict(kwargs))
|
||||
payload: Final = details.get("standard_logging_object")
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ def build_agentic_followup_kwargs(
|
|||
fingerprint: str,
|
||||
) -> Mapping[str, object]:
|
||||
"""Kwargs for an agentic follow-up call: the request's kwargs overlaid by the plan's, never repeating a key already sent as a request param"""
|
||||
seen: Final = [*fingerprints, fingerprint] # mutable-ok: the chat loop's settings reader only accepts a list
|
||||
seen: Final = [*fingerprints, fingerprint]
|
||||
return MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue