scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-14 04:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 97

Merge remote-tracking branch 'origin/main' into litellm_mcp_semantic_index_redis

Browse source
This commit is contained in:
mrinal 2026-10-10 23:18:11 +00:00
commit 325b4653b5
369 changed files with 34338 additions and 20056 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

400
.circleci/config.yml
View file

@ -473,7 +473,7 @@ jobs:
command: |
python3 tests/windows_tests/check_windows_wheel_install.py --lengths-only
local_testing_part1:
auth_ui_unit_tests:
docker:
- &python312_image
image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
@ -489,219 +489,6 @@ jobs:
working_directory: ~/project
environment: &sidecar_database_url
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
parallelism: 4
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
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
- save_cargo_target
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
# Run pytest and generate JUnit XML report
- run:
name: Run tests (Part 1 - A-M)
command: |
mkdir test-results
# Discover test files (A-M)
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_[a-mA-M]*.py")
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 --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
-k \"not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache\" \
-n 4 \
--timeout=300 \
--timeout_method=thread"
no_output_timeout: 15m
- run:
name: Rename the coverage files
command: |
# When CI reruns only the failed tests, a parallel node can receive
# zero tests and pytest never writes coverage. Emit empty placeholders
# so persist_to_workspace and the downstream coverage combine stay green.
if [ -f coverage.xml ]; then
mv coverage.xml local_testing_part1_coverage.xml
mv .coverage local_testing_part1_coverage
else
touch local_testing_part1_coverage.xml local_testing_part1_coverage
fi
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- local_testing_part1_coverage.xml
- local_testing_part1_coverage
local_testing_part2:
docker:
- *python312_image
- *postgres16_sidecar
working_directory: ~/project
environment: *sidecar_database_url
parallelism: 4
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
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
# Run pytest and generate JUnit XML report
- run:
name: Run tests (Part 2 - N-Z)
command: |
mkdir test-results
# Discover test files (N-Z)
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_[n-zN-Z]*.py")
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 --tb=short \
-vv \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
-k \"not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache\" \
-n 4 \
--timeout=300 \
--timeout_method=thread"
no_output_timeout: 15m
- run:
name: Rename the coverage files
command: |
# When CI reruns only the failed tests, a parallel node can receive
# zero tests and pytest never writes coverage. Emit empty placeholders
# so persist_to_workspace and the downstream coverage combine stay green.
if [ -f coverage.xml ]; then
mv coverage.xml local_testing_part2_coverage.xml
mv .coverage local_testing_part2_coverage
else
touch local_testing_part2_coverage.xml local_testing_part2_coverage
fi
# Store test results
- store_test_results:
path: test-results
- persist_to_workspace:
root: .
paths:
- local_testing_part2_coverage.xml
- local_testing_part2_coverage
langfuse_logging_unit_tests:
docker:
- *python312_image
- *postgres16_sidecar
working_directory: ~/project
environment: *sidecar_database_url
resource_class: medium
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
- setup_litellm_enterprise_pip
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Run prisma ./docker/entrypoint.sh
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
set -e
# Run pytest and generate JUnit XML report
- run:
name: Run tests
command: |
mkdir -p test-results
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 --tb=short \
-v \
--junitxml=test-results/junit.xml \
--durations=5 \
-k \"langfuse\""
no_output_timeout: 15m
# Store test results
- store_test_results:
path: test-results
auth_ui_unit_tests:
docker:
- *python312_image
- *postgres16_sidecar
working_directory: ~/project
environment: *sidecar_database_url
steps:
- checkout
@ -759,59 +546,6 @@ jobs:
- auth_ui_unit_tests_coverage.xml
- auth_ui_unit_tests_coverage
litellm_router_testing: # Runs all tests with the "router" keyword
docker:
- *python312_image
working_directory: ~/project
resource_class: large
parallelism: 1
environment:
FAKE_OPENAI_API_BASE: http://127.0.0.1:8190
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- 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" }}
- start_fake_openai_endpoint
# Run pytest and generate JUnit XML report
- setup_litellm_enterprise_pip
- run:
name: Run tests
command: |
# On a "rerun failed tests" build a parallel node can receive no
# tests, so the test command never creates test-results. Pre-create it
# so store_test_results doesn't fail the node on a missing path.
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/local_testing/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--split-by=timings \
--verbose \
--command="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 \
--junitxml=test-results/junit.xml \
--durations=5 \
--timeout=300 --timeout_method=thread"
no_output_timeout: 15m
# Store test results
- store_test_results:
path: test-results
llm_translation_testing:
docker:
- *python312_image
@ -2222,7 +1956,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 google_generate_content_endpoint_coverage litellm_utils_coverage
uv tool run --from 'coverage[toml]==7.10.6' coverage combine realtime_translation_coverage ocr_coverage search_coverage logging_coverage audio_coverage pass_through_unit_tests_coverage batches_coverage google_generate_content_endpoint_coverage litellm_utils_coverage
uv tool run --from 'coverage[toml]==7.10.6' coverage xml
- codecov/upload:
file: ./coverage.xml
@ -2747,33 +2481,6 @@ jobs:
- store_test_results:
path: test-results/provider-replay-harness
integration_timing_snapshot:
machine:
image: ubuntu-2204:2024.04.1
resource_class: medium
working_directory: ~/project
steps:
- checkout
- run:
name: Prepare optional integration timing snapshots
command: mkdir -p /tmp/integration-timings/{extensions,providers,mcp}
- restore_cache:
keys:
- integration-timings-v1-extensions-{{ .Branch }}-{{ .Revision }}
- integration-timings-v1-extensions-{{ .Branch }}-
- restore_cache:
keys:
- integration-timings-v1-providers-{{ .Branch }}-{{ .Revision }}
- integration-timings-v1-providers-{{ .Branch }}-
- restore_cache:
keys:
- integration-timings-v1-mcp-{{ .Branch }}-{{ .Revision }}
- integration-timings-v1-mcp-{{ .Branch }}-
- persist_to_workspace:
root: /tmp
paths:
- integration-timings
integration_contracts:
parameters:
suite:
@ -2796,15 +2503,6 @@ jobs:
parallelism: << parameters.parallelism >>
steps:
- setup_litellm_test_deps
- when:
condition:
or:
- equal: [extensions, << parameters.suite >>]
- equal: [providers, << parameters.suite >>]
- equal: [mcp, << parameters.suite >>]
steps:
- attach_workspace:
at: /tmp
- when:
condition:
equal: [browser, << parameters.suite >>]
@ -2850,24 +2548,26 @@ jobs:
test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)"
- when:
condition:
or:
- equal: [extensions, << parameters.suite >>]
- equal: [providers, << parameters.suite >>]
- equal: [mcp, << parameters.suite >>]
not:
equal: [1, << parameters.parallelism >>]
steps:
- run:
name: Prepare this shard's execution record
command: |
mkdir -p "/tmp/integration-shards/<< parameters.suite >>/$CIRCLE_NODE_INDEX"
cp test-results/integration-<< parameters.suite >>/execution.json "/tmp/integration-shards/<< parameters.suite >>/$CIRCLE_NODE_INDEX/execution.json"
results="test-results/integration-<< parameters.suite >>"
if [ "<< parameters.mode >>" = replica ]; then
results="$results-replica"
fi
mkdir -p "/tmp/integration-shards/<< parameters.suite >>-<< parameters.mode >>/$CIRCLE_NODE_INDEX"
cp "$results/execution.json" "/tmp/integration-shards/<< parameters.suite >>-<< parameters.mode >>/$CIRCLE_NODE_INDEX/execution.json"
if [ "<< parameters.suite >>" = mcp ]; then
mkdir -p "/tmp/integration-shards/mcp/$CIRCLE_NODE_INDEX/coverage"
cp test-results/integration-mcp/coverage/data "/tmp/integration-shards/mcp/$CIRCLE_NODE_INDEX/coverage/data.$CIRCLE_NODE_INDEX"
mkdir -p "/tmp/integration-shards/mcp-standard/$CIRCLE_NODE_INDEX/coverage"
cp test-results/integration-mcp/coverage/data "/tmp/integration-shards/mcp-standard/$CIRCLE_NODE_INDEX/coverage/data.$CIRCLE_NODE_INDEX"
fi
- persist_to_workspace:
root: /tmp/integration-shards
paths:
- << parameters.suite >>
- << parameters.suite >>-<< parameters.mode >>
- store_test_results:
path: test-results
- store_artifacts:
@ -2877,7 +2577,11 @@ jobs:
parameters:
suite:
type: enum
enum: [extensions, providers, mcp]
enum: [management, accounting, database, extensions, providers, mcp]
mode:
type: enum
enum: [standard, replica]
default: standard
shards:
type: integer
machine:
@ -2891,12 +2595,8 @@ jobs:
- run:
name: Verify every << parameters.suite >> case ran exactly once
command: |
.venv/bin/python -m tests.integration._support.sharding /tmp/integration-shards/<< parameters.suite >> \
--shards << parameters.shards >> --timings-output /tmp/integration-timings/<< parameters.suite >>/timings.json
- save_cache:
key: integration-timings-v1-<< parameters.suite >>-{{ .Branch }}-{{ .Revision }}
paths:
- /tmp/integration-timings/<< parameters.suite >>
.venv/bin/python -m tests.integration._support.sharding /tmp/integration-shards/<< parameters.suite >>-<< parameters.mode >> \
--shards << parameters.shards >>
- when:
condition:
equal: [mcp, << parameters.suite >>]
@ -2906,7 +2606,7 @@ jobs:
command: |
mkdir -p test-results/integration-mcp/coverage
export COVERAGE_FILE="$PWD/test-results/integration-mcp/coverage/data"
.venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml /tmp/integration-shards/mcp/*/coverage
.venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml /tmp/integration-shards/mcp-standard/*/coverage
.venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml > test-results/integration-mcp/coverage/coverage.txt
.venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml -d test-results/integration-mcp/coverage/html
- store_artifacts:
@ -3205,48 +2905,59 @@ workflows:
- not:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- integration_timing_snapshot:
name: integration-timing-snapshot
- integration_contracts:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [sdk, cost, browser]
parallelism: 1
- integration_contracts:
name: integration-security
suite: security
parallelism: 1
resource_class: xlarge
- integration_contracts:
name: integration-management
suite: management
parallelism: 8
parallelism: 4
- integration_shard_coverage:
name: integration-management-coverage
suite: management
shards: 4
requires: [integration-management]
- integration_contracts:
name: integration-accounting
suite: accounting
parallelism: 8
parallelism: 4
- integration_shard_coverage:
name: integration-accounting-coverage
suite: accounting
shards: 4
requires: [integration-accounting]
- integration_contracts:
name: integration-providers
suite: providers
parallelism: 16
requires:
- integration-timing-snapshot
parallelism: 8
- integration_shard_coverage:
name: integration-providers-coverage
suite: providers
shards: 16
shards: 8
requires:
- integration-providers
- integration_contracts:
name: integration-database
suite: database
parallelism: 2
- integration_shard_coverage:
name: integration-database-coverage
suite: database
shards: 2
requires: [integration-database]
- integration_contracts:
name: integration-mcp
suite: mcp
resource_class: xlarge
parallelism: 2
requires:
- integration-timing-snapshot
- integration_shard_coverage:
name: integration-mcp-coverage
suite: mcp
@ -3256,25 +2967,35 @@ workflows:
- integration_contracts:
name: integration-extensions
suite: extensions
parallelism: 32
requires:
- integration-timing-snapshot
parallelism: 8
- integration_shard_coverage:
name: integration-extensions-coverage
suite: extensions
shards: 32
shards: 8
requires:
- integration-extensions
- integration_contracts:
name: integration-management-replica
suite: management
mode: replica
parallelism: 8
parallelism: 4
- integration_shard_coverage:
name: integration-management-replica-coverage
suite: management
mode: replica
shards: 4
requires: [integration-management-replica]
- integration_contracts:
name: integration-database-replica
suite: database
mode: replica
parallelism: 2
- integration_shard_coverage:
name: integration-database-replica-coverage
suite: database
mode: replica
shards: 2
requires: [integration-database-replica]
- postgres_suite:
name: proxy-behavior
test_path: tests/proxy_behavior
@ -3319,10 +3040,6 @@ workflows:
equal: ["", << pipeline.parameters.routing_parity_base >>]
jobs:
- provider_replay_harness
- local_testing_part1
- local_testing_part2
- langfuse_logging_unit_tests
- litellm_router_testing
- auth_ui_unit_tests
- build_docker_database_image
- e2e_ui_testing
@ -3372,9 +3089,6 @@ workflows:
- image_gen_testing
- logging_testing
- audio_testing
- langfuse_logging_unit_tests
- local_testing_part1
- local_testing_part2
- db_migration_disable_update_check:
requires:
- build_docker_database_image

21
.circleci/scripts/run_integration.sh
View file

@ -216,18 +216,11 @@ if [ "$suite" = browser ]; then
exit 0
fi
node_files=()
case_shards=1
case_shard_index=0
if [ "$suite" = extensions ] || [ "$suite" = providers ] || [ "$suite" = mcp ]; then
case_shards="${CIRCLE_NODE_TOTAL:-1}"
case_shard_index="${CIRCLE_NODE_INDEX:-0}"
elif [ "${CIRCLE_NODE_TOTAL:-1}" -gt 1 ]; then
split="$(.venv/bin/python tests/integration/run.py "$suite" --list \
| circleci tests split --split-by=timings --timings-type=filename)"
read -r -a node_files <<< "$(printf '%s' "$split" | tr '\n' ' ')"
test "${#node_files[@]}" -gt 0
printf '%s\n' "${node_files[@]}" > "$results/node-files.txt"
shard_arguments=()
if [ "${CIRCLE_NODE_TOTAL:-1}" -gt 1 ]; then
.venv/bin/python tests/integration/run.py "$suite" --list \
| circleci tests split --split-by=timings --timings-type=filename > "$results/node-files.txt"
shard_arguments=(--shard-files "$results/node-files.txt")
fi
env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
@ -240,13 +233,13 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
INTEGRATION_CASE_SHARDS="$case_shards" INTEGRATION_CASE_SHARD_INDEX="$case_shard_index" \
INTEGRATION_CASE_SHARDS="${CIRCLE_NODE_TOTAL:-1}" INTEGRATION_CASE_SHARD_INDEX="${CIRCLE_NODE_INDEX:-0}" \
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
INTEGRATION_PROXY_DATABASE_URL="$INTEGRATION_PROXY_DATABASE_URL" \
INTEGRATION_PROXY_READ_REPLICA_URL="$INTEGRATION_PROXY_READ_REPLICA_URL" \
INTEGRATION_ROUTING="$INTEGRATION_ROUTING" \
.venv/bin/python tests/integration/run.py "$suite" --results "$results" \
--timings "/tmp/integration-timings/$suite/timings.json" "${node_files[@]}"
"${shard_arguments[@]}"
if [ "${INTEGRATION_COVERAGE:-0}" = 1 ]; then
for covered_pid in "$proxy_pid" "$peer_pid"; do

BIN
.github/assets/scaledown/provider-after.jpg vendored Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 108 KiB

BIN
.github/assets/scaledown/provider-before.jpg vendored Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 109 KiB

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

@ -18,9 +18,8 @@ test_paths:
paths:
- tests/rust-python-harness
- reason: >-
Live-provider caching cases in tests/local_testing that remain outside CI. Jobs that
glob that directory either deselect them (local_testing_part1 and part2 carry `-k "... and
not caching and not cache"`) or keep only another keyword (langfuse, router, assistants).
Live-provider caching cases in tests/local_testing that remain outside CI. No CircleCI job
globs that directory anymore, and the legacy jobs that did always deselected these files.
The two Redis Cluster IAM authentication tests moved to tests/integration/sdk and remain
selected by node ID in redis_compat, where redis-server is installed.
The five keyless files now run in the caching-local shard, leaving live cases in these two.
@ -30,14 +29,6 @@ test_paths:
paths:
- tests/local_testing/test_disk_cache_unit_tests.py
- tests/local_testing/test_gcs_cache_unit_tests.py
- reason: >-
Mixed file that no job has ever run in full. Every CircleCI job that globs
tests/local_testing deselects it by name ("caching") or keeps only another keyword, so only
its router test ran, and that test now lives in tests/unit/router_utils/pre_call_checks.
Its two offline tests now live in tests/unit/llms/anthropic/chat, and the five left all need
ANTHROPIC_API_KEY or Vertex credentials
paths:
- tests/local_testing/test_anthropic_prompt_caching.py
- reason: >-
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
from a pull request; it needs a live gateway and provider credentials no PR job holds

16
.github/workflows/publish-lint-base-counts.yml vendored
View file

@ -35,6 +35,8 @@ jobs:
script: scripts/type_discipline_gate.py
- checker: test-quality
script: scripts/test_quality_gate.py
- checker: openapi-docs
script: scripts/openapi_docs_gate.py
- checker: basedpyright
script: scripts/type_check_gate.py
@ -60,16 +62,24 @@ jobs:
uses: ./.github/actions/cache-cargo-build
- name: Cache Prisma binaries
if: matrix.checker == 'basedpyright'
if: matrix.checker == 'basedpyright' || matrix.checker == 'openapi-docs'
uses: ./.github/actions/cache-prisma-binaries
# The three source scanners only need the pinned dev tools (ruff and the
# The source scanners only need the pinned dev tools (ruff and the
# stdlib checkers), the same versions test-linting.yml's python job runs.
- name: Install the dev tools
if: matrix.checker != 'basedpyright'
if: matrix.checker != 'basedpyright' && matrix.checker != 'openapi-docs'
run: |
uv sync --frozen --only-group dev --no-install-project
# The OpenAPI gate imports the proxy app, so it needs the same install the
# test-linting.yml python job uses, Prisma client included.
- name: Install the proxy
if: matrix.checker == 'openapi-docs'
run: |
uv sync --frozen --group proxy-dev --group e2e-dev --extra proxy
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Emit ${{ matrix.checker }} counts for HEAD
if: matrix.checker != 'basedpyright'
run: |

8
.github/workflows/test-linting.yml vendored
View file

@ -156,6 +156,14 @@ jobs:
run: |
uv run --no-sync python scripts/test_quality_gate.py --base "$GATE_BASE_SHA"
- name: Check OpenAPI docs (POST routes without a request schema, description or 200 schema, delta vs merge-base counts)
if: steps.changes.outputs.decision != 'skip'
env:
GH_TOKEN: ${{ github.token }}
run: |
uv sync --frozen --group proxy-dev --group e2e-dev --extra proxy
uv run --no-sync python scripts/openapi_docs_gate.py --base "$GATE_BASE_SHA"
- name: Print OpenAI version
if: steps.changes.outputs.decision != 'skip'
run: |

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

@ -140,6 +140,7 @@ jobs:
tests/unit/batch_completion
tests/unit/decisions
tests/unit/litellm_core_utils
tests/unit/systemone
python-version: "3.12"
workers: 4
reruns: 0
@ -327,6 +328,7 @@ jobs:
tests/unit/llms/sagemaker
tests/unit/llms/sail
tests/unit/llms/sambanova
tests/unit/llms/scaledown
tests/unit/llms/sap
tests/unit/llms/scaleway
tests/unit/llms/searchapi

2
AGENTS.md
View file

@ -54,7 +54,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
The four lint gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/test_quality_gate.py`, `scripts/type_check_gate.py`) compare each rule's codebase count on your branch against the count at its merge-base with the default branch, and a rule may not grow. There are no budget files to edit or ratchet: when a gate fails, fix the violations the branch introduced or remove at least as many of that rule elsewhere in the tree. The one exception is `reportAny` / `reportExplicitAny`, which share a fixed codebase-wide cap in `ANY_CAPS` in `scripts/type_check_gate.py` because Any spreads past the lines you touch. A branch may add Anys while the total stays under that cap. Only lower `ANY_CAPS`, and only in its own PR to the default branch, never raise it on a feature branch
The five lint gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/test_quality_gate.py`, `scripts/openapi_docs_gate.py`, `scripts/type_check_gate.py`) compare each rule's codebase count on your branch against the count at its merge-base with the default branch, and a rule may not grow. There are no budget files to edit or ratchet: when a gate fails, fix the violations the branch introduced or remove at least as many of that rule elsewhere in the tree. The one exception is `reportAny` / `reportExplicitAny`, which share a fixed codebase-wide cap in `ANY_CAPS` in `scripts/type_check_gate.py` because Any spreads past the lines you touch. A branch may add Anys while the total stays under that cap. Only lower `ANY_CAPS`, and only in its own PR to the default branch, never raise it on a feature branch
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice

11
Makefile
View file

@ -7,7 +7,7 @@
test-rust-extension rust-sqlx-prepare \
info lint lint-inner lint-dev lint-checks format \
lint-basedpyright lint-e2e-basedpyright lint-type-discipline \
lint-ruff-strict lint-gate lint-test-quality \
lint-ruff-strict lint-gate lint-test-quality lint-openapi-docs \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \
lint-install lint-fetch-base bootstrap
@ -35,6 +35,7 @@ help:
@echo " make lint-ruff-strict - Gate each strict ruff rule's codebase total against its merge-base count"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)"
@echo " make lint-test-quality - Gate the test suite's TQ counts against their merge-base counts"
@echo " make lint-openapi-docs - Gate undocumented POST routes in the proxy OpenAPI document against their merge-base counts"
@echo " make check-circular-imports - Check for circular imports"
@echo " make check-import-safety - Check import safety"
@echo " make test - Run all tests"
@ -219,6 +220,12 @@ lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)"
# OpenAPI docs gate (POST routes without a request body schema, a description or a
# 200 schema in app.openapi()), counted the same delta-vs-base way so a new endpoint
# has to show up documented in Swagger UI.
lint-openapi-docs: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/openapi_docs_gate.py --base "$(BASE_REF)"
lint-format: format-check
lint-ruff-strict: install-dev
@ -250,7 +257,7 @@ lint-inner: lint-install
@base_ref=$$($(RESOLVE_BASE)) && \
$(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-openapi-docs lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed check-circular-imports check-import-safety

1
README.md
View file

@ -407,6 +407,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th
| [Sagemaker Chat (`sagemaker_chat`)](https://docs.litellm.ai/docs/providers/aws_sagemaker) | ✅ | ✅ | ✅ | | | | | | | |
| [Sail (`sail`)](https://docs.litellm.ai/docs/providers/sail) | ✅ | ✅ | ✅ | | | | | | | |
| [Sambanova (`sambanova`)](https://docs.litellm.ai/docs/providers/sambanova) | ✅ | ✅ | ✅ | | | | | | | |
| [ScaleDown (`scaledown`)](https://docs.litellm.ai/docs/providers/scaledown) | ✅ | | | | | | | | | |
| [Snowflake (`snowflake`)](https://docs.litellm.ai/docs/providers/snowflake) | ✅ | ✅ | ✅ | | | | | | | |
| [Strands Decider (`strands_decider`)](https://docs.litellm.ai/docs/providers) | | | | | | | | | | |
| [Text Completion Codestral (`text-completion-codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | |

16
litellm/__init__.py
View file

@ -674,6 +674,7 @@ nvidia_nim_models: Set = set()
nvidia_riva_models: Set = set()
soniox_models: Set = set()
sambanova_models: Set = set()
scaledown_models: Final[Set[str]] = set() # Price reloads update this registry through existing references.
sambanova_embedding_models: Set = set()
novita_models: Set = set()
assemblyai_models: Set = set()
@ -756,7 +757,7 @@ def is_openai_finetune_model(key: str) -> bool:
return key.startswith("ft:") and not key.count(":") > 1
def _populate_provider_model_sets(model_cost_map: Dict) -> None:
def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None:
for key, value in model_cost_map.items():
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
open_ai_chat_completion_models.add(key)
@ -918,6 +919,8 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None:
soniox_models.add(key)
elif value.get("litellm_provider") == "sambanova":
sambanova_models.add(key)
elif value.get("litellm_provider") == "scaledown":
scaledown_models.add(key)
elif value.get("litellm_provider") == "sambanova-embedding-models":
sambanova_embedding_models.add(key)
elif value.get("litellm_provider") == "novita":
@ -1106,6 +1109,7 @@ model_list = list(
| nvidia_riva_models
| soniox_models
| sambanova_models
| scaledown_models
| azure_text_models
| novita_models
| assemblyai_models
@ -1215,6 +1219,7 @@ def _build_models_by_provider() -> dict:
"nvidia_riva": nvidia_riva_models,
"soniox": soniox_models,
"sambanova": sambanova_models | sambanova_embedding_models,
"scaledown": scaledown_models,
"novita": novita_models,
"nebius": nebius_models | nebius_embedding_models,
"aiml": aiml_models,
@ -1485,6 +1490,7 @@ from .rust_bridge import rust
from .rag.main import *
from .sandbox.main import *
from .decisions.main import *
from .systemone.main import *
from .tool_loop import ToolLoopMaxRoundsExceeded, arun_tool_loop, run_tool_loop
from .search.main import *
from .realtime_api.main import (
@ -2054,6 +2060,9 @@ if TYPE_CHECKING:
from .llms.featherless_ai.chat.transformation import (
FeatherlessAIConfig as FeatherlessAIConfig,
)
from .llms.scaledown.chat.transformation import (
ScaleDownChatConfig as ScaleDownChatConfig,
)
from .llms.cerebras.chat import CerebrasConfig as CerebrasConfig
from .llms.nadir.chat.transformation import NadirConfig as NadirConfig
from .llms.baseten.chat import BasetenConfig as BasetenConfig
@ -2386,6 +2395,11 @@ def __getattr__(name: str) -> Any:
handler_func: Final = registry[name]
return handler_func(name)
if name == "proxy":
import importlib
return importlib.import_module("litellm.proxy")
if name == "harness" or name in _AGENT_EXPORTS:
import importlib

5
litellm/_lazy_imports_registry.py
View file

@ -277,6 +277,7 @@ LLM_CONFIG_NAMES: Final = (
"NvidiaNimConfig",
"NvidiaNimEmbeddingConfig",
"FeatherlessAIConfig",
"ScaleDownChatConfig",
"CerebrasConfig",
"NadirConfig",
"BasetenConfig",
@ -1103,6 +1104,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.featherless_ai.chat.transformation",
"FeatherlessAIConfig",
),
"ScaleDownChatConfig": (
".llms.scaledown.chat.transformation",
"ScaleDownChatConfig",
),
"CerebrasConfig": (".llms.cerebras.chat", "CerebrasConfig"),
"NadirConfig": (".llms.nadir.chat.transformation", "NadirConfig"),
"BasetenConfig": (".llms.baseten.chat", "BasetenConfig"),

2
litellm/_service_logger.py
View file

@ -16,7 +16,7 @@ from .types.services import ServiceLoggerPayload, ServiceTypes
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
Span = _Span | Any
OTELClass = OpenTelemetry

3
litellm/constants.py
View file

@ -234,6 +234,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS: Final = 5.0
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH: Final = 8
MCP_BYOK_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
@ -825,6 +826,7 @@ LITELLM_CHAT_PROVIDERS: Final = [
"lemonade",
"docker_model_runner",
"amazon_nova",
"scaledown",
]
# Resolving these providers runs an OAuth device flow (their provider info IS the login), so any
@ -1969,6 +1971,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
]
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
AUTH_OBJECTS_TARGET: Final = "auth_objects"
DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30

2
litellm/cost_calculator.py
View file

@ -2026,6 +2026,8 @@ def response_cost_calculator(
"asearch",
"decisions",
"adecisions",
"systemone",
"asystemone",
],
optional_params: dict,
cache_hit: bool | None = None,

347
litellm/decisions/call.py Normal file
View file

@ -0,0 +1,347 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from typing import Final, TypeAlias
import httpx
from pydantic import ConfigDict, TypeAdapter, ValidationError
from typing_extensions import assert_never
import litellm
from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.decisions.transformation import (
BaseDecisionsConfig,
ir_to_systemone_response,
systemone_request_to_ir,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.openai.decisions.transformation import ir_to_openai_response, openai_request_to_ir
from litellm.types.decisions import (
DecisionQuestion,
DecisionsIRRequest,
DecisionsIRResponse,
DecisionsJSON,
DecisionsRequestBody,
DecisionsResponse,
OpenAIDecisionInput,
OpenAIDecisionQuestion,
OpenAIDecisionRequestBody,
OpenAIDecisionResponse,
UnsupportedDecisionsRequest,
)
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
SystemOneQuestions: TypeAlias = Mapping[str, DecisionQuestion | Mapping[str, object]]
OpenAIDecisionQuestions: TypeAlias = Sequence[OpenAIDecisionQuestion | Mapping[str, object]]
DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequestBody
_SYSTEMONE_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody)
_OPENAI_REQUEST_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody)
_SAFETY_IDENTIFIER_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(
str | None, config=ConfigDict(title="safety_identifier")
)
_HANDLER: Final = BaseLLMHTTPHandler()
@dataclass(frozen=True, slots=True)
class SystemOneRequestFields:
state: DecisionsJSON | None
questions: SystemOneQuestions | None
def validate(self, safety_identifier: str | None) -> DecisionsRequestBody:
return _SYSTEMONE_REQUEST_ADAPTER.validate_python({"state": self.state, "questions": self.questions})
@dataclass(frozen=True, slots=True)
class OpenAIDecisionRequestFields:
input: OpenAIDecisionInput | None
questions: OpenAIDecisionQuestions | None
def validate(self, safety_identifier: str | None) -> OpenAIDecisionRequestBody:
return _OPENAI_REQUEST_ADAPTER.validate_python(
{"input": self.input, "questions": self.questions, "safety_identifier": safety_identifier}
)
RequestFields: TypeAlias = SystemOneRequestFields | OpenAIDecisionRequestFields
@dataclass(frozen=True, slots=True, repr=False)
class DecisionsCall:
model: str
requested_model: str
custom_llm_provider: str
provider_config: BaseDecisionsConfig
request: DecisionsRequestFormat
ir_request: DecisionsIRRequest
body: Mapping[str, object] = field(repr=False)
api_base: str
api_key: str | None = field(repr=False)
logging_obj: LiteLLMLoggingObj | None
headers: Mapping[str, str]
timeout: float | httpx.Timeout | None
litellm_params: Mapping[str, object] = field(repr=False)
def supported_decisions_providers() -> tuple[str, ...]:
return tuple(
provider.value
for provider in LlmProviders
if ProviderConfigManager.get_provider_decisions_config(model="", provider=provider) is not None
)
def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig:
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
provider_config: Final = (
None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider)
)
if provider_config is None:
supported: Final = ", ".join(supported_decisions_providers())
raise litellm.BadRequestError(
message=f"Unknown Decisions provider '{custom_llm_provider}'. Supported providers: {supported}",
model=model,
llm_provider=custom_llm_provider,
)
return provider_config
def _ir_request(request: DecisionsRequestFormat, safety_identifier: str | None) -> DecisionsIRRequest:
match request:
case DecisionsRequestBody():
return replace(systemone_request_to_ir(request), safety_identifier=safety_identifier)
case OpenAIDecisionRequestBody():
return openai_request_to_ir(request)
case _:
assert_never(request)
def _drops_params(kwargs: Mapping[str, object]) -> bool:
return litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True
def _request_safety_identifier(safety_identifier: object, kwargs: Mapping[str, object]) -> str | None:
if isinstance(safety_identifier, str) or not _drops_params(kwargs):
return _SAFETY_IDENTIFIER_ADAPTER.validate_python(safety_identifier)
return None
@dataclass(frozen=True, slots=True)
class _UnsupportedSafetyIdentifier:
pass
def _provider_ir_request(
request: DecisionsRequestFormat,
*,
safety_identifier: str | None,
provider_config: BaseDecisionsConfig,
kwargs: Mapping[str, object],
) -> DecisionsIRRequest | _UnsupportedSafetyIdentifier:
ir_request: Final = _ir_request(request, safety_identifier)
if ir_request.safety_identifier is None or provider_config.supports_safety_identifier:
return ir_request
if _drops_params(kwargs):
return replace(ir_request, safety_identifier=None)
return _UnsupportedSafetyIdentifier()
def prepare_call(
*,
model: str,
fields: RequestFields,
safety_identifier: str | None,
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> DecisionsCall:
upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
provider_config: Final = _provider_config(upstream_model, provider)
if not upstream_model:
raise litellm.BadRequestError(
message="A model name is required for the Decisions API",
model=model,
llm_provider=provider,
)
try:
canonical_model: Final = provider_config.canonical_model(upstream_model)
except ValueError as error:
raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
try:
request_safety_identifier: Final = _request_safety_identifier(safety_identifier, kwargs)
request: Final = fields.validate(request_safety_identifier)
except ValidationError as error:
raise litellm.BadRequestError(
message=f"Invalid Decisions request: {error}",
model=model,
llm_provider=provider,
) from error
resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
if resolved_api_base is None:
raise litellm.BadRequestError(
message=provider_config.missing_api_base_message(provider),
model=model,
llm_provider=provider,
)
resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
if resolved_api_key is None and provider_config.api_key_required:
raise litellm.AuthenticationError(
message=f"Missing API key for Decisions provider '{provider}'",
model=model,
llm_provider=provider,
)
ir_request: Final = _provider_ir_request(
request,
safety_identifier=request_safety_identifier,
provider_config=provider_config,
kwargs=kwargs,
)
if isinstance(ir_request, _UnsupportedSafetyIdentifier):
raise litellm.UnsupportedParamsError(
message=(
f"{provider} does not support parameters: ['safety_identifier'], for model={model}. "
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n"
"`litellm_settings:\n drop_params: true`\n"
),
model=model,
llm_provider=provider,
)
body: Final = provider_config.transform_decisions_request(
model=canonical_model, request=ir_request, custom_llm_provider=provider
)
if isinstance(body, UnsupportedDecisionsRequest):
raise litellm.BadRequestError(
message=f"Decisions provider '{provider}' cannot serve this request: {body.reason}",
model=model,
llm_provider=provider,
)
logging_obj: Final = kwargs.get("litellm_logging_obj")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_from_kwargs(
kwargs=dict(kwargs),
model=canonical_model,
litellm_params={
"litellm_call_id": kwargs.get("litellm_call_id"),
"api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
},
custom_llm_provider=provider,
)
return DecisionsCall(
model=canonical_model,
requested_model=model,
custom_llm_provider=provider,
provider_config=provider_config,
request=request,
ir_request=ir_request,
body=body,
api_base=resolved_api_base,
api_key=resolved_api_key,
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
headers=extra_headers or {},
timeout=timeout,
litellm_params=kwargs,
)
def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
formatted: Final = _formatted_response(response, call)
provider_cost: Final = call.provider_config.provider_reported_cost(response)
formatted.set_hidden_params(
{
"model": f"{call.custom_llm_provider}/{call.model}",
"custom_llm_provider": call.custom_llm_provider,
"provider_response_model": f"{call.custom_llm_provider}/{call.model}",
**({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}),
}
)
return formatted
def _formatted_response(
response: DecisionsIRResponse, call: DecisionsCall
) -> DecisionsResponse | OpenAIDecisionResponse:
match call.request:
case DecisionsRequestBody():
return ir_to_systemone_response(response, call.ir_request)
case OpenAIDecisionRequestBody():
return ir_to_openai_response(response, call.ir_request, call.requested_model)
case _:
assert_never(call.request)
def _map_upstream_exception(error: Exception, call: DecisionsCall) -> Exception:
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
provider_label: Final = f"{call.custom_llm_provider[0].upper()}{call.custom_llm_provider[1:]}Exception"
return litellm.APIConnectionError(
message=f"{provider_label} - {error.message}",
llm_provider=call.custom_llm_provider,
model=f"{call.custom_llm_provider}/{call.model}",
)
return litellm.exception_type(
model=f"{call.custom_llm_provider}/{call.model}",
custom_llm_provider=call.custom_llm_provider,
original_exception=error,
)
async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
try:
response: Final = await _HANDLER.adecisions(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
request=call.ir_request,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
)
except Exception as error:
raise _map_upstream_exception(error, call) from error
return _format_response(response, call)
def send(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
try:
response: Final = _HANDLER.decisions(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
request=call.ir_request,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
)
except Exception as error:
raise _map_upstream_exception(error, call) from error
return _format_response(response, call)
def reject_other_format(*, model: str, kwargs: Mapping[str, object], field_name: str, other_function: str) -> None:
if field_name not in kwargs:
return
raise litellm.BadRequestError(
message=f"Requests with '{field_name}' go through litellm.{other_function}",
model=model,
llm_provider="",
)

388
litellm/decisions/main.py
View file

@ -1,151 +1,56 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from typing import Final, TypeAlias
import warnings
from collections.abc import Mapping
from typing import Final
import httpx
from pydantic import ConfigDict, TypeAdapter, ValidationError
from typing_extensions import assert_never
from pydantic import TypeAdapter
import litellm
from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.decisions.transformation import (
BaseDecisionsConfig,
ir_to_systemone_response,
systemone_request_to_ir,
from litellm.decisions.call import (
DecisionsCall,
OpenAIDecisionQuestions,
OpenAIDecisionRequestFields,
RequestFields,
SystemOneQuestions,
SystemOneRequestFields,
asend,
prepare_call,
send,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.openai.decisions.transformation import ir_to_openai_response, openai_request_to_ir
from litellm.types.decisions import (
DecisionQuestion,
DecisionsIRRequest,
DecisionsIRResponse,
DecisionsJSON,
DecisionsRequestBody,
DecisionsResponse,
OpenAIDecisionInput,
OpenAIDecisionQuestion,
OpenAIDecisionRequestBody,
OpenAIDecisionResponse,
UnsupportedDecisionsRequest,
from litellm.types.decisions import DecisionsJSON, DecisionsResponse, OpenAIDecisionInput, OpenAIDecisionResponse
from litellm.utils import client
_STATE_ADAPTER: Final[TypeAdapter[DecisionsJSON | None]] = TypeAdapter(DecisionsJSON | None)
_SYSTEMONE_QUESTIONS_ADAPTER: Final[TypeAdapter[SystemOneQuestions | None]] = TypeAdapter(SystemOneQuestions | None)
_OPENAI_QUESTIONS_ADAPTER: Final[TypeAdapter[OpenAIDecisionQuestions | None]] = TypeAdapter(
OpenAIDecisionQuestions | None
)
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager, client
DecisionsQuestions: TypeAlias = (
Mapping[str, DecisionQuestion | Mapping[str, object]] | Sequence[OpenAIDecisionQuestion | Mapping[str, object]]
)
DecisionsRequestFormat: TypeAlias = DecisionsRequestBody | OpenAIDecisionRequestBody
_SYSTEMONE_REQUEST_ADAPTER: Final[TypeAdapter[DecisionsRequestBody]] = TypeAdapter(DecisionsRequestBody)
_OPENAI_REQUEST_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBody]] = TypeAdapter(OpenAIDecisionRequestBody)
_SAFETY_IDENTIFIER_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(
str | None, config=ConfigDict(title="safety_identifier")
)
_HANDLER: Final = BaseLLMHTTPHandler()
@dataclass(frozen=True, slots=True, repr=False)
class _DecisionsCall:
model: str
requested_model: str
custom_llm_provider: str
provider_config: BaseDecisionsConfig
request: DecisionsRequestFormat
ir_request: DecisionsIRRequest
body: Mapping[str, object] = field(repr=False)
api_base: str
api_key: str | None = field(repr=False)
logging_obj: LiteLLMLoggingObj | None
headers: Mapping[str, str]
timeout: float | httpx.Timeout | None
litellm_params: Mapping[str, object] = field(repr=False)
def _supported_providers() -> tuple[str, ...]:
return tuple(
provider.value
for provider in LlmProviders
if ProviderConfigManager.get_provider_decisions_config(model="", provider=provider) is not None
)
def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig:
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
provider_config: Final = (
None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider)
)
if provider_config is None:
supported: Final = ", ".join(_supported_providers())
raise litellm.BadRequestError(
message=f"Unknown Decisions provider '{custom_llm_provider}'. Supported providers: {supported}",
model=model,
llm_provider=custom_llm_provider,
)
return provider_config
def _validate_request(
*,
state: DecisionsJSON | None,
questions: DecisionsQuestions | None,
def _request_fields(
decision_input: OpenAIDecisionInput | None,
safety_identifier: str | None,
) -> DecisionsRequestFormat:
if decision_input is None:
return _SYSTEMONE_REQUEST_ADAPTER.validate_python({"state": state, "questions": questions})
return _OPENAI_REQUEST_ADAPTER.validate_python(
{"input": decision_input, "questions": questions, "safety_identifier": safety_identifier}
questions: OpenAIDecisionQuestions | SystemOneQuestions | None,
kwargs: Mapping[str, object],
) -> RequestFields:
if "state" not in kwargs:
return OpenAIDecisionRequestFields(
input=decision_input, questions=_OPENAI_QUESTIONS_ADAPTER.validate_python(questions)
)
warnings.warn(
"litellm.decisions(state=...) is deprecated, call litellm.systemone for the System One request format",
UserWarning,
stacklevel=5,
)
return SystemOneRequestFields(
state=_STATE_ADAPTER.validate_python(kwargs["state"]),
questions=_SYSTEMONE_QUESTIONS_ADAPTER.validate_python(questions),
)
def _ir_request(request: DecisionsRequestFormat, safety_identifier: str | None) -> DecisionsIRRequest:
match request:
case DecisionsRequestBody():
return replace(systemone_request_to_ir(request), safety_identifier=safety_identifier)
case OpenAIDecisionRequestBody():
return openai_request_to_ir(request)
case _:
assert_never(request)
def _drops_params(kwargs: Mapping[str, object]) -> bool:
return litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True
def _request_safety_identifier(safety_identifier: object, kwargs: Mapping[str, object]) -> str | None:
if isinstance(safety_identifier, str) or not _drops_params(kwargs):
return _SAFETY_IDENTIFIER_ADAPTER.validate_python(safety_identifier)
return None
@dataclass(frozen=True, slots=True)
class _UnsupportedSafetyIdentifier:
pass
def _provider_ir_request(
request: DecisionsRequestFormat,
*,
safety_identifier: str | None,
provider_config: BaseDecisionsConfig,
kwargs: Mapping[str, object],
) -> DecisionsIRRequest | _UnsupportedSafetyIdentifier:
ir_request: Final = _ir_request(request, safety_identifier)
if ir_request.safety_identifier is None or provider_config.supports_safety_identifier:
return ir_request
if _drops_params(kwargs):
return replace(ir_request, safety_identifier=None)
return _UnsupportedSafetyIdentifier()
def _prepare_call(
def _prepare_decisions_call(
*,
model: str,
state: DecisionsJSON | None,
questions: DecisionsQuestions | None,
decision_input: OpenAIDecisionInput | None,
questions: OpenAIDecisionQuestions | SystemOneQuestions | None,
safety_identifier: str | None,
api_key: str | None,
api_base: str | None,
@ -153,174 +58,37 @@ def _prepare_call(
custom_llm_provider: str | None,
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> _DecisionsCall:
upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider(
) -> DecisionsCall:
return prepare_call(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
fields=_request_fields(decision_input, questions, kwargs),
safety_identifier=safety_identifier,
api_key=api_key,
)
provider_config: Final = _provider_config(upstream_model, provider)
if not upstream_model:
raise litellm.BadRequestError(
message="A model name is required for the Decisions API",
model=model,
llm_provider=provider,
)
if state is not None and decision_input is not None:
raise litellm.BadRequestError(
message="Pass either state (System One format) or input (OpenAI format) to the Decisions API, not both",
model=model,
llm_provider=provider,
)
try:
canonical_model: Final = provider_config.canonical_model(upstream_model)
except ValueError as error:
raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
try:
request_safety_identifier: Final = _request_safety_identifier(safety_identifier, kwargs)
request: Final = _validate_request(
state=state,
questions=questions,
decision_input=decision_input,
safety_identifier=request_safety_identifier,
)
except ValidationError as error:
raise litellm.BadRequestError(
message=f"Invalid Decisions request: {error}",
model=model,
llm_provider=provider,
) from error
resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
if resolved_api_base is None:
raise litellm.BadRequestError(
message=provider_config.missing_api_base_message(provider),
model=model,
llm_provider=provider,
)
resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
if resolved_api_key is None and provider_config.api_key_required:
raise litellm.AuthenticationError(
message=f"Missing API key for Decisions provider '{provider}'",
model=model,
llm_provider=provider,
)
ir_request: Final = _provider_ir_request(
request,
safety_identifier=request_safety_identifier,
provider_config=provider_config,
kwargs=kwargs,
)
if isinstance(ir_request, _UnsupportedSafetyIdentifier):
raise litellm.UnsupportedParamsError(
message=(
f"{provider} does not support parameters: ['safety_identifier'], for model={model}. "
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n"
"`litellm_settings:\n drop_params: true`\n"
),
model=model,
llm_provider=provider,
)
body: Final = provider_config.transform_decisions_request(
model=canonical_model, request=ir_request, custom_llm_provider=provider
)
if isinstance(body, UnsupportedDecisionsRequest):
raise litellm.BadRequestError(
message=f"Decisions provider '{provider}' cannot serve this request: {body.reason}",
model=model,
llm_provider=provider,
)
logging_obj: Final = kwargs.get("litellm_logging_obj")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_from_kwargs(
kwargs=dict(kwargs),
model=canonical_model,
litellm_params={
"litellm_call_id": kwargs.get("litellm_call_id"),
"api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
},
custom_llm_provider=provider,
)
return _DecisionsCall(
model=canonical_model,
requested_model=model,
custom_llm_provider=provider,
provider_config=provider_config,
request=request,
ir_request=ir_request,
body=body,
api_base=resolved_api_base,
api_key=resolved_api_key,
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
headers=extra_headers or {},
api_base=api_base,
timeout=timeout,
litellm_params=kwargs,
)
def _format_response(response: DecisionsIRResponse, call: _DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
formatted: Final = _formatted_response(response, call)
provider_cost: Final = call.provider_config.provider_reported_cost(response)
formatted.set_hidden_params(
{
"model": f"{call.custom_llm_provider}/{call.model}",
"custom_llm_provider": call.custom_llm_provider,
"provider_response_model": f"{call.custom_llm_provider}/{call.model}",
**({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}),
}
)
return formatted
def _formatted_response(
response: DecisionsIRResponse, call: _DecisionsCall
) -> DecisionsResponse | OpenAIDecisionResponse:
match call.request:
case DecisionsRequestBody():
return ir_to_systemone_response(response, call.ir_request)
case OpenAIDecisionRequestBody():
return ir_to_openai_response(response, call.ir_request, call.requested_model)
case _:
assert_never(call.request)
def _map_upstream_exception(error: Exception, call: _DecisionsCall) -> Exception:
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
provider_label: Final = f"{call.custom_llm_provider[0].upper()}{call.custom_llm_provider[1:]}Exception"
return litellm.APIConnectionError(
message=f"{provider_label} - {error.message}",
llm_provider=call.custom_llm_provider,
model=f"{call.custom_llm_provider}/{call.model}",
)
return litellm.exception_type(
model=f"{call.custom_llm_provider}/{call.model}",
custom_llm_provider=call.custom_llm_provider,
original_exception=error,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs={key: value for key, value in kwargs.items() if key != "state"},
)
@client
async def adecisions(
model: str,
state: DecisionsJSON | None = None,
questions: DecisionsQuestions | None = None,
input: OpenAIDecisionInput | None = None,
questions: OpenAIDecisionQuestions | SystemOneQuestions | None = None,
safety_identifier: str | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
input: OpenAIDecisionInput | None = None,
safety_identifier: str | None = None,
**kwargs: object,
) -> DecisionsResponse | OpenAIDecisionResponse:
call: Final = _prepare_call(
) -> OpenAIDecisionResponse | DecisionsResponse:
call: Final = _prepare_decisions_call(
model=model,
state=state,
questions=questions,
decision_input=input,
questions=questions,
safety_identifier=safety_identifier,
api_key=api_key,
api_base=api_base,
@ -329,44 +97,26 @@ async def adecisions(
extra_headers=extra_headers,
kwargs=kwargs,
)
try:
response: Final = await _HANDLER.adecisions(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
request=call.ir_request,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
)
except Exception as error:
raise _map_upstream_exception(error, call) from error
return _format_response(response, call)
return await asend(call)
@client
def decisions(
model: str,
state: DecisionsJSON | None = None,
questions: DecisionsQuestions | None = None,
input: OpenAIDecisionInput | None = None,
questions: OpenAIDecisionQuestions | SystemOneQuestions | None = None,
safety_identifier: str | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
input: OpenAIDecisionInput | None = None,
safety_identifier: str | None = None,
**kwargs: object,
) -> DecisionsResponse | OpenAIDecisionResponse:
call: Final = _prepare_call(
) -> OpenAIDecisionResponse | DecisionsResponse:
call: Final = _prepare_decisions_call(
model=model,
state=state,
questions=questions,
decision_input=input,
questions=questions,
safety_identifier=safety_identifier,
api_key=api_key,
api_base=api_base,
@ -375,23 +125,7 @@ def decisions(
extra_headers=extra_headers,
kwargs=kwargs,
)
try:
response: Final = _HANDLER.decisions(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
request=call.ir_request,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
)
except Exception as error:
raise _map_upstream_exception(error, call) from error
return _format_response(response, call)
return send(call)
__all__ = ["adecisions", "decisions"]

407
litellm/experimental_mcp_client/client.py
View file

@ -8,12 +8,12 @@ import hashlib
import json
import os
import time
from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Sequence
from contextlib import AbstractAsyncContextManager
from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager, AsyncExitStack, nullcontext
from functools import partial
from importlib.metadata import version
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar, cast
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar, cast, overload
import anyio
import httpx2
@ -21,26 +21,42 @@ from httpx2._client import UseClientDefault
from httpx2._types import AuthTypes
from mcp import ClientSession, MCPError, ReadResourceResult, Resource, StdioServerParameters
from mcp.client._input_required import run_input_required_driver
from mcp.client.session import ClientRequestContext, ElicitationFnT, LoggingFnT, SamplingFnT
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.client.streamable_http import streamable_http_client
from mcp.shared._stream_protocols import ReadStream, WriteStream
from mcp.shared.message import SessionMessage
from typing_extensions import ReadOnly, TypedDict, assert_never
_TransportStreams: TypeAlias = tuple[
ReadStream[SessionMessage | Exception],
WriteStream[SessionMessage],
]
_TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams]
_SessionCallbackName: TypeAlias = Literal["sampling_callback", "elicitation_callback", "logging_callback"]
class _SessionCallbackMap(TypedDict):
sampling_callback: ReadOnly[SamplingFnT | None]
elicitation_callback: ReadOnly[ElicitationFnT | None]
logging_callback: ReadOnly[LoggingFnT | None]
from mcp.types import (
CONNECTION_CLOSED,
METHOD_NOT_FOUND,
REQUEST_TIMEOUT,
CacheableResult,
ClientCapabilities,
CreateMessageRequestParams,
CreateMessageResult,
CreateMessageResultWithTools,
DiscoverResult,
ElicitationCapability,
ElicitRequestParams,
ElicitResult,
ErrorData,
FormElicitationCapability,
GetPromptRequestParams,
GetPromptResult,
@ -58,6 +74,7 @@ from mcp.types import (
ListResourceTemplatesResult,
ListToolsRequest,
ListToolsResult,
LoggingMessageNotificationParams,
PaginatedRequestParams,
Prompt,
ResourceTemplate,
@ -302,6 +319,12 @@ class MCPSigV4Auth(httpx2.Auth):
raise ImportError("Missing botocore to use AWS SigV4 authentication. Run 'pip install boto3'.")
self.service_name = aws_service_name or "bedrock-agentcore"
self.region_name = aws_region_name or "us-east-1"
self.identity: Final[tuple[str, str, str | None, str | None]] = (
self.service_name,
self.region_name,
aws_role_name,
aws_access_key_id,
)
# Note: os.environ/ prefixed values are already resolved by
# ProxyConfig._check_for_os_environ_vars() at config load time.
# Values arrive here as plain strings.
@ -346,7 +369,7 @@ class MCPSigV4Auth(httpx2.Auth):
import time
import boto3
from botocore.credentials import Credentials
from botocore.credentials import RefreshableCredentials
session_name: Final = aws_session_name or f"litellm-mcp-{int(time.time())}"
sts_kwargs: Final[dict] = {"region_name": aws_region_name}
@ -356,15 +379,23 @@ class MCPSigV4Auth(httpx2.Auth):
if aws_session_token:
sts_kwargs["aws_session_token"] = aws_session_token
sts_client: Final = boto3.client("sts", **sts_kwargs)
sts_response: Final = sts_client.assume_role(
RoleArn=aws_role_name,
RoleSessionName=session_name,
)
sts_creds: Final = sts_response["Credentials"]
return Credentials(
access_key=sts_creds["AccessKeyId"],
secret_key=sts_creds["SecretAccessKey"],
token=sts_creds["SessionToken"],
def assume() -> Mapping[str, str]:
sts_creds: Final = sts_client.assume_role(
RoleArn=aws_role_name,
RoleSessionName=session_name,
)["Credentials"]
return MappingProxyType(
{
"access_key": sts_creds["AccessKeyId"],
"secret_key": sts_creds["SecretAccessKey"],
"token": sts_creds["SessionToken"],
"expiry_time": sts_creds["Expiration"].isoformat(),
}
)
return RefreshableCredentials.create_from_metadata(
metadata=dict(assume()), refresh_using=assume, method="sts-assume-role"
)
def auth_flow(self, request: httpx2.Request) -> Generator[httpx2.Request, httpx2.Response, None]:
@ -390,6 +421,19 @@ class MCPSigV4Auth(httpx2.Auth):
yield request
def _notify_stream_error(
callback: Callable[[asyncio.Future[Exception]], None] | None,
stream_error: asyncio.Future[Exception],
) -> None:
if callback is not None:
callback(stream_error)
def _notify_cleanup(callback: Callable[[], None] | None) -> None:
if callback is not None:
callback()
class MCPClient:
"""
MCP Client supporting:
@ -443,12 +487,37 @@ class MCPClient:
self._sampling_callback: Callable | None = sampling_callback
self._elicitation_callback: ElicitationCallback | None = elicitation_callback
self._logging_callback: Callable | None = logging_callback
self._session_callbacks_override: _SessionCallbackMap | None = None
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
def session_callbacks(self) -> _SessionCallbackMap:
return {
"sampling_callback": self._sampling_callback,
"elicitation_callback": self._elicitation_callback,
"logging_callback": self._logging_callback,
}
def session_settings(self) -> tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool]:
callbacks: Final = self.session_callbacks()
return (
self.transport_type,
self.protocol_version,
self.timeout,
callbacks["sampling_callback"] is not None,
callbacks["elicitation_callback"] is not None,
callbacks["logging_callback"] is not None,
)
def route_session_callbacks(self, callbacks: _SessionCallbackMap) -> None:
self._session_callbacks_override = callbacks
async def discovery_auth_fingerprint(self) -> str:
return self._hash_discovery_auth(await self.prepare_request_auth())
return self._hash_discovery_auth(await self.prepare_request_auth(), self._aws_identity())
def _aws_identity(self) -> tuple[str, str, str | None, str | None] | None:
return self._aws_auth.identity if isinstance(self._aws_auth, MCPSigV4Auth) else None
async def prepare_request_auth(self) -> httpx2.Request:
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
@ -465,8 +534,10 @@ class MCPClient:
await flow.aclose()
@staticmethod
def _hash_discovery_auth(request: httpx2.Request) -> str:
material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items()))))
def _hash_discovery_auth(
request: httpx2.Request, aws_identity: tuple[str, str, str | None, str | None] | None
) -> str:
material: Final = json.dumps((str(request.url), tuple(sorted(request.headers.multi_items())), aws_identity))
return hashlib.sha256(material.encode()).hexdigest()
def _create_transport_context(
@ -605,6 +676,9 @@ class MCPClient:
transport_ctx: _TransportContext,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
http_client: httpx2.AsyncClient | None = None,
*,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> TSessionResult:
"""
Execute an operation within a transport and session context.
@ -623,6 +697,7 @@ class MCPClient:
read_stream: Final = transport[0]
write_stream: Final = transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
_notify_stream_error(on_stream_error, stream_error)
async def receive_message(
message: ServerNotification | Exception,
@ -634,6 +709,11 @@ class MCPClient:
# The SDK closes pending requests when its message handler raises.
raise RuntimeError("MCP response stream failed")
session_kwargs: Final = (
self._session_callbacks_override
if self._session_callbacks_override is not None
else self.session_callbacks()
)
# The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else
# ever fails the request.
session_ctx: Final = ClientSession(
@ -641,9 +721,7 @@ class MCPClient:
write_stream,
read_timeout_seconds=self.timeout,
message_handler=receive_message,
sampling_callback=self._sampling_callback,
elicitation_callback=self._elicitation_callback,
logging_callback=self._logging_callback,
**session_kwargs,
)
session: Final = await session_ctx.__aenter__()
try:
@ -659,6 +737,7 @@ class MCPClient:
raise stream_error.result()
raise
finally:
_notify_cleanup(on_cleanup)
cleanup_scope.shield = True
cleanup_scope.deadline = anyio.current_time() + 5
try:
@ -671,6 +750,7 @@ class MCPClient:
in_flight_error = e
raise
finally:
_notify_cleanup(on_cleanup)
cleanup_scope.shield = True
cleanup_scope.deadline = min(cleanup_scope.deadline, anyio.current_time() + 5)
try:
@ -703,6 +783,8 @@ class MCPClient:
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
*,
quiet_on_error: bool = False,
on_stream_error: Callable[[asyncio.Future[Exception]], None] | None = None,
on_cleanup: Callable[[], None] | None = None,
) -> TSessionResult:
"""Open a session, run the provided coroutine, and clean up.
@ -714,7 +796,18 @@ class MCPClient:
try:
self._last_initialize_instructions = None
transport_ctx, http_client = self._create_transport_context()
result: Final = await self._execute_session_operation(transport_ctx, operation, http_client=http_client)
execution: Final = (
self._execute_session_operation(
transport_ctx,
operation,
http_client=http_client,
on_stream_error=on_stream_error,
on_cleanup=on_cleanup,
)
if on_stream_error is not None or on_cleanup is not None
else self._execute_session_operation(transport_ctx, operation, http_client=http_client)
)
result: Final = await execution
except Exception as e:
read_timeout: Final = as_mcp_read_timeout(e)
if read_timeout is not None:
@ -741,6 +834,12 @@ class MCPClient:
await anyio.lowlevel.checkpoint_if_cancelled()
return result
def open_persistent_session(
self,
admission: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> "PersistentMCPSession":
return PersistentMCPSession(self, admission)
def update_auth_value(self, mcp_auth_value: str | dict[str, str]) -> None:
"""
Set the authentication header for the MCP client.
@ -972,10 +1071,15 @@ class MCPClient:
host_progress_callback: Callable | None = None,
raise_on_error: bool = False,
allow_input_required: bool = False,
persistent_session: "PersistentMCPSession | None" = None,
) -> MCPCallToolResult | InputRequiredResult:
"""
Call an MCP Tool.
persistent_session runs the call inside an already-open upstream session (one
upstream mcp-session-id shared by every call of a stateful gateway session) instead of
the per-call initialize + teardown that run_with_session performs.
Args:
raise_on_error: When True, re-raise the underlying exception instead of returning an
``isError=True`` result. The token-exchange (OBO) tool-call path uses this to detect
@ -1030,7 +1134,15 @@ class MCPClient:
)
try:
tool_result: Final = await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error)
tool_result: Final = (
await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error)
if persistent_session is None
else await persistent_session.run(
_call_tool_operation,
quiet_on_error=raise_on_error,
callbacks=self.session_callbacks(),
)
)
verbose_logger.info("MCP client tool call '%s' completed successfully", call_tool_request_params.name)
return tool_result
except asyncio.CancelledError:
@ -1395,3 +1507,256 @@ class MCPClient:
"the MCP server may have crashed, disconnected, or timed out."
)
raise
_PendingOperation: TypeAlias = tuple[
Callable[[ClientSession], Awaitable[object]],
asyncio.Future[object],
_SessionCallbackMap | None,
]
_MAX_PENDING_OPERATIONS: Final = 64
def _ends_session(outcome: BaseException) -> bool:
if isinstance(
outcome,
(httpx2.HTTPError, OSError, anyio.BrokenResourceError, anyio.ClosedResourceError, anyio.EndOfStream),
):
return True
return isinstance(outcome, MCPError) and outcome.error.code == CONNECTION_CLOSED
def _deliver_persistent_outcome(
future: asyncio.Future[object],
outcome: object,
stream_error: asyncio.Future[Exception] | None,
) -> bool:
mapped_outcome: Final = (
stream_error.result()
if isinstance(outcome, MCPError) and stream_error is not None and stream_error.done()
else outcome
)
if not future.done():
if isinstance(mapped_outcome, Exception):
future.set_exception(mapped_outcome)
elif isinstance(mapped_outcome, BaseException):
future.set_exception(RuntimeError("upstream MCP operation was cancelled"))
else:
future.set_result(mapped_outcome)
return (
isinstance(outcome, BaseException)
and _ends_session(outcome)
or isinstance(mapped_outcome, BaseException)
and _ends_session(mapped_outcome)
or stream_error is not None
and stream_error.done()
)
class UpstreamSessionClosedError(RuntimeError):
def __init__(self) -> None:
super().__init__("upstream MCP session closed")
class PersistentMCPSession:
"""One upstream MCP session kept open across operations.
The SDK transport owns anyio cancel scopes that must be entered and exited by the same
task, so a dedicated task holds the session open and serves queued operations in order.
"""
def __init__(
self,
client: MCPClient,
admission: Callable[[], AbstractAsyncContextManager[None]] | None = None,
) -> None:
loop: Final = asyncio.get_running_loop()
self._client: Final = client
self._admission: Final = admission
self._queue: Final[asyncio.Queue[_PendingOperation]] = asyncio.Queue(maxsize=_MAX_PENDING_OPERATIONS)
self._ready: Final[asyncio.Future[None]] = loop.create_future()
self._active: asyncio.Future[object] | None = None
self._stream_error: asyncio.Future[Exception] | None = None
self._retiring: bool = False
self._ended: bool = False
self._initial_callbacks: Final[_SessionCallbackMap] = client.session_callbacks()
self._current_callbacks: _SessionCallbackMap = self._initial_callbacks
delegated_callbacks: Final[_SessionCallbackMap] = {
"sampling_callback": (
self._delegate("sampling_callback")
if self._initial_callbacks["sampling_callback"] is not None
else None
),
"elicitation_callback": (
self._delegate("elicitation_callback")
if self._initial_callbacks["elicitation_callback"] is not None
else None
),
"logging_callback": (
self._delegate("logging_callback") if self._initial_callbacks["logging_callback"] is not None else None
),
}
client.route_session_callbacks(delegated_callbacks)
self._task: Final = loop.create_task(self._serve())
@overload
def _delegate(self, name: Literal["sampling_callback"]) -> SamplingFnT: ...
@overload
def _delegate(self, name: Literal["elicitation_callback"]) -> ElicitationFnT: ...
@overload
def _delegate(self, name: Literal["logging_callback"]) -> LoggingFnT: ...
def _delegate(self, name: _SessionCallbackName) -> SamplingFnT | ElicitationFnT | LoggingFnT:
match name:
case "sampling_callback":
async def sampling_callback(
context: ClientRequestContext, params: CreateMessageRequestParams
) -> CreateMessageResult | CreateMessageResultWithTools | ErrorData:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("sampling callback is unavailable")
return await callback(context, params)
return sampling_callback
case "elicitation_callback":
async def elicitation_callback(
context: ClientRequestContext, params: ElicitRequestParams
) -> ElicitResult | ErrorData:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("elicitation callback is unavailable")
return await callback(context, params)
return elicitation_callback
case "logging_callback":
async def logging_callback(params: LoggingMessageNotificationParams) -> None:
current_callback: Final = self._current_callbacks.get(name)
initial_callback: Final = self._initial_callbacks.get(name)
callback: Final = current_callback if current_callback is not None else initial_callback
if callback is None:
raise RuntimeError("logging callback is unavailable")
await callback(params)
return logging_callback
assert_never(name)
@property
def closed(self) -> bool:
return self._task.done()
@property
def reusable(self) -> bool:
return not self._ended and not self._retiring and not self._task.done()
def _record_stream_error(self, stream_error: asyncio.Future[Exception]) -> None:
self._stream_error = stream_error
def _mark_ended(self) -> None:
self._ended = True
async def _serve(self) -> None:
admission_stack: Final = AsyncExitStack()
async def drain(session: ClientSession) -> None:
await admission_stack.aclose()
self._ready.set_result(None)
while True:
operation, future, callbacks = await self._queue.get()
if not future.done():
self._active = future
if callbacks is not None:
self._current_callbacks = callbacks
async with self._admission() if self._admission is not None else nullcontext():
if not future.done():
(outcome,) = await asyncio.gather(operation(session), return_exceptions=True)
if _deliver_persistent_outcome(future, outcome, self._stream_error):
self._active = None
self._ended = True
return
self._active = None
if self._retiring and self._queue.empty():
return
try:
if self._admission is not None:
await admission_stack.enter_async_context(self._admission())
(ended,) = await asyncio.gather(
self._client.run_with_session(
drain,
quiet_on_error=True,
on_stream_error=self._record_stream_error,
on_cleanup=self._mark_ended,
),
return_exceptions=True,
)
self._ended = True
except asyncio.CancelledError:
self._ended = True
self._fail_waiters(None)
raise
finally:
await admission_stack.aclose()
self._fail_waiters(ended if isinstance(ended, Exception) else None)
def _fail_waiters(self, cause: Exception | None) -> None:
pending: Final = (self._ready, self._active, *(future for _, future, _ in self._drained()))
for future in pending:
if future is not None and not future.done():
future.set_exception(UpstreamSessionClosedError() if cause is None else cause)
def _drained(self) -> tuple[_PendingOperation, ...]:
return tuple(self._queue.get_nowait() for _ in range(self._queue.qsize()))
async def run(
self,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
*,
quiet_on_error: bool = False,
callbacks: _SessionCallbackMap | None = None,
) -> TSessionResult:
if self.closed:
await asyncio.shield(self._ready)
raise UpstreamSessionClosedError()
future: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future()
await self._queue.put((operation, future, callbacks))
try:
_ = await asyncio.wait((future, self._task), return_when=asyncio.FIRST_COMPLETED)
except asyncio.CancelledError:
_ = future.cancel()
raise
if not future.done():
raise UpstreamSessionClosedError()
return cast(TSessionResult, future.result()) # cast-ok: one queue serves operations of every result type
def close(self) -> None:
_ = self._task.cancel()
def retire(self) -> None:
if self._retiring:
return
self._retiring = True
if self.closed:
return
async def retired_operation(_: ClientSession) -> None:
return None
loop: Final = asyncio.get_running_loop()
future: Final[asyncio.Future[object]] = loop.create_future()
future.set_result(None)
try:
self._queue.put_nowait((retired_operation, future, None))
except asyncio.QueueFull:
return
async def wait_closed(self) -> None:
await asyncio.gather(self._task, return_exceptions=True)

2
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -38,7 +38,6 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.repositories.table_repositories import InvitationLinkRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
@ -52,6 +51,7 @@ from litellm.types.integrations.slack_alerting import (
VirtualKeyEvent,
WebhookEvent,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.model_deprecation import (
DEFAULT_DEPRECATION_CHECK_INTERVAL_SECONDS,
DEPRECATION_IDLE_POLL_SECONDS,

65
litellm/integrations/azure_storage/azure_storage.py
View file

@ -2,7 +2,7 @@ import asyncio
import os
import time
from collections.abc import Callable
from datetime import datetime, timedelta
from datetime import datetime, timedelta, timezone
from functools import cache
from typing import Final
from urllib.parse import unquote
@ -28,7 +28,7 @@ from litellm.secret_managers.get_azure_ad_token_provider import (
from litellm.types.secret_managers.get_azure_ad_token_provider import (
AzureCredentialType,
)
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
_ADLS_SAFE_NAME: Final = str.maketrans("/", "_", "=")
@ -160,6 +160,24 @@ class AzureBlobStorageLogger(CustomBatchLogger):
except Exception as e:
verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e)
async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None:
now: Final = datetime.now(timezone.utc)
file_path: Final = f"audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log['id']}.json"
await self._upload_json_to_file_path(file_path=file_path, json_payload=safe_dumps(audit_log))
async def _upload_json_to_file_path(self, file_path: str, json_payload: str) -> None:
payload_bytes: Final = json_payload.encode("utf-8")
if self.azure_storage_account_key:
await self._upload_bytes_with_account_key(file_path=file_path, content=payload_bytes)
return
await self.set_valid_azure_ad_token()
async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
base_url: Final = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{file_path}"
await self._create_file(async_client, base_url)
await self._append_data(async_client, base_url, json_payload)
await self._flush_data(async_client, base_url, len(payload_bytes))
async def async_send_batch(self):
"""
Sends the in memory logs queue to Azure Blob Storage
@ -366,41 +384,22 @@ class AzureBlobStorageLogger(CustomBatchLogger):
This is used when Azure Storage Account Key is set - Azure Storage Account Key does not work directly with Azure Rest API
"""
# Create an async service client
service_client: Final = await self.get_service_client()
# Get file system client
file_system_client: Final = service_client.get_file_system_client(file_system=self.azure_storage_file_system)
try:
# Create directory with today's date
from datetime import datetime
today: Final = datetime.now().strftime("%Y-%m-%d")
directory_client: Final = file_system_client.get_directory_client(today)
# check if the directory exists
if not await directory_client.exists():
await directory_client.create_directory()
verbose_logger.debug("Created directory: %s", today)
# Create a file client
file_name: Final = adls_safe_file_name(payload.get("id"))
file_client: Final = directory_client.get_file_client(file_name)
# Create the file
await file_client.create_file()
# Content to append
content: Final = safe_dumps(payload).encode("utf-8")
# Append content to the file
await file_client.append_data(data=content, offset=0, length=len(content))
# Flush the content to finalize the file
await file_client.flush_data(position=len(content), offset=0)
await self._upload_bytes_with_account_key(
file_path=f"{today}/{file_name}", content=safe_dumps(payload).encode("utf-8")
)
verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name)
except Exception as e:
verbose_logger.exception("Error occurred: %s", e)
async def _upload_bytes_with_account_key(self, file_path: str, content: bytes) -> None:
service_client: Final = await self.get_service_client()
file_client: Final = service_client.get_file_system_client(
file_system=self.azure_storage_file_system
).get_file_client(file_path)
await file_client.create_file()
await file_client.append_data(data=content, offset=0, length=len(content))
await file_client.flush_data(position=len(content), offset=0)

4
litellm/integrations/custom_guardrail.py
View file

@ -44,7 +44,7 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
dc: Final = DualCache()
@ -119,7 +119,7 @@ def is_guardrail_intervention(e: Exception) -> bool:
def _user_api_key_auth_from_request(request_data: Mapping[str, object]) -> "UserAPIKeyAuth":
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
metadata: Final = request_data.get(get_metadata_variable_name_from_kwargs(request_data))
stamped: Final[Mapping[str, object]] = metadata if isinstance(metadata, dict) else {}

2
litellm/integrations/custom_logger.py
View file

@ -35,12 +35,12 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import (
MCPPostCallResponseObject,
MCPPreCallRequestObject,
MCPPreCallResponseObject,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.router import PreRoutingHookResponse
Span = _Span

2
litellm/integrations/datadog/datadog.py
View file

@ -66,7 +66,7 @@ from ..additional_logging_utils import AdditionalLoggingUtils
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _DatadogLoggingKwargs(TypedDict, total=False):

4
litellm/integrations/otel/README.md
View file

@ -295,7 +295,9 @@ nothing here imports outside it:
- [`config.py`](./model/config.py) — `OpenTelemetryV2Config`, a pydantic-settings
model that reads `OTEL_*` / `LITELLM_OTEL_*` env vars, plus the feature gate.
`capture_span_content` gates whether prompt/response bodies may be written as
span attributes; it defaults **off** (`no_content`). The Baggage allowlists are
span attributes; it defaults **off** (`no_content`). A team or key can override
the global mode per OTel v2 destination with `capture_message_content`; omitted
values follow the proxy's own callback for that backend, else the global mode. The Baggage allowlists are
configurable, not hard-coded: set `LITELLM_OTEL_BAGGAGE_PROMOTED_KEYS` /
`LITELLM_OTEL_BAGGAGE_METADATA_KEYS` /
`LITELLM_OTEL_BAGGAGE_TEAM_METADATA_KEYS` (comma-separated) as env vars, or

8
litellm/integrations/otel/langfuse_logger.py
View file

@ -13,7 +13,7 @@ from litellm.integrations.otel.model.trace_controls import caller_trace_controls
from litellm.integrations.otel.plumbing.context import request_root_span
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import ModelResponseStream
@ -51,6 +51,10 @@ class LangfuseContentOpenTelemetryV2(LangfuseOpenTelemetryV2):
response: "AsyncIterator[ModelResponseStream]",
request_data: Mapping[str, object],
) -> "AsyncGenerator[ModelResponseStream, None]":
if not self._capture_span_content():
async for chunk in response:
yield chunk
return
relayed: Final[list[ModelResponseStream]] = [] # mutable-ok: relayed as they arrive, assembled at end of stream
async for chunk in response:
relayed.append(chunk)
@ -58,6 +62,8 @@ class LangfuseContentOpenTelemetryV2(LangfuseOpenTelemetryV2):
self._stamp_root_io(request_data, lambda: stream_output(tuple(relayed), request_data))
def _stamp_root_io(self, data: Mapping[str, object], render_output: Callable[[], str | None]) -> None:
if not self._capture_span_content():
return
root: Final = request_root_span()
if root is None or not root.is_recording():
return

19
litellm/integrations/otel/logger.py
View file

@ -58,6 +58,7 @@ from litellm.integrations.otel.plumbing.context import (
is_recordable_span,
mcp_message_transport_span,
post_response_root,
request_destinations,
request_root_http_route,
request_root_span,
resolve_internal_call_span_context,
@ -82,12 +83,13 @@ from litellm.integrations.otel.plumbing.providers import (
resolve_meter_provider,
)
from litellm.integrations.otel.plumbing.routing import TenantTracerCache
from litellm.types.utils import captures_span_content
if TYPE_CHECKING:
from opentelemetry.metrics import MeterProvider
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.services import ServiceLoggerPayload
from litellm.types.utils import (
CallTypesLiteral,
@ -254,6 +256,11 @@ class OpenTelemetryV2(CustomLogger):
return None
return GenAIEventRecorder(get_event_logger(provider, LITELLM_TRACER_NAME), provider.resource)
def _capture_span_content(self) -> bool:
return self.config.capture_span_content or any(
captures_span_content(destination.capture_message_content) for destination in request_destinations()
)
# ====================================================================== #
# Proxy global registration
# ====================================================================== #
@ -433,7 +440,7 @@ class OpenTelemetryV2(CustomLogger):
return False
payload: Final = cast("StandardLoggingPayload", raw_payload)
data: Final = MCPToolCallSpanData.from_standard_logging_payload(
payload, capture_content=self.config.capture_span_content
payload, capture_content=self._capture_span_content()
)
# A stray LLM carrier from a ``pre_call`` that mis-fired for this id would
# otherwise linger until evicted; drop it so it's neither leaked nor closed
@ -477,7 +484,7 @@ class OpenTelemetryV2(CustomLogger):
return False
payload: Final = cast("StandardLoggingPayload", raw_payload)
data: Final = MCPListToolsSpanData.from_standard_logging_payload(
payload, capture_content=self.config.capture_span_content
payload, capture_content=self._capture_span_content()
)
if data.identity.call_id:
self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None))
@ -569,7 +576,7 @@ class OpenTelemetryV2(CustomLogger):
return None
data: Final = LLMCallSpanData.from_standard_logging_payload(
payload,
capture_content=self.config.capture_span_content,
capture_content=self._capture_span_content(),
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
request_route=request_root_http_route(),
request_purpose=call.purpose,
@ -1076,6 +1083,6 @@ def build_otel_v2_logger(
def _logger_class(config: OpenTelemetryV2Config) -> type[OpenTelemetryV2]:
if "langfuse" not in config.mapper_names:
return OpenTelemetryV2
from litellm.integrations.otel.langfuse_logger import LangfuseContentOpenTelemetryV2, LangfuseOpenTelemetryV2
from litellm.integrations.otel.langfuse_logger import LangfuseContentOpenTelemetryV2
return LangfuseContentOpenTelemetryV2 if config.capture_span_content else LangfuseOpenTelemetryV2
return LangfuseContentOpenTelemetryV2

59
litellm/integrations/otel/model/config.py
View file

@ -16,19 +16,17 @@ from litellm.integrations.otel.model.baggage import (
)
from litellm.integrations.otel.model.spans import POSTGRESQL, db_system
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import OtelSpanScope
from litellm.types.utils import (
CAPTURE_MESSAGE_CONTENT_VALUES,
CaptureMessageContent,
OtelSpanScope,
captures_span_content,
)
#: Master feature-flag env var. The logger is inert until this is truthy.
OTEL_V2_ENV: Final = "LITELLM_OTEL_V2"
class CaptureMessageContent(str):
NO_CONTENT = "no_content"
SPAN_ONLY = "span_only"
EVENT_ONLY = "event_only"
SPAN_AND_EVENT = "span_and_event"
class ExporterOwner(str, Enum):
"""The preset that contributed an exporter. Values match the callback names
in ``presets.PRESET_BY_CALLBACK`` so per-request dynamic-credential routing
@ -297,9 +295,7 @@ class OpenTelemetryV2Config(BaseSettings):
here are lower_snake_case; normalizing at the boundary keeps both
spellings working and lets every downstream comparison stay exact.
"""
if isinstance(value, str):
return value.lower()
return value
return _normalized_capture_message_content(value)
@field_validator("langfuse_span_scope", mode="before")
@classmethod
@ -371,22 +367,39 @@ class OpenTelemetryV2Config(BaseSettings):
@property
def capture_span_content(self) -> bool:
"""Whether prompt/response content may be stamped as span attributes.
Defaults off (``no_content``): an operator must opt in before message
bodies leave the process, so a user request can never force its prompt
or completion into the configured backend while capture is disabled.
"""
return self.capture_message_content in (
CaptureMessageContent.SPAN_ONLY,
CaptureMessageContent.SPAN_AND_EVENT,
)
"""Whether the global setting permits prompt/response span attributes."""
return captures_span_content(self.capture_message_content)
@classmethod
def from_env(cls) -> "OpenTelemetryV2Config":
return cls()
def _normalized_capture_message_content(value: object) -> object:
if isinstance(value, str):
return value.lower()
return value
class _CaptureMessageContentSettings(BaseSettings):
model_config = SettingsConfigDict(extra="ignore", frozen=True)
capture_message_content: str = Field(
default=CaptureMessageContent.NO_CONTENT,
validation_alias=AliasChoices("OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT"),
)
@field_validator("capture_message_content", mode="before")
@classmethod
def _normalize_capture_message_content(cls, value: object) -> object:
return _normalized_capture_message_content(value)
def capture_message_content_from_env() -> str:
"""The global capture mode, read without building a settings model that rereads every other env var"""
return _CaptureMessageContentSettings().capture_message_content
_EXCLUDED_SERVICES_INPUT: Final[TypeAdapter[str | tuple[object, ...]]] = TypeAdapter(str | tuple[object, ...])
@ -395,6 +408,10 @@ def excluded_db_systems_from(value: object) -> frozenset[str]:
return _normalize_excluded_services(excluded_service_names(value))
def parse_capture_message_content(value: object) -> str | None:
return value if isinstance(value, str) and value in CAPTURE_MESSAGE_CONTENT_VALUES else None
def excluded_service_names(value: object) -> frozenset[str]:
"""Read a YAML list or comma-separated string of service names, logging and dropping unusable input
so a malformed value cannot stop the OTel logger from being built"""

19
litellm/integrations/otel/model/destination.py
View file

@ -11,7 +11,7 @@ from urllib.parse import quote
from pydantic import ConfigDict, Field
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import OtelSpanScope
from litellm.types.utils import OtelSpanScope, captures_span_content
class OtelDestination(LiteLLMBaseModel):
@ -50,6 +50,17 @@ class OtelDestination(LiteLLMBaseModel):
default=None,
description="Same, for the requests with a failed span in their tree; ``None`` forwards every one.",
)
capture_message_content: str | None = Field(
default=None,
description=(
"An explicit mode overrides the global capture policy for this destination. "
"Omitted follows the global setting."
),
)
def captures_content(self, default: str) -> bool:
setting: Final = self.capture_message_content
return captures_span_content(default if setting is None else setting)
def header_string(self) -> str:
"""Render headers as the ``k=v,k2=v2`` form an ``ExporterSpec`` expects.
@ -64,9 +75,9 @@ class OtelDestination(LiteLLMBaseModel):
def cache_key(self) -> tuple[str, tuple[tuple[str, str], ...], tuple[tuple[str, str], ...], str | None]:
"""Identity for processor reuse, so one destination means one exporter.
``span_scope`` and the sampling rates are left out on purpose: they decide which
spans reach the processor, not how the processor exports them, so a full and an
``llm_only`` view of the same account share one exporter.
``span_scope``, the sampling rates and ``capture_message_content`` are left out on
purpose: they decide which spans reach the processor and what they carry, not how
the processor exports them, so two views of the same account share one exporter.
"""
return (
self.endpoint,

2
litellm/integrations/otel/model/request_io.py
View file

@ -6,7 +6,7 @@ from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.integrations.otel.mappers.utils import json_or_none
from litellm.proxy.guardrails.anthropic_sse import assemble_anthropic_sse_stream, is_raw_sse_stream
from litellm.llms.anthropic.pass_through.stream_assembly import assemble_anthropic_sse_stream, is_raw_sse_stream
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import ResponseCompletedEvent, ResponsesAPIResponse
from litellm.types.utils import ModelResponse, ModelResponseStream

187
litellm/integrations/otel/plumbing/providers.py
View file

@ -5,7 +5,7 @@ import random
import threading
import time
from collections import OrderedDict
from collections.abc import Callable, Iterable, Mapping, Sequence
from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
@ -40,8 +40,19 @@ from opentelemetry.util.types import Attributes, AttributeValue
from litellm._logging import verbose_logger
from litellm._version import version as litellm_version
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_TRACE_NAME
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, OpenTelemetryV2Config
from litellm.integrations.otel.mappers.langfuse import (
LANGFUSE_OBSERVATION_INPUT,
LANGFUSE_OBSERVATION_OUTPUT,
LANGFUSE_TRACE_NAME,
)
from litellm.integrations.otel.model.config import (
CaptureMessageContent,
ExporterOwner,
ExporterSpec,
OpenTelemetryV2Config,
capture_message_content_from_env,
captures_span_content,
)
from litellm.integrations.otel.model.semconv import (
DB,
MCP,
@ -385,6 +396,26 @@ _CAPTURED_HEADER_PREFIXES: Final = ("http.request.header.", "http.response.heade
# ``?key=`` query parameter.
_URL_KEYS: Final = frozenset({"http.url", "http.target", "url.full"})
_URL_QUERY_KEY: Final = "url.query"
# Prompt and response content, in every vocabulary the mappers write it in once the
# global capture policy collects it: GenAI semconv (MCP tool-call arguments and results
# included), Langfuse, OpenInference's blobs and per-index messages, Langtrace, Weave.
_MESSAGE_CONTENT_KEYS: Final = frozenset(
{
GenAI.INPUT_MESSAGES,
GenAI.OUTPUT_MESSAGES,
GenAI.SYSTEM_INSTRUCTIONS,
GenAI.TOOL_CALL_ARGUMENTS,
GenAI.TOOL_CALL_RESULT,
LANGFUSE_OBSERVATION_INPUT,
LANGFUSE_OBSERVATION_OUTPUT,
"input.value",
"output.value",
"llm.prompts",
"llm.completions",
"weave.output",
}
)
_INDEXED_MESSAGE_PREFIXES: Final = ("llm.input_messages.", "llm.output_messages.")
class _SpanView(ReadableSpan):
@ -415,6 +446,19 @@ class _SpanView(ReadableSpan):
)
def _is_message_content(key: str) -> bool:
return key in _MESSAGE_CONTENT_KEYS or key.startswith(_INDEXED_MESSAGE_PREFIXES)
def without_message_content(span: ReadableSpan) -> ReadableSpan:
"""``span`` without its prompt and response content, everything else left as it was."""
attributes: Final = span.attributes or _NO_ATTRIBUTES
if not any(_is_message_content(key) for key in attributes):
return span
kept: Final = MappingProxyType({key: value for key, value in attributes.items() if not _is_message_content(key)})
return _SpanView(span, span.resource, kept, span.events, span.status, parent=span.parent)
def _is_database_span(attributes: Mapping[str, AttributeValue]) -> bool:
return any(key in attributes for key in _DB_SYSTEM_KEYS)
@ -553,7 +597,11 @@ def _tenant_resource(resource: Resource, defaults: Mapping[str, str], extra: Map
return filled.merge(Resource(extra)) if extra else filled
def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> ReadableSpan:
def _for_destination(
span: ReadableSpan,
destination: "OtelDestination",
default_capture: str,
) -> ReadableSpan:
"""The view of ``span`` a tenant destination receives.
A span the tenant's own call produced keeps its error text. Every other span is
@ -564,8 +612,10 @@ def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> Read
exception it raised and names the operator's guardrail endpoint. Stack traces walk
the operator's install and come off every span, as do the headers the operator
captures on the server span, whose request side holds the caller's bearer token,
and the query string of the request URL, which can hold the same key. The span
itself stays, so the tenant still gets the whole trace tree.
and the query string of the request URL, which can hold the same key. A
destination whose effective setting disables span content loses the prompt and
response content too. The span itself stays, so the tenant still gets the whole
trace tree.
"""
extra: Final = destination.resource_attributes
defaults: Final = destination.resource_defaults
@ -573,11 +623,12 @@ def _for_destination(span: ReadableSpan, destination: "OtelDestination") -> Read
database: Final = _is_database_span(attributes)
owned: Final = _is_tenant_owned_span(attributes)
unreachable: Final = _guardrail_unreachable(attributes)
redacted: Final = not destination.captures_content(default_capture)
kept: Final = MappingProxyType(
{
key: _without_query(key, value)
for key, value in attributes.items()
if _tenant_visible(key, database, owned, unreachable)
if _tenant_visible(key, database, owned, unreachable) and not (redacted and _is_message_content(key))
}
)
recorded: Final = span.events
@ -612,6 +663,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
shutdown_drain_seconds: float = _SHUTDOWN_DRAIN_SECONDS,
operator_sinks: 'Mapping[_SinkKey, "OtelSpanScope"]' = MappingProxyType({}),
excluded_db_systems: frozenset[str] = frozenset(),
default_capture: str = CaptureMessageContent.NO_CONTENT,
default_capture_by_backend: Mapping[str, str] = MappingProxyType({}),
pending_drains: int = _MAX_PENDING_DRAINS,
drain_pool: _DrainPool | None = None,
sampling_draw: Callable[[], float] | None = None,
@ -619,6 +672,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
self._operator_sinks: Final = operator_sinks
self._draw: Final = sampling_draw if sampling_draw is not None else random.random
self._excluded_db_systems: Final = excluded_db_systems
self._default_capture: Final = default_capture
self._default_capture_by_backend: Final = default_capture_by_backend
self._drain_seconds: Final = shutdown_drain_seconds
self._lock: Final = threading.Condition()
self._closed = False # guarded by ``_lock``: an unlocked read races the teardown it gates
@ -661,12 +716,22 @@ class TenantFanOutSpanProcessor(SpanProcessor):
for held_span, held_destination in self._settle(trace_id, failed):
self._forward(held_span, held_destination)
def _default_capture_for(self, destination: "OtelDestination") -> str:
backend: Final = destination.callback_name
return (
self._default_capture_by_backend.get(backend, self._default_capture) if backend else self._default_capture
)
def _forward(self, span: ReadableSpan, destination: "OtelDestination") -> None:
processor: Final = self._acquire(destination)
if processor is None:
return
try:
processor.on_end(_scoped(_for_destination(span, destination), destination.span_scope))
processor.on_end(
_scoped(
_for_destination(span, destination, self._default_capture_for(destination)), destination.span_scope
)
)
except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span
verbose_logger.debug("OTel V2 fan-out: forwarding to %s failed: %s", destination.endpoint, exc)
finally:
@ -963,14 +1028,15 @@ class _OverriddenBackendFilter(SpanProcessor):
Wrapping is the only place this works: ``SynchronousMultiSpanProcessor.on_end``
ignores return values, so a sibling processor can never veto the export.
Under ``additive`` mode nothing is suppressed, so the wrapper passes every span
straight through and the operator keeps its copy.
Under ``additive`` mode spans are not suppressed, though global content policy
still strips message content when disabled.
``scope`` narrows what the exporter receives independently of that: under
``llm_only`` the model-call spans go through as trace roots and the rest of the
tree is held back, unless a destination of the request names ``sink``, the account
this exporter writes to, with a wider scope: the fan-out then delivers the rest of
the tree there and the model call keeps its place in it.
the tree there and the model call keeps its place in it. A destination naming
``sink`` sets this exporter's content policy; otherwise it follows global capture.
"""
def __init__(
@ -979,11 +1045,15 @@ class _OverriddenBackendFilter(SpanProcessor):
owner: str | None,
scope: "OtelSpanScope" = "full",
sink: _SinkKey | None = None,
global_captures: bool = True,
filter_content: bool = True,
) -> None:
self._inner: Final = inner
self._owner: Final = owner
self._scope: Final = scope
self._sink: Final = sink
self._global_captures: Final = global_captures
self._filter_content: Final = filter_content
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
self._inner.on_start(span, parent_context)
@ -991,7 +1061,28 @@ class _OverriddenBackendFilter(SpanProcessor):
def on_end(self, span: ReadableSpan) -> None:
if self._owner in suppressed_backends() or not _in_scope(span, self._scope):
return
self._inner.on_end(_scoped(span, self._account_scope()))
scoped: Final = _scoped(span, self._account_scope())
self._inner.on_end(
without_message_content(scoped) if self._filter_content and not self._account_captures_content() else scoped
)
def _account_captures_content(self) -> bool:
"""Whether a destination of this request that names ``sink`` captures content.
The fan-out skips a destination this exporter already writes to, so this copy is
the one that account gets and has to honor the destination's content mode.
"""
explicit: Final = next(
(
destination.capture_message_content
for destination in request_destinations()
if self._sink is not None
and destination.capture_message_content is not None
and _sink_key(destination.endpoint, destination.headers) == self._sink
),
None,
)
return captures_span_content(explicit) if explicit is not None else self._global_captures
def _account_scope(self) -> "OtelSpanScope":
if self._scope == "full" or self._sink is None:
@ -1010,6 +1101,25 @@ class _OverriddenBackendFilter(SpanProcessor):
return self._inner.force_flush(timeout_millis)
class _MessageContentFilter(SpanProcessor):
"""Strip prompt and response content from every span before ``inner`` exports it."""
def __init__(self, inner: SpanProcessor) -> None:
self._inner: Final = inner
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
self._inner.on_start(span, parent_context)
def on_end(self, span: ReadableSpan) -> None:
self._inner.on_end(without_message_content(span))
def shutdown(self) -> None:
self._inner.shutdown()
def force_flush(self, timeout_millis: int = 30000) -> bool:
return self._inner.force_flush(timeout_millis)
def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter:
"""Build a single exporter from the top-level config fields.
@ -1261,12 +1371,23 @@ def build_resource(config: OpenTelemetryV2Config) -> Resource:
return Resource.create(attributes)
def _spec_processor(
spec: ExporterSpec, use_simple_processor: bool | None, strip_message_content: bool
) -> SpanProcessor:
exporting: Final = _processor_for(
_exporter_from_spec(spec),
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
)
return _MessageContentFilter(exporting) if strip_message_content else exporting
def build_tracer_provider(
config: OpenTelemetryV2Config,
exporter: SpanExporter | None = None,
baggage_processor: SpanProcessor | None = None,
use_simple_processor: bool | None = None,
tenant_overrides: bool = False,
strip_message_content: bool = False,
) -> TracerProvider:
"""Build the shared :class:`TracerProvider`.
@ -1285,6 +1406,9 @@ def build_tracer_provider(
``config.langfuse_span_scope`` narrows the exporter owned by ``langfuse_otel``
alone; a collector or any other backend in the same config keeps the full tree.
``strip_message_content`` removes prompt and response content before every
exporter, for a routed provider whose spans can carry content another
destination asked for.
"""
provider: Final = TracerProvider(resource=build_resource(config))
if baggage_processor is None:
@ -1300,17 +1424,23 @@ def build_tracer_provider(
for spec in config.exporters:
if spec.requires_headers and not spec.headers:
continue
exp = _exporter_from_spec(spec)
processor = _processor_for(
exp,
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
)
processor = _spec_processor(spec, use_simple_processor, strip_message_content)
owner = spec.owner.value if tenant_overrides and spec.owner is not None else None
scope = _operator_scope(config, spec)
sink = _sink_key(spec.endpoint, parse_headers(spec.headers)) if _exports_to_the_wire(spec) else None
provider.add_span_processor(
_OverriddenBackendFilter(processor, owner, scope, sink)
if owner is not None or scope != "full"
_OverriddenBackendFilter(
processor,
owner,
scope,
sink,
config.capture_span_content,
tenant_overrides,
)
if owner is not None
or scope != "full"
or sink is not None
or (tenant_overrides and not config.capture_span_content)
else processor
)
return provider
@ -1340,10 +1470,29 @@ def attach_tenant_fan_out(
TenantFanOutSpanProcessor(
operator_sinks=operator_sink_scopes(*configs),
excluded_db_systems=excluded_db_systems,
default_capture=_operator_capture_mode(configs),
default_capture_by_backend=MappingProxyType(dict(_owned_capture_modes(configs))),
)
)
def _operator_capture_mode(configs: Sequence[OpenTelemetryV2Config]) -> str:
"""The mode of the operator's generic ``otel`` callback, so ``callback_settings.otel`` counts as the global default"""
return next(
(config.capture_message_content for config in configs if any(spec.owner is None for spec in config.exporters)),
capture_message_content_from_env(),
)
def _owned_capture_modes(
configs: Sequence[OpenTelemetryV2Config],
) -> Iterator[tuple[str, str]]:
for config in configs:
for spec in config.exporters:
if spec.owner is not None:
yield spec.owner.value, config.capture_message_content
def deliverable_destinations(
destinations: Iterable["OtelDestination"],
provider: trace.TracerProvider | None = None,

3
litellm/integrations/otel/plumbing/routing.py
View file

@ -282,7 +282,8 @@ class TenantTracerCache:
self._providers.move_to_end(cache_key)
return cached
built: Final = build_tracer_provider(
self._routed_config(credential_headers, project_headers, endpoint, service_name)
self._routed_config(credential_headers, project_headers, endpoint, service_name),
strip_message_content=not self._config.capture_span_content,
)
self._providers[cache_key] = built
return built

2
litellm/integrations/otel/presets/destinations.py
View file

@ -187,6 +187,7 @@ def destination_for(
callback_name: str,
params: StandardCallbackDynamicParams,
service_name: str | None = None,
capture_message_content: str | None = None,
) -> OtelDestination | None:
"""The destination ``params`` names for ``callback_name``, or ``None``.
@ -217,6 +218,7 @@ def destination_for(
callback_name=callback_name,
protocol=protocol,
span_scope=_span_scope(callback_name, params),
capture_message_content=capture_message_content,
success_sampling_rate=success_sampling_rate,
error_sampling_rate=error_sampling_rate,
)

7
litellm/integrations/otel/presets/newrelic.py
View file

@ -9,6 +9,7 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
from litellm._logging import verbose_logger
from litellm.integrations.otel.model.config import (
CaptureMessageContent,
ExporterOwner,
ExporterSpec,
OpenTelemetryV2Config,
@ -66,7 +67,11 @@ def newrelic_preset(
# New Relic ingests the OTLP GenAI semantic conventions natively.
"mapper_names": ensure_mappers(base.mapper_names, "genai"),
**(
{"capture_message_content": ("span_only" if settings.record_content else "no_content")}
{
"capture_message_content": (
CaptureMessageContent.SPAN_ONLY if settings.record_content else CaptureMessageContent.NO_CONTENT
)
}
if settings.record_content is not None
else {}
),

4
litellm/integrations/prometheus.py
View file

@ -21,6 +21,7 @@ import litellm
from litellm._internal_context import with_service_target
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import (
AUTH_OBJECTS_TARGET,
PROMETHEUS_OVERFLOW_SERIES_LABEL_VALUE,
PROXY_LLM_PROVIDER_FALLBACK,
PROXY_REJECTED_BEFORE_ROUTING_KEY,
@ -52,8 +53,6 @@ from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.user import LiteLLM_UserTable
from litellm.models.verification_token import LiteLLM_DeletedVerificationToken
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.organization_repository import OrganizationRepository
@ -67,6 +66,7 @@ from litellm.types.integrations.prometheus import (
sanitize_prometheus_label_value,
validate_prometheus_deployment_and_latency_caller_identity,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.carried_budget_state import (
KeyBudgetSnapshot,
OrgBudgetSnapshot,

2
litellm/integrations/rubrik.py
View file

@ -44,7 +44,7 @@ if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_WEBHOOK_PATH_RESPONSE_MODERATION: Final = "/v1/after_completion/openai/v1"
_WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1"

2
litellm/integrations/shadow_eval_logger.py
View file

@ -631,13 +631,13 @@ async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
"""
try:
from litellm.exceptions import BudgetExceededError
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
get_team_object,
team_max_budget_check,
virtual_key_max_budget_check,
)
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
except ImportError:
return False

2
litellm/integrations/websearch_interception/handler.py
View file

@ -77,10 +77,10 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper

4
litellm/litellm_core_utils/health_check_helpers.py
View file

@ -122,8 +122,8 @@ class HealthCheckHelpers:
2. `user_api_key_auth`: This helps identify health check calls in the DB.
We need this since the DB requires an API Key to track a log in the SpendLogs Table
"""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_metadata_variable_name: Final = "litellm_metadata"
litellm_metadata: Final = HealthCheckHelpers._get_metadata_for_health_check_call()
@ -311,7 +311,7 @@ class HealthCheckHelpers:
**_filter_model_params(model_params=model_params),
document=_ocr_health_check_document(model=model, custom_llm_provider=custom_llm_provider),
),
"evaluation": lambda: litellm.adecisions(
"evaluation": lambda: litellm.asystemone(
**DECISIONS_CALL_PARAMS.validate_python(
{
"state": prompt or "health check",

8
litellm/litellm_core_utils/initialize_dynamic_callback_params.py
View file

@ -6,6 +6,7 @@ from typing import Any, Final
from litellm.types.utils import (
ARIZE_OTLP_PROTOCOLS,
CAPTURE_MESSAGE_CONTENT_VALUES,
OTEL_SPAN_SCOPES,
TRUSTED_CALLBACK_VARS_FIELD,
StandardCallbackDynamicParams,
@ -91,6 +92,13 @@ def validate_arize_otlp_protocol_value(value: str) -> None:
raise ValueError(f"Invalid arize_otlp_protocol {value!r}: must be one of {sorted(ARIZE_OTLP_PROTOCOLS)}")
def validate_capture_message_content_value(value: str) -> None:
if value not in CAPTURE_MESSAGE_CONTENT_VALUES:
raise ValueError(
f"Invalid capture_message_content {value!r}: must be one of {sorted(CAPTURE_MESSAGE_CONTENT_VALUES)}"
)
# Hardcoded list of supported callback params to avoid runtime inspection issues with TypedDict
_supported_callback_params: Final[tuple[str, ...]] = (
"langfuse_public_key",

93
litellm/litellm_core_utils/resource_ownership.py Normal file
View file

@ -0,0 +1,93 @@
from typing import Final
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
def is_proxy_admin(user_api_key_dict: UserAPIKeyAuth | None) -> bool:
if user_api_key_dict is None:
return False
return (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
def get_resource_owner_scopes(
user_api_key_dict: UserAPIKeyAuth | None,
) -> list[str]:
"""
Return ownership scopes that may access a user-created proxy resource.
Raw user_id is included for rows created before scope prefixes existed.
Prefixes avoid collisions when falling back to team/org/key ownership
for keys that do not have a user_id.
Identity-less callers (no user_id, team_id, org_id, api_key, or token)
return ``[]`` — they share no scope with any other caller, so access
checks against an existing owner always fail and creates that depend
on a primary scope must reject up front. Returning a shared sentinel
here would let any two identity-less callers see each other's data.
"""
if user_api_key_dict is None:
return []
scopes: Final[list[str]] = []
def _add(scope: str | None) -> None:
if scope and scope not in scopes:
scopes.append(scope)
if user_api_key_dict.user_id:
_add(user_api_key_dict.user_id)
_add(f"user:{user_api_key_dict.user_id}")
if user_api_key_dict.team_id:
_add(f"team:{user_api_key_dict.team_id}")
if user_api_key_dict.org_id:
_add(f"org:{user_api_key_dict.org_id}")
if user_api_key_dict.api_key:
_add(f"key:{user_api_key_dict.api_key}")
if user_api_key_dict.token:
_add(f"key:{user_api_key_dict.token}")
return scopes
def get_primary_resource_owner_scope(
user_api_key_dict: UserAPIKeyAuth | None,
) -> str | None:
"""Return the canonical owner scope to stamp on newly-created rows.
``None`` for identity-less callers — callers that depend on a primary
scope to record ownership must surface that as a hard error rather
than fall back to a shared sentinel (which would collapse every
identity-less caller into the same logical owner).
"""
if user_api_key_dict is None:
return None
if user_api_key_dict.user_id:
return user_api_key_dict.user_id
if user_api_key_dict.team_id:
return f"team:{user_api_key_dict.team_id}"
if user_api_key_dict.org_id:
return f"org:{user_api_key_dict.org_id}"
if user_api_key_dict.api_key:
return f"key:{user_api_key_dict.api_key}"
if user_api_key_dict.token:
return f"key:{user_api_key_dict.token}"
return None
def user_can_access_resource_owner(
owner: str | None,
user_api_key_dict: UserAPIKeyAuth | None,
) -> bool:
if user_api_key_dict is None:
return True
if is_proxy_admin(user_api_key_dict):
return True
if owner is None:
return False
return owner in get_resource_owner_scopes(user_api_key_dict)

5
litellm/litellm_core_utils/served_output_texts.py
View file

@ -110,7 +110,10 @@ def served_output_texts(response: object) -> ServedTexts | None:
def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks):
return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)))
from litellm.proxy.guardrails.anthropic_sse import assemble_anthropic_sse_stream, is_anthropic_sse_stream
from litellm.llms.anthropic.pass_through.stream_assembly import (
assemble_anthropic_sse_stream,
is_anthropic_sse_stream,
)
if not is_anthropic_sse_stream(chunks):
return None

2
litellm/llms/a2a/chat/guardrail_translation/handler.py
View file

@ -26,7 +26,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _A2ATextPart(TypedDict, total=False):

2
litellm/llms/anthropic/cache_aware_routing.py
View file

@ -24,7 +24,7 @@ from litellm.llms.anthropic.prompt_cache_prediction import (
resolve_prediction_target,
supported_prediction_headers,
)
from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
from litellm.llms.anthropic.prompt_cache_pricing import price_cache_tokens
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.management_endpoints.prompt_cache_prediction import (
CacheCostScenario,

2
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -76,10 +76,10 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@dataclass(frozen=True, slots=True)

2
litellm/llms/anthropic/pass_through/adapters/handler.py
View file

@ -33,8 +33,8 @@ from litellm.types.utils import ModelResponse
from litellm.utils import get_model_info
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
# Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config", "safeguards"})

2
litellm/llms/anthropic/pass_through/context_management/dispatcher.py
View file

@ -13,8 +13,8 @@ from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112
from .result import PolyfillResult
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
EditorResult: TypeAlias = "PolyfillResult | tuple[list[dict[str, object]], AppliedEdit | None]"

68
litellm/llms/anthropic/pass_through/context_management/editors/compact.py
View file

@ -29,7 +29,6 @@ from litellm.types.llms.anthropic import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RateLimitDescriptor,
RateLimitDescriptorRateLimitObject,
@ -41,6 +40,7 @@ if TYPE_CHECKING:
AllAnthropicToolsValues,
)
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import ModelResponse
from ..constants import (
@ -84,6 +84,7 @@ _PROPAGATED_METADATA_KEYS: Final = (
"user_api_key_model_max_budget",
"user_api_key_team_model_max_budget",
"user_api_key_user_model_max_budget",
"user_api_key_team_member_model_max_budget",
"user_api_key_end_user_model_max_budget",
"litellm_call_id",
"litellm_parent_otel_span",
@ -125,6 +126,23 @@ class _SummaryAcompletion(Protocol):
) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ...
class _SummaryModelBudgetLimiter(Protocol):
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
async def is_end_user_within_model_budget(
self,
end_user_id: str,
end_user_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
class _CreateRateLimitDescriptors(Protocol):
def __call__(
self,
@ -392,6 +410,42 @@ async def _check_summary_model_access(
return True
async def _check_summary_team_member_model_budget(
user_api_key_auth: "UserAPIKeyAuth",
model_max_budget_limiter: _SummaryModelBudgetLimiter,
summary_model: str,
) -> bool:
team_member_model_max_budget: Final = user_api_key_auth.team_member_model_max_budget
user_id: Final = user_api_key_auth.user_id
team_id: Final = user_api_key_auth.team_id
if (
not isinstance(team_member_model_max_budget, Mapping)
or not team_member_model_max_budget
or not isinstance(user_id, str)
or not isinstance(team_id, str)
):
return True
try:
await model_max_budget_limiter.is_team_member_within_model_budget(
user_id=user_id,
team_id=team_id,
team_member_model_max_budget=team_member_model_max_budget,
model=summary_model,
)
except litellm.BudgetExceededError:
return False
except Exception as e: # noqa: BLE001 # unexpected budget-check errors must deny compaction
verbose_logger.warning(
"compact_20260112: unexpected error during team member model-budget "
"check for summary_model=%s; denying: %s",
summary_model,
e,
)
return False
return True
async def _check_summary_model_budget(
user_api_key_auth: Optional["UserAPIKeyAuth"],
summary_model: str,
@ -407,9 +461,9 @@ async def _check_summary_model_budget(
per-model budget is configured.
Every scope is checked because the summary's spend is charged to every
scope: this file propagates the key, team, user and end-user budgets into the
subrequest's metadata, so skipping one of them would let compaction
increment a counter it can never be refused by.
scope: this file propagates the key, team, user, team-member and end-user
budgets into the subrequest's metadata, so skipping one of them would let
compaction increment a counter it can never be refused by.
"""
if user_api_key_auth is None:
return True
@ -497,7 +551,11 @@ async def _check_summary_model_budget(
)
return False
return True
return await _check_summary_team_member_model_budget(
user_api_key_auth=user_api_key_auth,
model_max_budget_limiter=model_max_budget_limiter,
summary_model=summary_model,
)
def _without_parallel_request_gauges(

124
litellm/llms/anthropic/pass_through/stream_assembly.py
View file

@ -1,5 +1,6 @@
import codecs
import json
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import Final
import litellm
@ -395,3 +396,124 @@ def build_usage_only_response_from_chunks(
],
usage=usage_obj,
)
_ANTHROPIC_EVENT_TYPES: Final = frozenset(
{
"message_start",
"message_delta",
"message_stop",
"content_block_start",
"content_block_delta",
"content_block_stop",
"ping",
"error",
}
)
def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool:
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None:
raw: Final = b"".join(
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
for chunk in all_chunks
if isinstance(chunk, (str, bytes))
)
try:
return codecs.getincrementaldecoder("utf-8")().decode(raw, final=False)
except UnicodeDecodeError:
return None
def _parsed_sse_events(sse_stream: str) -> tuple[Mapping[str, object], ...]:
return tuple(
event_data
for event in split_sse_chunk_into_events(sse_stream)
if (event_data := extract_sse_data(event)) is not None
)
def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
return next(
(
message
for event_data in _parsed_sse_events(sse_stream)
if event_data.get("type") == "message_start" and isinstance(message := event_data.get("message"), dict)
),
None,
)
def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool:
"""Whether raw SSE frames are Anthropic Messages events.
``is_raw_sse_stream`` only says the chunks are unparsed bytes, and ``/v1/messages`` is not the
only endpoint that streams those: the Google ``:streamGenerateContent`` route marks its own
stream raw too. Reading its frames as Anthropic ones would refuse the response in a wire format
its client cannot parse, so the surface is decided on the event types actually present.
"""
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return False
return any(event.get("type") in _ANTHROPIC_EVENT_TYPES for event in _parsed_sse_events(sse_stream))
def assemble_anthropic_sse_stream(
all_chunks: Sequence[object], *, restore_identity: bool = False
) -> ModelResponse | None:
"""Assemble raw Anthropic SSE frames into a ModelResponse.
``restore_identity`` stamps the upstream message id and model onto the result, which the
assembler does not carry through. It is off by default so callers that re-emit the assembled
response keep the wire shape they had before this helper was shared. The writes land on a
freshly built object that is unreachable from caller state until returned.
"""
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return None
message_start: Final = _anthropic_message_start(sse_stream)
if message_start is None:
return None
model: Final = message_start.get("model") if restore_identity else None
try:
assembled: Final = build_complete_streaming_response(
all_chunks=(sse_stream,),
litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None
model=model if isinstance(model, str) else "",
)
except Exception: # noqa: BLE001 # stream_chunk_builder re-raises every assembly failure as litellm.APIError
return None
if not isinstance(assembled, ModelResponse):
return None
if not restore_identity:
return assembled
message_id: Final = message_start.get("id")
if isinstance(message_id, str):
assembled.id = message_id
if isinstance(model, str) and model:
assembled.model = model
return assembled
def is_sse_error_stream(all_chunks: Sequence[object]) -> bool:
"""Whether the buffered stream carries nothing but error frames.
post_call guardrails run in a chain, so a hook can be handed the terminal error frames an
earlier guardrail emitted when it blocked. Those carry no message to assemble, and replacing
them would hide the refusal the client is owed. Covers both wire forms a guardrail emits: the
Anthropic ``error`` event and the chat-completions ``{"error": ...}`` payload.
"""
if not all(isinstance(chunk, (str, bytes)) for chunk in all_chunks):
# A stream mixing typed chunks with an error frame still carries content to scan, and the
# frames-only join below would drop exactly the part that has to be scanned
return False
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return False
events: Final = _parsed_sse_events(sse_stream)
return len(events) > 0 and all(
event.get("type") == "error" or isinstance(event.get("error"), Mapping) for event in events
)

100
litellm/llms/anthropic/prompt_cache_pricing.py Normal file
View file

@ -0,0 +1,100 @@
from collections.abc import Mapping
from datetime import datetime, timezone
from math import isfinite
from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm.cost_calculator import (
completion_cost, # pyright: ignore[reportUnknownVariableType] # legacy optional parameters are untyped
select_model_name_for_cost_calc,
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
_PRICE_ENTRY: Final = TypeAdapter(Mapping[str, object])
def _valid_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and isfinite(value) and value >= 0
def _has_required_prices(prices: Mapping[str, object], tokens: CacheTokenBuckets, completion_tokens: int = 0) -> bool:
required: Final = (
("input_cost_per_token", True),
("cache_read_input_token_cost", tokens.cache_read_input_tokens > 0),
("cache_creation_input_token_cost", tokens.cache_creation_5m_input_tokens > 0),
("cache_creation_input_token_cost_above_1hr", tokens.cache_creation_1h_input_tokens > 0),
("output_cost_per_token", completion_tokens > 0),
)
if any(needed and not _valid_price(prices.get(key)) for key, needed in required):
return False
return all(
_valid_price(value)
for key, value in prices.items()
if value is not None and any(needed and key.startswith(f"{base}_above_") for base, needed in required)
)
def price_cache_tokens(
model: str, deployment_id: str, tokens: CacheTokenBuckets, completion_tokens: int = 0
) -> float | None:
try:
selected_model: Final = select_model_name_for_cost_calc(
model=model,
completion_response=None,
custom_pricing=True,
custom_llm_provider="anthropic",
router_model_id=deployment_id,
)
if selected_model is None:
return None
model_info: Final = litellm.get_model_info(model=selected_model, custom_llm_provider="anthropic")
registry: Final = _PRICE_ENTRY.validate_python(litellm.model_cost) # pyright: ignore[reportUnknownMemberType] # legacy registry is validated at this boundary
price_entry: Final = registry.get(model_info["key"])
if price_entry is None:
return None
prices: Final = _PRICE_ENTRY.validate_python(price_entry)
if completion_tokens < 0 or not _has_required_prices(prices, tokens, completion_tokens):
return None
usage: Final = Usage(
prompt_tokens=tokens.total_tokens,
completion_tokens=completion_tokens,
total_tokens=tokens.total_tokens + completion_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=tokens.cache_read_input_tokens,
cache_creation_tokens=tokens.cache_creation_5m_input_tokens + tokens.cache_creation_1h_input_tokens,
cache_creation_token_details=CacheCreationTokenDetails(
ephemeral_5m_input_tokens=tokens.cache_creation_5m_input_tokens,
ephemeral_1h_input_tokens=tokens.cache_creation_1h_input_tokens,
),
),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.now(timezone.utc),
litellm_call_id="prompt-cache-prediction",
function_id="prompt-cache-prediction",
)
completion_cost(
completion_response=ModelResponse(model=model, usage=usage),
model=model,
custom_llm_provider="anthropic",
custom_pricing=True,
router_model_id=deployment_id,
litellm_logging_obj=logging_obj,
)
breakdown: Final = logging_obj.cost_breakdown
input_cost: Final = breakdown.get("input_cost") if breakdown is not None else None
output_cost: Final = breakdown.get("output_cost") if breakdown is not None else None
if input_cost is None or output_cost is None:
return None
cost: Final = input_cost + output_cost
return cost if _valid_price(cost) else None
except Exception: # noqa: BLE001 # the shared pricing owners raise plain Exception for unpriceable models
return None

2
litellm/llms/base_llm/files/transformation.py
View file

@ -22,9 +22,9 @@ from ..chat.transformation import BaseConfig
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
from litellm.proxy._types import UserAPIKeyAuth as _UserAPIKeyAuth
from litellm.router import Router as _Router
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth as _UserAPIKeyAuth
LiteLLMLoggingObj = _LiteLLMLoggingObj
Span = Any

2
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -11,8 +11,8 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@dataclass(slots=True)

2
litellm/llms/base_llm/managed_resources/base_managed_resource.py
View file

@ -16,7 +16,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import SpecialEnums
MANAGED_RESOURCES_TARGET: Final = "managed_resources"

6
litellm/llms/base_llm/managed_resources/isolation.py
View file

@ -12,10 +12,10 @@ are denied so an empty user_id can never select an unscoped query.
from typing import Any, Final
from litellm.proxy._types import (
from litellm.types.proxy.auth.user_api_key_auth import (
UserAPIKeyAuth,
)
from litellm.proxy._types import (
from litellm.types.proxy.auth.user_api_key_auth import (
user_api_key_has_admin_view as _user_has_admin_view,
)
@ -29,7 +29,7 @@ def resolve_resource_owner_id(
A key with neither a user_id nor a team_id would otherwise stamp
``created_by=None`` and be locked out of its own resources, so it owns
them under its hashed token instead, using the ``key:`` scope prefix
already used by ``proxy/common_utils/resource_ownership.py``. ``None``
already used by ``litellm_core_utils/resource_ownership.py``. ``None``
means the caller has no usable identity of its own and must fall back
to team scoping, or be denied.
"""

18
litellm/llms/bedrock/passthrough/guardrail_translation/handler.py
View file

@ -17,8 +17,8 @@ if TYPE_CHECKING:
from litellm.llms.pass_through.guardrail_translation.handler import (
PassThroughEndpointHandler,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"})
_EVENT_STREAM_CONTENT_TYPE: Final = "vnd.amazon.eventstream"
@ -206,12 +206,14 @@ class _TextContainer(Protocol):
_DeltaHolder = tuple[_GroupKey, _TextContainer, str]
_StreamGroupKey: TypeAlias = tuple[int | None, _GroupKey]
class _StreamFrame(TypedDict):
"""One raw event-stream frame plus the guardrailable texts it carries."""
raw: ReadOnly[bytes]
texts: ReadOnly[Sequence[tuple[_GroupKey, str]]]
texts: ReadOnly[Sequence[tuple[_StreamGroupKey, str]]]
def _unpack_uint32(buffer: bytes) -> int:
@ -302,12 +304,14 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
frames.append({"raw": frame_raw, "texts": []})
continue
texts: list[tuple[_GroupKey, str]] = []
texts: list[tuple[_StreamGroupKey, str]] = []
if event_type == "contentBlockDelta":
try:
payload_dict: dict[str, object] = _json.loads(payload_bytes)
block_index = payload_dict.get("contentBlockIndex")
# Group per content block: a guardrail edit that changes length must not move text between blocks.
texts = [
(group_key, container[key])
((block_index if isinstance(block_index, int) else None, group_key), container[key])
for group_key, container, key in _collect_stream_delta_text_holders(payload_dict.get("delta"))
]
except Exception as e:
@ -321,9 +325,9 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
trailing_bytes: Final = body_bytes[offset:]
group_order: Final[list[_GroupKey]] = []
group_members: Final[dict[_GroupKey, list[tuple[int, int]]]] = {}
group_texts: Final[dict[_GroupKey, list[str]]] = {}
group_order: Final[list[_StreamGroupKey]] = []
group_members: Final[dict[_StreamGroupKey, list[tuple[int, int]]]] = {}
group_texts: Final[dict[_StreamGroupKey, list[str]]] = {}
for frame_idx, frame in enumerate(frames):
for local_idx, (group_key, text) in enumerate(frame["texts"]):
if group_key not in group_members:

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -213,7 +213,6 @@ if TYPE_CHECKING:
FakeAnthropicMessagesStreamIterator,
)
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.google_genai.main import GenerateContentResponse
from litellm.types.llms.openai_evals import (
CancelEvalResponse,
@ -225,6 +224,7 @@ if TYPE_CHECKING:
Run,
RunDeleteResponse,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:

2
litellm/llms/gemini/google_genai/guardrail_translation/handler.py
View file

@ -25,7 +25,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_EMPTY_REQUEST_DATA: Final[Mapping[str, object]] = MappingProxyType({})

14
litellm/llms/litellm_proxy/skills/handler.py
View file

@ -11,19 +11,19 @@ from typing import Final
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.llms.litellm_proxy.skills.constants import (
LITELLM_SKILL_ID_PREFIX,
MAX_SKILLS_PER_SEARCH,
)
from litellm.models.skills import LiteLLM_SkillsTable
from litellm.proxy._types import NewSkillRequest, UserAPIKeyAuth
from litellm.proxy.common_utils.resource_ownership import (
from litellm.litellm_core_utils.resource_ownership import (
get_primary_resource_owner_scope,
get_resource_owner_scopes,
is_proxy_admin,
user_can_access_resource_owner,
)
from litellm.llms.litellm_proxy.skills.constants import (
LITELLM_SKILL_ID_PREFIX,
MAX_SKILLS_PER_SEARCH,
)
from litellm.models.skills import LiteLLM_SkillsTable, NewSkillRequest
from litellm.repositories.table_repositories import SkillsRepository
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
# Skills are looked up on every chat completion that has skills enabled
# (`SkillsInjectionHook` calls ``fetch_skill_from_db``). 60s LRU/TTL cache

2
litellm/llms/litellm_proxy/skills/skill_search.py
View file

@ -21,9 +21,9 @@ from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.models.skills import LiteLLM_SkillsTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
DEFAULT_SKILL_SEARCH_TOP_K: Final = 5
MAX_SKILL_SEARCH_TOP_K: Final = 100

4
litellm/llms/litellm_proxy/skills/transformation.py
View file

@ -22,7 +22,7 @@ from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.models.skills import LiteLLM_SkillsTable
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class LiteLLMSkillsTransformationHandler:
@ -137,7 +137,7 @@ class LiteLLMSkillsTransformationHandler:
"""Async implementation of create_skill."""
# Lazy import to avoid SDK dependency on proxy
from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler
from litellm.proxy._types import NewSkillRequest
from litellm.models.skills import NewSkillRequest
skill_request: Final = NewSkillRequest(
display_title=display_title,

2
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -69,7 +69,7 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class OpenAIChatCompletionsHandler(BaseTranslation):

2
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -102,8 +102,8 @@ if TYPE_CHECKING:
ModifyResponseException,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import ResponseInputParam
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class _ToolCallShape(NamedTuple):

2
litellm/llms/openai/videos/guardrail_translation/handler.py
View file

@ -6,7 +6,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class OpenAIVideoGenerationHandler(BaseTranslation):

2
litellm/llms/pass_through/guardrail_translation/handler.py
View file

@ -18,8 +18,8 @@ from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class PassThroughEndpointHandler(BaseTranslation):

9
litellm/llms/sagemaker/chat/transformation.py
View file

@ -69,17 +69,20 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM):
litellm_params: dict, # mutable-ok: matches the base chat transform signature
headers: dict, # mutable-ok: matches the base chat transform signature
) -> dict: # mutable-ok: the handler sends this body straight to httpx
request: Final = super().transform_request(
request: Final[dict[str, object]] = super().transform_request( # mutable-ok: base transform returns a dict
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
)
container_request: Final = {
key: value for key, value in request.items() if key not in self.aws_authentication_params
}
served_model_name: Final = litellm_params.get("hf_model_name")
if not isinstance(served_model_name, str):
return request
return {**request, "model": served_model_name}
return container_request
return {**container_request, "model": served_model_name}
def get_complete_url(
self,

3
litellm/llms/scaledown/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .chat.transformation import ScaleDownChatConfig, ScaleDownError
__all__ = ["ScaleDownChatConfig", "ScaleDownError"]

3
litellm/llms/scaledown/chat/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .transformation import ScaleDownChatConfig, ScaleDownError
__all__ = ["ScaleDownChatConfig", "ScaleDownError"]

648
litellm/llms/scaledown/chat/transformation.py Normal file
View file

@ -0,0 +1,648 @@
import json
import time
import uuid
from collections.abc import Mapping, Sequence
from functools import reduce
from typing import TYPE_CHECKING, Final, NoReturn
from urllib.parse import unquote
import httpx
from pydantic import JsonValue, TypeAdapter
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse, Usage
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
JSON_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]])
DEFAULT_API_BASE: Final = "https://api.scaledown.xyz"
DOMAIN_MODELS: Final = frozenset({"classify", "extract", "summarize", "compress"})
DECISIONS_MODELS: Final = frozenset({"decisions"})
DECISIONS_UPSTREAM_MODEL: Final = "classify-1"
NATIVE_PATHS: Final[Mapping[str, str]] = {
"classify": "/classify",
"extract": "/extract",
"summarize": "/summarization/abstractive",
"compress": "/compress/raw/",
}
DECISION_QUESTION_TYPES: Final = frozenset({"choice", "noul", "score"})
EXTRA_BODY_KEYS: Final[Mapping[str, frozenset[str]]] = {
"extract": frozenset({"threshold", "top_n"}),
"summarize": frozenset(),
"compress": frozenset({"compression_rate"}),
"classify": frozenset({"labels"}),
"decisions": frozenset({"questions"}),
}
# Set by the LiteLLM router and proxy on every call; handled by LiteLLM, never sent upstream.
ROUTER_PARAMS: Final = frozenset({"max_retries", "stream_options"})
SCORE_MIN_LEVELS: Final = 2
SCORE_MAX_LEVELS: Final = 10
MAX_ENTITIES: Final = 1000
MAX_REF_CHAIN: Final = 32
MAX_SCHEMA_DEPTH: Final = 32
class ScaleDownError(BaseLLMException):
pass
def _reject(message: str, status_code: int = 400) -> NoReturn:
"""The single place this adapter raises its own errors."""
raise ScaleDownError(status_code=status_code, message=message)
class ScaleDownChatConfig(BaseConfig):
@property
def custom_llm_provider(self) -> str | None:
return "scaledown"
def validate_environment(
self,
headers: dict[str, str], # mutable-ok: BaseConfig signature
model: str,
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature
resolved_key: Final = api_key or get_secret_str("SCALEDOWN_API_KEY")
if resolved_key is None:
_reject("Missing ScaleDown API key. Set SCALEDOWN_API_KEY or pass api_key to the call.", 401)
if not api_key and api_base is not None and _root(api_base) != _trusted_root():
_reject(
"A custom api_base needs its own api_key: the SCALEDOWN_API_KEY from the environment is only "
"sent to the default host or the host in SCALEDOWN_API_BASE."
)
return {**headers, "x-api-key": resolved_key, "content-type": "application/json"}
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
stream: bool | None = None,
) -> str:
base: Final = _root(api_base or _trusted_root())
operation: Final = _operation(model)
if operation in DECISIONS_MODELS:
return f"{base}/v1/scaledown"
return f"{base}{NATIVE_PATHS[operation]}"
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig signature
operation: Final = _operation(model)
if operation == "extract":
return ["response_format", "stream"]
if operation == "summarize":
return ["max_tokens", "max_completion_tokens", "stream"]
return ["stream"]
def map_openai_params(
self,
non_default_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
model: str,
drop_params: bool,
) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature
supported: Final = frozenset(self.get_supported_openai_params(model))
unsupported: Final = sorted(key for key in non_default_params if key not in supported | ROUTER_PARAMS)
if unsupported and not drop_params:
_reject(
f"ScaleDown model '{model}' does not support {unsupported}. "
f"Supported: {sorted(supported)}. Set litellm.drop_params=True to ignore."
)
mapped: Final = {
("max_tokens" if key == "max_completion_tokens" else key): value
for key, value in non_default_params.items()
if key in supported
}
return {**optional_params, **mapped}
def should_fake_stream(
self, model: str | None, stream: bool | None, custom_llm_provider: str | None = None
) -> bool:
return bool(stream)
def transform_extra_body(
self,
extra_body: Mapping[str, object],
request: Mapping[str, object],
model: str,
litellm_params: Mapping[str, object],
) -> Mapping[str, JsonValue]:
# extra_body is merged after guardrails ran on the messages, so it may only add options.
# Anything that carries prompt text or picks the model is refused.
options: Final = JSON_OBJECT.validate_python(extra_body)
operation: Final = _operation(model)
allowed: Final = EXTRA_BODY_KEYS[operation]
unexpected: Final = sorted(set(options) - allowed)
if unexpected:
allowance: Final = f"may only set {sorted(allowed)}" if allowed else "is not accepted"
_reject(
f"extra_body for scaledown/{operation} {allowance}, got {unexpected}. "
"Text and instructions must be passed as messages."
)
if operation == "compress":
return {"scaledown": {"rate": options["compression_rate"]}} if "compression_rate" in options else {}
if "questions" in options:
_validate_questions(options["questions"])
if "labels" in options:
_validate_labels(options["labels"])
return dict(options)
def sign_request(
self,
headers: dict[str, str], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature
api_base: str,
api_key: str | None = None,
model: str | None = None,
stream: bool | None = None,
fake_stream: bool | None = None,
) -> tuple[dict[str, str], bytes | None]: # mutable-ok: BaseConfig signature
# Runs after extra_body is merged, so this is where a decisions body is known to be complete.
if model is not None and _operation(model) in DECISIONS_MODELS:
_require_complete_decisions(request_data)
if model is not None and _operation(model) == "classify":
_validate_labels(request_data.get("labels"))
return headers, None
def transform_request(
self,
model: str,
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
headers: dict[str, str], # mutable-ok: BaseConfig signature
) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature
checked_messages: Final = JSON_MESSAGES.validate_python(messages)
operation: Final = _operation(model)
_validate_text_input(operation, checked_messages, optional_params)
if operation == "classify":
return {
"text": _last_user_text(operation, checked_messages),
**({"labels": optional_params["labels"]} if "labels" in optional_params else {}),
}
if operation in DECISIONS_MODELS:
return dict(_decisions_request(checked_messages, optional_params))
if operation == "extract":
return dict(_extract_request(checked_messages, optional_params))
if operation == "summarize":
return dict(_summarize_request(checked_messages, optional_params))
return dict(_compress_request(checked_messages, optional_params))
def transform_response(
self,
model: str,
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: "LiteLLMLoggingObj",
request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature
messages: list[AllMessageValues], # mutable-ok: BaseConfig signature
optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature
encoding: "Tokenizer | None",
api_key: str | None = None,
json_mode: bool | None = None,
) -> ModelResponse:
try:
raw: Final = JSON_OBJECT.validate_json(raw_response.content)
except ValueError:
_reject(
f"ScaleDown returned a non-JSON response: {raw_response.text[:500]}",
raw_response.status_code if raw_response.is_error else 502,
)
operation: Final = _operation(model)
if operation in DECISIONS_MODELS:
return _decisions_response(operation, raw, model_response)
return _native_response(operation, raw, model_response, request_data.get("entities"))
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseConfig signature
) -> BaseLLMException:
return ScaleDownError(status_code=status_code, message=error_message, headers=headers)
def _root(api_base: str) -> str:
return api_base.rstrip("/").removesuffix("/v1")
def _trusted_root() -> str:
return _root(get_secret_str("SCALEDOWN_API_BASE") or DEFAULT_API_BASE)
def _operation(model: str) -> str:
operation: Final = model.split("/", 1)[1] if "/" in model else model
if operation not in DOMAIN_MODELS | DECISIONS_MODELS:
_reject(f"Unknown ScaleDown model '{model}'. Choose from {sorted(DOMAIN_MODELS | DECISIONS_MODELS)}.")
return operation
def _text_of(message: Mapping[str, JsonValue]) -> str | None:
content: Final = message.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
parts: Final = [
str(part.get("text", "")) for part in content if isinstance(part, dict) and part.get("type") == "text"
]
return "\n".join(parts) if parts else None
return None
def _validate_text_input(
operation: str, messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> None:
if any(key in optional_params for key in ("state", "document", "document_mime_type", "text")):
_reject("ScaleDown accepts text through messages only. Native state and document inputs are not supported.")
for message in messages:
if operation in {"classify", "decisions", "extract"} and message.get("role") in {"system", "developer"}:
_reject(
f"ScaleDown {operation} does not support system messages. Put instructions in the field definitions."
)
if _has_non_text_content(message):
_reject("ScaleDown currently supports text only. Image and document inputs are not supported.")
def _has_non_text_content(message: Mapping[str, JsonValue]) -> bool:
content: Final = message.get("content")
return isinstance(content, list) and any(
not isinstance(part, dict) or part.get("type") != "text" for part in content
)
def _last_user_text(operation: str, messages: Sequence[Mapping[str, JsonValue]]) -> str:
message: Final = next((message for message in reversed(messages) if message.get("role") == "user"), None)
text: Final = _text_of(message) if message is not None else None
if not text:
_reject(f"ScaleDown {operation} requires text in the last user message.")
return text
def _validate_labels(labels: JsonValue) -> None:
if not isinstance(labels, list) or not labels:
_reject("ScaleDown classify requires a non-empty 'labels' list of objects with 'name' and 'rubric' strings.")
for label in labels:
if not isinstance(label, Mapping) or not all(
isinstance(label.get(key), str) and label[key] for key in ("name", "rubric")
):
_reject("Each classify label needs non-empty 'name' and 'rubric' strings.")
def _child(node: JsonValue, key: str) -> JsonValue:
return node.get(key.replace("~1", "/").replace("~0", "~")) if isinstance(node, Mapping) else None
def _extend_ref_path(path: tuple[str, ...], ref: str) -> tuple[str, ...]:
return (*path, ref)
def _nullable_branch(schema: Mapping[str, JsonValue]) -> Mapping[str, JsonValue] | None:
if "allOf" in schema:
_reject("ScaleDown extraction does not support allOf. Supply the object's properties directly.")
variants: Final = schema.get("anyOf", schema.get("oneOf"))
if variants is None:
return None
if any(key in schema for key in ("properties", "items", "$ref", "type")) or (
"anyOf" in schema and "oneOf" in schema
):
_reject("ScaleDown extraction cannot combine nullable branches with sibling schema definitions.")
if not isinstance(variants, list) or len(variants) != 2:
_reject("ScaleDown extraction supports anyOf/oneOf only for one schema plus null.")
branches: Final = tuple(branch for branch in variants if isinstance(branch, dict))
non_null: Final = tuple(branch for branch in branches if branch.get("type") != "null")
if len(branches) != 2 or len(non_null) != 1:
_reject("ScaleDown extraction supports anyOf/oneOf only for one schema plus null.")
return {
**non_null[0],
**{key: value for key, value in schema.items() if key not in {"anyOf", "oneOf"}},
}
def _deref(
prop: Mapping[str, JsonValue], root: Mapping[str, JsonValue], path: tuple[str, ...]
) -> tuple[Mapping[str, JsonValue], tuple[str, ...], bool]:
"""Follow nullable branches and local `$ref`s (`#/$defs/X`, `#/definitions/X`).
Returns the schema, the refs followed so far, and whether a ref repeated (a cycle).
The walk is a bounded loop, so a long chain of distinct refs cannot exhaust the stack.
"""
current: Mapping[str, JsonValue] = prop # rebind-ok: bounded walk along a ref chain
followed: tuple[str, ...] = path # rebind-ok: bounded walk along a ref chain
for _ in range(MAX_REF_CHAIN):
nullable: Mapping[str, JsonValue] | None = _nullable_branch(current)
ref: JsonValue = current.get("$ref")
if nullable is not None:
current = nullable
continue
if "$ref" not in current:
return current, followed, False
if not isinstance(ref, str) or not ref.startswith("#"):
_reject("ScaleDown extraction supports only local JSON Pointer $refs.")
pointer: str = unquote(ref[1:])
if not pointer.startswith("/"):
_reject("ScaleDown extraction supports only local JSON Pointer $refs beginning with #/.")
if any(key in current for key in ("properties", "items", "type")):
_reject("ScaleDown extraction cannot combine $ref with sibling schema definitions.")
if ref in followed:
return current, followed, True
target: JsonValue = reduce(_child, pointer[1:].split("/"), dict(root))
if not isinstance(target, Mapping):
_reject(f"ScaleDown extraction cannot resolve $ref '{ref}' to a schema object.")
current, followed = target, _extend_ref_path(followed, ref)
_reject(f"The response_format schema follows more than {MAX_REF_CHAIN} chained $refs or nullable branches.")
def _build_properties(
properties: Mapping[str, JsonValue],
root: Mapping[str, JsonValue],
path: tuple[str, ...],
budget: int,
depth: int = 0,
) -> tuple[dict[str, JsonValue], int]: # mutable-ok: serialized JSON entity map
"""Turn JSON-schema properties into the /extract entity map, spending one budget unit per entity.
Each property becomes an entity whose value is its description (or its name when there is
none). Nested objects stay nested, arrays of objects become a one-element list holding the
item's entity map, local `$ref`s are followed, and a ref that points back into itself is cut
off at that point instead of expanded.
"""
if depth > MAX_SCHEMA_DEPTH:
_reject(f"The response_format schema nests deeper than {MAX_SCHEMA_DEPTH} levels.")
built: Final[dict[str, JsonValue]] = {} # mutable-ok: filled once while threading the budget through
remaining = budget # rebind-ok: the budget is threaded through the loop
for name, prop in properties.items():
if remaining <= 0:
_reject(
f"The response_format schema expands to more than {MAX_ENTITIES} entities; "
"remove recursive or heavily repeated $refs."
)
built[name], remaining = _build_entity(name, prop, root, path, remaining - 1, depth)
return built, remaining
def _build_entity(
name: str, prop: JsonValue, root: Mapping[str, JsonValue], path: tuple[str, ...], budget: int, depth: int
) -> tuple[JsonValue, int]:
schema, followed, cyclic = _deref(prop if isinstance(prop, Mapping) else {}, root, path)
leaf: Final = schema.get("description") or name
nested: Final = schema.get("properties")
if cyclic:
return leaf, budget
if isinstance(nested, Mapping):
return _build_properties(nested, root, followed, budget, depth + 1)
raw_items: Final = schema.get("items")
schema_type: Final = schema.get("type")
is_array: Final = schema_type == "array" or (isinstance(schema_type, list) and "array" in schema_type)
if is_array and isinstance(raw_items, Mapping):
items, item_path, item_cyclic = _deref(raw_items, root, followed)
item_properties: Final = items.get("properties")
if not item_cyclic and isinstance(item_properties, Mapping):
built, remaining = _build_properties(item_properties, root, item_path, budget, depth + 1)
return [built], remaining
return leaf, budget
def _extract_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
response_format: Final = optional_params.get("response_format")
json_schema: Final = response_format.get("json_schema") if isinstance(response_format, Mapping) else None
if isinstance(json_schema, Mapping) and json_schema.get("strict") is True:
_reject(
"ScaleDown extract does not support strict JSON schemas. Use strict=false and validate the returned fields."
)
schema: Final = json_schema.get("schema") if isinstance(json_schema, Mapping) else None
properties: Final = _deref(schema, schema, ())[0].get("properties") if isinstance(schema, Mapping) else None
if not isinstance(schema, dict) or not isinstance(properties, dict) or not properties:
_reject(
"ScaleDown extract needs response_format={'type': 'json_schema', ...} with a non-empty "
"'properties' map; each property becomes an entity and its description the extraction hint."
)
return {
"text": _last_user_text("extract", messages),
"entities": _build_properties(properties, schema, (), MAX_ENTITIES)[0],
**{key: optional_params[key] for key in ("threshold", "top_n") if key in optional_params},
}
def _summarize_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
instructions: Final = [
text for message in messages if message.get("role") in {"system", "developer"} and (text := _text_of(message))
]
return {
"text": _last_user_text("summarize", messages),
**({"instructions": "\n".join(instructions)} if instructions else {}),
**({"max_tokens": optional_params["max_tokens"]} if "max_tokens" in optional_params else {}),
}
def _compress_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
prompt: Final = _last_user_text("compress", messages)
last_user_index: Final = max(i for i, m in enumerate(messages) if m.get("role") == "user" and _text_of(m))
context: Final = "\n\n".join(text for i, m in enumerate(messages) if i != last_user_index and (text := _text_of(m)))
return {
"context": context,
"prompt": prompt,
"scaledown": {"rate": optional_params.get("compression_rate", "auto")},
}
def _decisions_request(
messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue]
) -> Mapping[str, JsonValue]:
questions: Final = optional_params.get("questions")
if questions is not None:
_validate_questions(questions)
return {
"model": DECISIONS_UPSTREAM_MODEL,
"state": {"text": _last_user_text("decisions", messages)},
**({"questions": questions} if questions is not None else {}),
}
def _require_complete_decisions(body: Mapping[str, JsonValue]) -> None:
if not body.get("questions"):
_reject(
"ScaleDown decisions requires a non-empty 'questions' map, passed as an extra parameter, "
'e.g. {"questions": {"category": {"type": "choice", "criteria": {"billing": "..."}}}}. '
"Question types are 'choice', 'noul', and 'score'."
)
def _validate_questions(questions: JsonValue) -> None:
if not isinstance(questions, Mapping):
_reject(f"'questions' must be a map of name to question, got {type(questions).__name__}.")
for name, question in questions.items():
if not isinstance(question, Mapping):
_reject(f"Question '{name}' must be an object.")
question_type = question.get("type")
if not isinstance(question_type, str) or question_type not in DECISION_QUESTION_TYPES:
_reject(f"Question '{name}' has type {question_type!r}; expected one of {sorted(DECISION_QUESTION_TYPES)}.")
criteria = question.get("criteria")
if question_type == "choice" and (not isinstance(criteria, Mapping) or not criteria):
_reject(f"Choice question '{name}' needs a non-empty 'criteria' map of option key to description.")
if question_type == "score" and (
not isinstance(criteria, list) or not SCORE_MIN_LEVELS <= len(criteria) <= SCORE_MAX_LEVELS
):
_reject(
f"Score question '{name}' needs 'criteria' as an ordered list of "
f"{SCORE_MIN_LEVELS} to {SCORE_MAX_LEVELS} level descriptions, lowest to highest."
)
def _unwrap(value: JsonValue) -> JsonValue:
wrapper_keys: Final = {"_value", "_span_anchor"}
if isinstance(value, Mapping) and "_value" in value and set(value) <= wrapper_keys:
return value["_value"]
return value
def _clean_extraction(value: JsonValue, requested: JsonValue) -> JsonValue:
"""Reduce an extraction result to the fields the caller's schema asked for.
The request's entity map says which keys are real fields, so provider-added keys such as
`<field>_span_anchor` and `{"_value": ..., "_span_anchor": ...}` wrappers are dropped without
touching a requested field that happens to share one of those names. The untouched payload
stays on `_hidden_params["scaledown_response"]`.
"""
if isinstance(requested, Mapping) and isinstance(value, Mapping):
return {key: _clean_extraction(value[key], requested[key]) for key in requested if key in value}
if isinstance(requested, list) and requested and isinstance(value, list):
return [_clean_extraction(item, requested[0]) for item in value]
return _unwrap(value)
def _extracted_fields(raw: Mapping[str, JsonValue], requested: JsonValue) -> JsonValue:
entities: Final = raw.get("entities")
matches: Final = (
tuple(entity for entity in entities if isinstance(entity, Mapping)) if isinstance(entities, list) else ()
)
scalar_names: Final = (
tuple(name for name, hint in requested.items() if isinstance(hint, str))
if isinstance(requested, Mapping)
else ()
)
scalars: Final = {
name: next(entity["text"] for entity in matches if entity.get("type") == name and "text" in entity)
for name in scalar_names
if any(entity.get("type") == name and "text" in entity for entity in matches)
}
structured: Final = raw.get("structured_result")
return _clean_extraction({**scalars, **(structured if isinstance(structured, Mapping) else {})}, requested)
def _native_response(
operation: str, raw: JsonValue, model_response: ModelResponse, requested: JsonValue = None
) -> ModelResponse:
"""Wrap a native /extract, /summarization/abstractive or /compress/raw/ payload.
Summarize and compress return the upstream payload as JSON on choices[0].message.content.
Extract returns the fields requested by response_format without enforcing schema types.
The native APIs report input tokens only; they return no output token
count, so completion_tokens stays 0 because it is unmeasured, not because nothing was generated.
"""
if not isinstance(raw, Mapping) or not raw:
_reject(
f"ScaleDown '{operation}' returned an empty or malformed response: {json.dumps(raw)[:500]}",
500,
)
results: Final = raw.get("results")
input_tokens: Final = raw.get("input_tokens") or (
(results.get("original_prompt_tokens") if isinstance(results, Mapping) else None)
or raw.get("original_prompt_tokens")
if operation == "compress"
else None
)
content: Final = _extracted_fields(raw, requested) if operation == "extract" else raw
model_response.id = f"chatcmpl-{uuid.uuid4().hex}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.created = int(time.time()) # rebind-ok: LiteLLM fills the response object it passes in
model_response.model = f"scaledown/{operation}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.object = "chat.completion" # rebind-ok: LiteLLM fills the response object it passes in
model_response.choices = [ # rebind-ok: LiteLLM fills the response object it passes in
Choices(
index=0,
message=Message(role="assistant", content=json.dumps(content, separators=(",", ":"))),
finish_reason="stop",
)
]
_set_usage(model_response, {"prompt_tokens": input_tokens or 0})
model_response.hidden_params["scaledown_response"] = raw # pyright: ignore[reportUnknownMemberType] # Inherited hidden_params has untyped values. # rebind-ok: supplied response
return model_response
def _decisions_response(operation: str, raw: JsonValue, model_response: ModelResponse) -> ModelResponse:
answers: Final = raw.get("answers") if isinstance(raw, Mapping) else None
if answers is None:
_reject(f"ScaleDown decisions response contained no answers: {json.dumps(raw)[:500]}", 500)
model_response.id = f"chatcmpl-{uuid.uuid4().hex}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.created = int(time.time()) # rebind-ok: LiteLLM fills the response object it passes in
model_response.model = f"scaledown/{operation}" # rebind-ok: LiteLLM fills the response object it passes in
model_response.object = "chat.completion" # rebind-ok: LiteLLM fills the response object it passes in
model_response.choices = [ # rebind-ok: LiteLLM fills the response object it passes in
Choices(
index=0,
message=Message(role="assistant", content=json.dumps(answers, separators=(",", ":"))),
finish_reason="stop",
)
]
usage: Final = raw.get("usage") if isinstance(raw, dict) else None
counts: Final = usage if isinstance(usage, dict) else {}
_set_usage(
model_response,
{
"prompt_tokens": counts.get("input_tokens", 0),
"completion_tokens": counts.get("output_tokens", 0),
},
)
model_response.hidden_params["scaledown_response"] = raw # pyright: ignore[reportUnknownMemberType] # Inherited hidden_params has untyped values. # rebind-ok: supplied response
return model_response
def _set_usage(model_response: ModelResponse, usage: Mapping[str, JsonValue]) -> None:
prompt_tokens: Final = _token_count(usage.get("prompt_tokens"))
completion_tokens: Final = _token_count(usage.get("completion_tokens"))
total_tokens: Final = usage.get("total_tokens")
model_response.usage = Usage( # pyright: ignore[reportAttributeAccessIssue] # ModelResponse stores usage as a dynamic Pydantic field. # rebind-ok: fills the supplied response
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=_token_count(total_tokens) if total_tokens is not None else prompt_tokens + completion_tokens,
)
def _token_count(value: JsonValue) -> int:
if value is None:
return 0
if isinstance(value, int) and value >= 0:
return value
_reject("ScaleDown returned an invalid token count.", 500)

15
litellm/llms/vertex_ai/common_utils.py
View file

@ -1,6 +1,7 @@
import re
from collections.abc import Mapping
from copy import deepcopy
from datetime import datetime, timezone
from enum import Enum
from functools import lru_cache
from types import MappingProxyType
@ -1000,18 +1001,18 @@ def strip_field(schema, field_name: str):
strip_field(items, field_name)
_VERTEX_DATETIME_FRACTION: Final = re.compile(r"\.\d+(?=Z$)", re.IGNORECASE)
def convert_vertex_datetime_to_openai_datetime(vertex_datetime: str) -> int:
"""
Converts a Vertex AI datetime string to an OpenAI datetime integer
Converts a Vertex AI / GCS RFC 3339 UTC timestamp to Unix epoch seconds
vertex_datetime: str = "2024-12-04T21:53:12.120184Z"
returns: int = 1722729192
returns: int = 1733349192
"""
from datetime import datetime
# Parse the ISO format string to datetime object
dt: Final = datetime.strptime(vertex_datetime, "%Y-%m-%dT%H:%M:%S.%fZ")
# Convert to Unix timestamp (seconds since epoch)
whole_seconds: Final = _VERTEX_DATETIME_FRACTION.sub("", vertex_datetime)
dt: Final = datetime.strptime(whole_seconds, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc)
return int(dt.timestamp())

35
litellm/main.py
View file

@ -1938,6 +1938,39 @@ def _complete_ragflow(ctx: CompletionDispatchContext) -> _CompletionDispatchResu
return response
def _complete_scaledown(ctx: CompletionDispatchContext) -> _CompletionDispatchResult:
try:
return base_llm_http_handler.completion(
model=ctx.model,
messages=ctx.messages,
headers=ctx.headers,
model_response=ctx.model_response,
api_key=ctx.api_key,
api_base=ctx.api_base,
acompletion=ctx.acompletion,
logging_obj=ctx.logging,
optional_params=ctx.optional_params,
litellm_params=ctx.litellm_params,
shared_session=ctx.shared_session,
timeout=cast( # cast-ok: CompletionTimeout.resolve normalizes this value
float | httpx.Timeout, ctx.timeout
),
client=_dispatch_client_http(ctx),
custom_llm_provider=ctx.custom_llm_provider,
encoding=_get_encoding(),
stream=ctx.stream,
provider_config=ctx.provider_config,
)
except Exception as exc:
ctx.logging.post_call(
input=ctx.messages,
api_key=ctx.api_key,
original_response=str(exc),
additional_args={"headers": ctx.headers},
)
raise
def _complete_xai(ctx: CompletionDispatchContext) -> _CompletionDispatchResult:
acompletion: Final = ctx.acompletion
api_base: Final = ctx.api_base
@ -5859,6 +5892,8 @@ def completion(
elif custom_llm_provider == "ragflow":
## COMPLETION CALL - RAGFlow uses HTTP handler to support custom URL paths
response = _complete_ragflow(_dispatch_ctx)
elif custom_llm_provider == "scaledown":
response = _complete_scaledown(_dispatch_ctx) # rebind-ok: shared provider return
elif custom_llm_provider == "xai":
## COMPLETION CALL
response = _complete_xai(_dispatch_ctx)

48
litellm/model_prices_and_context_window_backup.json
View file

@ -66876,7 +66876,7 @@
"supports_audio_input": true
},
"azure/gpt-realtime-translate": {
"deprecation_date": "2027-05-06",
"deprecation_date": "2027-11-09",
"input_cost_per_second": 0.000566666666667,
"litellm_provider": "azure",
"max_input_tokens": 32000,
@ -71842,5 +71842,51 @@
"image"
],
"supports_embedding_image_input": true
},
"scaledown/extract": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/summarize": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": true,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/compress": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": true,
"supports_function_calling": false,
"supports_vision": false
},
"scaledown/classify": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
},
"scaledown/decisions": {
"input_cost_per_token": 5e-08,
"output_cost_per_token": 0.0,
"litellm_provider": "scaledown",
"mode": "chat",
"supports_system_messages": false,
"supports_vision": false,
"supports_function_calling": false
}
}

15
litellm/models/skills.py
View file

@ -28,3 +28,18 @@ class LiteLLM_SkillsTable(LiteLLMPydanticObjectBase):
created_by: str | None = None
updated_at: datetime | None = None
updated_by: str | None = None
class NewSkillRequest(LiteLLMPydanticObjectBase):
"""Request to create a new skill in LiteLLM database"""
display_title: str | None = None
description: str | None = None
instructions: str | None = None
file_content: bytes | None = None # Binary content of skill files (zip)
file_name: str | None = None # Original filename
file_type: str | None = None # MIME type (e.g., "application/zip")
metadata: dict[str, Any] | None = None
authorization_url: str | None = None
token_url: str | None = None
registration_url: str | None = None

18
litellm/provider_endpoints_support_backup.json
View file

@ -2064,6 +2064,24 @@
"search": true
}
},
"scaledown": {
"display_name": "ScaleDown (`scaledown`)",
"url": "https://docs.litellm.ai/docs/providers/scaledown",
"endpoints": {
"chat_completions": true,
"messages": false,
"responses": false,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_speech": false,
"moderations": false,
"batches": false,
"rerank": false,
"a2a": false,
"interactions": false
}
},
"sambanova": {
"display_name": "Sambanova (`sambanova`)",
"url": "https://docs.litellm.ai/docs/providers/sambanova",

4
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -95,9 +95,7 @@ def _normalize_caller_admission_credential(value: str) -> str:
def _admission_request(scope: Scope) -> Request:
"""Request whose ``body()`` serves the routing layer's lazily peeked JSON-RPC
bytes (``b"{}"`` when no peek callable was stashed) instead of the ASGI
receive channel."""
"""Request whose ``body()`` serves the peeked JSON-RPC method (``b"{}"`` when nothing is stashed)."""
request: Final = Request(scope=scope)
peeked_body: Final[Callable[[], Awaitable[bytes]] | None] = scope.get(MCP_PEEKED_BODY_SCOPE_KEY)

4
litellm/proxy/_experimental/mcp_server/legacy_callbacks.py
View file

@ -27,6 +27,10 @@ class ElicitationCallback(Protocol):
async def __call__(self, context: object, params: ElicitRequestParams) -> ElicitResult | ErrorData: ...
def has_active_mcp_request_context() -> bool:
return get_active_mcp_request_ctx() is not None
def create_sampling_callback(
user_api_key_auth: UserAPIKeyAuth | None = None,
raw_headers: Mapping[str, str] | None = None,

202
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -24,7 +24,7 @@ from collections.abc import (
MutableMapping,
Sequence,
)
from contextlib import ExitStack, asynccontextmanager
from contextlib import AbstractAsyncContextManager, ExitStack, asynccontextmanager, nullcontext
from dataclasses import dataclass, replace
from functools import lru_cache
from itertools import chain, groupby
@ -69,9 +69,16 @@ from litellm.constants import (
MCP_HEALTH_CHECK_TIMEOUT,
MCP_METADATA_TIMEOUT,
MCP_TOOL_LISTING_TIMEOUT,
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS,
)
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
from litellm.experimental_mcp_client.client import (
MCPClient,
PersistentMCPSession,
UpstreamSessionClosedError,
strip_auth_scheme,
to_basic_credentials,
)
from litellm.integrations.custom_guardrail import (
_sync_guardrail_info_to_logging_obj, # pyright: ignore[reportPrivateUsage] - the same bridge @log_guardrail_information uses; reimplementing it here would fork the metadata-key logic
)
@ -98,6 +105,7 @@ from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
upstream_auth_challenge,
)
from litellm.proxy._experimental.mcp_server.interactions import bind_target
from litellm.proxy._experimental.mcp_server.legacy_callbacks import has_active_mcp_request_context
from litellm.proxy._experimental.mcp_server.mcp_debug import describe_upstream_http_failure, record_auth_resolution
from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
MCPPerUserTokenCache,
@ -343,6 +351,12 @@ _ToolParamMap: TypeAlias = dict[str, list[str]]
_EnvVarList: TypeAlias = list[dict[str, object]]
_InMemoryCacheDict: TypeAlias = dict[str, object]
_ToolArguments: TypeAlias = dict[str, object]
_UpstreamSessionKey: TypeAlias = tuple[
str,
str,
str,
tuple[MCPTransport, MCPUpstreamProtocol, float, bool, bool, bool],
]
@dataclass(frozen=True, slots=True)
@ -1927,6 +1941,9 @@ class MCPServerManager:
self.catalog = CatalogSnapshots(self)
self.registry: dict[str, MCPServer] = {}
self._upstream_sessions: dict[_UpstreamSessionKey, PersistentMCPSession] = {} # mutable-ok: session registry
self._retiring_upstream_sessions: list[tuple[str, PersistentMCPSession]] = [] # mutable-ok: session registry
self._live_gateway_sessions: frozenset[str] = frozenset()
self._openapi_health_probes: Callable[[str], _OpenAPIHealthProbe] = lru_cache(maxsize=128)(_OpenAPIHealthProbe)
self.config_mcp_servers: dict[str, MCPServer] = {}
"""
@ -6123,6 +6140,108 @@ class MCPServerManager:
async with semaphore:
yield
def _outbound_limit(
self,
mcp_server: MCPServer,
persistent_session: PersistentMCPSession | None,
) -> AbstractAsyncContextManager[None]:
if persistent_session is not None:
return nullcontext()
return self._limit_outbound_concurrency(mcp_server)
async def _upstream_session_for(
self,
client: MCPClient,
mcp_server: MCPServer,
raw_headers: Mapping[str, str] | None,
) -> PersistentMCPSession | None:
gateway_session_id: Final = next(
(
value
for key, value in (raw_headers.items() if raw_headers else ())
if key.lower() == "mcp-session-id" and value
),
None,
)
if (
not has_active_mcp_request_context()
or gateway_session_id is None
or gateway_session_id not in self._live_gateway_sessions
or mcp_server.transport == MCPTransport.stdio
):
return None
key: Final[_UpstreamSessionKey] = (
gateway_session_id,
mcp_server.server_id,
await client.discovery_auth_fingerprint(),
client.session_settings(),
)
if gateway_session_id not in self._live_gateway_sessions:
return None
existing: Final = self._upstream_sessions.get(key)
if existing is not None and existing.reusable:
return existing
opened: Final = client.open_persistent_session(
lambda: self._limit_outbound_concurrency(self.get_mcp_server_by_id(mcp_server.server_id) or mcp_server)
)
self._retiring_upstream_sessions = [
(gateway_session_id, session)
for gateway_session_id, session in self._retiring_upstream_sessions
if not session.closed
]
previous_keys: Final = tuple(
registered_key for registered_key in self._upstream_sessions if registered_key[:2] == key[:2]
)
previous_sessions: Final = tuple(
(key[0], self._upstream_sessions.pop(previous_key)) for previous_key in previous_keys
)
for retired_gateway_session_id, previous_session in previous_sessions:
previous_session.retire()
self._retiring_upstream_sessions.append((retired_gateway_session_id, previous_session))
self._upstream_sessions[key] = opened
return opened
def track_gateway_session(self, gateway_session_id: str) -> None:
self._live_gateway_sessions = self._live_gateway_sessions | frozenset((gateway_session_id,))
def release_upstream_sessions(self, gateway_session_id: str) -> None:
self._live_gateway_sessions = self._live_gateway_sessions - frozenset((gateway_session_id,))
for key in tuple(key for key in self._upstream_sessions if key[0] == gateway_session_id):
self._upstream_sessions.pop(key).close()
for retiring_gateway_session_id, session in self._retiring_upstream_sessions:
if retiring_gateway_session_id == gateway_session_id:
session.close()
self._retiring_upstream_sessions = [
(retiring_gateway_session_id, session)
for retiring_gateway_session_id, session in self._retiring_upstream_sessions
if retiring_gateway_session_id != gateway_session_id and not session.closed
]
async def close_all_upstream_sessions(self) -> None:
sessions: Final = (
*self._upstream_sessions.values(),
*(session for _, session in self._retiring_upstream_sessions),
)
self._upstream_sessions.clear()
self._retiring_upstream_sessions.clear()
self._live_gateway_sessions = frozenset()
for session in sessions:
session.close()
try:
await asyncio.wait_for(
asyncio.gather(*(session.wait_closed() for session in sessions)),
MCP_UPSTREAM_SESSION_SHUTDOWN_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
verbose_logger.warning("Timed out waiting for persistent MCP sessions to close during shutdown")
def _drop_upstream_session(self, session: PersistentMCPSession | None) -> None:
matching_keys: Final = tuple(key for key, value in self._upstream_sessions.items() if value is session)
dropped_sessions: Final = tuple((key[0], self._upstream_sessions.pop(key)) for key in matching_keys)
for gateway_session_id, dropped in dropped_sessions:
dropped.close()
self._retiring_upstream_sessions.append((gateway_session_id, dropped))
async def _obo_call_tool_with_retry(
self,
*,
@ -6138,27 +6257,33 @@ class MCPServerManager:
raw_headers: Mapping[str, str] | None = None,
client_ip: str | None = None,
allow_input_required: bool = False,
persistent_session: PersistentMCPSession | None = None,
) -> CallToolResult | InputRequiredResult:
"""Call a token_exchange (OBO) tool; on an upstream 401/403 re-mint the token once and retry.
The exchanged token is baked into the client at build time, so the retry invalidates the
cached exchange and rebuilds the client (which re-exchanges). One retry only: a non-auth
failure or a second auth failure degrades to the normal ``isError`` result, and a re-exchange
that now fails surfaces its own 401 challenge from ``_create_mcp_client``.
that now fails surfaces its own 401 challenge from ``create_mcp_client``.
"""
try:
return await client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
)
async with self._outbound_limit(mcp_server, persistent_session):
return await client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
except Exception as exc:
if _extract_upstream_auth_failure(exc) is None:
auth_failure: Final = _extract_upstream_auth_failure(exc)
if auth_failure is None and not isinstance(exc, UpstreamSessionClosedError):
return MCPClient.error_tool_result(exc)
spec: Final = to_server_spec(mcp_server)
if spec is not None:
await self._cred_provider.invalidate_credentials(to_subject(user_api_key_auth, subject_token), spec)
if auth_failure is not None:
spec: Final = to_server_spec(mcp_server)
if spec is not None:
await self._cred_provider.invalidate_credentials(to_subject(user_api_key_auth, subject_token), spec)
self._drop_upstream_session(persistent_session)
retry_client: Final = await self.create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
@ -6169,11 +6294,14 @@ class MCPServerManager:
raw_headers=raw_headers,
client_ip=client_ip,
)
return await retry_client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
)
retry_persistent_session: Final = await self._upstream_session_for(retry_client, mcp_server, raw_headers)
async with self._outbound_limit(mcp_server, retry_persistent_session):
return await retry_client.call_tool(
call_tool_params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
persistent_session=retry_persistent_session,
)
async def _call_regular_mcp_tool(
self,
@ -6330,6 +6458,7 @@ class MCPServerManager:
raw_headers=raw_headers,
client_ip=client_ip,
)
persistent_session: Final = await self._upstream_session_for(client, mcp_server, raw_headers)
call_tool_params: Final = MCPCallToolRequestParams(
name=original_tool_name,
@ -6342,24 +6471,21 @@ class MCPServerManager:
# OBO / ID-JAG: the exchanged token may have been revoked/rotated upstream since it was
# cached, so an upstream 401 gets one invalidate + re-mint + retry. Gated to these modes;
# all others keep the plain single call below.
async def _obo_call_tool_limited():
async with self._limit_outbound_concurrency(mcp_server):
return await self._obo_call_tool_with_retry(
client=client,
call_tool_params=call_tool_params,
host_progress_callback=host_progress_callback,
mcp_server=mcp_server,
server_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
allow_input_required=allow_input_required,
)
tool_call_coro = _obo_call_tool_limited()
tool_call_coro = self._obo_call_tool_with_retry(
client=client,
call_tool_params=call_tool_params,
host_progress_callback=host_progress_callback,
mcp_server=mcp_server,
server_auth_header=server_auth_header,
extra_headers=extra_headers,
stdio_env=stdio_env,
subject_token=subject_token,
user_api_key_auth=user_api_key_auth,
raw_headers=raw_headers,
client_ip=client_ip,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
else:
# Scoped to the two client-forwarded token modes this stack introduced; legacy
# oauth2 + delegate_auth_to_upstream (is_oauth_passthrough) is being removed, so it is not
@ -6368,12 +6494,13 @@ class MCPServerManager:
server_label: Final = mcp_server.name or mcp_server.server_name or mcp_server.alias or ""
async def _call_tool_via_client(client, params):
async with self._limit_outbound_concurrency(mcp_server):
async with self._outbound_limit(mcp_server, persistent_session):
if not relays_upstream_auth:
return await client.call_tool(
params,
host_progress_callback=host_progress_callback,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
# The client-forwarded modes carry the caller's own upstream token, so an upstream
# 401 (expired/invalid token) is the caller's to resolve: relay it as
@ -6390,6 +6517,7 @@ class MCPServerManager:
host_progress_callback=host_progress_callback,
raise_on_error=True,
allow_input_required=allow_input_required,
persistent_session=persistent_session,
)
except Exception as e:
auth_info: Final = _extract_upstream_auth_failure(e)

24
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -658,6 +658,7 @@ if MCP_AVAILABLE:
_stateful_session_locks.pop(session_id, None)
_stateful_session_active_request_counts.pop(session_id, None)
_stateful_session_client_info.pop(session_id, None)
operations.global_mcp_server_manager.release_upstream_sessions(session_id)
# Keep this alias so existing references to session_manager still work
session_manager: Final = session_manager_stateless
@ -793,6 +794,7 @@ if MCP_AVAILABLE:
_stateful_auth_context_cleanup_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await _stateful_auth_context_cleanup_task
await operations.global_mcp_server_manager.close_all_upstream_sessions()
if _session_manager_stateful_cm:
await _session_manager_stateful_cm.__aexit__(None, None, None)
if _session_manager_cm:
@ -1264,13 +1266,16 @@ if MCP_AVAILABLE:
return f"ip:{hashlib.sha256(client_ip.encode('utf-8')).hexdigest()}"
return "anonymous"
def _peeked_json_object_body(body: bytes) -> bytes | None:
"""The peeked bytes when they form a complete JSON-RPC object, else None so a
truncated prefix or a batch array fails closed and stays budget-enforced."""
def _admission_auth_body(body: bytes) -> bytes:
"""Return only the JSON-RPC method for admission auth."""
try:
return body if isinstance(json.loads(body), dict) else None
except (json.JSONDecodeError, TypeError):
return None
data: Final = TypeAdapter(dict[str, object]).validate_json(body)
except (ValueError, TypeError):
return b"{}"
method: Final = data.get("method")
if not isinstance(method, str):
return b"{}"
return json.dumps({"method": method}).encode()
def _is_initialize_request(body: bytes) -> bool:
"""
@ -2193,10 +2198,10 @@ if MCP_AVAILABLE:
if peek is not None:
receive = peek.receive # rebind-ok: replay peeked ASGI messages to the downstream handler
async def peeked_json_object_body() -> bytes:
return _peeked_json_object_body(await peek.admission_body()) or b"{}"
async def admission_auth_body() -> bytes:
return _admission_auth_body(await peek.admission_body())
scope[MCP_PEEKED_BODY_SCOPE_KEY] = peeked_json_object_body
scope[MCP_PEEKED_BODY_SCOPE_KEY] = admission_auth_body
(
user_api_key_auth,
mcp_auth_header,
@ -2856,6 +2861,7 @@ if MCP_AVAILABLE:
_stateful_session_auth_contexts[session_id] = auth_user
_stateful_session_auth_context_last_seen[session_id] = time.monotonic()
_stateful_session_owners[session_id] = owner_fingerprint
operations.global_mcp_server_manager.track_gateway_session(session_id)
if client_info is not None:
_stateful_session_client_info[session_id] = client_info
break

1924
litellm/proxy/_lazy_openapi_snapshot.json
View file

File diff suppressed because it is too large Load diff

671
litellm/proxy/_types.py
View file

@ -3,7 +3,7 @@ import json
import os
from collections.abc import Callable, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from types import MappingProxyType as MappingProxyType # noqa: PLC0414 # public re-export
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, TypeAlias
import httpx
@ -16,7 +16,6 @@ from pydantic import (
Json,
JsonValue,
PositiveInt,
PrivateAttr,
TypeAdapter,
field_validator,
model_validator,
@ -26,19 +25,24 @@ from typing_extensions import (
ReadOnly,
Required,
TypedDict,
TypeIs, # noqa: TID251 # narrows the before-validator input without copying it
)
from litellm._uuid import uuid
from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_arize_otlp_protocol_value,
validate_capture_message_content_value,
validate_langfuse_environment_value,
validate_langfuse_span_scope_value,
validate_no_callback_env_reference,
)
from litellm.proxy._experimental.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled
from litellm.types.agents import AgentCaller, AgentResponse
from litellm.types.agents import (
AgentCaller as AgentCaller, # noqa: PLC0414 # public re-export
)
from litellm.types.agents import (
AgentResponse as AgentResponse, # noqa: PLC0414 # public re-export
)
from litellm.types.integrations.compression_interception import (
CompressionSavingsMetadata,
)
@ -67,11 +71,17 @@ from litellm.types.proxy.agent_identity import ManagedAgentContext
from litellm.types.proxy.auth.special_headers import (
SpecialHeaders as SpecialHeaders, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_api_key_auth import is_jwt
from litellm.types.proxy.auth.user_api_key_auth import (
is_jwt as is_jwt, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.carried_budget_state import (
OrgBudgetSnapshot,
TeamBudgetSnapshot,
UserBudgetSnapshot,
OrgBudgetSnapshot as OrgBudgetSnapshot, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.carried_budget_state import (
TeamBudgetSnapshot as TeamBudgetSnapshot, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.carried_budget_state import (
UserBudgetSnapshot as UserBudgetSnapshot, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings
@ -79,6 +89,7 @@ from litellm.types.router import AllowedModelRegion, RouterErrors, UpdateRouterC
from litellm.types.router_weights import validate_router_settings_dict
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (
CAPTURE_MESSAGE_CONTENT_VAR,
AzureSpillover,
CallTypes,
CostBreakdown,
@ -196,17 +207,9 @@ class LitellmTableNames(str, enum.Enum):
from litellm.types.integrations.slack_alerting import (
Litellm_EntityType as Litellm_EntityType, # noqa: E402, PLC0414 # public re-export
)
def hash_token(token: str) -> str:
import hashlib
# Hash the string using SHA-256
hashed_token: Final = hashlib.sha256(token.encode()).hexdigest()
return hashed_token
from litellm.types.proxy.auth.user_api_key_auth import ( # noqa: E402 # re-export after the definitions above
hash_token as hash_token, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_roles import ( # noqa: E402 # re-export after the definitions above
KeyManagementRoutes as KeyManagementRoutes, # noqa: PLC0414 # public re-export
)
@ -1159,136 +1162,28 @@ class ModelParams(LiteLLMPydanticObjectBase):
return values
class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase):
mcp_servers: list[str] | None = None
mcp_access_groups: list[str] | None = None
mcp_tool_permissions: dict[str, list[str]] | None = None
mcp_toolsets: list[str] | None = None
blocked_tools: list[str] | None = None
vector_stores: list[str] | None = None
agents: list[str] | None = None
agent_access_groups: list[str] | None = None
models: list[str] | None = None
search_tools: list[str] | None = None
mcp_tool_search_enabled: bool | None = None
skills: list[str] | None = None
from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402
from litellm.types.object_permission import ( # noqa: E402
ObjectPermissionDict as ObjectPermissionDict,
)
class GenerateRequestBase(LiteLLMPydanticObjectBase):
"""
Overlapping schema between key and user generate/update requests
"""
key_alias: str | None = None
duration: str | None = None
models: list[Any] | None = []
spend: float | None = 0
max_budget: float | None = None
user_id: str | None = None
team_id: str | None = None
agent_id: str | None = None
max_parallel_requests: int | None = None
metadata: dict[Any, Any] | None = {}
tpm_limit: int | None = None
rpm_limit: int | None = None
budget_duration: str | None = None
budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows
allowed_cache_controls: list[object] | None = []
config: dict[object, object] | None = {}
permissions: dict[object, object] | None = {}
model_max_budget: dict[object, object] | None = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {}
budget_fallbacks: dict[str, list[str]] | None = None
model_config = ConfigDict(protected_namespaces=())
model_rpm_limit: dict[object, object] | None = None
model_tpm_limit: dict[object, object] | None = None
mcp_rpm_limit: dict[str, int] | None = None
tag_rpm_limit: dict[str, int] | None = None
guardrails: list[str] | None = None
policies: list[str] | None = None
prompts: list[str] | None = None
blocked: bool | None = None
aliases: dict[object, object] | None = {}
object_permission: LiteLLM_ObjectPermissionBase | None = None
@field_validator("max_budget", mode="before")
@classmethod
def check_max_budget(cls, v: object) -> object:
if v == "":
return None
return v
class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase):
index_name: str
index_permissions: list[Literal["read", "write"]]
class KeyRequestBase(GenerateRequestBase):
key: str | None = None
tpd_limit: int | None = None
default_estimated_output_tokens: PositiveInt | None = None
default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None
budget_id: str | None = None
end_user_budget_id: str | None = None
tags: list[str] | None = None
disable_global_guardrails: bool | None = None
enable_prompt_caching: bool | None = None
throttle_on_budget_exceeded: bool | None = None
enforced_params: list[str] | None = None
allowed_routes: list | None = []
allowed_passthrough_routes: list | None = None
denied_passthrough_routes: list[str] | None = None
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
rpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm
)
tpm_limit_type: Literal["guaranteed_throughput", "best_effort_throughput", "dynamic"] | None = (
None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm
)
router_settings: UpdateRouterConfig | None = None
access_group_ids: list[str] | None = None
class LiteLLMKeyType(str, enum.Enum):
"""
Enum for key types that determine what routes a key can access
"""
LLM_API = "llm_api" # Can call LLM API routes (chat/completions, embeddings, etc.)
MANAGEMENT = "management" # Can call management routes (user/team/key management)
READ_ONLY = "read_only" # Can only call info/read routes
DEFAULT = "default" # Uses default allowed routes
class GenerateKeyRequest(KeyRequestBase):
soft_budget: float | None = None
send_invite_email: bool | None = None
key_type: LiteLLMKeyType | None = Field(
default=LiteLLMKeyType.DEFAULT,
description="Type of key that determines default allowed routes.",
)
auto_rotate: bool | None = Field(default=False, description="Whether this key should be automatically rotated")
rotation_interval: str | None = Field(
default=None,
description="How often to rotate this key (e.g., '30d', '90d'). Required if auto_rotate=True",
)
organization_id: str | None = None
project_id: str | None = None
@field_validator("team_id", "organization_id", "project_id", mode="before")
@classmethod
def treat_cleared_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
AllowedVectorStoreIndexItem as AllowedVectorStoreIndexItem, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
GenerateKeyRequest as GenerateKeyRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
KeyRequestBase as KeyRequestBase, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
LiteLLMKeyType as LiteLLMKeyType, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.request_base import ( # noqa: E402 # re-export after the definitions above
GenerateRequestBase as GenerateRequestBase, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.request_base import ( # noqa: E402 # re-export after the definitions above
LiteLLM_ObjectPermissionBase as LiteLLM_ObjectPermissionBase, # noqa: PLC0414 # public re-export
)
class GenerateKeyResponse(KeyRequestBase):
@ -1333,61 +1228,12 @@ class GenerateKeyResponse(KeyRequestBase):
return values
class UpdateKeyRequest(KeyRequestBase):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
duration: str | None = None
spend: float | None = None
soft_budget: float | None = None
metadata: dict | None = None
temp_budget_increase: float | None = None
temp_budget_expiry: datetime | None = None
auto_rotate: bool | None = None
rotation_interval: str | None = None
organization_id: str | None = None
project_id: str | None = Field(
default=None,
description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.",
)
@model_validator(mode="before")
@classmethod
def drop_blank_team_id(cls, values: object) -> object:
if isinstance(values, Mapping) and values.get("team_id") == "":
return MappingProxyType({k: v for k, v in values.items() if k != "team_id"})
return values
@field_validator("organization_id", mode="before")
@classmethod
def treat_cleared_organization_id_as_unset(cls, v: object) -> object:
if v == "":
return None
return v
@model_validator(mode="after")
def validate_temp_budget(self) -> "UpdateKeyRequest":
if self.temp_budget_increase is not None or self.temp_budget_expiry is not None:
if self.temp_budget_increase is None or self.temp_budget_expiry is None:
raise ValueError("temp_budget_increase and temp_budget_expiry must be set together")
return self
@model_validator(mode="after")
def validate_key_identifier(self) -> "UpdateKeyRequest":
if self.key is None and self.key_alias is None:
raise ValueError("either key or key_alias must be provided")
return self
class RegenerateKeyRequest(GenerateKeyRequest):
# This needs to be different from UpdateKeyRequest, because "key" is optional for this
key: str | None = None
new_key: str | None = None
duration: str | None = None
spend: float | None = None
metadata: dict | None = None
new_master_key: str | None = None
grace_period: str | None = None # Duration to keep old key valid (e.g. "24h", "2d"); None = immediate revoke
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
RegenerateKeyRequest as RegenerateKeyRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.key_requests import ( # noqa: E402 # re-export after the definitions above
UpdateKeyRequest as UpdateKeyRequest, # noqa: PLC0414 # public re-export
)
class ResetSpendRequest(LiteLLMPydanticObjectBase):
@ -1809,19 +1655,9 @@ class MCPSubmissionsSummary(LiteLLMPydanticObjectBase):
######## Skills API Types ########
class NewSkillRequest(LiteLLMPydanticObjectBase):
"""Request to create a new skill in LiteLLM database"""
display_title: str | None = None
description: str | None = None
instructions: str | None = None
file_content: bytes | None = None # Binary content of skill files (zip)
file_name: str | None = None # Original filename
file_type: str | None = None # MIME type (e.g., "application/zip")
metadata: dict[str, Any] | None = None
authorization_url: str | None = None
token_url: str | None = None
registration_url: str | None = None
from litellm.models.skills import ( # noqa: E402 # public re-export
NewSkillRequest as NewSkillRequest, # noqa: PLC0414 # public re-export
)
class UpdateSkillRequest(LiteLLMPydanticObjectBase):
@ -1849,41 +1685,12 @@ class ListSkillsRequest(LiteLLMPydanticObjectBase):
offset: int | None = 0
class NewUserRequestTeam(LiteLLMPydanticObjectBase):
team_id: str
max_budget_in_team: float | None = None
user_role: Literal["user", "admin"] = "user"
class NewUserRequest(GenerateRequestBase):
max_budget: float | None = None
user_email: str | None = None
user_alias: str | None = None
user_role: (
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
| None
) = None
teams: list[str] | list[NewUserRequestTeam] | None = None
auto_create_key: bool = True # flag used for returning a key as part of the /user/new response
send_invite_email: bool | None = None
sso_user_id: str | None = None
organizations: list[str] | None = None
password: str | None = None
@field_validator("password")
@classmethod
def password_not_supported(cls, value: str | None) -> str | None:
if value is not None:
raise ValueError(
"password cannot be set via /user/new. Users set their own password through an "
"invitation link (POST /invitation/new)."
)
return value
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
NewUserRequest as NewUserRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
NewUserRequestTeam as NewUserRequestTeam, # noqa: PLC0414 # public re-export
)
class NewUserResponse(GenerateKeyResponse):
@ -1905,36 +1712,12 @@ class NewUserResponse(GenerateKeyResponse):
updated_at: datetime | None = None
class UpdateUserRequestNoUserIDorEmail(GenerateRequestBase): # shared with BulkUpdateUserRequest
# repr=False keeps the plaintext out of management-endpoint alerts, which str() the request model
password: str | None = Field(default=None, repr=False)
spend: float | None = None
metadata: dict[object, object] | None = None
user_alias: str | None = None
user_role: (
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
| None
) = None
max_budget: float | None = None
class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail):
# Note: the defaults of all Params here MUST BE NONE
# else they will get overwritten
user_id: str | None = None
user_email: str | None = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values: Mapping[str, object]) -> Mapping[str, object]:
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
UpdateUserRequest as UpdateUserRequest, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
UpdateUserRequestNoUserIDorEmail as UpdateUserRequestNoUserIDorEmail, # noqa: PLC0414 # public re-export
)
class ChangePasswordRequest(LiteLLMPydanticObjectBase):
@ -2116,6 +1899,7 @@ class NewTeamRequest(TeamBase):
team_member_tpm_limit: int | None = None # allow user to set TPM limit for all team members
team_member_key_duration: str | None = None # e.g. "1d", "1w", "1m"
team_member_budget_duration: str | None = None # e.g. "30d", "1mo"
team_member_model_max_budget: GenericBudgetConfigType | None = None
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
enforced_batch_output_expires_after: dict | None = None
enforced_file_expires_after: dict | None = None
@ -2175,6 +1959,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
require_trace_id: bool | None = None
team_member_budget: float | None = None
team_member_budget_duration: str | None = None
team_member_model_max_budget: GenericBudgetConfigType | None = None
team_member_rpm_limit: int | None = None
team_member_tpm_limit: int | None = None
team_member_key_duration: str | None = None
@ -2261,7 +2046,7 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
@classmethod
def validate_callback_vars(cls, values):
callback_vars: Final = values.get("callback_vars", {})
valid_keys: Final = set(StandardCallbackDynamicParams.__annotations__.keys())
valid_keys: Final = {*StandardCallbackDynamicParams.__annotations__, CAPTURE_MESSAGE_CONTENT_VAR}
for key, value in callback_vars.items():
if key not in valid_keys:
raise ValueError(f"Invalid callback variable: {key}. Must be one of {valid_keys}")
@ -2273,6 +2058,8 @@ class AddTeamCallback(LiteLLMPydanticObjectBase):
validate_langfuse_span_scope_value(callback_vars[key])
if key == "arize_otlp_protocol":
validate_arize_otlp_protocol_value(callback_vars[key])
if key == CAPTURE_MESSAGE_CONTENT_VAR:
validate_capture_message_content_value(callback_vars[key])
return values
@ -3272,295 +3059,15 @@ from litellm.models.verification_token import ( # noqa: E402
from litellm.models.verification_token import ( # noqa: E402
LiteLLM_VerificationToken as LiteLLM_VerificationToken,
)
_UntypedDict: TypeAlias = dict[Any, Any]
def _is_str_keyed_dict(value: object) -> TypeIs[dict[str, object]]: # guard-ok: pydantic input and JSON keys are str
return isinstance(value, dict)
class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
"""
Combined view of litellm verification token + litellm team table (select values)
"""
team_spend: float | None = None
team_alias: str | None = None
team_tpm_limit: int | None = None
team_rpm_limit: int | None = None
team_tpd_limit: int | None = None
team_max_budget: float | None = None
team_soft_budget: float | None = None
team_model_max_budget: dict[str, object] | None = None
team_models: list[Any] = []
team_blocked: bool = False
soft_budget: float | None = None
team_model_aliases: _UntypedDict | None = None
team_member: Member | None = None
team_metadata: _UntypedDict | None = None
team_object_permission_id: str | None = None
# Team Member Specific Params
team_member_spend: float | None = None
team_member_tpm_limit: int | None = None
team_member_rpm_limit: int | None = None
# End User Params
end_user_id: str | None = None
end_user_tpm_limit: int | None = None
end_user_rpm_limit: int | None = None
end_user_tpd_limit: int | None = None
end_user_max_budget: float | None = None
end_user_model_max_budget: _UntypedDict | None = None
# Organization Params
organization_alias: str | None = None
organization_max_budget: float | None = None
organization_tpm_limit: int | None = None
organization_rpm_limit: int | None = None
organization_metadata: _UntypedDict | None = None
# Project Params
project_alias: str | None = None
project_metadata: _UntypedDict | None = None
# Time stamps
last_refreshed_at: float | None = None # last time joint view was pulled from db
def __init__(self, **kwargs: object) -> None:
# Handle litellm_budget_table_* keys (budget table overrides when key value is None or empty)
for key, value in list(kwargs.items()):
if key.startswith("litellm_budget_table_") and value is not None:
# Extract the corresponding attribute name
attr_name = key.replace("litellm_budget_table_", "")
# Use key's value from kwargs (from DB view), not class default
current = kwargs.get(attr_name)
if current is None:
current = getattr(self, attr_name, None)
# Apply budget value when key has no value, or for model_max_budget when key has empty dict
should_apply = current is None or (
attr_name == "model_max_budget" and _is_str_keyed_dict(current) and len(current) == 0
)
if should_apply:
kwargs[attr_name] = value
if key == "end_user_id" and value is not None and isinstance(value, int):
kwargs[key] = str(value)
if kwargs.get("organization_id") is not None:
kwargs["org_id"] = kwargs.pop("organization_id")
# Initialize the superclass
super().__init__(**kwargs) # pyright: ignore[reportArgumentType, reportUnknownMemberType] # pydantic validates raw values
class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response object for user api key auth
"""
Return the row in the db
"""
api_key: str | None = None
user_role: LitellmUserRoles | None = None
allowed_model_region: AllowedModelRegion | None = None
parent_otel_span: Span | None = None
rpm_limit_per_model: dict[str, int] | None = None
tpm_limit_per_model: dict[str, int] | None = None
user_tpm_limit: int | None = None
user_rpm_limit: int | None = None
user_email: str | None = None
user_spend: float | None = None
user_max_budget: float | None = None
# Values stay `object` rather than BudgetConfig: this is the raw JSON column,
# and validating it here would make one malformed row fail auth outright.
# resolve_model_budget validates the single entry a request actually needs.
user_model_max_budget: Mapping[str, object] | None = None
request_route: str | None = None
is_session_token: bool = False
# Server-only marker set exclusively by the MCP gateway admission path
# (reload_admitted_user) for a keyless user-subject admitted via a gateway DCR session
# bearer or bridge envelope. Not a DB column and never populated from caller-controlled key
# metadata or JWT claims, so it cannot be forged to gain the team-inherited MCP grant union
# or to escape the caller-Authorization egress scrub. exclude=True keeps it out of serialization.
mcp_admitted_user_subject: bool = Field(default=False, exclude=True)
requires_fresh_policy: bool = Field(default=False, exclude=True)
mcp_explicit_grants_only: bool = Field(default=False, exclude=True)
# team_id -> that team's mcp_rpm_limit map, for a keyless admitted subject that reaches MCP
# servers through several teams at once and therefore has no single team_id for the limiter to
# key off. Server-only and stripped from validated input for the same reason as the marker
# above: a forged entry would let a caller pick which team's rpm bucket it is charged against.
mcp_source_team_rpm_limits: dict[str, dict[str, int]] | None = Field(default=None, exclude=True)
# The single MCP server_id a gateway session bearer was scoped to at authorize time (RFC 8707
# resource), or None for an aggregate-scope session. A RESTRICTION intersected against the live
# grant resolution, never a grant. Server-only, set exclusively by the MCP gateway admission
# path via post-construction assignment and stripped from validated input like the markers
# above; a forged value could at most narrow, but the stripping keeps the field's provenance
# single-owner so its meaning stays trustworthy.
mcp_session_resource_server_id: str | None = Field(default=None, exclude=True)
mcp_toolset_id: str | None = Field(default=None, exclude=True)
authenticated_by_custom_auth: bool = Field(default=False, exclude=True)
via_virtual_key: bool = Field(
default=False,
exclude=True,
description=(
"Server-only marker set exclusively by the DB virtual-key and master-key auth paths via "
"post-construction assignment. Stripped from validated input so custom auth handlers, JWT "
"claims, or key metadata cannot forge it. Gates overwrite_user_with_key_hash stamping: only "
"a credential the proxy itself validated as a key may be forwarded as the provider-facing "
"user id."
),
)
invoked_agent_id: str | None = Field(default=None, exclude=True)
invoked_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
agent_invocation_cost: float | None = Field(default=None, exclude=True)
billing_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
_managed_delegation_verified: bool = PrivateAttr(default=False)
managed_agent_policy: AgentResponse | None = Field(default=None, exclude=True)
managed_agent_context: ManagedAgentContext | None = Field(default=None, exclude=True)
agent_caller: AgentCaller | None = Field(
default=None,
exclude=True,
description=(
"Set per request from the x-litellm-user-id / x-litellm-team-id headers an agent echoes back on "
"calls made with its own key. Every check treats it as a ceiling, so a forged value can only "
"narrow the agent's access."
),
)
budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True)
team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True)
user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)
org_budget_snapshot: OrgBudgetSnapshot | None = Field(default=None, exclude=True)
matched_model_access_groups: list[str] | None = Field(default=None, exclude=True)
budget_throttle_pct: float | None = Field(default=None, exclude=True)
user: Any | None = None # Expanded user object when expand=user is used
created_by_user: Any | None = None # Expanded created_by user when expand=user is used
end_user_object_permission: LiteLLM_ObjectPermissionTable | None = None
# Team object_permission preloaded in auth (e.g. get_team_object) to avoid
# per-request object_permission fetches in downstream checks (vector stores, etc.)
team_object_permission: LiteLLM_ObjectPermissionTable | None = None
# Decoded upstream IdP claims (groups, roles, etc.) propagated by JWT auth machinery
# and forwarded into outbound tokens by guardrails such as MCPJWTSigner.
jwt_claims: _UntypedDict | None = None
model_config = ConfigDict(arbitrary_types_allowed=True)
@model_validator(mode="before")
@classmethod
def check_api_key(cls, values: object) -> object:
# If values is already an instance (not a dict), return it as-is
if not _is_str_keyed_dict(values):
return values
# mcp_admitted_user_subject is a server-only marker, set ONLY by the MCP gateway admission
# path via post-construction assignment. Strip it from any validated input (constructor
# kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data.
values.pop("mcp_admitted_user_subject", None)
values.pop("requires_fresh_policy", None)
values.pop("mcp_explicit_grants_only", None)
values.pop("mcp_source_team_rpm_limits", None)
values.pop("mcp_session_resource_server_id", None)
values.pop("mcp_toolset_id", None)
values.pop("via_virtual_key", None)
values.pop("authenticated_by_custom_auth", None)
values.pop("agent_caller", None)
values.pop("managed_agent_context", None)
values.pop("managed_agent_policy", None)
values.pop("invoked_agent_id", None)
values.pop("invoked_agent_policy", None)
values.pop("agent_invocation_cost", None)
values.pop("billing_agent_policy", None)
api_key: Final = values.get("api_key")
if api_key is not None:
values.update({"token": cls._safe_hash_litellm_api_key(api_key)}) # pyright: ignore[reportArgumentType] # non-str fails in the hasher
if isinstance(api_key, str):
values.update({"api_key": cls._safe_hash_litellm_api_key(api_key)})
return values
@classmethod
def _safe_hash_litellm_api_key(cls, api_key: str) -> str:
"""
Helper to ensure all logged keys are hashed
Covers:
1. Regular API keys from LiteLLM DB
2. JWT tokens used for connecting to LiteLLM API
"""
normalized = api_key
if normalized[:7].lower() == "bearer ":
normalized = normalized[7:]
if normalized.startswith("sk-"):
return hash_token(normalized)
if is_jwt(token=normalized):
return f"hashed-jwt-{hash_token(token=normalized)}"
return normalized
@classmethod
def get_litellm_internal_health_check_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for the litellm internal health check service account.
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
team_id=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
)
@classmethod
def get_litellm_cli_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for the litellm internal health check service account.
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_CLI_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
team_id=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME,
)
@classmethod
def get_litellm_internal_jobs_user_api_key_auth(cls) -> "UserAPIKeyAuth":
"""
Returns a `UserAPIKeyAuth` object for internal LiteLLM jobs like key rotation.
This is used to track actions performed by automated system jobs.
"""
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
team_id="system",
key_alias=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
team_alias="system",
user_id="system",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
@property
def is_team_service_account(self) -> bool:
return (
self.user_id is None
and self.team_id is not None
and bool(self.metadata) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped base field
and self.metadata.get("service_account_id") is not None # pyright: ignore[reportUnknownMemberType] # untyped base field
)
def user_api_key_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool:
"""Return True if the caller's role grants unscoped read access to all
tenant resources (managed files, batches, vector stores, spend rows, etc).
Lives on _types.py so leaf modules (e.g. litellm.llms.base_llm.managed_resources)
can use it without pulling in litellm.proxy.utils via management_endpoints.
"""
return user_api_key_dict.user_role in (
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
from litellm.types.proxy.auth.user_api_key_auth import ( # noqa: E402 # re-export after the definitions above
LiteLLM_VerificationTokenView as LiteLLM_VerificationTokenView, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_api_key_auth import (
UserAPIKeyAuth as UserAPIKeyAuth, # noqa: PLC0414 # public re-export
)
from litellm.types.proxy.auth.user_api_key_auth import (
user_api_key_has_admin_view as user_api_key_has_admin_view, # noqa: PLC0414 # public re-export
)
class UserInfoResponse(LiteLLMPydanticObjectBase):
@ -3792,10 +3299,6 @@ class LiteLLM_UserTableFiltered(LiteLLMBaseModel): # done to avoid exposing sen
user_email: str | None = None
class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable):
key_count: int = 0
from litellm.models.access_group import ( # noqa: E402
LiteLLM_AccessGroupTable as LiteLLM_AccessGroupTable,
)
@ -3809,6 +3312,9 @@ from litellm.models.spend_logs import ( # noqa: E402
LiteLLM_SpendLogs as LiteLLM_SpendLogs,
)
from litellm.models.tag import LiteLLM_TagTable as LiteLLM_TagTable # noqa: E402
from litellm.types.proxy.management_endpoints.user_requests import ( # noqa: E402 # re-export after the definitions above
LiteLLM_UserTableWithKeyCount as LiteLLM_UserTableWithKeyCount, # noqa: PLC0414 # public re-export
)
AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked", "unblocked", "rotated", "kill_switch_fired"]
@ -4584,16 +4090,9 @@ class OrganizationAddMemberResponse(LiteLLMPydanticObjectBase):
updated_organization_memberships: list[LiteLLM_OrganizationMembershipTable]
class MemberDeleteRequest(LiteLLMPydanticObjectBase):
user_id: str | None = None
user_email: str | None = None
@model_validator(mode="before")
@classmethod
def check_user_info(cls, values: Mapping[str, object]) -> Mapping[str, object]:
if values.get("user_id") is None and values.get("user_email") is None:
raise ValueError("Either user id or user email must be provided")
return values
from litellm.types.proxy.management_endpoints.team_endpoints import ( # noqa: E402 # re-export after the definitions above
MemberDeleteRequest as MemberDeleteRequest, # noqa: PLC0414 # public re-export
)
class MemberUpdateResponse(LiteLLMPydanticObjectBase):
@ -4661,6 +4160,13 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
default=None,
description="UTC expiry for temp_budget_increase",
)
model_max_budget: GenericBudgetConfigType | None = Field(
default=None,
description=(
"Per-model spend caps for this team member, each with its own budget_duration. "
"Overrides the team's default per-model member budget. Pass an empty dict to fall back to the team default."
),
)
@model_validator(mode="after")
def validate_temp_budget(self) -> "TeamMemberUpdateRequest":
@ -4679,6 +4185,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse):
allowed_models: list[str] | None = None
temp_budget_increase: float | None = None
temp_budget_expiry: datetime | None = None
model_max_budget: GenericBudgetConfigType | None = None
class TeamModelAddRequest(LiteLLMBaseModel):

4
litellm/proxy/auth/auth_checks.py
View file

@ -1739,6 +1739,7 @@ async def get_team_member_default_budget(
budget_id: str,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
raise_on_lookup_error: bool = False,
) -> LiteLLM_BudgetTable | None:
"""
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
@ -1773,6 +1774,8 @@ async def get_team_member_default_budget(
)
except Exception:
verbose_proxy_logger.exception("Error fetching team-default member budget %s", budget_id)
if raise_on_lookup_error:
raise
return None
if budget_record is None:
@ -5605,6 +5608,7 @@ def can_project_access_model(
model=model,
llm_router=llm_router,
models=project_object.models if project_object else [],
team_id=project_object.team_id if project_object else None,
key_model_aliases=key_model_aliases,
object_type="project",
)

72
litellm/proxy/auth/fallback_budget.py
View file

@ -11,13 +11,13 @@ attempted, mirroring `fallback_model_access.py`. It deliberately leaves the prim
a zero-cost model is never blocked by budget, and only the paid fallback is refused. On by default;
set `general_settings.enforce_fallback_budget: false` to restore the unguarded behaviour.
Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation
path before it can be: team, team-member, end-user, org, global and per-model budgets, whose
auth-path functions enforce rather than report (they raise), so reusing them would fire threshold
alerts and take spend reservations for a target that is then skipped; and the key's rolling
`budget_limits` windows, whose accumulated spend lives only in per-window counters
(`spend:key:{token}:window:{budget_duration}`), so enforcing them means more counter reads on the
fallback path rather than reusing state auth already loaded.
Scope: the key's and the user's `max_budget`, plus the team member's per-model caps. Not covered yet,
and each needs a read-only evaluation path before it can be: team and team-member overall, end-user,
org, global and other per-model budgets, whose auth-path functions enforce rather than report (they
raise), so reusing them would fire threshold alerts and take spend reservations for a target that is
then skipped; and the key's rolling `budget_limits` windows, whose accumulated spend lives only in
per-window counters (`spend:key:{token}:window:{budget_duration}`), so enforcing them means more
counter reads on the fallback path rather than reusing state auth already loaded.
Two known limitations of that narrow scope, both shared with `fallback_model_access.py`:
@ -33,10 +33,11 @@ Two known limitations of that narrow scope, both shared with `fallback_model_acc
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Final
from typing import Final, Protocol
from pydantic import ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module exports
@ -55,6 +56,16 @@ class _FallbackBudgetSettings(LiteLLMBaseModel):
enforce_fallback_budget: bool = True
class _TeamMemberModelBudgetLimiter(Protocol):
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None:
try:
return _RequestMetadata.model_validate(metadata).user_api_key_auth
@ -103,7 +114,13 @@ async def _counter_spend(counter_key: str, fallback_spend: float, max_budget: fl
)
async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool:
async def is_token_within_budget_for_model(
*,
model: str,
valid_token: UserAPIKeyAuth,
llm_router: Router,
team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None,
) -> bool:
"""
True when the key and the user behind it can still pay for `model`.
@ -139,6 +156,23 @@ async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKe
if user_spend >= user_budget:
return False
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
if (
team_member_model_budget_limiter is not None
and team_member_model_max_budget
and isinstance(valid_token.user_id, str)
and isinstance(valid_token.team_id, str)
):
try:
await team_member_model_budget_limiter.is_team_member_within_model_budget(
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_member_model_max_budget=team_member_model_max_budget,
model=model,
)
except litellm.BudgetExceededError:
return False
return True
@ -151,6 +185,7 @@ class RouterFallbackBudgetCheck:
"""
is_enforced: Callable[[], bool]
team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool:
if not self.is_enforced():
@ -158,11 +193,20 @@ class RouterFallbackBudgetCheck:
valid_token: Final = _user_api_key_auth_from_request(request_kwargs)
if valid_token is None:
return True
try:
return await is_token_within_budget_for_model(model=model, valid_token=valid_token, llm_router=llm_router)
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller
verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
return False
budget_models: Final = tuple(dict.fromkeys((model, llm_router.routable_model_group(model))))
for budget_model in budget_models:
try:
if not await is_token_within_budget_for_model(
model=budget_model,
valid_token=valid_token,
llm_router=llm_router,
team_member_model_budget_limiter=self.team_member_model_budget_limiter,
):
return False
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller
verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
return False
return True
router_fallback_budget_check: Final = RouterFallbackBudgetCheck(is_enforced=_enforced_by_general_settings)

709
litellm/proxy/auth/user_api_key_auth.py
View file

@ -11,12 +11,11 @@ import asyncio
import fnmatch
import re
import secrets
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Final, NamedTuple, Protocol, Union, cast
import fastapi
import orjson
from fastapi import HTTPException, Request, WebSocket, status
from fastapi.security.api_key import APIKeyHeader
from starlette.exceptions import WebSocketException
@ -51,13 +50,18 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
_get_user_role, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_is_user_proxy_admin, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # reuse common-check group budget enforcement
_team_member_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # reuse common-check pricing enforcement
_virtual_key_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_virtual_key_max_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_virtual_key_soft_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
cache_key_object,
can_customer_access_model,
can_key_call_model,
can_object_call_model,
can_project_access_model,
can_team_access_model,
check_end_user_budget,
common_checks,
delete_cache_key_object,
@ -67,6 +71,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
get_object_permission,
get_org_object_for_request,
get_project_object,
get_team_member_default_budget,
get_team_membership,
get_team_object,
get_user_object,
@ -76,9 +81,12 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
is_user_proxy_admin,
is_valid_fallback_model,
jwt_key_mapping_cache_key,
key_access_group_grants_model,
key_model_aliases_for_auth_check,
resolve_and_validate_end_user_id,
resolve_default_end_user_budget,
route_skips_budget_checks,
stamp_matched_model_access_groups,
team_member_max_budget_alert_check,
virtual_key_max_budget_alert_check,
virtual_key_max_budget_check,
@ -169,7 +177,13 @@ from litellm.proxy.utils import (
)
from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
from litellm.router_utils.fallback_event_handlers import (
PROVIDER_SCOPED_RESOURCE_KEYS,
fallback_lookup_groups,
fallbacks_disabled_for_request,
get_fallback_model_group_for_lookup_groups,
)
from litellm.secret_managers.main import get_secret_bool
from litellm.types.services import ServiceTypes
@ -332,6 +346,18 @@ def _get_model_names_for_budget_checks(
return model
def _budget_fallback_target_model(target: object) -> str | None:
if isinstance(target, str):
return target
if not isinstance(target, Mapping):
return None
target_mapping: Final = cast(Mapping[str, object], target) # cast-ok: isinstance cannot type the container's items
if set(target_mapping) != {"model"}:
return None
model: Final = target_mapping.get("model")
return model if isinstance(model, str) else None
class _KeyModelBudgetLimiter(Protocol):
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
@ -354,6 +380,34 @@ class _TeamModelBudgetLimiter(Protocol):
) -> bool: ...
class _EveryModelBudgetLimiter(Protocol):
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
async def is_user_within_model_budget(
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
) -> bool: ...
async def is_end_user_within_model_budget(
self, end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str
) -> bool: ...
async def is_team_within_model_budget(
self,
team_id: str,
team_model_max_budget: Mapping[str, object],
key_model_max_budget: Mapping[str, object] | None,
model: str,
) -> bool: ...
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
class _TokenTeamModels(Protocol):
@property
def team_models(self) -> list[str]: ...
@ -374,6 +428,30 @@ def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
return valid_token.team_models
class _TokenModelBudgetFields(Protocol):
@property
def model_max_budget(self) -> Mapping[str, object] | None: ...
@property
def end_user_model_max_budget(self) -> Mapping[str, object] | None: ...
@property
def team_model_aliases(self) -> dict[str, str] | None: ... # mutable-ok: mirrors _can_object_call_model's contract
def _token_model_budget_fields(valid_token: _TokenModelBudgetFields) -> _TokenModelBudgetFields:
return valid_token
class _RouterFallbacks(Protocol):
@property
def fallbacks(self) -> Sequence[object] | None: ...
def _router_fallbacks(router: _RouterFallbacks) -> Sequence[object] | None:
return router.fallbacks
async def _read_user_model_max_budget(
user_id: str | None,
prisma_client: PrismaClient | None,
@ -443,13 +521,455 @@ async def _check_team_model_budget(
)
async def _is_within_every_model_budget(
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _EveryModelBudgetLimiter,
model: str,
) -> bool:
budget_fields: Final = _token_model_budget_fields(valid_token)
key_model_max_budget: Final = budget_fields.model_max_budget
user_model_max_budget: Final = valid_token.user_model_max_budget
end_user_model_max_budget: Final = budget_fields.end_user_model_max_budget
team_model_max_budget: Final = valid_token.team_model_max_budget
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
try:
if valid_token.token is not None and isinstance(key_model_max_budget, Mapping) and key_model_max_budget:
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=model,
)
if valid_token.user_id is not None and isinstance(user_model_max_budget, Mapping) and user_model_max_budget:
await model_max_budget_limiter.is_user_within_model_budget(
user_id=valid_token.user_id,
user_model_max_budget=user_model_max_budget,
model=model,
)
if (
valid_token.end_user_id is not None
and isinstance(end_user_model_max_budget, Mapping)
and end_user_model_max_budget
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_model_max_budget,
model=model,
)
if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget:
await model_max_budget_limiter.is_team_within_model_budget(
team_id=valid_token.team_id,
team_model_max_budget=team_model_max_budget,
key_model_max_budget=key_model_max_budget,
model=model,
)
if (
valid_token.user_id is not None
and valid_token.team_id is not None
and isinstance(team_member_model_max_budget, Mapping)
and team_member_model_max_budget
):
await model_max_budget_limiter.is_team_member_within_model_budget(
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_member_model_max_budget=team_member_model_max_budget,
model=model,
)
except litellm.BudgetExceededError:
return False
return True
async def _is_fallback_model_authorized(
model: str,
valid_token: UserAPIKeyAuth,
llm_model_list: Sequence[object] | None,
llm_router: litellm.Router | None,
prisma_client: PrismaClient | None = None,
team_membership: LiteLLM_TeamMembership | None = None,
project_object: LiteLLM_ProjectTableCachedObj | None = None,
refuse_agent_keys: bool = False,
team_object: LiteLLM_TeamTable | None = None,
) -> bool:
if refuse_agent_keys and valid_token.agent_id is not None:
return False
try:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=valid_token,
llm_router=llm_router,
prisma_client=prisma_client,
)
if team_object is not None:
try:
await can_team_access_model(
model=model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
prisma_client=prisma_client,
)
except ProxyException as team_denial:
if team_denial.type != ProxyErrorTypes.team_model_access_denied:
raise
if not await key_access_group_grants_model(
model=model,
valid_token=valid_token,
team_object=team_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
):
raise
else:
team_models: Final = _token_team_models(valid_token)
if team_models:
can_object_call_model(
model=model,
llm_router=llm_router,
models=team_models,
team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="team",
)
member_allowed_models: Final = (
team_membership.litellm_budget_table.allowed_models
if team_membership is not None and team_membership.litellm_budget_table is not None
else None
)
if member_allowed_models:
can_object_call_model(
model=model,
llm_router=llm_router,
models=member_allowed_models,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="team",
)
if project_object is not None and len(project_object.models) > 0:
can_project_access_model(
model=model,
project_object=project_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)
except ProxyException:
return False
return True
async def _is_within_all_fallback_model_budgets(
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _EveryModelBudgetLimiter,
models: Sequence[str],
) -> bool:
for model in models:
if not await _is_within_every_model_budget(
valid_token=valid_token,
model_max_budget_limiter=model_max_budget_limiter,
model=model,
):
return False
return True
async def _check_team_member_model_budget_with_fallback(
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _EveryModelBudgetLimiter,
models: Sequence[str],
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
request: Request,
llm_model_list: Sequence[object] | None,
llm_router: litellm.Router | None,
project_object: LiteLLM_ProjectTableCachedObj | None = None,
team_object: LiteLLM_TeamTable | None = None,
end_user_object: LiteLLM_EndUserTable | None = None,
route: str = "",
prisma_client: PrismaClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> None:
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
if (
valid_token.user_id is None
or valid_token.team_id is None
or not isinstance(team_member_model_max_budget, Mapping)
or not team_member_model_max_budget
):
return
user_id: Final = valid_token.user_id
team_id: Final = valid_token.team_id
for model_name in models:
await _check_team_member_model_budget_for_model(
user_id=user_id,
team_id=team_id,
team_member_model_max_budget=team_member_model_max_budget,
model_name=model_name,
valid_token=valid_token,
model_max_budget_limiter=model_max_budget_limiter,
request_data=request_data,
request=request,
llm_model_list=llm_model_list,
llm_router=llm_router,
project_object=project_object,
team_object=team_object,
end_user_object=end_user_object,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
async def _check_team_member_model_budget_for_model(
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model_name: str,
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _EveryModelBudgetLimiter,
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
request: Request,
llm_model_list: Sequence[object] | None,
llm_router: litellm.Router | None,
project_object: LiteLLM_ProjectTableCachedObj | None = None,
team_object: LiteLLM_TeamTable | None = None,
end_user_object: LiteLLM_EndUserTable | None = None,
route: str = "",
prisma_client: PrismaClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> None:
try:
await model_max_budget_limiter.is_team_member_within_model_budget(
user_id=user_id,
team_id=team_id,
team_member_model_max_budget=team_member_model_max_budget,
model=model_name,
)
except litellm.BudgetExceededError:
if (
request_data.get("model") != model_name
or llm_router is None
or fallbacks_disabled_for_request(kwargs=request_data)
):
raise
fallback_succeeded: Final = await _try_team_member_model_budget_fallback(
user_id=user_id,
team_id=team_id,
model_name=model_name,
valid_token=valid_token,
model_max_budget_limiter=model_max_budget_limiter,
request_data=request_data,
request=request,
llm_model_list=llm_model_list,
llm_router=llm_router,
project_object=project_object,
team_object=team_object,
end_user_object=end_user_object,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if not fallback_succeeded:
raise
async def _try_team_member_model_budget_fallback(
user_id: str,
team_id: str,
model_name: str,
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _EveryModelBudgetLimiter,
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
request: Request,
llm_model_list: Sequence[object] | None,
llm_router: litellm.Router,
project_object: LiteLLM_ProjectTableCachedObj | None = None,
team_object: LiteLLM_TeamTable | None = None,
end_user_object: LiteLLM_EndUserTable | None = None,
route: str = "",
prisma_client: PrismaClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> bool:
if request.headers.get("content-type", "").lower().startswith("multipart/form-data") or any(
request_data.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS
):
return False
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import proxy_config
resolved_prisma_client: Final = prisma_client if prisma_client is not None else proxy_server.prisma_client
resolved_user_api_key_cache: Final = (
user_api_key_cache if user_api_key_cache is not None else proxy_server.user_api_key_cache
)
resolved_proxy_logging_obj: Final = (
proxy_logging_obj if proxy_logging_obj is not None else proxy_server.proxy_logging_obj
)
try:
router_settings: Final[dict[str, object] | None] = (
cast( # cast-ok: proxy settings return an unparameterized dict
dict[str, object] | None,
await proxy_config.get_hierarchical_router_settings(
user_api_key_dict=valid_token,
prisma_client=resolved_prisma_client,
proxy_logging_obj=resolved_proxy_logging_obj,
),
)
)
except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback.
return False
configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None
router_fallbacks: Final[Sequence[object] | None] = (
cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below
if isinstance(configured_fallbacks, list)
else _router_fallbacks(llm_router)
)
if router_fallbacks is None:
return False
fallback_models, _ = get_fallback_model_group_for_lookup_groups(
fallbacks=list(router_fallbacks),
lookup_groups=fallback_lookup_groups(kwargs=request_data, model_group=model_name),
)
fallback_targets: Final[Sequence[object]] = fallback_models or ()
fallback_candidates: Final[tuple[str, ...]] = tuple(
model for target in fallback_targets if (model := _budget_fallback_target_model(target)) is not None
)
if not fallback_candidates:
return False
model_group_aliases: Final[object] = (
router_settings.get("model_group_alias") if router_settings is not None else None
)
key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime
object, valid_token.aliases
)
resolved_fallback_candidates: Final[tuple[tuple[str, str], ...]] = tuple(
(
candidate,
target if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None else candidate,
)
for candidate in fallback_candidates
)
fallback_candidates_with_served_models: Final = (
(
fallback_candidate,
fallback_model,
resolve_served_model(
fallback_model,
litellm.model_alias_map,
key_aliases,
model_group_aliases,
),
)
for fallback_candidate, fallback_model in resolved_fallback_candidates
)
try:
team_membership: Final = await get_team_membership(
user_id=user_id,
team_id=team_id,
prisma_client=resolved_prisma_client,
user_api_key_cache=resolved_user_api_key_cache,
proxy_logging_obj=resolved_proxy_logging_obj,
)
except Exception: # noqa: BLE001 # membership lookup failures must not authorize a fallback
return False
for fallback_candidate, fallback_model, served_model in fallback_candidates_with_served_models:
if fallback_model == model_name:
continue
try:
if (
litellm.block_requests_for_models_without_pricing
and RouteChecks.is_llm_api_route(route=route)
and _unpriced_models_in_request(model=served_model, llm_router=llm_router)
):
continue
if end_user_object is not None and end_user_object.models:
can_customer_access_model(
model=fallback_model,
end_user_object=end_user_object,
llm_router=llm_router,
valid_token=valid_token,
)
if not is_model_cost_zero(
model=served_model, llm_router=llm_router
) and not await _is_within_all_fallback_model_budgets(
valid_token=valid_token,
model_max_budget_limiter=model_max_budget_limiter,
models=tuple(
dict.fromkeys(
(
fallback_candidate,
fallback_model,
served_model,
llm_router.routable_model_group(served_model),
)
)
),
):
continue
if not await _is_fallback_model_authorized(
model=fallback_model,
valid_token=valid_token,
llm_model_list=llm_model_list,
llm_router=llm_router,
prisma_client=prisma_client,
team_object=team_object,
team_membership=team_membership,
project_object=project_object,
refuse_agent_keys=True,
):
continue
valid_token.matched_model_access_groups = [] # rebind-ok: clear prior groups before evaluating this candidate
matched_groups = await stamp_matched_model_access_groups(
model=fallback_model,
valid_token=valid_token,
team_object=team_object,
project_object=project_object,
llm_router=llm_router,
prisma_client=resolved_prisma_client,
user_api_key_cache=resolved_user_api_key_cache,
proxy_logging_obj=resolved_proxy_logging_obj,
team_membership=team_membership,
team_membership_loaded=True,
)
valid_token.matched_model_access_groups = list(matched_groups) # rebind-ok: spend uses the candidate groups
if matched_groups:
await _model_access_group_max_budget_check(
matched_model_access_groups=list(matched_groups),
prisma_client=resolved_prisma_client,
user_api_key_cache=resolved_user_api_key_cache,
)
except Exception: # noqa: BLE001 # A candidate check failure must not escape the fail-closed walk.
verbose_proxy_logger.warning(
"Skipping fallback candidate %s after a budget or authorization check error",
fallback_model.replace("\r", "").replace("\n", ""),
exc_info=True,
)
continue
rewrite_request_model(request_data=request_data, request=request, model=fallback_model)
if "path_params" in request.scope:
request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model
return True
return False
async def _check_key_model_budget_with_fallback(
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _KeyModelBudgetLimiter,
model_name: str,
request_data: dict,
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
request: Request,
llm_model_list: list | None = None,
llm_model_list: Sequence[object] | None = None,
llm_router: litellm.Router | None = None,
) -> None:
"""
@ -489,32 +1009,16 @@ async def _check_key_model_budget_with_fallback(
)
if fallback_model is None:
raise e
try:
await can_key_call_model(
model=fallback_model,
llm_model_list=llm_model_list,
valid_token=valid_token,
llm_router=llm_router,
)
if valid_token.team_models:
can_object_call_model(
model=fallback_model,
llm_router=llm_router,
models=valid_token.team_models,
team_model_aliases=valid_token.team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="team",
)
except ProxyException:
raise e
request_data["model"] = fallback_model
safe_set_request_parsed_body(request=request, parsed_body=request_data)
request._json = request_data # pyright: ignore[reportPrivateUsage] # Starlette JSON cache
request._body = orjson.dumps(request_data) # pyright: ignore[reportPrivateUsage] # Starlette body cache
path_params: Final = request.scope.get("path_params")
if isinstance(path_params, dict) and "model" in path_params:
path_params["model"] = fallback_model
if not await _is_fallback_model_authorized(
model=fallback_model,
valid_token=valid_token,
llm_model_list=llm_model_list,
llm_router=llm_router,
):
raise
rewrite_request_model(request_data=request_data, request=request, model=fallback_model)
if "path_params" in request.scope:
request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model
def _get_bearer_token_or_received_api_key(api_key: str) -> str:
@ -2543,7 +3047,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
model_name=model_name,
request_data=request_data,
request=request,
llm_model_list=llm_model_list,
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
Sequence[object] | None, llm_model_list
),
llm_router=llm_router,
)
@ -2836,6 +3342,79 @@ def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, o
)
async def _load_team_member_default_model_budget(
user_api_key_auth_obj: UserAPIKeyAuth,
team_object: LiteLLM_TeamTable | None,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
skip_budget_checks: bool,
proxy_logging_obj: ProxyLogging | None,
) -> None:
team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None
team_metadata: Final[Mapping[str, object] | None] = (
cast( # cast-ok: team metadata is a string-keyed JSON object
Mapping[str, object], team_metadata_value
)
if isinstance(team_metadata_value, Mapping)
else None
)
team_member_budget_id: Final[object | None] = (
team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None
)
if (
not isinstance(user_api_key_auth_obj.user_id, str)
or not isinstance(user_api_key_auth_obj.team_id, str)
or prisma_client is None
):
return
try:
membership: Final = await get_team_membership(
user_id=user_api_key_auth_obj.user_id,
team_id=user_api_key_auth_obj.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
member_budget_row: Final = membership.litellm_budget_table if membership is not None else None
member_budget_value: Final[object] = (
getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None
)
if (
membership is not None
and membership.budget_id is not None
and membership.budget_id != team_member_budget_id
and isinstance(member_budget_value, Mapping)
and member_budget_value
):
user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream
member_budget_value
)
return
if not isinstance(team_member_budget_id, str):
return
team_member_default_budget: Final = await get_team_member_default_budget(
budget_id=team_member_budget_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
raise_on_lookup_error=True,
)
except Exception:
if not skip_budget_checks:
raise
verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True)
return
default_budget_value: Final[object] = (
getattr(team_member_default_budget, "model_max_budget", None)
if team_member_default_budget is not None
else None
)
user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream
default_budget_value if isinstance(default_budget_value, Mapping) else None
)
@tracer.wrap()
async def run_centralized_common_checks(
user_api_key_auth_obj: UserAPIKeyAuth,
@ -2864,6 +3443,7 @@ async def run_centralized_common_checks(
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
llm_router,
master_key,
model_max_budget_limiter,
@ -3124,6 +3704,17 @@ async def run_centralized_common_checks(
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
typed_team_object: Final[LiteLLM_TeamTable | None] = (
team_object if isinstance(team_object, LiteLLM_TeamTable) else None
)
await _load_team_member_default_model_budget(
user_api_key_auth_obj=user_api_key_auth_obj,
team_object=typed_team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
skip_budget_checks=skip_budget_checks or route_skips_budget_checks(route=route),
proxy_logging_obj=proxy_logging_obj,
)
# Pin the metadata variable name (litellm_metadata vs metadata) before
# any tag merge runs. Without this, header tags from
@ -3164,19 +3755,42 @@ async def run_centralized_common_checks(
project_object=project_object,
)
if not skip_budget_checks:
models: Final = _get_model_names_for_budget_checks(
model=_get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
)
await _check_team_model_budget(
valid_token=user_api_key_auth_obj,
model_max_budget_limiter=model_max_budget_limiter,
models=_get_model_names_for_budget_checks(
model=_get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
),
models=models,
)
typed_project_object: Final[LiteLLM_ProjectTableCachedObj | None] = (
project_object if isinstance(project_object, LiteLLM_ProjectTableCachedObj) else None
)
if not route_skips_budget_checks(route=route):
await _check_team_member_model_budget_with_fallback(
valid_token=user_api_key_auth_obj,
model_max_budget_limiter=model_max_budget_limiter,
models=models,
request_data=request_data,
request=request,
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
Sequence[object] | None, llm_model_list
),
llm_router=llm_router,
project_object=typed_project_object,
team_object=typed_team_object,
end_user_object=end_user_object if isinstance(end_user_object, LiteLLM_EndUserTable) else None,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
@ -3184,7 +3798,10 @@ async def run_centralized_common_checks(
request_data=request_data,
route=route,
llm_router=llm_router,
team_object=team_object,
team_object=cast( # cast-ok: preserve plain team-table objects through reservation
LiteLLM_TeamTableCachedObj | None,
team_object,
),
user_object=user_object,
end_user_id=end_user_id,
end_user_object=end_user_object,
@ -3855,13 +4472,13 @@ async def enforce_key_and_fallback_model_access(
)
_enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access
_enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access # pyright: ignore[reportUnknownVariableType] # existing tests use this alias
async def _run_post_custom_auth_checks(
valid_token: UserAPIKeyAuth,
request: Request,
request_data: dict,
request_data: dict[str, object],
route: str,
parent_otel_span: Span | None,
) -> UserAPIKeyAuth:
@ -3956,7 +4573,9 @@ async def _run_post_custom_auth_checks(
model_name=model_name,
request_data=request_data,
request=request,
llm_model_list=llm_model_list,
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
Sequence[object] | None, llm_model_list
),
llm_router=llm_router,
)

78
litellm/proxy/common_request_processing.py
View file

@ -83,6 +83,7 @@ from litellm.litellm_core_utils.streaming_handler import (
from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
can_key_call_resolved_model,
is_model_cost_zero,
request_skips_budget_checks,
tag_max_budget_check_for_tags,
)
@ -127,7 +128,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
)
from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
from litellm.types.router_weights import validate_router_weights
@ -183,6 +184,7 @@ ProxyRouteType: TypeAlias = Literal[
"aocr",
"asearch",
"adecisions",
"asystemone",
"avideo_generation",
"avideo_list",
"avideo_status",
@ -250,6 +252,7 @@ from litellm.proxy.litellm_pre_call_utils import (
add_litellm_data_to_request,
refresh_proxy_server_request_body_snapshot,
reject_url_valued_destination,
sanitize_for_log,
)
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
from litellm.types.utils import (
@ -737,6 +740,7 @@ async def _resolve_per_request_model_group_alias(
_REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
_ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
def _request_model(data: Mapping[str, object]) -> str | list[str] | None:
@ -1733,6 +1737,69 @@ def _timing_values(
return getattr(logging_obj, "response_timing_metrics", None) or {}
async def _team_member_within_fallback_model_budget(
user_api_key_dict: UserAPIKeyAuth,
model: str,
llm_router: Router | None,
proxy_config: ProxyConfig | None,
proxy_logging_obj: ProxyLogging,
) -> bool:
member_caps: Final = user_api_key_dict.team_member_model_max_budget
user_id: Final = user_api_key_dict.user_id
team_id: Final = user_api_key_dict.team_id
if not isinstance(member_caps, Mapping) or not member_caps or user_id is None or team_id is None:
return True
from litellm.proxy.proxy_server import model_max_budget_limiter
try:
from litellm.proxy.proxy_server import prisma_client
router_settings: Final = (
_ROUTER_SETTINGS_ADAPTER.validate_python(
await proxy_config.get_hierarchical_router_settings( # pyright: ignore[reportUnknownMemberType] # legacy bare-dict return
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
)
if proxy_config is not None and llm_router is not None
else None
)
raw_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime
object, user_api_key_dict.aliases
)
served_model: Final = resolve_served_model(
model,
litellm.model_alias_map,
raw_aliases,
router_settings.get("model_group_alias") if router_settings is not None else None,
)
if is_model_cost_zero(model=served_model, llm_router=llm_router):
return True
budget_models: Final = (
tuple(dict.fromkeys((model, served_model, llm_router.routable_model_group(served_model))))
if llm_router is not None
else tuple(dict.fromkeys((model, served_model)))
)
for budget_model in budget_models:
await model_max_budget_limiter.is_team_member_within_model_budget(
user_id=user_id,
team_id=team_id,
team_member_model_max_budget=member_caps,
model=budget_model,
)
except litellm.BudgetExceededError:
return False
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the member
verbose_proxy_logger.warning(
"Skipping rate limit fallback to model=%s: member budget lookup failed: %s",
sanitize_for_log(model),
sanitize_for_log(e),
)
return False
return True
class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict):
self.data = data
@ -2008,6 +2075,7 @@ class ProxyBaseLLMRequestProcessing:
"aocr",
"asearch",
"adecisions",
"asystemone",
"avideo_generation",
"avideo_list",
"avideo_status",
@ -2371,6 +2439,14 @@ class ProxyBaseLLMRequestProcessing:
for fallback_model in fallback_models:
if fallback_model == original_model:
continue
if not await _team_member_within_fallback_model_budget(
user_api_key_dict=user_api_key_dict,
model=fallback_model,
llm_router=llm_router,
proxy_config=proxy_config,
proxy_logging_obj=proxy_logging_obj,
):
continue
self.data = independent_snapshot(pristine)
self.data["model"] = fallback_model
try:

81
litellm/proxy/common_utils/callback_config_validation.py
View file

@ -10,6 +10,8 @@ from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VAR
_NEWRELIC_CALLBACK: Final = "newrelic"
_NEWRELIC_VAR_PREFIX: Final = "newrelic_"
_LANGFUSE_OTEL_CALLBACK: Final = "langfuse_otel"
@ -36,6 +38,9 @@ def callback_config_error(
)
if langfuse_error is not None:
return langfuse_error
capture_error: Final = _capture_message_content_error(callback_name, callback_vars, callback_type)
if capture_error is not None:
return capture_error
if callback_name != _NEWRELIC_CALLBACK:
return None
return _newrelic_config_error(callback_vars)
@ -80,6 +85,32 @@ def _langfuse_span_scope_error(callback_name: str | None, callback_vars: Mapping
return None
def _capture_message_content_error(
callback_name: str | None, callback_vars: Mapping[str, str], callback_type: str | None
) -> str | None:
value: Final = callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)
if value is None:
return None
from litellm.integrations.otel.model.config import is_otel_v2_enabled
from litellm.integrations.otel.presets.destinations import destination_capable_backends
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_capture_message_content_value,
)
supported: Final = sorted(destination_capable_backends())
if callback_name not in supported:
return f"{CAPTURE_MESSAGE_CONTENT_VAR} applies to the OTel v2 callbacks {supported} only, not {callback_name!r}"
if not is_otel_v2_enabled():
return f"Per-destination {CAPTURE_MESSAGE_CONTENT_VAR} requires the proxy to run with LITELLM_OTEL_V2=true."
try:
validate_capture_message_content_value(value)
except ValueError as e:
return str(e)
if callback_type == "failure":
return f"{CAPTURE_MESSAGE_CONTENT_VAR} needs callback_type 'success' or 'success_and_failure'"
return None
# Which credential family a dynamic variable belongs to. The families are the
# integrations that share one account: every langfuse_* variable configures the
# same Langfuse project whether it rides the classic callback or the OTel one,
@ -190,20 +221,41 @@ def conflicting_span_scope_error(
)
def conflicting_capture_error(
callback_name: str,
callback_vars: Mapping[str, str] | None,
stored_entries: Sequence[tuple[str | None, Mapping[str, str]]],
) -> str | None:
incoming: Final = None if callback_vars is None else callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)
if incoming is None:
return None
return next(
(
f"{CAPTURE_MESSAGE_CONTENT_VAR} is already set to {stored!r} by another {callback_name} entry. "
f"Every {callback_name} entry shares one value: remove that entry or send the same value."
for name, entry_vars in stored_entries
if name == callback_name and (stored := entry_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)) not in (None, incoming)
),
None,
)
def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str | None:
"""Validate every ``logging`` entry of a team/key metadata payload."""
if not metadata:
return None
entries: Final = metadata.get("logging")
if not isinstance(entries, Sequence) or isinstance(entries, (str, bytes)):
entries: Final = _object_sequence(metadata.get("logging"))
if entries is None:
return None
entry_vars: Final = tuple(_entry_callback_vars(entry) for entry in entries)
named_vars: Final = tuple(zip((_entry_callback_name(entry) for entry in entries), entry_vars))
return next(
(
error
for error in (
*(_logging_entry_error(entry) for entry in entries),
*(conflicting_span_scope_error(entry_vars[i], entry_vars[:i]) for i in range(len(entry_vars))),
*(_conflicting_entry_capture_error(named_vars[i], named_vars[:i]) for i in range(len(named_vars))),
)
if error is not None
),
@ -211,6 +263,31 @@ def logging_metadata_config_error(metadata: Mapping[str, object] | None) -> str
)
def stored_capture_entries(logging_entries: object) -> tuple[tuple[str | None, Mapping[str, str]], ...]:
"""The callback name and vars of every stored ``logging`` entry, for ``conflicting_capture_error``."""
entries: Final = _object_sequence(logging_entries) or ()
return tuple((_entry_callback_name(entry), _entry_callback_vars(entry)) for entry in entries)
def _object_sequence(value: object) -> Sequence[object] | None:
return value if isinstance(value, Sequence) and not isinstance(value, (str, bytes)) else None
def _entry_callback_name(entry: object) -> str | None:
match entry:
case {"callback_name": str() as callback_name}:
return callback_name
case _:
return None
def _conflicting_entry_capture_error(
entry: tuple[str | None, Mapping[str, str]], earlier: Sequence[tuple[str | None, Mapping[str, str]]]
) -> str | None:
callback_name, callback_vars = entry
return None if callback_name is None else conflicting_capture_error(callback_name, callback_vars, earlier)
def _entry_callback_vars(entry: object) -> Mapping[str, str]:
callback_vars: Final = entry.get("callback_vars") if isinstance(entry, Mapping) else None
if not isinstance(callback_vars, Mapping):

38
litellm/proxy/common_utils/custom_openapi_spec.py
View file

@ -437,3 +437,41 @@ class CustomOpenAPISpec:
# Add responses API request schema
return CustomOpenAPISpec.add_responses_api_request_schema(with_embeddings)
_MAX_INLINE_DEPTH: Final = 32
def _inlined(node: JsonValue, defs: Mapping[str, JsonValue], depth: int = 0) -> JsonValue:
if depth >= _MAX_INLINE_DEPTH or not isinstance(node, (list, dict)):
return node
if isinstance(node, list):
return [_inlined(item, defs, depth + 1) for item in node]
ref: Final = node.get("$ref")
if isinstance(ref, str) and ref.startswith("#/$defs/"):
target: Final = _inlined(defs[ref.removeprefix("#/$defs/")], defs, depth + 1)
siblings: Final = {key: _inlined(value, defs, depth + 1) for key, value in node.items() if key != "$ref"}
return {**target, **siblings} if isinstance(target, dict) else siblings
if "propertyName" in node and "mapping" in node:
return {"propertyName": node["propertyName"]}
return {key: _inlined(value, defs, depth + 1) for key, value in node.items() if key != "$defs"}
def inline_request_body(model_class: type, example: Mapping[str, object]) -> JsonObject:
"""
Build an ``openapi_extra["requestBody"]`` for a route that reads its body with ``request.body()``.
Swagger UI resolves ``$ref`` against the whole document, so a schema placed inline in an operation
cannot keep Pydantic's ``#/$defs/...`` references. Every reference is expanded in place instead,
which also keeps the operation self-contained inside the lazy OpenAPI snapshot.
"""
schema: Final = CustomOpenAPISpec.get_pydantic_schema(model_class)
media: Final[JsonObject] = {"example": cast(JsonValue, dict(example))} # cast-ok: JSON example literal
if schema is None:
return {"required": True, "content": {"application/json": media}}
raw_defs: Final = schema.get("$defs")
defs: Final[Mapping[str, JsonValue]] = raw_defs if isinstance(raw_defs, dict) else MappingProxyType({})
return {
"required": True,
"content": {"application/json": {"schema": _inlined(schema, defs), **media}},
}

101
litellm/proxy/common_utils/prompt_cache_pricing.py
View file

@ -1,100 +1,3 @@
from collections.abc import Mapping
from datetime import datetime, timezone
from math import isfinite
from typing import Final
from litellm.llms.anthropic.prompt_cache_pricing import price_cache_tokens
from pydantic import TypeAdapter
import litellm
from litellm.cost_calculator import (
completion_cost, # pyright: ignore[reportUnknownVariableType] # legacy optional parameters are untyped
select_model_name_for_cost_calc,
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
_PRICE_ENTRY: Final = TypeAdapter(Mapping[str, object])
def _valid_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and isfinite(value) and value >= 0
def _has_required_prices(prices: Mapping[str, object], tokens: CacheTokenBuckets, completion_tokens: int = 0) -> bool:
required: Final = (
("input_cost_per_token", True),
("cache_read_input_token_cost", tokens.cache_read_input_tokens > 0),
("cache_creation_input_token_cost", tokens.cache_creation_5m_input_tokens > 0),
("cache_creation_input_token_cost_above_1hr", tokens.cache_creation_1h_input_tokens > 0),
("output_cost_per_token", completion_tokens > 0),
)
if any(needed and not _valid_price(prices.get(key)) for key, needed in required):
return False
return all(
_valid_price(value)
for key, value in prices.items()
if value is not None and any(needed and key.startswith(f"{base}_above_") for base, needed in required)
)
def price_cache_tokens(
model: str, deployment_id: str, tokens: CacheTokenBuckets, completion_tokens: int = 0
) -> float | None:
try:
selected_model: Final = select_model_name_for_cost_calc(
model=model,
completion_response=None,
custom_pricing=True,
custom_llm_provider="anthropic",
router_model_id=deployment_id,
)
if selected_model is None:
return None
model_info: Final = litellm.get_model_info(model=selected_model, custom_llm_provider="anthropic")
registry: Final = _PRICE_ENTRY.validate_python(litellm.model_cost) # pyright: ignore[reportUnknownMemberType] # legacy registry is validated at this boundary
price_entry: Final = registry.get(model_info["key"])
if price_entry is None:
return None
prices: Final = _PRICE_ENTRY.validate_python(price_entry)
if completion_tokens < 0 or not _has_required_prices(prices, tokens, completion_tokens):
return None
usage: Final = Usage(
prompt_tokens=tokens.total_tokens,
completion_tokens=completion_tokens,
total_tokens=tokens.total_tokens + completion_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=tokens.cache_read_input_tokens,
cache_creation_tokens=tokens.cache_creation_5m_input_tokens + tokens.cache_creation_1h_input_tokens,
cache_creation_token_details=CacheCreationTokenDetails(
ephemeral_5m_input_tokens=tokens.cache_creation_5m_input_tokens,
ephemeral_1h_input_tokens=tokens.cache_creation_1h_input_tokens,
),
),
)
logging_obj: Final = Logging(
model=model,
messages=[],
stream=False,
call_type="completion",
start_time=datetime.now(timezone.utc),
litellm_call_id="prompt-cache-prediction",
function_id="prompt-cache-prediction",
)
completion_cost(
completion_response=ModelResponse(model=model, usage=usage),
model=model,
custom_llm_provider="anthropic",
custom_pricing=True,
router_model_id=deployment_id,
litellm_logging_obj=logging_obj,
)
breakdown: Final = logging_obj.cost_breakdown
input_cost: Final = breakdown.get("input_cost") if breakdown is not None else None
output_cost: Final = breakdown.get("output_cost") if breakdown is not None else None
if input_cost is None or output_cost is None:
return None
cost: Final = input_cost + output_cost
return cost if _valid_price(cost) else None
except Exception: # noqa: BLE001 # the shared pricing owners raise plain Exception for unpriceable models
return None
__all__ = ("price_cache_tokens",)

106
litellm/proxy/common_utils/resource_ownership.py
View file

@ -1,92 +1,20 @@
from typing import Final
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.litellm_core_utils.resource_ownership import (
get_primary_resource_owner_scope,
get_resource_owner_scopes,
is_proxy_admin,
user_can_access_resource_owner,
)
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.proxy.auth.user_roles import LitellmUserRoles
def is_proxy_admin(user_api_key_dict: UserAPIKeyAuth | None) -> bool:
if user_api_key_dict is None:
return False
return (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
)
def get_resource_owner_scopes(
user_api_key_dict: UserAPIKeyAuth | None,
) -> list[str]:
"""
Return ownership scopes that may access a user-created proxy resource.
Raw user_id is included for rows created before scope prefixes existed.
Prefixes avoid collisions when falling back to team/org/key ownership
for keys that do not have a user_id.
Identity-less callers (no user_id, team_id, org_id, api_key, or token)
return ``[]`` — they share no scope with any other caller, so access
checks against an existing owner always fail and creates that depend
on a primary scope must reject up front. Returning a shared sentinel
here would let any two identity-less callers see each other's data.
"""
if user_api_key_dict is None:
return []
scopes: Final[list[str]] = []
def _add(scope: str | None) -> None:
if scope and scope not in scopes:
scopes.append(scope)
if user_api_key_dict.user_id:
_add(user_api_key_dict.user_id)
_add(f"user:{user_api_key_dict.user_id}")
if user_api_key_dict.team_id:
_add(f"team:{user_api_key_dict.team_id}")
if user_api_key_dict.org_id:
_add(f"org:{user_api_key_dict.org_id}")
if user_api_key_dict.api_key:
_add(f"key:{user_api_key_dict.api_key}")
if user_api_key_dict.token:
_add(f"key:{user_api_key_dict.token}")
return scopes
def get_primary_resource_owner_scope(
user_api_key_dict: UserAPIKeyAuth | None,
) -> str | None:
"""Return the canonical owner scope to stamp on newly-created rows.
``None`` for identity-less callers — callers that depend on a primary
scope to record ownership must surface that as a hard error rather
than fall back to a shared sentinel (which would collapse every
identity-less caller into the same logical owner).
"""
if user_api_key_dict is None:
return None
if user_api_key_dict.user_id:
return user_api_key_dict.user_id
if user_api_key_dict.team_id:
return f"team:{user_api_key_dict.team_id}"
if user_api_key_dict.org_id:
return f"org:{user_api_key_dict.org_id}"
if user_api_key_dict.api_key:
return f"key:{user_api_key_dict.api_key}"
if user_api_key_dict.token:
return f"key:{user_api_key_dict.token}"
return None
def user_can_access_resource_owner(
owner: str | None,
user_api_key_dict: UserAPIKeyAuth | None,
) -> bool:
if user_api_key_dict is None:
return True
if is_proxy_admin(user_api_key_dict):
return True
if owner is None:
return False
return owner in get_resource_owner_scopes(user_api_key_dict)
__all__ = [
"Final",
"LitellmUserRoles",
"UserAPIKeyAuth",
"get_primary_resource_owner_scope",
"get_resource_owner_scopes",
"is_proxy_admin",
"user_can_access_resource_owner",
]

2
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -21,7 +21,7 @@ if TYPE_CHECKING:
T = TypeVar("T", bound=BaseModel)
AUTH_OBJECTS_TARGET: Final = "auth_objects"
from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa: E402, PLC0414 # public re-export
_HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}")

96
litellm/proxy/decisions_endpoints/endpoints.py
View file

@ -1,6 +1,6 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import Annotated, Final
from typing import Annotated, Final, Literal
from fastapi import APIRouter, Depends, Request, Response
from fastapi.responses import ORJSONResponse # pyright: ignore[reportDeprecated] # required endpoint contract
@ -9,7 +9,15 @@ from pydantic import TypeAdapter, ValidationError
from litellm.exceptions import BadRequestError
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.decisions import DecisionsRequestBody, OpenAIDecisionRequestBody
from litellm.proxy.common_utils.custom_openapi_spec import inline_request_body
from litellm.types.decisions import (
DecisionsRequest,
DecisionsRequestBody,
DecisionsResponse,
OpenAIDecisionRequest,
OpenAIDecisionRequestBody,
OpenAIDecisionResponse,
)
router: Final = APIRouter()
_REQUEST_DATA_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
@ -20,6 +28,51 @@ _OPENAI_DECISION_REQUEST_BODY_ADAPTER: Final[TypeAdapter[OpenAIDecisionRequestBo
_GENERAL_SETTINGS_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_OPTIONAL_STRING_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None)
_OPTIONAL_FLOAT_ADAPTER: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
_SYSTEMONE_REQUEST_BODY: Final = inline_request_body(
DecisionsRequest,
{
"model": "jev",
"state": "Customer wrote: I was charged twice for order #4411 and want one charge refunded today.",
"questions": {
"is_refund_request": {"type": "noul", "instructions": "Is the customer asking for a refund?"},
"urgency": {
"type": "choice",
"instructions": "How urgent is this?",
"criteria": {"low": "can wait a week", "high": "needs action today"},
},
"frustration": {
"type": "score",
"instructions": "How frustrated is the customer?",
"criteria": ["calm", "mildly annoyed", "angry"],
},
},
},
)
_DECISIONS_REQUEST_BODY: Final = inline_request_body(
OpenAIDecisionRequest,
{
"model": "luna",
"input": "Customer wrote: I was charged twice for order #4411 and want one charge refunded today.",
"questions": [
{"type": "predicate", "name": "is_refund_request", "instructions": "Is the customer asking for a refund?"},
{
"type": "choice",
"name": "urgency",
"instructions": "How urgent is this?",
"choices": [
{"value": "low", "description": "can wait a week"},
{"value": "high", "description": "needs action today"},
],
},
{
"type": "score",
"name": "frustration",
"instructions": "How frustrated is the customer?",
"levels": [{"label": "calm"}, {"label": "mildly annoyed"}, {"label": "angry"}],
},
],
},
)
async def _invalid_request(
@ -56,6 +109,7 @@ async def _process_decisions(
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
body_adapter: TypeAdapter[DecisionsRequestBody] | TypeAdapter[OpenAIDecisionRequestBody],
route_type: Literal["adecisions", "asystemone"],
) -> object:
from litellm.proxy.proxy_server import (
general_settings as proxy_general_settings,
@ -93,7 +147,7 @@ async def _process_decisions(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
route_type="adecisions",
route_type=route_type,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
general_settings=general_settings,
@ -121,23 +175,41 @@ async def _process_decisions(
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask named questions about a state (System One format)",
responses={200: {"model": DecisionsResponse}},
openapi_extra={"requestBody": _SYSTEMONE_REQUEST_BODY},
)
@router.post(
"/systemone",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask named questions about a state (System One format)",
responses={200: {"model": DecisionsResponse}},
openapi_extra={"requestBody": _SYSTEMONE_REQUEST_BODY},
)
async def systemone(
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Judge a `state` against 1 to 128 named `questions` and get one calibrated answer per question name.
Question types are `noul` (probability the answer is yes), `choice` (one label out of `criteria`, with
`confidence` and `probabilities`) and `score` (an index into the `criteria` levels, with `confidence`,
`legend` and `probabilities`). Requests that break these rules are rejected with a 400 before any provider
is called. `model` is any decision model in the proxy model list. OpenAI decision models accept this format
too, LiteLLM translates it and keeps the answers keyed by question name. Streaming is not supported.
[Docs](https://docs.litellm.ai/docs/decisions)
"""
return await _process_decisions(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
body_adapter=_DECISIONS_REQUEST_BODY_ADAPTER,
route_type="asystemone",
)
@ -146,21 +218,39 @@ async def systemone(
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask questions about an input (OpenAI Decisions format)",
responses={200: {"model": OpenAIDecisionResponse}},
openapi_extra={"requestBody": _DECISIONS_REQUEST_BODY},
)
@router.post(
"/decisions",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse, # pyright: ignore[reportDeprecated] # required endpoint contract
tags=["decisions"],
summary="Ask questions about an input (OpenAI Decisions format)",
responses={200: {"model": OpenAIDecisionResponse}},
openapi_extra={"requestBody": _DECISIONS_REQUEST_BODY},
)
async def decisions(
request: Request,
fastapi_response: Response,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Judge an `input` against 1 to 128 `questions` and get the answers back in the same order.
Question types are `predicate` (yes/no with `probability`), `choice` (one of the `choices`, with
`probabilities`) and `score` (an index into `levels`, with `probabilities`). Requests that break these
rules are rejected with a 400 before any provider is called. `model` is any decision model in the proxy
model list. System One decision models accept this format too, LiteLLM translates the request and the
answers. Streaming is not supported.
[Docs](https://docs.litellm.ai/docs/decisions)
"""
return await _process_decisions(
request=request,
fastapi_response=fastapi_response,
user_api_key_dict=user_api_key_dict,
body_adapter=_OPENAI_DECISION_REQUEST_BODY_ADAPTER,
route_type="adecisions",
)

136
litellm/proxy/guardrails/anthropic_sse.py
View file

@ -7,115 +7,19 @@ hook scan such a stream, and re-emit it when the guardrail rewrote the response.
from __future__ import annotations
import codecs
import json
from collections.abc import Mapping, Sequence
from typing import Final
from litellm.types.utils import Choices, ModelResponse
_ANTHROPIC_EVENT_TYPES: Final = frozenset(
{
"message_start",
"message_delta",
"message_stop",
"content_block_start",
"content_block_delta",
"content_block_stop",
"ping",
"error",
}
from litellm.llms.anthropic.pass_through.stream_assembly import (
assemble_anthropic_sse_stream as assemble_anthropic_sse_stream, # noqa: PLC0414 # public re-export
)
def is_raw_sse_stream(all_chunks: Sequence[object]) -> bool:
return any(isinstance(chunk, (str, bytes)) for chunk in all_chunks)
def _joined_sse_stream(all_chunks: Sequence[object]) -> str | None:
raw: Final = b"".join(
chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
for chunk in all_chunks
if isinstance(chunk, (str, bytes))
)
try:
return codecs.getincrementaldecoder("utf-8")().decode(raw, final=False)
except UnicodeDecodeError:
return None
def _parsed_sse_events(sse_stream: str) -> tuple[Mapping[str, object], ...]:
from litellm.llms.anthropic.pass_through.stream_assembly import extract_sse_data, split_sse_chunk_into_events
return tuple(
event_data
for event in split_sse_chunk_into_events(sse_stream)
if (event_data := extract_sse_data(event)) is not None
)
def _anthropic_message_start(sse_stream: str) -> Mapping[str, object] | None:
return next(
(
message
for event_data in _parsed_sse_events(sse_stream)
if event_data.get("type") == "message_start" and isinstance(message := event_data.get("message"), dict)
),
None,
)
def is_anthropic_sse_stream(all_chunks: Sequence[object]) -> bool:
"""Whether raw SSE frames are Anthropic Messages events.
``is_raw_sse_stream`` only says the chunks are unparsed bytes, and ``/v1/messages`` is not the
only endpoint that streams those: the Google ``:streamGenerateContent`` route marks its own
stream raw too. Reading its frames as Anthropic ones would refuse the response in a wire format
its client cannot parse, so the surface is decided on the event types actually present.
"""
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return False
return any(event.get("type") in _ANTHROPIC_EVENT_TYPES for event in _parsed_sse_events(sse_stream))
def assemble_anthropic_sse_stream(
all_chunks: Sequence[object], *, restore_identity: bool = False
) -> ModelResponse | None:
"""Assemble raw Anthropic SSE frames into a ModelResponse.
``restore_identity`` stamps the upstream message id and model onto the result, which the
assembler does not carry through. It is off by default so callers that re-emit the assembled
response keep the wire shape they had before this helper was shared. The writes land on a
freshly built object that is unreachable from caller state until returned.
"""
from litellm.llms.anthropic.pass_through.stream_assembly import build_complete_streaming_response
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return None
message_start: Final = _anthropic_message_start(sse_stream)
if message_start is None:
return None
model: Final = message_start.get("model") if restore_identity else None
try:
assembled: Final = build_complete_streaming_response(
all_chunks=(sse_stream,),
litellm_logging_obj=None, # pyright: ignore[reportArgumentType] # only forwarded to stream_chunk_builder, which accepts None
model=model if isinstance(model, str) else "",
)
except Exception: # noqa: BLE001 # stream_chunk_builder re-raises every assembly failure as litellm.APIError
return None
if not isinstance(assembled, ModelResponse):
return None
if not restore_identity:
return assembled
message_id: Final = message_start.get("id")
if isinstance(message_id, str):
assembled.id = message_id
if isinstance(model, str) and model:
assembled.model = model
return assembled
from litellm.llms.anthropic.pass_through.stream_assembly import (
is_anthropic_sse_stream as is_anthropic_sse_stream, # noqa: PLC0414 # public re-export
)
from litellm.llms.anthropic.pass_through.stream_assembly import (
is_raw_sse_stream as is_raw_sse_stream, # noqa: PLC0414 # public re-export
)
from litellm.types.utils import Choices, ModelResponse
def model_response_text(response: ModelResponse) -> str:
@ -141,25 +45,9 @@ def anthropic_sse_error_frames(message: str) -> tuple[bytes, ...]:
)
def is_sse_error_stream(all_chunks: Sequence[object]) -> bool:
"""Whether the buffered stream carries nothing but error frames.
post_call guardrails run in a chain, so a hook can be handed the terminal error frames an
earlier guardrail emitted when it blocked. Those carry no message to assemble, and replacing
them would hide the refusal the client is owed. Covers both wire forms a guardrail emits: the
Anthropic ``error`` event and the chat-completions ``{"error": ...}`` payload.
"""
if not all(isinstance(chunk, (str, bytes)) for chunk in all_chunks):
# A stream mixing typed chunks with an error frame still carries content to scan, and the
# frames-only join below would drop exactly the part that has to be scanned
return False
sse_stream: Final = _joined_sse_stream(all_chunks)
if sse_stream is None:
return False
events: Final = _parsed_sse_events(sse_stream)
return len(events) > 0 and all(
event.get("type") == "error" or isinstance(event.get("error"), Mapping) for event in events
)
from litellm.llms.anthropic.pass_through.stream_assembly import ( # noqa: E402 # re-export after the definitions above
is_sse_error_stream as is_sse_error_stream, # noqa: PLC0414 # public re-export
)
def anthropic_sse_chunks_from_response(assembled: ModelResponse) -> tuple[bytes, ...]:

2
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -1496,6 +1496,7 @@ async def get_guardrail_ui_settings() -> GuardrailUIAddGuardrailSettings:
- PII entity categories for UI organization
- Content filter settings (patterns and categories)
"""
from litellm.decisions.call import supported_decisions_providers
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.patterns import (
PATTERN_CATEGORIES,
get_available_content_categories,
@ -1533,6 +1534,7 @@ async def get_guardrail_ui_settings() -> GuardrailUIAddGuardrailSettings:
supported_modes_by_provider=supported_modes_by_provider,
providers_without_directional_logging_only_scope=providers_without_directional_logging_only_scope,
pii_entity_categories=category_maps,
decision_model_providers=supported_decisions_providers(),
content_filter_settings={
"prebuilt_patterns": get_pattern_metadata(),
"pattern_categories": list(PATTERN_CATEGORIES.keys()),

109
litellm/proxy/guardrails/guardrail_hooks/decision_model/__init__.py Normal file
View file

@ -0,0 +1,109 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter, ValidationError
from litellm.types.guardrails import (
GuardrailEventHooks,
Mode,
SupportedGuardrailIntegrations,
)
from litellm.types.proxy.guardrails.guardrail_hooks.decision_model import (
DecisionModelCheck,
)
from .decision_model import DecisionModelGuardrail
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
_CHECKS_ADAPTER: Final = TypeAdapter(list[DecisionModelCheck])
def _coerce_event_hook(
mode: str | Sequence[str] | Mode,
) -> GuardrailEventHooks | tuple[GuardrailEventHooks, ...] | Mode:
if isinstance(mode, Mode):
return mode
if isinstance(mode, str):
return GuardrailEventHooks(mode)
return tuple(GuardrailEventHooks(item) for item in mode)
def _parse_checks(raw_checks: object) -> tuple[DecisionModelCheck, ...]:
try:
checks: Final = _CHECKS_ADAPTER.validate_python(raw_checks)
except ValidationError as error:
raise ValueError(f"decision_model guardrail has invalid checks: {error}") from error
return tuple(checks)
def _streaming_end_of_stream_only(litellm_params: LitellmParams) -> bool:
value: Final[object] = getattr(litellm_params, "streaming_end_of_stream_only", None)
if value is None:
return False
if not isinstance(value, bool):
raise TypeError(f"decision_model guardrail streaming_end_of_stream_only must be a boolean, got {value!r}")
return value
def _streaming_sampling_rate(litellm_params: LitellmParams) -> int:
value: Final[object] = getattr(litellm_params, "streaming_sampling_rate", None)
if value is None:
return 5
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(f"decision_model guardrail streaming_sampling_rate must be an integer, got {value!r}")
if value < 1:
raise ValueError(f"decision_model guardrail streaming_sampling_rate must be at least 1, got {value}")
return value
def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> DecisionModelGuardrail:
import litellm
decision_model: Final = litellm_params.decision_model
if not decision_model:
raise ValueError("decision_model guardrail requires decision_model in litellm_params")
checks: Final = _parse_checks(litellm_params.checks)
_callback: Final = DecisionModelGuardrail(
guardrail_name=guardrail["guardrail_name"],
decision_model=decision_model,
checks=checks,
event_hook=_coerce_event_hook(litellm_params.mode),
default_on=litellm_params.default_on or False,
unreachable_fallback=(
litellm_params.unreachable_fallback
if "unreachable_fallback" in litellm_params.model_fields_set
else "fail_closed"
),
max_input_chars=litellm_params.max_input_chars or 24000,
timeout=litellm_params.timeout,
max_concurrent_decision_calls=litellm_params.max_concurrent_decision_calls or 8,
streaming_end_of_stream_only=_streaming_end_of_stream_only(litellm_params),
streaming_sampling_rate=_streaming_sampling_rate(litellm_params),
)
litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] # callback manager is untyped
_callback
)
return _callback
guardrail_initializer_registry: Final = {
SupportedGuardrailIntegrations.DECISION_MODEL.value: initialize_guardrail,
}
guardrail_class_registry: Final = {
SupportedGuardrailIntegrations.DECISION_MODEL.value: DecisionModelGuardrail,
}
__all__ = [
"DecisionModelGuardrail",
"guardrail_class_registry",
"guardrail_initializer_registry",
"initialize_guardrail",
]

331
litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py Normal file
View file

@ -0,0 +1,331 @@
"""Decision-model guardrail: a decisions-API model (e.g. Jev) scores the request or
response against predicate checks and blocks when any check crosses its threshold."""
from __future__ import annotations
import asyncio
from asyncio import AbstractEventLoop, Semaphore
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import chain
from time import time
from typing import TYPE_CHECKING, ClassVar, Final, Literal
from fastapi import HTTPException
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped upstream
)
from litellm.litellm_core_utils.llm_judge import default_router_provider, judge_target
from litellm.types.decisions import MAX_DECISION_QUESTIONS, OpenAIDecisionResponse, OpenAIPredicateAnswer
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.llms.openai import ChatCompletionAssistantToolCall, ChatCompletionToolCallChunk
from litellm.types.proxy.guardrails.guardrail_hooks.decision_model import (
DecisionModelCheck,
DecisionModelGuardrailConfigModel,
)
from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs, GuardrailStatus
if TYPE_CHECKING:
from litellm import Router
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
class _CheckVerdict(TypedDict):
name: ReadOnly[str]
probability: ReadOnly[float | None]
threshold: ReadOnly[float]
action: ReadOnly[Literal["block", "log"]]
flagged: ReadOnly[bool]
reason: ReadOnly[str | None]
def _resolve_checks(checks: tuple[DecisionModelCheck, ...]) -> tuple[DecisionModelCheck, ...]:
"""Enforce the check invariants (non-empty list, unique names, bounded count)
that decide what a run sends, wherever the guardrail is built."""
if not checks:
raise ValueError("decision_model guardrail requires at least one check")
if len(checks) > MAX_DECISION_QUESTIONS:
raise ValueError(
f"decision_model guardrail supports at most {MAX_DECISION_QUESTIONS} checks, got {len(checks)}"
)
names: Final = [check.name for check in checks]
if len(set(names)) != len(names):
raise ValueError("decision_model guardrail check names must be unique")
return checks
def _tool_call_text(
tool_call: ChatCompletionToolCallChunk | ChatCompletionAssistantToolCall | ChatCompletionMessageToolCall,
) -> str | None:
if isinstance(tool_call, dict):
function: Final = tool_call.get("function")
dict_arguments: Final = function.get("arguments") if function else None
if not isinstance(dict_arguments, str) or not dict_arguments.strip():
return None
dict_name: Final = function.get("name") if function else None
return f"{dict_name or ''}({dict_arguments})"
arguments: Final = tool_call.function.arguments
if not arguments or not arguments.strip():
return None
return f"{tool_call.function.name or ''}({arguments})"
def _structured_tool_calls(inputs: GenericGuardrailAPIInputs) -> Iterator[ChatCompletionAssistantToolCall]:
for message in inputs.get("structured_messages") or ():
if "role" in message and message["role"] == "assistant":
yield from message.get("tool_calls") or ()
def _tool_call_texts(inputs: GenericGuardrailAPIInputs) -> tuple[str, ...]:
tool_calls: Final = chain(inputs.get("tool_calls") or (), _structured_tool_calls(inputs))
return tuple(dict.fromkeys(text for text in map(_tool_call_text, tool_calls) if text is not None))
def _chunk_text(text: str, max_chars: int) -> tuple[str, ...]:
"""Split text into chunks of at most max_chars that overlap by
min(2000, max_chars // 4) and together cover the whole text."""
if len(text) <= max_chars:
return (text,)
step: Final = max_chars - min(2000, max_chars // 4)
return tuple(text[start : start + max_chars] for start in range(0, len(text) - max_chars + step, step))
class DecisionModelGuardrail(CustomGuardrail):
"""Runs predicate checks on request/response text via the Decisions API and blocks on flagged checks."""
records_own_guardrail_information: ClassVar[bool] = True
def __init__(
self,
guardrail_name: str,
decision_model: str,
checks: Sequence[DecisionModelCheck],
event_hook: GuardrailEventHooks | tuple[GuardrailEventHooks, ...] | Mode | None = None,
default_on: bool = False,
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
max_input_chars: int = 24000,
router_provider: Callable[[], Router | None] | None = None,
timeout: float | None = None,
max_concurrent_decision_calls: int = 8,
streaming_end_of_stream_only: bool = False,
streaming_sampling_rate: int = 5,
) -> None:
super().__init__( # pyright: ignore[reportUnknownMemberType] # base init takes untyped **kwargs
guardrail_name=guardrail_name,
supported_event_hooks=self.get_supported_event_hooks(),
event_hook=list(event_hook)
if isinstance(event_hook, tuple)
else (event_hook or GuardrailEventHooks.pre_call),
default_on=default_on,
)
self.decision_model = decision_model
self.checks = _resolve_checks(tuple(checks))
self.unreachable_fallback = unreachable_fallback
self.max_input_chars = max_input_chars
self._router_provider = router_provider or default_router_provider
self.timeout = timeout
self._max_concurrent_decision_calls = max_concurrent_decision_calls
self.streaming_end_of_stream_only = streaming_end_of_stream_only
self.streaming_sampling_rate = streaming_sampling_rate
self._semaphores_by_loop: dict[AbstractEventLoop, Semaphore] = {} # mutable-ok: per-loop semaphore registry
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: # mutable-ok: base signature returns list
return [GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call, GuardrailEventHooks.post_call]
@staticmethod
def get_config_model() -> type[DecisionModelGuardrailConfigModel] | None:
return DecisionModelGuardrailConfigModel
async def _call_decisions(self, text: str, questions: Sequence[Mapping[str, object]]) -> object:
router: Final = self._router_provider()
if judge_target(router, self.decision_model).via == "router" and router is not None:
return await router.adecisions( # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportUnknownArgumentType] # factory_function-generated member
model=self.decision_model,
input=text,
questions=questions,
num_retries=0,
fallbacks=[],
timeout=self.timeout,
)
return await litellm.adecisions( # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportUnknownArgumentType] # SDK dispatch returns a broad union
model=self.decision_model,
input=text,
questions=questions,
num_retries=0,
timeout=self.timeout,
)
def _get_decision_semaphore(self) -> asyncio.Semaphore:
"""Per-event-loop semaphore shared by every decisions call on this instance."""
loop: Final = asyncio.get_running_loop()
existing: Final = self._semaphores_by_loop.get(loop)
if existing is not None:
return existing
created: Final = asyncio.Semaphore(self._max_concurrent_decision_calls)
self._semaphores_by_loop[loop] = created
return created
async def _run_checks(
self, text: str, questions: Sequence[Mapping[str, object]], semaphore: asyncio.Semaphore
) -> OpenAIDecisionResponse:
async with semaphore:
raw_response: Final = await self._call_decisions(text, questions)
if not isinstance(raw_response, OpenAIDecisionResponse):
raise TypeError(f"decisions call returned {type(raw_response).__name__}, expected OpenAI format")
return raw_response
async def _screen_text(
self, text: str, questions: Sequence[Mapping[str, object]], semaphore: asyncio.Semaphore
) -> tuple[tuple[OpenAIDecisionResponse, ...], tuple[Exception, ...]]:
results: Final = await asyncio.gather(
*(self._run_checks(segment, questions, semaphore) for segment in _chunk_text(text, self.max_input_chars)),
return_exceptions=True,
)
for result in results:
if isinstance(result, BaseException) and not isinstance(result, Exception):
raise result
return (
tuple(result for result in results if isinstance(result, OpenAIDecisionResponse)),
tuple(result for result in results if isinstance(result, Exception)),
)
def _check_verdict(
self, check: DecisionModelCheck, answers_by_call: Sequence[Mapping[str | None, object]]
) -> _CheckVerdict:
"""Aggregate one check across calls: probability is the max over predicate answers;
a call with no predicate answer counts as unanswered, which blocks a Block check under
fail_closed and only warns under fail_open or for log-only checks."""
answers: Final = tuple(call_answers.get(check.name) for call_answers in answers_by_call)
probabilities: Final = tuple(
answer.probability for answer in answers if isinstance(answer, OpenAIPredicateAnswer)
)
missing: Final = len(probabilities) < len(answers)
if missing:
verbose_logger.warning(
"decision_model guardrail %s: check '%s' got no predicate answer in %d of %d decisions calls",
self.guardrail_name,
check.name,
len(answers) - len(probabilities),
len(answers),
)
probability: Final = max(probabilities) if probabilities else None
unanswered_block: Final = missing and check.action == "block" and self.unreachable_fallback == "fail_closed"
return {
"name": check.name,
"probability": probability,
"threshold": check.threshold,
"action": check.action,
"flagged": unanswered_block or (probability is not None and probability >= check.threshold),
"reason": "no_answer" if missing else None,
}
def _verdicts(self, responses: Sequence[OpenAIDecisionResponse]) -> tuple[_CheckVerdict, ...]:
answers_by_call: Final = tuple({answer.name: answer for answer in response.answers} for response in responses)
return tuple(self._check_verdict(check, answers_by_call) for check in self.checks)
def _handle_call_failure(self, error: Exception) -> None:
"""fail_open logs and returns; fail_closed raises a generic 502 (upstream detail stays in server logs)."""
if self.unreachable_fallback == "fail_open":
verbose_logger.warning(
"decision_model guardrail %s: decisions call failed; fail_open configured, allowing the request. Error: %s",
self.guardrail_name,
error,
)
return
verbose_logger.error(
"decision_model guardrail %s: decisions call failed. Error: %s", self.guardrail_name, error
)
raise HTTPException(
status_code=502,
detail={"error": "Decision model guardrail failed to respond"},
)
def _log(
self,
request_data: dict[str, object], # mutable-ok: base helper writes into the live request dict
verdicts: tuple[_CheckVerdict, ...],
status: GuardrailStatus,
start_time: float,
) -> None:
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType] # untyped base helper
guardrail_provider="decision_model",
guardrail_json_response={
"model": self.decision_model,
"checks": [
{
"name": v["name"],
"probability": v["probability"],
"threshold": v["threshold"],
"action": v["action"],
"flagged": v["flagged"],
"reason": v["reason"],
}
for v in verdicts
],
},
request_data=request_data,
guardrail_status=status,
start_time=start_time,
end_time=time(),
)
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object], # mutable-ok: base signature passes the live request dict
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
texts: Final = tuple(text for text in inputs.get("texts") or [] if text and text.strip()) + _tool_call_texts(
inputs
)
if not texts:
return inputs
start_time: Final = time()
questions: Final = [
{"type": "predicate", "name": check.name, "instructions": check.instructions or ""} for check in self.checks
]
semaphore: Final = self._get_decision_semaphore()
responses: Final[list[OpenAIDecisionResponse]] = [] # mutable-ok: sequential per-text accumulator
errors: Final[list[Exception]] = [] # mutable-ok: sequential per-text accumulator
for text in texts:
text_responses, text_errors = await self._screen_text(text, questions, semaphore)
responses.extend(text_responses)
errors.extend(text_errors)
if any(v["flagged"] and v["action"] == "block" for v in self._verdicts(responses)):
break
if errors and not responses:
self._log(request_data, (), "guardrail_failed_to_respond", start_time)
self._handle_call_failure(errors[0])
return inputs
verdicts: Final = self._verdicts(responses)
flagged: Final = tuple(v for v in verdicts if v["flagged"])
intervened: Final = any(v["action"] == "block" for v in flagged)
if errors and not intervened:
self._log(request_data, verdicts, "guardrail_failed_to_respond", start_time)
self._handle_call_failure(errors[0])
return inputs
status: Final = "guardrail_intervened" if intervened else ("guardrail_flagged" if flagged else "success")
self._log(request_data, verdicts, status, start_time)
if intervened:
raise HTTPException(
status_code=400,
detail={
"error": "Violated decision model guardrail policy",
"guardrail_name": self.guardrail_name,
"flagged_checks": [v["name"] for v in flagged if v["action"] == "block"],
},
)
return inputs
__all__ = ["DecisionModelGuardrail"]

2
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -18,7 +18,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
event_hook=litellm_params.mode,
guardrailIdentifier=litellm_params.guardrailIdentifier,
guardrailVersion=litellm_params.guardrailVersion,
checks=litellm_params.checks,
checks=litellm_params.checks if isinstance(litellm_params.checks, BedrockChecksConfigModel) else None,
content_filter_threshold=litellm_params.content_filter_threshold,
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
pii_confidence_threshold=litellm_params.pii_confidence_threshold,

60
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -4,7 +4,7 @@ import time
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from typing import Final, cast # noqa: TID251 # untyped callback payload requires explicit narrowing
from openai.types import Batch
@ -25,6 +25,7 @@ VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend"
END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend"
USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend"
TEAM_SPEND_CACHE_KEY_PREFIX: Final = "team_model_spend"
TEAM_MEMBER_SPEND_CACHE_KEY_PREFIX: Final = "team_member_model_spend"
_SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
{
@ -32,6 +33,7 @@ _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX,
Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX,
Litellm_EntityType.TEAM: TEAM_SPEND_CACHE_KEY_PREFIX,
Litellm_EntityType.TEAM_MEMBER: TEAM_MEMBER_SPEND_CACHE_KEY_PREFIX,
}
)
@ -45,6 +47,7 @@ _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType(
Litellm_EntityType.USER: "user_model_budget_start_time",
Litellm_EntityType.END_USER: "end_user_budget_start_time",
Litellm_EntityType.TEAM: "team_model_budget_start_time",
Litellm_EntityType.TEAM_MEMBER: "team_member_model_budget_start_time",
}
)
@ -63,6 +66,10 @@ class ResolvedModelBudget:
budget_config: BudgetConfig
def team_member_budget_entity_id(user_id: str, team_id: str) -> str:
return f"{len(user_id)}:{user_id}:{team_id}"
def model_budget_spend_cache_key(
entity_type: Litellm_EntityType,
entity_id: str | None,
@ -290,6 +297,18 @@ def _resolve_entity_model_budgets(
)
def _metadata_value(metadata: object, key: str) -> object | None:
metadata_mapping: Final[Mapping[str, object]] = cast( # cast-ok: callback metadata has string keys
Mapping[str, object], metadata
)
return metadata_mapping.get(key)
def _metadata_string_value(metadata: object, key: str) -> str | None:
value: Final = _metadata_value(metadata, key)
return value if isinstance(value, str) else None
class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
"""
Handles budgets for model + virtual key
@ -410,6 +429,24 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}",
)
@with_service_target("model_budgets")
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool:
return await self._is_entity_within_model_budget(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=team_member_budget_entity_id(user_id=user_id, team_id=team_id),
model_max_budget=team_member_model_max_budget,
model=model,
exceeded_message=(
f"LiteLLM Team Member: user={user_id}, team={team_id}, exceeded budget for model={model}"
),
)
async def _is_entity_within_model_budget(
self,
entity_type: Litellm_EntityType,
@ -523,6 +560,19 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
key_model_max_budget: Final = _metadata.get("user_api_key_model_max_budget")
team_member_user_id: Final = _metadata_string_value(
payload_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
"user_api_key_user_id",
)
team_member_team_id: Final = _metadata_string_value(
payload_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
"user_api_key_team_id",
)
team_member_entity_id: Final = (
team_member_budget_entity_id(user_id=team_member_user_id, team_id=team_member_team_id)
if team_member_user_id is not None and team_member_team_id is not None
else None
)
entity_budgets: Final = (
(
Litellm_EntityType.KEY,
@ -553,6 +603,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
standard_logging_payload.get("end_user") or payload_metadata.get("user_api_key_end_user_id"),
_metadata.get("user_api_key_end_user_model_max_budget"),
),
(
Litellm_EntityType.TEAM_MEMBER,
team_member_entity_id,
_metadata_value(
_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
"user_api_key_team_member_model_max_budget",
),
),
)
resolved_budgets: Final = _resolve_entity_model_budgets(model=model, entity_budgets=entity_budgets)

42
litellm/proxy/litellm_pre_call_utils.py
View file

@ -227,6 +227,7 @@ from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
CAPTURE_MESSAGE_CONTENT_VAR,
CustomPricingLiteLLMParams,
LlmProviders,
ProviderSpecificHeader,
@ -975,6 +976,8 @@ def convert_key_logging_metadata_to_callback(
team_callback_settings_obj.callbacks.append(data.callback_name)
for var, value in data.callback_vars.items():
if var == CAPTURE_MESSAGE_CONTENT_VAR:
continue
# New Relic routing reads these from the trusted-vars overlay with no
# callback-name check, so scope them to the newrelic entry: a team that
# put newrelic_* under a different callback never asked for New Relic and
@ -1162,7 +1165,6 @@ def resolve_tenant_otel_destinations(
span still routes to the tenant's credentials the way it did then.
"""
from litellm.integrations.otel.model.config import is_otel_v2_enabled
from litellm.integrations.otel.presets.destinations import destination_for
if not is_otel_v2_enabled():
return ()
@ -1186,23 +1188,28 @@ def resolve_tenant_otel_destinations(
for name in dict.fromkeys(
callback.callback_name for callback in callbacks if callback.callback_type != "failure"
)
if (
destination := destination_for(
name,
_tenant_otel_params(
MappingProxyType(
{
var: value
for callback in callbacks
if callback.callback_name == name
for var, value in callback.callback_vars.items()
}
)
),
_tenant_service_name(user_api_key_dict),
if (destination := _tenant_destination(name, callbacks, _tenant_service_name(user_api_key_dict))) is not None
)
def _tenant_destination(
name: str, callbacks: Sequence[AddTeamCallback], service_name: str | None
) -> "OtelDestination | None":
from litellm.integrations.otel.model.config import parse_capture_message_content
from litellm.integrations.otel.presets.destinations import destination_for
callback_vars: Final = MappingProxyType(
dict(
itertools.chain.from_iterable(
callback.callback_vars.items() for callback in callbacks if callback.callback_name == name
)
)
is not None
)
return destination_for(
name,
_tenant_otel_params(callback_vars),
service_name,
parse_capture_message_content(callback_vars.get(CAPTURE_MESSAGE_CONTENT_VAR)),
)
@ -2508,6 +2515,9 @@ async def add_litellm_data_to_request(
data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget
user_model_budget: Final = user_api_key_dict.user_model_max_budget
data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param
data[_metadata_variable_name][ # rebind-ok: adds member budget for spend tracking
"user_api_key_team_member_model_max_budget"
] = user_api_key_dict.team_member_model_max_budget
data[_metadata_variable_name].update(carried_budget_metadata(user_api_key_dict))
data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)

9
litellm/proxy/management_endpoints/budget_management_endpoints.py
View file

@ -19,6 +19,8 @@ from typing import Final
from fastapi import APIRouter, Depends, HTTPException
from litellm._internal_context import service_target
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
@ -141,7 +143,7 @@ async def update_budget(
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(
@ -199,6 +201,11 @@ async def update_budget(
where={"budget_id": budget_obj.budget_id},
data=budget_obj_jsonified,
)
with service_target("team_member_budgets"):
try:
await user_api_key_cache.async_delete_cache(key=f"team_member_default_budget:{budget_obj.budget_id}")
except Exception: # noqa: BLE001 # Cache failure must not fail a committed budget update.
verbose_proxy_logger.warning("Failed to evict the team member default budget cache", exc_info=True)
return response

55
litellm/proxy/management_endpoints/common_utils.py
View file

@ -1,3 +1,4 @@
import json
import math
from collections.abc import Mapping
from types import MappingProxyType
@ -580,21 +581,31 @@ MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType(
"allowed_models": "allowed_models",
"temp_budget_increase": "temp_budget_increase",
"temp_budget_expiry": "temp_budget_expiry",
"model_max_budget": "model_max_budget",
}
)
def _prisma_value(value: object) -> object:
_JSON_BUDGET_COLUMNS: Final = frozenset({"model_max_budget"})
def _prisma_value(column: str, value: object) -> object:
if column in _JSON_BUDGET_COLUMNS:
return json.dumps(value if value is not None else {})
return list(value) if isinstance(value, tuple) else value
def _prisma_payload(data: Mapping[str, object]) -> dict[str, object]:
return {column: _prisma_value(column, value) for column, value in data.items()}
def member_budget_patch(source: BaseModel) -> Mapping[str, object]:
"""Map the per-member limit fields a request actually set to their budget-table
columns (merge-patch: a sent value updates, an explicit null clears, an absent
field is left untouched)."""
provided: Final = source.model_dump(exclude_unset=True)
return {
column: _prisma_value(provided[request_field])
column: provided[request_field]
for request_field, column in MEMBER_BUDGET_PATCH_FIELDS.items()
if request_field in provided
}
@ -603,7 +614,7 @@ def member_budget_patch(source: BaseModel) -> Mapping[str, object]:
def _is_set_budget_value(value: object) -> bool:
if value is None:
return False
if isinstance(value, list) and len(value) == 0:
if isinstance(value, (list, dict)) and len(value) == 0:
return False
return True
@ -645,10 +656,26 @@ async def upsert_budget_and_membership(
A patch that only touches the temporary budget pair never copies permanent
limits into a new row, so the member keeps inheriting the live team default.
"""
if not budget_patch:
ignore_default_model_budget_clear: Final = (
existing_budget_id is not None
and existing_budget_id == team_default_budget_id
and "model_max_budget" in budget_patch
and (
budget_patch["model_max_budget"] is None
or (isinstance(budget_patch["model_max_budget"], Mapping) and not budget_patch["model_max_budget"])
)
)
effective_patch: Final = MappingProxyType(
{
field: value
for field, value in budget_patch.items()
if not (ignore_default_model_budget_clear and field == "model_max_budget")
}
)
if not effective_patch:
return
write_data: Final = dict(budget_patch)
write_data: Final = dict(effective_patch)
if "budget_duration" in write_data:
duration: Final = write_data["budget_duration"]
write_data["budget_reset_at"] = (
@ -675,7 +702,7 @@ async def upsert_budget_and_membership(
return
await tx.litellm_budgettable.update(
where={"budget_id": existing_budget_id},
data={"updated_by": user_api_key_dict.user_id or "", **write_data},
data={"updated_by": user_api_key_dict.user_id or "", **_prisma_payload(write_data)},
)
return
@ -689,14 +716,24 @@ async def upsert_budget_and_membership(
create_data: Final[dict[str, object]] = { # mutable-ok: Prisma create payloads are dict-shaped
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
"model_max_budget": {},
**MappingProxyType(
{f: source[f] for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS if _is_set_budget_value(source.get(f))}
{
f: source[f]
for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS
if (
f != "model_max_budget"
or team_default_budget_id is None
or source.get("budget_id") != team_default_budget_id
)
and _is_set_budget_value(source.get(f))
}
),
**write_data,
}
# Restarting an inherited window on an unrelated edit hands the member a free period.
carried: Final = source.get("budget_reset_at") if "budget_duration" not in budget_patch else None
carried: Final = source.get("budget_reset_at") if "budget_duration" not in effective_patch else None
if carried is not None:
create_data["budget_reset_at"] = carried
if create_data.get("budget_reset_at") is None:
@ -708,7 +745,7 @@ async def upsert_budget_and_membership(
return
new_budget: Final = await tx.litellm_budgettable.create(
data=create_data,
data=_prisma_payload(create_data),
include={"team_membership": True},
)
await tx.litellm_teammembership.upsert(

5
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -8040,7 +8040,10 @@ def _validate_key_alias_format(key_alias: str | None) -> None:
raise_if_unsafe_secret_name(key_alias)
except ValueError:
raise ProxyException(
message="Invalid key_alias",
message=(
"Invalid key_alias: must not contain control characters (such as tab or newline) "
+ 'or a ".." path segment'
),
type=ProxyErrorTypes.bad_request_error,
param="key_alias",
code=400,

13
litellm/proxy/management_endpoints/team_callback_endpoints.py
View file

@ -31,8 +31,10 @@ from litellm.proxy._types import (
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.callback_config_validation import (
callback_config_error,
conflicting_capture_error,
conflicting_span_scope_error,
cross_entry_family_error,
stored_capture_entries,
)
from litellm.proxy.common_utils.callback_utils import ( # noqa: F401 # legacy module exports
_CALLBACK_VAR_ENCRYPTED_PREFIX, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -288,6 +290,9 @@ async def add_team_callbacks(
- langfuse_host: The host for the Langfuse callback
- langfuse_environment: The tracing environment for the Langfuse callback (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT)
- langfuse_span_scope: For langfuse_otel, "full" (default) sends the whole request trace, "llm_only" sends only the model-call spans
- capture_message_content: For OTel v2 callbacks, "no_content" or "span_only".
An explicit value overrides the global setting for this destination; omitted follows the proxy's own callback for that backend, else the global setting.
"span_only" puts prompt and response content on the destination's spans, "no_content" leaves it out
- gcs_bucket_name: The name of the GCS bucket
- gcs_path_service_account: The path to the GCS service account
- langsmith_api_key: The API key for the Langsmith callback
@ -351,11 +356,19 @@ async def add_team_callbacks(
# Decrypted, because the checks compare the incoming values against
# the stored ones and the credentials are encrypted at rest.
decrypted_logging: Final = decrypt_callback_vars(team_metadata).get("logging")
stored_capture: Final = stored_capture_entries(decrypted_logging)
stored_entries: Final = decrypted_logging if isinstance(decrypted_logging, list) else ()
stored_entry_vars: Final = [entry.get("callback_vars") or {} for entry in stored_entries]
scope_error: Final = conflicting_span_scope_error(data.callback_vars, stored_entry_vars)
if scope_error is not None:
raise _callback_config_error(scope_error)
capture_error: Final = conflicting_capture_error(
data.callback_name,
data.callback_vars,
stored_capture,
)
if capture_error is not None:
raise _callback_config_error(capture_error)
# One entry has to own a credential family end to end. The entries are
# flattened into one dict before a request reads them, so an entry
# naming only a destination would pair with a key written on another

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