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

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

Browse source
This commit is contained in:
mrinal 2026-10-10 23:22:27 +00:00
commit cf5832bf84
552 changed files with 46583 additions and 22508 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

436
.circleci/config.yml
View file

@ -243,7 +243,7 @@ commands:
mkdir -p "$HOME/.local/bin"
mv /tmp/codecov "$HOME/.local/bin/codecov"
start_openai_record_replay_proxy:
description: "Start the record/replay proxy (tests/_openai_record_replay_proxy.py) on host port 8090 and wait until healthy. Models whose api_base points here replay recorded provider responses, so the E2E run neither pays for nor depends on the live provider. The default upstream is OpenAI; a non-OpenAI model must point its api_base at /__recorder_upstream/<host>/ so the recorder forwards there instead of defaulting to OpenAI. Run after uv deps are synced."
description: "Start the record/replay proxy (tests/_openai_record_replay_proxy.py) on host port 8090 and wait until healthy. Models whose api_base points here replay Redis cache hits. Cache misses call the live provider and store its response, so these legacy E2E jobs still require provider availability and may incur provider charges. The default upstream is OpenAI; a non-OpenAI model must point its api_base at /__recorder_upstream/<host>/ so the recorder forwards there instead of defaulting to OpenAI. Run after uv deps are synced."
steps:
- run:
name: Start record/replay proxy
@ -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
@ -2812,11 +2546,72 @@ jobs:
docker logs redis-cache > test-results/services-<< parameters.suite >>-<< parameters.mode >>/redis.log 2>&1 || true
docker rm -f postgres-db redis-cache
test -z "$(docker ps -aq --filter name=postgres-db --filter name=redis-cache)"
- when:
condition:
not:
equal: [1, << parameters.parallelism >>]
steps:
- run:
name: Prepare this shard's execution record
command: |
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-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.mode >>
- store_test_results:
path: test-results
- store_artifacts:
path: test-results
integration_shard_coverage:
parameters:
suite:
type: enum
enum: [management, accounting, database, extensions, providers, mcp]
mode:
type: enum
enum: [standard, replica]
default: standard
shards:
type: integer
machine:
image: ubuntu-2204:2024.04.1
resource_class: medium
working_directory: ~/project
steps:
- setup_litellm_test_deps
- attach_workspace:
at: /tmp/integration-shards
- run:
name: Verify every << parameters.suite >> case ran exactly once
command: |
.venv/bin/python -m tests.integration._support.sharding /tmp/integration-shards/<< parameters.suite >>-<< parameters.mode >> \
--shards << parameters.shards >>
- when:
condition:
equal: [mcp, << parameters.suite >>]
steps:
- run:
name: Combine coverage from every MCP shard
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-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:
path: test-results/integration-mcp/coverage
routing_parity:
parameters:
suite:
@ -2898,6 +2693,9 @@ jobs:
default: ""
timeout_minutes:
type: integer
workers:
type: integer
default: 1
machine:
image: ubuntu-2204:2024.04.1
resource_class: large
@ -2925,12 +2723,16 @@ jobs:
command: |
mkdir -p test-results
coverage_args=()
worker_args=()
if [ "<< parameters.workers >>" -gt 1 ]; then
worker_args=(-n << parameters.workers >> --dist=worksteal)
fi
if [ -n "<< parameters.coverage_flag >>" ]; then
coverage_args=(--cov=./litellm --cov-report=xml:coverage.xml)
fi
timeout --signal=TERM << parameters.timeout_minutes >>m uv run --no-sync pytest \
<< parameters.test_path >> -vv --tb=short --durations=10 -o junit_family=xunit1 \
--junitxml=test-results/junit.xml "${coverage_args[@]}"
--junitxml=test-results/junit.xml "${coverage_args[@]}" "${worker_args[@]}"
- when:
condition: << parameters.coverage_flag >>
steps:
@ -3107,21 +2909,93 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, sdk, cost, security, browser]
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: 4
- integration_shard_coverage:
name: integration-management-coverage
suite: management
shards: 4
requires: [integration-management]
- integration_contracts:
name: integration-accounting
suite: accounting
parallelism: 4
- integration_shard_coverage:
name: integration-accounting-coverage
suite: accounting
shards: 4
requires: [integration-accounting]
- integration_contracts:
name: integration-providers
suite: providers
parallelism: 8
- integration_shard_coverage:
name: integration-providers-coverage
suite: providers
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
- integration_shard_coverage:
name: integration-mcp-coverage
suite: mcp
shards: 2
requires:
- integration-mcp
- integration_contracts:
name: integration-extensions
suite: extensions
parallelism: 4
parallelism: 8
- integration_shard_coverage:
name: integration-extensions-coverage
suite: extensions
shards: 8
requires:
- integration-extensions
- integration_contracts:
name: integration-<< matrix.suite >>-replica
matrix:
parameters:
suite: [management, database]
mode: [replica]
name: integration-management-replica
suite: management
mode: replica
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
@ -3138,6 +3012,7 @@ workflows:
test_path: tests/proxy_migration_tests
seed: false
timeout_minutes: 20
workers: 2
- postgres_suite:
name: roi-database
test_path: tests/integration/database/test_roi_observed.py
@ -3165,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
@ -3218,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

17
.circleci/scripts/run_integration.sh
View file

@ -134,6 +134,9 @@ upstream_pid=$!
if [ "$suite" = cost ]; then
export INTEGRATION_WORKERS=8
fi
if [ "$suite" = security ]; then
export INTEGRATION_WORKERS=4
fi
if [ "$suite" = mcp ]; then
export INTEGRATION_WORKERS=4 INTEGRATION_COVERAGE=1
fi
@ -213,13 +216,11 @@ if [ "$suite" = browser ]; then
exit 0
fi
node_files=()
shard_arguments=()
if [ "${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"
.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" \
@ -232,11 +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="${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" "${node_files[@]}"
.venv/bin/python tests/integration/run.py "$suite" --results "$results" \
"${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

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

@ -9,6 +9,8 @@ import sys
import warnings
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from functools import lru_cache
from types import MappingProxyType
from typing import Final
import yaml
@ -316,6 +318,7 @@ def _keyword_terms(
def _slices() -> tuple[Slice, ...]:
if not CIRCLECI_CONFIG.exists():
return ()
test_files: Final = frozenset(_test_files())
jobs: Final = yaml.safe_load(CIRCLECI_CONFIG.read_text()).get("jobs", {})
return tuple(
Slice(job=job, globs=globs, named=named, required=required, excluded=excluded, understood=understood)
@ -323,7 +326,7 @@ def _slices() -> tuple[Slice, ...]:
for text in ("\n".join(_strings(body)),)
if "pytest" in text
for globs in (tuple(GLOB_CALL_RE.findall(text)),)
for named in (frozenset(TEST_TOKEN_RE.findall(text)) & frozenset(_test_files()),)
for named in (frozenset(TEST_TOKEN_RE.findall(text)) & test_files,)
for required, excluded, understood in (
_keyword_terms(tuple(KEYWORD_RE.findall(text)), attributable=len(globs) < 2),
)
@ -363,6 +366,7 @@ def _workflow_named_tokens() -> frozenset[str]:
def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
matchable_names: Final = lru_cache(maxsize=None)(_matchable_names)
slices: Final = _slices()
named_by_workflow: Final = _workflow_named_tokens()
globbed: Final = tuple(
@ -376,7 +380,7 @@ def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
for path in globbed
if not allowlist.covers_test(path)
and not any(_token_covers(token, path) for token in named_by_workflow)
and not any(slice_.claims(path, _matchable_names(path)) for slice_ in slices)
and not any(slice_.claims(path, matchable_names(path)) for slice_ in slices)
)

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: |

10
.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
@ -927,6 +929,7 @@ jobs:
timeout-minutes: 3
with:
persist-credentials: false
fetch-depth: ${{ matrix.shard == 'db-and-spend' && '0' || '1' }}
- name: Detect relevant changes
id: changes
@ -959,11 +962,13 @@ jobs:
with:
artifact: ${{ needs.rust-bridge.outputs.artifact }}
- name: Install dependencies
- name: Install dependencies and check schema
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
env:
RUST_BRIDGE_ARTIFACT: ${{ needs.rust-bridge.outputs.artifact }}
UNIT_SHARD: ${{ matrix.shard }}
SCHEMA_BASE_REF: origin/${{ github.event.repository.default_branch }}
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
if [ -z "$RUST_BRIDGE_ARTIFACT" ]; then
@ -975,6 +980,9 @@ jobs:
uv run --no-sync python -c "import importlib.metadata; import litellm.rust_bridge._native; print(importlib.metadata.version('litellm'))"
fi
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
if [ "$UNIT_SHARD" = "db-and-spend" ]; then
uv run --no-sync python scripts/check_schema_compatibility.py --base-ref "$SCHEMA_BASE_REF"
fi
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'

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) | ✅ | ✅ | ✅ | | | | | | | |

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

@ -1,4 +1,5 @@
import random
import re
import time
from collections.abc import Generator, Mapping
from contextlib import contextmanager
@ -7,25 +8,50 @@ from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.prisma_toolchain import MIGRATION_LOCK_TIMEOUT_ENV_VAR, migration_lock_timeout
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
migration_ddl_lock_timeout,
migration_lock_timeout,
)
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
_LOCK_TIMEOUT_PINNED_RE: Final = re.compile(r"(?:-c\s*|--)lock_timeout=")
if TYPE_CHECKING:
import psycopg
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
database_url: Final = environment.get("DATABASE_URL")
direct_url: Final = environment.get("DIRECT_URL")
if not database_url or not direct_url:
return environment
def _migration_url(database_url: str, direct_url: "str | None") -> str:
if not direct_url:
return database_url
schema: Final = next((value for key, value in parse_qsl(urlsplit(database_url).query) if key == "schema"), "public")
direct: Final = urlsplit(direct_url)
parameters: Final = tuple((key, value) for key, value in parse_qsl(direct.query) if key != "schema")
return urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema)))))
def _with_ddl_lock_timeout(url: str) -> str:
split: Final = urlsplit(url)
pairs: Final = tuple(parse_qsl(split.query))
if any(key == "pgbouncer" and value == "true" for key, value in pairs):
return url
existing: Final = next((value for key, value in pairs if key == "options"), "")
if _LOCK_TIMEOUT_PINNED_RE.search(existing):
return url
timeout_ms: Final = max(1, int(migration_ddl_lock_timeout() * 1000))
options: Final = f"{existing} -c lock_timeout={timeout_ms}".strip()
rewritten: Final = tuple((key, value) for key, value in pairs if key != "options") + (("options", options),)
return urlunsplit(split._replace(query=urlencode(rewritten)))
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
database_url: Final = environment.get("DATABASE_URL")
if not database_url:
return environment
migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
return {
**environment,
"DATABASE_URL": urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema))))),
"DATABASE_URL": _with_ddl_lock_timeout(migration_url),
}

7
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
View file

@ -60,11 +60,13 @@ PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_MIGRATE_DEPLOY_TIMEOUT"
MIGRATION_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_LOCK_TIMEOUT"
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR = "LITELLM_MIGRATION_DDL_LOCK_TIMEOUT"
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0
DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT = 10.0
BOOTSTRAP_ARG = "--version"
PRISMA_CONSOLE_SCRIPT = "prisma"
@ -111,6 +113,11 @@ def migration_lock_timeout() -> float:
return _timeout_from_env(MIGRATION_LOCK_TIMEOUT_ENV_VAR, 600.0)
def migration_ddl_lock_timeout() -> float:
"""Seconds one migration statement may wait on a table lock before Postgres cancels it."""
return _timeout_from_env(MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR, DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT)
def prisma_bootstrap_timeout() -> float:
"""Seconds the one-time Node toolchain install may run for."""
return _timeout_from_env(

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

@ -18,6 +18,7 @@ from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR,
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
ensure_prisma_toolchain,
@ -62,6 +63,7 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
_MIGRATION_LOCK_TIMEOUT_MARKER = "canceling statement due to lock timeout"
INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
_INVALID_LITELLM_INDEXES_SQL: Final = (
@ -1162,9 +1164,11 @@ class ProxyExtrasDBManager:
raise RuntimeError(
f"Database migration failed after {MAX_MIGRATE_DEPLOY_ATTEMPTS} "
"attempts that made no progress (timeouts or deadlock retries). Check database connectivity, "
"load, and _prisma_migrations ledger state, and raise "
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out."
"attempts that made no progress (timeouts, lock timeouts or deadlock retries). "
"Check database connectivity, load, long-running transactions on the migrated tables, "
"and _prisma_migrations ledger state. Raise "
f"{PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR} if the attempts timed out, or "
f"{MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if they timed out waiting for a table lock."
)
finally:
os.chdir(original_dir)
@ -1225,6 +1229,13 @@ class ProxyExtrasDBManager:
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
return budget.spend()
if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs:
logger.info(
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
return budget.spend()
if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
logger.info(
"Migration %s started at %s was already rolled back or completed by a concurrent "
@ -1270,6 +1281,16 @@ class ProxyExtrasDBManager:
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
return budget.spend()
if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
ProxyExtrasDBManager._log_migration_lock_holders(migration_name)
logger.warning(
"Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying. "
f"Raise {MIGRATION_DDL_LOCK_TIMEOUT_ENV_VAR} if the database needs longer.",
migration_name,
)
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
return budget.spend()
raise RuntimeError(
"Database migration failed and cannot be auto-recovered. "
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
@ -1282,6 +1303,13 @@ class ProxyExtrasDBManager:
)
return budget.spend()
if _MIGRATION_LOCK_TIMEOUT_MARKER in stderr:
logger.info(
"Waiting for the advisory lock held by another Prisma migration; "
"contention does not spend a migration failure attempt"
)
return budget.after_contention(attempt_seconds)
if "P1002" in stderr and "advisory lock" in stderr:
logger.info(
"Waiting for the advisory lock held by another Prisma migration; "
@ -1294,6 +1322,68 @@ class ProxyExtrasDBManager:
f"Manual intervention required.\n\nPrisma error:\n{stderr}"
) from error
@staticmethod
def _log_migration_lock_holders(migration_name: str) -> None:
"""Best-effort log of the sessions holding locks on the tables a timed-out migration touches."""
try:
migration_sql: Final = (
Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" / migration_name / "migration.sql"
).read_text()
except OSError:
return
relations: Final = sorted(set(re.findall(r'"(LiteLLM_\w+)"', migration_sql)))
database_url: Final = os.getenv("DATABASE_URL")
if not relations or not database_url:
return
try:
import psycopg
except ImportError:
return
try:
with psycopg.connect(
ProxyExtrasDBManager._strip_prisma_query_params(database_url),
connect_timeout=10,
autocommit=True,
options="-c statement_timeout=5000",
) as conn:
rows: Final = conn.execute(
"SELECT l.pid, c.relname, l.mode, a.state, a.application_name, "
"date_trunc('second', now() - a.xact_start) "
"FROM pg_locks l "
"JOIN pg_class c ON c.oid = l.relation "
"JOIN pg_namespace n ON n.oid = c.relnamespace "
"JOIN pg_stat_activity a ON a.pid = l.pid "
"WHERE l.granted AND n.nspname = %s AND c.relname = ANY(%s) AND l.pid <> pg_backend_pid() "
"ORDER BY a.xact_start NULLS LAST "
"LIMIT 10",
(
ProxyExtrasDBManager._prisma_schema_param(database_url) or "public",
relations,
),
).fetchall()
except psycopg.Error:
return
if not rows:
logger.warning(
"Migration %s timed out waiting for a lock, but no current lock holder "
"was found on %s; the holder has likely since committed or rolled back",
migration_name,
", ".join(relations),
)
return
for pid, relname, mode, state, application_name, age in rows:
logger.warning(
"Migration %s timed out waiting for a lock on %s held by pid %s "
"(mode %s, state %s, application %s, transaction open for %s)",
migration_name,
relname,
pid,
mode,
state,
application_name,
age,
)
@staticmethod
def _mark_migration_applied(name: str) -> None:
"""Roll a failed ledger row back if it is still there, then mark it applied."""

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

8
litellm/caching/redis_cache.py
View file

@ -797,6 +797,8 @@ class RedisCache(BaseCache):
def init_async_client(
self,
) -> async_redis_client | async_redis_cluster_client:
from redis.asyncio import RedisCluster
from litellm import in_memory_llm_clients_cache
from .._redis import get_redis_async_client, get_redis_connection_pool
@ -809,7 +811,11 @@ class RedisCache(BaseCache):
# Create new connection pool and client for current event loop
self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs)
redis_async_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs)
in_memory_llm_clients_cache.set_cache(key=cache_key, value=redis_async_client)
in_memory_llm_clients_cache.set_cache(
key=cache_key,
value=redis_async_client,
litellm_owned_client=isinstance(redis_async_client, RedisCluster),
)
self.redis_async_client = redis_async_client
return redis_async_client

18
litellm/caching/redis_cluster_cache.py
View file

@ -11,7 +11,7 @@ from litellm.caching.redis_cache import RedisCache
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from redis.asyncio import Redis, RedisCluster
from redis.asyncio import Redis
from redis.asyncio.client import Pipeline
pipeline = Pipeline
@ -26,22 +26,6 @@ else:
class RedisClusterCache(RedisCache):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.redis_async_redis_cluster_client: RedisCluster | None = None
self.redis_sync_redis_cluster_client: RedisCluster | None = None
def init_async_client(self):
from redis.asyncio import RedisCluster
from .._redis import get_redis_async_client
if self.redis_async_redis_cluster_client:
return self.redis_async_redis_cluster_client
_redis_client: Final = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs)
if isinstance(_redis_client, RedisCluster):
self.redis_async_redis_cluster_client = _redis_client
return _redis_client
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""

7
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
@ -387,6 +388,8 @@ BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime"
BEDROCK_REALTIME_SDK_SUPPORTED_RANGE: Final = ">=0.10.0,<0.12.0"
CLIENT_REQUESTED_MODEL_SCOPE_KEY: Final = "litellm.client_requested_model"
MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolved"
MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body"
MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS: Final = 5.0
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
@ -823,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
@ -1967,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
@ -2215,6 +2220,8 @@ USAGE_TOP_API_KEYS_DEFAULT: Final[int] = 100
USAGE_TOP_API_KEYS_MAX: Final[int] = 1000
USAGE_KEY_PAGE_DEFAULT: Final[int] = 50
USAGE_KEY_PAGE_MAX: Final[int] = 100
USAGE_USER_PAGE_DEFAULT: Final[int] = 50
USAGE_USER_PAGE_MAX: Final[int] = 100
USAGE_KEY_SEARCH_DEFAULT: Final[int] = 100
USAGE_KEY_SEARCH_MAX: Final[int] = 100
USAGE_MODEL_TOP_KEYS_DEFAULT: Final[int] = 5

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/budget_alert_types.py
View file

@ -1,7 +1,7 @@
from abc import ABC, abstractmethod
from typing import Final, Literal
from litellm.proxy._types import CallInfo, Litellm_EntityType
from litellm.types.integrations.slack_alerting import CallInfo, Litellm_EntityType
class BaseBudgetAlertType(ABC):

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

@ -38,20 +38,20 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import (
CallInfo,
InvitationModel,
InvitationNew,
Litellm_EntityType,
UserAPIKeyAuth,
VirtualKeyEvent,
WebhookEvent,
)
from litellm.repositories.table_repositories import InvitationLinkRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.types.integrations.slack_alerting import *
from litellm.types.integrations.slack_alerting import AlertType
from litellm.types.integrations.slack_alerting import (
AlertType,
CallInfo,
InvitationModel,
InvitationNew,
Litellm_EntityType,
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):

5
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -33,6 +33,10 @@ from litellm.integrations.datadog.datadog_mock_client import (
should_use_datadog_mock,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.llm_cost_calc.cache_tokens import (
extract_cache_creation_tokens,
extract_cache_read_tokens,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
handle_any_messages_to_chat_completion_str_messages_conversion,
@ -44,7 +48,6 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
from litellm.types.integrations.datadog_llm_obs import *
from litellm.types.utils import (
AUDIT_GUARDRAIL_FIELDS,

2
litellm/integrations/email_alerting.py
View file

@ -6,8 +6,8 @@ import os
from typing import Final
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.proxy._types import WebhookEvent
from litellm.repositories.team_repository import TeamRepository
from litellm.types.integrations.slack_alerting import WebhookEvent
# we use this for the email header, please send a test email if you change this. verify it looks good on email
LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"

4
litellm/integrations/gcs_bucket/gcs_bucket.py
View file

@ -26,8 +26,8 @@ else:
class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
def __init__(self, bucket_name: str | None = None) -> None:
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import premium_user
from litellm.types.proxy.common_proxy_errors import CommonProxyErrors
self.batch_size = int(os.getenv("GCS_BATCH_SIZE", GCS_DEFAULT_BATCH_SIZE))
self.flush_interval = int(os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS))
@ -52,8 +52,8 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
#### ASYNC ####
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import premium_user
from litellm.types.proxy.common_proxy_errors import CommonProxyErrors
if premium_user is not True:
raise ValueError(

38
litellm/integrations/generic_api/generic_api_callback.py
View file

@ -10,7 +10,9 @@ import asyncio
import json
import os
import re
import time
import traceback
from collections.abc import Mapping
from typing import Final, Literal
import httpx
@ -21,6 +23,7 @@ from litellm._uuid import uuid
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
httpxSpecialProvider,
)
@ -104,6 +107,7 @@ class GenericAPILogger(CustomBatchLogger):
max_retries: int = 0,
retry_delay: float = 1.0,
timeout: float | httpx.Timeout | None = None,
async_httpx_client: AsyncHTTPHandler | None = None,
**kwargs,
):
"""
@ -151,7 +155,11 @@ class GenericAPILogger(CustomBatchLogger):
#########################################################
# Init httpx client
#########################################################
self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
self.async_httpx_client = (
async_httpx_client
if async_httpx_client is not None
else get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
)
endpoint = endpoint or os.getenv("GENERIC_LOGGER_ENDPOINT")
if endpoint is None:
raise ValueError(
@ -191,7 +199,7 @@ class GenericAPILogger(CustomBatchLogger):
self.flush_lock = asyncio.Lock()
super().__init__(**kwargs, flush_lock=self.flush_lock)
asyncio.create_task(self.periodic_flush())
self.log_queue: list[dict | StandardLoggingPayload] = []
self.log_queue: list[Mapping[str, object] | StandardLoggingPayload] = []
def _get_headers(self, headers: dict | None = None):
"""
@ -343,7 +351,14 @@ class GenericAPILogger(CustomBatchLogger):
except Exception as e:
verbose_logger.exception("Generic API Logger Error - %s\n%s", e, traceback.format_exc())
async def async_send_batch(self):
async def flush_queue(self) -> None:
await self.async_send_batch()
async def async_send_batch(self) -> None:
async with self.flush_lock:
await self._send_batch()
async def _send_batch(self) -> None:
"""
Sends the batch of messages to Generic API Endpoint
@ -352,18 +367,18 @@ class GenericAPILogger(CustomBatchLogger):
- ndjson: Sends logs as newline-delimited JSON
- single: Sends each log as individual HTTP request in parallel
"""
batch: Final = tuple(self.log_queue)
if not batch:
return
try:
if not self.log_queue:
return
verbose_logger.debug(
"Generic API Logger - about to flush %s events in '%s' format", len(self.log_queue), self.log_format
"Generic API Logger - about to flush %s events in '%s' format", len(batch), self.log_format
)
if self.log_format == "single":
# Send each log as individual HTTP request in parallel
tasks: Final = []
for log_entry in self.log_queue:
for log_entry in batch:
task = self._post_with_retries(data=safe_dumps(log_entry))
tasks.append(task)
@ -384,9 +399,9 @@ class GenericAPILogger(CustomBatchLogger):
else:
# Format the payload based on log_format
if self.log_format == "json_array":
data = safe_dumps(self.log_queue)
data = safe_dumps(batch)
elif self.log_format == "ndjson":
data = "\n".join(safe_dumps(log) for log in self.log_queue)
data = "\n".join(safe_dumps(log) for log in batch)
else:
raise ValueError(f"Unknown log_format: {self.log_format}")
@ -403,7 +418,8 @@ class GenericAPILogger(CustomBatchLogger):
except Exception as e:
verbose_logger.exception("Generic API Logger Error sending batch - %s\n%s", e, traceback.format_exc())
finally:
self.log_queue.clear()
del self.log_queue[: len(batch)]
self.last_flush_time = time.time()
def _get_v1_logging_payload(self, kwargs, response_obj, start_time, end_time) -> dict:
"""

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 {}
),

12
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,
@ -49,13 +50,9 @@ from litellm.litellm_core_utils.service_tier_utils import (
get_service_tier_from_standard_logging_payload,
)
from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.proxy._types import (
LiteLLM_DeletedVerificationToken,
LiteLLM_TeamTable,
LiteLLM_UserTable,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
from litellm.models.team import LiteLLM_TeamTable
from litellm.models.user import LiteLLM_UserTable
from litellm.models.verification_token import LiteLLM_DeletedVerificationToken
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.organization_repository import OrganizationRepository
@ -69,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

61
litellm/litellm_core_utils/cold_storage_handler.py Normal file
View file

@ -0,0 +1,61 @@
"""
This module is responsible for handling Getting/Setting the proxy server request from cold storage.
It allows fetching a dict of the proxy server request from s3 or GCS bucket.
"""
from typing import Final
import litellm
from litellm import _custom_logger_compatible_callbacks_literal
from litellm.integrations.custom_logger import CustomLogger
class ColdStorageHandler:
"""
This class is responsible for handling Getting/Setting the proxy server request from cold storage.
It allows fetching a dict of the proxy server request from s3 or GCS bucket.
The cold storage logger can be injected for testing; when omitted it is
resolved from the configured ``litellm.cold_storage_custom_logger``.
"""
def __init__(self, cold_storage_logger: CustomLogger | None = None):
self._injected_cold_storage_logger = cold_storage_logger
async def get_proxy_server_request_from_cold_storage_with_object_key(
self,
object_key: str,
) -> dict | None:
"""
Get the proxy server request from cold storage using the object key directly.
Args:
object_key: The S3/GCS object key to retrieve
Returns:
Optional[dict]: The proxy server request dict or None if not found
"""
custom_logger: Final = self._injected_cold_storage_logger or self._resolve_cold_storage_logger()
if custom_logger is None:
return None
return await custom_logger.get_proxy_server_request_from_cold_storage_with_object_key(
object_key=object_key,
)
def _resolve_cold_storage_logger(self) -> CustomLogger | None:
custom_logger_name: Final = self._select_custom_logger_for_cold_storage()
if custom_logger_name is None:
return None
return litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(custom_logger_name)
def _select_custom_logger_for_cold_storage(
self,
) -> _custom_logger_compatible_callbacks_literal | None:
cold_storage_custom_logger: Final[_custom_logger_compatible_callbacks_literal | None] = (
litellm.cold_storage_custom_logger
)
return cold_storage_custom_logger

27
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -7,6 +7,7 @@ from types import MappingProxyType
from typing import Final, Protocol, cast
import httpx
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import _ENABLE_SECRET_REDACTION, redact_string, verbose_logger
@ -256,6 +257,23 @@ class _ProviderHTTPException(Protocol):
llm_provider: str
_ERROR_BODY_ADAPTER: Final = TypeAdapter(dict[str, object])
def _openai_cyber_policy_body(original_exception: _ProviderHTTPException) -> Mapping[str, object] | None:
raw_body: Final[object] = getattr(original_exception, "body", None)
try:
body: Final = (
_ERROR_BODY_ADAPTER.validate_python(raw_body)
if raw_body is not None
else _ERROR_BODY_ADAPTER.validate_json(getattr(original_exception, "message", ""))
)
error: Final = _ERROR_BODY_ADAPTER.validate_python(body.get("error", body))
except ValidationError:
return None
return error if error.get("code") == "cyber_policy" else None
def _litellm_proxy_response(
original_exception: _ProviderHTTPException, custom_llm_provider: str
) -> httpx.Response | None:
@ -284,6 +302,15 @@ def _map_openai_exception(
extra_information: str,
) -> None:
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
if (cyber_policy_body := _openai_cyber_policy_body(original_exception)) is not None:
raise ContentPolicyViolationError(
message=f"ContentPolicyViolationError: {exception_provider} - {original_exception}",
llm_provider=custom_llm_provider,
model=model,
response=response,
litellm_debug_info=extra_information,
body=dict(cyber_policy_body),
)
# custom_llm_provider is openai, make it OpenAI
message = get_error_message(error_obj=original_exception)
if message is None:

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",

49
litellm/litellm_core_utils/llm_cost_calc/cache_tokens.py Normal file
View file

@ -0,0 +1,49 @@
from collections.abc import Mapping
from typing import Final
def extract_cache_read_tokens(usage_object: Mapping[str, object] | None) -> int:
"""Cache-read tokens from a logged usage object, whatever shape recorded them.
Anthropic writes a top-level ``cache_read_input_tokens``; OpenAI-compatible
providers (moonshotai, openai, deepseek, etc.) write
``prompt_tokens_details.cached_tokens``. This is the one owner of that
normalization: callers hand over the usage object rather than threading a
count that could disagree with it.
"""
if not usage_object:
return 0
explicit: Final = usage_object.get("cache_read_input_tokens")
if isinstance(explicit, (int, float)) and explicit:
return int(explicit)
details: Final = usage_object.get("prompt_tokens_details")
if not isinstance(details, Mapping):
return 0
cached: Final = details.get("cached_tokens")
return int(cached) if isinstance(cached, (int, float)) else 0
def extract_cache_creation_tokens(usage_object: Mapping[str, object] | None) -> int:
"""Cache-write tokens from a logged usage object, whatever shape recorded them.
Anthropic writes a top-level ``cache_creation_input_tokens``; OpenAI-compatible
providers (kimi-k2 etc.) write ``prompt_tokens_details.cache_write_tokens`` or
``prompt_tokens_details.cache_creation_tokens``.
"""
if not usage_object:
return 0
explicit: Final = usage_object.get("cache_creation_input_tokens")
if isinstance(explicit, (int, float)) and explicit:
return int(explicit)
details: Final = usage_object.get("prompt_tokens_details")
if not isinstance(details, Mapping):
return 0
written: Final = next(
(
value
for value in (details.get("cache_write_tokens"), details.get("cache_creation_tokens"))
if isinstance(value, (int, float)) and value
),
0,
)
return int(written)

96
litellm/litellm_core_utils/path_utils.py Normal file
View file

@ -0,0 +1,96 @@
"""
Safe filesystem path construction for user-controlled inputs.
Use safe_join() instead of os.path.join() whenever a path component
comes from user input (request parameters, uploaded filenames, etc.)
to prevent directory traversal attacks.
"""
import os
from typing import Final
def safe_join(base_dir: str, *parts: str) -> str:
"""
Join path components and verify the result stays within base_dir.
Resolves symlinks and '..' sequences, then checks the final path
is a descendant of base_dir. Raises ValueError if traversal is
detected.
Args:
base_dir: The trusted base directory.
*parts: User-controlled path components to append.
Returns:
The resolved absolute path as a string.
Raises:
ValueError: If the resolved path escapes base_dir.
"""
for part in parts:
if "\x00" in part:
raise ValueError("Path contains null byte")
base: Final = os.path.realpath(base_dir)
resolved: Final = os.path.realpath(os.path.join(base, *parts))
if not (resolved.startswith(base + os.sep) or resolved == base):
raise ValueError(f"Path {resolved!r} escapes base directory {base!r}")
return resolved
def try_safe_join(base_dir: str, *parts: str) -> str | None:
"""safe_join, with None instead of ValueError when the path escapes base_dir."""
try:
return safe_join(base_dir, *parts)
except ValueError:
return None
def is_within(path: str, base_dir: str) -> bool:
"""True when path, with symlinks resolved, is base_dir or sits inside it."""
base: Final = os.path.realpath(base_dir)
resolved: Final = os.path.realpath(path)
return resolved.startswith(base + os.sep) or resolved == base
def join_within(base_dir: str, *parts: str) -> str | None:
"""Join without following symlinks; None when the joined path leaves base_dir.
Only the supplied components are checked (``..`` and absolute parts are
rejected), so a symlink stored inside base_dir that points elsewhere is
still returned. Use safe_join when the target itself must stay inside.
"""
for part in parts:
if "\x00" in part:
return None
base: Final = os.path.normpath(os.path.abspath(base_dir))
joined: Final = os.path.normpath(os.path.join(base, *parts))
if not joined.startswith(base + os.sep):
return None
return joined
def safe_filename(filename: str) -> str:
"""
Extract a safe filename from a user-supplied path.
Strips all directory components (both Unix and Windows separators),
returning only the final name. Use this for uploaded file names
before writing to disk.
Args:
filename: User-supplied filename (may contain path separators).
Returns:
The basename only, with no directory components.
Raises:
ValueError: If the resulting filename is empty or contains null bytes.
"""
if "\x00" in filename:
raise ValueError("Filename contains null byte")
# Normalize backslash separators for cross-platform safety
name: Final = filename.replace("\\", "/").rsplit("/", 1)[-1]
if not name or name in (".", ".."):
raise ValueError("Empty or unsafe filename")
return name

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

427
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -3,12 +3,17 @@ import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast
from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict, Union, cast
from pydantic import BaseModel
from typing_extensions import ReadOnly, Required
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict
)
from litellm.types.llms.openai import (
ChatCompletionAssistantContentValue,
ChatCompletionAudioDelta,
@ -28,11 +33,13 @@ from litellm.types.utils import (
Delta,
Function,
FunctionCall,
LlmProviders,
ModelResponse,
ModelResponseStream,
PromptTokensDetailsWrapper,
ServerToolUse,
StreamingChoices,
TextCompletionResponse,
Usage,
)
from litellm.utils import print_verbose, token_counter
@ -1194,3 +1201,421 @@ def concatenate_base64_list(base64_strings: list[str]) -> str:
# Encode the concatenated bytes back to base64
return base64.b64encode(combined_bytes).decode("utf-8")
def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse:
id: Final = chunks[0]["id"]
object: Final = chunks[0]["object"]
created: Final = chunks[0]["created"]
model: Final = chunks[0]["model"]
system_fingerprint: Final = chunks[0].get("system_fingerprint", None)
finish_reason: Final = chunks[-1]["choices"][0]["finish_reason"]
logprobs: Final = chunks[-1]["choices"][0]["logprobs"]
content_list: Final = []
for chunk in chunks:
choices = chunk["choices"]
for choice in choices:
if choice is not None and hasattr(choice, "text") and choice.get("text") is not None:
_choice = choice.get("text")
content_list.append(_choice)
# Combine the "content" strings into a single string || combine the 'function' strings into a single string
combined_content: Final = "".join(content_list)
try:
prompt_tokens = token_counter(model=model, messages=messages)
except Exception: # don't allow this failing to block a complete streaming response from being returned
print_verbose("token_counter failed, assuming prompt tokens is 0")
prompt_tokens = 0
completion_tokens: Final = token_counter(
model=model,
text=combined_content,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
)
response: Final = {
"id": id,
"object": object,
"created": created,
"model": model,
"system_fingerprint": system_fingerprint,
"choices": [
{
"text": combined_content,
"index": 0,
"logprobs": logprobs,
"finish_reason": finish_reason,
}
],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
}
return TextCompletionResponse(**response)
_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value})
def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool:
if logging_obj is None:
return False
provider: Final[object] = logging_obj.model_call_details.get("custom_llm_provider")
return provider in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS
def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None:
usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None)
if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj):
return float(usage_cost)
if logging_obj is not None:
return None
provider_hint: Final = response._hidden_params.get( # pyright: ignore[reportPrivateUsage] # no public accessor
"custom_llm_provider"
)
try:
return litellm.completion_cost(completion_response=response, custom_llm_provider=provider_hint)
except Exception:
return _stream_builder_model_map_cost(response)
def _joined_streamed_citations(streamed_citations: "tuple[object, ...]") -> "list[object]":
if all(isinstance(citation, list) for citation in streamed_citations):
return list(streamed_citations)
return [list(streamed_citations)]
def _stream_builder_model_map_cost(response: ModelResponse) -> float | None:
model_name: Final = response.model
usage: Final = getattr(response, "usage", None)
if not model_name or not isinstance(usage, Usage):
return None
try:
prompt_cost, completion_tokens_cost = litellm.cost_per_token(model=model_name, usage_object=usage)
return prompt_cost + completion_tokens_cost
except Exception: # noqa: BLE001 # cost_per_token raises bare Exception for unpriceable models
return None
def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
response_cost: Final = _stream_builder_response_cost(response, logging_obj)
if response_cost is None:
return
hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor
hidden_params["response_cost"] = response_cost
def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
if logging_obj is None:
return
if isinstance(getattr(usage, "cost", None), (int, float)):
return
computed_cost: Final = logging_obj.response_cost_calculator(result=response)
if isinstance(computed_cost, (int, float)) and computed_cost > 0:
setattr(usage, "cost", computed_cost)
_NON_TEXT_DELTA_FIELDS: Final = (
"tool_calls",
"function_call",
"reasoning_content",
"thinking_blocks",
"annotations",
"audio",
"images",
"provider_specific_fields",
)
def _stream_choice_delta(choice: object) -> Mapping[str, object]:
delta: Final = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {})
if isinstance(delta, Mapping):
return delta
if isinstance(delta, BaseModel):
return delta.model_dump()
return {}
def _delta_carries_more_than_text(delta: Mapping[str, object]) -> bool:
return any(delta.get(field) is not None for field in _NON_TEXT_DELTA_FIELDS)
def _simple_text_part(choices: Sequence[object]) -> str | None:
deltas: Final = tuple(_stream_choice_delta(choice) for choice in choices)
if any(_delta_carries_more_than_text(delta) for delta in deltas):
return None
content: Final = deltas[0].get("content")
return content if isinstance(content, str) else ""
def stream_chunk_builder(
chunks: list,
messages: Sequence | None = None,
start_time=None,
end_time=None,
logging_obj: Optional["Logging"] = None,
count_prompt_tokens: Callable[[], int] | None = None,
) -> ModelResponse | TextCompletionResponse | None:
try:
if chunks is None:
raise litellm.APIError(
status_code=500,
message="Error building chunks for logging/streaming usage calculation",
llm_provider="",
model="",
)
if not chunks:
return None
processor: Final = ChunkProcessor(chunks, messages)
chunks = processor.chunks
### BASE-CASE ###
if len(chunks) == 0:
return None
## Route to the text completion logic
first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), None)
if first_chunk_with_choices is not None and isinstance(
first_chunk_with_choices["choices"][0], litellm.utils.TextChoices
): # route to the text completion logic
return stream_chunk_builder_text_completion(chunks=chunks, messages=messages)
model: Final = chunks[0]["model"]
# Initialize the response dictionary
response: Final = processor.build_base_response(chunks)
# Fast path for the common text-only streaming case:
# avoid repeated multi-pass list scans over chunks.
simple_content_parts: Final[list[str]] = []
is_simple_text_stream = True
for chunk in chunks:
if not chunk.get("choices"):
continue
if (part := _simple_text_part(chunk["choices"])) is None:
is_simple_text_stream = False
break
if part:
simple_content_parts.append(part)
if is_simple_text_stream:
if simple_content_parts:
response["choices"][0]["message"]["content"] = "".join(simple_content_parts)
completion_output = get_content_from_model_response(response)
usage = processor.calculate_usage(
chunks=chunks,
model=model,
completion_output=completion_output,
messages=messages,
reasoning_tokens=0,
count_prompt_tokens=count_prompt_tokens,
)
setattr(response, "usage", usage)
# Propagate provider_specific_fields from chunk hidden params when present.
for chunk in reversed(chunks):
if isinstance(chunk, dict):
hidden = chunk.get("_hidden_params")
else:
hidden = getattr(chunk, HIDDEN_PARAMS_ATTR, None)
if isinstance(hidden, dict) and "provider_specific_fields" in hidden:
response.hidden_params.setdefault("provider_specific_fields", {}).update(
hidden["provider_specific_fields"]
)
break
_stamp_streaming_usage_cost(usage, response, logging_obj)
_set_stream_builder_response_cost(response, logging_obj)
processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
return response
tool_call_chunks: Final = [
chunk
for chunk in chunks
if any(
"tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None
for choice in chunk.get("choices") or ()
)
]
if len(tool_call_chunks) > 0:
tool_calls_list: Final = processor.get_combined_tool_content(tool_call_chunks)
_choice = response.choices[0]
_choice.message.content = None
_choice.message.tool_calls = tool_calls_list
function_call_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "function_call" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["function_call"] is not None
]
if len(function_call_chunks) > 0:
_choice = response.choices[0]
_choice.message.content = None
_choice.message.function_call = processor.get_combined_function_call_content(function_call_chunks)
content_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["content"] is not None
]
if len(content_chunks) > 0:
response["choices"][0]["message"]["content"] = processor.get_combined_content(content_chunks)
thinking_blocks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "thinking_blocks" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["thinking_blocks"] is not None
]
if len(thinking_blocks) > 0:
response["choices"][0]["message"]["thinking_blocks"] = processor.get_combined_thinking_content(
thinking_blocks
)
reasoning_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "reasoning_content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["reasoning_content"] is not None
]
if len(reasoning_chunks) > 0:
response["choices"][0]["message"]["reasoning_content"] = processor.get_combined_reasoning_content(
reasoning_chunks
)
annotation_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "annotations" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["annotations"] is not None
]
if len(annotation_chunks) > 0:
# Merge annotations from ALL chunks — providers may spread
# them across multiple streaming chunks or send them only in
# the final chunk.
all_annotations: Final[list] = []
for ac in annotation_chunks:
all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
response["choices"][0]["message"]["annotations"] = all_annotations
audio_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "audio" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["audio"] is not None
]
if len(audio_chunks) > 0:
_choice = response.choices[0]
_choice.message.audio = processor.get_combined_audio_content(audio_chunks)
# Handle image chunks from models like gemini-2.5-flash-image
# See: https://github.com/BerriAI/litellm/issues/19478
image_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "images" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["images"] is not None
]
if len(image_chunks) > 0:
# Images come complete in a single chunk, collect all images from all chunks
all_images: Final = []
for chunk in image_chunks:
all_images.extend(chunk["choices"][0]["delta"]["images"])
response["choices"][0]["message"]["images"] = all_images
# Combine provider_specific_fields from streaming chunks (e.g., web_search_results, citations)
# See: https://github.com/BerriAI/litellm/issues/17737
provider_specific_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "provider_specific_fields" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["provider_specific_fields"] is not None
]
if len(provider_specific_chunks) > 0:
provider_field_dicts: Final = tuple(
fields
for chunk in provider_specific_chunks
for fields in (chunk["choices"][0]["delta"]["provider_specific_fields"],)
if isinstance(fields, dict)
)
streamed_citations: Final = tuple(
fields["citation"] for fields in provider_field_dicts if fields.get("citation") is not None
)
citation_fields: Final = (
{"citations": _joined_streamed_citations(streamed_citations)} if streamed_citations else {}
)
combined_provider_fields: Final = {
key: value
for fields in (citation_fields, *provider_field_dicts)
for key, value in fields.items()
if key != "citation"
}
if combined_provider_fields:
_choice = response.choices[0]
_choice.message.provider_specific_fields = combined_provider_fields
completion_output = get_content_from_model_response(response)
reasoning_tokens: Final = processor.count_reasoning_tokens(response)
usage = processor.calculate_usage(
chunks=chunks,
model=model,
completion_output=completion_output,
messages=messages,
reasoning_tokens=reasoning_tokens,
count_prompt_tokens=count_prompt_tokens,
)
setattr(response, "usage", usage)
# Propagate provider_specific_fields from the last chunk (contains provider
# metadata like traffic_type set during streaming)
for chunk in reversed(chunks):
if isinstance(chunk, dict):
hidden = chunk.get("_hidden_params")
else:
hidden = getattr(chunk, HIDDEN_PARAMS_ATTR, None)
if isinstance(hidden, dict) and "provider_specific_fields" in hidden:
response.hidden_params.setdefault("provider_specific_fields", {}).update(
hidden["provider_specific_fields"]
)
break
_stamp_streaming_usage_cost(usage, response, logging_obj)
_set_stream_builder_response_cost(response, logging_obj)
processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
return response
except Exception as e:
verbose_logger.exception("litellm.main.py::stream_chunk_builder() - Exception occurred - %s", e)
raise litellm.APIError(
status_code=500,
message="Error building chunks for logging/streaming usage calculation",
llm_provider="",
model="",
)

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):

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

@ -19,12 +19,12 @@ from litellm.llms.anthropic.prompt_cache_prediction import (
UnsupportedPredictionTarget,
cache_scope,
count_prompt_tokens,
lookup,
parse_prompt,
resolve_prediction_target,
supported_prediction_headers,
)
from litellm.proxy.common_utils.prompt_cache_pricing import price_cache_tokens
from litellm.proxy.hooks.prompt_cache_prediction import lookup
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,

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

@ -29,6 +29,10 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import (
LiteLLMAnthropicMessagesAdapter,
is_provider_native_tool_dict,
)
from litellm.llms.anthropic.pass_through.stream_assembly import (
build_complete_streaming_response,
build_usage_only_response_from_chunks,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
@ -47,9 +51,6 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
stream_item_fingerprint,
unappliable_request_rewrite,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
from litellm.types.llms.anthropic import (
AllAnthropicToolsValues,
AnthropicMessagesRequest,
@ -75,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)
@ -322,7 +323,7 @@ class AnthropicMessagesHandler(BaseTranslation):
if not chunks:
return None
try:
return AnthropicPassthroughLoggingHandler.build_usage_only_response_from_chunks(
return build_usage_only_response_from_chunks(
all_chunks=chunks,
model=str((request_data or {}).get("model") or ""),
)
@ -1226,7 +1227,7 @@ class AnthropicMessagesHandler(BaseTranslation):
has_ended: Final = self._check_streaming_has_ended(responses_so_far)
if has_ended:
# build the model response from the responses_so_far
built_response: Final = AnthropicPassthroughLoggingHandler.build_complete_streaming_response(
built_response: Final = build_complete_streaming_response(
all_chunks=responses_so_far,
litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj),
model="",

4
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"})
@ -412,7 +412,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
return None
if resolved_provider == "litellm_proxy":
return None
from litellm.main import responses_api_bridge_check
from litellm.responses.bridge_check import responses_api_bridge_check
web_search_options: Final = completion_kwargs.get("web_search_options")
tools: Final = completion_kwargs.get("tools")

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(

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

@ -45,7 +45,7 @@ _RESPONSES_API_PROVIDERS: Final = frozenset({"openai"})
def _bridges_to_responses_api(model: str, custom_llm_provider: str) -> bool:
from litellm.main import responses_api_bridge_check
from litellm.responses.bridge_check import responses_api_bridge_check
model_info, _ = responses_api_bridge_check(model=model, custom_llm_provider=custom_llm_provider)
return model_info.get("mode") == "responses"

519
litellm/llms/anthropic/pass_through/stream_assembly.py Normal file
View file

@ -0,0 +1,519 @@
import codecs
import json
from collections.abc import Mapping, Sequence
from typing import Final
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.chat.handler import (
ModelResponseIterator as AnthropicModelResponseIterator,
)
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.types.utils import (
Choices,
Message,
ModelResponse,
TextCompletionResponse,
)
def split_sse_chunk_into_events(chunk: str | bytes) -> list[str]:
"""
Split a chunk that may contain multiple SSE events into individual events.
SSE format: "event: type\ndata: {...}\n\n"
Multiple events in a single chunk are separated by double newlines.
Args:
chunk: Raw chunk string that may contain multiple SSE events
Returns:
List of individual SSE event strings (each containing "event: X\ndata: {...}")
"""
# Handle bytes input
if isinstance(chunk, bytes):
chunk = chunk.decode("utf-8")
# Split on double newlines to separate SSE events
# Filter out empty strings
events: Final = [event.strip() for event in chunk.split("\n\n") if event.strip()]
return events
def build_complete_streaming_response(
all_chunks: Sequence[str | bytes],
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
speed: str | None = None,
) -> ModelResponse | TextCompletionResponse | None:
"""
Builds complete response from raw Anthropic chunks.
Fast path: for the dominant case of a pure-text streaming response
(no tool_use / thinking / non-text content blocks), the long run of
``content_block_delta`` text deltas is collapsed into a single
equivalent SSE event before conversion. ``chunk_parser`` and
``stream_chunk_builder`` remain the single source of truth for chunk
shape, usage math and finish-reason mapping, so the rebuilt response
(and therefore the logged/billed payload) is identical -- this is
asserted by a parity test. Anything non-trivial falls back to the
unchanged legacy reconstruction.
Per-event Pydantic ``ModelResponseStream`` construction dominated
event-loop CPU under concurrent streaming; collapsing the homogeneous
text run removes O(num_output_tokens) of it.
"""
collapsed: Final = collapse_pure_text_chunks(all_chunks)
if collapsed is not None:
return build_complete_streaming_response_legacy(
all_chunks=collapsed,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
return build_complete_streaming_response_legacy(
all_chunks=all_chunks,
litellm_logging_obj=litellm_logging_obj,
model=model,
speed=speed,
)
# Anthropic SSE block/delta types that the fast path is NOT allowed to
# collapse -- their presence forces the unchanged legacy path so tool
# calls, thinking, citations, etc. keep byte-identical reconstruction.
_FAST_PATH_DISALLOWED_DELTA_TYPES: Final = frozenset(
{
"input_json_delta",
"thinking_delta",
"signature_delta",
"citations_delta",
}
)
def collapse_pure_text_chunks(
all_chunks: Sequence[str | bytes],
) -> list[str] | None:
"""
Return a new chunk list with the contiguous run of text-only
``content_block_delta`` events replaced by a single equivalent event,
or ``None`` if the stream is not a pure single-text-block response
(in which case the caller uses the legacy path unchanged).
Only ``message_start`` / ``content_block_start(text)`` /
``content_block_delta(text_delta)`` / ``content_block_stop`` /
``message_delta`` / ``message_stop`` / ``ping`` events are accepted.
Any other content-block type or delta type returns ``None``.
"""
normalized: Final[list[str]] = []
for raw in all_chunks:
line = raw.decode("utf-8") if isinstance(raw, bytes) else raw
for ev in line.split("\n\n"):
ev = ev.strip()
if ev:
normalized.append(ev)
text_block_indexes: Final[set] = set()
out: Final[list[str]] = []
pending_text: list[str] = []
pending_index: int | None = None
saw_any_text_delta = False
def flush() -> None:
nonlocal pending_text, pending_index
if pending_text:
merged: Final = {
"type": "content_block_delta",
"index": pending_index if pending_index is not None else 0,
"delta": {"type": "text_delta", "text": "".join(pending_text)},
}
out.append("data: " + json.dumps(merged))
pending_text = []
pending_index = None
for ev in normalized:
idx = ev.find("data:")
if idx == -1:
# Bare "event: <name>" line. The legacy converter turns this
# into an empty ModelResponseStream that contributes nothing
# to stream_chunk_builder. Drop the high-frequency interior
# markers (content_block_delta / ping); keep every other
# bare event line verbatim so chunk ordering and the
# load-bearing chunks[0] (event: message_start) are retained.
name = ev[len("event:") :].strip() if ev.startswith("event:") else ""
if name in ("content_block_delta", "ping"):
continue
flush()
out.append(ev)
continue
json_str = ev[idx + len("data:") :].strip()
try:
data = json.loads(json_str)
except (json.JSONDecodeError, ValueError):
return None
etype = data.get("type")
if etype == "content_block_start":
block = data.get("content_block") or {}
if block.get("type") != "text":
return None
text_block_indexes.add(data.get("index"))
flush()
out.append(ev)
elif etype == "content_block_delta":
delta = data.get("delta") or {}
dtype = delta.get("type")
if dtype in _FAST_PATH_DISALLOWED_DELTA_TYPES:
return None
if dtype != "text_delta":
return None
cur_index = data.get("index")
if cur_index not in text_block_indexes:
return None
# Defensive: Anthropic sends blocks strictly sequentially
# (start/deltas/stop, then next block), so pending_text from
# block N must be flushed by content_block_stop before block
# N+1's deltas arrive. If we ever see a delta whose index
# disagrees with the current pending buffer, the stream is
# interleaved -- fall back to legacy rather than risk merging
# text from different blocks under a single index.
if pending_text and pending_index is not None and cur_index != pending_index:
return None
saw_any_text_delta = True
pending_index = cur_index
pending_text.append(delta.get("text") or "")
elif etype == "ping":
# Interior no-op; legacy maps it to an empty chunk.
continue
else:
# message_start / content_block_stop / message_delta /
# message_stop / error: pass through unchanged.
flush()
out.append(ev)
flush()
if not saw_any_text_delta:
return None
return out
def build_complete_streaming_response_legacy(
all_chunks: Sequence[str | bytes],
litellm_logging_obj: LiteLLMLoggingObj,
model: str,
speed: str | None = None,
) -> ModelResponse | TextCompletionResponse | None:
"""
Original reconstruction: convert every SSE event to a generic chunk
and assemble via stream_chunk_builder. Kept verbatim as the fallback
/ source of truth for the fast path's parity test.
- Splits multi-event chunks into individual SSE events
- Converts str chunks to generic chunks
- Converts generic chunks to litellm chunks (OpenAI format)
- Builds complete response from litellm chunks
"""
verbose_proxy_logger.debug("Building complete streaming response from %d chunks", len(all_chunks))
anthropic_model_response_iterator: Final = AnthropicModelResponseIterator(
streaming_response=None,
sync_stream=False,
speed=speed,
)
all_openai_chunks: Final = []
# Process each chunk - a chunk may contain multiple SSE events
for _chunk_str in all_chunks:
# Split chunk into individual SSE events
individual_events = split_sse_chunk_into_events(_chunk_str)
# Process each individual event
for event_str in individual_events:
try:
# Skip OpenAI-style [DONE] sentinels some Anthropic-compatible
# providers emit. Match the whole SSE line so a valid chunk whose
# text payload happens to contain "[DONE]" is not dropped.
if any(line.strip() == "data: [DONE]" for line in event_str.split("\n")):
continue
transformed_openai_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk(
chunk=event_str
)
if transformed_openai_chunk is not None:
all_openai_chunks.append(transformed_openai_chunk)
except (StopIteration, StopAsyncIteration):
break
except json.JSONDecodeError:
# Some upstreams emit non-JSON SSE lines; skip them so the
# logging pipeline is not broken by a single bad frame.
verbose_proxy_logger.debug(
"Skipping non-JSON SSE event: %s",
event_str[:200],
)
continue
complete_streaming_response: Final = litellm.stream_chunk_builder(
chunks=all_openai_chunks,
logging_obj=litellm_logging_obj,
)
verbose_proxy_logger.debug("Complete streaming response built: %s", complete_streaming_response)
return complete_streaming_response
def extract_sse_data(event_str: str) -> dict | None:
"""Parse the JSON object from the ``data:`` line of an Anthropic SSE event."""
for line in event_str.splitlines():
stripped = line.strip()
if stripped.startswith("data:"):
payload = stripped[len("data:") :].strip()
if not payload or payload == "[DONE]":
return None
try:
return json.loads(payload)
except (ValueError, TypeError):
return None
return None
def build_usage_only_response_from_chunks(
all_chunks: Sequence[str | bytes],
model: str,
speed: str | None = None,
) -> ModelResponse | None:
"""
Build a usage-bearing ModelResponse from Anthropic SSE token-usage events, for
cost tracking when stream_chunk_builder cannot reassemble the stream.
Anthropic emits usage in ``message_start`` (uncached input + cache tokens, and an
initial output_tokens) and the final ``message_delta`` (cumulative output_tokens)
regardless of the content/tool shape, so cost is recoverable even when full
content assembly fails. Returns ``None`` if no usage event is found.
"""
input_tokens = 0
cache_read = 0
cache_creation = 0
cache_creation_5m: int | None = None
cache_creation_1h: int | None = None
output_tokens = 0
web_search_requests: int | None = None
tool_search_requests: int | None = None
inference_geo: str | None = None
speed_from_stream: str | None = None
stop_reason: str | None = None
found_usage = False
resolved_model = model
for _chunk_str in all_chunks:
for event_str in split_sse_chunk_into_events(_chunk_str):
data = extract_sse_data(event_str)
if not data:
continue
event_type = data.get("type")
if event_type == "message_start":
message = data.get("message") or {}
if not resolved_model or resolved_model == "unknown":
resolved_model = message.get("model") or resolved_model
usage = message.get("usage") or {}
input_tokens = usage.get("input_tokens") or input_tokens
cache_read = usage.get("cache_read_input_tokens") or cache_read
cache_creation = usage.get("cache_creation_input_tokens") or cache_creation
_cc = usage.get("cache_creation")
if isinstance(_cc, dict):
cache_creation_5m = _cc.get("ephemeral_5m_input_tokens")
cache_creation_1h = _cc.get("ephemeral_1h_input_tokens")
if usage.get("inference_geo") is not None:
inference_geo = usage.get("inference_geo")
if isinstance(usage.get("speed"), str):
speed_from_stream = usage.get("speed")
if usage.get("output_tokens") is not None:
output_tokens = usage.get("output_tokens")
found_usage = True
elif event_type == "message_delta":
_delta_stop = (data.get("delta") or {}).get("stop_reason")
if _delta_stop:
stop_reason = _delta_stop
usage = data.get("usage") or {}
if usage.get("output_tokens") is not None:
output_tokens = usage.get("output_tokens")
_stu = usage.get("server_tool_use")
if isinstance(_stu, dict):
if _stu.get("web_search_requests") is not None:
web_search_requests = _stu.get("web_search_requests")
if _stu.get("tool_search_requests") is not None:
tool_search_requests = _stu.get("tool_search_requests")
if usage.get("cache_read_input_tokens") is not None:
cache_read = usage.get("cache_read_input_tokens")
if usage.get("inference_geo") is not None:
inference_geo = usage.get("inference_geo")
if isinstance(usage.get("speed"), str):
speed_from_stream = usage.get("speed")
found_usage = True
if not found_usage:
return None
# If only the 5m/1h split was provided, derive the cache_creation total from it.
if not cache_creation and (cache_creation_5m or cache_creation_1h):
cache_creation = (cache_creation_5m or 0) + (cache_creation_1h or 0)
# build usage via the same AnthropicConfig.calculate_usage path the success
# cases use, so prompt_tokens are cache-inclusive and cache / server_tool_use /
# inference_geo tokens are priced instead of left at $0
usage_object: Final[dict] = {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
}
if cache_read:
usage_object["cache_read_input_tokens"] = cache_read
if cache_creation:
usage_object["cache_creation_input_tokens"] = cache_creation
if cache_creation_5m is not None or cache_creation_1h is not None:
usage_object["cache_creation"] = {
"ephemeral_5m_input_tokens": cache_creation_5m or 0,
"ephemeral_1h_input_tokens": cache_creation_1h or 0,
}
if web_search_requests is not None or tool_search_requests is not None:
_server_tool_use: Final[dict] = {}
if web_search_requests is not None:
_server_tool_use["web_search_requests"] = web_search_requests
if tool_search_requests is not None:
_server_tool_use["tool_search_requests"] = tool_search_requests
usage_object["server_tool_use"] = _server_tool_use
if inference_geo is not None:
usage_object["inference_geo"] = inference_geo
if speed_from_stream is not None:
usage_object["speed"] = speed_from_stream
usage_obj: Final = AnthropicConfig().calculate_usage(usage_object=usage_object, reasoning_content=None, speed=speed)
return ModelResponse(
model=resolved_model,
choices=[
Choices(
finish_reason=(map_finish_reason(stop_reason) if stop_reason else "stop"),
index=0,
message=Message(role="assistant", content=""),
)
],
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
)

49
litellm/llms/anthropic/prompt_cache_prediction.py
View file

@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio
import hashlib
import json
import time
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from functools import reduce
@ -14,6 +15,7 @@ import httpx
from pydantic import ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
from litellm.llms.anthropic.common_utils import (
AnthropicModelInfo,
@ -771,3 +773,50 @@ def parse_observed_cache(
cached_tokens=cache_tokens,
cache_creation_tokens=usage.cache_creation_tokens,
)
_RETENTION_SECONDS: Final = 86_400
class CacheObservation(LiteLLMBaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
fingerprint: str = Field(pattern=r"^[0-9a-f]{64}$")
cached_tokens: int = Field(gt=0)
observed_at: float = Field(ge=0, allow_inf_nan=False)
expires_at: float = Field(ge=0, allow_inf_nan=False)
_CACHE_ENTRY: Final[TypeAdapter[CacheObservation | str | None]] = TypeAdapter(CacheObservation | str | None)
def _cache_key(scope: str, fingerprint: str) -> str:
return f"prompt-cache-observation:{scope}:{fingerprint}"
async def lookup(
cache: DualCache, scope: str, prefix: PromptPrefix, now: float | None = None
) -> CacheObservation | None:
checked_at: Final = time.time() if now is None else now
exact: Final = await _read_exact(cache, scope, prefix.fingerprint)
if exact is not None and exact.expires_at > checked_at:
return exact
older: Final = await asyncio.gather(
*(_read_exact(cache, scope, fingerprint) for fingerprint in prefix.fingerprints[1:])
)
observations: Final = tuple(observation for observation in (exact, *older) if observation is not None)
return next(
(observation for observation in observations if observation.expires_at > checked_at),
next(iter(observations), None),
)
async def _read_exact(cache: DualCache, scope: str, fingerprint: str) -> CacheObservation | None:
try:
value: Final = _CACHE_ENTRY.validate_python(await cache.async_get_cache(_cache_key(scope, fingerprint), ttl=1)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # validate the legacy cache's untyped result at the I/O boundary
if value is None:
return None
observation: Final = CacheObservation.model_validate_json(value) if isinstance(value, str) else value
except ValidationError:
return None
return observation if observation.fingerprint == fingerprint else None

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/azure_blob_storage_backend.py
View file

@ -14,7 +14,7 @@ from urllib.parse import quote, urlparse
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
from litellm.proxy.common_utils.path_utils import safe_filename
from litellm.litellm_core_utils.path_utils import safe_filename
from .storage_backend import BaseFileStorageBackend

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/bedrock/passthrough/transformation.py
View file

@ -92,7 +92,7 @@ def _translate_message(decoder: "AWSEventStreamDecoder", message: str) -> ModelR
def _build_logged_response(
chunks: Sequence[ModelResponseStream], litellm_logging_obj: "LiteLLMLoggingObj"
) -> Optional["CostResponseTypes"]:
from litellm.main import stream_chunk_builder
from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder
if len(chunks) == 0:
return None

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({})

2
litellm/llms/gigachat/passthrough/transformation.py
View file

@ -146,11 +146,11 @@ class GigaChatPassthroughConfig(BasePassthroughConfig):
3. Return the model_response
"""
from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder
from litellm.litellm_core_utils.streaming_handler import (
convert_generic_chunk_to_model_response_stream,
generic_chunk_has_all_required_fields,
)
from litellm.main import stream_chunk_builder
from litellm.types.utils import ModelResponseStream
all_translated_chunks: Final[list[object]] = [] # mutable-ok: accumulator

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

@ -11,18 +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.proxy._types import LiteLLM_SkillsTable, 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/prompt_injection.py
View file

@ -11,7 +11,7 @@ from io import BytesIO
from typing import Any, Final
from litellm._logging import verbose_logger
from litellm.proxy._types import LiteLLM_SkillsTable
from litellm.models.skills import LiteLLM_SkillsTable
class SkillPromptInjectionHandler:

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

@ -20,9 +20,10 @@ from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth
from litellm.models.skills import LiteLLM_SkillsTable
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

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

@ -21,7 +21,8 @@ from litellm.types.utils import LlmProviders
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth
from litellm.models.skills import LiteLLM_SkillsTable
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
class LiteLLMSkillsTransformationHandler:
@ -136,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,

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

@ -25,6 +25,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
@ -45,7 +46,6 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
stream_item_items,
unappliable_request_rewrite,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
coerce_stream_holdback_value,
@ -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):

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

@ -12,14 +12,14 @@ from typing import TYPE_CHECKING, Any, Final, Optional
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import without_server_streaming_classification
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.proxy._types import PassThroughGuardrailSettings
from litellm.types.passthrough_endpoints.pass_through_endpoints import PassThroughGuardrailSettings
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):

127
litellm/llms/pass_through/typesafe_logging_handler.py Normal file
View file

@ -0,0 +1,127 @@
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import (
get_standard_logging_object_payload, # pyright: ignore[reportUnknownVariableType] # legacy helper has an untyped signature
)
from litellm.llms.laya.common_utils import laya_response_model
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import ModelResponse, StandardPassThroughResponseObject, Usage
if TYPE_CHECKING:
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
class _TypeSafeUsage(LiteLLMBaseModel):
input_tokens: int = 0
output_tokens: int = 0
cost: float | None = None
class _TypeSafeResponse(LiteLLMBaseModel):
model: str | None = None
usage: _TypeSafeUsage | None = None
class _RegistryPricing(LiteLLMBaseModel):
input_cost_per_token: float = 0.0
output_cost_per_token: float = 0.0
_TYPESAFE_RESPONSE_ADAPTER: Final = TypeAdapter(_TypeSafeResponse)
_REGISTRY_PRICING_ADAPTER: Final = TypeAdapter(_RegistryPricing)
def _parse_typesafe_response(response_body: Mapping[str, object]) -> _TypeSafeResponse:
try:
return _TYPESAFE_RESPONSE_ADAPTER.validate_python(response_body)
except ValidationError:
return _TypeSafeResponse()
def _pricing_for(model_keys: tuple[str, ...]) -> _RegistryPricing:
for model_key in model_keys:
if model_key not in litellm.model_cost: # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed
continue
try:
return _REGISTRY_PRICING_ADAPTER.validate_python(
litellm.model_cost[model_key] # pyright: ignore[reportUnknownMemberType] # registry is dynamically typed
)
except ValidationError:
continue
return _RegistryPricing()
class TypeSafePassthroughLoggingHandler:
@staticmethod
def typesafe_passthrough_handler(
httpx_response: httpx.Response,
response_body: Mapping[str, object],
logging_obj: LiteLLMLoggingObj,
url_route: str,
result: str,
start_time: datetime,
end_time: datetime,
cache_hit: bool,
request_body: Mapping[str, object],
custom_llm_provider: str,
**kwargs: object,
) -> "PassThroughEndpointLoggingTypedDict":
response: Final = _parse_typesafe_response(response_body)
request_model_value: Final = request_body.get("model")
request_model: Final = request_model_value if isinstance(request_model_value, str) else None
response_model: Final = (
laya_response_model(response_body, request_model) if custom_llm_provider == "laya" else response.model
)
logged_model: Final = response_model or request_model or "unknown"
model_name: Final = f"{custom_llm_provider}/{logged_model}"
usage: Final = response.usage or _TypeSafeUsage()
input_tokens: Final = usage.input_tokens
output_tokens: Final = usage.output_tokens
candidate_model_keys: Final = tuple(
f"{custom_llm_provider}/{model}" for model in (response_model, request_model) if model is not None
)
pricing: Final = _pricing_for(candidate_model_keys)
response_cost: Final = (
usage.cost
if custom_llm_provider == "openrouter" and usage.cost is not None
else input_tokens * pricing.input_cost_per_token + output_tokens * pricing.output_cost_per_token
)
usage_object: Final = Usage(
prompt_tokens=input_tokens,
completion_tokens=output_tokens,
total_tokens=input_tokens + output_tokens,
)
updated_kwargs: Final = {
**kwargs,
"model": model_name,
"custom_llm_provider": custom_llm_provider,
"response_cost": response_cost,
"combined_usage_object": usage_object,
}
logging_obj.model_call_details.update(
model=model_name,
custom_llm_provider=custom_llm_provider,
response_cost=response_cost,
)
standard_logging_object: Final = get_standard_logging_object_payload(
kwargs=updated_kwargs,
init_response_obj=ModelResponse(model=model_name, usage=usage_object),
start_time=start_time,
end_time=end_time,
logging_obj=logging_obj,
status="success",
)
return {
"result": StandardPassThroughResponseObject(response=result),
"kwargs": {
**updated_kwargs,
"standard_logging_object": standard_logging_object,
},
}

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())

616
litellm/main.py
View file

@ -26,7 +26,6 @@ from copy import deepcopy
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
from urllib.parse import urlsplit
from litellm._logging import redact_string
from litellm._uuid import uuid
@ -56,9 +55,7 @@ from litellm.utils import (
strip_reasoning_summary_aliases_from_optional_params,
)
# Logging is imported lazily when needed to avoid loading litellm_logging at import time
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.router import Router
from litellm.types.utils import TokenCountResponse
@ -104,17 +101,13 @@ from litellm.litellm_core_utils.mock_functions import (
mock_image_generation,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_content_from_model_response,
get_content_from_model_response as get_content_from_model_response, # noqa: PLC0414 # public re-export
)
from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header
from litellm.litellm_core_utils.request_timeout_resolver import (
get_configured_request_timeout,
)
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
from litellm.llms.azure_ai.common_utils import (
azure_ai_supports_native_responses,
foundry_chat_rejects_function_tools_while_reasoning,
)
from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig
from litellm.llms.base_llm.base_model_iterator import (
convert_model_response_to_streaming,
@ -169,7 +162,6 @@ from litellm.utils import (
create_tokenizer,
get_llm_provider,
get_model_info,
get_model_info_helper,
get_non_default_completion_params,
get_non_default_transcription_params,
get_optional_params_embeddings,
@ -183,13 +175,15 @@ from litellm.utils import (
read_config_args,
should_run_mock_completion,
supports_httpx_timeout,
token_counter,
validate_and_fix_openai_messages,
validate_and_fix_openai_tools,
validate_and_fix_thinking_param,
validate_chat_completion_tool_choice,
validate_openai_optional_params,
)
from litellm.utils import (
token_counter as token_counter, # noqa: PLC0414 # public re-export
)
from ._logging import verbose_logger
from .caching.caching import disable_cache, enable_cache, update_cache
@ -211,7 +205,9 @@ from .litellm_core_utils.prompt_templates.factory import (
prompt_factory,
stringify_json_tool_call_content,
)
from .litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from .litellm_core_utils.streaming_chunk_builder_utils import (
ChunkProcessor as ChunkProcessor, # noqa: PLC0414 # public re-export
)
from .llms.anthropic.chat import AnthropicChatCompletion
from .llms.azure.audio_transcriptions import AzureAudioTranscription
from .llms.azure.azure import AzureChatCompletion, check_dynamic_azure_params
@ -1071,142 +1067,12 @@ def mock_completion(
raise Exception(f"Mock completion response failed - {e}")
_OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
_OPENAI_API_HOST: Final = "api.openai.com"
def _is_openai_backed_api_base(api_base: str) -> bool:
hostname: Final = urlsplit(api_base).hostname
return hostname is not None and (hostname == _OPENAI_API_HOST or hostname.endswith(f".{_OPENAI_API_HOST}"))
def _resolve_openai_api_base(api_base: str | None) -> str:
"""Effective OpenAI base a chat request will hit: arg > global > env > default. The bridge gate
and the ``_complete_custom_openai`` chat handler MUST resolve this identically, or a custom base
set via ``litellm.api_base`` or ``OPENAI_BASE_URL``/``OPENAI_API_BASE`` is invisible to the gate,
which then misreads it as the default OpenAI endpoint and bridges a request the backend can't serve."""
return (
api_base
or litellm.api_base
or get_secret_str("OPENAI_BASE_URL")
or get_secret_str("OPENAI_API_BASE")
or _OPENAI_DEFAULT_API_BASE
)
def responses_api_bridge_check(
model: str,
custom_llm_provider: str,
web_search_options: OpenAIWebSearchOptions | None = None,
tools: Sequence[Mapping[str, object]] | None = None,
reasoning_effort: str | Mapping[str, object] | None = None,
reasoning_summary: object | None = None,
api_base: str | None = None,
) -> tuple[dict, str]:
model_info: dict[str, object] = {}
# Global flag: route ALL OpenAI chat completions through Responses API.
# Returns early with minimal model_info; callers only inspect the "mode" key.
if litellm.route_all_chat_openai_to_responses and custom_llm_provider == "openai":
model = model.replace("responses/", "")
model_info["mode"] = "responses"
return model_info, model
try:
model_info = cast(
dict,
get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider),
)
if model_info.get("mode") is None and model.startswith("responses/"):
model = model.replace("responses/", "")
mode = "responses"
model_info["mode"] = mode
except Exception as e:
verbose_logger.debug("Error getting model info: %s", e)
if model.startswith("responses/"): # handle azure models - `azure/responses/<deployment-name>`
model = model.replace("responses/", "")
mode = "responses"
model_info["mode"] = mode
if web_search_options is not None and custom_llm_provider == "xai":
model_info["mode"] = "responses"
model = model.replace("responses/", "")
# OpenAI/Azure GPT-5 chat-completions that need Responses-only fields (e.g.
# ``reasoningSummary`` in ``extra_body``) must be bridged; Chat Completions rejects
# those keys.
#
# - gpt-5.4+: FUNCTION tools with reasoning active must be bridged. OpenAI enables
# reasoning by default for these models (unset reasoning_effort means medium
# server-side), and Chat Completions rejects function tools whenever reasoning is
# on ("Function tools with reasoning_effort are not supported ... use
# /v1/responses or set reasoning_effort to 'none'"), so only an explicit
# ``"none"`` keeps the request chat-servable. Custom (grammar) tools are served
# natively by Chat Completions with reasoning on, so custom-only requests stay on
# chat and keep their native custom tool_call response shape.
# - The UNSET-effort arm only fires against endpoints known to enforce that
# constraint (any api.openai.com host, or Azure OpenAI where api_base is
# always set): chat-only OpenAI-compatible backends registered under the openai
# provider with a custom api_base and gpt-5.4+ model names serve tools without
# reasoning fine and have no /responses route, so they keep pre-existing
# behavior (bridge only on an explicit reasoning_effort).
# - Azure AI Foundry's OpenAI v1 hosts (azure_ai provider) enforce it later in the series:
# an explicit effort with function tools is rejected from gpt-5.6 on, and the unset
# effort only from gpt-6 on (gpt-5.6 serves tools with reasoning silently off), so the
# azure_ai gate keys on those measured boundaries instead of gpt-5.4+.
# - Older GPT-5 names (e.g. ``gpt-5``, ``gpt-5.1``): bridge only when a reasoning
# summary alias is present with ``reasoning_effort`` (tools alone stay on chat).
has_function_tool: Final = any(
(
tool.get("type") == "function" and (isinstance(tool.get("function"), dict) or "name" in tool)
if isinstance(tool, dict)
else getattr(tool, "type", None) == "function"
)
for tool in (tools or ())
)
if isinstance(reasoning_effort, dict):
reasoning_active = reasoning_effort.get("effort") != "none" or reasoning_effort.get("summary") is not None
else:
reasoning_active = reasoning_effort != "none"
# The reasoning+tools constraint is enforced by the real OpenAI backend behind any api.openai.com
# host (the default URL or a PrivateLink hostname such as <region>.privatelink.api.openai.com) and
# by Azure OpenAI through the azure provider. Resolve the effective OpenAI base arg>global>env>default
# exactly as the chat handler does, so a custom base set via litellm.api_base or
# OPENAI_BASE_URL/OPENAI_API_BASE isn't misread as the default and bridged to a /responses route it
# lacks. A whitespace-only base collapses to the default too.
resolved_api_base: Final = _resolve_openai_api_base(api_base).strip()
on_foundry_openai_endpoint: Final = custom_llm_provider == "azure_ai" and azure_ai_supports_native_responses(
model, api_base
)
on_constraint_enforcing_endpoint: Final = (
custom_llm_provider == "azure" or resolved_api_base == "" or _is_openai_backed_api_base(resolved_api_base)
)
chat_rejects_function_tools: Final = (
has_function_tool
and reasoning_active
and (
foundry_chat_rejects_function_tools_while_reasoning(model, reasoning_effort)
if on_foundry_openai_endpoint
else (
OpenAIGPT5Config.is_model_gpt_5_4_plus_model(model)
and (reasoning_effort is not None or on_constraint_enforcing_endpoint)
)
)
)
if (
(custom_llm_provider in ("openai", "azure") or on_foundry_openai_endpoint)
and model_info.get("mode") != "responses"
and OpenAIGPT5Config.is_model_gpt_5_model(model)
and not OpenAIGPT5Config.is_model_gpt_5_search_model(model)
and ((reasoning_effort is not None and reasoning_summary is not None) or chat_rejects_function_tools)
):
model_info["mode"] = "responses"
model = model.replace("responses/", "")
return model_info, model
from litellm.responses.bridge_check import ( # noqa: E402 # re-export at the old definition site
_resolve_openai_api_base as _resolve_openai_api_base, # pyright: ignore[reportPrivateUsage] # noqa: PLC0414 # re-export of moved helper
)
from litellm.responses.bridge_check import ( # noqa: E402 # re-export at the old definition site
responses_api_bridge_check as responses_api_bridge_check, # pyright: ignore[reportUnknownVariableType] # noqa: PLC0414 # re-export of moved function
)
_ANTHROPIC_ONLY_TOOL_KEYS: Final = frozenset({"input_examples", "eager_input_streaming"})
@ -2072,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
@ -5993,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)
@ -8966,423 +8867,12 @@ def config_completion(**kwargs):
)
def stream_chunk_builder_text_completion(chunks: list, messages: Sequence | None = None) -> TextCompletionResponse:
id: Final = chunks[0]["id"]
object: Final = chunks[0]["object"]
created: Final = chunks[0]["created"]
model: Final = chunks[0]["model"]
system_fingerprint: Final = chunks[0].get("system_fingerprint", None)
finish_reason: Final = chunks[-1]["choices"][0]["finish_reason"]
logprobs: Final = chunks[-1]["choices"][0]["logprobs"]
content_list: Final = []
for chunk in chunks:
choices = chunk["choices"]
for choice in choices:
if choice is not None and hasattr(choice, "text") and choice.get("text") is not None:
_choice = choice.get("text")
content_list.append(_choice)
# Combine the "content" strings into a single string || combine the 'function' strings into a single string
combined_content: Final = "".join(content_list)
try:
prompt_tokens = token_counter(model=model, messages=messages)
except Exception: # don't allow this failing to block a complete streaming response from being returned
print_verbose("token_counter failed, assuming prompt tokens is 0")
prompt_tokens = 0
completion_tokens: Final = token_counter(
model=model,
text=combined_content,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
)
response: Final = {
"id": id,
"object": object,
"created": created,
"model": model,
"system_fingerprint": system_fingerprint,
"choices": [
{
"text": combined_content,
"index": 0,
"logprobs": logprobs,
"finish_reason": finish_reason,
}
],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
}
return TextCompletionResponse(**response)
_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value})
def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool:
if logging_obj is None:
return False
provider: Final[object] = logging_obj.model_call_details.get("custom_llm_provider")
return provider in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS
def _stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> float | None:
usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None)
if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj):
return float(usage_cost)
if logging_obj is not None:
return None
provider_hint: Final = response._hidden_params.get( # pyright: ignore[reportPrivateUsage] # no public accessor
"custom_llm_provider"
)
try:
return litellm.completion_cost(completion_response=response, custom_llm_provider=provider_hint)
except Exception:
return _stream_builder_model_map_cost(response)
def _joined_streamed_citations(streamed_citations: "tuple[object, ...]") -> "list[object]":
if all(isinstance(citation, list) for citation in streamed_citations):
return list(streamed_citations)
return [list(streamed_citations)]
def _stream_builder_model_map_cost(response: ModelResponse) -> float | None:
model_name: Final = response.model
usage: Final = getattr(response, "usage", None)
if not model_name or not isinstance(usage, Usage):
return None
try:
prompt_cost, completion_tokens_cost = litellm.cost_per_token(model=model_name, usage_object=usage)
return prompt_cost + completion_tokens_cost
except Exception: # noqa: BLE001 # cost_per_token raises bare Exception for unpriceable models
return None
def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
response_cost: Final = _stream_builder_response_cost(response, logging_obj)
if response_cost is None:
return
hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor
hidden_params["response_cost"] = response_cost
def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
if logging_obj is None:
return
if isinstance(getattr(usage, "cost", None), (int, float)):
return
computed_cost: Final = logging_obj.response_cost_calculator(result=response)
if isinstance(computed_cost, (int, float)) and computed_cost > 0:
setattr(usage, "cost", computed_cost)
_NON_TEXT_DELTA_FIELDS: Final = (
"tool_calls",
"function_call",
"reasoning_content",
"thinking_blocks",
"annotations",
"audio",
"images",
"provider_specific_fields",
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( # noqa: E402 # re-export after the move
stream_chunk_builder as stream_chunk_builder, # noqa: PLC0414 # pyright: ignore[reportUnknownVariableType] # chunks is an untyped list
)
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( # noqa: E402 # re-export after the move
stream_chunk_builder_text_completion as stream_chunk_builder_text_completion, # noqa: PLC0414 # pyright: ignore[reportUnknownVariableType] # chunks is an untyped list
)
def _stream_choice_delta(choice: object) -> Mapping[str, object]:
delta: Final = choice.get("delta", {}) if isinstance(choice, dict) else getattr(choice, "delta", {})
if isinstance(delta, Mapping):
return delta
if isinstance(delta, BaseModel):
return delta.model_dump()
return {}
def _delta_carries_more_than_text(delta: Mapping[str, object]) -> bool:
return any(delta.get(field) is not None for field in _NON_TEXT_DELTA_FIELDS)
def _simple_text_part(choices: Sequence[object]) -> str | None:
deltas: Final = tuple(_stream_choice_delta(choice) for choice in choices)
if any(_delta_carries_more_than_text(delta) for delta in deltas):
return None
content: Final = deltas[0].get("content")
return content if isinstance(content, str) else ""
def stream_chunk_builder(
chunks: list,
messages: Sequence | None = None,
start_time=None,
end_time=None,
logging_obj: Optional["Logging"] = None,
count_prompt_tokens: Callable[[], int] | None = None,
) -> ModelResponse | TextCompletionResponse | None:
try:
if chunks is None:
raise litellm.APIError(
status_code=500,
message="Error building chunks for logging/streaming usage calculation",
llm_provider="",
model="",
)
if not chunks:
return None
processor: Final = ChunkProcessor(chunks, messages)
chunks = processor.chunks
### BASE-CASE ###
if len(chunks) == 0:
return None
## Route to the text completion logic
first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), None)
if first_chunk_with_choices is not None and isinstance(
first_chunk_with_choices["choices"][0], litellm.utils.TextChoices
): # route to the text completion logic
return stream_chunk_builder_text_completion(chunks=chunks, messages=messages)
model: Final = chunks[0]["model"]
# Initialize the response dictionary
response: Final = processor.build_base_response(chunks)
# Fast path for the common text-only streaming case:
# avoid repeated multi-pass list scans over chunks.
simple_content_parts: Final[list[str]] = []
is_simple_text_stream = True
for chunk in chunks:
if not chunk.get("choices"):
continue
if (part := _simple_text_part(chunk["choices"])) is None:
is_simple_text_stream = False
break
if part:
simple_content_parts.append(part)
if is_simple_text_stream:
if simple_content_parts:
response["choices"][0]["message"]["content"] = "".join(simple_content_parts)
completion_output = get_content_from_model_response(response)
usage = processor.calculate_usage(
chunks=chunks,
model=model,
completion_output=completion_output,
messages=messages,
reasoning_tokens=0,
count_prompt_tokens=count_prompt_tokens,
)
setattr(response, "usage", usage)
# Propagate provider_specific_fields from chunk hidden params when present.
for chunk in reversed(chunks):
if isinstance(chunk, dict):
hidden = chunk.get("_hidden_params")
else:
hidden = getattr(chunk, HIDDEN_PARAMS_ATTR, None)
if isinstance(hidden, dict) and "provider_specific_fields" in hidden:
response.hidden_params.setdefault("provider_specific_fields", {}).update(
hidden["provider_specific_fields"]
)
break
_stamp_streaming_usage_cost(usage, response, logging_obj)
_set_stream_builder_response_cost(response, logging_obj)
processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
return response
tool_call_chunks: Final = [
chunk
for chunk in chunks
if any(
"tool_calls" in choice["delta"] and choice["delta"]["tool_calls"] is not None
for choice in chunk.get("choices") or ()
)
]
if len(tool_call_chunks) > 0:
tool_calls_list: Final = processor.get_combined_tool_content(tool_call_chunks)
_choice = response.choices[0]
_choice.message.content = None
_choice.message.tool_calls = tool_calls_list
function_call_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "function_call" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["function_call"] is not None
]
if len(function_call_chunks) > 0:
_choice = response.choices[0]
_choice.message.content = None
_choice.message.function_call = processor.get_combined_function_call_content(function_call_chunks)
content_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["content"] is not None
]
if len(content_chunks) > 0:
response["choices"][0]["message"]["content"] = processor.get_combined_content(content_chunks)
thinking_blocks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "thinking_blocks" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["thinking_blocks"] is not None
]
if len(thinking_blocks) > 0:
response["choices"][0]["message"]["thinking_blocks"] = processor.get_combined_thinking_content(
thinking_blocks
)
reasoning_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "reasoning_content" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["reasoning_content"] is not None
]
if len(reasoning_chunks) > 0:
response["choices"][0]["message"]["reasoning_content"] = processor.get_combined_reasoning_content(
reasoning_chunks
)
annotation_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "annotations" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["annotations"] is not None
]
if len(annotation_chunks) > 0:
# Merge annotations from ALL chunks — providers may spread
# them across multiple streaming chunks or send them only in
# the final chunk.
all_annotations: Final[list] = []
for ac in annotation_chunks:
all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
response["choices"][0]["message"]["annotations"] = all_annotations
audio_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "audio" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["audio"] is not None
]
if len(audio_chunks) > 0:
_choice = response.choices[0]
_choice.message.audio = processor.get_combined_audio_content(audio_chunks)
# Handle image chunks from models like gemini-2.5-flash-image
# See: https://github.com/BerriAI/litellm/issues/19478
image_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "images" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["images"] is not None
]
if len(image_chunks) > 0:
# Images come complete in a single chunk, collect all images from all chunks
all_images: Final = []
for chunk in image_chunks:
all_images.extend(chunk["choices"][0]["delta"]["images"])
response["choices"][0]["message"]["images"] = all_images
# Combine provider_specific_fields from streaming chunks (e.g., web_search_results, citations)
# See: https://github.com/BerriAI/litellm/issues/17737
provider_specific_chunks: Final = [
chunk
for chunk in chunks
if chunk.get("choices")
and "provider_specific_fields" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["provider_specific_fields"] is not None
]
if len(provider_specific_chunks) > 0:
provider_field_dicts: Final = tuple(
fields
for chunk in provider_specific_chunks
for fields in (chunk["choices"][0]["delta"]["provider_specific_fields"],)
if isinstance(fields, dict)
)
streamed_citations: Final = tuple(
fields["citation"] for fields in provider_field_dicts if fields.get("citation") is not None
)
citation_fields: Final = (
{"citations": _joined_streamed_citations(streamed_citations)} if streamed_citations else {}
)
combined_provider_fields: Final = {
key: value
for fields in (citation_fields, *provider_field_dicts)
for key, value in fields.items()
if key != "citation"
}
if combined_provider_fields:
_choice = response.choices[0]
_choice.message.provider_specific_fields = combined_provider_fields
completion_output = get_content_from_model_response(response)
reasoning_tokens: Final = processor.count_reasoning_tokens(response)
usage = processor.calculate_usage(
chunks=chunks,
model=model,
completion_output=completion_output,
messages=messages,
reasoning_tokens=reasoning_tokens,
count_prompt_tokens=count_prompt_tokens,
)
setattr(response, "usage", usage)
# Propagate provider_specific_fields from the last chunk (contains provider
# metadata like traffic_type set during streaming)
for chunk in reversed(chunks):
if isinstance(chunk, dict):
hidden = chunk.get("_hidden_params")
else:
hidden = getattr(chunk, HIDDEN_PARAMS_ATTR, None)
if isinstance(hidden, dict) and "provider_specific_fields" in hidden:
response.hidden_params.setdefault("provider_specific_fields", {}).update(
hidden["provider_specific_fields"]
)
break
_stamp_streaming_usage_cost(usage, response, logging_obj)
_set_stream_builder_response_cost(response, logging_obj)
processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj)
return response
except Exception as e:
verbose_logger.exception("litellm.main.py::stream_chunk_builder() - Exception occurred - %s", e)
raise litellm.APIError(
status_code=500,
message="Error building chunks for logging/streaming usage calculation",
llm_provider="",
model="",
)
########## Token Counting API ##########

64
litellm/model_prices_and_context_window_backup.json
View file

@ -13213,28 +13213,30 @@
},
"bedrock/moonshotai.kimi-k2-thinking": {
"deprecation_date": "2027-03-30",
"input_cost_per_token": 7.3e-07,
"input_cost_per_token_flex": 3.6e-07,
"input_cost_per_token": 6e-07,
"input_cost_per_token_flex": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 262144,
"max_output_tokens": 16000,
"max_tokens": 16000,
"mode": "chat",
"output_cost_per_token": 3.03e-06,
"output_cost_per_token_flex": 1.52e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-thinking.html",
"output_cost_per_token": 2.5e-06,
"output_cost_per_token_flex": 1.25e-06,
"source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json",
"supports_function_calling": true,
"supports_reasoning": true
},
"bedrock/moonshotai.kimi-k2.5": {
"input_cost_per_token": 6e-07,
"input_cost_per_token_flex": 3e-07,
"litellm_provider": "bedrock",
"max_input_tokens": 262144,
"max_output_tokens": 16000,
"max_tokens": 16000,
"mode": "chat",
"output_cost_per_token": 3.03e-06,
"source": "https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-moonshot-ai-kimi-k2-5.html",
"output_cost_per_token": 3e-06,
"output_cost_per_token_flex": 1.5e-06,
"source": "https://b0.p.awsstatic.com/pricing/2.0/meteredUnitMaps/bedrock/USD/current/bedrock.json",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_system_messages": true,
@ -66874,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,
@ -71840,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",

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

@ -1,6 +1,6 @@
import re
import secrets
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Set as AbstractSet
from dataclasses import dataclass
from datetime import datetime, timezone
@ -17,7 +17,7 @@ from typing_extensions import assert_never
import litellm
from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
from litellm.constants import MCP_ALL_TOOLS_WILDCARD
from litellm.constants import MCP_ALL_TOOLS_WILDCARD, MCP_PEEKED_BODY_SCOPE_KEY
from litellm.experimental_mcp_client.client import strip_auth_scheme
from litellm.proxy._experimental.mcp_server.catalog import global_manager
from litellm.proxy._experimental.mcp_server.oauth_utils import (
@ -94,6 +94,20 @@ def _normalize_caller_admission_credential(value: str) -> str:
return _get_bearer_token_or_received_api_key(strip_auth_scheme(value, "Bearer")).strip()
def _admission_request(scope: Scope) -> Request:
"""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)
async def mock_body() -> bytes:
if peeked_body is None:
return b"{}"
return await peeked_body()
request.body = mock_body
return request
def _as_list(values: Sequence[str] | None) -> list[str] | None: # mutable-ok: resolver returns a list
"""Widen a read-only allowlist back to the mutable list the resolver's own contract returns,
preserving the ``None`` that means "no restriction"."""
@ -505,12 +519,7 @@ class MCPRequestHandler:
if mcp_servers_header == "" or (mcp_servers is not None and len(mcp_servers) == 0):
mcp_servers = []
# Create a proper Request object with mock body method to avoid ASGI receive channel issues
request: Final = Request(scope=scope)
async def mock_body():
return b"{}"
request.body = mock_body
request: Final = _admission_request(scope)
# Inline import — auth_utils participates in a proxy import cycle.
from litellm.proxy.auth.auth_utils import ( # noqa: PLC0415
get_request_route,
@ -1614,7 +1623,7 @@ class MCPRequestHandler:
)
litellm_api_key: Final = (
headers.get(custom_key_header_name)
if custom_key_header_name is not None
if custom_key_header_name is not None and custom_key_header_name in headers
else next(
(header_value for header_name in admission_header_names if (header_value := headers.get(header_name))),
None,

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)

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

@ -26,7 +26,9 @@ from starlette.types import Message, Receive, Scope, Send
from litellm._logging import verbose_logger
from litellm.constants import (
MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS,
MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH,
MCP_PEEKED_BODY_SCOPE_KEY,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -656,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
@ -791,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:
@ -1262,6 +1266,17 @@ if MCP_AVAILABLE:
return f"ip:{hashlib.sha256(client_ip.encode('utf-8')).hexdigest()}"
return "anonymous"
def _admission_auth_body(body: bytes) -> bytes:
"""Return only the JSON-RPC method for admission auth."""
try:
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:
"""
Check if the request body is a JSON-RPC initialize method.
@ -1402,6 +1417,7 @@ if MCP_AVAILABLE:
receive: Receive,
*,
full_body: bool = False,
already_read: int = 0,
) -> tuple[list[Message], bytes]:
"""
Read just enough of the request body to decide whether this is a
@ -1412,14 +1428,14 @@ if MCP_AVAILABLE:
Stops reading from the wire as soon as either (a) we have peeked
``_MCP_ROUTING_PEEK_MAX_BYTES`` of body, or (b) the body is complete.
The remainder of an oversized body is streamed lazily through
``wrapped_receive`` in the caller — so an authenticated client cannot
``_LazyPeekedBody.receive`` in the caller — so an authenticated client cannot
force the proxy to buffer an arbitrarily large payload just to make a
routing decision. Modern interaction preflight requests the full body,
which the SDK's single-exchange transport also requires.
"""
consumed_messages: Final[list[Message]] = []
body_chunks: Final[list[bytes]] = []
peeked_bytes = 0
peeked_bytes = already_read
while True:
message = await receive()
@ -1457,6 +1473,60 @@ if MCP_AVAILABLE:
return consumed_messages, b"".join(body_chunks)
_PeekedRead = tuple[list[Message], bytes]
class _LazyPeekedBody:
"""Reads the POST body for routing and auth only when first asked, so an
allowlist rejection never touches the receive channel. Whatever it
consumed is replayed to the downstream handler through ``receive``,
which falls through to the wire once the replay buffer is drained."""
__slots__ = ("_consumed_messages", "_peeked_body", "_read", "_receive")
def __init__(self, receive: Receive) -> None:
self._receive = receive
self._read: asyncio.Future[_PeekedRead] | None = None
self._consumed_messages: list[Message] | None = None
self._peeked_body: bytes = b""
def _start_read(self) -> asyncio.Future[_PeekedRead]:
if self._read is None:
self._read = asyncio.ensure_future(_read_request_body_for_routing(self._receive))
return self._read
async def body(self) -> bytes:
if self._consumed_messages is None:
self._consumed_messages, self._peeked_body = await self._start_read()
return self._peeked_body
async def admission_body(self) -> bytes:
try:
await asyncio.wait_for(
asyncio.shield(self._start_read()), timeout=MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS
)
except asyncio.TimeoutError:
return b""
return await self.body()
async def full_body(self) -> bytes:
await self.body()
consumed: Final = self._consumed_messages or []
consumed_body: Final = b"".join(message.get("body", b"") or b"" for message in consumed)
if len(consumed_body) > session_manager_stateless.max_request_body_size:
raise HTTPException(status_code=413, detail="Request body too large")
if not consumed or consumed[-1].get("type") != "http.request" or not consumed[-1].get("more_body", False):
return consumed_body
rest_messages, rest_body = await _read_request_body_for_routing(
self._receive, full_body=True, already_read=len(consumed_body)
)
self._consumed_messages = [*consumed, *rest_messages]
return consumed_body + rest_body
async def receive(self) -> Message:
if self._consumed_messages:
return self._consumed_messages.pop(0)
return await self._receive()
async def _handle_stale_mcp_session(
scope: Scope,
receive: Receive,
@ -2124,6 +2194,14 @@ if MCP_AVAILABLE:
)(scope, receive, send)
return
path: Final[str] = scope.get("path", "")
peek: Final = _LazyPeekedBody(receive) if scope.get("method") == "POST" else None
if peek is not None:
receive = peek.receive # rebind-ok: replay peeked ASGI messages to the downstream handler
async def admission_auth_body() -> bytes:
return _admission_auth_body(await peek.admission_body())
scope[MCP_PEEKED_BODY_SCOPE_KEY] = admission_auth_body
(
user_api_key_auth,
mcp_auth_header,
@ -2133,6 +2211,8 @@ if MCP_AVAILABLE:
raw_headers,
) = await extract_mcp_auth_context(scope, path)
reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth)
body: Final = await peek.body() if peek is not None else b""
is_initialize: Final = _is_initialize_request(body)
scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1
request_headers: Final = StarletteRequest(scope).headers
defer_upstream_probes: Final = request_headers.get(
@ -2217,8 +2297,6 @@ if MCP_AVAILABLE:
# - No session ID + initialize → stateful (so client gets mcp-session-id)
# - No session ID + other → stateless (curl, Inspector, Notion)
session_id = _get_session_id_from_scope(scope)
is_initialize = False
consumed_messages: list[Message] = []
# Owner-binding: a live stateful session may only be driven by the
# caller that created it. Reject mismatches with 403 so a leaked
@ -2226,8 +2304,7 @@ if MCP_AVAILABLE:
#
# Run before ``_handle_stale_mcp_session`` so a non-owner cannot
# force-clean another caller's residual tracking entries via a
# stale DELETE, and before peeking the request body so the 403
# response sees a pristine ``receive`` channel.
# stale DELETE.
if session_id:
expected_owner: Final = _stateful_session_owners.get(session_id)
request_owner = _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip)
@ -2256,27 +2333,23 @@ if MCP_AVAILABLE:
return
session_id = _get_session_id_from_scope(scope)
body = b""
if scope.get("method") == "POST":
consumed_messages, body = await _read_request_body_for_routing(receive, full_body=defer_upstream_probes)
if defer_upstream_probes:
rejection: Final = await _preflight_modern_interaction(
scope,
body,
OperationContext(
_caller=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_client_ip,
),
)
if rejection is not None:
await rejection(scope, receive, send)
return
is_initialize = _is_initialize_request(body)
if peek is not None and defer_upstream_probes:
rejection: Final = await _preflight_modern_interaction(
scope,
await peek.full_body(),
OperationContext(
_caller=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
client_ip=_client_ip,
),
)
if rejection is not None:
await rejection(scope, receive, send)
return
use_stateful: Final = bool(session_id or is_initialize)
target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless
@ -2306,17 +2379,6 @@ if MCP_AVAILABLE:
await too_many_response(scope, receive, send)
return
# Replay body messages if we consumed them for peeking
original_receive: Final = receive
if consumed_messages:
async def wrapped_receive():
if consumed_messages:
return consumed_messages.pop(0)
return await original_receive()
receive = wrapped_receive
# Serialize requests on the same stateful session so concurrent
# callers don't clobber each other's auth context mid-flight.
#
@ -2799,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

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