scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

merge(litellm_internal_staging): bring vendor strategy e2e branch up to date

Browse source
This commit is contained in:
mubashir1osmani 2026-08-04 13:05:52 -07:00
commit 2e5f7138e5
2488 changed files with 67813 additions and 43512 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

136
.circleci/config.yml
View file

@ -88,6 +88,36 @@ commands:
rm -f /tmp/uv-install.sh
echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV"
export PATH="$HOME/.local/bin:$PATH"
install_rust:
description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself."
steps:
- run:
name: Install Rust (rustup 1.28.2, toolchain 1.97.1)
command: |
case "$(uname -m)" in
x86_64)
RUSTUP_TRIPLE=x86_64-unknown-linux-gnu
RUSTUP_SHA256=20a06e644b0d9bd2fbdbfd52d42540bdde820ea7df86e92e533c073da0cdd43c
;;
aarch64)
RUSTUP_TRIPLE=aarch64-unknown-linux-gnu
RUSTUP_SHA256=e3853c5a252fca15252d07cb23a1bdd9377a8c6f3efa01531109281ae47f841c
;;
*)
echo "install_rust: unsupported architecture $(uname -m)" >&2
exit 1
;;
esac
curl -sSLf -o /tmp/rustup-init \
"https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init"
echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c -
chmod +x /tmp/rustup-init
/tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1
rm -f /tmp/rustup-init
echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV"
export PATH="$HOME/.cargo/bin:$PATH"
rustc --version
cargo --version
start_postgres:
description: "Start a postgres-db container on port 5432 and wait until it accepts connections."
parameters:
@ -163,6 +193,26 @@ commands:
done
echo "fake OpenAI endpoint did not become ready" >&2
exit 1
start_cost_center_service:
description: "Start the stand-in cost center validation service (tests/store_model_in_db_tests/cost_center_service.py) on host port 9414 and wait until healthy. The proxy's team-metadata validator (team_metadata_validator_e2e.py, impl 'http') reaches it via TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate. Run after uv deps are synced."
steps:
- run:
name: Start cost center validation service
background: true
command: |
uv run --no-sync python tests/store_model_in_db_tests/cost_center_service.py --host 0.0.0.0 --port 9414
- run:
name: Wait for cost center validation service
command: |
for i in $(seq 1 30); do
if curl -sf http://localhost:9414/health >/dev/null 2>&1; then
echo "cost center validation service is up"
exit 0
fi
sleep 1
done
echo "cost center validation service did not become ready" >&2
exit 1
setup_litellm_enterprise_pip:
steps:
- run:
@ -178,6 +228,7 @@ commands:
- checkout
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -281,6 +332,33 @@ jobs:
uv build --wheel --out-dir dist
uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py
base_sdk_install:
docker:
- image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
steps:
- checkout
- setup_google_dns
- install_uv
- install_rust
- run:
name: Build the wheel
environment:
UV_HTTP_TIMEOUT: "300"
command: |
uv build --wheel --out-dir dist
- run:
name: Install the wheel with no extras and smoke-check it
environment:
UV_HTTP_TIMEOUT: "300"
command: |
uv venv /tmp/base-sdk --python 3.12
VIRTUAL_ENV=/tmp/base-sdk uv pip install dist/*.whl
/tmp/base-sdk/bin/python tests/base_sdk_tests/check_base_sdk_install.py
local_testing_part1:
docker:
- &python312_image
@ -298,6 +376,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -371,6 +450,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -445,6 +525,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -496,6 +577,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -562,6 +644,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -602,6 +685,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -643,6 +727,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -676,6 +761,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -726,6 +812,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -777,6 +864,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -810,6 +898,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -856,6 +945,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -902,6 +992,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -944,6 +1035,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -990,6 +1082,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1037,6 +1130,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -1077,6 +1171,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1122,6 +1217,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1166,6 +1262,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1198,6 +1295,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1241,6 +1339,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1285,6 +1384,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1329,6 +1429,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1360,6 +1461,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1406,6 +1508,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1451,6 +1554,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1501,6 +1605,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1525,6 +1630,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1551,6 +1657,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1652,6 +1759,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1747,6 +1855,7 @@ jobs:
at: ~/project
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1835,6 +1944,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1918,6 +2028,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2050,6 +2161,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2136,6 +2248,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2232,12 +2345,14 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- start_postgres
- start_fake_openai_endpoint
- start_cost_center_service
- attach_workspace:
at: ~/project
- run:
@ -2257,11 +2372,13 @@ jobs:
-e STORE_MODEL_IN_DB="True" \
-e LITELLM_MASTER_KEY="sk-1234" \
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
-e TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate \
-e LITELLM_LICENSE=$LITELLM_LICENSE \
-e LITELLM_LOG=ERROR \
--add-host host.docker.internal:host-gateway \
--name my-app \
-v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \
-v $(pwd)/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py:/app/team_metadata_validator_e2e.py \
litellm-docker-database:ci \
--config /app/config.yaml \
--port 4000
@ -2307,6 +2424,7 @@ jobs:
- setup_google_dns
# Remove Docker CLI installation since it's already available in machine executor
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2388,6 +2506,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2473,14 +2592,14 @@ jobs:
bundle exec rspec
no_output_timeout: 30m
# Install Node.js directly from nodejs.org with SHA256 verification,
# instead of piping NodeSource's setup_18.x apt-repo installer into
# instead of piping NodeSource's setup_24.x apt-repo installer into
# sudo bash (which runs a mutable upstream script unattended).
- run:
name: Install Node.js 18.20.8
name: Install Node.js 24.19.0
command: |
NODE_VERSION="18.20.8"
NODE_VERSION="24.19.0"
NODE_TARBALL="node-v${NODE_VERSION}-linux-x64.tar.xz"
NODE_EXPECTED_SHA="5467ee62d6af1411d46b6a10e3fb5cacc92734dbcef465fea14e7b90993001c9"
NODE_EXPECTED_SHA="14b342e71204f811bde6153be8e04b62aef63c236fef92b55f9c83154b409647"
curl -sSLf -o "/tmp/${NODE_TARBALL}" "https://nodejs.org/dist/v${NODE_VERSION}/${NODE_TARBALL}"
echo "${NODE_EXPECTED_SHA} /tmp/${NODE_TARBALL}" | sha256sum -c -
sudo tar -xJf "/tmp/${NODE_TARBALL}" -C /usr/local --strip-components=1
@ -2527,6 +2646,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2614,7 +2734,7 @@ jobs:
ui_build:
docker:
- image: cimg/node:20.19@sha256:35e64883e8d21bc345b0a7b04c35ee46442c127607ed1d8d7d37d8a1ed76db81
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
@ -2658,7 +2778,7 @@ jobs:
ui_unit_tests:
docker:
- image: cimg/node:20.19@sha256:35e64883e8d21bc345b0a7b04c35ee46442c127607ed1d8d7d37d8a1ed76db81
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
@ -2717,6 +2837,7 @@ jobs:
category: client
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -2859,6 +2980,7 @@ jobs:
category: client
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -3031,6 +3153,8 @@ workflows:
only:
- main
- /litellm_.*/
- base_sdk_install:
filters: *main_branches
- local_testing_part1:
filters: *main_branches
- local_testing_part2:

2
.github/workflows/check-ui-api-types.yml vendored
View file

@ -56,7 +56,7 @@ jobs:
- name: Set up Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version: "20"
node-version-file: ui/litellm-dashboard/.nvmrc
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json

33
.github/workflows/image-scan.yml vendored
View file

@ -9,6 +9,8 @@ on:
- "litellm_**"
paths:
- docker/Dockerfile.non_root
- migrations/Dockerfile
- migrations/run.py
- tests/proxy_migration_tests/test_offline_image_migration.py
- uv.lock
- ui/litellm-dashboard/package-lock.json
@ -83,3 +85,34 @@ jobs:
--only-fixed \
--fail-on high \
--output table
migrations-image:
name: migrations-image
runs-on: ubuntu-latest
if: >-
github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository
timeout-minutes: 30
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build migrations image
run: docker build -f migrations/Dockerfile -t litellm-migrations-scan:${{ github.sha }} .
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Verify offline migration as a non-root uid
env:
LITELLM_IMAGE: litellm-migrations-scan:${{ github.sha }}
LITELLM_MIGRATION_INTERPRETER: python3
LITELLM_MIGRATION_SCRIPT: /app/run.py
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v

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

@ -7,13 +7,17 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
code-quality:

2
.github/workflows/test-litellm-ui-build.yml vendored
View file

@ -27,7 +27,7 @@ jobs:
- name: Setup Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version: "20"
node-version-file: ui/litellm-dashboard/.nvmrc
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json

2
.github/workflows/test-litellm-ui-lint.yml vendored
View file

@ -61,7 +61,7 @@ jobs:
if: steps.changed.outputs.has_files == 'true'
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version: "20"
node-version-file: ui/litellm-dashboard/.nvmrc
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json

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

@ -35,7 +35,7 @@ jobs:
- name: Setup Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version: "20"
node-version-file: ui/litellm-dashboard/.nvmrc
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json

8
.github/workflows/test-unit-core-utils.yml vendored
View file

@ -7,6 +7,10 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -14,8 +18,8 @@ permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
core-utils:

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

@ -7,13 +7,17 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
documentation:

8
.github/workflows/test-unit-enterprise-routing.yml vendored
View file

@ -7,6 +7,10 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -14,8 +18,8 @@ permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
enterprise-routing:

8
.github/workflows/test-unit-integrations.yml vendored
View file

@ -7,6 +7,10 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -14,8 +18,8 @@ permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
integrations:

8
.github/workflows/test-unit-llm-providers.yml vendored
View file

@ -7,13 +7,17 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
vertex-ai:

8
.github/workflows/test-unit-misc.yml vendored
View file

@ -7,6 +7,10 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -14,8 +18,8 @@ permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
misc:

8
.github/workflows/test-unit-proxy-auth.yml vendored
View file

@ -7,6 +7,10 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -14,8 +18,8 @@ permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-auth:

8
.github/workflows/test-unit-proxy-db.yml vendored
View file

@ -7,13 +7,17 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
# Semantic matrix: each shard groups tests by concern (auth, server, logging, …)
# rather than alphabetical letter ranges. Adding a new test file means adding it

8
.github/workflows/test-unit-proxy-endpoints.yml vendored
View file

@ -7,14 +7,18 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-endpoints:

8
.github/workflows/test-unit-proxy-infra.yml vendored
View file

@ -7,6 +7,10 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -14,8 +18,8 @@ permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-infra:

8
.github/workflows/test-unit-proxy-legacy.yml vendored
View file

@ -7,13 +7,17 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
test:

8
.github/workflows/test-unit-responses-caching-types.yml vendored
View file

@ -7,6 +7,10 @@ on:
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
@ -14,8 +18,8 @@ permissions:
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
responses-caching-types:

2
Dockerfile
View file

@ -7,7 +7,7 @@ ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
FROM $UV_IMAGE AS uvbin

4
Makefile
View file

@ -75,7 +75,7 @@ install-dev:
bootstrap:
$(UV) sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
cd ui/litellm-dashboard && npm ci --no-audit --no-fund
cd ui/litellm-dashboard && npm install --no-audit --no-fund
@main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \
if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \
cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \
@ -239,7 +239,7 @@ lint-dev: lint-format-changed check-circular-imports check-import-safety
# test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and
# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage.
# Not auto-installed as a git hook so it never slows an unrelated human commit.
pre-commit:
pre-commit: bootstrap
./scripts/pre_commit_lint.sh
# Testing targets

4
backend/Dockerfile
View file

@ -63,6 +63,8 @@ RUN mkdir -p /home/nonroot && \
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
chown -R nonroot:nonroot /home/nonroot/.cache
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
# ---------- Runtime ----------
FROM $LITELLM_RUNTIME_IMAGE AS runtime
@ -93,5 +95,5 @@ USER nonroot
EXPOSE 4001/tcp
ENTRYPOINT ["uvicorn", "backend.main:app"]
ENTRYPOINT ["/app/docker/component_entrypoint.sh", "uvicorn", "backend.main:app"]
CMD ["--host", "0.0.0.0", "--port", "4001"]

30
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 31903
"limit": 29813
},
"reportArgumentType": {
"limit": 2645
@ -21,10 +21,10 @@
"limit": 325
},
"reportDuplicateImport": {
"limit": 42
"limit": 38
},
"reportExplicitAny": {
"limit": 10214
"limit": 9473
},
"reportFunctionMemberAccess": {
"limit": 11
@ -33,7 +33,7 @@
"limit": 227
},
"reportIncompatibleMethodOverride": {
"limit": 78
"limit": 77
},
"reportIncompatibleVariableOverride": {
"limit": 12
@ -42,7 +42,7 @@
"limit": 18
},
"reportIndexIssue": {
"limit": 37
"limit": 35
},
"reportInvalidTypeForm": {
"limit": 35
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5869
"limit": 5855
},
"reportMissingTypeArgument": {
"limit": 15861
"limit": 15849
},
"reportMissingTypeStubs": {
"limit": 41
@ -84,7 +84,7 @@
"limit": 77
},
"reportPrivateUsage": {
"limit": 2437
"limit": 2436
},
"reportRedeclaration": {
"limit": 12
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45366
"limit": 45269
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40477
"limit": 40452
},
"reportUnknownParameterType": {
"limit": 20338
"limit": 20309
},
"reportUnknownVariableType": {
"limit": 32047
"limit": 31978
},
"reportUnnecessaryCast": {
"limit": 177
"limit": 173
},
"reportUnnecessaryComparison": {
"limit": 1021
"limit": 1017
},
"reportUnnecessaryContains": {
"limit": 7
},
"reportUnnecessaryIsInstance": {
"limit": 1205
"limit": 1203
},
"reportUntypedBaseClass": {
"limit": 165

2
docker/Dockerfile.database
View file

@ -7,7 +7,7 @@ ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b17a5ee8bc23b532600a44d705acef2409e0933c1251b45f
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
FROM $UV_IMAGE AS uvbin

2
docker/Dockerfile.non_root
View file

@ -6,7 +6,7 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:42df77a9974d6ec8b
ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:20.18-alpine3.20@sha256:3488b10bf958af7125a176419d2d8a9937d895bf124012aae811651988d2ffe6
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
FROM $UV_IMAGE AS uvbin

5
docker/build_admin_ui.sh
View file

@ -54,8 +54,9 @@ else
fi || { echo "nvm checksum verification failed"; exit 1; }
bash "$NVM_SCRIPT"
source ~/.nvm/nvm.sh
nvm install v18.17.0
nvm use v18.17.0
NODE_VERSION="$(cat ui/litellm-dashboard/.nvmrc)"
nvm install "v${NODE_VERSION}"
nvm use "v${NODE_VERSION}"
# cd in to /ui/litellm-dashboard

8
docker/component_entrypoint.sh Executable file
View file

@ -0,0 +1,8 @@
#!/bin/sh
if [ "$USE_DDTRACE" = "true" ]; then
export DD_TRACE_OPENAI_ENABLED="False"
exec ddtrace-run "$@"
fi
exec "$@"

33
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -17,6 +17,7 @@ if TYPE_CHECKING:
from litellm.proxy._types import LiteLLM_ManagedObjectTable
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
from litellm.types.router import Deployment
from litellm.types.utils import LiteLLMBatch
@ -281,6 +282,32 @@ class CheckBatchCost:
return deployment_id
return None
@classmethod
def _get_managed_file_model_name(
cls,
job: "LiteLLM_ManagedObjectTable",
deployment_info: "Deployment",
) -> Optional[str]:
"""
Public model group name to encode as ``target_model_names`` on unified output file ids.
Key model-access checks resolve a managed file id back to a model via its
``target_model_names``, so this must be the model group the caller requested, never the
underlying provider model (e.g. ``gpt-5.5``), which no key is allowed to call.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
convert_b64_uid_to_unified_uid,
get_models_from_unified_file_id,
)
input_file_id = cls._get_input_file_id(job)
target_model_names = (
get_models_from_unified_file_id(convert_b64_uid_to_unified_uid(input_file_id)) if input_file_id else []
)
if target_model_names:
return ",".join(target_model_names)
return deployment_info.model_name or None
@staticmethod
def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]:
import json
@ -406,6 +433,10 @@ class CheckBatchCost:
managed_files_hook = self.proxy_logging_obj.get_proxy_hook("managed_files")
if managed_files_hook is not None:
from litellm.proxy._types import UserAPIKeyAuth
managed_file_model_name = self._get_managed_file_model_name(
job=job, deployment_info=deployment_info
)
_minimal_auth = UserAPIKeyAuth(
user_id=job.created_by or "default-user-id",
team_id=getattr(job, "team_id", None),
@ -417,7 +448,7 @@ class CheckBatchCost:
_unified_file_id = managed_files_hook.get_unified_output_file_id(
output_file_id=_raw_file_id,
model_id=model_id,
model_name=str(model_name) if model_name else deployment_info.model_name or None,
model_name=managed_file_model_name,
)
await managed_files_hook.store_unified_file_id(
file_id=_unified_file_id,

8
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -215,7 +215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
if result:
return LiteLLM_ManagedFileTable(**result)
return LiteLLM_ManagedFileTable.model_validate(result)
## CHECK DB
db_object = await self.prisma_client.db.litellm_managedfiletable.find_first(
@ -223,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
if db_object:
return LiteLLM_ManagedFileTable(**db_object.model_dump())
return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump())
return None
async def delete_unified_file_id(
@ -349,7 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if isinstance(batch.file_object, str)
else batch.file_object
)
batch_obj = LiteLLMBatch(**batch_data)
batch_obj = LiteLLMBatch.model_validate(batch_data)
batch_obj.id = batch.unified_object_id
batch_objects.append(batch_obj)
@ -382,7 +382,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"flat_model_file_ids": {"hasSome": model_object_ids},
}
)
return [OpenAIFileObject(**file_object.file_object) for file_object in file_ids]
return [OpenAIFileObject.model_validate(file_object.file_object) for file_object in file_ids]
async def check_managed_file_id_access(
self, data: Dict, user_api_key_dict: UserAPIKeyAuth

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.52"
version = "0.1.53"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.52"
version = "0.1.53"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
gateway/Dockerfile
View file

@ -65,6 +65,8 @@ RUN mkdir -p /home/nonroot && \
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
chown -R nonroot:nonroot /home/nonroot/.cache
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
# ---------- Runtime ----------
FROM $LITELLM_RUNTIME_IMAGE AS runtime
@ -95,5 +97,5 @@ USER nonroot
EXPOSE 4000/tcp
ENTRYPOINT ["sh", "-c", "exec uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
CMD ["--host", "0.0.0.0", "--port", "4000"]

2
helm/litellm-helm/templates/migrations-job.yaml
View file

@ -35,6 +35,8 @@ spec:
{{- toYaml . | nindent 8 }}
{{- end }}
serviceAccountName: {{ include "litellm.migrationServiceAccountName" . }}
securityContext:
{{- toYaml .Values.podSecurityContext | nindent 8 }}
{{- with .Values.migrationJob.extraInitContainers }}
initContainers:
{{- tpl (toYaml .) $ | nindent 8 }}

36
helm/litellm-helm/tests/migrations-job_tests.yaml
View file

@ -254,3 +254,39 @@ tests:
content:
name: sidecar-tpl
image: "ghcr.io/berriai/litellm-database:test"
- it: should render the pod-level securityContext from podSecurityContext
template: migrations-job.yaml
set:
migrationJob:
enabled: true
podSecurityContext:
fsGroup: 10000
runAsUser: 10000
runAsNonRoot: true
asserts:
- equal:
path: spec.template.spec.securityContext
value:
fsGroup: 10000
runAsUser: 10000
runAsNonRoot: true
- it: should keep the pod-level and container-level securityContext separate
template: migrations-job.yaml
set:
migrationJob:
enabled: true
podSecurityContext:
fsGroup: 10000
securityContext:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
asserts:
- equal:
path: spec.template.spec.securityContext
value:
fsGroup: 10000
- equal:
path: spec.template.spec.containers[0].securityContext
value:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true

53
helm/litellm/templates/_helpers.tpl
View file

@ -138,6 +138,59 @@ is false the chart uses the provided name, or the namespace `default` SA.
{{- end -}}
{{- end -}}
{{/*
ServiceAccount name for the migrations Job.
The Job is a pre-install / pre-upgrade hook, so it is created before the
chart's ordinary resources. A ServiceAccount the chart creates is one of
those ordinary resources, which makes borrowing the backend name a cycle:
the hook pod is rejected because the account does not exist yet. So when
`serviceAccounts.backend.create` is true the Job falls back to the namespace
`default` account unless the operator names one that already exists. With
`create` false the backend name is either an operator-supplied existing
account or `default`, both of which are safe for the hook, so the Job keeps
sharing it.
`migrationJob.serviceAccountName` always wins when set, which is how a Job
that needs credentials of its own (IRSA / Workload Identity for IAM database
auth) gets them.
*/}}
{{- define "litellm.migrations.serviceAccountName" -}}
{{- if .Values.migrationJob.serviceAccountName -}}
{{ .Values.migrationJob.serviceAccountName }}
{{- else if .Values.serviceAccounts.backend.create -}}
default
{{- else -}}
{{ include "litellm.backend.serviceAccountName" . }}
{{- end -}}
{{- end -}}
{{/*
Extra pod labels for a component's Deployment, validated against its selector.
Invoke with a dict:
(dict "podLabels" .Values.gateway.podLabels "componentName" "gateway")
The three selector keys are also emitted on the pod template, so a podLabels
entry reusing one renders a duplicate YAML key whose later value wins. That
leaves the pod template no longer matching the (immutable) selector and the
apiserver rejects the Deployment. Fail at template time naming the key
instead, so the operator gets the reason here rather than an opaque
`selector does not match template labels` from the apiserver.
The migrations Job takes podLabels unvalidated: a Job's selector is generated
by the controller rather than declared, so nothing there can collide.
*/}}
{{- define "litellm.podLabels" -}}
{{- $componentName := .componentName -}}
{{- range $key, $value := .podLabels }}
{{- if has $key (list "app.kubernetes.io/name" "app.kubernetes.io/instance" "app.kubernetes.io/component") }}
{{- fail (printf "%s.podLabels cannot set %s: it is part of the Deployment's immutable selector" $componentName $key) }}
{{- end }}
{{- end }}
{{- toYaml .podLabels }}
{{- end -}}
{{/*
Master-key + database + redis env block — shared by gateway, backend, and the
migrations Job.

22
helm/litellm/templates/backend/deployment.yaml
View file

@ -23,9 +23,16 @@ spec:
{{- end }}
labels:
{{- include "litellm.backend.selectorLabels" . | nindent 8 }}
{{- with .Values.backend.podLabels }}
{{- include "litellm.podLabels" (dict "podLabels" . "componentName" "backend") | nindent 8 }}
{{- end }}
spec:
serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }}
automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }}
{{- with .Values.backend.podSecurityContext }}
securityContext:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
@ -34,6 +41,10 @@ spec:
- name: backend
image: "{{ .Values.backend.image.repository }}:{{ .Values.backend.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.backend.image.pullPolicy }}
{{- with .Values.backend.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
ports:
- name: http
containerPort: 4001
@ -70,8 +81,15 @@ spec:
readinessProbe:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- with .Values.backend.lifecycle }}
lifecycle:
{{- toYaml . | nindent 12 }}
{{- end }}
resources:
{{- toYaml .Values.backend.resources | nindent 12 }}
{{- with .Values.backend.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
@ -102,4 +120,8 @@ spec:
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- $gracePeriod := .Values.backend.terminationGracePeriodSeconds }}
{{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }}
terminationGracePeriodSeconds: {{ $gracePeriod }}
{{- end }}
{{- end }}

22
helm/litellm/templates/gateway/deployment.yaml
View file

@ -21,9 +21,16 @@ spec:
{{- end }}
labels:
{{- include "litellm.gateway.selectorLabels" . | nindent 8 }}
{{- with .Values.gateway.podLabels }}
{{- include "litellm.podLabels" (dict "podLabels" . "componentName" "gateway") | nindent 8 }}
{{- end }}
spec:
serviceAccountName: {{ include "litellm.gateway.serviceAccountName" . }}
automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }}
{{- with .Values.gateway.podSecurityContext }}
securityContext:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
@ -32,6 +39,10 @@ spec:
- name: gateway
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
ports:
- name: http
containerPort: 4000
@ -72,8 +83,15 @@ spec:
readinessProbe:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- with .Values.gateway.lifecycle }}
lifecycle:
{{- toYaml . | nindent 12 }}
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
@ -104,4 +122,8 @@ spec:
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- $gracePeriod := .Values.gateway.terminationGracePeriodSeconds }}
{{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }}
terminationGracePeriodSeconds: {{ $gracePeriod }}
{{- end }}
{{- end }}

27
helm/litellm/templates/migrations-job.yaml
View file

@ -23,12 +23,21 @@ spec:
ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }}
template:
metadata:
{{- /* The Job's selector is generated by the controller rather than
declared, so podLabels may override a chart label here. Merge
instead of appending so an override replaces the key rather than
rendering it twice. */}}
{{- $chartLabels := merge (dict "app.kubernetes.io/component" "migrations") (fromYaml (include "litellm.commonLabels" .)) }}
labels:
{{- include "litellm.commonLabels" . | nindent 8 }}
app.kubernetes.io/component: migrations
{{- toYaml (merge (deepCopy .Values.migrationJob.podLabels) $chartLabels) | nindent 8 }}
spec:
restartPolicy: Never
serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }}
serviceAccountName: {{ include "litellm.migrations.serviceAccountName" . }}
automountServiceAccountToken: {{ .Values.migrationJob.automountServiceAccountToken }}
{{- with .Values.migrationJob.podSecurityContext }}
securityContext:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
@ -37,10 +46,22 @@ spec:
- name: prisma-migrations
image: "{{ .Values.migrationJob.image.repository }}:{{ .Values.migrationJob.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.migrationJob.image.pullPolicy }}
{{- with .Values.migrationJob.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
env:
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.migrationJob) | nindent 12 }}
{{- with .Values.migrationJob.volumeMounts }}
volumeMounts:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- with .Values.migrationJob.resources }}
resources:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- with .Values.migrationJob.volumes }}
volumes:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

22
helm/litellm/templates/ui/deployment.yaml
View file

@ -18,9 +18,16 @@ spec:
{{- end }}
labels:
{{- include "litellm.ui.selectorLabels" . | nindent 8 }}
{{- with .Values.ui.podLabels }}
{{- include "litellm.podLabels" (dict "podLabels" . "componentName" "ui") | nindent 8 }}
{{- end }}
spec:
serviceAccountName: {{ include "litellm.ui.serviceAccountName" . }}
automountServiceAccountToken: {{ .Values.serviceAccounts.ui.automount }}
{{- with .Values.ui.podSecurityContext }}
securityContext:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.imagePullSecrets }}
imagePullSecrets:
{{- toYaml . | nindent 8 }}
@ -29,6 +36,10 @@ spec:
- name: ui
image: "{{ .Values.ui.image.repository }}:{{ .Values.ui.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.ui.image.pullPolicy }}
{{- with .Values.ui.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
ports:
- name: http
containerPort: 3000
@ -58,8 +69,15 @@ spec:
readinessProbe:
{{- toYaml . | nindent 12 }}
{{- end }}
{{- with .Values.ui.lifecycle }}
lifecycle:
{{- toYaml . | nindent 12 }}
{{- end }}
resources:
{{- toYaml .Values.ui.resources | nindent 12 }}
{{- with .Values.ui.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- with .Values.ui.volumes }}
volumes:
{{- toYaml . | nindent 8 }}
@ -80,4 +98,8 @@ spec:
topologySpreadConstraints:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- $gracePeriod := .Values.ui.terminationGracePeriodSeconds }}
{{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }}
terminationGracePeriodSeconds: {{ $gracePeriod }}
{{- end }}
{{- end }}

169
helm/litellm/tests/migration_job_tests.yaml Normal file
View file

@ -0,0 +1,169 @@
suite: test migrations Job ServiceAccount resolution and pod hardening
templates:
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
- it: borrows the namespace default account when no ServiceAccount is configured
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: default
- it: falls back to the namespace default account when the chart creates the backend ServiceAccount
set:
serviceAccounts.backend.create: true
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: default
- notEqual:
path: spec.template.spec.serviceAccountName
value: RELEASE-NAME-litellm-backend
- it: keeps sharing an existing backend ServiceAccount the chart does not create
set:
serviceAccounts.backend.create: false
serviceAccounts.backend.name: existing-backend-sa
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: existing-backend-sa
- it: prefers an explicit migration ServiceAccount over the created backend one
set:
serviceAccounts.backend.create: true
migrationJob.serviceAccountName: migrations-sa
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: migrations-sa
- it: prefers an explicit migration ServiceAccount over an existing backend one
set:
serviceAccounts.backend.create: false
serviceAccounts.backend.name: existing-backend-sa
migrationJob.serviceAccountName: migrations-sa
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: migrations-sa
- it: mounts no ServiceAccount token by default
asserts:
- equal:
path: spec.template.spec.automountServiceAccountToken
value: false
- it: mounts a ServiceAccount token when the operator asks for one
set:
migrationJob.automountServiceAccountToken: true
asserts:
- equal:
path: spec.template.spec.automountServiceAccountToken
value: true
- it: keeps the token off the Job when the backend disables automounting
set:
serviceAccounts.backend.create: true
serviceAccounts.backend.automount: false
asserts:
- equal:
path: spec.template.spec.serviceAccountName
value: default
- equal:
path: spec.template.spec.automountServiceAccountToken
value: false
- it: renders no hardening fields by default
asserts:
- isNull:
path: spec.template.spec.securityContext
- isNull:
path: spec.template.spec.containers[0].securityContext
- isNull:
path: spec.template.spec.volumes
- isNull:
path: spec.template.spec.containers[0].volumeMounts
- equal:
path: spec.template.metadata.labels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/managed-by: Helm
helm.sh/chart: litellm-0.1.0
app.kubernetes.io/component: migrations
- it: renders pod-level and container-level securityContext in their own scopes
set:
migrationJob.podSecurityContext:
runAsNonRoot: true
runAsUser: 65532
seccompProfile:
type: RuntimeDefault
migrationJob.securityContext:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
capabilities:
drop:
- ALL
asserts:
- equal:
path: spec.template.spec.securityContext
value:
runAsNonRoot: true
runAsUser: 65532
seccompProfile:
type: RuntimeDefault
- equal:
path: spec.template.spec.containers[0].securityContext
value:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
capabilities:
drop:
- ALL
- it: renders volumes on the pod and volumeMounts on the migration container
set:
migrationJob.volumes:
- name: tmp
emptyDir:
sizeLimit: 64Mi
migrationJob.volumeMounts:
- name: tmp
mountPath: /tmp
asserts:
- equal:
path: spec.template.spec.volumes
value:
- name: tmp
emptyDir:
sizeLimit: 64Mi
- equal:
path: spec.template.spec.containers[0].volumeMounts
value:
- name: tmp
mountPath: /tmp
- it: merges podLabels with the chart labels on the Job pod
set:
migrationJob.podLabels:
egress-policy: restricted
asserts:
- equal:
path: spec.template.metadata.labels['egress-policy']
value: restricted
- equal:
path: spec.template.metadata.labels['app.kubernetes.io/component']
value: migrations
- it: accepts a podLabel that reuses a chart label, since the Job selector is controller-generated
set:
migrationJob.podLabels:
app.kubernetes.io/component: batch-migrations
asserts:
- notFailedTemplate: {}
- equal:
path: spec.template.metadata.labels['app.kubernetes.io/component']
value: batch-migrations

298
helm/litellm/tests/pod_hardening_tests.yaml Normal file
View file

@ -0,0 +1,298 @@
suite: test pod hardening knobs on the component deployments
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- ui/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: gateway renders no hardening fields by default
template: gateway/deployment.yaml
asserts:
- isNull:
path: spec.template.spec.securityContext
- isNull:
path: spec.template.spec.containers[0].securityContext
- isNull:
path: spec.template.spec.containers[0].lifecycle
- isNull:
path: spec.template.spec.terminationGracePeriodSeconds
- lengthEqual:
path: spec.template.spec.containers
count: 1
- equal:
path: spec.template.metadata.labels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
- it: gateway renders pod-level and container-level securityContext in their own scopes
template: gateway/deployment.yaml
set:
gateway.podSecurityContext:
runAsNonRoot: true
runAsUser: 65532
fsGroup: 65532
seccompProfile:
type: RuntimeDefault
gateway.securityContext:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
capabilities:
drop:
- ALL
asserts:
- equal:
path: spec.template.spec.securityContext
value:
runAsNonRoot: true
runAsUser: 65532
fsGroup: 65532
seccompProfile:
type: RuntimeDefault
- equal:
path: spec.template.spec.containers[0].securityContext
value:
allowPrivilegeEscalation: false
readOnlyRootFilesystem: true
capabilities:
drop:
- ALL
- it: gateway merges podLabels with the selector labels
template: gateway/deployment.yaml
set:
gateway.podLabels:
egress-policy: restricted
team: platform
asserts:
- equal:
path: spec.template.metadata.labels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
egress-policy: restricted
team: platform
- equal:
path: spec.selector.matchLabels
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
- it: gateway rejects a podLabel that collides with the selector
template: gateway/deployment.yaml
set:
gateway.podLabels:
app.kubernetes.io/component: not-gateway
asserts:
- failedTemplate:
errorMessage: "gateway.podLabels cannot set app.kubernetes.io/component: it is part of the Deployment's immutable selector"
- it: backend rejects a podLabel that collides with the selector
template: backend/deployment.yaml
set:
backend.podLabels:
app.kubernetes.io/name: not-litellm
asserts:
- failedTemplate:
errorMessage: "backend.podLabels cannot set app.kubernetes.io/name: it is part of the Deployment's immutable selector"
- it: ui rejects a podLabel that collides with the selector
template: ui/deployment.yaml
set:
ui.podLabels:
app.kubernetes.io/instance: not-the-release
asserts:
- failedTemplate:
errorMessage: "ui.podLabels cannot set app.kubernetes.io/instance: it is part of the Deployment's immutable selector"
- it: gateway renders lifecycle hooks on the container
template: gateway/deployment.yaml
set:
gateway.lifecycle:
preStop:
httpGet:
path: /health/drain
port: 4000
asserts:
- equal:
path: spec.template.spec.containers[0].lifecycle
value:
preStop:
httpGet:
path: /health/drain
port: 4000
- it: gateway renders terminationGracePeriodSeconds on the pod spec
template: gateway/deployment.yaml
set:
gateway.terminationGracePeriodSeconds: 90
asserts:
- equal:
path: spec.template.spec.terminationGracePeriodSeconds
value: 90
- it: gateway honors an explicit terminationGracePeriodSeconds of zero
template: gateway/deployment.yaml
set:
gateway.terminationGracePeriodSeconds: 0
asserts:
- equal:
path: spec.template.spec.terminationGracePeriodSeconds
value: 0
- it: gateway appends extraContainers after the gateway container
template: gateway/deployment.yaml
set:
gateway.extraContainers:
- name: auth-sidecar
image: registry.example.com/auth-proxy:1.2.3
args:
- --upstream
- http://127.0.0.1:4000
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 2
- equal:
path: spec.template.spec.containers[0].name
value: gateway
- equal:
path: spec.template.spec.containers[1]
value:
name: auth-sidecar
image: registry.example.com/auth-proxy:1.2.3
args:
- --upstream
- http://127.0.0.1:4000
- it: gateway templates chart context inside extraContainers
template: gateway/deployment.yaml
set:
gateway.extraContainers:
- name: auth-sidecar
image: registry.example.com/auth-proxy:1.2.3
env:
- name: RELEASE
value: "{{ .Release.Name }}"
asserts:
- equal:
path: spec.template.spec.containers[1].env[0].value
value: RELEASE-NAME
- it: backend renders every hardening knob in the right scope
template: backend/deployment.yaml
set:
backend.podLabels:
egress-policy: restricted
backend.podSecurityContext:
runAsNonRoot: true
backend.securityContext:
readOnlyRootFilesystem: true
backend.lifecycle:
preStop:
exec:
command:
- sleep
- "5"
backend.terminationGracePeriodSeconds: 60
backend.extraContainers:
- name: auth-sidecar
image: registry.example.com/auth-proxy:1.2.3
asserts:
- equal:
path: spec.template.metadata.labels['egress-policy']
value: restricted
- equal:
path: spec.template.spec.securityContext
value:
runAsNonRoot: true
- equal:
path: spec.template.spec.containers[0].securityContext
value:
readOnlyRootFilesystem: true
- equal:
path: spec.template.spec.containers[0].lifecycle
value:
preStop:
exec:
command:
- sleep
- "5"
- equal:
path: spec.template.spec.terminationGracePeriodSeconds
value: 60
- equal:
path: spec.template.spec.containers[1].name
value: auth-sidecar
- it: ui renders every hardening knob in the right scope
template: ui/deployment.yaml
set:
ui.podLabels:
egress-policy: restricted
ui.podSecurityContext:
runAsNonRoot: true
fsGroup: 101
ui.securityContext:
readOnlyRootFilesystem: true
ui.lifecycle:
preStop:
exec:
command:
- /bin/sh
- -c
- nginx -s quit
ui.terminationGracePeriodSeconds: 30
ui.extraContainers:
- name: auth-sidecar
image: registry.example.com/auth-proxy:1.2.3
asserts:
- equal:
path: spec.template.metadata.labels['egress-policy']
value: restricted
- equal:
path: spec.template.spec.securityContext
value:
runAsNonRoot: true
fsGroup: 101
- equal:
path: spec.template.spec.containers[0].securityContext
value:
readOnlyRootFilesystem: true
- equal:
path: spec.template.spec.containers[0].lifecycle
value:
preStop:
exec:
command:
- /bin/sh
- -c
- nginx -s quit
- equal:
path: spec.template.spec.terminationGracePeriodSeconds
value: 30
- equal:
path: spec.template.spec.containers[1].name
value: auth-sidecar
- it: backend and ui render no hardening fields by default
templates:
- backend/deployment.yaml
- ui/deployment.yaml
asserts:
- isNull:
path: spec.template.spec.securityContext
- isNull:
path: spec.template.spec.containers[0].securityContext
- isNull:
path: spec.template.spec.containers[0].lifecycle
- isNull:
path: spec.template.spec.terminationGracePeriodSeconds
- lengthEqual:
path: spec.template.spec.containers
count: 1

106
helm/litellm/tests/probe_tests.yaml Normal file
View file

@ -0,0 +1,106 @@
suite: test liveness and readiness probe timeouts
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: gateway probes set an explicit timeout that outlasts a saturated event loop
template: gateway/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].livenessProbe
value:
httpGet:
path: /health/liveliness
port: http
initialDelaySeconds: 10
periodSeconds: 15
timeoutSeconds: 10
failureThreshold: 6
- equal:
path: spec.template.spec.containers[0].readinessProbe
value:
httpGet:
path: /health/readiness
port: http
initialDelaySeconds: 5
periodSeconds: 10
timeoutSeconds: 10
- it: backend probes set an explicit timeout that outlasts a saturated event loop
template: backend/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].livenessProbe
value:
httpGet:
path: /health/liveliness
port: http
initialDelaySeconds: 10
periodSeconds: 15
timeoutSeconds: 10
failureThreshold: 6
- equal:
path: spec.template.spec.containers[0].readinessProbe
value:
httpGet:
path: /health/readiness
port: http
initialDelaySeconds: 5
periodSeconds: 10
timeoutSeconds: 10
- it: no single-event-loop component is left on the kubernetes default 1s probe timeout
templates:
- gateway/deployment.yaml
- backend/deployment.yaml
asserts:
- isNotNullOrEmpty:
path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds
- isNotNullOrEmpty:
path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds
- equal:
path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds
value: 10
- equal:
path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds
value: 10
- it: gateway liveness tolerates a longer outage than readiness before acting
template: gateway/deployment.yaml
asserts:
- equal:
path: spec.template.spec.containers[0].livenessProbe.failureThreshold
value: 6
- notExists:
path: spec.template.spec.containers[0].readinessProbe.failureThreshold
- it: probe timeouts and thresholds stay overridable per component
template: gateway/deployment.yaml
set:
gateway.readinessProbe.timeoutSeconds: 3
gateway.readinessProbe.periodSeconds: 20
gateway.livenessProbe.timeoutSeconds: 4
gateway.livenessProbe.failureThreshold: 3
asserts:
- equal:
path: spec.template.spec.containers[0].readinessProbe
value:
httpGet:
path: /health/readiness
port: http
initialDelaySeconds: 5
periodSeconds: 20
timeoutSeconds: 3
- equal:
path: spec.template.spec.containers[0].livenessProbe
value:
httpGet:
path: /health/liveliness
port: http
initialDelaySeconds: 10
periodSeconds: 15
timeoutSeconds: 4
failureThreshold: 3

90
helm/litellm/values.yaml
View file

@ -57,6 +57,42 @@ migrationJob:
backoffLimit: 4
ttlSecondsAfterFinished: 120
resources: {}
# ServiceAccount for the Job pod only.
#
# The Job is a pre-install / pre-upgrade hook, so it runs before the chart's
# ordinary resources exist. With `serviceAccounts.backend.create: true` the
# backend ServiceAccount is one of those ordinary resources, so a Job that
# borrowed its name would reference an account that does not exist yet and
# the first install would fail with a forbidden pod creation. The name set
# here always wins; when it is empty the Job falls back to `default` if the
# chart creates the backend ServiceAccount, and to the backend
# ServiceAccount name otherwise (that name is either an existing account you
# supplied or `default`).
#
# Point this at a pre-existing ServiceAccount when the Job needs credentials
# of its own, e.g. the IRSA / Workload Identity annotations that
# `database.writer.useIAMAuth` relies on. That is also the upgrade path to
# watch: a release already running with `serviceAccounts.backend.create:
# true` used to hand the Job the created backend account on every upgrade,
# and now hands it `default` unless you name an account here.
serviceAccountName: ""
# The Job runs `prisma migrate deploy` against Postgres and never calls the
# K8s API, so it defaults to no projected ServiceAccount token, the same
# reasoning the ui SA above uses. Flip to true if your Job genuinely needs
# one; IAM database auth does not, since EKS Pod Identity injects its own
# projected token volume and GKE Workload Identity goes through the
# metadata server, neither of which is the default token mount.
automountServiceAccountToken: false
# Standard k8s pod-level and container-level securityContext for the Job
# pod. Same shape as gateway.podSecurityContext / gateway.securityContext.
podSecurityContext: {}
securityContext: {}
# Extra pod labels on the Job pod, merged into the chart's common labels.
podLabels: {}
# Additional volumes on the Job pod and volumeMounts on its container, e.g.
# the writable scratch space a read-only root filesystem needs.
volumes: []
volumeMounts: []
image:
repository: ghcr.io/berriai/litellm-migrations
tag: "" # defaults to .Chart.AppVersion
@ -180,10 +216,13 @@ gateway:
httpGet: { path: /health/liveliness, port: http }
initialDelaySeconds: 10
periodSeconds: 15
timeoutSeconds: 10
failureThreshold: 6
readinessProbe:
httpGet: { path: /health/readiness, port: http }
initialDelaySeconds: 5
periodSeconds: 10
timeoutSeconds: 10
hpa:
enabled: true
minReplicas: 1
@ -200,6 +239,37 @@ gateway:
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
# Extra pod labels, merged into the chart's selector labels. Do not
# re-declare `app.kubernetes.io/name` / `instance` / `component` here: they
# form the Deployment's immutable selector.
podLabels: {}
# Pod-level securityContext, applied to every container in the pod
# (runAsNonRoot, runAsUser, fsGroup, seccompProfile, ...). Empty by default
# so the cluster's own defaults keep applying to existing installs; clusters
# enforcing a restricted Pod Security Standard usually want at least
# `runAsNonRoot: true` and `seccompProfile.type: RuntimeDefault`.
podSecurityContext: {}
# Container-level securityContext for the gateway container. Empty by
# default for the same reason. Example:
# allowPrivilegeEscalation: false
# readOnlyRootFilesystem: true
# capabilities:
# drop:
# - ALL
# `readOnlyRootFilesystem: true` needs writable scratch space; supply it
# through `volumes` / `volumeMounts` above rather than expecting the chart
# to guess the paths your workload writes to.
securityContext: {}
# Extra sidecar containers appended to the gateway pod, e.g. an auth or
# egress proxy. Rendered through `tpl`, so entries may reference chart
# values and release metadata.
extraContainers: []
# Container lifecycle hooks (postStart / preStop) for the gateway container.
lifecycle: {}
# Grace period the kubelet allows between SIGTERM and SIGKILL. Leave empty
# to inherit the Kubernetes default of 30s. Set it a few seconds above the
# proxy's GRACEFUL_SHUTDOWN_TIMEOUT when you use a draining preStop hook.
terminationGracePeriodSeconds: ""
nodeSelector: {}
tolerations: []
affinity: {}
@ -242,10 +312,13 @@ backend:
httpGet: { path: /health/liveliness, port: http }
initialDelaySeconds: 10
periodSeconds: 15
timeoutSeconds: 10
failureThreshold: 6
readinessProbe:
httpGet: { path: /health/readiness, port: http }
initialDelaySeconds: 5
periodSeconds: 10
timeoutSeconds: 10
hpa:
enabled: true
minReplicas: 1
@ -257,6 +330,13 @@ backend:
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
# Same shape as the gateway blocks of the same name.
podLabels: {}
podSecurityContext: {}
securityContext: {}
extraContainers: []
lifecycle: {}
terminationGracePeriodSeconds: ""
nodeSelector: {}
tolerations: []
affinity: {}
@ -310,6 +390,16 @@ ui:
minAvailable: ""
maxUnavailable: ""
podAnnotations: {}
# Same shape as the gateway blocks of the same name. The nginx runtime
# writes its pid, cache, and proxy temp files under the image's root
# filesystem, so `securityContext.readOnlyRootFilesystem: true` here needs
# emptyDir volumes mounted over those paths.
podLabels: {}
podSecurityContext: {}
securityContext: {}
extraContainers: []
lifecycle: {}
terminationGracePeriodSeconds: ""
nodeSelector: {}
tolerations: []
affinity: {}

17
litellm-proxy-extras/litellm_proxy_extras/migrations/20260731000000_add_autorouter_savings_spend/migration.sql Normal file
View file

@ -0,0 +1,17 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

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

@ -748,6 +748,7 @@ model LiteLLM_DailyUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -782,6 +783,7 @@ model LiteLLM_DailyOrganizationSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -816,6 +818,7 @@ model LiteLLM_DailyEndUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -849,6 +852,7 @@ model LiteLLM_DailyAgentSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -882,6 +886,7 @@ model LiteLLM_DailyTeamSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -917,6 +922,7 @@ model LiteLLM_DailyTagSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)

3
litellm/__init__.py
View file

@ -264,6 +264,7 @@ databricks_key: Optional[str] = None
openai_like_key: Optional[str] = None
azure_key: Optional[str] = None
anthropic_key: Optional[str] = None
autorouter_savings_baseline_model: Optional[str] = None
replicate_key: Optional[str] = None
bytez_key: Optional[str] = None
gdc_key: Optional[str] = None
@ -449,6 +450,8 @@ enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
custom_prometheus_metadata_labels: List[str] = []
custom_prometheus_tags: List[str] = []
prometheus_metrics_config: Optional[List] = None
prometheus_exclude_metrics: Optional[List[str]] = None
prometheus_exclude_labels: Optional[List[str]] = None
prometheus_emit_stream_label: bool = False
# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on
# `litellm_proxy_failed_requests_metric`. Off by default to preserve the

53
litellm/_lazy_imports.py
View file

@ -17,39 +17,40 @@ until they're actually needed.
import importlib
import sys
from typing import Any, Optional, cast, Callable
from collections.abc import Callable
from typing import Any, cast
# Import all the data structures that define what can be lazy-loaded
# These are just lists of names and maps of where to find them
from ._lazy_imports_registry import (
# Name tuples
COST_CALCULATOR_NAMES,
LITELLM_LOGGING_NAMES,
UTILS_NAMES,
TOKEN_COUNTER_NAMES,
LLM_CLIENT_CACHE_NAMES,
BEDROCK_TYPES_NAMES,
TYPES_UTILS_NAMES,
CACHING_NAMES,
HTTP_HANDLER_NAMES,
DOTPROMPT_NAMES,
LLM_CONFIG_NAMES,
TYPES_NAMES,
LLM_PROVIDER_LOGIC_NAMES,
UTILS_MODULE_NAMES,
# Import maps
_UTILS_IMPORT_MAP,
_COST_CALCULATOR_IMPORT_MAP,
_TYPES_UTILS_IMPORT_MAP,
_TOKEN_COUNTER_IMPORT_MAP,
_BEDROCK_TYPES_IMPORT_MAP,
_CACHING_IMPORT_MAP,
_LITELLM_LOGGING_IMPORT_MAP,
_COST_CALCULATOR_IMPORT_MAP,
_DOTPROMPT_IMPORT_MAP,
_TYPES_IMPORT_MAP,
_LITELLM_LOGGING_IMPORT_MAP,
_LLM_CONFIGS_IMPORT_MAP,
_LLM_PROVIDER_LOGIC_IMPORT_MAP,
_TOKEN_COUNTER_IMPORT_MAP,
_TYPES_IMPORT_MAP,
_TYPES_UTILS_IMPORT_MAP,
_UTILS_IMPORT_MAP,
_UTILS_MODULE_IMPORT_MAP,
# Name tuples
BEDROCK_TYPES_NAMES,
CACHING_NAMES,
COST_CALCULATOR_NAMES,
DOTPROMPT_NAMES,
HTTP_HANDLER_NAMES,
LITELLM_LOGGING_NAMES,
LLM_CLIENT_CACHE_NAMES,
LLM_CONFIG_NAMES,
LLM_PROVIDER_LOGIC_NAMES,
TOKEN_COUNTER_NAMES,
TYPES_NAMES,
TYPES_UTILS_NAMES,
UTILS_MODULE_NAMES,
UTILS_NAMES,
)
@ -77,7 +78,7 @@ def _get_utils_globals() -> dict:
# They're separate from the main lazy import system because they have specific use cases
# Lazy loader for default encoding - avoids importing heavy tiktoken library at startup
_default_encoding: Optional[Any] = None
_default_encoding: Any | None = None
def _get_default_encoding() -> Any:
@ -99,7 +100,7 @@ def _get_default_encoding() -> Any:
# Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time
_get_modified_max_tokens_func: Optional[Any] = None
_get_modified_max_tokens_func: Any | None = None
def _get_modified_max_tokens() -> Any:
@ -123,7 +124,7 @@ def _get_modified_max_tokens() -> Any:
# Lazy loader for token_counter to avoid importing token_counter module at module import time
_token_counter_new_func: Optional[Any] = None
_token_counter_new_func: Any | None = None
def _get_token_counter_new() -> Any:
@ -153,7 +154,7 @@ def _get_token_counter_new() -> Any:
# This registry maps attribute names (like "ModelResponse") to handler functions
# It's built once the first time someone accesses a lazy-loaded attribute
# Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...}
_LAZY_IMPORT_REGISTRY: Optional[dict[str, Callable[[str], Any]]] = None
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:

12
litellm/_logging.py
View file

@ -4,11 +4,11 @@ import os
import sys
from datetime import datetime
from logging import Formatter
from typing import Any, Dict, Optional
from typing import Any
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.secret_redaction import redact_string
set_verbose = False
@ -86,7 +86,7 @@ handler.setLevel(numeric_level)
handler.addFilter(_secret_filter)
def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
def _try_parse_json_message(message: str) -> dict[str, Any] | None:
"""
Try to parse a log message as JSON. Returns parsed dict if valid, else None.
Handles messages that are entirely valid JSON (e.g. json.dumps output).
@ -103,7 +103,7 @@ def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
return parsed
def _try_parse_embedded_python_dict(message: str) -> Optional[Dict[str, Any]]:
def _try_parse_embedded_python_dict(message: str) -> dict[str, Any] | None:
"""
Try to find and parse a Python dict repr (e.g. str(d) or repr(d)) embedded in
the message. Handles patterns like:
@ -149,7 +149,7 @@ _STANDARD_RECORD_ATTRS = _get_standard_record_attrs()
class JsonFormatter(Formatter):
def __init__(self):
super(JsonFormatter, self).__init__()
super().__init__()
def formatTime(self, record, datefmt=None):
# Use datetime to format the timestamp in ISO 8601 format
@ -158,7 +158,7 @@ class JsonFormatter(Formatter):
def format(self, record):
message_str = record.getMessage()
json_record: Dict[str, Any] = {
json_record: dict[str, Any] = {
"message": message_str,
"level": record.levelname,
"timestamp": self.formatTime(record),

59
litellm/_redis.py
View file

@ -12,7 +12,7 @@ import json
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
import os
from typing import Callable, List, Optional, Union
from collections.abc import Callable
import redis # type: ignore
import redis.asyncio as async_redis # type: ignore
@ -76,7 +76,7 @@ def _init_arg_names(cls: type) -> frozenset[str]:
)
def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]:
def _get_redis_url_kwargs(client: type | None = None) -> tuple[str, ...]:
"""Connection kwargs that redis-py forwards from ``from_url`` down to the connection.
``from_url`` is declared as ``(cls, url, **kwargs)``, so introspecting it yields no
@ -160,7 +160,7 @@ def _redis_kwargs_from_environment():
def create_gcp_iam_redis_connect_func(
service_account: str,
ssl_ca_certs: Optional[str] = None,
ssl_ca_certs: str | None = None,
) -> Callable:
"""
Creates a custom Redis connection function for GCP IAM authentication.
@ -203,9 +203,9 @@ def create_gcp_iam_redis_connect_func(
def _build_azure_credential(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
azure_client_id: str | None = None,
azure_tenant_id: str | None = None,
azure_client_secret: str | None = None,
):
"""
Build a long-lived Azure credential object.
@ -241,9 +241,9 @@ def _build_azure_credential(
def _generate_azure_ad_redis_token(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
azure_client_id: str | None = None,
azure_tenant_id: str | None = None,
azure_client_secret: str | None = None,
) -> str:
"""
One-shot helper that builds a credential and fetches a single Azure AD
@ -263,9 +263,9 @@ def _generate_azure_ad_redis_token(
def create_azure_ad_redis_connect_func(
azure_client_id: Optional[str] = None,
azure_tenant_id: Optional[str] = None,
azure_client_secret: Optional[str] = None,
azure_client_id: str | None = None,
azure_tenant_id: str | None = None,
azure_client_secret: str | None = None,
) -> Callable:
"""
Creates a custom Redis connection function for Azure AD authentication.
@ -369,7 +369,7 @@ def _get_redis_client_logic(**env_overrides):
**env_overrides,
}
_startup_nodes: Optional[Union[str, list]] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
_startup_nodes: str | list | None = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
"REDIS_CLUSTER_NODES"
)
@ -380,21 +380,21 @@ def _get_redis_client_logic(**env_overrides):
elif _startup_nodes is None:
redis_kwargs.pop("startup_nodes", None)
_sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
_sentinel_nodes: str | list | None = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore
"REDIS_SENTINEL_NODES"
)
if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str):
redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes)
_sentinel_password: Optional[str] = redis_kwargs.get("sentinel_password", None) or get_secret_str(
_sentinel_password: str | None = redis_kwargs.get("sentinel_password", None) or get_secret_str(
"REDIS_SENTINEL_PASSWORD"
)
if _sentinel_password is not None:
redis_kwargs["sentinel_password"] = _sentinel_password
_service_name: Optional[str] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
_service_name: str | None = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
"REDIS_SERVICE_NAME"
)
@ -465,9 +465,12 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs.pop("port", None)
redis_kwargs.pop("db", None)
redis_kwargs.pop("password", None)
elif "startup_nodes" in redis_kwargs and redis_kwargs["startup_nodes"] is not None:
pass
elif "sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None:
elif (
"startup_nodes" in redis_kwargs
and redis_kwargs["startup_nodes"] is not None
or "sentinel_nodes" in redis_kwargs
and redis_kwargs["sentinel_nodes"] is not None
):
pass
elif "host" not in redis_kwargs or redis_kwargs["host"] is None:
raise ValueError("Either 'host' or 'url' must be specified for redis.")
@ -477,7 +480,7 @@ def _get_redis_client_logic(**env_overrides):
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
_redis_cluster_nodes_in_env: Optional[str] = get_secret("REDIS_CLUSTER_NODES") # type: ignore
_redis_cluster_nodes_in_env: str | None = get_secret("REDIS_CLUSTER_NODES") # type: ignore
if _redis_cluster_nodes_in_env is not None:
try:
redis_kwargs["startup_nodes"] = json.loads(_redis_cluster_nodes_in_env)
@ -495,7 +498,7 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
if arg in args:
cluster_kwargs[arg] = redis_kwargs[arg]
new_startup_nodes: List[ClusterNode] = []
new_startup_nodes: list[ClusterNode] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
@ -587,9 +590,9 @@ def get_redis_client(**env_overrides):
def get_redis_async_client(
connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
connection_pool: async_redis.BlockingConnectionPool | None = None,
**env_overrides,
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
) -> async_redis.Redis | async_redis.RedisCluster:
redis_kwargs = _get_redis_client_logic(**env_overrides)
if "startup_nodes" in redis_kwargs:
@ -618,7 +621,7 @@ def get_redis_async_client(
username=os.environ.get("REDIS_USERNAME") or None,
)
new_startup_nodes: List[ClusterNode] = []
new_startup_nodes: list[ClusterNode] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
@ -649,7 +652,7 @@ def get_redis_async_client(
url_kwargs[arg] = redis_kwargs[arg]
else:
verbose_logger.debug(
"REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(arg)
"REDIS: ignoring argument: %s. Not an allowed async_redis.Redis.from_url arg.", arg
)
return async_redis.Redis.from_url(**url_kwargs)
@ -682,7 +685,7 @@ def get_redis_async_client(
def get_redis_connection_pool(
**env_overrides,
) -> Optional[async_redis.BlockingConnectionPool]:
) -> async_redis.BlockingConnectionPool | None:
redis_kwargs = _get_redis_client_logic(**env_overrides)
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
@ -804,6 +807,6 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
# Fallback to simple logging if rich is not available
masker = SensitiveDataMasker()
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
verbose_logger.info(f"Redis configuration: {masked_redis_kwargs}")
verbose_logger.info("Redis configuration: %s", masked_redis_kwargs)
except Exception as e:
verbose_logger.error(f"Error pretty printing Redis configuration: {e}")
verbose_logger.error("Error pretty printing Redis configuration: %s", e)

14
litellm/_redis_credential_provider.py
View file

@ -1,7 +1,7 @@
import asyncio
import threading
import time
from typing import Any, Dict, Optional, Tuple, Union
from typing import Any
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
@ -14,7 +14,7 @@ _GCP_IAM_TOKEN_TTL_SECONDS = 3300
# Module-level cache shared across all GCPIAMCredentialProvider instances for the
# same service account, so multiple Redis connections on the same pod share one token.
# Keyed by service_account → (token, expiry_monotonic_timestamp).
_token_cache: Dict[str, Tuple[str, float]] = {}
_token_cache: dict[str, tuple[str, float]] = {}
_token_cache_lock = threading.Lock()
@ -95,11 +95,11 @@ class GCPIAMCredentialProvider(CredentialProvider):
def __init__(self, gcp_service_account: str) -> None:
self._gcp_service_account = gcp_service_account
def get_credentials(self) -> Tuple[str]:
def get_credentials(self) -> tuple[str]:
token = _get_cached_gcp_iam_token(self._gcp_service_account)
return (token,)
async def get_credentials_async(self) -> Tuple[str]:
async def get_credentials_async(self) -> tuple[str]:
token = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account)
return (token,)
@ -115,17 +115,17 @@ class AzureADCredentialProvider(CredentialProvider):
fail authentication after the initial token expired (~1 hour TTL).
"""
def __init__(self, credential: Any, username: Optional[str] = None) -> None:
def __init__(self, credential: Any, username: str | None = None) -> None:
self._credential = credential
self._username = username
def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]:
def get_credentials(self) -> tuple[str] | tuple[str, str]:
token = self._credential.get_token(AZURE_REDIS_SCOPE).token
if self._username:
return (self._username, token)
return (token,)
async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]:
async def get_credentials_async(self) -> tuple[str] | tuple[str, str]:
token_obj = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE)
if self._username:
return (self._username, token_obj.token)

38
litellm/_service_logger.py
View file

@ -1,6 +1,6 @@
import asyncio
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Union
import litellm
from litellm._logging import verbose_logger
@ -24,7 +24,7 @@ else:
UserAPIKeyAuth = Any
def _get_otel_v2_class() -> Optional[type]:
def _get_otel_v2_class() -> type | None:
"""Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent.
Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry
@ -54,7 +54,7 @@ class ServiceLogging(CustomLogger):
if "prometheus_system" in litellm.service_callback:
self.prometheusServicesLogger = PrometheusServicesLogger()
def _resolve_otel_service_logger(self, callback: Any) -> Optional[Any]:
def _resolve_otel_service_logger(self, callback: Any) -> Any | None:
"""Resolve the OTel logger (legacy or V2) to emit a service span on.
Returns the logger instance whose ``async_service_*_hook`` should fire for
@ -88,9 +88,9 @@ class ServiceLogging(CustomLogger):
service: ServiceTypes,
duration: float,
call_type: str,
parent_otel_span: Optional[Span] = None,
start_time: Optional[Union[datetime, float]] = None,
end_time: Optional[Union[float, datetime]] = None,
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: float | datetime | None = None,
):
"""
Handles both sync and async monitoring by checking for existing event loop.
@ -152,10 +152,10 @@ class ServiceLogging(CustomLogger):
service: ServiceTypes,
call_type: str,
duration: float,
parent_otel_span: Optional[Span] = None,
start_time: Optional[Union[datetime, float]] = None,
end_time: Optional[Union[datetime, float]] = None,
event_metadata: Optional[dict] = None,
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: datetime | float | None = None,
event_metadata: dict | None = None,
):
"""
- For counting if the redis, postgres call is successful
@ -218,7 +218,6 @@ class ServiceLogging(CustomLogger):
self.prometheusServicesLogger = PrometheusServicesLogger()
elif self.prometheusServicesLogger is None:
self.prometheusServicesLogger = self.prometheusServicesLogger()
return
async def init_datadog_logger_if_none(self):
"""
@ -230,8 +229,6 @@ class ServiceLogging(CustomLogger):
if not hasattr(self, "dd_logger"):
self.dd_logger: DataDogLogger = DataDogLogger()
return
async def init_otel_logger_if_none(self):
"""
initializes otel_logger if it is None or no attribute exists on ServiceLogging Object
@ -246,18 +243,17 @@ class ServiceLogging(CustomLogger):
verbose_logger.warning(
"ServiceLogger: open_telemetry_logger is None or not an instance of OpenTelemetry"
)
return
async def async_service_failure_hook(
self,
service: ServiceTypes,
duration: float,
error: Union[str, Exception],
error: str | Exception,
call_type: str,
parent_otel_span: Optional[Span] = None,
start_time: Optional[Union[datetime, float]] = None,
end_time: Optional[Union[float, datetime]] = None,
event_metadata: Optional[dict] = None,
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: float | datetime | None = None,
event_metadata: dict | None = None,
):
"""
- For counting if the redis, postgres call is unsuccessful
@ -324,7 +320,7 @@ class ServiceLogging(CustomLogger):
request_data: dict,
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: Optional[str] = None,
traceback_str: str | None = None,
):
"""
Hook to track failed litellm-service calls
@ -347,7 +343,7 @@ class ServiceLogging(CustomLogger):
pass
else:
raise Exception(
"Duration={} is not a float or timedelta object. type={}".format(_duration, type(_duration))
f"Duration={_duration} is not a float or timedelta object. type={type(_duration)}"
) # invalid _duration value
# Batch polling callbacks (check_batch_cost) don't include call_type in kwargs.
# Use .get() to avoid KeyError.

8
litellm/a2a_protocol/card_resolver.py
View file

@ -4,7 +4,7 @@ Custom A2A Card Resolver for LiteLLM.
Extends the A2A SDK's card resolver to support multiple well-known paths.
"""
from typing import TYPE_CHECKING, Any, Dict
from typing import TYPE_CHECKING, Any
from litellm._logging import verbose_logger
from litellm.constants import LOCALHOST_URL_PATTERNS
@ -114,7 +114,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
async def get_agent_card(
self,
relative_card_path: str | None = None,
http_kwargs: Dict[str, Any] | None = None,
http_kwargs: dict[str, Any] | None = None,
) -> "AgentCard":
"""
Fetch the agent card, trying multiple well-known paths.
@ -148,13 +148,13 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
last_error = None
for path in paths:
try:
verbose_logger.debug(f"Attempting to fetch agent card from {self.base_url}{path}")
verbose_logger.debug("Attempting to fetch agent card from %s%s", self.base_url, path)
return await super().get_agent_card(
relative_card_path=path,
http_kwargs=http_kwargs,
)
except Exception as e:
verbose_logger.debug(f"Failed to fetch agent card from {self.base_url}{path}: {e}")
verbose_logger.debug("Failed to fetch agent card from %s%s: %s", self.base_url, path, e)
last_error = e
continue

7
litellm/a2a_protocol/client.py
View file

@ -4,7 +4,8 @@ LiteLLM A2A Client class.
Provides a class-based interface for A2A agent invocation.
"""
from typing import TYPE_CHECKING, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING
from litellm.types.agents import LiteLLMSendMessageResponse
@ -50,7 +51,7 @@ class A2AClient:
self,
base_url: str,
timeout: float = 60.0,
extra_headers: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
):
"""
Initialize the A2A client wrapper.
@ -63,7 +64,7 @@ class A2AClient:
self.base_url = base_url
self.timeout = timeout
self.extra_headers = extra_headers
self._a2a_client: Optional["A2AClientType"] = None
self._a2a_client: A2AClientType | None = None
async def _get_client(self) -> "A2AClientType":
"""Get or create the underlying A2A client."""

8
litellm/a2a_protocol/cost_calculator.py
View file

@ -5,7 +5,7 @@ Supports dynamic cost parameters that allow platform owners
to define custom costs per agent query or per token.
"""
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import (
@ -18,7 +18,7 @@ else:
class A2ACostCalculator:
@staticmethod
def calculate_a2a_cost(
litellm_logging_obj: Optional[LitellmLoggingObject],
litellm_logging_obj: LitellmLoggingObject | None,
) -> float:
"""
Calculate the cost of an A2A send_message call.
@ -73,8 +73,8 @@ class A2ACostCalculator:
@staticmethod
def _calculate_token_based_cost(
model_call_details: dict,
input_cost_per_token: Optional[float],
output_cost_per_token: Optional[float],
input_cost_per_token: float | None,
output_cost_per_token: float | None,
) -> float:
"""
Calculate cost based on token usage and per-token pricing.

18
litellm/a2a_protocol/exception_mapping_utils.py
View file

@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils.
Maps A2A SDK exceptions to LiteLLM A2A exception types.
"""
from typing import TYPE_CHECKING, Any, Optional
from typing import TYPE_CHECKING, Any
from litellm._logging import verbose_logger
from litellm.a2a_protocol.card_resolver import (
@ -57,7 +57,7 @@ class A2AExceptionCheckers:
return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS)
@staticmethod
def is_localhost_url(url: Optional[str]) -> bool:
def is_localhost_url(url: str | None) -> bool:
"""
Check if a URL is a localhost/internal URL.
@ -96,9 +96,9 @@ class A2AExceptionCheckers:
def map_a2a_exception(
original_exception: Exception,
card_url: Optional[str] = None,
api_base: Optional[str] = None,
model: Optional[str] = None,
card_url: str | None = None,
api_base: str | None = None,
model: str | None = None,
) -> Exception:
"""
Map an A2A SDK exception to a LiteLLM A2A exception type.
@ -192,9 +192,11 @@ async def handle_a2a_localhost_retry(
request_type = "streaming " if is_streaming else ""
verbose_logger.warning(
f"A2A {request_type}request to '{error.localhost_url}' failed: {error.original_error}. "
f"Agent card contains localhost/internal URL. "
f"Retrying with base_url '{error.base_url}'."
"A2A %srequest to '%s' failed: %s. Agent card contains localhost/internal URL. Retrying with base_url '%s'.",
request_type,
error.localhost_url,
error.original_error,
error.base_url,
)
# Fix the agent card URL

36
litellm/a2a_protocol/exceptions.py
View file

@ -4,8 +4,6 @@ A2A Protocol Exceptions.
Custom exception types for A2A protocol operations, following LiteLLM's exception pattern.
"""
from typing import Optional
import httpx
@ -21,11 +19,11 @@ class A2AError(Exception):
message: str,
status_code: int = 500,
llm_provider: str = "a2a_agent",
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
model: str | None = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.status_code = status_code
self.message = f"litellm.A2AError: {message}"
@ -65,12 +63,12 @@ class A2AConnectionError(A2AError):
def __init__(
self,
message: str,
url: Optional[str] = None,
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
max_retries: Optional[int] = None,
num_retries: Optional[int] = None,
url: str | None = None,
model: str | None = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
):
self.url = url
super().__init__(
@ -98,10 +96,10 @@ class A2AAgentCardError(A2AError):
def __init__(
self,
message: str,
url: Optional[str] = None,
model: Optional[str] = None,
response: Optional[httpx.Response] = None,
litellm_debug_info: Optional[str] = None,
url: str | None = None,
model: str | None = None,
response: httpx.Response | None = None,
litellm_debug_info: str | None = None,
):
self.url = url
super().__init__(
@ -132,8 +130,8 @@ class A2ALocalhostURLError(A2AConnectionError):
self,
localhost_url: str,
base_url: str,
original_error: Optional[Exception] = None,
model: Optional[str] = None,
original_error: Exception | None = None,
model: str | None = None,
):
self.localhost_url = localhost_url
self.base_url = base_url

2
litellm/a2a_protocol/litellm_completion_bridge/__init__.py
View file

@ -16,8 +16,8 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
)
__all__ = [
"A2ACompletionBridgeTransformation",
"A2ACompletionBridgeHandler",
"A2ACompletionBridgeTransformation",
"handle_a2a_completion",
"handle_a2a_completion_streaming",
]

61
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -10,7 +10,8 @@ A2A Streaming Events (in order):
4. Status update (kind: "status-update") - Final status "completed" with final=true
"""
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
import litellm
from litellm._logging import verbose_logger
@ -46,13 +47,13 @@ class A2ACompletionBridgeHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Handle non-streaming A2A request via litellm.acompletion.
@ -75,7 +76,7 @@ class A2ACompletionBridgeHandler:
)
if a2a_provider_config is not None:
verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}")
verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider)
return await a2a_provider_config.handle_non_streaming(
request_id=request_id,
@ -102,10 +103,10 @@ class A2ACompletionBridgeHandler:
else:
full_model = model
verbose_logger.info(f"A2A completion bridge: model={full_model}, api_base={api_base}")
verbose_logger.info("A2A completion bridge: model=%s, api_base=%s", full_model, api_base)
# Build completion params dict
completion_params: Dict[str, Any] = {
completion_params: dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
@ -142,20 +143,20 @@ class A2ACompletionBridgeHandler:
request_id=request_id,
)
verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}")
verbose_logger.info("A2A completion bridge completed: request_id=%s", request_id)
return a2a_response
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
*,
_skip_a2a_provider_routing: bool = False,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming A2A request via litellm.acompletion with stream=True.
@ -184,7 +185,7 @@ class A2ACompletionBridgeHandler:
)
if a2a_provider_config is not None:
verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider} (streaming)")
verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider)
async for chunk in a2a_provider_config.handle_streaming(
request_id=request_id,
@ -220,10 +221,10 @@ class A2ACompletionBridgeHandler:
else:
full_model = model
verbose_logger.info(f"A2A completion bridge streaming: model={full_model}, api_base={api_base}")
verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
# Build completion params dict
completion_params: Dict[str, Any] = {
completion_params: dict[str, Any] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
@ -299,17 +300,19 @@ class A2ACompletionBridgeHandler:
)
yield completed_event
verbose_logger.info(f"A2A completion bridge streaming completed: request_id={request_id}, chunks={chunk_count}")
verbose_logger.info(
"A2A completion bridge streaming completed: request_id=%s, chunks=%s", request_id, chunk_count
)
# Convenience functions that delegate to the class methods
async def handle_a2a_completion(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""Convenience function for non-streaming A2A completion."""
return await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=request_id,
@ -322,11 +325,11 @@ async def handle_a2a_completion(
async def handle_a2a_completion_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
api_base: Optional[str] = None,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
params: dict[str, Any],
litellm_params: dict[str, Any],
api_base: str | None = None,
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
"""Convenience function for streaming A2A completion."""
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
request_id=request_id,

54
litellm/a2a_protocol/litellm_completion_bridge/transformation.py
View file

@ -18,7 +18,7 @@ A2A Streaming Events:
"""
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from typing import Any
from uuid import uuid4
from litellm._logging import verbose_logger
@ -30,7 +30,7 @@ class A2AStreamingContext:
Tracks task_id, context_id, and message accumulation.
"""
def __init__(self, request_id: str, input_message: Dict[str, Any]):
def __init__(self, request_id: str, input_message: dict[str, Any]):
self.request_id = request_id
self.task_id = str(uuid4())
self.context_id = str(uuid4())
@ -46,9 +46,9 @@ class A2ACompletionBridgeTransformation:
"""
@staticmethod
def _extract_text_from_a2a_parts(parts: List[Dict[str, Any]]) -> str:
def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str:
"""Extract text from A2A parts (with or without explicit ``kind``)."""
content_parts: List[str] = []
content_parts: list[str] = []
for part in parts:
if not isinstance(part, dict):
continue
@ -62,16 +62,16 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def get_forward_metadata(
a2a_message: Dict[str, Any],
params: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, Any]]:
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
"""
Merge A2A metadata from MessageSendParams and the message for downstream providers.
Forwarded once on the LangGraph run payload (``metadata``), not duplicated on
each input message — see ``apply_forward_metadata_to_completion_params``.
"""
merged: Dict[str, Any] = {}
merged: dict[str, Any] = {}
if params and isinstance(params.get("metadata"), dict):
merged.update(params["metadata"])
message_metadata = a2a_message.get("metadata")
@ -81,9 +81,9 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def apply_forward_metadata_to_completion_params(
completion_params: Dict[str, Any],
a2a_message: Dict[str, Any],
params: Optional[Dict[str, Any]] = None,
completion_params: dict[str, Any],
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
) -> None:
"""
Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph).
@ -104,17 +104,17 @@ class A2ACompletionBridgeTransformation:
# ``extra_body.metadata`` so the configured keys remain authoritative
# and an A2A caller cannot overwrite server-set run metadata.
existing_metadata = extra_body.get("metadata")
existing_dict: Dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: Dict[str, Any] = {**forward_metadata, **existing_dict}
existing_dict: dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: dict[str, Any] = {**forward_metadata, **existing_dict}
extra_body = {**extra_body, "metadata": merged_metadata}
completion_params["extra_body"] = extra_body
verbose_logger.debug(f"A2A -> completion forward metadata keys={list(forward_metadata.keys())}")
verbose_logger.debug("A2A -> completion forward metadata keys=%s", list(forward_metadata.keys()))
@staticmethod
def a2a_message_to_openai_messages(
a2a_message: Dict[str, Any],
) -> List[Dict[str, Any]]:
a2a_message: dict[str, Any],
) -> list[dict[str, Any]]:
"""
Transform an A2A message to OpenAI message format.
@ -143,17 +143,19 @@ class A2ACompletionBridgeTransformation:
# Do not attach A2A message.metadata here — the completion bridge forwards it
# once at run level via extra_body.metadata (LangGraph POST /runs/wait shape).
openai_message: Dict[str, Any] = {"role": openai_role, "content": content}
openai_message: dict[str, Any] = {"role": openai_role, "content": content}
verbose_logger.debug(f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}")
verbose_logger.debug(
"A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content)
)
return [openai_message]
@staticmethod
def openai_response_to_a2a_response(
response: Any,
request_id: Optional[str] = None,
) -> Dict[str, Any]:
request_id: str | None = None,
) -> dict[str, Any]:
"""
Transform a LiteLLM ModelResponse to A2A SendMessageResponse format.
@ -186,7 +188,7 @@ class A2ACompletionBridgeTransformation:
"result": a2a_message,
}
verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}")
verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(content))
return a2a_response
@ -198,7 +200,7 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def create_task_event(
ctx: A2AStreamingContext,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Create the initial task event with status 'submitted'.
@ -232,8 +234,8 @@ class A2ACompletionBridgeTransformation:
ctx: A2AStreamingContext,
state: str,
final: bool = False,
message_text: Optional[str] = None,
) -> Dict[str, Any]:
message_text: str | None = None,
) -> dict[str, Any]:
"""
Create a status update event.
@ -243,7 +245,7 @@ class A2ACompletionBridgeTransformation:
final: Whether this is the final event
message_text: Optional message text for 'working' status
"""
status: Dict[str, Any] = {
status: dict[str, Any] = {
"state": state,
"timestamp": A2ACompletionBridgeTransformation._get_timestamp(),
}
@ -275,7 +277,7 @@ class A2ACompletionBridgeTransformation:
def create_artifact_update_event(
ctx: A2AStreamingContext,
text: str,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Create an artifact update event with content.

61
litellm/a2a_protocol/main.py
View file

@ -12,14 +12,11 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
import asyncio
import datetime
import uuid
from collections.abc import AsyncIterator, Coroutine
from typing import (
TYPE_CHECKING,
Any,
AsyncIterator,
Coroutine,
Dict,
Optional,
Union,
cast,
)
@ -87,7 +84,7 @@ A2ACardResolver = LiteLLMA2ACardResolver
def _set_usage_on_logging_obj(
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
prompt_tokens: int,
completion_tokens: int,
) -> None:
@ -110,7 +107,7 @@ def _set_usage_on_logging_obj(
def _set_agent_id_on_logging_obj(
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
agent_id: str | None,
) -> None:
"""
@ -156,7 +153,7 @@ def _set_litellm_params_on_logging_obj(
logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params}
def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str:
"""
Extract agent info and set model/custom_llm_provider for cost tracking.
@ -199,15 +196,15 @@ async def _send_message_via_completion_bridge(
request: "SendMessageRequest",
custom_llm_provider: str,
api_base: str | None,
litellm_params: Dict[str, Any],
agent_extra_headers: Dict[str, str] | None = None,
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> LiteLLMSendMessageResponse:
"""
Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
Requires request; api_base is optional for providers that derive endpoint from model.
"""
verbose_logger.info(f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}")
verbose_logger.info("A2A using completion bridge: provider=%s, api_base=%s", custom_llm_provider, api_base)
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
@ -370,9 +367,9 @@ async def asend_message(
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendMessageRequest"] = None,
api_base: str | None = None,
litellm_params: Dict[str, Any] | None = None,
litellm_params: dict[str, Any] | None = None,
agent_id: str | None = None,
agent_extra_headers: Dict[str, str] | None = None,
agent_extra_headers: dict[str, str] | None = None,
**kwargs: Any,
) -> LiteLLMSendMessageResponse:
"""
@ -453,7 +450,7 @@ async def asend_message(
if api_base is None:
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
trace_id = trace_id or str(uuid.uuid4())
extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
extra_headers: dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
if agent_id:
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
# Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones)
@ -466,7 +463,7 @@ async def asend_message(
agent_name = _get_a2a_model_info(a2a_client, kwargs)
verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}")
verbose_logger.info("A2A send_message request_id=%s, agent=%s", request.id, agent_name)
# Get agent card URL for localhost retry logic
agent_card = _get_a2a_client_agent_card(a2a_client)
@ -481,7 +478,7 @@ async def asend_message(
agent_name=agent_name,
)
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
verbose_logger.info("A2A send_message completed, request_id=%s", request.id)
# Wrap in LiteLLM response type for _hidden_params support
response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
@ -518,7 +515,7 @@ def send_message(
a2a_client: "A2AClientType",
request: "SendMessageRequest",
**kwargs: Any,
) -> Union[LiteLLMSendMessageResponse, Coroutine[Any, Any, LiteLLMSendMessageResponse]]:
) -> LiteLLMSendMessageResponse | Coroutine[Any, Any, LiteLLMSendMessageResponse]:
"""
Sync: Send a message to an A2A agent.
@ -547,9 +544,9 @@ def _build_streaming_logging_obj(
request: "SendStreamingMessageRequest",
agent_name: str,
agent_id: str | None,
litellm_params: Dict[str, Any] | None,
metadata: Dict[str, Any] | None,
proxy_server_request: Dict[str, Any] | None,
litellm_params: dict[str, Any] | None,
metadata: dict[str, Any] | None,
proxy_server_request: dict[str, Any] | None,
) -> Logging:
"""Build logging object for streaming A2A requests."""
start_time = datetime.datetime.now()
@ -590,11 +587,11 @@ async def asend_message_streaming(
a2a_client: Optional["A2AClientType"] = None,
request: Optional["SendStreamingMessageRequest"] = None,
api_base: str | None = None,
litellm_params: Dict[str, Any] | None = None,
litellm_params: dict[str, Any] | None = None,
agent_id: str | None = None,
metadata: Dict[str, Any] | None = None,
proxy_server_request: Dict[str, Any] | None = None,
agent_extra_headers: Dict[str, str] | None = None,
metadata: dict[str, Any] | None = None,
proxy_server_request: dict[str, Any] | None = None,
agent_extra_headers: dict[str, str] | None = None,
**kwargs: object,
) -> AsyncIterator[Any]:
"""
@ -643,7 +640,7 @@ async def asend_message_streaming(
raise ValueError("request is required for completion bridge")
# api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore)
verbose_logger.info(f"A2A streaming using completion bridge: provider={custom_llm_provider}")
verbose_logger.info("A2A streaming using completion bridge: provider=%s", custom_llm_provider)
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
@ -700,7 +697,7 @@ async def asend_message_streaming(
proxy_server_request=proxy_server_request,
)
verbose_logger.info(f"A2A send_message_streaming request_id={request.id}, agent={agent_name}")
verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name)
agent_card = _get_a2a_client_agent_card(a2a_client)
card_url = get_agent_card_url(agent_card) if agent_card else None
@ -728,7 +725,7 @@ async def asend_message_streaming(
async def create_a2a_client(
base_url: str,
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
extra_headers: Dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
streaming: bool = False,
) -> "A2AClientType":
"""
@ -762,7 +759,7 @@ async def create_a2a_client(
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
)
verbose_logger.info(f"Creating A2A client for {base_url}")
verbose_logger.info("Creating A2A client for %s", base_url)
# Use get_async_httpx_client with per-agent params so that different agents
# (with different extra_headers) get separate cached clients. The params
@ -784,7 +781,7 @@ async def create_a2a_client(
httpx_client = _async_handler.client
if extra_headers:
httpx_client.headers.update(extra_headers)
verbose_proxy_logger.debug(f"A2A client created with extra_headers={list(extra_headers.keys())}")
verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
a2a_client = await create_client( # pyright: ignore[reportOptionalCall]
base_url,
@ -801,7 +798,7 @@ async def create_a2a_client(
if agent_card is not None:
a2a_client._litellm_agent_card = agent_card # type: ignore[attr-defined]
verbose_logger.info(f"A2A client created for {base_url}")
verbose_logger.info("A2A client created for %s", base_url)
return a2a_client
@ -809,7 +806,7 @@ async def create_a2a_client(
async def aget_agent_card(
base_url: str,
timeout: float = DEFAULT_A2A_AGENT_TIMEOUT,
extra_headers: Dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
) -> "AgentCard":
"""
Fetch the agent card from an A2A agent.
@ -827,7 +824,7 @@ async def aget_agent_card(
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
)
verbose_logger.info(f"Fetching agent card from {base_url}")
verbose_logger.info("Fetching agent card from %s", base_url)
# Use LiteLLM's cached httpx client
http_handler = get_async_httpx_client(
@ -842,5 +839,5 @@ async def aget_agent_card(
)
agent_card = await resolver.get_agent_card()
verbose_logger.info(f"Fetched agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}")
verbose_logger.info("Fetched agent card: %s", agent_card.name if hasattr(agent_card, "name") else "unknown")
return agent_card

2
litellm/a2a_protocol/providers/__init__.py
View file

@ -7,4 +7,4 @@ This module contains provider-specific implementations for the A2A protocol.
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
__all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"]
__all__ = ["A2AProviderConfigManager", "BaseA2AProviderConfig"]

16
litellm/a2a_protocol/providers/base.py
View file

@ -3,7 +3,8 @@ Base configuration for A2A protocol providers.
"""
from abc import ABC, abstractmethod
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
class BaseA2AProviderConfig(ABC):
@ -18,10 +19,10 @@ class BaseA2AProviderConfig(ABC):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Handle non-streaming A2A request.
@ -34,16 +35,15 @@ class BaseA2AProviderConfig(ABC):
Returns:
A2A SendMessageResponse dict
"""
pass
@abstractmethod
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming A2A request.

15
litellm/a2a_protocol/providers/bedrock_agentcore/config.py
View file

@ -2,7 +2,8 @@
Bedrock AgentCore A2A provider configuration.
"""
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.bedrock_agentcore.handler import (
@ -22,10 +23,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Handle non-streaming request to AgentCore A2A agent."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
@ -42,10 +43,10 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""Handle streaming request to AgentCore A2A agent."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:

25
litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
View file

@ -6,7 +6,8 @@ completion bridge that would otherwise strip the envelope.
"""
import json
from typing import Any, AsyncIterator, Dict, Optional, cast
from collections.abc import AsyncIterator
from typing import Any, cast
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
@ -27,10 +28,10 @@ class BedrockAgentCoreA2AHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
params: dict[str, Any],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Handle non-streaming A2A request to AgentCore.
@ -52,7 +53,7 @@ class BedrockAgentCoreA2AHandler:
agent_extra_headers=agent_extra_headers,
)
verbose_logger.info(f"BedrockAgentCore A2A: Sending non-streaming request to {url}")
verbose_logger.info("BedrockAgentCore A2A: Sending non-streaming request to %s", url)
client = get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
@ -66,17 +67,17 @@ class BedrockAgentCoreA2AHandler:
response_data = response.json()
if "error" in response_data:
verbose_logger.warning(f"BedrockAgentCore A2A: Agent returned error: {response_data['error']}")
verbose_logger.warning("BedrockAgentCore A2A: Agent returned error: %s", response_data["error"])
return response_data
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
params: dict[str, Any],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming A2A request to AgentCore.
@ -99,7 +100,7 @@ class BedrockAgentCoreA2AHandler:
agent_extra_headers=agent_extra_headers,
)
verbose_logger.info(f"BedrockAgentCore A2A: Sending streaming request to {url}")
verbose_logger.info("BedrockAgentCore A2A: Sending streaming request to %s", url)
client = get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),

23
litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py
View file

@ -6,7 +6,8 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT).
"""
import json
from typing import Any, AsyncIterator, Dict, Mapping, Optional, Tuple
from collections.abc import AsyncIterator, Mapping
from typing import Any
from litellm._logging import verbose_logger
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
@ -28,15 +29,15 @@ _RESERVED_EXACT_HEADERS = frozenset(
"host",
}
)
_RESERVED_PREFIX_HEADERS: Tuple[str, ...] = (
_RESERVED_PREFIX_HEADERS: tuple[str, ...] = (
"x-amzn-bedrock-agentcore-runtime-",
"x-amz-",
)
def _filter_reserved_headers(
agent_extra_headers: Optional[Mapping[str, str]],
) -> Optional[Dict[str, str]]:
agent_extra_headers: Mapping[str, str] | None,
) -> dict[str, str] | None:
"""
Strip reserved AWS / AgentCore headers from caller-supplied
``agent_extra_headers`` before they are merged into the signed request.
@ -46,7 +47,7 @@ def _filter_reserved_headers(
if not agent_extra_headers:
return None
filtered: Dict[str, str] = {}
filtered: dict[str, str] = {}
dropped: list = []
for k, v in agent_extra_headers.items():
k_lower = k.lower()
@ -76,12 +77,12 @@ class BedrockAgentCoreA2ATransformation:
@staticmethod
def get_url_and_signed_request(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
params: dict[str, Any],
litellm_params: dict[str, Any],
method: str = "message/send",
stream: bool = False,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Tuple[str, dict, bytes]:
agent_extra_headers: dict[str, str] | None = None,
) -> tuple[str, dict, bytes]:
"""
Build the AgentCore URL, construct a JSON-RPC envelope, and sign the request.
@ -169,7 +170,7 @@ class BedrockAgentCoreA2ATransformation:
return url, signed_headers, signed_body
@staticmethod
async def parse_sse_events(response: Any) -> AsyncIterator[Dict[str, Any]]:
async def parse_sse_events(response: Any) -> AsyncIterator[dict[str, Any]]:
"""
Parse SSE events from an httpx streaming response.
@ -194,5 +195,5 @@ class BedrockAgentCoreA2ATransformation:
event = json.loads(data_str)
yield event
except json.JSONDecodeError:
verbose_logger.debug(f"BedrockAgentCore A2A: Skipping non-JSON SSE line: {data_str[:100]}")
verbose_logger.debug("BedrockAgentCore A2A: Skipping non-JSON SSE line: %s", data_str[:100])
continue

8
litellm/a2a_protocol/providers/config_manager.py
View file

@ -4,8 +4,6 @@ A2A Provider Config Manager.
Manages provider-specific configurations for A2A protocol.
"""
from typing import Optional
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
@ -18,9 +16,9 @@ class A2AProviderConfigManager:
@staticmethod
def get_provider_config(
custom_llm_provider: Optional[str],
model: Optional[str] = None,
) -> Optional[BaseA2AProviderConfig]:
custom_llm_provider: str | None,
model: str | None = None,
) -> BaseA2AProviderConfig | None:
"""
Get the provider configuration for a given custom_llm_provider.

15
litellm/a2a_protocol/providers/langflow/config.py
View file

@ -1,4 +1,5 @@
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2A_USER_API_KEY_HASH_PARAM,
@ -15,10 +16,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> Dict[str, Any]:
) -> dict[str, Any]:
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
@ -38,10 +39,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(

2
litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py
View file

@ -13,4 +13,4 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
PydanticAITransformation,
)
__all__ = ["PydanticAIHandler", "PydanticAITransformation", "PydanticAIProviderConfig"]
__all__ = ["PydanticAIHandler", "PydanticAIProviderConfig", "PydanticAITransformation"]

15
litellm/a2a_protocol/providers/pydantic_ai_agents/config.py
View file

@ -2,7 +2,8 @@
Pydantic AI provider configuration.
"""
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAIHandler
@ -19,10 +20,10 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs: Any,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Handle non-streaming request to Pydantic AI agent."""
if api_base is None:
raise ValueError("api_base is required for PydanticAIProviderConfig")
@ -37,10 +38,10 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""Handle streaming request with fake streaming."""
if not api_base:
raise ValueError("api_base is required for Pydantic AI agents")

23
litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py
View file

@ -5,7 +5,8 @@ Pydantic AI agents follow A2A protocol but don't support streaming natively.
This handler provides fake streaming by converting non-streaming responses into streaming chunks.
"""
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
@ -25,11 +26,11 @@ class PydanticAIHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Handle non-streaming request to Pydantic AI agent.
@ -46,7 +47,7 @@ class PydanticAIHandler:
"""
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
verbose_logger.info(f"Pydantic AI: Routing to Pydantic AI agent at {api_base}")
verbose_logger.info("Pydantic AI: Routing to Pydantic AI agent at %s", api_base)
# Send request directly to Pydantic AI agent
response_data = await PydanticAITransformation.send_non_streaming_request(
@ -62,13 +63,13 @@ class PydanticAIHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
timeout: float = 60.0,
chunk_size: int = 50,
delay_ms: int = 10,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> AsyncIterator[Dict[str, Any]]:
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
"""
Handle streaming request to Pydantic AI agent with fake streaming.
@ -91,7 +92,7 @@ class PydanticAIHandler:
"""
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
verbose_logger.info(f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}")
verbose_logger.info("Pydantic AI: Faking streaming for Pydantic AI agent at %s", api_base)
# Get raw task response first (not the transformed A2A format)
raw_response = await PydanticAITransformation.send_and_get_raw_response(

41
litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py
View file

@ -6,7 +6,8 @@ This module provides fake streaming by converting non-streaming responses into s
"""
import asyncio
from typing import Any, AsyncIterator, Dict, Optional, cast
from collections.abc import AsyncIterator
from typing import Any, cast
from uuid import uuid4
from litellm._logging import verbose_logger
@ -48,7 +49,7 @@ class PydanticAITransformation:
return obj
@staticmethod
def _params_to_dict(params: Any) -> Dict[str, Any]:
def _params_to_dict(params: Any) -> dict[str, Any]:
"""
Convert params to a dict, handling Pydantic models.
@ -78,8 +79,8 @@ class PydanticAITransformation:
request_id: str,
max_attempts: int = 30,
poll_interval: float = 0.5,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Poll for task completion using tasks/get method.
@ -117,7 +118,7 @@ class PydanticAITransformation:
status = result.get("status", {})
state = status.get("state", "")
verbose_logger.debug(f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}")
verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state)
if state == "completed":
return poll_data
@ -134,8 +135,8 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -172,7 +173,7 @@ class PydanticAITransformation:
# FastA2A uses root endpoint (/) not /messages
endpoint = api_base.rstrip("/")
verbose_logger.info(f"Pydantic AI: Sending non-streaming request to {endpoint}")
verbose_logger.info("Pydantic AI: Sending non-streaming request to %s", endpoint)
# Send request to Pydantic AI agent using shared async HTTP client
client = get_async_httpx_client(
@ -199,7 +200,7 @@ class PydanticAITransformation:
# Need to poll for completion
task_id = result.get("id")
if task_id:
verbose_logger.info(f"Pydantic AI: Task {task_id} submitted, polling for completion...")
verbose_logger.info("Pydantic AI: Task %s submitted, polling for completion...", task_id)
response_data = await PydanticAITransformation._poll_for_completion(
client=client,
endpoint=endpoint,
@ -208,7 +209,7 @@ class PydanticAITransformation:
agent_extra_headers=agent_extra_headers,
)
verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}")
verbose_logger.info("Pydantic AI: Received completed response for request_id=%s", request_id)
return response_data
@ -218,8 +219,8 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
@ -254,8 +255,8 @@ class PydanticAITransformation:
request_id: str,
params: Any,
timeout: float = 60.0,
agent_extra_headers: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
@ -281,9 +282,9 @@ class PydanticAITransformation:
@staticmethod
def _transform_to_a2a_response(
response_data: Dict[str, Any],
response_data: dict[str, Any],
request_id: str,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""
Transform Pydantic AI task response to standard A2A non-streaming format.
@ -327,7 +328,7 @@ class PydanticAITransformation:
}
@staticmethod
def _extract_response_text(response_data: Dict[str, Any]) -> tuple[str, str, list]:
def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]:
"""
Extract response text from completed task response.
@ -382,11 +383,11 @@ class PydanticAITransformation:
@staticmethod
async def fake_streaming_from_response(
response_data: Dict[str, Any],
response_data: dict[str, Any],
request_id: str,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Convert a non-streaming A2A response into fake streaming chunks.
@ -517,4 +518,4 @@ class PydanticAITransformation:
}
yield completed_event
verbose_logger.info(f"Pydantic AI: Fake streaming completed for request_id={request_id}")
verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id)

15
litellm/a2a_protocol/providers/watsonx_orchestrate/config.py
View file

@ -2,7 +2,8 @@
A2A provider configuration for IBM watsonx Orchestrate (WXO).
"""
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
@ -16,10 +17,10 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
async def handle_non_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs: Any,
) -> Dict[str, Any]:
) -> dict[str, Any]:
"""Handle a non-streaming A2A request via WXO runs API."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:
@ -36,10 +37,10 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
async def handle_streaming(
self,
request_id: str,
params: Dict[str, Any],
api_base: Optional[str] = None,
params: dict[str, Any],
api_base: str | None = None,
**kwargs: Any,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""Handle a streaming A2A request via WXO streaming runs API."""
litellm_params = kwargs.get("litellm_params")
if not litellm_params:

51
litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py
View file

@ -6,7 +6,8 @@ import asyncio
import hashlib
import json
import time
from typing import Any, AsyncIterator, Dict, NamedTuple, Optional, Tuple, cast
from collections.abc import AsyncIterator
from typing import Any, NamedTuple, cast
import httpx
@ -24,7 +25,7 @@ _IBM_CLOUD_IAM_URL = "https://iam.cloud.ibm.com/identity/token"
_POLL_INTERVAL_S = 2.0
_MAX_POLL_ATTEMPTS = 90
_TOKEN_CACHE_TTL_BUFFER_S = 60
_token_cache: Dict[str, Tuple[str, float]] = {}
_token_cache: dict[str, tuple[str, float]] = {}
class WXORequestParams(NamedTuple):
@ -32,9 +33,9 @@ class WXORequestParams(NamedTuple):
instance_id: str
wxo_agent_id: str
api_key: str
username: Optional[str]
username: str | None
auth_mode: str
thread_id: Optional[str]
thread_id: str | None
class WatsonxOrchestrateHandler:
@ -50,13 +51,13 @@ class WatsonxOrchestrateHandler:
auth_mode: str,
cp4d_host: str,
api_key: str,
username: Optional[str],
username: str | None,
) -> str:
material = f"{auth_mode}:{cp4d_host}:{username or ''}:{api_key}"
return hashlib.sha256(material.encode()).hexdigest()
@staticmethod
def _cp4d_token_ttl_seconds(expiration: Any, now_wall: Optional[float] = None) -> int:
def _cp4d_token_ttl_seconds(expiration: Any, now_wall: float | None = None) -> int:
# CP4D returns expiration as absolute Unix epoch seconds, not a duration.
expires_at = int(expiration)
wall = now_wall if now_wall is not None else time.time()
@ -67,8 +68,8 @@ class WatsonxOrchestrateHandler:
cp4d_host: str,
auth_mode: str,
api_key: str,
username: Optional[str] = None,
client: Optional[AsyncHTTPHandler] = None,
username: str | None = None,
client: AsyncHTTPHandler | None = None,
) -> str:
cache_key = WatsonxOrchestrateHandler._token_cache_key(auth_mode, cp4d_host, api_key, username)
now = time.monotonic()
@ -121,20 +122,20 @@ class WatsonxOrchestrateHandler:
async def _poll_run(
base_url: str,
run_id: str,
auth_headers: Dict[str, str],
auth_headers: dict[str, str],
client: AsyncHTTPHandler,
max_attempts: int = _MAX_POLL_ATTEMPTS,
interval_s: float = _POLL_INTERVAL_S,
) -> Dict[str, Any]:
) -> dict[str, Any]:
url = f"{base_url}/v1/orchestrate/runs/{run_id}"
for attempt in range(max_attempts):
await asyncio.sleep(interval_s)
response = await client.get(url, headers=auth_headers)
response.raise_for_status()
result: Dict[str, Any] = response.json()
result: dict[str, Any] = response.json()
status = result.get("status", "")
verbose_logger.debug(f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'")
verbose_logger.debug("WXO: Poll %s/%s run='%s' status='%s'", attempt + 1, max_attempts, run_id, status)
if status in WatsonxOrchestrateTransformation.TERMINAL_STATES:
return result
@ -144,11 +145,11 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def _get_successful_run_data(
run_data: Dict[str, Any],
run_data: dict[str, Any],
base_url: str,
auth_headers: Dict[str, str],
auth_headers: dict[str, str],
client: AsyncHTTPHandler,
) -> Dict[str, Any]:
) -> dict[str, Any]:
status = run_data.get("status", "")
if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
run_id = run_data.get("run_id") or run_data.get("id") or ""
@ -186,7 +187,7 @@ class WatsonxOrchestrateHandler:
return accumulated_text
@staticmethod
def _extract_litellm_params(litellm_params: Dict[str, Any]) -> WXORequestParams:
def _extract_litellm_params(litellm_params: dict[str, Any]) -> WXORequestParams:
cp4d_host = litellm_params.get("cp4d_host") or ""
instance_id = litellm_params.get("instance_id") or ""
wxo_agent_id = litellm_params.get("wxo_agent_id") or ""
@ -214,9 +215,9 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
) -> Dict[str, Any]:
params: dict[str, Any],
litellm_params: dict[str, Any],
) -> dict[str, Any]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=90.0)
@ -245,7 +246,7 @@ class WatsonxOrchestrateHandler:
headers=auth_headers,
)
run_response.raise_for_status()
run_data: Dict[str, Any] = run_response.json()
run_data: dict[str, Any] = run_response.json()
run_data = await WatsonxOrchestrateHandler._get_successful_run_data(
run_data=run_data,
@ -260,11 +261,11 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
params: dict[str, Any],
litellm_params: dict[str, Any],
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=120.0)
@ -296,8 +297,8 @@ class WatsonxOrchestrateHandler:
response.raise_for_status()
except httpx.TransportError as exc:
verbose_logger.warning(
f"WXO: Streaming request failed before a run was submitted "
f"({exc!r}), falling back to non-streaming + fake streaming",
"WXO: Streaming request failed before a run was submitted (%r), falling back to non-streaming + fake streaming",
exc,
exc_info=True,
)
result = await WatsonxOrchestrateHandler.handle_non_streaming(

19
litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py
View file

@ -8,7 +8,8 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model:
"""
import asyncio
from typing import Any, AsyncIterator, Dict, Optional
from collections.abc import AsyncIterator
from typing import Any
from uuid import uuid4
from litellm._logging import verbose_logger
@ -28,7 +29,7 @@ class WatsonxOrchestrateTransformation:
return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}"
@staticmethod
def extract_text_from_a2a_params(params: Dict[str, Any]) -> str:
def extract_text_from_a2a_params(params: dict[str, Any]) -> str:
"""
Extract user message text from A2A MessageSendParams.
@ -49,10 +50,10 @@ class WatsonxOrchestrateTransformation:
def build_wxo_run_body(
wxo_agent_id: str,
text: str,
thread_id: Optional[str] = None,
) -> Dict[str, Any]:
thread_id: str | None = None,
) -> dict[str, Any]:
"""Build the WXO POST /v1/orchestrate/runs request body."""
body: Dict[str, Any] = {
body: dict[str, Any] = {
"agent_id": wxo_agent_id,
"message": {
"role": "user",
@ -102,7 +103,7 @@ class WatsonxOrchestrateTransformation:
return ""
@staticmethod
def extract_text_from_a2a_message_response(a2a_response: Dict[str, Any]) -> str:
def extract_text_from_a2a_message_response(a2a_response: dict[str, Any]) -> str:
result = a2a_response.get("result")
if not isinstance(result, dict):
verbose_logger.warning("WXO: A2A response missing result object")
@ -118,7 +119,7 @@ class WatsonxOrchestrateTransformation:
return ""
@staticmethod
def build_a2a_message_response(request_id: str, text: str) -> Dict[str, Any]:
def build_a2a_message_response(request_id: str, text: str) -> dict[str, Any]:
"""
Build a standard A2A non-streaming SendMessageResponse (kind=message).
"""
@ -139,7 +140,7 @@ class WatsonxOrchestrateTransformation:
request_id: str,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[Dict[str, Any]]:
) -> AsyncIterator[dict[str, Any]]:
"""
Emit standard A2A streaming events from a completed text response.
@ -213,4 +214,4 @@ class WatsonxOrchestrateTransformation:
},
}
verbose_logger.debug(f"WXO: Fake streaming completed for request_id={request_id}")
verbose_logger.debug("WXO: Fake streaming completed for request_id=%s", request_id)

23
litellm/a2a_protocol/streaming_iterator.py
View file

@ -3,8 +3,9 @@ A2A Streaming Iterator with token tracking and logging support.
"""
import asyncio
from collections.abc import AsyncIterator
from datetime import datetime
from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional
from typing import TYPE_CHECKING, Any
import litellm
from litellm._logging import verbose_logger
@ -37,9 +38,9 @@ class A2AStreamingIterator:
self.start_time = datetime.now()
# Collect chunks for token counting
self.chunks: List[Any] = []
self.collected_text_parts: List[str] = []
self.final_chunk: Optional[Any] = None
self.chunks: list[Any] = []
self.collected_text_parts: list[str] = []
self.final_chunk: Any | None = None
def __aiter__(self):
return self
@ -137,17 +138,19 @@ class A2AStreamingIterator:
)
verbose_logger.info(
f"A2A streaming completed: prompt_tokens={prompt_tokens}, "
f"completion_tokens={completion_tokens}, total_tokens={total_tokens}, "
f"response_cost={response_cost}"
"A2A streaming completed: prompt_tokens=%s, completion_tokens=%s, total_tokens=%s, response_cost=%s",
prompt_tokens,
completion_tokens,
total_tokens,
response_cost,
)
except Exception as e:
verbose_logger.debug(f"Error in A2A streaming completion handler: {e}")
verbose_logger.debug("Error in A2A streaming completion handler: %s", e)
def _build_logging_result(self, usage: litellm.Usage) -> Dict[str, Any]:
def _build_logging_result(self, usage: litellm.Usage) -> dict[str, Any]:
"""Build a result dict for logging."""
result: Dict[str, Any] = {
result: dict[str, Any] = {
"id": getattr(self.request, "id", "unknown"),
"jsonrpc": "2.0",
"usage": (usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)),

16
litellm/a2a_protocol/utils.py
View file

@ -2,7 +2,7 @@
Utility functions for A2A protocol.
"""
from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union
from typing import TYPE_CHECKING, Any
import litellm
from litellm._logging import verbose_logger
@ -34,7 +34,7 @@ class A2ARequestUtils:
else:
parts = getattr(message, "parts", []) or []
text_parts: List[str] = []
text_parts: list[str] = []
for part in parts:
if isinstance(part, dict):
if part.get("kind") == "text":
@ -46,7 +46,7 @@ class A2ARequestUtils:
return " ".join(text_parts)
@staticmethod
def extract_text_from_response(response_dict: Dict[str, Any]) -> str:
def extract_text_from_response(response_dict: dict[str, Any]) -> str:
"""
Extract text content from A2A response result.
@ -71,7 +71,7 @@ class A2ARequestUtils:
@staticmethod
def get_input_message_from_request(
request: "Union[SendMessageRequest, SendStreamingMessageRequest]",
request: "SendMessageRequest | SendStreamingMessageRequest",
) -> Any:
"""
Extract the input message from an A2A request.
@ -108,9 +108,9 @@ class A2ARequestUtils:
@staticmethod
def calculate_usage_from_request_response(
request: "Union[SendMessageRequest, SendStreamingMessageRequest]",
response_dict: Dict[str, Any],
) -> Tuple[int, int, int]:
request: "SendMessageRequest | SendStreamingMessageRequest",
response_dict: dict[str, Any],
) -> tuple[int, int, int]:
"""
Calculate token usage from A2A request and response.
@ -145,5 +145,5 @@ def extract_text_from_a2a_message(message: Any) -> str:
return A2ARequestUtils.extract_text_from_message(message)
def extract_text_from_a2a_response(response_dict: Dict[str, Any]) -> str:
def extract_text_from_a2a_response(response_dict: dict[str, Any]) -> str:
return A2ARequestUtils.extract_text_from_response(response_dict)

27
litellm/anthropic_beta_headers_manager.py
View file

@ -25,14 +25,13 @@ Environment Variables:
import json
import os
from importlib.resources import files
from typing import Dict, List, Optional, Set
import httpx
from litellm.litellm_core_utils.litellm_logging import verbose_logger
# Cache for the loaded configuration
_BETA_HEADERS_CONFIG: Optional[Dict] = None
_BETA_HEADERS_CONFIG: dict | None = None
class GetAnthropicBetaHeadersConfig:
@ -44,7 +43,7 @@ class GetAnthropicBetaHeadersConfig:
"""
@staticmethod
def load_local_beta_headers_config() -> Dict:
def load_local_beta_headers_config() -> dict:
"""Load the local backup beta headers config bundled with the package."""
try:
content = json.loads(
@ -52,7 +51,7 @@ class GetAnthropicBetaHeadersConfig:
)
return content
except Exception as e:
verbose_logger.error(f"Failed to load local beta headers config: {e}")
verbose_logger.error("Failed to load local beta headers config: %s", e)
# Return empty config as fallback
return {
"anthropic": {},
@ -159,7 +158,7 @@ def get_beta_headers_config(url: str) -> dict:
return content
def _load_beta_headers_config() -> Dict:
def _load_beta_headers_config() -> dict:
"""
Load the beta headers configuration.
Uses caching to avoid repeated fetches/file reads.
@ -183,7 +182,7 @@ def _load_beta_headers_config() -> Dict:
return _BETA_HEADERS_CONFIG
def reload_beta_headers_config() -> Dict:
def reload_beta_headers_config() -> dict:
"""
Force reload the beta headers configuration from source (remote or local).
Clears the cache and fetches fresh configuration.
@ -213,9 +212,9 @@ def get_provider_name(provider: str) -> str:
def filter_and_transform_beta_headers(
beta_headers: List[str],
beta_headers: list[str],
provider: str,
) -> List[str]:
) -> list[str]:
"""
Filter and transform beta headers based on provider's mapping configuration.
@ -240,14 +239,16 @@ def filter_and_transform_beta_headers(
# Get the header mapping for this provider
provider_mapping = config.get(provider, {})
filtered_headers: Set[str] = set()
filtered_headers: set[str] = set()
for header in beta_headers:
header = header.strip()
# Check if header is in the mapping
if header not in provider_mapping:
verbose_logger.debug(f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)")
verbose_logger.debug(
"Dropping unknown beta header '%s' for provider '%s' (not in mapping)", header, provider
)
continue
# Get the mapped header value
@ -255,7 +256,7 @@ def filter_and_transform_beta_headers(
# Skip if header is unsupported (null value)
if mapped_header is None:
verbose_logger.debug(f"Dropping unsupported beta header '{header}' for provider '{provider}'")
verbose_logger.debug("Dropping unsupported beta header '%s' for provider '%s'", header, provider)
continue
# Add the mapped header
@ -289,7 +290,7 @@ def is_beta_header_supported(
def get_provider_beta_header(
anthropic_beta_header: str,
provider: str,
) -> Optional[str]:
) -> str | None:
"""
Get the provider-specific beta header name for a given Anthropic beta header.
@ -390,7 +391,7 @@ def update_request_with_filtered_beta(
return headers, request_data
def get_unsupported_headers(provider: str) -> List[str]:
def get_unsupported_headers(provider: str) -> list[str]:
"""
Get all beta headers that are unsupported by a provider (have null values in mapping).

4
litellm/anthropic_interface/exceptions/__init__.py
View file

@ -11,9 +11,9 @@ from .exceptions import (
)
__all__ = [
"AnthropicErrorType",
"ANTHROPIC_ERROR_TYPE_MAP",
"AnthropicErrorDetail",
"AnthropicErrorResponse",
"ANTHROPIC_ERROR_TYPE_MAP",
"AnthropicErrorType",
"AnthropicExceptionMapping",
]

9
litellm/anthropic_interface/exceptions/exception_mapping_utils.py
View file

@ -5,13 +5,12 @@ Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthrop
"""
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from typing import Dict, Optional
from .exceptions import AnthropicErrorResponse, AnthropicErrorType
# HTTP status code -> Anthropic error type
# Source: https://docs.anthropic.com/en/api/errors
ANTHROPIC_ERROR_TYPE_MAP: Dict[int, AnthropicErrorType] = {
ANTHROPIC_ERROR_TYPE_MAP: dict[int, AnthropicErrorType] = {
400: "invalid_request_error",
401: "authentication_error",
403: "permission_error",
@ -39,7 +38,7 @@ class AnthropicExceptionMapping:
def create_error_response(
status_code: int,
message: str,
request_id: Optional[str] = None,
request_id: str | None = None,
) -> AnthropicErrorResponse:
"""
Create an Anthropic-formatted error response dict.
@ -124,7 +123,7 @@ class AnthropicExceptionMapping:
def transform_to_anthropic_error(
status_code: int,
raw_message: str,
request_id: Optional[str] = None,
request_id: str | None = None,
) -> AnthropicErrorResponse:
"""
Transform an error message to Anthropic format.
@ -143,7 +142,7 @@ class AnthropicExceptionMapping:
AnthropicErrorResponse dict
"""
# Try to parse as JSON once
parsed: Optional[dict] = safe_json_loads(raw_message)
parsed: dict | None = safe_json_loads(raw_message)
if not isinstance(parsed, dict):
parsed = None

4
litellm/anthropic_interface/exceptions/exceptions.py
View file

@ -1,6 +1,8 @@
"""Anthropic error format type definitions."""
from typing_extensions import Literal, Required, TypedDict
from typing import Literal
from typing_extensions import Required, TypedDict
# Known Anthropic error types
# Source: https://docs.anthropic.com/en/api/errors

65
litellm/anthropic_interface/messages/__init__.py
View file

@ -10,7 +10,8 @@ This is an __init__.py file to allow the following interface
"""
from typing import Any, AsyncIterator, Coroutine, Dict, Iterator, List, Optional, Union
from collections.abc import AsyncIterator, Coroutine, Iterator
from typing import Any
from litellm.llms.anthropic.experimental_pass_through.messages.handler import (
anthropic_messages as _async_anthropic_messages,
@ -25,21 +26,21 @@ from litellm.types.llms.anthropic_messages.anthropic_response import (
async def acreate(
max_tokens: int,
messages: List[Dict],
messages: list[dict],
model: str,
metadata: Optional[Dict] = None,
stop_sequences: Optional[List[str]] = None,
stream: Optional[bool] = False,
system: Optional[str] = None,
temperature: Optional[float] = None,
thinking: Optional[Dict] = None,
tool_choice: Optional[Dict] = None,
tools: Optional[List[Dict]] = None,
top_k: Optional[int] = None,
top_p: Optional[float] = None,
container: Optional[Dict] = None,
metadata: dict | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
top_k: int | None = None,
top_p: float | None = None,
container: dict | None = None,
**kwargs,
) -> Union[AnthropicMessagesResponse, AsyncIterator]:
) -> AnthropicMessagesResponse | AsyncIterator:
"""
Async wrapper for Anthropic's messages API
@ -84,26 +85,26 @@ async def acreate(
def create(
max_tokens: int,
messages: List[Dict],
messages: list[dict],
model: str,
metadata: Optional[Dict] = None,
stop_sequences: Optional[List[str]] = None,
stream: Optional[bool] = False,
system: Optional[str] = None,
temperature: Optional[float] = None,
thinking: Optional[Dict] = None,
tool_choice: Optional[Dict] = None,
tools: Optional[List[Dict]] = None,
top_k: Optional[int] = None,
top_p: Optional[float] = None,
container: Optional[Dict] = None,
metadata: dict | None = None,
stop_sequences: list[str] | None = None,
stream: bool | None = False,
system: str | None = None,
temperature: float | None = None,
thinking: dict | None = None,
tool_choice: dict | None = None,
tools: list[dict] | None = None,
top_k: int | None = None,
top_p: float | None = None,
container: dict | None = None,
**kwargs,
) -> Union[
AnthropicMessagesResponse,
Iterator[bytes],
AsyncIterator[Any],
Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]],
]:
) -> (
AnthropicMessagesResponse
| Iterator[bytes]
| AsyncIterator[Any]
| Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]]
):
"""
Async wrapper for Anthropic's messages API

215
litellm/assistants/main.py
View file

@ -3,8 +3,9 @@
import asyncio
import contextvars
import os
from collections.abc import Coroutine, Iterable
from functools import partial
from typing import Any, Coroutine, Dict, Iterable, List, Literal, Optional, Union
from typing import Any, Literal
import httpx
from openai import AsyncOpenAI, OpenAI
@ -36,7 +37,7 @@ azure_assistants_api = AzureAssistantsAPI()
async def aget_assistants(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> AsyncCursorPage[Assistant]:
loop = asyncio.get_event_loop()
@ -73,13 +74,13 @@ async def aget_assistants(
def get_assistants(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[Any] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
client: Any | None = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
**kwargs,
) -> SyncCursorPage[Assistant]:
aget_assistants: Optional[bool] = kwargs.pop("aget_assistants", None)
aget_assistants: bool | None = kwargs.pop("aget_assistants", None)
if aget_assistants is not None and not isinstance(aget_assistants, bool):
raise Exception("Invalid value passed in for aget_assistants. Only bool or None allowed")
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
@ -101,7 +102,7 @@ def get_assistants(
elif timeout is None:
timeout = 600.0
response: Optional[SyncCursorPage[Assistant]] = None
response: SyncCursorPage[Assistant] | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -147,7 +148,7 @@ def get_assistants(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -166,9 +167,7 @@ def get_assistants(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_assistants'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_assistants'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -180,9 +179,7 @@ def get_assistants(
if response is None:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_assistants'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_assistants'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -197,7 +194,7 @@ def get_assistants(
async def acreate_assistants(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> Assistant:
loop = asyncio.get_event_loop()
@ -237,22 +234,22 @@ async def acreate_assistants(
def create_assistants(
custom_llm_provider: Literal["openai", "azure"],
model: str,
name: Optional[str] = None,
description: Optional[str] = None,
instructions: Optional[str] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_resources: Optional[Dict[str, Any]] = None,
metadata: Optional[Dict[str, str]] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
response_format: Optional[Union[str, Dict[str, str]]] = None,
client: Optional[Any] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
name: str | None = None,
description: str | None = None,
instructions: str | None = None,
tools: list[dict[str, Any]] | None = None,
tool_resources: dict[str, Any] | None = None,
metadata: dict[str, str] | None = None,
temperature: float | None = None,
top_p: float | None = None,
response_format: str | dict[str, str] | None = None,
client: Any | None = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
**kwargs,
) -> Union[Assistant, Coroutine[Any, Any, Assistant]]:
async_create_assistants: Optional[bool] = kwargs.pop("async_create_assistants", None)
) -> Assistant | Coroutine[Any, Any, Assistant]:
async_create_assistants: bool | None = kwargs.pop("async_create_assistants", None)
if async_create_assistants is not None and not isinstance(async_create_assistants, bool):
raise ValueError("Invalid value passed in for async_create_assistants. Only bool or None allowed")
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
@ -290,7 +287,7 @@ def create_assistants(
# only send params that are not None
create_assistant_data = {k: v for k, v in create_assistant_data.items() if v is not None}
response: Optional[Union[Coroutine[Any, Any, Assistant], Assistant]] = None
response: Coroutine[Any, Any, Assistant] | Assistant | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -337,7 +334,7 @@ def create_assistants(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -360,9 +357,7 @@ def create_assistants(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_assistants'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_assistants'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -382,7 +377,7 @@ def create_assistants(
async def adelete_assistant(
custom_llm_provider: Literal["openai", "azure"],
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> AssistantDeleted:
loop = asyncio.get_event_loop()
@ -421,17 +416,17 @@ async def adelete_assistant(
def delete_assistant(
custom_llm_provider: Literal["openai", "azure"],
assistant_id: str,
client: Optional[Any] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
api_version: Optional[str] = None,
client: Any | None = None,
api_key: str | None = None,
api_base: str | None = None,
api_version: str | None = None,
**kwargs,
) -> Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]:
) -> AssistantDeleted | Coroutine[Any, Any, AssistantDeleted]:
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
async_delete_assistants: Optional[bool] = kwargs.pop("async_delete_assistants", None)
async_delete_assistants: bool | None = kwargs.pop("async_delete_assistants", None)
if async_delete_assistants is not None and not isinstance(async_delete_assistants, bool):
raise ValueError("Invalid value passed in for async_delete_assistants. Only bool or None allowed")
@ -451,7 +446,7 @@ def delete_assistant(
elif timeout is None:
timeout = 600.0
response: Optional[Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]] = None
response: AssistantDeleted | Coroutine[Any, Any, AssistantDeleted] | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base
@ -490,7 +485,7 @@ def delete_assistant(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -513,9 +508,7 @@ def delete_assistant(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'delete_assistant'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'delete_assistant'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -571,10 +564,10 @@ async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwar
def create_thread(
custom_llm_provider: Literal["openai", "azure"],
messages: Optional[Iterable[OpenAICreateThreadParamsMessage]] = None,
metadata: Optional[dict] = None,
tool_resources: Optional[OpenAICreateThreadParamsToolResources] = None,
client: Optional[OpenAI] = None,
messages: Iterable[OpenAICreateThreadParamsMessage] | None = None,
metadata: dict | None = None,
tool_resources: OpenAICreateThreadParamsToolResources | None = None,
client: OpenAI | None = None,
**kwargs,
) -> Thread:
"""
@ -619,10 +612,10 @@ def create_thread(
elif timeout is None:
timeout = 600.0
api_base: Optional[str] = None
api_key: Optional[str] = None
api_base: str | None = None
api_key: str | None = None
response: Optional[Thread] = None
response: Thread | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -666,12 +659,10 @@ def create_thread(
or get_secret("AZURE_API_KEY")
) # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -695,9 +686,7 @@ def create_thread(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -712,7 +701,7 @@ def create_thread(
async def aget_thread(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> Thread:
loop = asyncio.get_event_loop()
@ -772,9 +761,9 @@ def get_thread(
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
api_base: Optional[str] = None
api_key: Optional[str] = None
response: Optional[Thread] = None
api_base: str | None = None
api_key: str | None = None
response: Thread | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -810,9 +799,7 @@ def get_thread(
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_key = (
optional_params.api_key
@ -823,7 +810,7 @@ def get_thread(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -846,9 +833,7 @@ def get_thread(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -868,8 +853,8 @@ async def a_add_message(
thread_id: str,
role: Literal["user", "assistant"],
content: str,
attachments: Optional[List[Attachment]] = None,
metadata: Optional[dict] = None,
attachments: list[Attachment] | None = None,
metadata: dict | None = None,
client=None,
**kwargs,
) -> OpenAIMessage:
@ -921,8 +906,8 @@ def add_message(
thread_id: str,
role: Literal["user", "assistant"],
content: str,
attachments: Optional[List[Attachment]] = None,
metadata: Optional[dict] = None,
attachments: list[Attachment] | None = None,
metadata: dict | None = None,
client=None,
**kwargs,
) -> OpenAIMessage:
@ -955,9 +940,9 @@ def add_message(
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
api_key: Optional[str] = None
api_base: Optional[str] = None
response: Optional[OpenAIMessage] = None
api_key: str | None = None
api_base: str | None = None
response: OpenAIMessage | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -993,9 +978,7 @@ def add_message(
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_key = (
optional_params.api_key
@ -1006,7 +989,7 @@ def add_message(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -1027,9 +1010,7 @@ def add_message(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -1045,7 +1026,7 @@ def add_message(
async def aget_messages(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
client: Optional[AsyncOpenAI] = None,
client: AsyncOpenAI | None = None,
**kwargs,
) -> AsyncCursorPage[OpenAIMessage]:
loop = asyncio.get_event_loop()
@ -1090,7 +1071,7 @@ async def aget_messages(
def get_messages(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
client: Optional[Any] = None,
client: Any | None = None,
**kwargs,
) -> SyncCursorPage[OpenAIMessage]:
aget_messages = kwargs.pop("aget_messages", None)
@ -1113,9 +1094,9 @@ def get_messages(
elif timeout is None:
timeout = 600.0
response: Optional[SyncCursorPage[OpenAIMessage]] = None
api_key: Optional[str] = None
api_base: Optional[str] = None
response: SyncCursorPage[OpenAIMessage] | None = None
api_key: str | None = None
api_base: str | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -1150,9 +1131,7 @@ def get_messages(
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_version: Optional[str] = (
optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
) # type: ignore
api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_key = (
optional_params.api_key
@ -1163,7 +1142,7 @@ def get_messages(
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
azure_ad_token: str | None = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
else:
@ -1183,9 +1162,7 @@ def get_messages(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'get_messages'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_messages'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -1201,7 +1178,7 @@ def get_messages(
### RUNS ###
def arun_thread_stream(
*,
event_handler: Optional[AssistantEventHandler] = None,
event_handler: AssistantEventHandler | None = None,
**kwargs,
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
kwargs["arun_thread"] = True
@ -1212,13 +1189,13 @@ async def arun_thread(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
assistant_id: str,
additional_instructions: Optional[str] = None,
instructions: Optional[str] = None,
metadata: Optional[dict] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
tools: Optional[Iterable[AssistantToolParam]] = None,
client: Optional[Any] = None,
additional_instructions: str | None = None,
instructions: str | None = None,
metadata: dict | None = None,
model: str | None = None,
stream: bool | None = None,
tools: Iterable[AssistantToolParam] | None = None,
client: Any | None = None,
**kwargs,
) -> Run:
loop = asyncio.get_event_loop()
@ -1269,7 +1246,7 @@ async def arun_thread(
def run_thread_stream(
*,
event_handler: Optional[AssistantEventHandler] = None,
event_handler: AssistantEventHandler | None = None,
**kwargs,
) -> AssistantStreamManager[AssistantEventHandler]:
return run_thread(stream=True, event_handler=event_handler, **kwargs) # type: ignore
@ -1279,14 +1256,14 @@ def run_thread(
custom_llm_provider: Literal["openai", "azure"],
thread_id: str,
assistant_id: str,
additional_instructions: Optional[str] = None,
instructions: Optional[str] = None,
metadata: Optional[dict] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
tools: Optional[Iterable[AssistantToolParam]] = None,
client: Optional[Any] = None,
event_handler: Optional[AssistantEventHandler] = None, # for stream=True calls
additional_instructions: str | None = None,
instructions: str | None = None,
metadata: dict | None = None,
model: str | None = None,
stream: bool | None = None,
tools: Iterable[AssistantToolParam] | None = None,
client: Any | None = None,
event_handler: AssistantEventHandler | None = None, # for stream=True calls
**kwargs,
) -> Run:
"""Run a given thread + assistant."""
@ -1310,7 +1287,7 @@ def run_thread(
elif timeout is None:
timeout = 600.0
response: Optional[Run] = None
response: Run | None = None
if custom_llm_provider == "openai":
api_base = (
optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
@ -1392,9 +1369,7 @@ def run_thread(
) # type: ignore
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'run_thread'. Only 'openai' is supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'run_thread'. Only 'openai' is supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(

43
litellm/assistants/utils.py
View file

@ -1,5 +1,3 @@
from typing import Optional, Union
import litellm
from ..exceptions import UnsupportedParamsError
@ -7,21 +5,10 @@ from ..types.llms.openai import *
def get_optional_params_add_message(
role: Optional[str],
content: Optional[
Union[
str,
List[
Union[
MessageContentTextObject,
MessageContentImageFileObject,
MessageContentImageURLObject,
]
],
]
],
attachments: Optional[List[Attachment]],
metadata: Optional[dict],
role: str | None,
content: str | List[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None,
attachments: List[Attachment] | None,
metadata: dict | None,
custom_llm_provider: str,
**kwargs,
):
@ -56,9 +43,7 @@ def get_optional_params_add_message(
elif k not in supported_params:
raise litellm.utils.UnsupportedParamsError(
status_code=500,
message="k={}, not supported by {}. Supported params={}. To drop it from the call, set `litellm.drop_params = True`.".format(
k, custom_llm_provider, supported_params
),
message=f"k={k}, not supported by {custom_llm_provider}. Supported params={supported_params}. To drop it from the call, set `litellm.drop_params = True`.",
)
return non_default_params
@ -71,19 +56,19 @@ def get_optional_params_add_message(
non_default_params=non_default_params, optional_params=optional_params
)
for k in passed_params.keys():
if k not in default_params.keys():
if k not in default_params:
optional_params[k] = passed_params[k]
return optional_params
def get_optional_params_image_gen(
n: Optional[int] = None,
quality: Optional[str] = None,
response_format: Optional[str] = None,
size: Optional[str] = None,
style: Optional[str] = None,
user: Optional[str] = None,
custom_llm_provider: Optional[str] = None,
n: int | None = None,
quality: str | None = None,
response_format: str | None = None,
size: str | None = None,
style: str | None = None,
user: str | None = None,
custom_llm_provider: str | None = None,
**kwargs,
):
# retrieve all parameters passed to the function
@ -142,6 +127,6 @@ def get_optional_params_image_gen(
optional_params["sampleCount"] = int(n)
for k in passed_params.keys():
if k not in default_params.keys():
if k not in default_params:
optional_params[k] = passed_params[k]
return optional_params

35
litellm/batch_completion/main.py
View file

@ -1,5 +1,4 @@
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from typing import List, Optional
import litellm
from litellm._logging import print_verbose
@ -11,23 +10,23 @@ from ..llms.vllm.completion import handler as vllm_handler
def batch_completion(
model: str,
# Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create
messages: List = [],
functions: Optional[List] = None,
function_call: Optional[str] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
n: Optional[int] = None,
stream: Optional[bool] = None,
messages: list = [],
functions: list | None = None,
function_call: str | None = None,
temperature: float | None = None,
top_p: float | None = None,
n: int | None = None,
stream: bool | None = None,
stop=None,
max_tokens: Optional[int] = None,
presence_penalty: Optional[float] = None,
frequency_penalty: Optional[float] = None,
logit_bias: Optional[dict] = None,
user: Optional[str] = None,
max_tokens: int | None = None,
presence_penalty: float | None = None,
frequency_penalty: float | None = None,
logit_bias: dict | None = None,
user: str | None = None,
deployment_id=None,
request_timeout: Optional[int] = None,
timeout: Optional[int] = 600,
max_workers: Optional[int] = 100,
request_timeout: int | None = None,
timeout: int | None = 600,
max_workers: int | None = 100,
# Optional liteLLM function params
**kwargs,
):
@ -164,7 +163,7 @@ def batch_completion_models(*args, **kwargs):
futures = {}
with ThreadPoolExecutor(max_workers=len(deployments)) as executor:
for deployment in deployments:
for key in kwargs.keys():
for key in kwargs:
if key not in deployment: # don't override deployment values e.g. model name, api base, etc.
deployment[key] = kwargs[key]
kwargs = {**deployment, **nested_kwargs}
@ -250,7 +249,7 @@ def batch_completion_models_all_responses(*args, **kwargs):
if result is not None:
responses.append(result)
except Exception as e:
print_verbose(f"batch_completion_models_all_responses: model request failed: {str(e)}")
print_verbose(f"batch_completion_models_all_responses: model request failed: {e}")
continue
return responses

47
litellm/batches/batch_utils.py
View file

@ -1,6 +1,7 @@
import json
from collections.abc import Iterable, Iterator
from dataclasses import dataclass
from typing import Any, Iterable, Iterator, List, Literal, Optional, Tuple
from typing import Any, Literal
import litellm
from litellm._logging import verbose_logger
@ -11,11 +12,11 @@ from litellm.utils import token_counter
async def calculate_batch_cost_and_usage(
file_content_dictionary: List[dict],
file_content_dictionary: list[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> Tuple[float, Usage, List[str]]:
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
"""
Calculate the cost and usage of a batch.
@ -44,9 +45,9 @@ async def calculate_batch_cost_and_usage(
async def _handle_completed_batch(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: Optional[str] = None,
litellm_params: Optional[dict] = None,
) -> Tuple[float, Usage, List[str]]:
model_name: str | None = None,
litellm_params: dict | None = None,
) -> tuple[float, Usage, list[str]]:
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
never materialized in memory.
@ -84,14 +85,14 @@ class _BatchOutputLineStats:
total_tokens: int
cache_read_tokens: int
cache_creation_tokens: int
model: Optional[str]
model: str | None
def _iter_successful_output_line_stats(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: Optional[str],
model_info: Optional[ModelInfo],
model_name: str | None,
model_info: ModelInfo | None,
) -> Iterator[_BatchOutputLineStats]:
from litellm.cost_calculator import batch_cost_calculator
@ -135,9 +136,9 @@ def _iter_successful_output_line_stats(
def _aggregate_batch_cost_usage_models(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> Tuple[float, Usage, List[str]]:
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
"""Aggregate cost, usage, and models from batch output entries in a single
pass, holding one small stats record per line instead of the parsed file."""
line_stats = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
@ -163,9 +164,9 @@ def _aggregate_batch_cost_usage_models(
def calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses: List[dict],
model_name: Optional[str] = None,
) -> Tuple[float, Usage]:
vertex_ai_batch_responses: list[dict],
model_name: str | None = None,
) -> tuple[float, Usage]:
"""
Calculate both cost and usage from raw Vertex AI batch responses.
@ -233,7 +234,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
async def _fetch_batch_output_file_content(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
litellm_params: Optional[dict] = None,
litellm_params: dict | None = None,
) -> bytes:
"""
Fetch the batch output file and return its raw JSONL bytes
@ -257,10 +258,10 @@ async def _fetch_batch_output_file_content(
if is_base64_unified_file_id:
try:
file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0]
verbose_logger.debug(f"Extracted LLM output file ID from unified file ID: {file_id}")
verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id)
except (IndexError, AttributeError) as e:
verbose_logger.error(
f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}"
"Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e
)
# Build kwargs for afile_content with credentials from litellm_params
@ -277,7 +278,7 @@ async def _fetch_batch_output_file_content(
return _file_content.content
def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
"""
Extract credentials from litellm_params for file access operations.
@ -316,7 +317,7 @@ def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
return credentials
def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]:
"""
Get the file content as a list of dictionaries from JSON Lines format
"""
@ -366,7 +367,7 @@ def _estimate_batch_entry_tokens(raw_line: bytes) -> int:
def _count_entry_tokens(
entry: dict,
model_name: Optional[str] = None,
model_name: str | None = None,
) -> int:
"""Token-count a single batch input entry's body (chat / text / embedding)."""
body = entry.get("body", {}) or {}

105
litellm/batches/main.py
View file

@ -13,8 +13,9 @@ https://platform.openai.com/docs/api-reference/batch
import asyncio
import contextvars
import os
from collections.abc import Coroutine
from functools import partial
from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast
from typing import Any, Literal, cast
import httpx
from openai.types.batch import BatchRequestCounts
@ -63,7 +64,7 @@ base_llm_http_handler = BaseLLMHTTPHandler()
def _resolve_timeout(
optional_params: GenericLiteLLMParams,
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
custom_llm_provider: str,
default_timeout: float = 600.0,
) -> float:
@ -106,10 +107,10 @@ async def acreate_batch(
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
output_expires_after: Optional[Dict[str, Any]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
output_expires_after: dict[str, Any] | None = None,
**kwargs,
) -> LiteLLMBatch:
"""
@ -156,12 +157,12 @@ def create_batch(
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
output_expires_after: Optional[Dict[str, Any]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
output_expires_after: dict[str, Any] | None = None,
**kwargs,
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
"""
Creates and executes a batch from an uploaded file of request
@ -172,7 +173,7 @@ def create_batch(
litellm_call_id = kwargs.get("litellm_call_id", None)
proxy_server_request = kwargs.get("proxy_server_request", None)
model_info = kwargs.get("model_info", None)
model: Optional[str] = kwargs.get("model", None)
model: str | None = kwargs.get("model", None)
try:
if model is not None:
model, _, _, _ = get_llm_provider(
@ -181,7 +182,7 @@ def create_batch(
)
except Exception as e:
verbose_logger.exception(
f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}"
"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - %s", e
)
_is_async = kwargs.pop("acreate_batch", False) is True
@ -237,7 +238,7 @@ def create_batch(
model=model,
)
return response
api_base: Optional[str] = None
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
@ -320,7 +321,7 @@ def create_batch(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support custom_llm_provider={} for 'create_batch'".format(custom_llm_provider),
message=f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'create_batch'",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -338,9 +339,9 @@ def create_batch(
async def aretrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMBatch:
"""
@ -379,14 +380,14 @@ async def aretrieve_batch(
def _handle_retrieve_batch_providers_without_provider_config(
batch_id: str,
optional_params: GenericLiteLLMParams,
timeout: Union[float, httpx.Timeout],
timeout: float | httpx.Timeout,
litellm_params: dict,
_retrieve_batch_request: RetrieveBatchRequest,
_is_async: bool,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
logging_obj: Optional[Any] = None,
logging_obj: Any | None = None,
):
api_base: Optional[str] = None
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
api_base = (
@ -488,10 +489,10 @@ def _handle_retrieve_batch_providers_without_provider_config(
else:
raise litellm.exceptions.BadRequestError(
message=(
"LiteLLM doesn't support custom_llm_provider={} for 'retrieve_batch' without a `model` kwarg. "
f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'retrieve_batch' without a `model` kwarg. "
"Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. "
"'bedrock' is supported but requires `model` to be passed so the provider config can be loaded."
).format(custom_llm_provider),
),
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(
@ -507,11 +508,11 @@ def _handle_retrieve_batch_providers_without_provider_config(
def retrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
"""
Retrieves a batch.
@ -519,7 +520,7 @@ def retrieve_batch(
"""
try:
optional_params = GenericLiteLLMParams(**kwargs)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None)
litellm_logging_obj: LiteLLMLoggingObj | None = kwargs.get("litellm_logging_obj", None)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
litellm_params = get_litellm_params(
@ -588,7 +589,7 @@ def retrieve_batch(
)
# Try to use provider config first (for providers like bedrock)
model: Optional[str] = kwargs.get("model", None)
model: str | None = kwargs.get("model", None)
if model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
@ -642,12 +643,12 @@ def retrieve_batch(
@client
async def alist_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
after: str | None = None,
limit: int | None = None,
custom_llm_provider: ListBatchesSupportedProvider = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -685,11 +686,11 @@ async def alist_batches(
@client
def list_batches(
after: Optional[str] = None,
limit: Optional[int] = None,
after: str | None = None,
limit: int | None = None,
custom_llm_provider: ListBatchesSupportedProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
):
"""
@ -822,11 +823,11 @@ def list_batches(
async def acancel_batch(
batch_id: str,
model: Optional[str] = None,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMBatch:
"""
@ -868,13 +869,13 @@ async def acancel_batch(
def cancel_batch(
batch_id: str,
model: Optional[str] = None,
custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] | str = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]:
"""
Cancels a batch.
@ -889,7 +890,7 @@ def cancel_batch(
)
except Exception as e:
verbose_logger.exception(
f"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - {str(e)}"
"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - %s", e
)
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params = get_litellm_params(
@ -919,7 +920,7 @@ def cancel_batch(
)
_is_async = kwargs.pop("acancel_batch", False) is True
api_base: Optional[str] = None
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
api_base = (
optional_params.api_base
@ -992,9 +993,7 @@ def cancel_batch(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format(
custom_llm_provider
),
message=f"LiteLLM doesn't support {custom_llm_provider} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(

16
litellm/budget_manager.py
View file

@ -11,7 +11,7 @@ import json
import os
import threading
import time
from typing import Literal, Optional
from typing import Literal
import litellm
from litellm.constants import (
@ -28,8 +28,8 @@ class BudgetManager:
self,
project_name: str,
client_type: str = "local",
api_base: Optional[str] = None,
headers: Optional[dict] = None,
api_base: str | None = None,
headers: dict | None = None,
):
self.client_type = client_type
self.project_name = project_name
@ -73,7 +73,7 @@ class BudgetManager:
self,
total_budget: float,
user: str,
duration: Optional[Literal["daily", "weekly", "monthly", "yearly"]] = None,
duration: Literal["daily", "weekly", "monthly", "yearly"] | None = None,
created_at: float = time.time(),
):
self.user_dict[user] = {"total_budget": total_budget}
@ -113,10 +113,10 @@ class BudgetManager:
def update_cost(
self,
user: str,
completion_obj: Optional[ModelResponse] = None,
model: Optional[str] = None,
input_text: Optional[str] = None,
output_text: Optional[str] = None,
completion_obj: ModelResponse | None = None,
model: str | None = None,
input_text: str | None = None,
output_text: str | None = None,
):
if model and input_text and output_text:
prompt_tokens = litellm.token_counter(model=model, messages=[{"role": "user", "content": input_text}])

2
litellm/caching/__init__.py
View file

@ -2,10 +2,10 @@ from .azure_blob_cache import AzureBlobCache
from .caching import Cache, LiteLLMCacheType
from .disk_cache import DiskCache
from .dual_cache import DualCache
from .gcs_cache import GCSCache
from .in_memory_cache import InMemoryCache
from .qdrant_semantic_cache import QdrantSemanticCache
from .redis_cache import RedisCache
from .redis_cluster_cache import RedisClusterCache
from .redis_semantic_cache import RedisSemanticCache
from .s3_cache import S3Cache
from .gcs_cache import GCSCache

5
litellm/caching/_internal_lru_cache.py
View file

@ -1,11 +1,12 @@
from collections.abc import Callable
from functools import lru_cache
from typing import Callable, Optional, TypeVar
from typing import TypeVar
T = TypeVar("T")
def lru_cache_wrapper(
maxsize: Optional[int] = None,
maxsize: int | None = None,
) -> Callable[[Callable[..., T]], Callable[..., T]]:
"""
Wrapper for lru_cache that caches success and exceptions

12
litellm/caching/azure_blob_cache.py
View file

@ -19,12 +19,12 @@ from .base_cache import BaseCache
class AzureBlobCache(BaseCache):
def __init__(self, account_url, container) -> None:
from azure.storage.blob import BlobServiceClient
from azure.core.exceptions import ResourceExistsError
from azure.identity import DefaultAzureCredential
from azure.identity.aio import (
DefaultAzureCredential as AsyncDefaultAzureCredential,
)
from azure.storage.blob import BlobServiceClient
from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
self.container_client = BlobServiceClient(
@ -67,7 +67,10 @@ class AzureBlobCache(BaseCache):
cached_response = json.loads(as_str)
verbose_logger.debug(
f"Got Azure Blob Cache: key: {key}, cached_response {cached_response}. Type Response {type(cached_response)}"
"Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
key,
cached_response,
type(cached_response),
)
return cached_response
@ -84,7 +87,10 @@ class AzureBlobCache(BaseCache):
as_str = as_bytes.decode("utf-8")
cached_response = json.loads(as_str)
verbose_logger.debug(
f"Got Azure Blob Cache: key: {key}, cached_response {cached_response}. Type Response {type(cached_response)}"
"Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s",
key,
cached_response,
type(cached_response),
)
return cached_response
except ResourceNotFoundError:

6
litellm/caching/base_cache.py
View file

@ -9,7 +9,7 @@ Has 4 methods:
"""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Union
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -23,8 +23,8 @@ class BaseCache(ABC):
def __init__(self, default_ttl: int = 60):
self.default_ttl = default_ttl
def get_ttl(self, **kwargs) -> Optional[int]:
kwargs_ttl: Optional[int] = kwargs.get("ttl")
def get_ttl(self, **kwargs) -> int | None:
kwargs_ttl: int | None = kwargs.get("ttl")
if kwargs_ttl is not None:
try:
return int(kwargs_ttl)

159
litellm/caching/caching.py
View file

@ -13,7 +13,7 @@ import json
import time
import traceback
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union
from typing import Any
from pydantic import BaseModel
@ -55,19 +55,18 @@ class CacheMode(str, Enum):
class Cache:
def __init__(
self,
type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL,
mode: Optional[
CacheMode
] = CacheMode.default_on, # when default_on cache is always on, when default_off cache is opt in
host: Optional[str] = None,
port: Optional[str] = None,
password: Optional[str] = None,
namespace: Optional[str] = None,
ttl: Optional[float] = None,
default_in_memory_ttl: Optional[float] = None,
default_in_redis_ttl: Optional[float] = None,
similarity_threshold: Optional[float] = None,
supported_call_types: Optional[List[CachingSupportedCallTypes]] = [
type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL,
mode: CacheMode
| None = CacheMode.default_on, # when default_on cache is always on, when default_off cache is opt in
host: str | None = None,
port: str | None = None,
password: str | None = None,
namespace: str | None = None,
ttl: float | None = None,
default_in_memory_ttl: float | None = None,
default_in_redis_ttl: float | None = None,
similarity_threshold: float | None = None,
supported_call_types: list[CachingSupportedCallTypes] | None = [
"completion",
"acompletion",
"embedding",
@ -82,38 +81,38 @@ class Cache:
"aresponses",
],
# s3 Bucket, boto3 configuration
azure_account_url: Optional[str] = None,
azure_blob_container: Optional[str] = None,
s3_bucket_name: Optional[str] = None,
s3_region_name: Optional[str] = None,
s3_api_version: Optional[str] = None,
s3_use_ssl: Optional[bool] = True,
s3_verify: Optional[Union[bool, str]] = None,
s3_endpoint_url: Optional[str] = None,
s3_aws_access_key_id: Optional[str] = None,
s3_aws_secret_access_key: Optional[str] = None,
s3_aws_session_token: Optional[str] = None,
s3_config: Optional[Any] = None,
s3_path: Optional[str] = None,
gcs_bucket_name: Optional[str] = None,
gcs_path_service_account: Optional[str] = None,
gcs_path: Optional[str] = None,
azure_account_url: str | None = None,
azure_blob_container: str | None = None,
s3_bucket_name: str | None = None,
s3_region_name: str | None = None,
s3_api_version: str | None = None,
s3_use_ssl: bool | None = True,
s3_verify: bool | str | None = None,
s3_endpoint_url: str | None = None,
s3_aws_access_key_id: str | None = None,
s3_aws_secret_access_key: str | None = None,
s3_aws_session_token: str | None = None,
s3_config: Any | None = None,
s3_path: str | None = None,
gcs_bucket_name: str | None = None,
gcs_path_service_account: str | None = None,
gcs_path: str | None = None,
redis_semantic_cache_embedding_model: str = "text-embedding-ada-002",
redis_semantic_cache_index_name: Optional[str] = None,
redis_semantic_cache_index_name: str | None = None,
valkey_semantic_cache_embedding_model: str = "text-embedding-ada-002",
valkey_semantic_cache_index_name: str | None = None,
redis_flush_size: Optional[int] = None,
redis_startup_nodes: Optional[List] = None,
disk_cache_dir: Optional[str] = None,
qdrant_api_base: Optional[str] = None,
qdrant_api_key: Optional[str] = None,
qdrant_collection_name: Optional[str] = None,
qdrant_quantization_config: Optional[str] = None,
redis_flush_size: int | None = None,
redis_startup_nodes: list | None = None,
disk_cache_dir: str | None = None,
qdrant_api_base: str | None = None,
qdrant_api_key: str | None = None,
qdrant_collection_name: str | None = None,
qdrant_quantization_config: str | None = None,
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
qdrant_semantic_cache_vector_size: Optional[int] = None,
qdrant_semantic_cache_vector_size: int | None = None,
# GCP IAM authentication parameters
gcp_service_account: Optional[str] = None,
gcp_ssl_ca_certs: Optional[str] = None,
gcp_service_account: str | None = None,
gcp_ssl_ca_certs: str | None = None,
**kwargs,
):
"""
@ -352,15 +351,15 @@ class Cache:
if param in scope_excluded_params:
continue
if param in combined_kwargs:
param_value: Optional[str] = self._get_param_value(param, kwargs)
param_value: str | None = self._get_param_value(param, kwargs)
if param_value is not None:
cache_key += f"{str(param)}: {str(param_value)}"
cache_key += f"{param}: {param_value}"
elif param not in litellm_param_kwargs: # check if user passed in optional param - e.g. top_k
if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now
if kwargs[param] is None:
continue # ignore None params
param_value = kwargs[param]
cache_key += f"{str(param)}: {str(param_value)}"
cache_key += f"{param}: {param_value}"
if is_semantic_cache:
cache_key += self._get_semantic_cache_tenant_scope(kwargs)
@ -382,7 +381,7 @@ class Cache:
self,
param: str,
kwargs: dict,
) -> Optional[str]:
) -> str | None:
"""
Get the value for the given param from kwargs
"""
@ -400,15 +399,15 @@ class Cache:
2. Else if a model_group is set, then return the model_group as the model. This is used for all requests sent through the litellm.Router()
3. Else use the `model` passed in kwargs
"""
metadata: Dict = kwargs.get("metadata", {}) or {}
litellm_params: Dict = kwargs.get("litellm_params", {}) or {}
metadata_in_litellm_params: Dict = litellm_params.get("metadata", {}) or {}
model_group: Optional[str] = metadata.get("model_group") or metadata_in_litellm_params.get("model_group")
metadata: dict = kwargs.get("metadata", {}) or {}
litellm_params: dict = kwargs.get("litellm_params", {}) or {}
metadata_in_litellm_params: dict = litellm_params.get("metadata", {}) or {}
model_group: str | None = metadata.get("model_group") or metadata_in_litellm_params.get("model_group")
caching_group = self._get_caching_group(metadata, model_group)
return caching_group or model_group or kwargs["model"]
def _get_caching_group(self, metadata: dict, model_group: Optional[str]) -> Optional[str]:
caching_groups: Optional[List] = metadata.get("caching_groups", [])
def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None:
caching_groups: list | None = metadata.get("caching_groups", [])
if caching_groups:
for group in caching_groups:
if model_group in group:
@ -429,7 +428,7 @@ class Cache:
or litellm_params.get("file_name")
)
def _get_preset_cache_key_from_kwargs(self, **kwargs) -> Optional[str]:
def _get_preset_cache_key_from_kwargs(self, **kwargs) -> str | None:
"""
Get the preset cache key from kwargs["litellm_params"]
@ -510,8 +509,8 @@ class Cache:
def _get_cache_logic(
self,
cached_result: Optional[Any],
max_age: Optional[float],
cached_result: Any | None,
max_age: float | None,
):
"""
Common get cache logic across sync + async implementations
@ -544,8 +543,8 @@ class Cache:
return cached_result
@staticmethod
def _get_safe_cache_lookup_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
cache_lookup_kwargs: Dict[str, Any] = {}
def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
cache_lookup_kwargs: dict[str, Any] = {}
for prompt_kwarg in ("messages", "input"):
if prompt_kwarg in kwargs:
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
@ -558,7 +557,7 @@ class Cache:
@staticmethod
def _update_metadata_from_cache_lookup_kwargs(
original_kwargs: Dict[str, Any], cache_lookup_kwargs: Dict[str, Any]
original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any]
) -> None:
original_metadata = original_kwargs.get("metadata")
cache_lookup_metadata = cache_lookup_kwargs.get("metadata")
@ -568,7 +567,7 @@ class Cache:
if "semantic-similarity" in cache_lookup_metadata:
original_metadata["semantic-similarity"] = cache_lookup_metadata["semantic-similarity"]
def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
def get_cache(self, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Retrieves the cached result for the given arguments.
@ -603,7 +602,7 @@ class Cache:
print_verbose(f"An exception occurred: {traceback.format_exc()}")
return None
async def async_get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
async def async_get_cache(self, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Async get cache implementation.
@ -677,9 +676,9 @@ class Cache:
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
self.cache.set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
async def async_add_cache(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Async implementation of add_cache
"""
@ -696,14 +695,14 @@ class Cache:
else:
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
def _convert_to_cached_embedding(
self,
embedding_response: Any,
model: Optional[str],
prompt_tokens: Optional[int] = None,
prompt_tokens_details: Optional[dict] = None,
model: str | None,
prompt_tokens: int | None = None,
prompt_tokens_details: dict | None = None,
) -> CachedEmbedding:
"""
Convert any embedding response into the standardized CachedEmbedding TypedDict format.
@ -745,7 +744,7 @@ class Cache:
self,
result: EmbeddingResponse,
idx_in_result_data: int,
) -> Optional[dict]:
) -> dict | None:
"""
Extract per-item prompt_tokens_details from a response for caching.
@ -788,7 +787,7 @@ class Cache:
self,
result: EmbeddingResponse,
idx_in_result_data: int,
) -> Optional[int]:
) -> int | None:
"""
Extract the per-item prompt_tokens from a response for caching.
@ -813,7 +812,7 @@ class Cache:
input: str,
kwargs: dict,
idx_in_result_data: int = 0,
) -> Tuple[str, dict, dict]:
) -> tuple[str, dict, dict]:
preset_cache_key = self.get_cache_key(**{**kwargs, "input": input})
kwargs["cache_key"] = preset_cache_key
embedding_response = result.data[idx_in_result_data]
@ -843,7 +842,7 @@ class Cache:
)
return cache_key, cached_data, kwargs
async def async_add_cache_pipeline(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
async def async_add_cache_pipeline(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
"""
Async implementation of add_cache for Embedding calls
@ -875,7 +874,7 @@ class Cache:
else:
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
def should_use_cache(self, **kwargs):
"""
@ -926,11 +925,11 @@ class Cache:
def enable_cache(
type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL,
host: Optional[str] = None,
port: Optional[str] = None,
password: Optional[str] = None,
supported_call_types: Optional[List[CachingSupportedCallTypes]] = [
type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL,
host: str | None = None,
port: str | None = None,
password: str | None = None,
supported_call_types: list[CachingSupportedCallTypes] | None = [
"completion",
"acompletion",
"embedding",
@ -986,11 +985,11 @@ def enable_cache(
def update_cache(
type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL,
host: Optional[str] = None,
port: Optional[str] = None,
password: Optional[str] = None,
supported_call_types: Optional[List[CachingSupportedCallTypes]] = [
type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL,
host: str | None = None,
port: str | None = None,
password: str | None = None,
supported_call_types: list[CachingSupportedCallTypes] | None = [
"completion",
"acompletion",
"embedding",

122
litellm/caching/caching_handler.py
View file

@ -18,17 +18,11 @@ import asyncio
import datetime
import inspect
import time
from collections.abc import AsyncGenerator, Callable, Generator
from typing import (
TYPE_CHECKING,
Any,
AsyncGenerator,
Callable,
Dict,
Generator,
List,
Optional,
Tuple,
Union,
)
from pydantic import BaseModel
@ -77,8 +71,8 @@ class CachingHandlerResponse(BaseModel):
For embeddings there can be a cache hit for some of the inputs in the list and a cache miss for others
"""
cached_result: Optional[Any] = None
final_embedding_cached_response: Optional[EmbeddingResponse] = None
cached_result: Any | None = None
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
@ -111,7 +105,7 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
return "choices" in cached_result
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool:
def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bool:
"""
When stream=True, do not run success callbacks at cache-hit time.
@ -127,25 +121,24 @@ class LLMCachingHandler:
def __init__(
self,
original_function: Callable,
request_kwargs: Dict[str, Any],
request_kwargs: dict[str, Any],
start_time: datetime.datetime,
):
from litellm.caching import DualCache, RedisCache
self.async_streaming_chunks: List[ModelResponse] = []
self.sync_streaming_chunks: List[ModelResponse] = []
self.async_streaming_chunks: list[ModelResponse] = []
self.sync_streaming_chunks: list[ModelResponse] = []
self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs)
self.preset_cache_key: Optional[str] = None
self.preset_cache_key: str | None = None
self.original_function = original_function
self.start_time = start_time
if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
self.dual_cache: Optional[DualCache] = DualCache(
self.dual_cache: DualCache | None = DualCache(
redis_cache=litellm.cache.cache,
in_memory_cache=in_memory_cache_obj,
)
else:
self.dual_cache = None
pass
async def _async_get_cache(
self,
@ -154,9 +147,9 @@ class LLMCachingHandler:
logging_obj: LiteLLMLoggingObj,
start_time: datetime.datetime,
call_type: str,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
) -> Optional[CachingHandlerResponse]:
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
) -> CachingHandlerResponse | None:
"""
Internal method to get from the cache.
Handles different call types (embeddings, chat/completions, text_completion, transcription)
@ -184,15 +177,15 @@ class LLMCachingHandler:
kwargs.get("cache", {}).get("no-cache", False) is not True
): # allow users to control returning cached responses from the completion function
args = args or ()
final_embedding_cached_response: Optional[EmbeddingResponse] = None
final_embedding_cached_response: EmbeddingResponse | None = None
embedding_all_elements_cache_hit: bool = False
cached_result: Optional[Any] = None
cached_result: Any | None = None
kwargs = kwargs.copy()
#########################################################
# Init cache timing metrics
#########################################################
cache_check_start_time = time.perf_counter()
cache_check_end_time: Optional[float] = None
cache_check_end_time: float | None = None
#########################################################
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
kwargs["parent_otel_span"] = parent_otel_span
@ -278,7 +271,7 @@ class LLMCachingHandler:
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
)
verbose_logger.debug(f"CACHE RESULT: {cached_result}")
verbose_logger.debug("CACHE RESULT: %s", cached_result)
return CachingHandlerResponse(
cached_result=cached_result,
final_embedding_cached_response=final_embedding_cached_response,
@ -293,10 +286,10 @@ class LLMCachingHandler:
logging_obj: LiteLLMLoggingObj,
start_time: datetime.datetime,
call_type: str,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
) -> CachingHandlerResponse:
cached_result: Optional[Any] = None
cached_result: Any | None = None
# Check if caching should be performed BEFORE doing expensive kwargs copy
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
@ -371,7 +364,7 @@ class LLMCachingHandler:
return CachingHandlerResponse(cached_result=cached_result)
return CachingHandlerResponse(cached_result=cached_result)
def handle_kwargs_input_list_or_str(self, kwargs: Dict[str, Any]) -> List[str]:
def handle_kwargs_input_list_or_str(self, kwargs: dict[str, Any]) -> list[str]:
"""
Handles the input of kwargs['input'] being a list or a string
"""
@ -382,7 +375,7 @@ class LLMCachingHandler:
else:
raise ValueError("input must be a string or a list")
def _extract_model_from_cached_results(self, non_null_list: List[Tuple[int, CachedEmbedding]]) -> Optional[str]:
def _extract_model_from_cached_results(self, non_null_list: list[tuple[int, CachedEmbedding]]) -> str | None:
"""
Helper method to extract the model name from cached results.
@ -399,13 +392,13 @@ class LLMCachingHandler:
def _process_async_embedding_cached_response(
self,
final_embedding_cached_response: Optional[EmbeddingResponse],
cached_result: List[Optional[CachedEmbedding]],
kwargs: Dict[str, Any],
final_embedding_cached_response: EmbeddingResponse | None,
cached_result: list[CachedEmbedding | None],
kwargs: dict[str, Any],
logging_obj: LiteLLMLoggingObj,
start_time: datetime.datetime,
model: str,
) -> Tuple[Optional[EmbeddingResponse], bool]:
) -> tuple[EmbeddingResponse | None, bool]:
"""
Returns the final embedding cached response and a boolean indicating if all elements in the list have a cache hit
@ -448,7 +441,7 @@ class LLMCachingHandler:
final_embedding_cached_response._hidden_params["cache_hit"] = True
prompt_tokens = 0
aggregated_details: Optional[dict] = None
aggregated_details: dict | None = None
for val in non_null_list:
idx, cr = val # (idx, cr) tuple
if cr is not None:
@ -478,10 +471,10 @@ class LLMCachingHandler:
aggregated_details[key] = value
## USAGE
prompt_tokens_details: Optional["PromptTokensDetailsWrapper"] = None
if aggregated_details:
from litellm.types.utils import PromptTokensDetailsWrapper
from litellm.types.utils import PromptTokensDetailsWrapper
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
if aggregated_details:
try:
prompt_tokens_details = PromptTokensDetailsWrapper(**aggregated_details)
except Exception:
@ -676,9 +669,7 @@ class LLMCachingHandler:
cache_hit=cache_hit,
)
async def _retrieve_from_cache(
self, call_type: str, kwargs: Dict[str, Any], args: Tuple[Any, ...]
) -> Optional[Any]:
async def _retrieve_from_cache(self, call_type: str, kwargs: dict[str, Any], args: tuple[Any, ...]) -> Any | None:
"""
Internal method to
- get cache key
@ -711,7 +702,7 @@ class LLMCachingHandler:
if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs)
cached_result: Optional[Any] = None
cached_result: Any | None = None
if call_type == CallTypes.aembedding.value:
if isinstance(new_kwargs["input"], str):
new_kwargs["input"] = [new_kwargs["input"]]
@ -756,21 +747,20 @@ class LLMCachingHandler:
self,
cached_result: Any,
call_type: str,
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
logging_obj: LiteLLMLoggingObj,
model: str,
args: Tuple[Any, ...],
custom_llm_provider: Optional[str] = None,
) -> Optional[
Union[
ModelResponse,
TextCompletionResponse,
EmbeddingResponse,
RerankResponse,
TranscriptionResponse,
CustomStreamWrapper,
]
]:
args: tuple[Any, ...],
custom_llm_provider: str | None = None,
) -> (
ModelResponse
| TextCompletionResponse
| EmbeddingResponse
| RerankResponse
| TranscriptionResponse
| CustomStreamWrapper
| None
):
"""
Internal method to process the cached result
@ -923,7 +913,7 @@ class LLMCachingHandler:
convert_to_streaming_response_async,
)
_stream_cached_result: Union[AsyncGenerator, Generator]
_stream_cached_result: AsyncGenerator | Generator
if call_type == CallTypes.acompletion.value or call_type == CallTypes.atext_completion.value:
_stream_cached_result = convert_to_streaming_response_async(
response_object=cached_result,
@ -943,8 +933,8 @@ class LLMCachingHandler:
self,
result: Any,
original_function: Callable,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
):
"""
Internal method to check the type of the result & cache used and adds the result to the cache accordingly
@ -1009,8 +999,8 @@ class LLMCachingHandler:
def sync_set_cache(
self,
result: Any,
kwargs: Dict[str, Any],
args: Optional[Tuple[Any, ...]] = None,
kwargs: dict[str, Any],
args: tuple[Any, ...] | None = None,
):
"""
Sync internal method to add the result to the cache
@ -1031,7 +1021,7 @@ class LLMCachingHandler:
return
def _should_store_result_in_cache(self, original_function: Callable, kwargs: Dict[str, Any]) -> bool:
def _should_store_result_in_cache(self, original_function: Callable, kwargs: dict[str, Any]) -> bool:
"""
Helper function to determine if the result should be stored in the cache.
@ -1077,7 +1067,7 @@ class LLMCachingHandler:
"""
complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = (
complete_streaming_response: ModelResponse | TextCompletionResponse | None = (
_assemble_complete_response_from_streaming_chunks(
result=processed_chunk,
start_time=self.start_time,
@ -1099,7 +1089,7 @@ class LLMCachingHandler:
"""
Sync internal method to add the streaming response to the cache
"""
complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = (
complete_streaming_response: ModelResponse | TextCompletionResponse | None = (
_assemble_complete_response_from_streaming_chunks(
result=processed_chunk,
start_time=self.start_time,
@ -1121,12 +1111,12 @@ class LLMCachingHandler:
self,
logging_obj: LiteLLMLoggingObj,
model: str,
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
cached_result: Any,
is_async: bool,
is_embedding: bool = False,
custom_llm_provider: Optional[str] = None,
cache_duration_ms: Optional[float] = None,
custom_llm_provider: str | None = None,
cache_duration_ms: float | None = None,
):
"""
Helper function to update the LiteLLMLoggingObj environment variables.
@ -1180,8 +1170,8 @@ class LLMCachingHandler:
def convert_args_to_kwargs(
original_function: Callable,
args: Optional[Tuple[Any, ...]] = None,
) -> Dict[str, Any]:
args: tuple[Any, ...] | None = None,
) -> dict[str, Any]:
# Get the signature of the original function
signature = inspect.signature(original_function)

4
litellm/caching/disk_cache.py
View file

@ -1,5 +1,5 @@
import json
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Union
from .base_cache import BaseCache
@ -12,7 +12,7 @@ else:
class DiskCache(BaseCache):
def __init__(self, disk_cache_dir: Optional[str] = None):
def __init__(self, disk_cache_dir: str | None = None):
try:
import diskcache as dc
except ModuleNotFoundError as e:

64
litellm/caching/dual_cache.py
View file

@ -13,7 +13,7 @@ import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from threading import Lock
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Union
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
@ -57,11 +57,11 @@ class DualCache(BaseCache):
def __init__(
self,
in_memory_cache: Optional[InMemoryCache] = None,
redis_cache: Optional[RedisCache] = None,
default_in_memory_ttl: Optional[float] = None,
default_redis_ttl: Optional[float] = None,
default_redis_batch_cache_expiry: Optional[float] = None,
in_memory_cache: InMemoryCache | None = None,
redis_cache: RedisCache | None = None,
default_in_memory_ttl: float | None = None,
default_redis_ttl: float | None = None,
default_redis_batch_cache_expiry: float | None = None,
default_max_redis_batch_cache_size: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE,
) -> None:
super().__init__()
@ -77,7 +77,7 @@ class DualCache(BaseCache):
self.default_in_memory_ttl = default_in_memory_ttl or litellm.default_in_memory_ttl
self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl
def update_cache_ttl(self, default_in_memory_ttl: Optional[float], default_redis_ttl: Optional[float]):
def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None):
if default_in_memory_ttl is not None:
self.default_in_memory_ttl = default_in_memory_ttl
@ -86,9 +86,9 @@ class DualCache(BaseCache):
def attach_redis_cache(
self,
redis_cache: Optional[RedisCache] = None,
redis_cache: RedisCache | None = None,
*,
default_redis_ttl: Optional[float] = None,
default_redis_ttl: float | None = None,
) -> None:
"""
Attach a Redis backend if this DualCache does not already have one.
@ -147,13 +147,13 @@ class DualCache(BaseCache):
return result
except Exception as e:
verbose_logger.error(f"LiteLLM Cache: Excepton async add_cache: {str(e)}")
verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e)
raise e
def get_cache(
self,
key,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -184,7 +184,7 @@ class DualCache(BaseCache):
def batch_get_cache(
self,
keys: list,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -217,7 +217,7 @@ class DualCache(BaseCache):
async def async_get_cache(
self,
key,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -250,15 +250,15 @@ class DualCache(BaseCache):
def _reserve_redis_batch_keys(
self,
current_time: float,
keys: List[str],
result: List[Any],
) -> Tuple[List[str], Dict[str, Optional[float]]]:
keys: list[str],
result: list[Any],
) -> tuple[list[str], dict[str, float | None]]:
"""
Atomically choose keys to fetch from Redis and reserve their access time.
This prevents check-then-act races under concurrent async callers.
"""
sublist_keys: List[str] = []
previous_access_times: Dict[str, Optional[float]] = {}
sublist_keys: list[str] = []
previous_access_times: dict[str, float | None] = {}
with self._last_redis_batch_access_time_lock:
for key, value in zip(keys, result):
@ -275,7 +275,7 @@ class DualCache(BaseCache):
return sublist_keys, previous_access_times
def _rollback_redis_batch_key_reservations(self, previous_access_times: Dict[str, Optional[float]]) -> None:
def _rollback_redis_batch_key_reservations(self, previous_access_times: dict[str, float | None]) -> None:
with self._last_redis_batch_access_time_lock:
for key, previous_time in previous_access_times.items():
if previous_time is None:
@ -286,7 +286,7 @@ class DualCache(BaseCache):
async def async_batch_get_cache(
self,
keys: list,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
**kwargs,
):
@ -347,7 +347,7 @@ class DualCache(BaseCache):
if self.redis_cache is not None and local_only is False:
await self.redis_cache.async_set_cache(key, value, **kwargs)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {str(e)}")
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
# async_batch_set_cache
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs):
@ -366,17 +366,17 @@ class DualCache(BaseCache):
cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs
)
except Exception as e:
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {str(e)}")
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
async def async_increment_cache(
self,
key,
value: float,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
local_only: bool = False,
refresh_ttl: bool = False,
**kwargs,
) -> Optional[float]:
) -> float | None:
"""
Key - the key in cache
@ -388,7 +388,7 @@ class DualCache(BaseCache):
Returns - the incremented value, or None if no cache backend is
available (in_memory_cache is None and Redis failed/is absent).
"""
result: Optional[float] = None
result: float | None = None
try:
if self.in_memory_cache is not None:
result = await self.in_memory_cache.async_increment(key, value, **kwargs)
@ -412,12 +412,12 @@ class DualCache(BaseCache):
async def async_increment_cache_pipeline(
self,
increment_list: List["RedisPipelineIncrementOperation"],
increment_list: list["RedisPipelineIncrementOperation"],
local_only: bool = False,
parent_otel_span: Optional[Span] = None,
parent_otel_span: Span | None = None,
**kwargs,
) -> Optional[List[float]]:
result: Optional[List[float]] = None
) -> list[float] | None:
result: list[float] | None = None
try:
if self.in_memory_cache is not None:
result = await self.in_memory_cache.async_increment_pipeline(
@ -439,7 +439,7 @@ class DualCache(BaseCache):
)
return result
async def async_set_cache_sadd(self, key, value: List, local_only: bool = False, **kwargs) -> None:
async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None:
"""
Add value to a set
@ -456,7 +456,7 @@ class DualCache(BaseCache):
if self.redis_cache is not None and local_only is False:
_ = await self.redis_cache.async_set_cache_sadd(key, value, ttl=kwargs.get("ttl", None))
return None
return
except Exception as e:
raise e # don't log, if exception is raised
@ -484,7 +484,7 @@ class DualCache(BaseCache):
if self.redis_cache is not None:
await self.redis_cache.async_delete_cache(key)
async def async_get_ttl(self, key: str) -> Optional[int]:
async def async_get_ttl(self, key: str) -> int | None:
"""
Get the remaining TTL of a key in in-memory cache or redis
"""

276
litellm/caching/evicted_client_closer.py Normal file
View file

@ -0,0 +1,276 @@
"""
Deferred close of HTTP/SDK clients that the LLM client cache has evicted.
Eviction only drops the cache's reference to a client. Every OpenAI/Azure SDK
client is a reference cycle (each resource namespace holds the client back), so
an evicted client and its pooled TCP connections survive until a generational
collection runs, which under load is thousands of requests later.
Closing at eviction time is not an option: a request that was handed the client
just before it was evicted is still using it, and closing it underneath that
request raises ``RuntimeError: Cannot send a request, as the client has been
closed.``
So an evicted client is closed once two conditions hold. A grace window must
have passed since its eviction, which covers a request that holds the client
but is momentarily not on the wire, and the client must report no connection in
flight. The second condition is what keeps the first honest: a request may run
for ``litellm.request_timeout`` seconds, 6000 by default, and a streaming
response is bounded only by how long the upstream keeps sending, so no deadline
on its own can promise that a request has finished.
Only clients litellm itself created are closed; a client the caller supplied is
left alone because litellm does not own its lifecycle.
A client that closes synchronously is closed from wherever the cache is next
used. One whose close is a coroutine needs the event loop it was evicted on, so
it waits for a call from that loop rather than having work scheduled onto a loop
it does not belong to. Queued clients are therefore bucketed by what it takes to
close them, and each bucket is ordered by deadline, so a reap walks the entries
that are due rather than the whole queue.
The queue holds its clients weakly, so waiting out a grace window never keeps
alive anything the collector would have reclaimed first.
"""
import asyncio
import inspect
import threading
import time
import weakref
from collections import deque
from collections.abc import Awaitable, Callable, Iterator
from dataclasses import dataclass, replace
from litellm.constants import (
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
)
_CLOSABLE_ANYWHERE = "closable-anywhere"
_CLOSABLE_ON_ANY_LOOP = "closable-on-any-loop"
_BucketKey = str | int
@dataclass(frozen=True, slots=True)
class _PendingClose:
"""A queued close.
The client is held weakly, so queueing one never keeps alive anything the
collector would otherwise have reclaimed first.
``needs_loop`` is set for a client whose close is a coroutine; those can only
be closed from the event loop they were evicted on, recorded in ``loop_id``.
A client that closes synchronously carries neither constraint.
"""
client_ref: "weakref.ref[object]"
loop_id: int | None
needs_loop: bool
close_after: float
def _bucket_key(pending: _PendingClose) -> _BucketKey:
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
if not pending.needs_loop:
return _CLOSABLE_ANYWHERE
if pending.loop_id is None:
return _CLOSABLE_ON_ANY_LOOP
return pending.loop_id
def _running_loop_id() -> int | None:
try:
return id(asyncio.get_running_loop())
except RuntimeError:
return None
def _close_function(client: object) -> Callable[[], object] | None:
close_fn: Callable[[], object] | None = getattr(client, "aclose", None) or getattr(client, "close", None)
return close_fn
def _transport_of(client: object) -> object:
"""The httpx transport behind an SDK wrapper, a litellm handler, or a bare client."""
for holder in (getattr(client, "_client", None), getattr(client, "client", None), client):
transport: object = getattr(holder, "_transport", None)
if transport is not None:
return transport
return None
def _connection_is_idle(connection: object) -> bool:
"""A pooled connection is idle unless it is servicing a request."""
is_idle: object = getattr(connection, "is_idle", None)
return bool(is_idle()) if callable(is_idle) else True
def _pool_has_busy_connection(transport: object) -> bool | None:
"""Whether the httpcore pool behind the transport is servicing a request.
``None`` when there is no such pool, so the caller can ask the other backend.
"""
pooled: object = getattr(getattr(transport, "_pool", None), "connections", None)
if not isinstance(pooled, (list, tuple)):
return None
return any(
not _connection_is_idle(connection) # pyright: ignore[reportUnknownArgumentType] # untyped pool list
for connection in pooled # pyright: ignore[reportUnknownVariableType] # untyped pool list
)
def _has_connection_in_flight(client: object) -> bool:
"""Whether the client is servicing a request right now.
Both connection backends litellm uses already account for the connections
they have handed out, so this reads the client's own lease accounting rather
than inferring it from elapsed time: httpcore reports a non-idle connection
for the whole of a response including a stream, and aiohttp holds the
connection in ``_acquired`` over the same span.
A client that cannot answer is reported as idle, which leaves the grace
window as the only guard, exactly as it was before this check existed.
"""
try:
transport = _transport_of(client)
pooled_busy = _pool_has_busy_connection(transport)
if pooled_busy is not None:
return pooled_busy
session: object = getattr(transport, "client", None)
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle
return False
async def _close_quietly(closing: Awaitable[object]) -> None:
try:
await closing
except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers
pass
class EvictedClientCloser:
"""Closes evicted, litellm-owned clients once they are idle and out of grace."""
def __init__(
self,
grace_seconds: float = EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
max_pending: int = EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
clock: Callable[[], float] = time.monotonic,
) -> None:
self._grace_seconds = grace_seconds
self._max_pending = max_pending
self._clock = clock
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
self._pending_count = 0
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
def mark_owned(self, client: object) -> None:
"""Record that litellm created this client, so it may be closed on eviction."""
try:
self._owned.add(client)
except TypeError:
pass # values that cannot be weak-referenced are never litellm clients
def _is_owned(self, client: object) -> bool:
try:
return client in self._owned
except TypeError:
return False # unhashable values are never litellm clients
def schedule(self, client: object) -> None:
"""Queue an evicted client for closing once it is idle and out of grace.
Past ``max_pending`` the client is left to the collector instead, so a
workload that churns the cache cannot grow this queue without bound.
Every queued entry comes due within one grace window, so the capacity it
occupies is returned within that window rather than held.
"""
if client is None or not self._is_owned(client):
return
close_fn = _close_function(client)
if close_fn is None:
return
if self._pending_count >= self._max_pending:
return
self._enqueue(
_PendingClose(
client_ref=weakref.ref(client),
loop_id=_running_loop_id(),
needs_loop=inspect.iscoroutinefunction(close_fn),
close_after=self._clock() + self._grace_seconds,
)
)
def reap(self) -> None:
"""Close every queued client that is due, idle, and closable from here.
Called from the cache's read path, so the empty-queue exit comes first and
the work done past it is proportional to what is due, not to the queue.
"""
if not self._pending_count:
return
now = self._clock()
for pending in self._take_due(_running_loop_id(), now):
client = pending.client_ref()
if client is None:
continue
if _has_connection_in_flight(client):
self._enqueue(replace(pending, close_after=now + self._grace_seconds))
continue
self._close(client)
@property
def pending_count(self) -> int:
return self._pending_count
def _enqueue(self, pending: _PendingClose) -> None:
"""Append to the entry's bucket, dropping any dead entries it queues behind.
Deadlines only ever move forward, so appending keeps each bucket ordered
by deadline, and entries whose client the collector already took sit at
the front rather than having to be searched for.
"""
with self._queue_lock:
bucket = self._buckets.setdefault(_bucket_key(pending), deque()) # mutable-ok: FIFO by design
while bucket and bucket[0].client_ref() is None:
bucket.popleft()
self._pending_count -= 1
bucket.append(pending)
self._pending_count += 1
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
with self._queue_lock:
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:
bucket = self._buckets.get(key)
if bucket is None:
return
while bucket and bucket[0].close_after <= now:
self._pending_count -= 1
yield bucket.popleft()
if not bucket:
del self._buckets[key]
def _close(self, client: object) -> None:
close_fn = _close_function(client)
if close_fn is None:
return
try:
closing = close_fn()
except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers
return
if not inspect.isawaitable(closing):
return
task = asyncio.get_running_loop().create_task(_close_quietly(closing))
self._close_tasks.add(task)
task.add_done_callback(self._close_tasks.discard)
default_evicted_client_closer = EvictedClientCloser()

21
litellm/caching/gcs_cache.py
View file

@ -2,27 +2,27 @@
Supports syncing responses to Google Cloud Storage Buckets using HTTP requests.
"""
import json
import asyncio
from typing import Optional
import json
from urllib.parse import quote
from litellm._logging import print_verbose, verbose_logger
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
_get_httpx_client,
get_async_httpx_client,
httpxSpecialProvider,
)
from .base_cache import BaseCache
class GCSCache(BaseCache):
def __init__(
self,
bucket_name: Optional[str] = None,
path_service_account: Optional[str] = None,
gcs_path: Optional[str] = None,
bucket_name: str | None = None,
path_service_account: str | None = None,
gcs_path: str | None = None,
) -> None:
super().__init__()
self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME
@ -71,12 +71,15 @@ class GCSCache(BaseCache):
if response.status_code == 200:
cached_response = json.loads(response.text)
verbose_logger.debug(
f"Got GCS Cache: key: {key}, cached_response {cached_response}. Type Response {type(cached_response)}"
"Got GCS Cache: key: %s, cached_response %s. Type Response %s",
key,
cached_response,
type(cached_response),
)
return cached_response
return None
except Exception as e:
verbose_logger.error(f"GCS Caching: get_cache() - Got exception from GCS: {e}")
verbose_logger.error("GCS Caching: get_cache() - Got exception from GCS: %s", e)
async def async_get_cache(self, key, **kwargs):
try:
@ -89,7 +92,7 @@ class GCSCache(BaseCache):
return json.loads(response.text)
return None
except Exception as e:
verbose_logger.error(f"GCS Caching: async_get_cache() - Got exception from GCS: {e}")
verbose_logger.error("GCS Caching: async_get_cache() - Got exception from GCS: %s", e)
def flush_cache(self):
pass

29
litellm/caching/in_memory_cache.py
View file

@ -8,12 +8,12 @@ Has 4 methods:
- async_get_cache
"""
import heapq
import json
import sys
import time
import heapq
import threading
from typing import TYPE_CHECKING, Any, List, Optional
import time
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
@ -28,11 +28,10 @@ from .base_cache import BaseCache
class InMemoryCache(BaseCache):
def __init__(
self,
max_size_in_memory: Optional[int] = 200,
default_ttl: Optional[
int
] = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute
max_size_per_item: Optional[int] = 1024, # 1MB = 1024KB
max_size_in_memory: int | None = 200,
default_ttl: int
| None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute
max_size_per_item: int | None = 1024, # 1MB = 1024KB
):
"""
max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default
@ -146,9 +145,7 @@ class InMemoryCache(BaseCache):
Check if ttl is set for a key
"""
ttl_time = self.ttl_dict.get(key)
if ttl_time is None: # if ttl is not set, allow override
return True
elif float(ttl_time) < time.time(): # if ttl is expired, allow override
if ttl_time is None or float(ttl_time) < time.time(): # if ttl is not set, allow override
return True
else:
return False
@ -184,7 +181,7 @@ class InMemoryCache(BaseCache):
else:
self.set_cache(key=cache_key, value=cache_value)
async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float]):
async def async_set_cache_sadd(self, key, value: list, ttl: float | None):
"""
Add value to set
"""
@ -247,8 +244,8 @@ class InMemoryCache(BaseCache):
return self.increment_cache(key=key, value=value, **kwargs)
async def async_increment_pipeline(
self, increment_list: List["RedisPipelineIncrementOperation"], **kwargs
) -> Optional[List[float]]:
self, increment_list: list["RedisPipelineIncrementOperation"], **kwargs
) -> list[float] | None:
results = []
for increment in increment_list:
result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs)
@ -266,13 +263,13 @@ class InMemoryCache(BaseCache):
def delete_cache(self, key):
self._remove_key(key)
async def async_get_ttl(self, key: str) -> Optional[int]:
async def async_get_ttl(self, key: str) -> int | None:
"""
Get the remaining TTL of a key in in-memory cache
"""
return self.ttl_dict.get(key, None)
async def async_get_oldest_n_keys(self, n: int) -> List[str]:
async def async_get_oldest_n_keys(self, n: int) -> list[str]:
"""
Get the oldest n keys in the cache
"""

45
litellm/caching/llm_caching_handler.py
View file

@ -4,21 +4,44 @@ Add the event loop to the cache key, to prevent event loop closed errors.
import asyncio
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
from .in_memory_cache import InMemoryCache
class LLMClientCache(InMemoryCache):
"""Cache for LLM HTTP clients (OpenAI, Azure, httpx, etc.).
IMPORTANT: This cache intentionally does NOT close clients on eviction.
Evicted clients may still be in use by in-flight requests. Closing them
eagerly causes ``RuntimeError: Cannot send a request, as the client has
been closed.`` errors in production after the TTL (1 hour) expires.
An evicted client is never closed on the spot: a request handed the client
just before eviction is still using it, and closing it there raises
``RuntimeError: Cannot send a request, as the client has been closed.``
Clients that are no longer referenced will be garbage-collected normally.
For explicit shutdown cleanup, use ``close_litellm_async_clients()``.
Nor can eviction be left to rely on garbage collection. The SDK clients are
reference cycles, so an evicted client and its open TCP connections survive
until a generational collection runs. Instead a client litellm created is
handed to ``EvictedClientCloser``, which closes it once a grace window has
passed. Clients the caller supplied are left untouched.
"""
def __init__(
self,
max_size_in_memory: int | None = 200,
default_ttl: int | None = 600,
max_size_per_item: int | None = 1024,
evicted_client_closer: EvictedClientCloser | None = None,
):
super().__init__(
max_size_in_memory=max_size_in_memory,
default_ttl=default_ttl,
max_size_per_item=max_size_per_item,
)
self.evicted_client_closer = evicted_client_closer or default_evicted_client_closer
def _remove_key(self, key: str) -> None:
evicted: object = self.cache_dict.get(key)
super()._remove_key(key)
self.evicted_client_closer.schedule(evicted)
self.evicted_client_closer.reap()
def update_cache_key_with_event_loop(self, key):
"""
Add the event loop to the cache key, to prevent event loop closed errors.
@ -31,16 +54,22 @@ class LLMClientCache(InMemoryCache):
except RuntimeError: # handle no current running event loop
return key
def set_cache(self, key, value, **kwargs):
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
"""``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return super().set_cache(key, value, **kwargs)
async def async_set_cache(self, key, value, **kwargs):
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return await super().async_set_cache(key, value, **kwargs)
def get_cache(self, key, **kwargs):
key = self.update_cache_key_with_event_loop(key)
self.evicted_client_closer.reap()
return super().get_cache(key, **kwargs)

14
litellm/caching/qdrant_semantic_cache.py
View file

@ -12,7 +12,7 @@ import ast
import asyncio
import json
import os
from typing import Any, Dict, cast
from typing import Any, cast
import litellm
from litellm._logging import print_verbose
@ -104,7 +104,7 @@ class QdrantSemanticCache(BaseCache):
print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}")
self._ensure_cache_key_payload_index()
else:
quantization_params: Dict[str, Any]
quantization_params: dict[str, Any]
if quantization_config is None or quantization_config == "binary":
quantization_params = {
"binary": {
@ -178,7 +178,7 @@ class QdrantSemanticCache(BaseCache):
if response.status_code not in (200, 201):
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {response.text}")
except Exception as exc:
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {str(exc)}")
print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {exc}")
def _payload_matches_cache_key(self, payload: dict, key: str) -> bool:
# Pre-isolation points stored only prompt + response with no cache-key
@ -188,7 +188,7 @@ class QdrantSemanticCache(BaseCache):
cached_key = payload.get(self.CACHE_KEY_FIELD_NAME)
return cached_key is not None and str(cached_key) == str(key)
def _get_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse:
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
"""Embed via the proxy Router when it serves the model, else direct."""
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
@ -210,7 +210,7 @@ class QdrantSemanticCache(BaseCache):
cache={"no-store": True, "no-cache": True},
)
async def _get_async_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse:
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
except ImportError:
@ -270,7 +270,6 @@ class QdrantSemanticCache(BaseCache):
headers=self.headers,
json=data,
)
return
def get_cache(self, key, **kwargs):
print_verbose(f"sync qdrant semantic-cache get_cache, kwargs: {kwargs}")
@ -344,7 +343,6 @@ class QdrantSemanticCache(BaseCache):
else:
# cache miss !
return None
pass
async def async_set_cache(self, key, value, **kwargs):
from litellm._uuid import uuid
@ -381,7 +379,6 @@ class QdrantSemanticCache(BaseCache):
headers=self.headers,
json=data,
)
return
async def async_get_cache(self, key, **kwargs):
print_verbose(f"async qdrant semantic-cache get_cache, kwargs: {kwargs}")
@ -452,7 +449,6 @@ class QdrantSemanticCache(BaseCache):
else:
# cache miss !
return None
pass
async def _collection_info(self):
return self.collection_info

141
litellm/caching/redis_cache.py
View file

@ -16,9 +16,9 @@ import inspect
import json
import time
from collections.abc import Awaitable, Callable, Sequence
from datetime import timedelta
from contextvars import ContextVar
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union, cast
from datetime import timedelta
from typing import TYPE_CHECKING, Any, TypeVar, Union, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -127,7 +127,7 @@ class RedisCircuitBreaker:
self.recovery_timeout = recovery_timeout
self.enabled = enabled
self._failure_count = 0
self._opened_at: Optional[float] = None
self._opened_at: float | None = None
self._state = self.CLOSED
def is_open(self) -> bool:
@ -272,10 +272,10 @@ class RedisCache(BaseCache):
host=None,
port=None,
password=None,
redis_flush_size: Optional[int] = 100,
namespace: Optional[str] = None,
startup_nodes: Optional[List] = None, # for redis-cluster
socket_timeout: Optional[float] = 5.0, # default 5 second timeout
redis_flush_size: int | None = 100,
namespace: str | None = None,
startup_nodes: list | None = None, # for redis-cluster
socket_timeout: float | None = 5.0, # default 5 second timeout
**kwargs,
):
from litellm._service_logger import ServiceLogging
@ -304,7 +304,7 @@ class RedisCache(BaseCache):
redis_kwargs.update(kwargs)
self.redis_client = get_redis_client(**redis_kwargs)
self.redis_async_client: Optional[Union[async_redis_client, async_redis_cluster_client]] = None
self.redis_async_client: async_redis_client | async_redis_cluster_client | None = None
self.redis_kwargs = redis_kwargs
self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs)
@ -346,7 +346,8 @@ class RedisCache(BaseCache):
verbose_logger.debug("Ignoring async redis ping. No running event loop.")
else:
verbose_logger.error(
"Error connecting to Async Redis client - {}".format(str(e)),
"Error connecting to Async Redis client - %s",
e,
extra={"error": str(e)},
)
self._handle_async_ping_error(e)
@ -407,7 +408,7 @@ class RedisCache(BaseCache):
def init_async_client(
self,
) -> Union[async_redis_client, async_redis_cluster_client]:
) -> async_redis_client | async_redis_cluster_client:
from litellm import in_memory_llm_clients_cache
from .._redis import get_redis_async_client, get_redis_connection_pool
@ -415,7 +416,7 @@ class RedisCache(BaseCache):
cache_key = self._get_async_client_cache_key()
cached_client = in_memory_llm_clients_cache.get_cache(key=cache_key)
if cached_client is not None:
redis_async_client = cast(Union[async_redis_client, async_redis_cluster_client], cached_client)
redis_async_client = cast(async_redis_client | async_redis_cluster_client, cached_client)
else:
# Create new connection pool and client for current event loop
self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs)
@ -483,9 +484,9 @@ class RedisCache(BaseCache):
)
except Exception as e:
# NON blocking - notify users Redis is throwing an exception
print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {str(e)}")
print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}")
def increment_cache(self, key, value: int, ttl: Optional[float] = None, **kwargs) -> int:
def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int:
_redis_client = self.redis_client
start_time = time.time()
set_ttl = self.get_ttl(ttl=ttl)
@ -626,7 +627,7 @@ class RedisCache(BaseCache):
async def run_script(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any:
async def execute() -> object:
executor: Optional[Callable[..., Awaitable[Any]]] = litellm.in_memory_llm_clients_cache.get_cache(
executor: Callable[..., Awaitable[Any]] | None = litellm.in_memory_llm_clients_cache.get_cache(
key=script_cache_key
)
if executor is None:
@ -755,10 +756,10 @@ class RedisCache(BaseCache):
async def _pipeline_helper(
self,
pipe: Union[pipeline, cluster_pipeline],
cache_list: List[Tuple[Any, Any]],
ttl: Optional[float],
) -> List:
pipe: pipeline | cluster_pipeline,
cache_list: list[tuple[Any, Any]],
ttl: float | None,
) -> list:
"""
Helper function for executing a pipeline of set operations on Redis
"""
@ -769,7 +770,7 @@ class RedisCache(BaseCache):
print_verbose(f"Set ASYNC Redis Cache PIPELINE: key: {cache_key}\nValue {cache_value}\nttl={ttl}")
json_cache_value = json.dumps(cache_value)
# Set the value with a TTL if it's provided.
_td: Optional[timedelta] = None
_td: timedelta | None = None
if ttl is not None:
_td = timedelta(seconds=ttl)
pipe.set( # type: ignore
@ -782,7 +783,7 @@ class RedisCache(BaseCache):
return results
@_redis_circuit_breaker_guard
async def async_set_cache_pipeline(self, cache_list: List[Tuple[Any, Any]], ttl: Optional[float] = None, **kwargs):
async def async_set_cache_pipeline(self, cache_list: list[tuple[Any, Any]], ttl: float | None = None, **kwargs):
"""
Use Redis Pipelines for bulk write operations
"""
@ -814,7 +815,7 @@ class RedisCache(BaseCache):
parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs),
)
)
return None
return
except Exception as e:
## LOGGING ##
end_time = time.time()
@ -842,8 +843,8 @@ class RedisCache(BaseCache):
self,
redis_client: async_redis_client,
key: str,
value: List,
ttl: Optional[float],
value: list,
ttl: float | None,
) -> None:
"""Helper function for async_set_cache_sadd. Separated for testing."""
ttl = self.get_ttl(ttl=ttl)
@ -856,7 +857,7 @@ class RedisCache(BaseCache):
raise
@_redis_circuit_breaker_guard
async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float], **kwargs):
async def async_set_cache_sadd(self, key, value: list, ttl: float | None, **kwargs):
from redis.asyncio import Redis
start_time = time.time()
@ -938,8 +939,8 @@ class RedisCache(BaseCache):
self,
key,
value: float,
ttl: Optional[int] = None,
parent_otel_span: Optional[Span] = None,
ttl: int | None = None,
parent_otel_span: Span | None = None,
refresh_ttl: bool = False,
) -> float:
from redis.asyncio import Redis
@ -1051,7 +1052,7 @@ class RedisCache(BaseCache):
cached_response = ast.literal_eval(cached_response)
return cached_response
def get_cache(self, key, parent_otel_span: Optional[Span] = None, **kwargs):
def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
try:
key = self.check_and_fix_namespace(key=key)
print_verbose(f"Get Redis Cache: key: {key}")
@ -1073,7 +1074,7 @@ class RedisCache(BaseCache):
# NON blocking - notify users Redis is throwing an exception
verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e)
def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
Wrapper to call `mget` on the redis client
@ -1081,7 +1082,7 @@ class RedisCache(BaseCache):
"""
return self.redis_client.mget(keys=keys) # type: ignore
async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]:
async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]:
"""
Wrapper to call `mget` on the redis client
@ -1092,8 +1093,8 @@ class RedisCache(BaseCache):
def batch_get_cache(
self,
key_list: Union[List[str], List[Optional[str]]],
parent_otel_span: Optional[Span] = None,
key_list: list[str] | list[str | None],
parent_otel_span: Span | None = None,
) -> dict:
"""
Use Redis for bulk read operations
@ -1114,7 +1115,7 @@ class RedisCache(BaseCache):
cache_key = self.check_and_fix_namespace(key=cache_key or "")
_keys.append(cache_key)
start_time = time.time()
results: List = self._run_redis_mget_operation(keys=_keys)
results: list = self._run_redis_mget_operation(keys=_keys)
end_time = time.time()
_duration = end_time - start_time
self.service_logger_obj.service_success_hook(
@ -1139,11 +1140,11 @@ class RedisCache(BaseCache):
return decoded_results
except Exception as e:
verbose_logger.error(f"Error occurred in batch get cache - {str(e)}")
verbose_logger.error("Error occurred in batch get cache - %s", e)
return key_value_dict
@_redis_circuit_breaker_guard
async def async_get_cache(self, key, parent_otel_span: Optional[Span] = None, **kwargs):
async def async_get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
from redis.asyncio import Redis
_redis_client: Redis = self.init_async_client() # type: ignore
@ -1185,14 +1186,14 @@ class RedisCache(BaseCache):
event_metadata={"key": key},
)
)
print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {str(e)}")
print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}")
_record_swallowed_redis_failure(self._circuit_breaker, e)
@_redis_circuit_breaker_guard
async def async_batch_get_cache(
self,
key_list: Union[List[str], List[Optional[str]]],
parent_otel_span: Optional[Span] = None,
key_list: list[str] | list[str | None],
parent_otel_span: Span | None = None,
) -> dict:
"""
Use Redis for bulk read operations
@ -1257,7 +1258,7 @@ class RedisCache(BaseCache):
parent_otel_span=parent_otel_span,
)
)
verbose_logger.error(f"Error occurred in async batch get cache - {str(e)}")
verbose_logger.error("Error occurred in async batch get cache - %s", e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return key_value_dict
@ -1292,7 +1293,7 @@ class RedisCache(BaseCache):
error=e,
call_type=f"sync_ping <- {_get_call_stack_info()}",
)
verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}")
verbose_logger.error("LiteLLM Redis Cache PING: - Got exception from REDIS : %s", e)
raise e
async def ping(self) -> bool:
@ -1326,7 +1327,7 @@ class RedisCache(BaseCache):
call_type=f"async_ping <- {_get_call_stack_info()}",
)
)
verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}")
verbose_logger.error("LiteLLM Redis Cache PING: - Got exception from REDIS : %s", e)
raise e
@_redis_circuit_breaker_guard
@ -1337,8 +1338,8 @@ class RedisCache(BaseCache):
# keys is a list, unpack it so it gets passed as individual elements to delete
await _redis_client.delete(*keys)
def client_list(self) -> List:
client_list: List = self.redis_client.client_list() # type: ignore
def client_list(self) -> list:
client_list: list = self.redis_client.client_list() # type: ignore
return client_list
def info(self):
@ -1388,10 +1389,10 @@ class RedisCache(BaseCache):
else:
return {"status": "failed", "message": "Redis ping returned False"}
except Exception as e:
verbose_logger.error(f"Redis connection test failed: {str(e)}")
verbose_logger.error("Redis connection test failed: %s", e)
return {
"status": "failed",
"message": f"Redis connection failed: {str(e)}",
"message": f"Redis connection failed: {e}",
"error": str(e),
}
@ -1410,8 +1411,8 @@ class RedisCache(BaseCache):
async def _pipeline_increment_helper(
self,
pipe: pipeline,
increment_list: List[RedisPipelineIncrementOperation],
) -> Optional[List[float]]:
increment_list: list[RedisPipelineIncrementOperation],
) -> list[float] | None:
"""Helper function for pipeline increment operations"""
# Iterate through each increment operation and add commands to pipeline
for increment_op in increment_list:
@ -1426,13 +1427,13 @@ class RedisCache(BaseCache):
# Execute the pipeline and return results
results = await pipe.execute()
# only return float values
verbose_logger.debug(f"Increment ASYNC Redis Cache PIPELINE: results: {results}")
verbose_logger.debug("Increment ASYNC Redis Cache PIPELINE: results: %s", results)
return [r for r in results if isinstance(r, float)]
@_redis_circuit_breaker_guard
async def async_increment_pipeline(
self, increment_list: List[RedisPipelineIncrementOperation], **kwargs
) -> Optional[List[float]]:
self, increment_list: list[RedisPipelineIncrementOperation], **kwargs
) -> list[float] | None:
"""
Use Redis Pipelines for bulk increment operations
Args:
@ -1492,7 +1493,7 @@ class RedisCache(BaseCache):
raise e
@_redis_circuit_breaker_guard
async def async_get_ttl(self, key: str) -> Optional[int]:
async def async_get_ttl(self, key: str) -> int | None:
"""
Get the remaining TTL of a key in Redis
@ -1513,7 +1514,7 @@ class RedisCache(BaseCache):
return None
return ttl
except Exception as e:
verbose_logger.debug(f"Redis TTL Error: {e}")
verbose_logger.debug("Redis TTL Error: %s", e)
_record_swallowed_redis_failure(self._circuit_breaker, e)
return None
@ -1521,8 +1522,8 @@ class RedisCache(BaseCache):
async def async_rpush(
self,
key: str,
values: List[Any],
parent_otel_span: Optional[Span] = None,
values: list[Any],
parent_otel_span: Span | None = None,
**kwargs,
) -> int:
"""
@ -1565,14 +1566,14 @@ class RedisCache(BaseCache):
call_type=f"async_rpush <- {_get_call_stack_info()}",
)
)
verbose_logger.error(f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {str(e)}")
verbose_logger.error("LiteLLM Redis Cache RPUSH: - Got exception from REDIS : %s", e)
raise e
async def _pipeline_rpush_helper(
self,
pipe: pipeline,
rpush_list: List[RedisPipelineRpushOperation],
) -> List[int]:
rpush_list: list[RedisPipelineRpushOperation],
) -> list[int]:
"""Helper function for pipeline rpush operations"""
for rpush_op in rpush_list:
key = self.check_and_fix_namespace(key=rpush_op["key"])
@ -1587,8 +1588,8 @@ class RedisCache(BaseCache):
@_redis_circuit_breaker_guard
async def async_rpush_pipeline(
self,
rpush_list: List[RedisPipelineRpushOperation],
) -> List[int]:
rpush_list: list[RedisPipelineRpushOperation],
) -> list[int]:
"""
Use Redis Pipelines for bulk RPUSH operations
@ -1639,8 +1640,8 @@ class RedisCache(BaseCache):
)
raise e
async def handle_lpop_count_for_older_redis_versions(self, pipe: pipeline, key: str, count: int) -> List[bytes]:
result: List[bytes] = []
async def handle_lpop_count_for_older_redis_versions(self, pipe: pipeline, key: str, count: int) -> list[bytes]:
result: list[bytes] = []
for _ in range(count):
pipe.lpop(key)
results = await pipe.execute()
@ -1656,10 +1657,10 @@ class RedisCache(BaseCache):
async def async_lpop(
self,
key: str,
count: Optional[int] = None,
parent_otel_span: Optional[Span] = None,
count: int | None = None,
parent_otel_span: Span | None = None,
**kwargs,
) -> Union[Any, List[Any]]:
) -> Any | list[Any]:
_redis_client: Any = self.init_async_client()
key = self.check_and_fix_namespace(key=key)
start_time = time.time()
@ -1711,14 +1712,14 @@ class RedisCache(BaseCache):
call_type=f"async_lpop <- {_get_call_stack_info()}",
)
)
verbose_logger.error(f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}")
verbose_logger.error("LiteLLM Redis Cache LPOP: - Got exception from REDIS : %s", e)
raise e
async def _pipeline_lpop_helper(
self,
pipe: pipeline,
lpop_list: List[RedisPipelineLpopOperation],
) -> List[Optional[List[str]]]:
lpop_list: list[RedisPipelineLpopOperation],
) -> list[list[str] | None]:
"""Helper function for pipeline lpop operations.
For Redis >= 7, queues one LPOP(key, count) per operation.
@ -1734,7 +1735,7 @@ class RedisCache(BaseCache):
else:
# For Redis < 7, LPOP doesn't support count param.
# Issue `count` individual LPOP commands per key, all in one pipeline.
counts: List[int] = []
counts: list[int] = []
for lpop_op in lpop_list:
key = self.check_and_fix_namespace(key=lpop_op["key"])
count = lpop_op["count"] or 1
@ -1757,7 +1758,7 @@ class RedisCache(BaseCache):
raise r
# Decode bytes -> str for each result set
decoded_results: List[Optional[List[str]]] = []
decoded_results: list[list[str] | None] = []
for r in raw_results:
if r is None:
decoded_results.append(None)
@ -1776,8 +1777,8 @@ class RedisCache(BaseCache):
@_redis_circuit_breaker_guard
async def async_lpop_pipeline(
self,
lpop_list: List[RedisPipelineLpopOperation],
) -> List[Optional[List[str]]]:
lpop_list: list[RedisPipelineLpopOperation],
) -> list[list[str] | None]:
"""
Use Redis Pipelines for bulk LPOP operations

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