scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge branch 'litellm_internal_staging' into fix-managed-files-null-object

Browse source
This commit is contained in:
mateo-berri 2026-08-05 03:05:32 -07:00
commit f270b53144
2881 changed files with 118539 additions and 78518 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

176
.circleci/config.yml
View file

@ -88,6 +88,59 @@ commands:
rm -f /tmp/uv-install.sh
echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV"
export PATH="$HOME/.local/bin:$PATH"
install_node:
description: "Install the Node.js version pinned in ui/litellm-dashboard/.nvmrc (24.19.0, which bundles npm 11.17.0) with checksum verification, and prepend it to PATH. Run this on any executor whose image does not already ship that version, or `npm ci` in ui/litellm-dashboard fails EBADENGINE against the engines floor. Installs into /opt/node rather than over /usr/local on purpose: cimg/python:*-browsers ships its own node there, and unpacking the tarball on top of it leaves npm 11.17 files merged with the image's npm 11.9 tree, which reports the new version and then exits 1 on `npm ci` with no error text at all. Requires checkout, which the .nvmrc drift check reads."
steps:
- run:
name: Install Node.js 24.19.0
command: |
NODE_VERSION="24.19.0"
NODE_TARBALL="node-v${NODE_VERSION}-linux-x64.tar.xz"
NODE_EXPECTED_SHA="14b342e71204f811bde6153be8e04b62aef63c236fef92b55f9c83154b409647"
NVMRC_VERSION="$(tr -d '[:space:]' < ui/litellm-dashboard/.nvmrc)"
if [ "$NVMRC_VERSION" != "$NODE_VERSION" ]; then
echo "install_node: ui/litellm-dashboard/.nvmrc pins ${NVMRC_VERSION} but this command pins ${NODE_VERSION}; update NODE_VERSION and NODE_EXPECTED_SHA together" >&2
exit 1
fi
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 mkdir -p /opt/node
sudo tar -xJf "/tmp/${NODE_TARBALL}" -C /opt/node --strip-components=1
rm -f "/tmp/${NODE_TARBALL}"
echo 'export PATH="/opt/node/bin:$PATH"' >> "$BASH_ENV"
export PATH="/opt/node/bin:$PATH"
node --version
npm --version
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 +216,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 +251,7 @@ commands:
- checkout
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -281,6 +355,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 +399,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -371,6 +473,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -445,6 +548,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -496,6 +600,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -562,6 +667,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -602,6 +708,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -643,6 +750,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -676,6 +784,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -726,6 +835,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -777,6 +887,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -810,6 +921,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -856,6 +968,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -902,6 +1015,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -944,6 +1058,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -990,6 +1105,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1037,6 +1153,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -1077,6 +1194,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1122,6 +1240,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1166,6 +1285,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1198,6 +1318,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1241,6 +1362,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1285,6 +1407,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1329,6 +1452,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1360,6 +1484,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1406,6 +1531,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1451,6 +1577,7 @@ jobs:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1501,6 +1628,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1525,6 +1653,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1551,6 +1680,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1652,6 +1782,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1747,6 +1878,7 @@ jobs:
at: ~/project
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1835,6 +1967,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -1918,6 +2051,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2050,6 +2184,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2136,6 +2271,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2232,12 +2368,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 +2395,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 +2447,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 +2529,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2473,20 +2615,9 @@ 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
command: |
NODE_VERSION="18.20.8"
NODE_TARBALL="node-v${NODE_VERSION}-linux-x64.tar.xz"
NODE_EXPECTED_SHA="5467ee62d6af1411d46b6a10e3fb5cacc92734dbcef465fea14e7b90993001c9"
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
rm -f "/tmp/${NODE_TARBALL}"
node --version
npm --version
- install_node
- run:
name: Install Node.js test dependencies
@ -2527,6 +2658,7 @@ jobs:
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
@ -2614,7 +2746,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 +2790,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}
@ -2716,7 +2848,9 @@ jobs:
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- install_node
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -2731,7 +2865,7 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
# The cimg/python:3.12-browsers image already ships the Chromium system
@ -2746,7 +2880,7 @@ jobs:
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
@ -2858,7 +2992,9 @@ jobs:
- skip_if_unrelated_changes:
category: client
- setup_google_dns
- install_node
- install_uv
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
@ -2873,7 +3009,7 @@ jobs:
- ~/.cache/uv
- restore_cache:
keys:
- ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install Node dependencies and Playwright
command: |
@ -2883,7 +3019,7 @@ jobs:
npm ci
npx playwright install chromium
- save_cache:
key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ui/litellm-dashboard/node_modules
- tests/e2e/ui/node_modules
@ -3031,6 +3167,8 @@ workflows:
only:
- main
- /litellm_.*/
- base_sdk_install:
filters: *main_branches
- local_testing_part1:
filters: *main_branches
- local_testing_part2:

46
.flake8
View file

@ -1,46 +0,0 @@
[flake8]
ignore =
# The following ignores can be removed when formatting using black
W191,W291,W292,W293,W391,W504
E101,E111,E114,E116,E117,E121,E122,E123,E124,E125,E126,E127,E128,E129,E131,
E201,E202,E221,E222,E225,E226,E231,E241,E251,E252,E261,E265,E271,E272,E275,
E301,E302,E303,E305,E306,
# line break before binary operator
W503,
# inline comment should start with '# '
E262,
# too many leading '#' for block comment
E266,
# multiple imports on one line
E401,
# module level import not at top of file
E402,
# Line too long (82 > 79 characters)
E501,
# comparison to None should be 'if cond is None:'
E711,
# comparison to True should be 'if cond is True:' or 'if cond:'
E712,
# do not compare types, for exact checks use `is` / `is not`, for instance checks use `isinstance()`
E721,
# do not use bare 'except'
E722,
# x is imported but unused
F401,
# 'from . import *' used; unable to detect undefined names
F403,
# x may be undefined, or defined from star imports:
F405,
# f-string is missing placeholders
F541,
# dictionary key '' repeated with different values
F601,
# redefinition of unused x from line 123
F811,
# undefined name x
F821,
# local variable x is assigned to but never used
F841,
# https://black.readthedocs.io/en/stable/guides/using_black_with_other_tools.html#flake8
extend-ignore = E203

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:

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

@ -104,9 +104,8 @@ jobs:
- name: Check basedpyright budget (delta vs base)
env:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
NODE_OPTIONS: --max-old-space-size=12288
run: |
(uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA"
uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA"
- name: Check tests/e2e basedpyright (zero errors)
env:

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:

9
CLAUDE.md
View file

@ -29,7 +29,7 @@ Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We pref
If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ";", ".", etc.
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
@ -39,15 +39,11 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
On a fresh worktree or clone, run `make bootstrap` before anything else. It provisions everything tests, `make pre-commit`, and a local proxy need
Run tests before you commit. Also, run `make pre-commit` right before each commit, which generates types (as needed) and formats/lints your code. Any errors found must be fixed. It only runs when there are staged frontend and/or backend changes and calculates violations, generates types, etc. based on the worktree, so stage what you need or stash/delete unwanted files in litellm/ or ui/ (where backend and frontend lint run, respectively) before running it. If it fails because dashboard api types are stale, it already regenerated them for you. You just need to stage the schema.d.ts, re-run `make pre-commit` to confirm it passes, and commit
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
@ -75,6 +71,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>` explaining why
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match

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

10
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 && ../../scripts/with_dashboard_node.sh 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"; \
@ -176,10 +176,8 @@ lint-ruff-FULL-dev: install-dev
if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
else echo "No changed .py files to check."; fi
lint-basedpyright lint-basedpyright-budget-update: export NODE_OPTIONS := --max-old-space-size=12288
lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
$(UV_RUN) basedpyright tests/e2e
@ -192,7 +190,7 @@ lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
# --update lowers each limit by what this branch fixed since its branch point, so
# it needs the base ref fetched to resolve the merge-base.
lint-basedpyright-budget-update: install-dev lint-fetch-base
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --update
$(UV_RUN) python scripts/type_check_gate.py --update
lint-format: format-check
@ -239,7 +237,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"]

1
backend/routes/allowlist.py
View file

@ -44,6 +44,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
"/router/",
"/router_settings",
"/adaptive_router/",
"/auto_router/",
"/fallback",
"/fallbacks",
"/cache_settings",

50
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 31903
"limit": 29809
},
"reportArgumentType": {
"limit": 2645
@ -15,25 +15,25 @@
"limit": 123
},
"reportConstantRedefinition": {
"limit": 59
"limit": 40
},
"reportDeprecated": {
"limit": 325
},
"reportDuplicateImport": {
"limit": 42
"limit": 24
},
"reportExplicitAny": {
"limit": 10214
"limit": 9473
},
"reportFunctionMemberAccess": {
"limit": 11
},
"reportGeneralTypeIssues": {
"limit": 227
"limit": 157
},
"reportIncompatibleMethodOverride": {
"limit": 78
"limit": 77
},
"reportIncompatibleVariableOverride": {
"limit": 12
@ -42,7 +42,7 @@
"limit": 18
},
"reportIndexIssue": {
"limit": 37
"limit": 35
},
"reportInvalidTypeForm": {
"limit": 35
@ -54,13 +54,13 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5869
"limit": 5855
},
"reportMissingTypeArgument": {
"limit": 15861
"limit": 15849
},
"reportMissingTypeStubs": {
"limit": 41
"limit": 40
},
"reportOperatorIssue": {
"limit": 0
@ -81,13 +81,13 @@
"limit": 0
},
"reportPossiblyUnboundVariable": {
"limit": 77
"limit": 56
},
"reportPrivateUsage": {
"limit": 2437
"limit": 2436
},
"reportRedeclaration": {
"limit": 12
"limit": 8
},
"reportReturnType": {
"limit": 219
@ -99,31 +99,31 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 45366
"limit": 45262
},
"reportUnknownLambdaType": {
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40477
"limit": 40452
},
"reportUnknownParameterType": {
"limit": 20338
"limit": 20309
},
"reportUnknownVariableType": {
"limit": 32047
"limit": 31978
},
"reportUnnecessaryCast": {
"limit": 177
"limit": 124
},
"reportUnnecessaryComparison": {
"limit": 1021
"limit": 703
},
"reportUnnecessaryContains": {
"limit": 7
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 1205
"limit": 866
},
"reportUntypedBaseClass": {
"limit": 165
@ -132,15 +132,15 @@
"limit": 33
},
"reportUnusedClass": {
"limit": 33
"limit": 23
},
"reportUnusedFunction": {
"limit": 204
"limit": 139
},
"reportUnusedImport": {
"limit": 1003
"limit": 588
},
"reportUnusedVariable": {
"limit": 1297
"limit": 147
}
}

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)
@ -383,7 +383,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
}
)
return [
OpenAIFileObject(**file_object.file_object)
OpenAIFileObject.model_validate(file_object.file_object)
for file_object in file_ids
if file_object.file_object is not None
]

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/README.md
View file

@ -39,7 +39,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
| `livenessProbe.*` | Liveness probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` |
| `readinessProbe.*` | Readiness probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` |
| `startupProbe.*` | Startup probe settings for the LiteLLM container (`path`, `periodSeconds`, `timeoutSeconds`, thresholds, and initial delay). | See `values.yaml` |
| `resources.*` | CPU/memory requests and limits for the LiteLLM container. | `{}` |
| `resources.*` | CPU/memory requests and limits for the LiteLLM container. Unset by default; production deployments should set 1 CPU and 4Gi of memory per worker. | `{}` |
| `service.loadBalancerClass` | Optional LoadBalancer implementation class (only used when `service.type` is `LoadBalancer`) | `""` |
| `ingress.labels` | Additional labels for the Ingress resource | `{}` |
| `ingress.*` | See [values.yaml](./values.yaml) for example settings | N/A |

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

27
helm/litellm-helm/values.yaml
View file

@ -181,16 +181,19 @@ proxy_config:
resources:
{}
# We usually recommend not to specify default resources and to leave this as a conscious
# choice for the user. This also increases chances charts run on environments with little
# resources, such as Minikube. If you do want to specify resources, uncomment the following
# lines, adjust them as necessary, and remove the curly braces after 'resources:'.
# limits:
# cpu: 100m
# memory: 128Mi
# Unset by default so the chart installs on small clusters such as Minikube, and so an
# upgrade never leaves a running pod Pending. Production deployments should set these.
# A proxy at DB-connected steady state needs about 1 CPU and 4Gi of memory per worker;
# sizing below that gets the pod OOMKilled once traffic and DB connections ramp up.
# Scale both figures with --num_workers, then uncomment the lines below and remove the
# curly braces after 'resources:'. See "Recommended Machine Specifications" in
# https://docs.litellm.ai/docs/proxy/prod.
# requests:
# cpu: 100m
# memory: 128Mi
# cpu: "1"
# memory: 4Gi
# limits:
# cpu: "1"
# memory: 4Gi
autoscaling:
enabled: false
@ -432,9 +435,9 @@ migrationJob:
annotations: {}
ttlSecondsAfterFinished: 120
resources: {}
# requests:
# cpu: 100m
# memory: 100Mi
# Unset by default. This job runs the database migration and exits, so it does not
# need the steady-state headroom the proxy does; size it from your own migration
# runs rather than from the proxy figures above.
extraContainers: []
extraInitContainers: []

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

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql Normal file
View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_Config" ADD COLUMN IF NOT EXISTS "last_run_at" TIMESTAMP(3),
ADD COLUMN IF NOT EXISTS "reload_revision" BIGINT NOT NULL DEFAULT 0;

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;

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

@ -601,6 +601,8 @@ model LiteLLM_TagTable {
model LiteLLM_Config {
param_name String @id
param_value Json?
last_run_at DateTime?
reload_revision BigInt @default(0)
}
// View spend, model, api_key per request
@ -748,6 +750,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 +785,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 +820,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 +854,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 +888,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 +924,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)

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

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.81"
version = "0.4.83"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.81"
version = "0.4.83"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

36
litellm/__init__.py
View file

@ -27,18 +27,19 @@ if os.getenv("LITELLM_MODE", "DEV") == "DEV":
_dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
from typing import (
Callable,
List,
Optional,
Dict,
Union,
Any,
Literal,
Callable,
Dict,
Final,
get_args,
TYPE_CHECKING,
Tuple,
List,
Literal,
Optional,
overload,
Tuple,
Type,
TYPE_CHECKING,
Union,
)
from litellm.types.integrations.datadog import DatadogInitParams
from litellm.types.integrations.newrelic import NewRelicInitParams
@ -264,6 +265,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 +451,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
@ -678,12 +682,12 @@ def is_bedrock_pricing_only_model(key: str) -> bool:
bool: True if the key matches the Bedrock pattern, False otherwise.
"""
# Regex to match 'bedrock/<region>/<model>'
bedrock_pattern = re.compile(r"^bedrock/[a-zA-Z0-9_-]+/.+$")
bedrock_pattern: Final = re.compile(r"^bedrock/[a-zA-Z0-9_-]+/.+$")
if "month-commitment" in key:
return True
is_match = bedrock_pattern.match(key)
is_match: Final = bedrock_pattern.match(key)
return is_match is not None
@ -701,7 +705,7 @@ def is_openai_finetune_model(key: str) -> bool:
def add_known_models(model_cost_map: Optional[Dict] = None):
_map = model_cost_map if model_cost_map is not None else model_cost
_map: Final = model_cost_map if model_cost_map is not None else model_cost
for key, value in _map.items():
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
open_ai_chat_completion_models.add(key)
@ -2137,11 +2141,11 @@ def __getattr__(name: str) -> Any:
# Use cached registry from _lazy_imports instead of importing tuples every time
from ._lazy_imports import _get_lazy_import_registry
registry = _get_lazy_import_registry()
registry: Final = _get_lazy_import_registry()
# Check if name is in registry and call the cached handler function
if name in registry:
handler_func = registry[name]
handler_func: Final = registry[name]
return handler_func(name)
# Lazy load encoding from main.py to avoid heavy tiktoken import
@ -2194,7 +2198,7 @@ def __getattr__(name: str) -> Any:
return _globals["openaiOSeriesConfig"]
# Lazy load other config instances
_config_instances = {
_config_instances: Final = {
"openAIGPTConfig": "OpenAIGPTConfig",
"openAIGPTAudioConfig": "OpenAIGPTAudioConfig",
"openAIGPT5Config": "OpenAIGPT5Config",
@ -2236,7 +2240,7 @@ def __getattr__(name: str) -> Any:
# Check if already cached
if "priority_reservation_settings" not in _globals:
# Import the class and instantiate it
PriorityReservationSettings = __getattr__("PriorityReservationSettings")
PriorityReservationSettings: Final = __getattr__("PriorityReservationSettings")
_globals["priority_reservation_settings"] = PriorityReservationSettings()
return _globals["priority_reservation_settings"]
@ -2248,7 +2252,7 @@ def __getattr__(name: str) -> Any:
# Check if already cached
if "logging_callback_manager" not in _globals:
# Import the class and instantiate it
LoggingCallbackManager = __getattr__("LoggingCallbackManager")
LoggingCallbackManager: Final = __getattr__("LoggingCallbackManager")
_globals["logging_callback_manager"] = LoggingCallbackManager()
return _globals["logging_callback_manager"]

3
litellm/_internal_context.py
View file

@ -7,7 +7,8 @@ asyncio task and cannot be injected via HTTP request bodies.
"""
from contextvars import ContextVar
from typing import Final
# When True, suppresses async logging and billing for internal sub-calls
# (e.g., emulated file-search steps that make nested LLM calls).
is_internal_call: ContextVar[bool] = ContextVar("is_internal_call", default=False)
is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False)

79
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, Final, 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]]:
@ -232,7 +233,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
raise AttributeError(f"{category} lazy import: unknown attribute {name!r}")
# Step 2: Get the cache (where we store imported things)
_globals = _get_litellm_globals()
_globals: Final = _get_litellm_globals()
# Step 3: If we've already imported it, just return the cached version
if name in _globals:
@ -254,7 +255,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Step 6: Get the actual attribute from the module
# Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class
value = getattr(module, attr_name)
value: Final = getattr(module, attr_name)
# Step 7: Cache it so we don't have to import again next time
_globals[name] = value
@ -338,7 +339,7 @@ def _lazy_import_utils_module(name: str) -> Any:
raise AttributeError(f"Utils module lazy import: unknown attribute {name!r}")
# Get the cache (where we store imported things) - use utils globals
_globals = _get_utils_globals()
_globals: Final = _get_utils_globals()
# If we've already imported it, just return the cached version
if name in _globals:
@ -354,7 +355,7 @@ def _lazy_import_utils_module(name: str) -> Any:
module = importlib.import_module(module_path)
# Get the actual attribute from the module
value = getattr(module, attr_name)
value: Final = getattr(module, attr_name)
# Cache it so we don't have to import again next time
_globals[name] = value
@ -378,15 +379,15 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
- "in_memory_llm_clients_cache" is a singleton instance of that class
So we need custom logic to handle both cases.
"""
_globals = _get_litellm_globals()
_globals: Final = _get_litellm_globals()
# If already cached, return it
if name in _globals:
return _globals[name]
# Import the class
module = importlib.import_module("litellm.caching.llm_caching_handler")
LLMClientCache = getattr(module, "LLMClientCache")
module: Final = importlib.import_module("litellm.caching.llm_caching_handler")
LLMClientCache: Final = getattr(module, "LLMClientCache")
# If they want the class itself, return it
if name == "LLMClientCache":
@ -395,7 +396,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
# If they want the singleton instance, create it (only once)
if name == "in_memory_llm_clients_cache":
instance = LLMClientCache()
instance: Final = LLMClientCache()
_globals["in_memory_llm_clients_cache"] = instance
return instance
@ -411,7 +412,7 @@ def _lazy_import_http_handlers(name: str) -> Any:
- They need configuration (timeout, etc.) from the module globals
- They use factory functions instead of direct instantiation
"""
_globals = _get_litellm_globals()
_globals: Final = _get_litellm_globals()
if name == "module_level_aclient":
# Create an async HTTP client using the factory function
@ -419,11 +420,11 @@ def _lazy_import_http_handlers(name: str) -> Any:
# Get timeout from module config (if set)
timeout = _globals.get("request_timeout")
params = {"timeout": timeout, "client_alias": "module level aclient"}
params: Final = {"timeout": timeout, "client_alias": "module level aclient"}
# Create the client instance
provider_id = cast(Any, "litellm_module_level_client")
async_client = get_async_httpx_client(
provider_id: Final = cast(Any, "litellm_module_level_client")
async_client: Final = get_async_httpx_client(
llm_provider=provider_id,
params=params,
)
@ -437,7 +438,7 @@ def _lazy_import_http_handlers(name: str) -> Any:
from litellm.llms.custom_httpx.http_handler import HTTPHandler
timeout = _globals.get("request_timeout")
sync_client = HTTPHandler(timeout=timeout)
sync_client: Final = HTTPHandler(timeout=timeout)
# Cache it
_globals["module_level_client"] = sync_client

54
litellm/_lazy_imports_registry.py
View file

@ -5,21 +5,23 @@ This module contains all the name tuples and import maps used by the lazy import
Separated from the handler functions for better organization.
"""
from typing import Final
# Cost calculator names that support lazy loading via _lazy_import_cost_calculator
COST_CALCULATOR_NAMES = (
COST_CALCULATOR_NAMES: Final = (
"completion_cost",
"cost_per_token",
"response_cost_calculator",
)
# Litellm logging names that support lazy loading via _lazy_import_litellm_logging
LITELLM_LOGGING_NAMES = (
LITELLM_LOGGING_NAMES: Final = (
"Logging",
"modify_integration",
)
# Utils names that support lazy loading via _lazy_import_utils
UTILS_NAMES = (
UTILS_NAMES: Final = (
"exception_type",
"get_optional_params",
"get_response_string",
@ -66,20 +68,20 @@ UTILS_NAMES = (
)
# Token counter names that support lazy loading via _lazy_import_token_counter
TOKEN_COUNTER_NAMES = ("get_modified_max_tokens",)
TOKEN_COUNTER_NAMES: Final = ("get_modified_max_tokens",)
# LLM client cache names that support lazy loading via _lazy_import_llm_client_cache
LLM_CLIENT_CACHE_NAMES = (
LLM_CLIENT_CACHE_NAMES: Final = (
"LLMClientCache",
"in_memory_llm_clients_cache",
)
# Bedrock type names that support lazy loading via _lazy_import_bedrock_types
BEDROCK_TYPES_NAMES = ("COHERE_EMBEDDING_INPUT_TYPES",)
BEDROCK_TYPES_NAMES: Final = ("COHERE_EMBEDDING_INPUT_TYPES",)
# Common types from litellm.types.utils that support lazy loading via
# _lazy_import_types_utils
TYPES_UTILS_NAMES = (
TYPES_UTILS_NAMES: Final = (
"ImageObject",
"BudgetConfig",
"all_litellm_params",
@ -92,7 +94,7 @@ TYPES_UTILS_NAMES = (
)
# Caching / cache classes that support lazy loading via _lazy_import_caching
CACHING_NAMES = (
CACHING_NAMES: Final = (
"Cache",
"DualCache",
"RedisCache",
@ -100,20 +102,20 @@ CACHING_NAMES = (
)
# HTTP handler names that support lazy loading via _lazy_import_http_handlers
HTTP_HANDLER_NAMES = (
HTTP_HANDLER_NAMES: Final = (
"module_level_aclient",
"module_level_client",
)
# Dotprompt integration names that support lazy loading via _lazy_import_dotprompt
DOTPROMPT_NAMES = (
DOTPROMPT_NAMES: Final = (
"global_prompt_manager",
"global_prompt_directory",
"set_global_prompt_directory",
)
# LLM config classes that support lazy loading via _lazy_import_llm_configs
LLM_CONFIG_NAMES = (
LLM_CONFIG_NAMES: Final = (
"AmazonConverseConfig",
"OpenAILikeChatConfig",
"GaladrielChatConfig",
@ -328,7 +330,7 @@ LLM_CONFIG_NAMES = (
)
# Types that support lazy loading via _lazy_import_types
TYPES_NAMES = (
TYPES_NAMES: Final = (
"GuardrailItem",
"DefaultTeamSSOParams",
"LiteLLM_UpperboundKeyGenerateParams",
@ -344,14 +346,14 @@ TYPES_NAMES = (
)
# LLM provider logic names that support lazy loading via _lazy_import_llm_provider_logic
LLM_PROVIDER_LOGIC_NAMES = (
LLM_PROVIDER_LOGIC_NAMES: Final = (
"get_llm_provider",
"remove_index_from_tool_calls",
)
# Utils module names that support lazy loading via _lazy_import_utils_module
# These are attributes accessed from litellm.utils module
UTILS_MODULE_NAMES = (
UTILS_MODULE_NAMES: Final = (
"encoding",
"BaseVectorStore",
"CredentialAccessor",
@ -423,7 +425,7 @@ UTILS_MODULE_NAMES = (
)
# Import maps for registry pattern - reduces repetition
_UTILS_IMPORT_MAP = {
_UTILS_IMPORT_MAP: Final = {
"exception_type": (".utils", "exception_type"),
"get_optional_params": (".utils", "get_optional_params"),
"get_response_string": (".utils", "get_response_string"),
@ -478,13 +480,13 @@ _UTILS_IMPORT_MAP = {
),
}
_COST_CALCULATOR_IMPORT_MAP = {
_COST_CALCULATOR_IMPORT_MAP: Final = {
"completion_cost": (".cost_calculator", "completion_cost"),
"cost_per_token": (".cost_calculator", "cost_per_token"),
"response_cost_calculator": (".cost_calculator", "response_cost_calculator"),
}
_TYPES_UTILS_IMPORT_MAP = {
_TYPES_UTILS_IMPORT_MAP: Final = {
"ImageObject": (".types.utils", "ImageObject"),
"BudgetConfig": (".types.utils", "BudgetConfig"),
"all_litellm_params": (".types.utils", "all_litellm_params"),
@ -496,28 +498,28 @@ _TYPES_UTILS_IMPORT_MAP = {
"GenericStreamingChunk": (".types.utils", "GenericStreamingChunk"),
}
_TOKEN_COUNTER_IMPORT_MAP = {
_TOKEN_COUNTER_IMPORT_MAP: Final = {
"get_modified_max_tokens": (
"litellm.litellm_core_utils.token_counter",
"get_modified_max_tokens",
),
}
_BEDROCK_TYPES_IMPORT_MAP = {
_BEDROCK_TYPES_IMPORT_MAP: Final = {
"COHERE_EMBEDDING_INPUT_TYPES": (
"litellm.types.llms.bedrock",
"COHERE_EMBEDDING_INPUT_TYPES",
),
}
_CACHING_IMPORT_MAP = {
_CACHING_IMPORT_MAP: Final = {
"Cache": ("litellm.caching.caching", "Cache"),
"DualCache": ("litellm.caching.caching", "DualCache"),
"RedisCache": ("litellm.caching.caching", "RedisCache"),
"InMemoryCache": ("litellm.caching.caching", "InMemoryCache"),
}
_LITELLM_LOGGING_IMPORT_MAP = {
_LITELLM_LOGGING_IMPORT_MAP: Final = {
"Logging": ("litellm.litellm_core_utils.litellm_logging", "Logging"),
"modify_integration": (
"litellm.litellm_core_utils.litellm_logging",
@ -525,7 +527,7 @@ _LITELLM_LOGGING_IMPORT_MAP = {
),
}
_DOTPROMPT_IMPORT_MAP = {
_DOTPROMPT_IMPORT_MAP: Final = {
"global_prompt_manager": (
"litellm.integrations.dotprompt",
"global_prompt_manager",
@ -540,7 +542,7 @@ _DOTPROMPT_IMPORT_MAP = {
),
}
_TYPES_IMPORT_MAP = {
_TYPES_IMPORT_MAP: Final = {
"GuardrailItem": ("litellm.types.guardrails", "GuardrailItem"),
"DefaultTeamSSOParams": (
"litellm.types.proxy.management_endpoints.ui_sso",
@ -569,7 +571,7 @@ _TYPES_IMPORT_MAP = {
),
}
_LLM_PROVIDER_LOGIC_IMPORT_MAP = {
_LLM_PROVIDER_LOGIC_IMPORT_MAP: Final = {
"get_llm_provider": (
"litellm.litellm_core_utils.get_llm_provider_logic",
"get_llm_provider",
@ -580,7 +582,7 @@ _LLM_PROVIDER_LOGIC_IMPORT_MAP = {
),
}
_LLM_CONFIGS_IMPORT_MAP = {
_LLM_CONFIGS_IMPORT_MAP: Final = {
"AmazonConverseConfig": (
".llms.bedrock.chat.converse_transformation",
"AmazonConverseConfig",
@ -1215,7 +1217,7 @@ _LLM_CONFIGS_IMPORT_MAP = {
}
# Import map for utils module lazy imports
_UTILS_MODULE_IMPORT_MAP = {
_UTILS_MODULE_IMPORT_MAP: Final = {
"encoding": ("litellm.main", "encoding"),
"BaseVectorStore": (
"litellm.integrations.vector_store_integrations.base_vector_store",

66
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, Final
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
@ -17,7 +17,7 @@ if set_verbose is True:
"`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."
)
_ENABLE_SECRET_REDACTION = os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true"
_ENABLE_SECRET_REDACTION: Final = os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true"
def _redact_string(value: str) -> str:
@ -74,19 +74,19 @@ class SecretRedactionFilter(logging.Filter):
return True
_secret_filter = SecretRedactionFilter()
_secret_filter: Final = SecretRedactionFilter()
json_logs = bool(os.getenv("JSON_LOGS", False))
# Create a handler for the logger (you may need to adapt this based on your needs)
log_level = os.getenv("LITELLM_LOG", "DEBUG")
numeric_level: str = getattr(logging, log_level.upper())
handler = logging.StreamHandler()
log_level: Final = os.getenv("LITELLM_LOG", "DEBUG")
numeric_level: Final[str] = getattr(logging, log_level.upper())
handler: Final = logging.StreamHandler()
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).
@ -94,16 +94,16 @@ def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]:
"""
if not message or not isinstance(message, str):
return None
msg_stripped = message.strip()
msg_stripped: Final = message.strip()
if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")):
return None
parsed = safe_json_loads(message, default=None)
parsed: Final = safe_json_loads(message, default=None)
if parsed is None or not isinstance(parsed, dict):
return None
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:
@ -144,21 +144,21 @@ def _get_standard_record_attrs() -> frozenset:
return frozenset(logging.LogRecord("", 0, "", 0, "", (), None).__dict__.keys())
_STANDARD_RECORD_ATTRS = _get_standard_record_attrs()
_STANDARD_RECORD_ATTRS: Final = _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
dt = datetime.fromtimestamp(record.created)
dt: Final = datetime.fromtimestamp(record.created)
return dt.isoformat()
def format(self, record):
message_str = record.getMessage()
json_record: Dict[str, Any] = {
message_str: Final = record.getMessage()
json_record: Final[dict[str, Any]] = {
"message": message_str,
"level": record.levelname,
"timestamp": self.formatTime(record),
@ -193,13 +193,13 @@ class JsonFormatter(Formatter):
# Function to set up exception handlers for JSON logging
def _setup_json_exception_handlers(formatter):
# Create a handler with JSON formatting for exceptions
error_handler = logging.StreamHandler()
error_handler: Final = logging.StreamHandler()
error_handler.setFormatter(formatter)
error_handler.addFilter(_secret_filter)
# Setup excepthook for uncaught exceptions
def json_excepthook(exc_type, exc_value, exc_traceback):
record = logging.LogRecord(
record: Final = logging.LogRecord(
name="LiteLLM",
level=logging.ERROR,
pathname="",
@ -217,10 +217,10 @@ def _setup_json_exception_handlers(formatter):
import asyncio
def async_json_exception_handler(loop, context):
exception = context.get("exception")
exception: Final = context.get("exception")
if exception:
exc_type = type(exception)
record = logging.LogRecord(
exc_type: Final = type(exception)
record: Final = logging.LogRecord(
name="LiteLLM",
level=logging.ERROR,
pathname="",
@ -243,7 +243,7 @@ if json_logs:
handler.setFormatter(JsonFormatter())
_setup_json_exception_handlers(JsonFormatter())
else:
formatter = logging.Formatter(
formatter: Final = logging.Formatter(
"\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s",
datefmt="%H:%M:%S",
)
@ -263,20 +263,20 @@ verbose_logger.addHandler(handler)
def _suppress_loggers():
"""Suppress noisy loggers at INFO level"""
# Suppress httpx request logging at INFO level
httpx_logger = logging.getLogger("httpx")
httpx_logger: Final = logging.getLogger("httpx")
httpx_logger.setLevel(logging.WARNING)
# Suppress APScheduler logging at INFO level
apscheduler_executors_logger = logging.getLogger("apscheduler.executors.default")
apscheduler_executors_logger: Final = logging.getLogger("apscheduler.executors.default")
apscheduler_executors_logger.setLevel(logging.WARNING)
apscheduler_scheduler_logger = logging.getLogger("apscheduler.scheduler")
apscheduler_scheduler_logger: Final = logging.getLogger("apscheduler.scheduler")
apscheduler_scheduler_logger.setLevel(logging.WARNING)
# Call the suppression function
_suppress_loggers()
ALL_LOGGERS = [
ALL_LOGGERS: Final = [
logging.getLogger(),
verbose_logger,
verbose_router_logger,
@ -293,11 +293,11 @@ def _get_loggers_to_initialize():
"""
import litellm
loggers = list(ALL_LOGGERS)
loggers: Final = list(ALL_LOGGERS)
# Add langfuse logger if langfuse is being used as a callback
langfuse_callbacks = {"langfuse", "langfuse_otel"}
all_callbacks = set(litellm.success_callback + litellm.failure_callback)
langfuse_callbacks: Final = {"langfuse", "langfuse_otel"}
all_callbacks: Final = set(litellm.success_callback + litellm.failure_callback)
if langfuse_callbacks & all_callbacks:
loggers.append(logging.getLogger("langfuse"))
@ -325,12 +325,12 @@ def _get_uvicorn_json_log_config():
This ensures that uvicorn's access logs, error logs, and all application logs
are formatted as JSON when json_logs is enabled.
"""
json_formatter_class = "litellm._logging.JsonFormatter"
json_formatter_class: Final = "litellm._logging.JsonFormatter"
# Use the module-level log_level variable for consistency
uvicorn_log_level = log_level.upper()
uvicorn_log_level: Final = log_level.upper()
log_config = {
log_config: Final = {
"version": 1,
"disable_existing_loggers": False,
"formatters": {
@ -384,7 +384,7 @@ def _turn_on_json():
- Adds a JSON formatter to all loggers
"""
handler = logging.StreamHandler()
handler: Final = logging.StreamHandler()
handler.setFormatter(JsonFormatter())
_initialize_loggers_with_handler(handler)
# Set up exception handlers

186
litellm/_redis.py
View file

@ -12,7 +12,8 @@ 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
from typing import Final
import redis # type: ignore
import redis.asyncio as async_redis # type: ignore
@ -32,20 +33,20 @@ from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from ._logging import verbose_logger
AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
def _get_redis_kwargs():
arg_spec = inspect.getfullargspec(redis.Redis)
arg_spec: Final = inspect.getfullargspec(redis.Redis)
# Only allow primitive arguments
exclude_args = {
exclude_args: Final = {
"self",
"connection_pool",
"retry",
}
include_args = {
include_args: Final = {
"url",
"redis_connect_func",
"gcp_service_account",
@ -56,7 +57,7 @@ def _get_redis_kwargs():
"azure_client_secret",
}
available_args = {x for x in arg_spec.args if x not in exclude_args} | include_args
available_args: Final = {x for x in arg_spec.args if x not in exclude_args} | include_args
return available_args
@ -76,7 +77,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
@ -92,9 +93,9 @@ def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]:
"""
if client is None:
client = redis.Redis
connection_cls = async_redis.Connection if client is async_redis.Redis else redis.Connection
connection_cls: Final = async_redis.Connection if client is async_redis.Redis else redis.Connection
exclude_args = frozenset(
exclude_args: Final = frozenset(
{
"self",
"connection_pool",
@ -103,7 +104,7 @@ def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]:
)
# Only allow primitive arguments
include_args = ("url", "max_connections")
include_args: Final = ("url", "max_connections")
return tuple(x for x in _init_arg_names(connection_cls) if x not in exclude_args) + include_args
@ -111,10 +112,10 @@ def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]:
def _get_redis_cluster_kwargs(client=None):
if client is None:
client = redis.Redis.from_url
arg_spec = inspect.getfullargspec(redis.RedisCluster)
arg_spec: Final = inspect.getfullargspec(redis.RedisCluster)
# Only allow primitive arguments
exclude_args = {"self", "connection_pool", "retry", "host", "port", "startup_nodes"}
exclude_args: Final = {"self", "connection_pool", "retry", "host", "port", "startup_nodes"}
available_args = {x for x in arg_spec.args if x not in exclude_args}
available_args |= {
@ -142,15 +143,15 @@ def _get_redis_cluster_kwargs(client=None):
def _get_redis_env_kwarg_mapping():
PREFIX = "REDIS_"
PREFIX: Final = "REDIS_"
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs()}
def _redis_kwargs_from_environment():
mapping = _get_redis_env_kwarg_mapping()
mapping: Final = _get_redis_env_kwarg_mapping()
return_dict = {}
return_dict: Final = {}
for k, v in mapping.items():
value = get_secret(k, default_value=None) # type: ignore
if value is not None:
@ -160,7 +161,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.
@ -183,7 +184,7 @@ def create_gcp_iam_redis_connect_func(
self._parser.on_connect(self)
auth_args = (_generate_gcp_iam_access_token(service_account),)
auth_args: Final = (_generate_gcp_iam_access_token(service_account),)
self.send_command("AUTH", *auth_args, check_health=False)
try:
@ -203,9 +204,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.
@ -224,9 +225,9 @@ def _build_azure_credential(
"azure-identity is required for Azure AD Redis authentication. Install it with: pip install azure-identity"
)
_client_id = azure_client_id or os.environ.get("AZURE_CLIENT_ID")
_tenant_id = azure_tenant_id or os.environ.get("AZURE_TENANT_ID")
_client_secret = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET")
_client_id: Final = azure_client_id or os.environ.get("AZURE_CLIENT_ID")
_tenant_id: Final = azure_tenant_id or os.environ.get("AZURE_TENANT_ID")
_client_secret: Final = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET")
if _client_id and _tenant_id and _client_secret:
return ClientSecretCredential(
@ -241,9 +242,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
@ -253,19 +254,19 @@ def _generate_azure_ad_redis_token(
(``AzureADCredentialProvider``) keep the credential alive across
connections so the Azure SDK's internal cache + silent refresh apply.
"""
credential = _build_azure_credential(
credential: Final = _build_azure_credential(
azure_client_id=azure_client_id,
azure_tenant_id=azure_tenant_id,
azure_client_secret=azure_client_secret,
)
token = credential.get_token(AZURE_REDIS_SCOPE)
token: Final = credential.get_token(AZURE_REDIS_SCOPE)
return token.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.
@ -274,7 +275,7 @@ def create_azure_ad_redis_connect_func(
closure) and reused across connections — the Azure SDK handles token caching
and silent renewal internally. Only ``get_token`` is called per connection.
"""
credential = _build_azure_credential(
credential: Final = _build_azure_credential(
azure_client_id=azure_client_id,
azure_tenant_id=azure_tenant_id,
azure_client_secret=azure_client_secret,
@ -290,11 +291,11 @@ def create_azure_ad_redis_connect_func(
self._parser.on_connect(self)
access_token = credential.get_token(AZURE_REDIS_SCOPE).token
access_token: Final = credential.get_token(AZURE_REDIS_SCOPE).token
# Only include username when explicitly set — sending AUTH "" <token>
# is invalid for most ACL-configured Azure Redis instances.
username = os.environ.get("REDIS_USERNAME", "")
username: Final = os.environ.get("REDIS_USERNAME", "")
if username:
auth_args = (username, access_token)
else:
@ -353,23 +354,23 @@ def _get_redis_client_logic(**env_overrides):
value = get_secret(v) # type: ignore
env_overrides[k] = value
environment_kwargs = _redis_kwargs_from_environment()
environment_kwargs: Final = _redis_kwargs_from_environment()
# An explicitly configured connection target outranks REDIS_URL from the
# environment. Without this, the url branch below strips the caller's
# host/port/password and silently connects to whatever REDIS_URL names.
caller_named_a_target = any(
caller_named_a_target: Final = any(
env_overrides.get(key) is not None for key in ("host", "startup_nodes", "sentinel_nodes")
)
if caller_named_a_target and env_overrides.get("url") is None:
environment_kwargs.pop("url", None)
redis_kwargs = {
redis_kwargs: Final = {
**environment_kwargs,
**env_overrides,
}
_startup_nodes: Optional[Union[str, list]] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
_startup_nodes: Final[str | list | None] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore
"REDIS_CLUSTER_NODES"
)
@ -380,21 +381,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: Final[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: Final[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: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
"REDIS_SERVICE_NAME"
)
@ -402,8 +403,8 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs["service_name"] = _service_name
# Handle GCP IAM authentication
_gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
_gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
_gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
_gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
if _gcp_service_account is not None:
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
@ -422,9 +423,9 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
# Handle Azure AD authentication (after GCP IAM block)
_azure_redis_ad_token = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
_azure_ad_enabled = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
@ -433,9 +434,9 @@ def _get_redis_client_logic(**env_overrides):
)
if _azure_ad_enabled and _gcp_service_account is None:
_azure_client_id = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
_azure_tenant_id = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
_azure_client_secret = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET")
_azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
_azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
_azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET")
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
@ -465,9 +466,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 +481,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: Final[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)
@ -489,13 +493,13 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
verbose_logger.debug("init_redis_cluster: startup nodes are being initialized.")
from redis.cluster import ClusterNode
args = _get_redis_cluster_kwargs()
cluster_kwargs = {}
args: Final = _get_redis_cluster_kwargs()
cluster_kwargs: Final = {}
for arg in redis_kwargs:
if arg in args:
cluster_kwargs[arg] = redis_kwargs[arg]
new_startup_nodes: List[ClusterNode] = []
new_startup_nodes: Final[list[ClusterNode]] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
@ -505,8 +509,8 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
def _get_redis_sentinel_connection_kwargs(redis_kwargs: dict) -> dict:
connection_kwargs = {}
args = _get_redis_kwargs()
connection_kwargs: Final = {}
args: Final = _get_redis_kwargs()
for arg in redis_kwargs:
if arg in args:
connection_kwargs[arg] = redis_kwargs[arg]
@ -515,12 +519,12 @@ def _get_redis_sentinel_connection_kwargs(redis_kwargs: dict) -> dict:
def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
sentinel_nodes = redis_kwargs.get("sentinel_nodes")
sentinel_password = redis_kwargs.get("sentinel_password")
service_name = redis_kwargs.get("service_name")
connection_kwargs = _get_redis_sentinel_connection_kwargs(redis_kwargs)
sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes")
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs = dict(connection_kwargs)
sentinel_kwargs: Final = dict(connection_kwargs)
sentinel_kwargs["password"] = sentinel_password
if not sentinel_nodes or not service_name:
@ -529,7 +533,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.")
# Set up the Sentinel client
sentinel = redis.Sentinel(
sentinel: Final = redis.Sentinel(
sentinel_nodes,
sentinel_kwargs=sentinel_kwargs,
)
@ -540,12 +544,12 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
sentinel_nodes = redis_kwargs.get("sentinel_nodes")
sentinel_password = redis_kwargs.get("sentinel_password")
service_name = redis_kwargs.get("service_name")
connection_kwargs = _get_redis_sentinel_connection_kwargs(redis_kwargs)
sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes")
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs = dict(connection_kwargs)
sentinel_kwargs: Final = dict(connection_kwargs)
sentinel_kwargs["password"] = sentinel_password
if not sentinel_nodes or not service_name:
@ -554,7 +558,7 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.")
# Set up the Sentinel client
sentinel = async_redis.Sentinel(
sentinel: Final = async_redis.Sentinel(
sentinel_nodes,
sentinel_kwargs=sentinel_kwargs,
)
@ -565,14 +569,14 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
def get_redis_client(**env_overrides):
redis_kwargs = _get_redis_client_logic(**env_overrides)
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
if "startup_nodes" in redis_kwargs:
return init_redis_cluster(redis_kwargs)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
args = _get_redis_url_kwargs()
url_kwargs = {}
args: Final = _get_redis_url_kwargs()
url_kwargs: Final = {}
for arg in redis_kwargs:
if arg in args:
url_kwargs[arg] = redis_kwargs[arg]
@ -587,16 +591,16 @@ 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]:
redis_kwargs = _get_redis_client_logic(**env_overrides)
) -> async_redis.Redis | async_redis.RedisCluster:
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
if "startup_nodes" in redis_kwargs:
from redis.cluster import ClusterNode
args = _get_redis_cluster_kwargs()
cluster_kwargs = {}
cluster_kwargs: Final = {}
for arg in redis_kwargs:
if arg in args:
cluster_kwargs[arg] = redis_kwargs[arg]
@ -618,7 +622,7 @@ def get_redis_async_client(
username=os.environ.get("REDIS_USERNAME") or None,
)
new_startup_nodes: List[ClusterNode] = []
new_startup_nodes: Final[list[ClusterNode]] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
@ -632,7 +636,7 @@ def get_redis_async_client(
cluster_kwargs.setdefault("socket_keepalive", True)
# Create async RedisCluster with IAM token as password if available
cluster_client = async_redis.RedisCluster(
cluster_client: Final = async_redis.RedisCluster(
startup_nodes=new_startup_nodes,
**cluster_kwargs, # type: ignore
)
@ -643,13 +647,13 @@ def get_redis_async_client(
if connection_pool is not None:
return async_redis.Redis(connection_pool=connection_pool)
args = _get_redis_url_kwargs(client=async_redis.Redis)
url_kwargs = {}
url_kwargs: Final = {}
for arg in redis_kwargs:
if arg in args:
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,16 +686,16 @@ def get_redis_async_client(
def get_redis_connection_pool(
**env_overrides,
) -> Optional[async_redis.BlockingConnectionPool]:
redis_kwargs = _get_redis_client_logic(**env_overrides)
) -> async_redis.BlockingConnectionPool | None:
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
if "startup_nodes" in redis_kwargs:
return None
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
allowed_args = _get_redis_url_kwargs(client=async_redis.Redis)
pool_kwargs = {k: v for k, v in redis_kwargs.items() if k in allowed_args and k != "max_connections"}
allowed_args: Final = _get_redis_url_kwargs(client=async_redis.Redis)
pool_kwargs: Final = {k: v for k, v in redis_kwargs.items() if k in allowed_args and k != "max_connections"}
pool_kwargs["timeout"] = REDIS_CONNECTION_POOL_TIMEOUT
pool_kwargs["url"] = redis_kwargs["url"]
if "max_connections" in redis_kwargs:
@ -707,7 +711,7 @@ def get_redis_connection_pool(
# Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed
# connections re-fetch tokens via the SDK's internal cache + silent refresh
# rather than reusing a single token captured at pool creation.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
redis_connect_func: Final = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
@ -734,7 +738,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
if not verbose_logger.isEnabledFor(logging.DEBUG):
return
console = Console()
console: Final = Console()
# Initialize the sensitive data masker
masker = SensitiveDataMasker()
@ -743,10 +747,10 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
# Create main panel title
title = Text("Redis Configuration", style="bold blue")
title: Final = Text("Redis Configuration", style="bold blue")
# Create configuration table
config_table = Table(
config_table: Final = Table(
title="🔧 Redis Connection Parameters",
show_header=True,
header_style="bold magenta",
@ -783,7 +787,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
connection_type = "Redis (URL-based)"
# Create connection type info
info_table = Table(
info_table: Final = Table(
title="📊 Connection Info",
show_header=True,
header_style="bold green",
@ -804,6 +808,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)

34
litellm/_redis_credential_provider.py
View file

@ -1,21 +1,21 @@
import asyncio
import threading
import time
from typing import Any, Dict, Optional, Tuple, Union
from typing import Any, Final
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
# Azure AD scope for Redis Cache for Azure.
AZURE_REDIS_SCOPE = "https://redis.azure.com/.default"
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
# GCP IAM tokens are valid for 1 hour. Cache for 55 minutes to refresh before expiry.
_GCP_IAM_TOKEN_TTL_SECONDS = 3300
_GCP_IAM_TOKEN_TTL_SECONDS: Final = 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_lock = threading.Lock()
_token_cache: Final[dict[str, tuple[str, float]]] = {}
_token_cache_lock: Final = threading.Lock()
def _generate_gcp_iam_access_token(service_account: str) -> str:
@ -36,12 +36,12 @@ def _generate_gcp_iam_access_token(service_account: str) -> str:
"Install it with: pip install google-cloud-iam"
)
client = iam_credentials_v1.IAMCredentialsClient()
request = iam_credentials_v1.GenerateAccessTokenRequest(
client: Final = iam_credentials_v1.IAMCredentialsClient()
request: Final = iam_credentials_v1.GenerateAccessTokenRequest(
name=service_account,
scope=["https://www.googleapis.com/auth/cloud-platform"],
)
response = client.generate_access_token(request=request)
response: Final = client.generate_access_token(request=request)
return str(response.access_token)
@ -95,12 +95,12 @@ class GCPIAMCredentialProvider(CredentialProvider):
def __init__(self, gcp_service_account: str) -> None:
self._gcp_service_account = gcp_service_account
def get_credentials(self) -> Tuple[str]:
token = _get_cached_gcp_iam_token(self._gcp_service_account)
def get_credentials(self) -> tuple[str]:
token: Final = _get_cached_gcp_iam_token(self._gcp_service_account)
return (token,)
async def get_credentials_async(self) -> Tuple[str]:
token = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account)
async def get_credentials_async(self) -> tuple[str]:
token: Final = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account)
return (token,)
@ -115,18 +115,18 @@ 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]]:
token = self._credential.get_token(AZURE_REDIS_SCOPE).token
def get_credentials(self) -> tuple[str] | tuple[str, str]:
token: Final = 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]]:
token_obj = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE)
async def get_credentials_async(self) -> tuple[str] | tuple[str, str]:
token_obj: Final = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE)
if self._username:
return (self._username, token_obj.token)
return (token_obj.token,)

50
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, Final, 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
@ -67,7 +67,7 @@ class ServiceLogging(CustomLogger):
whether the callback is the logger instance itself or the ``"otel"`` string
(which routes to the proxy's registered ``open_telemetry_logger``).
"""
otel_v2_cls = _get_otel_v2_class()
otel_v2_cls: Final = _get_otel_v2_class()
def _is_otel_logger(obj: Any) -> bool:
if isinstance(obj, OpenTelemetry):
@ -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.
@ -101,7 +101,7 @@ class ServiceLogging(CustomLogger):
try:
# Try to get the current event loop
loop = asyncio.get_event_loop()
loop: Final = asyncio.get_event_loop()
# Check if the loop is running
if loop.is_running():
# If we're in a running loop, create a task
@ -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
@ -163,7 +163,7 @@ class ServiceLogging(CustomLogger):
if self.mock_testing:
self.mock_testing_async_success_hook += 1
payload = ServiceLoggerPayload(
payload: Final = ServiceLoggerPayload(
is_error=False,
error=None,
service=service,
@ -178,7 +178,7 @@ class ServiceLogging(CustomLogger):
# (the V2 logger self-registers its instance even when the string is
# present, unlike V1). Without this guard each such reference emits its own
# span, so a single DB call shows up as duplicate ``postgres ...`` spans.
emitted_otel_logger_ids: set = set()
emitted_otel_logger_ids: Final[set] = set()
for callback in litellm.service_callback:
if callback == "prometheus_system":
await self.init_prometheus_services_logger_if_none()
@ -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
@ -271,7 +267,7 @@ class ServiceLogging(CustomLogger):
elif isinstance(error, str):
error_message = error
payload = ServiceLoggerPayload(
payload: Final = ServiceLoggerPayload(
is_error=True,
error=error_message,
service=service,
@ -282,7 +278,7 @@ class ServiceLogging(CustomLogger):
# Dedupe OTel loggers per event — see ``async_service_success_hook`` for why
# the same logger can be referenced twice in ``service_callback``.
emitted_otel_logger_ids: set = set()
emitted_otel_logger_ids: Final[set] = set()
for callback in litellm.service_callback:
if callback == "prometheus_system":
await self.init_prometheus_services_logger_if_none()
@ -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.

28
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, Final
from litellm._logging import verbose_logger
from litellm.constants import LOCALHOST_URL_PATTERNS
@ -43,18 +43,18 @@ def is_localhost_or_internal_url(url: str | None) -> bool:
if not url:
return False
url_lower = url.lower()
url_lower: Final = url.lower()
return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS)
def get_agent_card_url(agent_card: "AgentCard") -> str | None:
"""Return the agent endpoint URL from the resolved SDK card."""
url = getattr(agent_card, "url", None)
url: Final = getattr(agent_card, "url", None)
if url:
return url
interfaces = getattr(agent_card, "supported_interfaces", None)
interfaces: Final = getattr(agent_card, "supported_interfaces", None)
if interfaces:
return getattr(interfaces[0], "url", None)
return None
@ -62,11 +62,11 @@ def get_agent_card_url(agent_card: "AgentCard") -> str | None:
def set_agent_card_url(agent_card: "AgentCard", url: str) -> None:
"""Set the agent endpoint URL on the resolved SDK card."""
normalized = url.rstrip("/") + "/"
normalized: Final = url.rstrip("/") + "/"
if hasattr(agent_card, "url"):
agent_card.url = normalized
interfaces = getattr(agent_card, "supported_interfaces", None)
interfaces: Final = getattr(agent_card, "supported_interfaces", None)
if interfaces:
interfaces[0].url = normalized
@ -86,16 +86,16 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard":
Returns:
The agent card with the URL fixed if necessary
"""
card_url = getattr(agent_card, "url", None)
card_url: Final = getattr(agent_card, "url", None)
if card_url and is_localhost_or_internal_url(card_url):
# Normalize base_url to ensure it ends with /
fixed_url = base_url.rstrip("/") + "/"
fixed_url: Final = base_url.rstrip("/") + "/"
agent_card.url = fixed_url
interfaces = getattr(agent_card, "supported_interfaces", None)
interfaces: Final = getattr(agent_card, "supported_interfaces", None)
if interfaces:
interface_url = getattr(interfaces[0], "url", None)
interface_url: Final = getattr(interfaces[0], "url", None)
if interface_url and is_localhost_or_internal_url(interface_url):
interfaces[0].url = base_url.rstrip("/") + "/"
@ -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.
@ -140,7 +140,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
)
# Try both well-known paths
paths = [
paths: Final = [
AGENT_CARD_WELL_KNOWN_PATH,
PREV_AGENT_CARD_WELL_KNOWN_PATH,
]
@ -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

11
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, Final
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."""
@ -91,7 +92,7 @@ class A2AClient:
"""Send a message to the A2A agent."""
from litellm.a2a_protocol.main import asend_message
a2a_client = await self._get_client()
a2a_client: Final = await self._get_client()
return await asend_message(a2a_client=a2a_client, request=request)
async def send_message_streaming(
@ -100,6 +101,6 @@ class A2AClient:
"""Send a streaming message to the A2A agent."""
from litellm.a2a_protocol.main import asend_message_streaming
a2a_client = await self._get_client()
a2a_client: Final = await self._get_client()
async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request):
yield chunk

28
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, Final
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.
@ -42,23 +42,23 @@ class A2ACostCalculator:
if litellm_logging_obj is None:
return 0.0
model_call_details = litellm_logging_obj.model_call_details
model_call_details: Final = litellm_logging_obj.model_call_details
# Check if user set a custom response cost (backward compatibility)
response_cost = model_call_details.get("response_cost", None)
response_cost: Final = model_call_details.get("response_cost", None)
if response_cost is not None:
return float(response_cost)
# Get litellm_params for cost parameters
litellm_params = model_call_details.get("litellm_params", {}) or {}
litellm_params: Final = model_call_details.get("litellm_params", {}) or {}
# Check for cost_per_query (fixed cost per query)
if litellm_params.get("cost_per_query") is not None:
return float(litellm_params["cost_per_query"])
# Check for token-based pricing
input_cost_per_token = litellm_params.get("input_cost_per_token")
output_cost_per_token = litellm_params.get("output_cost_per_token")
input_cost_per_token: Final = litellm_params.get("input_cost_per_token")
output_cost_per_token: Final = litellm_params.get("output_cost_per_token")
if input_cost_per_token is not None or output_cost_per_token is not None:
return A2ACostCalculator._calculate_token_based_cost(
@ -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.
@ -88,16 +88,16 @@ class A2ACostCalculator:
float: The calculated cost
"""
# Get usage from model_call_details
usage = model_call_details.get("usage")
usage: Final = model_call_details.get("usage")
if usage is None:
return 0.0
# Get token counts
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
prompt_tokens: Final = getattr(usage, "prompt_tokens", 0) or 0
completion_tokens: Final = getattr(usage, "completion_tokens", 0) or 0
# Calculate costs
input_cost = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0)
output_cost = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0)
input_cost: Final = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0)
output_cost: Final = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0)
return input_cost + output_cost

32
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, Final
from litellm._logging import verbose_logger
from litellm.a2a_protocol.card_resolver import (
@ -53,11 +53,11 @@ class A2AExceptionCheckers:
if not isinstance(error_str, str):
return False
error_str_lower = error_str.lower()
error_str_lower: Final = error_str.lower()
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.
@ -83,8 +83,8 @@ class A2AExceptionCheckers:
if not isinstance(error_str, str):
return False
error_str_lower = error_str.lower()
agent_card_patterns = [
error_str_lower: Final = error_str.lower()
agent_card_patterns: Final = [
"agent card",
"agent-card",
".well-known",
@ -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.
@ -118,7 +118,7 @@ def map_a2a_exception(
A2AAgentCardError: If the error is related to agent card issues
A2AError: For other A2A-related errors
"""
error_str = str(original_exception)
error_str: Final = str(original_exception)
# Check for localhost URL connection error (special case - retryable)
if (
@ -190,11 +190,13 @@ async def handle_a2a_localhost_retry(
"rewrite, so the upstream URL cannot be corrected."
)
request_type = "streaming " if is_streaming else ""
request_type: Final = "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
@ -203,14 +205,14 @@ async def handle_a2a_localhost_retry(
# Reuse the httpx client LiteLLM attached at creation. It carries this agent's
# trace-id and auth headers, so a fresh client would drop them. Only clients built
# by ``create_a2a_client`` have it; an externally-supplied client cannot be retried.
httpx_client = getattr(a2a_client, "_litellm_httpx_client", None)
httpx_client: Final = getattr(a2a_client, "_litellm_httpx_client", None)
if httpx_client is None:
raise RuntimeError(
"Cannot retry A2A localhost URL fix: the client was not created by "
"create_a2a_client, so no LiteLLM httpx client is attached."
)
new_client = await create_client( # pyright: ignore[reportOptionalCall]
new_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
agent_card,
client_config=ClientConfig( # pyright: ignore[reportOptionalCall]
httpx_client=httpx_client,

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

101
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, Final
import litellm
from litellm._logging import verbose_logger
@ -24,10 +25,10 @@ from litellm.interactions.agents.utils import merge_agent_headers
# litellm_params key carrying the authenticated principal (hashed virtual key) so
# A2A provider configs can scope provider-side state (e.g. LangFlow session memory)
# per key instead of trusting the client-supplied A2A contextId.
A2A_USER_API_KEY_HASH_PARAM = "litellm_a2a_user_api_key_hash"
A2A_USER_API_KEY_HASH_PARAM: Final = "litellm_a2a_user_api_key_hash"
# Agent metadata fields stored in litellm_params that are not valid litellm.acompletion() kwargs
_AGENT_ONLY_PARAMS = frozenset(
_AGENT_ONLY_PARAMS: Final = frozenset(
{
"is_public",
"agent_name",
@ -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.
@ -69,13 +70,13 @@ class A2ACompletionBridgeHandler:
"""
custom_llm_provider = litellm_params.get("custom_llm_provider")
if not _skip_a2a_provider_routing:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
a2a_provider_config: Final = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
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,
@ -86,14 +87,14 @@ class A2ACompletionBridgeHandler:
)
# Extract message from params
message = params.get("message", {})
message: Final = params.get("message", {})
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
model: Final = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
@ -102,17 +103,17 @@ 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: Final[dict[str, Any]] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": False,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
litellm_params_to_add: Final = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
@ -134,28 +135,28 @@ class A2ACompletionBridgeHandler:
)
# Call litellm.acompletion
response = await litellm.acompletion(**completion_params)
response: Final = await litellm.acompletion(**completion_params)
# Transform response to A2A format
a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
a2a_response: Final = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
response=response,
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.
@ -178,13 +179,13 @@ class A2ACompletionBridgeHandler:
"""
custom_llm_provider = litellm_params.get("custom_llm_provider")
if not _skip_a2a_provider_routing:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
a2a_provider_config: Final = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider,
model=litellm_params.get("model"),
)
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,
@ -198,20 +199,20 @@ class A2ACompletionBridgeHandler:
return
# Extract message from params
message = params.get("message", {})
message: Final = params.get("message", {})
# Create streaming context
ctx = A2AStreamingContext(
ctx: Final = A2AStreamingContext(
request_id=request_id,
input_message=message,
)
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
model: Final = litellm_params.get("model", "agent")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
@ -220,17 +221,17 @@ 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: Final[dict[str, Any]] = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": True,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
litellm_params_to_add: Final = {
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
@ -252,11 +253,11 @@ class A2ACompletionBridgeHandler:
)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)
task_event: Final = A2ACompletionBridgeTransformation.create_task_event(ctx)
yield task_event
# 2. Emit status update (kind: "status-update", status: "working")
working_event = A2ACompletionBridgeTransformation.create_status_update_event(
working_event: Final = A2ACompletionBridgeTransformation.create_status_update_event(
ctx=ctx,
state="working",
final=False,
@ -265,7 +266,7 @@ class A2ACompletionBridgeHandler:
yield working_event
# Call litellm.acompletion with streaming
response = await litellm.acompletion(**completion_params)
response: Final = await litellm.acompletion(**completion_params)
# 3. Accumulate content and emit artifact update
accumulated_text = ""
@ -285,31 +286,33 @@ class A2ACompletionBridgeHandler:
# Emit artifact update with accumulated content
if accumulated_text:
artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event(
artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=accumulated_text,
)
yield artifact_event
# 4. Emit final status update (kind: "status-update", status: "completed", final: true)
completed_event = A2ACompletionBridgeTransformation.create_status_update_event(
completed_event: Final = A2ACompletionBridgeTransformation.create_status_update_event(
ctx=ctx,
state="completed",
final=True,
)
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,

70
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, Final
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: Final[list[str]] = []
for part in parts:
if not isinstance(part, dict):
continue
@ -62,35 +62,35 @@ 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: Final[dict[str, Any]] = {}
if params and isinstance(params.get("metadata"), dict):
merged.update(params["metadata"])
message_metadata = a2a_message.get("metadata")
message_metadata: Final = a2a_message.get("metadata")
if isinstance(message_metadata, dict):
merged.update(message_metadata)
return merged or None
@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).
Uses ``extra_body`` so we do not collide with LiteLLM's spend-log ``metadata`` kwarg.
"""
forward_metadata = A2ACompletionBridgeTransformation.get_forward_metadata(
forward_metadata: Final = A2ACompletionBridgeTransformation.get_forward_metadata(
a2a_message=a2a_message,
params=params,
)
@ -103,18 +103,18 @@ class A2ACompletionBridgeTransformation:
# Layer client-supplied A2A metadata under any agent-owner-configured
# ``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_metadata: Final = extra_body.get("metadata")
existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: Final[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.
@ -124,7 +124,7 @@ class A2ACompletionBridgeTransformation:
Returns:
List of OpenAI-format messages
"""
role = a2a_message.get("role", "user")
role: Final = a2a_message.get("role", "user")
parts = a2a_message.get("parts", [])
# Map A2A roles to OpenAI roles
@ -139,21 +139,23 @@ class A2ACompletionBridgeTransformation:
if not isinstance(parts, list):
parts = []
content = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts)
content: Final = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts)
# 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: Final[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.
@ -167,12 +169,12 @@ class A2ACompletionBridgeTransformation:
# Extract content from response
content = ""
if hasattr(response, "choices") and response.choices:
choice = response.choices[0]
choice: Final = response.choices[0]
if hasattr(choice, "message") and choice.message:
content = choice.message.content or ""
# Build A2A message
a2a_message = {
a2a_message: Final = {
"kind": "message",
"role": "agent",
"parts": [{"kind": "text", "text": content}],
@ -180,13 +182,13 @@ class A2ACompletionBridgeTransformation:
}
# Build A2A response
a2a_response = {
a2a_response: Final = {
"jsonrpc": "2.0",
"id": request_id,
"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: Final[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.

162
litellm/a2a_protocol/main.py
View file

@ -12,16 +12,8 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
import asyncio
import datetime
import uuid
from typing import (
TYPE_CHECKING,
Any,
AsyncIterator,
Coroutine,
Dict,
Optional,
Union,
cast,
)
from collections.abc import AsyncIterator, Coroutine
from typing import TYPE_CHECKING, Any, Final, Optional, cast
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
@ -83,11 +75,11 @@ from litellm.a2a_protocol.exception_mapping_utils import (
from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
# Use our custom resolver instead of the default A2A SDK resolver
A2ACardResolver = LiteLLMA2ACardResolver
A2ACardResolver: Final = LiteLLMA2ACardResolver
def _set_usage_on_logging_obj(
kwargs: Dict[str, Any],
kwargs: dict[str, Any],
prompt_tokens: int,
completion_tokens: int,
) -> None:
@ -99,9 +91,9 @@ def _set_usage_on_logging_obj(
prompt_tokens: Number of input tokens
completion_tokens: Number of output tokens
"""
litellm_logging_obj = kwargs.get("litellm_logging_obj")
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
usage = litellm.Usage(
usage: Final = litellm.Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
@ -110,7 +102,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:
"""
@ -123,13 +115,13 @@ def _set_agent_id_on_logging_obj(
if agent_id is None:
return
litellm_logging_obj = kwargs.get("litellm_logging_obj")
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
# Set agent_id directly on model_call_details (same pattern as custom_llm_provider)
litellm_logging_obj.model_call_details["agent_id"] = agent_id
_A2A_COST_PARAM_KEYS = ("cost_per_query", "input_cost_per_token", "output_cost_per_token")
_A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output_cost_per_token")
def _set_litellm_params_on_logging_obj(
@ -144,7 +136,7 @@ def _set_litellm_params_on_logging_obj(
litellm_params already carries metadata / proxy_server_request / user-key
context, so merge the pricing keys in rather than replacing the dict.
"""
logging_obj = kwargs.get("litellm_logging_obj")
logging_obj: Final = kwargs.get("litellm_logging_obj")
if logging_obj is None:
return
@ -152,11 +144,11 @@ def _set_litellm_params_on_logging_obj(
if not cost_params:
return
existing = logging_obj.model_call_details.get("litellm_params") or {}
existing: Final = logging_obj.model_call_details.get("litellm_params") or {}
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.
@ -165,17 +157,17 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
"""
agent_name = "unknown"
agent_card = _get_a2a_client_agent_card(a2a_client)
agent_card: Final = _get_a2a_client_agent_card(a2a_client)
if agent_card is not None:
agent_name = getattr(agent_card, "name", "unknown") or "unknown"
# Build model string
model = f"a2a_agent/{agent_name}"
custom_llm_provider = "a2a_agent"
model: Final = f"a2a_agent/{agent_name}"
custom_llm_provider: Final = "a2a_agent"
# Set on litellm_logging_obj if available (for standard logging payload)
litellm_logging_obj = kwargs.get("litellm_logging_obj")
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
litellm_logging_obj.model = model
litellm_logging_obj.custom_llm_provider = custom_llm_provider
@ -199,15 +191,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,
@ -215,7 +207,7 @@ async def _send_message_via_completion_bridge(
params = request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params)
response_dict = await A2ACompletionBridgeHandler.handle_non_streaming(
response_dict: Final = await A2ACompletionBridgeHandler.handle_non_streaming(
request_id=str(request.id),
params=params,
litellm_params=litellm_params,
@ -233,18 +225,18 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
)
pb_request = _a2a_conversions.to_core_send_message_request(request)
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
last_event = None
async for event in a2a_client.send_message(pb_request):
last_event = event
if last_event is None:
raise RuntimeError("A2A send_message failed: no response received from agent.")
stream_compat = _a2a_conversions.to_compat_stream_response(
stream_compat: Final = _a2a_conversions.to_compat_stream_response(
last_event,
request_id=request.id,
)
result = stream_compat.result
result: Final = stream_compat.result
if not isinstance(result, (Message, Task)):
raise RuntimeError(
"A2A send_message failed: non-streaming message/send expects the "
@ -308,7 +300,7 @@ async def _stream_messages(
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
)
pb_request = _a2a_conversions.to_core_send_message_request(request)
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
async for event in a2a_client.send_message(pb_request):
compat_chunk = _a2a_conversions.to_compat_stream_response(
event,
@ -370,9 +362,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:
"""
@ -428,9 +420,9 @@ async def asend_message(
```
"""
litellm_params = litellm_params or {}
logging_obj = kwargs.get("litellm_logging_obj")
logging_obj: Final = kwargs.get("litellm_logging_obj")
trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None
custom_llm_provider = litellm_params.get("custom_llm_provider")
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
# Route through completion bridge if custom_llm_provider is set
if custom_llm_provider:
@ -453,7 +445,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: Final[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)
@ -464,15 +456,15 @@ async def asend_message(
# Type assertion: a2a_client is guaranteed to be non-None here
assert a2a_client is not None
agent_name = _get_a2a_model_info(a2a_client, kwargs)
agent_name: Final = _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)
card_url = get_agent_card_url(agent_card) if agent_card else None
agent_card: Final = _get_a2a_client_agent_card(a2a_client)
card_url: Final = get_agent_card_url(agent_card) if agent_card else None
a2a_response = await _execute_a2a_send_with_retry(
a2a_response: Final = await _execute_a2a_send_with_retry(
a2a_client=a2a_client,
request=request,
agent_card=agent_card,
@ -481,13 +473,13 @@ 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))
response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
# Calculate token usage from request and response
response_dict = a2a_response.model_dump(mode="json", exclude_none=True)
response_dict: Final = a2a_response.model_dump(mode="json", exclude_none=True)
(
prompt_tokens,
completion_tokens,
@ -518,7 +510,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,15 +539,15 @@ 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()
model = f"a2a_agent/{agent_name}"
start_time: Final = datetime.datetime.now()
model: Final = f"a2a_agent/{agent_name}"
logging_obj = Logging(
logging_obj: Final = Logging(
model=model,
messages=[{"role": "user", "content": "streaming-request"}],
stream=False,
@ -572,7 +564,7 @@ def _build_streaming_logging_obj(
if agent_id:
logging_obj.model_call_details["agent_id"] = agent_id
_litellm_params = litellm_params.copy() if litellm_params else {}
_litellm_params: Final = litellm_params.copy() if litellm_params else {}
if metadata:
_litellm_params["metadata"] = metadata
if proxy_server_request:
@ -590,11 +582,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]:
"""
@ -635,7 +627,7 @@ async def asend_message_streaming(
```
"""
litellm_params = litellm_params or {}
custom_llm_provider = litellm_params.get("custom_llm_provider")
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
# Route through completion bridge if custom_llm_provider is set
if custom_llm_provider:
@ -643,14 +635,14 @@ 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,
)
# Extract params from request
params = (
params: Final = (
request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params)
)
@ -667,15 +659,15 @@ async def asend_message_streaming(
if request is None:
raise ValueError("request is required")
_raw_logging_obj = kwargs.get("litellm_logging_obj")
_raw_logging_obj: Final = kwargs.get("litellm_logging_obj")
logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None
if a2a_client is None:
if api_base is None:
raise ValueError("Either a2a_client or api_base is required for standard A2A flow")
logging_trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None
trace_id = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4()))
extra_headers: dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id}
logging_trace_id: Final = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None
trace_id: Final = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4()))
extra_headers: Final[dict[str, str]] = {"X-LiteLLM-Trace-Id": trace_id}
if agent_id:
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
if agent_extra_headers:
@ -688,7 +680,7 @@ async def asend_message_streaming(
assert a2a_client is not None
agent_name = _get_a2a_model_info(a2a_client, kwargs)
agent_name: Final = _get_a2a_model_info(a2a_client, kwargs)
if logging_obj is None:
logging_obj = _build_streaming_logging_obj(
@ -700,12 +692,12 @@ 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
agent_card: Final = _get_a2a_client_agent_card(a2a_client)
card_url: Final = get_agent_card_url(agent_card) if agent_card else None
stream = _execute_a2a_stream_with_retry(
stream: Final = _execute_a2a_stream_with_retry(
a2a_client=a2a_client,
request=request,
agent_card=agent_card,
@ -728,7 +720,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 +754,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
@ -772,21 +764,21 @@ async def create_a2a_client(
# Only pass params that AsyncHTTPHandler.__init__ accepts (e.g. timeout).
# Use "disable_aiohttp_transport" key for cache-key-only data (it's
# filtered out before reaching the constructor).
_client_params: dict = {"timeout": timeout}
_client_params: Final[dict] = {"timeout": timeout}
if extra_headers:
# Encode headers into a cache-key-only param so each unique header
# set produces a distinct cache key.
_client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items()))
_async_handler = get_async_httpx_client(
_async_handler: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.A2AProvider,
params=_client_params,
)
httpx_client = _async_handler.client
httpx_client: Final = _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]
a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
base_url,
client_config=ClientConfig( # pyright: ignore[reportOptionalCall]
httpx_client=httpx_client,
@ -797,11 +789,11 @@ async def create_a2a_client(
# the configured httpx client (with this agent's trace-id/auth headers) without
# excavating a2a-sdk private internals.
a2a_client._litellm_httpx_client = httpx_client # type: ignore[attr-defined]
agent_card = getattr(a2a_client, "_card", None)
agent_card: Final = getattr(a2a_client, "_card", None)
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 +801,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,20 +819,20 @@ 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(
http_handler: Final = get_async_httpx_client(
llm_provider=httpxSpecialProvider.A2A,
params={"timeout": timeout},
)
httpx_client = http_handler.client
httpx_client: Final = http_handler.client
resolver = A2ACardResolver(
resolver: Final = A2ACardResolver(
httpx_client=httpx_client,
base_url=base_url,
)
agent_card = await resolver.get_agent_card()
agent_card: Final = 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.

19
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, Final
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.bedrock_agentcore.handler import (
@ -22,12 +23,12 @@ 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")
litellm_params: Final = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for BedrockAgentCoreA2AConfig (must contain model with AgentCore ARN)"
@ -42,12 +43,12 @@ 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")
litellm_params: Final = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for BedrockAgentCoreA2AConfig (must contain model with AgentCore ARN)"

41
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, Final, 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,31 +53,31 @@ 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(
client: Final = get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
)
response = await client.post(
response: Final = await client.post(
url,
headers=headers,
data=body,
)
response.raise_for_status()
response_data = response.json()
response_data: Final = 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,12 +100,12 @@ 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(
client: Final = get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
)
response = await client.post(
response: Final = await client.post(
url,
headers=headers,
data=body,
@ -113,15 +114,15 @@ class BedrockAgentCoreA2AHandler:
response.raise_for_status()
# Check content type — AgentCore may return JSON instead of SSE
content_type = response.headers.get("content-type", "").lower()
content_type: Final = response.headers.get("content-type", "").lower()
if "application/json" in content_type:
# Single JSON response fallback (not SSE)
verbose_logger.debug(
"BedrockAgentCore A2A streaming: received JSON instead of SSE, yielding as single event"
)
response_body = await response.aread()
response_data = json.loads(response_body)
response_body: Final = await response.aread()
response_data: Final = json.loads(response_body)
yield response_data
else:
# SSE stream — parse data: lines

45
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, Final
from litellm._logging import verbose_logger
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
@ -22,21 +23,21 @@ from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreCo
# ``runtimeSessionId`` / ``runtimeUserId`` in the agent's ``litellm_params``;
# ``authorization`` is set by the AgentCore signer (JWT or SigV4); ``host`` and
# the ``x-amz-*`` family are owned by SigV4 itself.
_RESERVED_EXACT_HEADERS = frozenset(
_RESERVED_EXACT_HEADERS: Final = frozenset(
{
"authorization",
"host",
}
)
_RESERVED_PREFIX_HEADERS: Tuple[str, ...] = (
_RESERVED_PREFIX_HEADERS: Final[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,8 +47,8 @@ def _filter_reserved_headers(
if not agent_extra_headers:
return None
filtered: Dict[str, str] = {}
dropped: list = []
filtered: Final[dict[str, str]] = {}
dropped: Final[list] = []
for k, v in agent_extra_headers.items():
k_lower = k.lower()
if k_lower in _RESERVED_EXACT_HEADERS or any(k_lower.startswith(prefix) for prefix in _RESERVED_PREFIX_HEADERS):
@ -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.
@ -106,19 +107,19 @@ class BedrockAgentCoreA2ATransformation:
"""
# Extract model and strip the "bedrock/" prefix
# "bedrock/agentcore/arn:aws:..." → "agentcore/arn:aws:..."
model = litellm_params.get("model", "")
model: Final = litellm_params.get("model", "")
if model.startswith("bedrock/"):
agentcore_model = model[len("bedrock/") :]
else:
agentcore_model = model
# Build optional_params from litellm_params (everything except model and custom_llm_provider)
optional_params = {k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider")}
optional_params: Final = {k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider")}
agentcore_config = AmazonAgentCoreConfig()
agentcore_config: Final = AmazonAgentCoreConfig()
# Derive URL from ARN
url = agentcore_config.get_complete_url(
url: Final = agentcore_config.get_complete_url(
api_base=optional_params.get("api_base"),
api_key=optional_params.get("api_key"),
model=agentcore_model,
@ -128,7 +129,7 @@ class BedrockAgentCoreA2ATransformation:
)
# Construct JSON-RPC 2.0 envelope
json_rpc_body = {
json_rpc_body: Final = {
"jsonrpc": "2.0",
"method": method,
"id": request_id,
@ -137,17 +138,17 @@ class BedrockAgentCoreA2ATransformation:
# Set required AgentCore session headers (normally set by transform_request,
# which we skip because it also builds {"prompt": "..."})
headers: dict = {}
session_id = agentcore_config._get_runtime_session_id(optional_params)
headers: Final[dict] = {}
session_id: Final = agentcore_config._get_runtime_session_id(optional_params)
headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = session_id
runtime_user_id = agentcore_config._get_runtime_user_id(optional_params)
runtime_user_id: Final = agentcore_config._get_runtime_user_id(optional_params)
if runtime_user_id:
headers["X-Amzn-Bedrock-AgentCore-Runtime-User-Id"] = runtime_user_id
# Merge per-request agent headers before signing so SigV4 covers them.
# Reserved headers are stripped first to prevent client-controlled values
# from spoofing the AgentCore runtime identity / SigV4 metadata.
safe_extra_headers = _filter_reserved_headers(agent_extra_headers)
safe_extra_headers: Final = _filter_reserved_headers(agent_extra_headers)
if safe_extra_headers:
headers.update(safe_extra_headers)
@ -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")

27
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, Final
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,10 +47,10 @@ 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(
response_data: Final = await PydanticAITransformation.send_non_streaming_request(
api_base=api_base,
request_id=request_id,
params=params,
@ -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,10 +92,10 @@ 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(
raw_response: Final = await PydanticAITransformation.send_and_get_raw_response(
api_base=api_base,
request_id=request_id,
params=params,

87
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, Final, 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.
@ -162,7 +163,7 @@ class PydanticAITransformation:
params_dict["message"]["kind"] = "message"
# Build A2A JSON-RPC request using message/send method for FastA2A compatibility
a2a_request = {
a2a_request: Final = {
"jsonrpc": "2.0",
"id": request_id,
"method": "message/send",
@ -170,16 +171,16 @@ class PydanticAITransformation:
}
# FastA2A uses root endpoint (/) not /messages
endpoint = api_base.rstrip("/")
endpoint: Final = 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(
client: Final = get_async_httpx_client(
llm_provider=cast(Any, "pydantic_ai_agent"),
params={"timeout": timeout},
)
response = await client.post(
response: Final = await client.post(
endpoint,
json=a2a_request,
headers={
@ -191,15 +192,15 @@ class PydanticAITransformation:
response_data = response.json()
# Check if task is already completed
result = response_data.get("result", {})
status = result.get("status", {})
state = status.get("state", "")
result: Final = response_data.get("result", {})
status: Final = result.get("status", {})
state: Final = status.get("state", "")
if state != "completed":
# Need to poll for completion
task_id = result.get("id")
task_id: Final = 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.
@ -234,7 +235,7 @@ class PydanticAITransformation:
Standard A2A non-streaming response format with message
"""
# Get raw task response
raw_response = await PydanticAITransformation._send_and_poll_raw(
raw_response: Final = await PydanticAITransformation._send_and_poll_raw(
api_base=api_base,
request_id=request_id,
params=params,
@ -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.
@ -312,7 +313,7 @@ class PydanticAITransformation:
full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data)
# Build standard A2A message
a2a_message = {
a2a_message: Final = {
"kind": "message",
"role": "agent",
"parts": parts if parts else [{"kind": "text", "text": full_text}],
@ -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.
@ -341,10 +342,10 @@ class PydanticAITransformation:
Returns:
Tuple of (full_text, message_id, parts)
"""
result = response_data.get("result", {})
result: Final = response_data.get("result", {})
# Try to extract from artifacts first (preferred for results)
artifacts = result.get("artifacts", [])
artifacts: Final = result.get("artifacts", [])
if artifacts:
for artifact in artifacts:
parts = artifact.get("parts", [])
@ -355,7 +356,7 @@ class PydanticAITransformation:
return text, str(uuid4()), parts
# Fall back to history - get the last agent message
history = result.get("history", [])
history: Final = result.get("history", [])
for msg in reversed(history):
if msg.get("role") == "agent":
parts = msg.get("parts", [])
@ -368,7 +369,7 @@ class PydanticAITransformation:
return full_text, message_id, parts
# Fall back to message field (original format)
message = result.get("message", {})
message: Final = result.get("message", {})
if message:
parts = message.get("parts", [])
message_id = message.get("messageId", str(uuid4()))
@ -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.
@ -409,8 +410,8 @@ class PydanticAITransformation:
full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data)
# Extract input message from raw response for history
result = response_data.get("result", {})
history = result.get("history", [])
result: Final = response_data.get("result", {})
history: Final = result.get("history", [])
input_message = {}
for msg in history:
if msg.get("role") == "user":
@ -418,14 +419,14 @@ class PydanticAITransformation:
break
# Generate IDs for streaming events
task_id = str(uuid4())
context_id = str(uuid4())
artifact_id = str(uuid4())
input_message_id = input_message.get("messageId", str(uuid4()))
task_id: Final = str(uuid4())
context_id: Final = str(uuid4())
artifact_id: Final = str(uuid4())
input_message_id: Final = input_message.get("messageId", str(uuid4()))
# 1. Emit initial task event (kind: "task", status: "submitted")
# Format matches A2ACompletionBridgeTransformation.create_task_event
task_event = {
task_event: Final = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
@ -451,7 +452,7 @@ class PydanticAITransformation:
# 2. Emit status update (kind: "status-update", status: "working")
# Format matches A2ACompletionBridgeTransformation.create_status_update_event
working_event = {
working_event: Final = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
@ -502,7 +503,7 @@ class PydanticAITransformation:
await asyncio.sleep(delay_ms / 1000.0)
# 4. Emit final status update (kind: "status-update", status: "completed", final: true)
completed_event = {
completed_event: Final = {
"jsonrpc": "2.0",
"id": request_id,
"result": {
@ -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)

19
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, Final
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
@ -16,12 +17,12 @@ 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")
litellm_params: Final = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for WatsonxOrchestrateA2AConfig "
@ -36,12 +37,12 @@ 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")
litellm_params: Final = kwargs.get("litellm_params")
if not litellm_params:
raise ValueError(
"litellm_params is required for WatsonxOrchestrateA2AConfig "

129
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, Final, NamedTuple, cast
import httpx
@ -20,11 +21,11 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.types.llms.custom_http import httpxSpecialProvider
_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]] = {}
_IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token"
_POLL_INTERVAL_S: Final = 2.0
_MAX_POLL_ATTEMPTS: Final = 90
_TOKEN_CACHE_TTL_BUFFER_S: Final = 60
_token_cache: Final[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,16 +51,16 @@ 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}"
material: Final = 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()
expires_at: Final = int(expiration)
wall: Final = now_wall if now_wall is not None else time.time()
return max(expires_at - int(wall), 0)
@staticmethod
@ -67,12 +68,12 @@ 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()
cached = _token_cache.get(cache_key)
cache_key: Final = WatsonxOrchestrateHandler._token_cache_key(auth_mode, cp4d_host, api_key, username)
now: Final = time.monotonic()
cached: Final = _token_cache.get(cache_key)
if cached and cached[1] > now:
return cached[0]
@ -95,7 +96,7 @@ class WatsonxOrchestrateHandler:
else:
if not username:
raise ValueError("'username' is required in litellm_params when auth_mode='cp4d'")
token_url = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize"
token_url: Final = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize"
response = await client.post(
token_url,
json={"username": username, "api_key": api_key},
@ -104,13 +105,13 @@ class WatsonxOrchestrateHandler:
response.raise_for_status()
payload = response.json()
token = str(payload["token"])
expiration = payload.get("expiration")
expiration: Final = payload.get("expiration")
if expiration is None:
ttl_s = 3600
else:
ttl_s = WatsonxOrchestrateHandler._cp4d_token_ttl_seconds(expiration)
expires_at = now + max(ttl_s - _TOKEN_CACHE_TTL_BUFFER_S, 0)
expires_at: Final = now + max(ttl_s - _TOKEN_CACHE_TTL_BUFFER_S, 0)
_token_cache[cache_key] = (token, expires_at)
for stale_key, (_, stale_expires_at) in list(_token_cache.items()):
if stale_expires_at <= now:
@ -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]:
url = f"{base_url}/v1/orchestrate/runs/{run_id}"
) -> dict[str, Any]:
url: Final = 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,14 +145,14 @@ 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 ""
run_id: Final = run_data.get("run_id") or run_data.get("id") or ""
if not run_id:
raise ValueError(f"WXO: No run_id in response: {run_data}")
run_data = await WatsonxOrchestrateHandler._poll_run(
@ -186,11 +187,11 @@ class WatsonxOrchestrateHandler:
return accumulated_text
@staticmethod
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 ""
api_key = litellm_params.get("api_key") or ""
def _extract_litellm_params(litellm_params: dict[str, Any]) -> WXORequestParams:
cp4d_host: Final = litellm_params.get("cp4d_host") or ""
instance_id: Final = litellm_params.get("instance_id") or ""
wxo_agent_id: Final = litellm_params.get("wxo_agent_id") or ""
api_key: Final = litellm_params.get("api_key") or ""
if not cp4d_host:
raise ValueError("'cp4d_host' is required in litellm_params for WXO agents")
@ -214,38 +215,38 @@ class WatsonxOrchestrateHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: Dict[str, Any],
litellm_params: Dict[str, Any],
) -> Dict[str, Any]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
params: dict[str, Any],
litellm_params: dict[str, Any],
) -> dict[str, Any]:
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=90.0)
token = await WatsonxOrchestrateHandler._get_bearer_token(
client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0)
token: Final = await WatsonxOrchestrateHandler._get_bearer_token(
cp4d_host=wxo.cp4d_host,
auth_mode=wxo.auth_mode,
api_key=wxo.api_key,
username=wxo.username,
client=client,
)
base_url = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id)
auth_headers = {
base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id)
auth_headers: Final = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"Accept": "application/json",
}
text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body = WatsonxOrchestrateTransformation.build_wxo_run_body(
text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body(
wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id
)
run_response = await client.post(
run_response: Final = await client.post(
f"{base_url}/v1/orchestrate/runs",
json=body,
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,
@ -254,40 +255,40 @@ class WatsonxOrchestrateHandler:
client=client,
)
response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data)
response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data)
return WatsonxOrchestrateTransformation.build_a2a_message_response(request_id=request_id, text=response_text)
@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]]:
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
) -> AsyncIterator[dict[str, Any]]:
wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
client = WatsonxOrchestrateHandler._http_client(timeout=120.0)
token = await WatsonxOrchestrateHandler._get_bearer_token(
client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0)
token: Final = await WatsonxOrchestrateHandler._get_bearer_token(
cp4d_host=wxo.cp4d_host,
auth_mode=wxo.auth_mode,
api_key=wxo.api_key,
username=wxo.username,
client=client,
)
base_url = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id)
auth_headers = {
base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id)
auth_headers: Final = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
"Accept": "text/event-stream, application/json",
}
text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body = WatsonxOrchestrateTransformation.build_wxo_run_body(
text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params)
body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body(
wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id
)
try:
response = await client.post(
response: Final = await client.post(
f"{base_url}/v1/orchestrate/runs/stream",
json=body,
headers=auth_headers,
@ -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(
@ -305,7 +306,7 @@ class WatsonxOrchestrateHandler:
params=params,
litellm_params=litellm_params,
)
response_text = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result)
response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result)
async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text(
text=response_text,
request_id=request_id,
@ -315,9 +316,9 @@ class WatsonxOrchestrateHandler:
yield chunk
return
content_type = response.headers.get("content-type", "").lower()
content_type: Final = response.headers.get("content-type", "").lower()
if "text/event-stream" not in content_type:
response_body = await response.aread()
response_body: Final = await response.aread()
result = json.loads(response_body)
result = await WatsonxOrchestrateHandler._get_successful_run_data(
run_data=result,

39
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, Final
from uuid import uuid4
from litellm._logging import verbose_logger
@ -28,15 +29,15 @@ 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.
A2A format: params.message.parts[*] where part.kind == "text"
"""
message = params.get("message", {})
parts = message.get("parts", [])
texts = []
message: Final = params.get("message", {})
parts: Final = message.get("parts", [])
texts: Final = []
for part in parts:
if not isinstance(part, dict):
continue
@ -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: Final[dict[str, Any]] = {
"agent_id": wxo_agent_id,
"message": {
"role": "user",
@ -95,19 +96,19 @@ class WatsonxOrchestrateTransformation:
pass
# Tertiary: results as a raw string
results = result.get("results")
results: Final = result.get("results")
if results and isinstance(results, str):
return results
return ""
@staticmethod
def extract_text_from_a2a_message_response(a2a_response: Dict[str, Any]) -> str:
result = a2a_response.get("result")
def extract_text_from_a2a_message_response(a2a_response: dict[str, Any]) -> str:
result: Final = a2a_response.get("result")
if not isinstance(result, dict):
verbose_logger.warning("WXO: A2A response missing result object")
return ""
parts = result.get("parts")
parts: Final = result.get("parts")
if not isinstance(parts, list):
verbose_logger.warning("WXO: A2A result has no parts list")
return ""
@ -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.
@ -149,9 +150,9 @@ class WatsonxOrchestrateTransformation:
3. artifact-update chunks
4. status-update (kind="status-update", state="completed", final=True)
"""
task_id = str(uuid4())
context_id = str(uuid4())
artifact_id = str(uuid4())
task_id: Final = str(uuid4())
context_id: Final = str(uuid4())
artifact_id: Final = str(uuid4())
# 1. Task submitted
yield {
@ -180,7 +181,7 @@ class WatsonxOrchestrateTransformation:
await asyncio.sleep(delay_ms / 1000.0)
# 3. Artifact chunks (always emit at least one chunk, even for empty text)
text_to_chunk = text or ""
text_to_chunk: Final = text or ""
for i in range(0, max(len(text_to_chunk), 1), chunk_size):
chunk_text = text_to_chunk[i : i + chunk_size]
is_last = (i + chunk_size) >= max(len(text_to_chunk), 1)
@ -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)

57
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, Final
import litellm
from litellm._logging import verbose_logger
@ -37,16 +38,16 @@ 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
async def __anext__(self) -> "SendStreamingMessageResponse":
try:
chunk = await self.stream.__anext__()
chunk: Final = await self.stream.__anext__()
# Store chunk
self.chunks.append(chunk)
@ -70,8 +71,8 @@ class A2AStreamingIterator:
def _collect_text_from_chunk(self, chunk: Any) -> None:
"""Extract text from a streaming chunk and add to collected parts."""
try:
chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
text = A2ARequestUtils.extract_text_from_response(chunk_dict)
chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
text: Final = A2ARequestUtils.extract_text_from_response(chunk_dict)
if text:
self.collected_text_parts.append(text)
except Exception:
@ -80,10 +81,10 @@ class A2AStreamingIterator:
def _is_completed_chunk(self, chunk: Any) -> bool:
"""Check if chunk indicates stream completion."""
try:
chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
result = chunk_dict.get("result", {})
chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
result: Final = chunk_dict.get("result", {})
if isinstance(result, dict):
status = result.get("status", {})
status: Final = result.get("status", {})
if isinstance(status, dict):
return status.get("state") == "completed"
except Exception:
@ -93,21 +94,21 @@ class A2AStreamingIterator:
async def _handle_stream_complete(self) -> None:
"""Handle logging and token counting when stream completes."""
try:
end_time = datetime.now()
end_time: Final = datetime.now()
# Calculate tokens from collected text
input_message = A2ARequestUtils.get_input_message_from_request(self.request)
input_text = A2ARequestUtils.extract_text_from_message(input_message)
prompt_tokens = A2ARequestUtils.count_tokens(input_text)
input_message: Final = A2ARequestUtils.get_input_message_from_request(self.request)
input_text: Final = A2ARequestUtils.extract_text_from_message(input_message)
prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text)
# Use the last (most complete) text from chunks
output_text = self.collected_text_parts[-1] if self.collected_text_parts else ""
completion_tokens = A2ARequestUtils.count_tokens(output_text)
output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else ""
completion_tokens: Final = A2ARequestUtils.count_tokens(output_text)
total_tokens = prompt_tokens + completion_tokens
total_tokens: Final = prompt_tokens + completion_tokens
# Create usage object
usage = litellm.Usage(
usage: Final = litellm.Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
@ -119,11 +120,11 @@ class A2AStreamingIterator:
self.logging_obj.model_call_details["stream"] = False
# Calculate cost using A2ACostCalculator
response_cost = A2ACostCalculator.calculate_a2a_cost(self.logging_obj)
response_cost: Final = A2ACostCalculator.calculate_a2a_cost(self.logging_obj)
self.logging_obj.model_call_details["response_cost"] = response_cost
# Build result for logging
result = self._build_logging_result(usage)
result: Final = self._build_logging_result(usage)
# Call success handlers - they will build standard_logging_object
asyncio.create_task(
@ -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: Final[dict[str, Any]] = {
"id": getattr(self.request, "id", "unknown"),
"jsonrpc": "2.0",
"usage": (usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)),
@ -156,7 +159,7 @@ class A2AStreamingIterator:
# Add final chunk result if available
if self.final_chunk:
try:
chunk_dict = self.final_chunk.model_dump(mode="json", exclude_none=True)
chunk_dict: Final = self.final_chunk.model_dump(mode="json", exclude_none=True)
result["result"] = chunk_dict.get("result", {})
except Exception:
pass

32
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, Final
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: Final[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.
@ -56,7 +56,7 @@ class A2ARequestUtils:
Returns:
Text from response message parts
"""
result = response_dict.get("result", {})
result: Final = response_dict.get("result", {})
if not isinstance(result, dict):
return ""
@ -66,12 +66,12 @@ class A2ARequestUtils:
if result.get("kind") == "message":
return A2ARequestUtils.extract_text_from_message(result)
message = result.get("message", {})
message: Final = result.get("message", {})
return A2ARequestUtils.extract_text_from_message(message)
@staticmethod
def get_input_message_from_request(
request: "Union[SendMessageRequest, SendStreamingMessageRequest]",
request: "SendMessageRequest | SendStreamingMessageRequest",
) -> Any:
"""
Extract the input message from an A2A request.
@ -82,7 +82,7 @@ class A2ARequestUtils:
Returns:
The message object/dict or None
"""
params = getattr(request, "params", None)
params: Final = getattr(request, "params", None)
if params is None:
return None
return getattr(params, "message", None)
@ -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.
@ -128,14 +128,14 @@ class A2ARequestUtils:
input_message = A2ARequestUtils.get_input_message_from_request(request)
if input_message is not None and hasattr(input_message, "model_dump"):
input_message = input_message.model_dump(mode="json")
input_text = A2ARequestUtils.extract_text_from_message(input_message)
prompt_tokens = A2ARequestUtils.count_tokens(input_text)
input_text: Final = A2ARequestUtils.extract_text_from_message(input_message)
prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text)
# Count output tokens
output_text = A2ARequestUtils.extract_text_from_response(response_dict)
completion_tokens = A2ARequestUtils.count_tokens(output_text)
output_text: Final = A2ARequestUtils.extract_text_from_response(response_dict)
completion_tokens: Final = A2ARequestUtils.count_tokens(output_text)
total_tokens = prompt_tokens + completion_tokens
total_tokens: Final = prompt_tokens + completion_tokens
return prompt_tokens, completion_tokens, total_tokens
@ -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)

68
litellm/anthropic_beta_headers_manager.py
View file

@ -25,14 +25,14 @@ Environment Variables:
import json
import os
from importlib.resources import files
from typing import Dict, List, Optional, Set
from typing import Final
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,15 +44,15 @@ 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(
content: Final = json.loads(
files("litellm").joinpath("anthropic_beta_headers_config.json").read_text(encoding="utf-8")
)
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": {},
@ -80,14 +80,14 @@ class GetAnthropicBetaHeadersConfig:
return False
# Check for at least one provider key
provider_keys = [
provider_keys: Final = [
"anthropic",
"azure_ai",
"bedrock",
"bedrock_converse",
"vertex_ai",
]
has_provider = any(key in fetched_config for key in provider_keys)
has_provider: Final = any(key in fetched_config for key in provider_keys)
if not has_provider:
verbose_logger.warning(
@ -114,7 +114,7 @@ class GetAnthropicBetaHeadersConfig:
Returns the parsed JSON dict. Raises on network/parse errors
(caller is expected to handle).
"""
response = httpx.get(url, timeout=timeout)
response: Final = httpx.get(url, timeout=timeout)
response.raise_for_status()
return response.json()
@ -139,7 +139,7 @@ def get_beta_headers_config(url: str) -> dict:
return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()
try:
content = GetAnthropicBetaHeadersConfig.fetch_remote_beta_headers_config(url)
content: Final = GetAnthropicBetaHeadersConfig.fetch_remote_beta_headers_config(url)
except Exception as e:
verbose_logger.warning(
"LiteLLM: Failed to fetch remote beta headers config from %s: %s. Falling back to local backup.",
@ -159,7 +159,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 +183,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.
@ -207,15 +207,15 @@ def get_provider_name(provider: str) -> str:
Returns:
Canonical provider name
"""
config = _load_beta_headers_config()
aliases = config.get("provider_aliases", {})
config: Final = _load_beta_headers_config()
aliases: Final = config.get("provider_aliases", {})
return aliases.get(provider, provider)
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.
@ -234,20 +234,22 @@ def filter_and_transform_beta_headers(
if not beta_headers:
return []
config = _load_beta_headers_config()
config: Final = _load_beta_headers_config()
provider = get_provider_name(provider)
# Get the header mapping for this provider
provider_mapping = config.get(provider, {})
provider_mapping: Final = config.get(provider, {})
filtered_headers: Set[str] = set()
filtered_headers: Final[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 +257,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
@ -278,9 +280,9 @@ def is_beta_header_supported(
Returns:
True if the header is in the mapping with a non-null value, False otherwise
"""
config = _load_beta_headers_config()
config: Final = _load_beta_headers_config()
provider = get_provider_name(provider)
provider_mapping = config.get(provider, {})
provider_mapping: Final = config.get(provider, {})
# Header is supported if it's in the mapping and has a non-null value
return beta_header in provider_mapping and provider_mapping[beta_header] is not None
@ -289,7 +291,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.
@ -302,11 +304,11 @@ def get_provider_beta_header(
Returns:
The provider-specific header name if supported, or None if unsupported/unknown
"""
config = _load_beta_headers_config()
config: Final = _load_beta_headers_config()
provider = get_provider_name(provider)
# Get the header mapping for this provider
provider_mapping = config.get(provider, {})
provider_mapping: Final = config.get(provider, {})
# Check if header is in the mapping
if anthropic_beta_header not in provider_mapping:
@ -331,15 +333,15 @@ def update_headers_with_filtered_beta(
Returns:
Updated headers dict
"""
existing_beta = headers.get("anthropic-beta")
existing_beta: Final = headers.get("anthropic-beta")
if not existing_beta:
return headers
# Parse existing beta headers
beta_values = [b.strip() for b in existing_beta.split(",") if b.strip()]
beta_values: Final = [b.strip() for b in existing_beta.split(",") if b.strip()]
# Filter and transform based on provider
filtered_beta_values = filter_and_transform_beta_headers(
filtered_beta_values: Final = filter_and_transform_beta_headers(
beta_headers=beta_values,
provider=provider,
)
@ -373,11 +375,11 @@ def update_request_with_filtered_beta(
"""
headers = update_headers_with_filtered_beta(headers=headers, provider=provider)
existing_body_betas = request_data.get("anthropic_beta")
existing_body_betas: Final = request_data.get("anthropic_beta")
if not existing_body_betas:
return headers, request_data
filtered_body_betas = filter_and_transform_beta_headers(
filtered_body_betas: Final = filter_and_transform_beta_headers(
beta_headers=existing_body_betas,
provider=provider,
)
@ -390,7 +392,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).
@ -400,9 +402,9 @@ def get_unsupported_headers(provider: str) -> List[str]:
Returns:
List of unsupported Anthropic beta header names
"""
config = _load_beta_headers_config()
config: Final = _load_beta_headers_config()
provider = get_provider_name(provider)
provider_mapping = config.get(provider, {})
provider_mapping: Final = config.get(provider, {})
# Return headers with null values
return [header for header, value in provider_mapping.items() if value is None]

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

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

@ -4,14 +4,15 @@ Utilities for mapping exceptions to Anthropic error format.
Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthropic response format.
"""
from typing import Final
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: Final[dict[int, AnthropicErrorType]] = {
400: "invalid_request_error",
401: "authentication_error",
403: "permission_error",
@ -39,7 +40,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.
@ -51,9 +52,9 @@ class AnthropicExceptionMapping:
"request_id": "req_..."
}
"""
error_type = AnthropicExceptionMapping.get_error_type(status_code)
error_type: Final = AnthropicExceptionMapping.get_error_type(status_code)
response: AnthropicErrorResponse = {
response: Final[AnthropicErrorResponse] = {
"type": "error",
"error": {
"type": error_type,
@ -77,7 +78,7 @@ class AnthropicExceptionMapping:
- Generic: {"message": "..."}
- Plain strings
"""
parsed = safe_json_loads(raw_message)
parsed: Final = safe_json_loads(raw_message)
if isinstance(parsed, dict):
# Bedrock format
if "detail" in parsed and isinstance(parsed["detail"], dict):
@ -124,7 +125,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 +144,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

395
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, Final, Literal
import httpx
from openai import AsyncOpenAI, OpenAI
@ -28,34 +29,34 @@ from ..types.router import *
from .utils import get_optional_params_add_message
####### ENVIRONMENT VARIABLES ###################
openai_assistants_api = OpenAIAssistantsAPI()
azure_assistants_api = AzureAssistantsAPI()
openai_assistants_api: Final = OpenAIAssistantsAPI()
azure_assistants_api: Final = AzureAssistantsAPI()
### ASSISTANTS ###
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()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["aget_assistants"] = True
try:
# Use a partial function to pass your keyword arguments
func = partial(get_assistants, custom_llm_provider, client, **kwargs)
func: Final = partial(get_assistants, custom_llm_provider, client, **kwargs)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -73,17 +74,17 @@ 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: Final[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)
litellm_params_dict = get_litellm_params(**kwargs)
optional_params: Final = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
@ -94,14 +95,14 @@ def get_assistants(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
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
@ -110,7 +111,7 @@ def get_assistants(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -146,8 +147,8 @@ def get_assistants(
or get_secret("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
extra_body: Final = optional_params.get("extra_body", {})
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,28 +194,28 @@ 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()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["async_create_assistants"] = True
model = kwargs.pop("model", None)
model: Final = kwargs.pop("model", None)
try:
kwargs["client"] = client
# Use a partial function to pass your keyword arguments
func = partial(create_assistants, custom_llm_provider, model, **kwargs)
func: Final = partial(create_assistants, custom_llm_provider, model, **kwargs)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model=model, custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -237,26 +234,26 @@ 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: Final[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)
litellm_params_dict = get_litellm_params(**kwargs)
optional_params: Final = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
@ -267,7 +264,7 @@ def create_assistants(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
@ -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
@ -299,7 +296,7 @@ def create_assistants(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -336,8 +333,8 @@ def create_assistants(
or get_secret("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
extra_body: Final = optional_params.get("extra_body", {})
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,27 +377,27 @@ 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()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["async_delete_assistants"] = True
try:
kwargs["client"] = client
# Use a partial function to pass your keyword arguments
func = partial(delete_assistant, custom_llm_provider, **kwargs)
func: Final = partial(delete_assistant, custom_llm_provider, **kwargs)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -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]]:
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
) -> AssistantDeleted | Coroutine[Any, Any, AssistantDeleted]:
optional_params: Final = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
async_delete_assistants: Optional[bool] = kwargs.pop("async_delete_assistants", None)
async_delete_assistants: Final[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")
@ -444,14 +439,14 @@ def delete_assistant(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
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
@ -460,7 +455,7 @@ def delete_assistant(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None
)
# set API KEY
@ -489,8 +484,8 @@ def delete_assistant(
or get_secret("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
extra_body: Final = optional_params.get("extra_body", {})
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(
@ -537,23 +530,23 @@ def delete_assistant(
async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwargs) -> Thread:
loop = asyncio.get_event_loop()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["acreate_thread"] = True
try:
# Use a partial function to pass your keyword arguments
func = partial(create_thread, custom_llm_provider, **kwargs)
func: Final = partial(create_thread, custom_llm_provider, **kwargs)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -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:
"""
@ -599,9 +592,9 @@ def create_thread(
)
```
"""
acreate_thread = kwargs.get("acreate_thread", None)
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
acreate_thread: Final = kwargs.get("acreate_thread", None)
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
@ -612,17 +605,17 @@ def create_thread(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
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
@ -631,7 +624,7 @@ def create_thread(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -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
extra_body: Final = optional_params.get("extra_body", {})
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,26 +701,26 @@ 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()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["aget_thread"] = True
try:
# Use a partial function to pass your keyword arguments
func = partial(get_thread, custom_llm_provider, thread_id, client, **kwargs)
func: Final = partial(get_thread, custom_llm_provider, thread_id, client, **kwargs)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -754,9 +743,9 @@ def get_thread(
**kwargs,
) -> Thread:
"""Get the thread object, given a thread_id"""
aget_thread = kwargs.pop("aget_thread", None)
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
aget_thread: Final = kwargs.pop("aget_thread", None)
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
# set timeout for 10 minutes by default
@ -766,15 +755,15 @@ def get_thread(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
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
@ -783,7 +772,7 @@ def get_thread(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -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
@ -822,8 +809,8 @@ def get_thread(
or get_secret("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
extra_body: Final = optional_params.get("extra_body", {})
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,17 +853,17 @@ 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:
loop = asyncio.get_event_loop()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["a_add_message"] = True
try:
# Use a partial function to pass your keyword arguments
func = partial(
func: Final = partial(
add_message,
custom_llm_provider,
thread_id,
@ -891,15 +876,15 @@ async def a_add_message(
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -921,18 +906,18 @@ 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:
### COMMON OBJECTS ###
a_add_message = kwargs.pop("a_add_message", None)
_message_data = MessageData(role=role, content=content, attachments=attachments, metadata=metadata)
litellm_params_dict = get_litellm_params(**kwargs)
optional_params = GenericLiteLLMParams(**kwargs)
a_add_message: Final = kwargs.pop("a_add_message", None)
_message_data: Final = MessageData(role=role, content=content, attachments=attachments, metadata=metadata)
litellm_params_dict: Final = get_litellm_params(**kwargs)
optional_params: Final = GenericLiteLLMParams(**kwargs)
message_data = get_optional_params_add_message(
message_data: Final = get_optional_params_add_message(
role=_message_data["role"],
content=_message_data["content"],
attachments=_message_data["attachments"],
@ -949,15 +934,15 @@ def add_message(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
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
@ -966,7 +951,7 @@ def add_message(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -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
@ -1005,8 +988,8 @@ def add_message(
or get_secret("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
extra_body: Final = optional_params.get("extra_body", {})
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,15 +1026,15 @@ 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()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["aget_messages"] = True
try:
# Use a partial function to pass your keyword arguments
func = partial(
func: Final = partial(
get_messages,
custom_llm_provider,
thread_id,
@ -1062,15 +1043,15 @@ async def aget_messages(
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -1090,12 +1071,12 @@ 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)
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
aget_messages: Final = kwargs.pop("aget_messages", None)
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
@ -1106,16 +1087,16 @@ def get_messages(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
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
@ -1124,7 +1105,7 @@ def get_messages(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -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
@ -1162,8 +1141,8 @@ def get_messages(
or get_secret("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
azure_ad_token: Optional[str] = None
extra_body: Final = optional_params.get("extra_body", {})
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,21 +1189,21 @@ 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()
loop: Final = asyncio.get_event_loop()
### PASS ARGS TO GET ASSISTANTS ###
kwargs["arun_thread"] = True
try:
# Use a partial function to pass your keyword arguments
func = partial(
func: Final = partial(
run_thread,
custom_llm_provider,
thread_id,
@ -1242,15 +1219,15 @@ async def arun_thread(
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
_, custom_llm_provider, _, _ = get_llm_provider( # type: ignore
model="", custom_llm_provider=custom_llm_provider
) # type: ignore
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -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,20 +1256,20 @@ 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."""
arun_thread = kwargs.pop("arun_thread", None)
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
arun_thread: Final = kwargs.pop("arun_thread", None)
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
@ -1303,14 +1280,14 @@ def run_thread(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
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
@ -1319,7 +1296,7 @@ def run_thread(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -1364,7 +1341,7 @@ def run_thread(
or get_secret("AZURE_API_KEY")
) # type: ignore
extra_body = optional_params.get("extra_body", {})
extra_body: Final = optional_params.get("extra_body", {})
azure_ad_token = None
if extra_body is not None:
azure_ad_token = extra_body.pop("azure_ad_token", None)
@ -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(

61
litellm/assistants/utils.py
View file

@ -1,4 +1,4 @@
from typing import Optional, Union
from typing import Final
import litellm
@ -7,21 +7,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,
):
@ -30,13 +19,13 @@ def get_optional_params_add_message(
Reference - https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-messages?tabs=python#create-message
"""
passed_params = locals()
passed_params: Final = locals()
custom_llm_provider = passed_params.pop("custom_llm_provider")
special_params = passed_params.pop("kwargs")
special_params: Final = passed_params.pop("kwargs")
for k, v in special_params.items():
passed_params[k] = v
default_params = {
default_params: Final = {
"role": None,
"content": None,
"attachments": None,
@ -49,51 +38,49 @@ def get_optional_params_add_message(
## raise exception if non-default value passed for non-openai/azure embedding calls
def _check_valid_arg(supported_params):
if len(non_default_params.keys()) > 0:
keys = list(non_default_params.keys())
keys: Final = list(non_default_params.keys())
for k in keys:
if litellm.drop_params is True and k not in supported_params: # drop the unsupported non-default values
non_default_params.pop(k, None)
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
if custom_llm_provider == "openai":
optional_params = non_default_params
elif custom_llm_provider == "azure":
supported_params = litellm.AzureOpenAIAssistantsAPIConfig().get_supported_openai_create_message_params()
supported_params: Final = litellm.AzureOpenAIAssistantsAPIConfig().get_supported_openai_create_message_params()
_check_valid_arg(supported_params=supported_params)
optional_params = litellm.AzureOpenAIAssistantsAPIConfig().map_openai_params_create_message_params(
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
passed_params = locals()
passed_params: Final = locals()
custom_llm_provider = passed_params.pop("custom_llm_provider")
special_params = passed_params.pop("kwargs")
special_params: Final = passed_params.pop("kwargs")
for k, v in special_params.items():
passed_params[k] = v
default_params = {
default_params: Final = {
"n": None,
"quality": None,
"response_format": None,
@ -108,7 +95,7 @@ def get_optional_params_image_gen(
## raise exception if non-default value passed for non-openai/azure embedding calls
def _check_valid_arg(supported_params):
if len(non_default_params.keys()) > 0:
keys = list(non_default_params.keys())
keys: Final = list(non_default_params.keys())
for k in keys:
if litellm.drop_params is True and k not in supported_params: # drop the unsupported non-default values
non_default_params.pop(k, None)
@ -142,6 +129,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

54
litellm/batch_completion/main.py
View file

@ -1,5 +1,5 @@
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from typing import List, Optional
from typing import Final
import litellm
from litellm._logging import print_verbose
@ -11,23 +11,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,
):
@ -56,17 +56,17 @@ def batch_completion(
Returns:
list: A list of completion results.
"""
args = locals()
args: Final = locals()
batch_messages = messages
completions = []
batch_messages: Final = messages
completions: Final = []
model = model
custom_llm_provider = None
if model.split("/", 1)[0] in litellm.provider_list:
custom_llm_provider = model.split("/", 1)[0]
model = model.split("/", 1)[1]
if custom_llm_provider == "vllm":
optional_params = get_optional_params(
optional_params: Final = get_optional_params(
functions=functions,
function_call=function_call,
temperature=temperature,
@ -146,7 +146,7 @@ def batch_completion_models(*args, **kwargs):
if "model" in kwargs:
kwargs.pop("model")
if "models" in kwargs:
models = kwargs["models"]
models: Final = kwargs["models"]
kwargs.pop("models")
futures = {}
with ThreadPoolExecutor(max_workers=len(models)) as executor:
@ -157,14 +157,14 @@ def batch_completion_models(*args, **kwargs):
if future.result() is not None:
return future.result()
elif "deployments" in kwargs:
deployments = kwargs["deployments"]
deployments: Final = kwargs["deployments"]
kwargs.pop("deployments")
kwargs.pop("model_list")
nested_kwargs = kwargs.pop("kwargs", {})
nested_kwargs: Final = kwargs.pop("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}
@ -239,10 +239,10 @@ def batch_completion_models_all_responses(*args, **kwargs):
if len(models) == 0:
return []
responses = []
responses: Final = []
with concurrent.futures.ThreadPoolExecutor(max_workers=len(models)) as executor:
futures = [executor.submit(litellm.completion, *args, model=model, **kwargs) for model in models]
futures: Final = [executor.submit(litellm.completion, *args, model=model, **kwargs) for model in models]
for future in futures:
try:
@ -250,7 +250,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

97
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, Final, 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,14 +136,14 @@ 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))
line_stats: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
cache_token_params = {
cache_token_params: Final = {
key: tokens
for key, tokens in (
("cache_read_input_tokens", sum(stats.cache_read_tokens for stats in line_stats)),
@ -150,22 +151,22 @@ def _aggregate_batch_cost_usage_models(
)
if tokens > 0
}
batch_usage = Usage(
batch_usage: Final = Usage(
total_tokens=sum(stats.total_tokens for stats in line_stats),
prompt_tokens=sum(stats.prompt_tokens for stats in line_stats),
completion_tokens=sum(stats.completion_tokens for stats in line_stats),
**cache_token_params,
)
batch_models = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
total_cost = sum((stats.cost for stats in line_stats), 0.0)
batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
total_cost: Final = sum((stats.cost for stats in line_stats), 0.0)
verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models)
return total_cost, batch_usage, batch_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.
@ -183,7 +184,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
total_tokens = 0
prompt_tokens = 0
completion_tokens = 0
actual_model_name = model_name or "gemini-2.0-flash-001"
actual_model_name: Final = model_name or "gemini-2.0-flash-001"
for response in vertex_ai_batch_responses:
response_body = response.get("response")
@ -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
@ -253,31 +254,31 @@ async def _fetch_batch_output_file_content(
raise ValueError("Output file id is None cannot retrieve file content")
file_id = batch.output_file_id
is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id)
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
file_content_kwargs = {
file_content_kwargs: Final = {
"file_id": file_id,
"custom_llm_provider": custom_llm_provider,
}
# Extract and add credentials for file access
credentials = _extract_file_access_credentials(litellm_params)
credentials: Final = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
_file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
_file_content: Final = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
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.
@ -290,11 +291,11 @@ def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
Returns:
Dictionary containing only the credentials needed for file access
"""
credentials = {}
credentials: Final = {}
if litellm_params:
# List of credential keys that should be passed to file operations
credential_keys = [
credential_keys: Final = [
"api_key",
"api_base",
"api_version",
@ -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
"""
@ -354,7 +355,7 @@ def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]:
# A batch request's input tokens scale roughly with its serialized size, so this
# is a conservative per-row fallback when the token counter cannot measure a row.
_BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN = 4
_BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN: Final = 4
def _estimate_batch_entry_tokens(raw_line: bytes) -> int:
@ -366,21 +367,21 @@ 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 {}
model = body.get("model", model_name or "")
body: Final = entry.get("body", {}) or {}
model: Final = body.get("model", model_name or "")
messages = body.get("messages")
messages: Final = body.get("messages")
if messages:
return token_counter(model=model, messages=messages)
prompt = body.get("prompt")
prompt: Final = body.get("prompt")
if prompt:
return _count_prompt_or_input_tokens(model=model, value=prompt)
input_data = body.get("input")
input_data: Final = body.get("input")
if input_data:
return _count_prompt_or_input_tokens(model=model, value=input_data)
@ -431,8 +432,12 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov
usage_object=response_body.get("usage", None) or {},
reasoning_content=None,
)
_usage_dict = response_body.get("usage", None) or {}
usage: Usage = Usage(**_usage_dict)
from litellm.responses.utils import ResponseAPILoggingUtils
_usage_dict: Final = response_body.get("usage", None) or {}
if ResponseAPILoggingUtils._is_response_api_usage(_usage_dict):
return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_usage_dict)
usage: Final[Usage] = Usage(**_usage_dict)
return usage
@ -454,8 +459,8 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("message", None) or {}
if custom_llm_provider == "bedrock":
return batch_job_output_file.get("modelOutput", None) or {}
_response: dict = batch_job_output_file.get("response", None) or {}
_response_body = _response.get("body", None) or {}
_response: Final[dict] = batch_job_output_file.get("response", None) or {}
_response_body: Final = _response.get("body", None) or {}
return _response_body
@ -471,5 +476,5 @@ def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provi
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded"
if custom_llm_provider == "bedrock":
return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None
_response: dict = batch_job_output_file.get("response", None) or {}
_response: Final[dict] = batch_job_output_file.get("response", None) or {}
return _response.get("status_code", None) == 200

269
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, Final, Literal, cast
import httpx
from openai.types.batch import BatchRequestCounts
@ -53,17 +54,17 @@ from litellm.utils import (
)
####### ENVIRONMENT VARIABLES ###################
openai_batches_instance = OpenAIBatchesAPI()
azure_batches_instance = AzureBatchesAPI()
vertex_ai_batches_instance = VertexAIBatchPrediction(gcs_bucket_name="")
anthropic_batches_instance = AnthropicBatchesHandler()
openai_batches_instance: Final = OpenAIBatchesAPI()
azure_batches_instance: Final = AzureBatchesAPI()
vertex_ai_batches_instance: Final = VertexAIBatchPrediction(gcs_bucket_name="")
anthropic_batches_instance: Final = AnthropicBatchesHandler()
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:
@ -79,13 +80,13 @@ def _resolve_timeout(
Returns:
Resolved timeout as float
"""
timeout = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout
timeout: Final = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout
# Handle httpx.Timeout objects
if isinstance(timeout, httpx.Timeout):
if supports_httpx_timeout(custom_llm_provider) is False:
# Extract read timeout for providers that don't support httpx.Timeout
read_timeout = timeout.read or default_timeout
read_timeout: Final = timeout.read or default_timeout
return float(read_timeout)
else:
# For providers that support httpx.Timeout, we still need to return a float
@ -103,13 +104,13 @@ def _resolve_timeout(
@client
async def acreate_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
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:
"""
@ -118,11 +119,11 @@ async def acreate_batch(
LiteLLM Equivalent of POST: https://api.openai.com/v1/batches
"""
try:
loop = asyncio.get_event_loop()
loop: Final = asyncio.get_event_loop()
kwargs["acreate_batch"] = True
# Use a partial function to pass your keyword arguments
func = partial(
func: Final = partial(
create_batch,
completion_window,
endpoint,
@ -136,9 +137,9 @@ async def acreate_batch(
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
@ -153,26 +154,26 @@ async def acreate_batch(
@client
def create_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"],
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
LiteLLM Equivalent of POST: https://api.openai.com/v1/batches
"""
try:
optional_params = GenericLiteLLMParams(**kwargs)
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)
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_call_id: Final = kwargs.get("litellm_call_id", None)
proxy_server_request: Final = kwargs.get("proxy_server_request", None)
model_info: Final = kwargs.get("model_info", None)
model: str | None = kwargs.get("model", None)
try:
if model is not None:
model, _, _, _ = get_llm_provider(
@ -181,14 +182,14 @@ 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
litellm_params = dict(GenericLiteLLMParams(**kwargs))
litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None))
_is_async: Final = kwargs.pop("acreate_batch", False) is True
litellm_params: Final = dict(GenericLiteLLMParams(**kwargs))
litellm_logging_obj: Final[LiteLLMLoggingObj] = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None))
### TIMEOUT LOGIC ###
timeout = _resolve_timeout(optional_params, kwargs, custom_llm_provider)
timeout: Final = _resolve_timeout(optional_params, kwargs, custom_llm_provider)
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
model=model,
@ -205,7 +206,7 @@ def create_batch(
custom_llm_provider=custom_llm_provider,
)
_create_batch_request = CreateBatchRequest(
_create_batch_request: Final = CreateBatchRequest(
completion_window=completion_window,
endpoint=endpoint,
input_file_id=input_file_id,
@ -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 = (
@ -247,7 +248,7 @@ def create_batch(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -300,13 +301,13 @@ def create_batch(
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or ""
vertex_ai_project = (
vertex_ai_project: Final = (
optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
vertex_ai_location: Final = (
optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
response = vertex_ai_batches_instance.create_batch(
_is_async=_is_async,
@ -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:
"""
@ -349,11 +350,11 @@ async def aretrieve_batch(
LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id}
"""
try:
loop = asyncio.get_event_loop()
loop: Final = asyncio.get_event_loop()
kwargs["aretrieve_batch"] = True
# Use a partial function to pass your keyword arguments
func = partial(
func: Final = partial(
retrieve_batch,
batch_id,
custom_llm_provider,
@ -363,9 +364,9 @@ async def aretrieve_batch(
**kwargs,
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -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 = (
@ -396,7 +397,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -421,7 +422,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
)
elif custom_llm_provider == "azure":
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE")
api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
api_version: Final = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION")
api_key = (
optional_params.api_key
@ -431,7 +432,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
or get_secret_str("AZURE_API_KEY")
)
extra_body = optional_params.get("extra_body", {})
extra_body: Final = optional_params.get("extra_body", {})
if extra_body is not None:
extra_body.pop("azure_ad_token", None)
else:
@ -449,13 +450,13 @@ def _handle_retrieve_batch_providers_without_provider_config(
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or ""
vertex_ai_project = (
vertex_ai_project: Final = (
optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
vertex_ai_location: Final = (
optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
response = vertex_ai_batches_instance.retrieve_batch(
_is_async=_is_async,
@ -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,22 +508,22 @@ 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.
LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id}
"""
try:
optional_params = GenericLiteLLMParams(**kwargs)
litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None)
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_logging_obj: Final[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(
litellm_params: Final = get_litellm_params(
custom_llm_provider=custom_llm_provider,
**kwargs,
)
@ -541,21 +542,21 @@ def retrieve_batch(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
_retrieve_batch_request = RetrieveBatchRequest(
_retrieve_batch_request: Final = RetrieveBatchRequest(
batch_id=batch_id,
extra_headers=extra_headers,
extra_body=extra_body,
)
_is_async = kwargs.pop("aretrieve_batch", False) is True
client = kwargs.get("client", None)
_is_async: Final = kwargs.pop("aretrieve_batch", False) is True
client: Final = kwargs.get("client", None)
# Bedrock has two distinct ARN families that need different APIs:
# * async-invoke ARNs (Twelve Labs Marengo embeddings) -> bedrock-runtime data plane
@ -567,7 +568,7 @@ def retrieve_batch(
if batch_id.startswith("arn:aws") and ":bedrock:" in batch_id:
if ":async-invoke/" in batch_id:
# Remove aws_region_name from kwargs to avoid duplicate parameter
async_kwargs = kwargs.copy()
async_kwargs: Final = kwargs.copy()
async_kwargs.pop("aws_region_name", None)
return BedrockBatchesHandler._handle_async_invoke_status(
@ -577,7 +578,7 @@ def retrieve_batch(
**async_kwargs,
)
if ":model-invocation-job/" in batch_id:
mij_kwargs = kwargs.copy()
mij_kwargs: Final = kwargs.copy()
mij_kwargs.pop("aws_region_name", None)
return BedrockBatchesHandler._handle_model_invocation_job_status(
@ -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: Final[str | None] = kwargs.get("model", None)
if model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
@ -598,7 +599,7 @@ def retrieve_batch(
provider_config = None
if provider_config is not None:
response = base_llm_http_handler.retrieve_batch(
response: Final = base_llm_http_handler.retrieve_batch(
batch_id=batch_id,
provider_config=provider_config,
litellm_params=litellm_params,
@ -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,
):
"""
@ -655,11 +656,11 @@ async def alist_batches(
"""
try:
loop = asyncio.get_event_loop()
loop: Final = asyncio.get_event_loop()
kwargs["alist_batches"] = True
# Use a partial function to pass your keyword arguments
func = partial(
func: Final = partial(
list_batches,
after,
limit,
@ -670,9 +671,9 @@ async def alist_batches(
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -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,
):
"""
@ -699,8 +700,8 @@ def list_batches(
"""
try:
# set API KEY
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params = get_litellm_params(
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_params: Final = get_litellm_params(
custom_llm_provider=custom_llm_provider,
**kwargs,
)
@ -719,14 +720,14 @@ def list_batches(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
_is_async = kwargs.pop("alist_batches", False) is True
_is_async: Final = kwargs.pop("alist_batches", False) is True
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 = (
@ -736,7 +737,7 @@ def list_batches(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization
or litellm.organization
or os.getenv("OPENAI_ORGANIZATION", None)
@ -782,13 +783,13 @@ def list_batches(
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or ""
vertex_ai_project = (
vertex_ai_project: Final = (
optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
vertex_ai_location: Final = (
optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
response = vertex_ai_batches_instance.list_batches(
_is_async=_is_async,
@ -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:
"""
@ -835,14 +836,14 @@ async def acancel_batch(
LiteLLM Equivalent of POST https://api.openai.com/v1/batches/{batch_id}/cancel
"""
try:
loop = asyncio.get_event_loop()
loop: Final = asyncio.get_event_loop()
kwargs["acancel_batch"] = True
# Preserve model parameter - only pop from kwargs if it exists there
# (to avoid passing it twice), otherwise keep the function parameter value
model = kwargs.pop("model", None) or model
# Use a partial function to pass your keyword arguments
func = partial(
func: Final = partial(
cancel_batch,
batch_id,
model,
@ -853,9 +854,9 @@ async def acancel_batch(
**kwargs,
)
# Add the context to the function
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
ctx: Final = contextvars.copy_context()
func_with_context: Final = partial(ctx.run, func)
init_response: Final = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
@ -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,10 +890,10 @@ 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(
optional_params: Final = GenericLiteLLMParams(**kwargs)
litellm_params: Final = get_litellm_params(
custom_llm_provider=custom_llm_provider,
**kwargs,
)
@ -905,21 +906,21 @@ def cancel_batch(
and isinstance(timeout, httpx.Timeout)
and supports_httpx_timeout(custom_llm_provider) is False
):
read_timeout = timeout.read or 600
read_timeout: Final = timeout.read or 600
timeout = read_timeout # default 10 min timeout
elif timeout is not None and not isinstance(timeout, httpx.Timeout):
timeout = float(timeout) # type: ignore
elif timeout is None:
timeout = 600.0
_cancel_batch_request = CancelBatchRequest(
_cancel_batch_request: Final = CancelBatchRequest(
batch_id=batch_id,
extra_headers=extra_headers,
extra_body=extra_body,
)
_is_async = kwargs.pop("acancel_batch", False) is True
api_base: Optional[str] = None
_is_async: Final = kwargs.pop("acancel_batch", False) is True
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
api_base = (
optional_params.api_base
@ -928,7 +929,7 @@ def cancel_batch(
or os.getenv("OPENAI_API_BASE")
or "https://api.openai.com/v1"
)
organization = (
organization: Final = (
optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None
)
api_key = optional_params.api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY")
@ -972,13 +973,13 @@ def cancel_batch(
)
elif custom_llm_provider == "vertex_ai":
api_base = optional_params.api_base or None
vertex_ai_project = (
vertex_ai_project: Final = (
optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
)
vertex_ai_location = (
vertex_ai_location: Final = (
optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION")
)
vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
response = vertex_ai_batches_instance.cancel_batch(
_is_async=_is_async,
@ -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(
@ -1026,10 +1025,10 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj
async def _async_get_status():
# Create embedding handler instance
embedding_handler = BedrockEmbedding()
embedding_handler: Final = BedrockEmbedding()
# Get the status of the async invoke job
status_response = await embedding_handler._get_async_invoke_status(
status_response: Final = await embedding_handler._get_async_invoke_status(
invocation_arn=batch_id,
aws_region_name=aws_region_name,
logging_obj=logging_obj,
@ -1041,16 +1040,16 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj
from litellm.types.utils import LiteLLMBatch
# Normalize status to lowercase (AWS returns 'Completed', 'Failed', etc.)
aws_status_raw = status_response.get("status", "")
aws_status_lower = aws_status_raw.lower()
aws_status_raw: Final = status_response.get("status", "")
aws_status_lower: Final = aws_status_raw.lower()
# Map AWS status values to LiteLLM expected values
status_mapping: dict[str, BatchJobStatus] = {
status_mapping: Final[dict[str, BatchJobStatus]] = {
"completed": "completed",
"failed": "failed",
"inprogress": "in_progress",
"in_progress": "in_progress",
}
normalized_status: BatchJobStatus = status_mapping.get(
normalized_status: Final[BatchJobStatus] = status_mapping.get(
aws_status_lower, "failed"
) # Default to "failed" if unknown status
@ -1074,7 +1073,7 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj
_,
_,
) = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw)
result = LiteLLMBatch(
result: Final = LiteLLMBatch(
id=status_response["invocationArn"],
object="batch",
status=normalized_status,
@ -1106,7 +1105,7 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj
import concurrent.futures
def run_in_thread():
new_loop = asyncio.new_event_loop()
new_loop: Final = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
try:
return new_loop.run_until_complete(_async_get_status())
@ -1114,5 +1113,5 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj
new_loop.close()
with concurrent.futures.ThreadPoolExecutor() as executor:
future = executor.submit(run_in_thread)
future: Final = executor.submit(run_in_thread)
return future.result()

40
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 Final, 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
@ -60,8 +60,8 @@ class BudgetManager:
self.print_verbose(f"user dict from local: {self.user_dict}")
elif self.client_type == "hosted":
# Load the user_dict from hosted db
url = self.api_base + "/get_budget"
data = {"project_name": self.project_name}
url: Final = self.api_base + "/get_budget"
data: Final = {"project_name": self.project_name}
response = litellm.module_level_client.post(url, headers=self.headers, json=data)
response = response.json()
if response["status"] == "error":
@ -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}
@ -100,11 +100,11 @@ class BudgetManager:
return self.user_dict[user]
def projected_cost(self, model: str, messages: list, user: str):
text = "".join(message["content"] for message in messages)
prompt_tokens = litellm.token_counter(model=model, text=text)
text: Final = "".join(message["content"] for message in messages)
prompt_tokens: Final = litellm.token_counter(model=model, text=text)
prompt_cost, _ = litellm.cost_per_token(model=model, prompt_tokens=prompt_tokens, completion_tokens=0)
current_cost = self.user_dict[user].get("current_cost", 0)
projected_cost = prompt_cost + current_cost
current_cost: Final = self.user_dict[user].get("current_cost", 0)
projected_cost: Final = prompt_cost + current_cost
return projected_cost
def get_total_budget(self, user: str):
@ -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}])
@ -178,11 +178,11 @@ class BudgetManager:
def reset_on_duration(self, user: str):
# Get current and creation time
last_updated_at = self.user_dict[user]["last_updated_at"]
current_time = time.time()
last_updated_at: Final = self.user_dict[user]["last_updated_at"]
current_time: Final = time.time()
# Convert duration from days to seconds
duration_in_seconds = self.user_dict[user]["duration"] * HOURS_IN_A_DAY * 60 * 60
duration_in_seconds: Final = self.user_dict[user]["duration"] * HOURS_IN_A_DAY * 60 * 60
# Check if duration has elapsed
if current_time - last_updated_at >= duration_in_seconds:
@ -197,7 +197,7 @@ class BudgetManager:
self.reset_on_duration(user)
def _save_data_thread(self):
thread = threading.Thread(target=self.save_data) # [Non-Blocking]: saves data without blocking execution
thread: Final = threading.Thread(target=self.save_data) # [Non-Blocking]: saves data without blocking execution
thread.start()
def save_data(self):
@ -209,8 +209,8 @@ class BudgetManager:
json.dump(self.user_dict, json_file, indent=4) # Indent for pretty formatting
return {"status": "success"}
elif self.client_type == "hosted":
url = self.api_base + "/set_budget"
data = {"project_name": self.project_name, "user_dict": self.user_dict}
url: Final = self.api_base + "/set_budget"
data: Final = {"project_name": self.project_name, "user_dict": self.user_dict}
response = litellm.module_level_client.post(url, headers=self.headers, json=data)
response = response.json()
return response

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

6
litellm/caching/_embedding_router.py
View file

@ -12,7 +12,7 @@ This module is dependency-injected: callers pass the proxy ``llm_router`` and
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, Final
if TYPE_CHECKING:
from litellm.router import Router
@ -26,7 +26,7 @@ def resolve_embedding_router(
"""Return ``llm_router`` iff it serves ``embedding_model`` as a deployment."""
if llm_router is None:
return None
router_model_names: list[str] = (
router_model_names: Final[list[str]] = (
[m["model_name"] for m in llm_model_list if "model_name" in m] if llm_model_list is not None else []
)
if embedding_model in router_model_names:
@ -38,6 +38,6 @@ def build_router_embedding_metadata(
request_metadata: dict[str, Any] | None,
) -> dict[str, Any]:
"""Forward the caller's full metadata, flagged as a semantic-cache embedding."""
metadata: dict[str, Any] = dict(request_metadata or {})
metadata: Final[dict[str, Any]] = dict(request_metadata or {})
metadata["semantic-cache-embedding"] = True
return metadata

7
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 Final, 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
@ -20,7 +21,7 @@ def lru_cache_wrapper(
return ("error", e)
def wrapped(*args, **kwargs):
result = wrapper(*args, **kwargs)
result: Final = wrapper(*args, **kwargs)
if result[0] == "error":
raise result[1]
return result[1]

33
litellm/caching/azure_blob_cache.py
View file

@ -11,6 +11,7 @@ Has 4 methods:
import asyncio
import json
from contextlib import suppress
from typing import Final
from litellm._logging import print_verbose, verbose_logger
@ -19,12 +20,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(
@ -41,7 +42,7 @@ class AzureBlobCache(BaseCache):
def set_cache(self, key, value, **kwargs) -> None:
print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}")
serialized_value = json.dumps(value)
serialized_value: Final = json.dumps(value)
try:
self.container_client.upload_blob(key, serialized_value)
except Exception as e:
@ -50,7 +51,7 @@ class AzureBlobCache(BaseCache):
async def async_set_cache(self, key, value, **kwargs) -> None:
print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}")
serialized_value = json.dumps(value)
serialized_value: Final = json.dumps(value)
try:
await self.async_container_client.upload_blob(key, serialized_value, overwrite=True)
except Exception as e:
@ -62,12 +63,15 @@ class AzureBlobCache(BaseCache):
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
as_bytes = self.container_client.download_blob(key).readall()
as_str = as_bytes.decode("utf-8")
cached_response = json.loads(as_str)
as_bytes: Final = self.container_client.download_blob(key).readall()
as_str: Final = as_bytes.decode("utf-8")
cached_response: Final = 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
@ -79,12 +83,15 @@ class AzureBlobCache(BaseCache):
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
blob = await self.async_container_client.download_blob(key)
as_bytes = await blob.readall()
as_str = as_bytes.decode("utf-8")
cached_response = json.loads(as_str)
blob: Final = await self.async_container_client.download_blob(key)
as_bytes: Final = await blob.readall()
as_str: Final = as_bytes.decode("utf-8")
cached_response: Final = 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:
@ -99,7 +106,7 @@ class AzureBlobCache(BaseCache):
await self.async_container_client.close()
async def async_set_cache_pipeline(self, cache_list, **kwargs) -> None:
tasks = []
tasks: Final = []
for val in cache_list:
tasks.append(self.async_set_cache(val[0], val[1], **kwargs))
await asyncio.gather(*tasks)

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