mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge branch 'litellm_internal_staging' into litellm_bedrock_batch_non_chat_records
This commit is contained in:
commit
7c2b709727
2217 changed files with 98271 additions and 54788 deletions
|
|
@ -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" }}
|
||||
|
|
@ -292,6 +366,7 @@ jobs:
|
|||
- checkout
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Build the wheel
|
||||
environment:
|
||||
|
|
@ -324,6 +399,7 @@ jobs:
|
|||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -397,6 +473,7 @@ jobs:
|
|||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -471,6 +548,7 @@ jobs:
|
|||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -522,6 +600,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -588,6 +667,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -628,6 +708,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -669,6 +750,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -702,6 +784,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- restore_cache:
|
||||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
|
|
@ -752,6 +835,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- restore_cache:
|
||||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
|
|
@ -803,6 +887,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -836,6 +921,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- restore_cache:
|
||||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
|
|
@ -882,6 +968,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -928,6 +1015,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -970,6 +1058,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1016,6 +1105,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1063,6 +1153,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- restore_cache:
|
||||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
|
|
@ -1103,6 +1194,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1148,6 +1240,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1192,6 +1285,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1224,6 +1318,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1267,6 +1362,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1311,6 +1407,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1355,6 +1452,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1386,6 +1484,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1432,6 +1531,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1477,6 +1577,7 @@ jobs:
|
|||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1527,6 +1628,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1551,6 +1653,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1577,6 +1680,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1678,6 +1782,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1773,6 +1878,7 @@ jobs:
|
|||
at: ~/project
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1861,6 +1967,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -1944,6 +2051,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -2076,6 +2184,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -2162,6 +2271,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -2258,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:
|
||||
|
|
@ -2283,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
|
||||
|
|
@ -2333,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: |
|
||||
|
|
@ -2414,6 +2529,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -2499,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
|
||||
|
|
@ -2553,6 +2658,7 @@ jobs:
|
|||
- skip_if_unrelated_changes
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- install_rust
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -2640,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}
|
||||
|
|
@ -2684,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}
|
||||
|
|
@ -2742,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" }}
|
||||
|
|
@ -2757,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
|
||||
|
|
@ -2772,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
|
||||
|
|
@ -2884,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" }}
|
||||
|
|
@ -2899,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: |
|
||||
|
|
@ -2909,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
|
||||
|
|
|
|||
46
.flake8
46
.flake8
|
|
@ -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
|
||||
159
.github/ci-coverage-allowlist.yml
vendored
Normal file
159
.github/ci-coverage-allowlist.yml
vendored
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
description: >-
|
||||
Paths deliberately outside CI coverage, each with the reason it is exempt.
|
||||
assert_ci_coverage.py fails when a test file or Dockerfile is neither invoked
|
||||
by a job nor listed here, so every entry below is a decision on the record.
|
||||
|
||||
test_paths:
|
||||
- reason: >-
|
||||
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
|
||||
from a pull request; it needs a live gateway and provider credentials no PR job holds
|
||||
paths:
|
||||
- tests/e2e
|
||||
- reason: >-
|
||||
The documentation and code-quality workflows execute four files in this directory by name as
|
||||
scripts and pytest never collects the directory, so these six run nowhere; listed individually
|
||||
so a seventh cannot inherit the exemption
|
||||
paths:
|
||||
- tests/documentation_tests/test_exception_types.py
|
||||
- tests/documentation_tests/test_general_setting_keys.py
|
||||
- tests/documentation_tests/test_optional_params.py
|
||||
- tests/documentation_tests/test_readme_providers.py
|
||||
- tests/documentation_tests/test_requests_lib_usage.py
|
||||
- tests/documentation_tests/test_standard_logging_payload.py
|
||||
- reason: >-
|
||||
Sibling files here are executed by name from the code-quality workflow; this one is referenced
|
||||
by no job
|
||||
paths:
|
||||
- tests/code_coverage_tests/test_aio_http_image_conversion.py
|
||||
- reason: >-
|
||||
A second mirror of the package tree living beside tests/test_litellm, which is the mirror the
|
||||
repo convention names; only test_no_hardcoded_secrets.py is invoked, from the linting
|
||||
workflow, and whether this directory should exist at all is unresolved
|
||||
paths:
|
||||
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py
|
||||
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
|
||||
- tests/litellm/integrations/helicone/test_helicone_gemini.py
|
||||
- tests/litellm/litellm_core_utils/test_json_schema_validation.py
|
||||
- tests/litellm/llms/anthropic/test_anthropic_reasoning_effort.py
|
||||
- tests/litellm/llms/anthropic/test_anthropic_schema_filter.py
|
||||
- tests/litellm/llms/azure/test_azure_embedding.py
|
||||
- tests/litellm/llms/bedrock/embed/test_embedding.py
|
||||
- tests/litellm/llms/bedrock/test_nova_imported_models.py
|
||||
- tests/litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py
|
||||
- tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py
|
||||
- tests/litellm/llms/oci/chat/test_oci_chat_transformation.py
|
||||
- tests/litellm/llms/openai_like/test_abliteration_provider.py
|
||||
- tests/litellm/llms/openai_like/test_assemblyai_provider.py
|
||||
- tests/litellm/llms/openai_like/test_empiriolabs_provider.py
|
||||
- tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py
|
||||
- tests/litellm/llms/vertex_ai/gemini/test_transformation.py
|
||||
- tests/litellm/llms/vertex_ai/text_to_speech/test_transformation.py
|
||||
- tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
|
||||
- tests/litellm/proxy/agent_endpoints/test_agent_rbac.py
|
||||
- tests/litellm/proxy/common_utils/test_rbac_utils.py
|
||||
- tests/litellm/proxy/management_endpoints/test_common_utils.py
|
||||
- tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py
|
||||
- tests/litellm/proxy/test_claude_code_marketplace.py
|
||||
- tests/litellm/proxy/test_init_litellm_callbacks.py
|
||||
- tests/litellm/proxy/test_prisma_engine_watchdog.py
|
||||
- tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py
|
||||
- tests/litellm/test_bedrock_extended_beta_models.py
|
||||
- tests/litellm/test_bedrock_nemotron_super.py
|
||||
- tests/litellm/test_proxy_auth.py
|
||||
- tests/litellm/test_router_retry_backoff_headers.py
|
||||
- tests/litellm/test_sambanova_model_metadata.py
|
||||
- tests/litellm/test_stream_chunk_builder_images.py
|
||||
- reason: >-
|
||||
Legacy proxy suite superseded by the proxy shards; no job invokes it and whether it still
|
||||
describes supported behaviour is unresolved
|
||||
paths:
|
||||
- tests/old_proxy_tests/tests/test_anthropic_context_caching.py
|
||||
- tests/old_proxy_tests/tests/test_anthropic_sdk.py
|
||||
- tests/old_proxy_tests/tests/test_async.py
|
||||
- tests/old_proxy_tests/tests/test_gemini_context_caching.py
|
||||
- tests/old_proxy_tests/tests/test_langchain_embedding.py
|
||||
- tests/old_proxy_tests/tests/test_langchain_request.py
|
||||
- tests/old_proxy_tests/tests/test_llamaindex.py
|
||||
- tests/old_proxy_tests/tests/test_mistral_sdk.py
|
||||
- tests/old_proxy_tests/tests/test_openai_embedding.py
|
||||
- tests/old_proxy_tests/tests/test_openai_exception_request.py
|
||||
- tests/old_proxy_tests/tests/test_openai_request.py
|
||||
- tests/old_proxy_tests/tests/test_openai_request_with_traceparent.py
|
||||
- tests/old_proxy_tests/tests/test_openai_simple_embedding.py
|
||||
- tests/old_proxy_tests/tests/test_openai_tts_request.py
|
||||
- tests/old_proxy_tests/tests/test_pass_through_langfuse.py
|
||||
- tests/old_proxy_tests/tests/test_q.py
|
||||
- tests/old_proxy_tests/tests/test_simple_traceparent_openai.py
|
||||
- tests/old_proxy_tests/tests/test_vertex_sdk_forward_headers.py
|
||||
- tests/old_proxy_tests/tests/test_vtx_embedding.py
|
||||
- tests/old_proxy_tests/tests/test_vtx_sdk_embedding.py
|
||||
- reason: >-
|
||||
No job invokes this suite and its files mix pure transformation tests with ones driving live
|
||||
vendor vector stores, so assigning them needs a per-file decision
|
||||
paths:
|
||||
- tests/vector_store_tests/rag/test_rag_bedrock.py
|
||||
- tests/vector_store_tests/rag/test_rag_openai.py
|
||||
- tests/vector_store_tests/rag/test_rag_s3_vectors.py
|
||||
- tests/vector_store_tests/rag/test_rag_vertex_ai.py
|
||||
- tests/vector_store_tests/test_azure_ai_vector_store.py
|
||||
- tests/vector_store_tests/test_azure_vector_store.py
|
||||
- tests/vector_store_tests/test_bedrock_vector_store.py
|
||||
- tests/vector_store_tests/test_gemini_vector_store.py
|
||||
- tests/vector_store_tests/test_milvus_vector_store.py
|
||||
- tests/vector_store_tests/test_openai_vector_store.py
|
||||
- tests/vector_store_tests/test_ragflow_vector_store.py
|
||||
- tests/vector_store_tests/test_s3_vectors_vector_store.py
|
||||
- tests/vector_store_tests/test_vertex_ai_search_api_vector_store.py
|
||||
- tests/vector_store_tests/test_vertex_ai_vector_store.py
|
||||
- reason: >-
|
||||
Throughput and memory-growth measurements whose runtime and variance make them unsuitable for
|
||||
a per-pull-request job
|
||||
paths:
|
||||
- tests/load_tests/test_datadog_load_test.py
|
||||
- tests/load_tests/test_langsmith_load_test.py
|
||||
- tests/load_tests/test_linear_memory_growth.py
|
||||
- tests/load_tests/test_memory_usage.py
|
||||
- tests/load_tests/test_otel_load_test.py
|
||||
- tests/load_tests/test_vertex_embeddings_load_test.py
|
||||
- tests/load_tests/test_vertex_load_tests.py
|
||||
- reason: >-
|
||||
Third-party integration tests that skip themselves without OCI configuration or sandbox
|
||||
credentials, neither of which a pull request job holds
|
||||
paths:
|
||||
- tests/integration/sandbox/test_e2b_sandbox.py
|
||||
- tests/integration/test_oci_integration.py
|
||||
- tests/integration/test_oci_proxy_integration.py
|
||||
- reason: >-
|
||||
Two prompt-factory tests sitting at the top level of tests/ instead of under the
|
||||
tests/test_litellm mirror the shards enumerate; they need moving rather than a shard entry
|
||||
paths:
|
||||
- tests/litellm_core_utils/test_anthropic_dedup_factory.py
|
||||
- tests/litellm_core_utils/test_bedrock_converse_dedup_factory.py
|
||||
- reason: >-
|
||||
A unit test for the proxy-extras package that no job invokes, while the package's other tests
|
||||
live under tests/proxy_migration_tests
|
||||
paths:
|
||||
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
|
||||
|
||||
dockerfiles:
|
||||
- reason: >-
|
||||
The componentized images the microservices chart deploys are built by no job; wiring both into
|
||||
the scan workflow costs a full image build each and is deferred to a change that prices the
|
||||
whole set
|
||||
paths:
|
||||
- backend/Dockerfile
|
||||
- gateway/Dockerfile
|
||||
- reason: >-
|
||||
The dashboard container is a static Next.js export served by nginx, and the dashboard build
|
||||
and lint workflows already exercise that output, so building the image adds no signal about it
|
||||
paths:
|
||||
- ui/Dockerfile
|
||||
- reason: >-
|
||||
The Rust gateway ships as its own chart and package with a separate release pipeline, so its
|
||||
image is not part of this repo's Python image set
|
||||
paths:
|
||||
- litellm-rust/crates/ai-gateway/Dockerfile
|
||||
- reason: >-
|
||||
An example image under cookbook/ that is documentation rather than a shipped artifact
|
||||
paths:
|
||||
- cookbook/litellm-ollama-docker-image/Dockerfile
|
||||
262
.github/scripts/assert_ci_coverage.py
vendored
Normal file
262
.github/scripts/assert_ci_coverage.py
vendored
Normal file
|
|
@ -0,0 +1,262 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import pathlib
|
||||
import re
|
||||
import sys
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
import yaml
|
||||
|
||||
REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
|
||||
WORKFLOW_DIR = REPO_ROOT / ".github" / "workflows"
|
||||
CIRCLECI_CONFIG = REPO_ROOT / ".circleci" / "config.yml"
|
||||
ALLOWLIST_FILE = REPO_ROOT / ".github" / "ci-coverage-allowlist.yml"
|
||||
TESTS_ROOT = REPO_ROOT / "tests"
|
||||
|
||||
ALLOWLIST_KEYS = frozenset({"description", "test_paths", "dockerfiles"})
|
||||
PATH_FILTER_KEYS = frozenset({"paths", "paths-ignore"})
|
||||
TEST_PATH_KEYS = frozenset({"test-path", "test-paths"})
|
||||
DOCKERFILE_INPUT_KEYS = frozenset({"file", "dockerfile"})
|
||||
TEST_RUNNER_RE = re.compile(r"\bpytest\b|\bcircleci tests\b|\bhelm unittest\b|\bplaywright test\b|\bpython[0-9.]*\s")
|
||||
IMAGE_BUILD_RE = re.compile(r"\bdocker\s+(?:buildx\s+)?build\b")
|
||||
TEST_TOKEN_RE = re.compile(r"tests/[A-Za-z0-9_./*?-]+")
|
||||
DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
|
||||
COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
|
||||
GLOB_CHARS = frozenset("*?")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AllowEntry:
|
||||
paths: tuple[str, ...]
|
||||
reason: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Allowlist:
|
||||
test_paths: tuple[AllowEntry, ...]
|
||||
dockerfiles: tuple[AllowEntry, ...]
|
||||
|
||||
def covers_test(self, relative_path: str) -> bool:
|
||||
return any(_token_covers(path, relative_path) for entry in self.test_paths for path in entry.paths)
|
||||
|
||||
def covers_dockerfile(self, relative_path: str) -> bool:
|
||||
return any(relative_path == path for entry in self.dockerfiles for path in entry.paths)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Scalar:
|
||||
key: str
|
||||
value: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Finding:
|
||||
subject: str
|
||||
detail: str
|
||||
|
||||
|
||||
def _scalars(node: object, key: str) -> tuple[Scalar, ...]:
|
||||
if isinstance(node, str):
|
||||
return (Scalar(key=key, value=node),)
|
||||
if isinstance(node, Mapping):
|
||||
return tuple(
|
||||
scalar
|
||||
for child_key, value in node.items()
|
||||
if child_key not in PATH_FILTER_KEYS
|
||||
for scalar in _scalars(value, str(child_key))
|
||||
)
|
||||
if isinstance(node, Sequence):
|
||||
return tuple(scalar for item in node for scalar in _scalars(item, key))
|
||||
return ()
|
||||
|
||||
|
||||
def _config_files() -> tuple[pathlib.Path, ...]:
|
||||
workflows = tuple(sorted(path for path in WORKFLOW_DIR.iterdir() if path.suffix in (".yml", ".yaml")))
|
||||
circleci = (CIRCLECI_CONFIG,) if CIRCLECI_CONFIG.is_file() else ()
|
||||
return workflows + circleci
|
||||
|
||||
|
||||
def _all_scalars() -> tuple[Scalar, ...]:
|
||||
return tuple(
|
||||
scalar
|
||||
for path in _config_files()
|
||||
for scalar in _scalars(yaml.safe_load(path.read_text(encoding="utf-8")), path.name)
|
||||
)
|
||||
|
||||
|
||||
def _uncommented(value: str) -> str:
|
||||
return COMMENT_RE.sub("", value)
|
||||
|
||||
|
||||
def _invoked_test_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
match.group(0).rstrip("/")
|
||||
for scalar in scalars
|
||||
if scalar.key in TEST_PATH_KEYS or TEST_RUNNER_RE.search(scalar.value)
|
||||
for match in TEST_TOKEN_RE.finditer(_uncommented(scalar.value))
|
||||
)
|
||||
|
||||
|
||||
def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
match.group(0)
|
||||
for scalar in scalars
|
||||
if scalar.key in DOCKERFILE_INPUT_KEYS or IMAGE_BUILD_RE.search(scalar.value)
|
||||
for match in DOCKERFILE_TOKEN_RE.finditer(_uncommented(scalar.value))
|
||||
)
|
||||
|
||||
|
||||
def _glob_to_regex(token: str) -> re.Pattern[str]:
|
||||
parts = re.split(r"(\*\*/|\*\*|\*|\?)", token)
|
||||
translated = "".join(
|
||||
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part, re.escape(part)) for part in parts
|
||||
)
|
||||
return re.compile(rf"{translated}(?:/.*)?$")
|
||||
|
||||
|
||||
def _token_covers(token: str, relative_path: str) -> bool:
|
||||
if GLOB_CHARS & set(token):
|
||||
return _glob_to_regex(token).match(relative_path) is not None
|
||||
return relative_path == token or relative_path.startswith(f"{token}/")
|
||||
|
||||
|
||||
def _test_files() -> tuple[str, ...]:
|
||||
return tuple(
|
||||
sorted(
|
||||
path.relative_to(REPO_ROOT).as_posix()
|
||||
for path in TESTS_ROOT.rglob("test_*.py")
|
||||
if path.is_file() and "node_modules" not in path.parts
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _dockerfiles() -> tuple[str, ...]:
|
||||
return tuple(
|
||||
sorted(
|
||||
path.relative_to(REPO_ROOT).as_posix()
|
||||
for path in REPO_ROOT.rglob("Dockerfile*")
|
||||
if path.is_file()
|
||||
and ".git" not in path.parts
|
||||
and "node_modules" not in path.parts
|
||||
and not path.name.endswith(".dockerignore")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _uncovered_tests(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
|
||||
uncovered = tuple(
|
||||
relative_path
|
||||
for relative_path in _test_files()
|
||||
if not any(_token_covers(token, relative_path) for token in tokens) and not allowlist.covers_test(relative_path)
|
||||
)
|
||||
directories = tuple(dict.fromkeys(path.rsplit("/", 1)[0] for path in uncovered))
|
||||
return tuple(
|
||||
Finding(
|
||||
subject=directory,
|
||||
detail=_describe(tuple(p for p in uncovered if p.rsplit("/", 1)[0] == directory)),
|
||||
)
|
||||
for directory in directories
|
||||
)
|
||||
|
||||
|
||||
def _describe(paths: tuple[str, ...]) -> str:
|
||||
names = ", ".join(path.rsplit("/", 1)[1] for path in paths[:3])
|
||||
suffix = f", +{len(paths) - 3} more" if len(paths) > 3 else ""
|
||||
return f"{len(paths)} test file(s) invoked by no job: {names}{suffix}"
|
||||
|
||||
|
||||
def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
|
||||
return tuple(
|
||||
Finding(subject=relative_path, detail="built by no job")
|
||||
for relative_path in _dockerfiles()
|
||||
if relative_path not in tokens and not allowlist.covers_dockerfile(relative_path)
|
||||
)
|
||||
|
||||
|
||||
def _parse_entry(item: object, section: str) -> AllowEntry:
|
||||
if not isinstance(item, dict):
|
||||
raise SystemExit(f"{ALLOWLIST_FILE.name}: '{section}' entries must be mappings")
|
||||
paths = item.get("paths")
|
||||
reason = item.get("reason")
|
||||
if (
|
||||
not isinstance(paths, list)
|
||||
or not paths
|
||||
or not all(isinstance(path, str) for path in paths)
|
||||
or not isinstance(reason, str)
|
||||
or not reason.strip()
|
||||
):
|
||||
raise SystemExit(
|
||||
f"{ALLOWLIST_FILE.name}: every '{section}' entry needs a non-empty 'paths' "
|
||||
"list of strings and a non-empty 'reason'"
|
||||
)
|
||||
return AllowEntry(paths=tuple(paths), reason=reason)
|
||||
|
||||
|
||||
def _parse_entries(raw: object, section: str) -> tuple[AllowEntry, ...]:
|
||||
if not isinstance(raw, list):
|
||||
raise SystemExit(f"{ALLOWLIST_FILE.name}: '{section}' must be a list")
|
||||
return tuple(_parse_entry(item, section) for item in raw)
|
||||
|
||||
|
||||
def _load_allowlist() -> Allowlist:
|
||||
if not ALLOWLIST_FILE.is_file():
|
||||
return Allowlist(test_paths=(), dockerfiles=())
|
||||
raw = yaml.safe_load(ALLOWLIST_FILE.read_text(encoding="utf-8")) or {}
|
||||
if not isinstance(raw, dict):
|
||||
raise SystemExit(f"{ALLOWLIST_FILE.name}: top level must be a mapping")
|
||||
unknown = sorted(str(key) for key in raw if key not in ALLOWLIST_KEYS)
|
||||
if unknown:
|
||||
raise SystemExit(
|
||||
f"{ALLOWLIST_FILE.name}: unknown top-level key(s) {unknown}; expected only {sorted(ALLOWLIST_KEYS)}"
|
||||
)
|
||||
return Allowlist(
|
||||
test_paths=_parse_entries(raw.get("test_paths", []), "test_paths"),
|
||||
dockerfiles=_parse_entries(raw.get("dockerfiles", []), "dockerfiles"),
|
||||
)
|
||||
|
||||
|
||||
def _write(message: str) -> None:
|
||||
sys.stdout.write(f"{message}\n")
|
||||
|
||||
|
||||
def _report(title: str, findings: tuple[Finding, ...], remedy: str) -> None:
|
||||
_write(f"ERROR: {title}")
|
||||
for finding in findings:
|
||||
_write(f" - {finding.subject}: {finding.detail}")
|
||||
_write("")
|
||||
_write(remedy)
|
||||
_write("")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
allowlist = _load_allowlist()
|
||||
scalars = _all_scalars()
|
||||
|
||||
test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars))
|
||||
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
|
||||
|
||||
if test_findings:
|
||||
_report(
|
||||
"test files that no CI job invokes",
|
||||
test_findings,
|
||||
"Add each to a job's test path, or list it in .github/ci-coverage-allowlist.yml with a reason.",
|
||||
)
|
||||
if dockerfile_findings:
|
||||
_report(
|
||||
"Dockerfiles that no CI job builds",
|
||||
dockerfile_findings,
|
||||
"Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
|
||||
)
|
||||
if test_findings or dockerfile_findings:
|
||||
return 1
|
||||
|
||||
_write(
|
||||
f"OK: {len(_test_files())} test files and {len(_dockerfiles())} Dockerfiles are each "
|
||||
"invoked by at least one job or carry an explicit allowlist entry."
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
13
.github/workflows/_test-unit-base.yml
vendored
13
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -154,6 +154,19 @@ jobs:
|
|||
merge-multiple: true
|
||||
|
||||
- name: Upload to Codecov
|
||||
id: codecov-upload
|
||||
continue-on-error: true
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
with:
|
||||
use_oidc: true
|
||||
directory: coverage-reports
|
||||
root_dir: ${{ github.workspace }}
|
||||
flags: ${{ inputs.artifact-name }}
|
||||
fail_ci_if_error: false
|
||||
|
||||
- name: Upload to Codecov (retry)
|
||||
if: steps.codecov-upload.outcome == 'failure'
|
||||
continue-on-error: true
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
with:
|
||||
use_oidc: true
|
||||
|
|
|
|||
2
.github/workflows/check-ui-api-types.yml
vendored
2
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
42
.github/workflows/ci-coverage.yml
vendored
Normal file
42
.github/workflows/ci-coverage.yml
vendored
Normal file
|
|
@ -0,0 +1,42 @@
|
|||
name: "CI Coverage"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- 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.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
assert-ci-coverage:
|
||||
name: assert-ci-coverage
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Assert every test file and Dockerfile is invoked by a job
|
||||
run: |
|
||||
python -m pip install "pyyaml==6.0.3"
|
||||
python .github/scripts/assert_ci_coverage.py
|
||||
|
|
@ -13,35 +13,16 @@ jobs:
|
|||
contents: write
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
|
||||
- name: Create daily oss-agent-shin branch
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
# Configure Git user
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "github-actions[bot]@users.noreply.github.com"
|
||||
|
||||
# Generate branch name with MM_DD_YYYY format
|
||||
BRANCH_NAME="litellm_oss_agent_shin_$(date +'%m_%d_%Y')"
|
||||
echo "Creating branch: $BRANCH_NAME"
|
||||
|
||||
# Fetch all branches
|
||||
git fetch --all
|
||||
|
||||
# Check if the branch already exists
|
||||
if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then
|
||||
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
|
||||
echo "Branch $BRANCH_NAME already exists. Skipping creation."
|
||||
else
|
||||
echo "Creating new branch: $BRANCH_NAME"
|
||||
# Create the new branch from main
|
||||
git checkout -b $BRANCH_NAME origin/main
|
||||
# Push the new branch
|
||||
git push origin $BRANCH_NAME
|
||||
echo "Successfully created and pushed branch: $BRANCH_NAME"
|
||||
exit 0
|
||||
fi
|
||||
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
|
||||
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
|
||||
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"
|
||||
|
|
|
|||
|
|
@ -13,38 +13,19 @@ jobs:
|
|||
contents: write
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
|
||||
- name: Create daily staging branch
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
# Configure Git user
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "github-actions[bot]@users.noreply.github.com"
|
||||
|
||||
# Generate branch name with MM_DD_YYYY format
|
||||
BRANCH_NAME="litellm_oss_staging_$(date +'%m_%d_%Y')"
|
||||
echo "Creating branch: $BRANCH_NAME"
|
||||
|
||||
# Fetch all branches
|
||||
git fetch --all
|
||||
|
||||
# Check if the branch already exists
|
||||
if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then
|
||||
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
|
||||
echo "Branch $BRANCH_NAME already exists. Skipping creation."
|
||||
else
|
||||
echo "Creating new branch: $BRANCH_NAME"
|
||||
# Create the new branch from main
|
||||
git checkout -b $BRANCH_NAME origin/main
|
||||
# Push the new branch
|
||||
git push origin $BRANCH_NAME
|
||||
echo "Successfully created and pushed branch: $BRANCH_NAME"
|
||||
exit 0
|
||||
fi
|
||||
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
|
||||
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
|
||||
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"
|
||||
|
||||
create-internal-dev-branch:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
|
|
@ -53,35 +34,16 @@ jobs:
|
|||
contents: write
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
|
||||
- name: Create internal dev branch
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
# Configure Git user
|
||||
git config user.name "github-actions[bot]"
|
||||
git config user.email "github-actions[bot]@users.noreply.github.com"
|
||||
|
||||
# Generate branch name with MM_DD_YYYY format
|
||||
BRANCH_NAME="litellm_internal_dev_$(date +'%m_%d_%Y')"
|
||||
echo "Creating branch: $BRANCH_NAME"
|
||||
|
||||
# Fetch all branches
|
||||
git fetch --all
|
||||
|
||||
# Check if the branch already exists
|
||||
if git show-ref --verify --quiet refs/remotes/origin/$BRANCH_NAME; then
|
||||
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
|
||||
echo "Branch $BRANCH_NAME already exists. Skipping creation."
|
||||
else
|
||||
echo "Creating new branch: $BRANCH_NAME"
|
||||
# Create the new branch from main
|
||||
git checkout -b $BRANCH_NAME origin/main
|
||||
# Push the new branch
|
||||
git push origin $BRANCH_NAME
|
||||
echo "Successfully created and pushed branch: $BRANCH_NAME"
|
||||
exit 0
|
||||
fi
|
||||
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
|
||||
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
|
||||
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"
|
||||
|
|
|
|||
31
.github/workflows/helm_unit_test.yml
vendored
31
.github/workflows/helm_unit_test.yml
vendored
|
|
@ -23,21 +23,28 @@ jobs:
|
|||
with:
|
||||
version: "3.11.1"
|
||||
|
||||
- name: Download and verify Helm Unit Test Plugin
|
||||
run: |
|
||||
curl -fsSLo "$RUNNER_TEMP/helm-unittest.tgz" https://github.com/helm-unittest/helm-unittest/releases/download/v0.8.2/helm-unittest-linux-amd64-0.8.2.tgz
|
||||
echo "56ab3091e6fa52a7c92ee951def9bed957f295d9ce98483aed404e748d7b3a94 $RUNNER_TEMP/helm-unittest.tgz" | sha256sum -c -
|
||||
|
||||
- name: Install Helm Unit Test Plugin
|
||||
run: |
|
||||
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4
|
||||
- name: Verify Helm Unit Test Plugin integrity
|
||||
run: |
|
||||
EXPECTED_SHA="e251ba198448629678ff2168e1a469249d998155"
|
||||
PLUGIN_DIR="$(helm env HELM_PLUGINS)/helm-unittest"
|
||||
ACTUAL_SHA="$(git -C "$PLUGIN_DIR" rev-parse HEAD)"
|
||||
if [ "$ACTUAL_SHA" != "$EXPECTED_SHA" ]; then
|
||||
echo "::error::Helm unittest plugin checksum mismatch! Expected $EXPECTED_SHA but got $ACTUAL_SHA"
|
||||
exit 1
|
||||
fi
|
||||
echo "Helm unittest plugin integrity verified: $ACTUAL_SHA"
|
||||
mkdir -p "$PLUGIN_DIR"
|
||||
tar -xzf "$RUNNER_TEMP/helm-unittest.tgz" -C "$PLUGIN_DIR"
|
||||
helm plugin list
|
||||
|
||||
- name: Run unit tests
|
||||
run: |
|
||||
helm unittest -f 'tests/*.yaml' helm/litellm-helm
|
||||
helm unittest -f 'tests/*.yaml' helm/litellm
|
||||
for chart in helm/litellm-helm helm/litellm; do
|
||||
declared="$(grep -h '^suite:' "$chart"/tests/*.yaml | wc -l | tr -d '[:space:]')"
|
||||
output="$(mktemp)"
|
||||
helm unittest -f 'tests/*.yaml' "$chart" | tee "$output"
|
||||
executed="$(sed -n 's/^Test Suites:.*[[:space:]]\([0-9][0-9]*\) total$/\1/p' "$output")"
|
||||
if [ "$declared" != "$executed" ]; then
|
||||
echo "::error::$chart declares $declared test suites but helm-unittest ran $executed. Suites are being skipped silently, so their assertions never execute."
|
||||
exit 1
|
||||
fi
|
||||
echo "$chart: all $declared declared test suites ran"
|
||||
done
|
||||
|
|
|
|||
33
.github/workflows/image-scan.yml
vendored
33
.github/workflows/image-scan.yml
vendored
|
|
@ -8,10 +8,12 @@ on:
|
|||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- Dockerfile
|
||||
- docker/Dockerfile.non_root
|
||||
- migrations/Dockerfile
|
||||
- migrations/run.py
|
||||
- tests/proxy_migration_tests/test_offline_image_migration.py
|
||||
- litellm-proxy-extras/**
|
||||
- tests/proxy_migration_tests/**
|
||||
- uv.lock
|
||||
- ui/litellm-dashboard/package-lock.json
|
||||
- .github/workflows/image-scan.yml
|
||||
|
|
@ -86,6 +88,35 @@ jobs:
|
|||
--fail-on high \
|
||||
--output table
|
||||
|
||||
runtime-image:
|
||||
name: runtime-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 runtime image
|
||||
run: docker build -f Dockerfile -t litellm-runtime-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-runtime-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
|
||||
migrations-image:
|
||||
name: migrations-image
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
70
.github/workflows/publish-basedpyright-base-counts.yml
vendored
Normal file
70
.github/workflows/publish-basedpyright-base-counts.yml
vendored
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
name: Publish basedpyright base counts
|
||||
|
||||
# Every commit on litellm_internal_staging is some branch's future merge-base.
|
||||
# Publishing its per-rule basedpyright counts as an artifact lets
|
||||
# scripts/type_check_gate.py download them in seconds instead of paying a
|
||||
# 60-110s second basedpyright pass on every fresh worktree or moved merge-base.
|
||||
# No concurrency group on purpose: runs must never cancel each other, because
|
||||
# every sha's artifact matters (any of them can become a merge-base).
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- litellm_internal_staging
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
ref:
|
||||
description: "Ref to compute and publish base counts for"
|
||||
required: false
|
||||
default: litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
publish:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ inputs.ref || github.sha }}
|
||||
clean: true
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
uv sync --frozen --group proxy-dev --group e2e-dev
|
||||
|
||||
# Mirrors test-linting.yml's lint job: basedpyright resolves Prisma's
|
||||
# generated client only after `prisma generate`, and the published counts
|
||||
# must match what that job would measure for the same tree.
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Emit basedpyright counts for HEAD
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts"
|
||||
counts_file=$(ls "$RUNNER_TEMP"/basedpyright-counts/basedpyright-counts-*.json)
|
||||
echo "COUNTS_ARTIFACT_NAME=$(basename "$counts_file" .json)" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Upload counts artifact
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: ${{ env.COUNTS_ARTIFACT_NAME }}
|
||||
path: ${{ runner.temp }}/basedpyright-counts/
|
||||
if-no-files-found: error
|
||||
53
.github/workflows/test-linting.yml
vendored
53
.github/workflows/test-linting.yml
vendored
|
|
@ -15,6 +15,12 @@ jobs:
|
|||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
# actions: read lets scripts/type_check_gate.py download the base-counts
|
||||
# artifact published by publish-basedpyright-base-counts.yml instead of
|
||||
# re-running basedpyright over the merge-base tree.
|
||||
permissions:
|
||||
contents: read
|
||||
actions: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
|
|
@ -23,10 +29,21 @@ jobs:
|
|||
# Any-discipline) would otherwise blame on this branch.
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha }}
|
||||
fetch-depth: 0
|
||||
fetch-depth: 1
|
||||
clean: true
|
||||
persist-credentials: false
|
||||
|
||||
- name: Fetch gate base (merge-base with target branch)
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
MERGE_BASE=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
|
||||
test -n "$MERGE_BASE"
|
||||
git fetch --no-tags --depth=1 origin "$MERGE_BASE"
|
||||
echo "GATE_BASE_SHA=$MERGE_BASE" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
|
|
@ -60,10 +77,8 @@ jobs:
|
|||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Check ruff format
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
git diff --name-only --diff-filter=ACMR "$BASE_SHA"...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
|
||||
git diff --name-only --diff-filter=ACMR "$GATE_BASE_SHA" HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' > "$RUNNER_TEMP/ruff_format_files.txt" || true
|
||||
if [ ! -s "$RUNNER_TEMP/ruff_format_files.txt" ]; then
|
||||
echo "No changed litellm Python files to check with ruff format."
|
||||
exit 0
|
||||
|
|
@ -86,16 +101,12 @@ jobs:
|
|||
cd ..
|
||||
|
||||
- name: Check strict-rule budget (delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
uv run --no-sync python scripts/ruff_strict_gate.py --base "$BASE_SHA"
|
||||
uv run --no-sync python scripts/ruff_strict_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Check type-discipline budget (mutable collections / casts / type guards / kwargs / unexplained suppressions, delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_discipline_gate.py --base "$BASE_SHA"
|
||||
uv run --no-sync python scripts/type_discipline_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
- name: Print OpenAI version
|
||||
run: |
|
||||
|
|
@ -103,16 +114,13 @@ jobs:
|
|||
|
||||
- name: Check basedpyright budget (delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
NODE_OPTIONS: --max-old-space-size=12288
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
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 "$GATE_BASE_SHA"
|
||||
|
||||
- name: Check tests/e2e basedpyright (zero errors)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
if git diff --name-only --diff-filter=ACMRD "$BASE_SHA"...HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
|
||||
if git diff --name-only --diff-filter=ACMRD "$GATE_BASE_SHA" HEAD -- 'tests/e2e/**/*.py' | grep -q .; then
|
||||
uv run --no-sync basedpyright tests/e2e
|
||||
else
|
||||
echo "No changed tests/e2e Python files; skipping."
|
||||
|
|
@ -141,9 +149,15 @@ jobs:
|
|||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- name: Fetch ratchet base
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
git fetch --no-tags --depth=1 origin "$BASE_SHA"
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
|
|
@ -164,7 +178,7 @@ jobs:
|
|||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
|
|
@ -179,13 +193,14 @@ jobs:
|
|||
|
||||
- name: Run secret scan test
|
||||
run: |
|
||||
uv run --frozen --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
|
||||
uv run --no-project --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
|
||||
|
||||
- name: Run ggshield secret scan
|
||||
env:
|
||||
GITGUARDIAN_API_KEY: ${{ secrets.GITGUARDIAN_API_KEY }}
|
||||
run: |
|
||||
if [ -n "$GITGUARDIAN_API_KEY" ]; then
|
||||
git fetch --no-tags --unshallow origin
|
||||
uv tool run --from 'ggshield==1.48.0' ggshield secret scan repo .
|
||||
else
|
||||
echo "GITGUARDIAN_API_KEY not set, skipping ggshield scan"
|
||||
|
|
|
|||
2
.github/workflows/test-litellm-ui-build.yml
vendored
2
.github/workflows/test-litellm-ui-build.yml
vendored
|
|
@ -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
|
||||
|
||||
|
|
|
|||
9
.github/workflows/test-litellm-ui-lint.yml
vendored
9
.github/workflows/test-litellm-ui-lint.yml
vendored
|
|
@ -22,12 +22,13 @@ jobs:
|
|||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- name: Collect changed files
|
||||
id: changed
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
|
|
@ -37,7 +38,9 @@ jobs:
|
|||
# landed since, so a PR that touches no UI file still gets linted
|
||||
# against hundreds of other people's files. Diff the PR head against its
|
||||
# own merge base instead, which is exactly what this PR changed.
|
||||
merge_base=$(git merge-base "$BASE_SHA" "$HEAD_SHA")
|
||||
merge_base=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
|
||||
test -n "$merge_base"
|
||||
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
|
||||
: > "$RUNNER_TEMP/prettier_files.txt"
|
||||
: > "$RUNNER_TEMP/eslint_files.txt"
|
||||
while IFS= read -r f; do
|
||||
|
|
@ -61,7 +64,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
|
||||
|
||||
|
|
|
|||
21
.github/workflows/test-litellm-ui-unit.yml
vendored
21
.github/workflows/test-litellm-ui-unit.yml
vendored
|
|
@ -29,13 +29,13 @@ jobs:
|
|||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
fetch-depth: 1
|
||||
persist-credentials: false
|
||||
|
||||
- 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
|
||||
|
||||
|
|
@ -45,11 +45,24 @@ jobs:
|
|||
- name: Run UI unit tests (Vitest)
|
||||
env:
|
||||
CI: "true"
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
if [ -n "$BASE_SHA" ]; then
|
||||
echo "Pull request: running only tests related to changes since $BASE_SHA"
|
||||
npm run test -- --run --changed "$BASE_SHA" --passWithNoTests \
|
||||
merge_base=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
|
||||
test -n "$merge_base"
|
||||
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
|
||||
changed_files=()
|
||||
while IFS= read -r f; do
|
||||
changed_files+=("$f")
|
||||
done < <(git diff --name-only --relative "$merge_base" "$HEAD_SHA" -- .)
|
||||
if [ ${#changed_files[@]} -eq 0 ]; then
|
||||
echo "No UI files changed in this PR; skipping unit tests."
|
||||
exit 0
|
||||
fi
|
||||
echo "Pull request: running tests related to ${#changed_files[@]} changed UI files"
|
||||
npm run test -- related "${changed_files[@]}" --run --passWithNoTests \
|
||||
--pool forks --poolOptions.forks.maxForks=14
|
||||
else
|
||||
echo "Push to $GITHUB_REF_NAME: running the full suite"
|
||||
|
|
|
|||
4
.github/workflows/test-unit-misc.yml
vendored
4
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -40,7 +40,11 @@ jobs:
|
|||
tests/test_litellm/interactions
|
||||
tests/test_litellm/ocr
|
||||
tests/test_litellm/passthrough
|
||||
tests/test_litellm/rag
|
||||
tests/test_litellm/realtime_api
|
||||
tests/test_litellm/rerank_api
|
||||
tests/test_litellm/sandbox
|
||||
tests/test_litellm/test_router
|
||||
tests/test_litellm/vector_stores
|
||||
tests/test_litellm/videos
|
||||
tests/test_litellm/test_*.py
|
||||
|
|
|
|||
|
|
@ -29,7 +29,9 @@ jobs:
|
|||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: >-
|
||||
tests/test_litellm/proxy/analytics_endpoints
|
||||
tests/test_litellm/proxy/management_endpoints
|
||||
tests/test_litellm/proxy/memory
|
||||
tests/test_litellm/proxy/guardrails
|
||||
tests/test_litellm/proxy/management_helpers
|
||||
tests/test_litellm/proxy/anthropic_endpoints
|
||||
|
|
|
|||
2
.github/workflows/test-unit-proxy-infra.yml
vendored
2
.github/workflows/test-unit-proxy-infra.yml
vendored
|
|
@ -33,6 +33,8 @@ jobs:
|
|||
tests/test_litellm/proxy/_experimental
|
||||
tests/test_litellm/proxy/experimental
|
||||
tests/test_litellm/proxy/common_utils
|
||||
tests/test_litellm/proxy/enterprise_billing
|
||||
tests/test_litellm/proxy/types_utils
|
||||
tests/test_litellm/proxy/logging_endpoints
|
||||
tests/test_litellm/proxy/test_*.py
|
||||
workers: 2
|
||||
|
|
|
|||
11
CLAUDE.md
11
CLAUDE.md
|
|
@ -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,13 @@ 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
|
||||
|
||||
`make pre-commit` always saves its complete output to a per-worktree log file and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice, and re-run only after the working tree actually changed
|
||||
|
||||
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 +73,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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -134,7 +134,8 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
|||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
|
||||
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
|
||||
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \
|
||||
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
|
||||
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
|
|
|
|||
13
Makefile
13
Makefile
|
|
@ -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 install --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"; \
|
||||
|
|
@ -99,7 +99,10 @@ install-test-deps: install-proxy-dev
|
|||
$(UV_RUN) prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
install-helm-unittest:
|
||||
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4 || echo "ignore error if plugin exists"
|
||||
@helm plugin list | grep -qE '^unittest[[:space:]]+0\.8\.2([[:space:]]|$$)' || { \
|
||||
helm plugin uninstall unittest >/dev/null 2>&1 || true; \
|
||||
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.8.2; \
|
||||
}
|
||||
|
||||
# Install git hooks that enforce Conventional Commits and Conventional Branches.
|
||||
# Opt-in: not chained into install-dev.
|
||||
|
|
@ -176,10 +179,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 +193,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
|
||||
|
||||
|
|
|
|||
|
|
@ -59,9 +59,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra semantic-router \
|
||||
--python python3
|
||||
|
||||
RUN mkdir -p /home/nonroot && \
|
||||
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
|
||||
chown -R nonroot:nonroot /home/nonroot/.cache
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
||||
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
|
||||
|
||||
|
|
@ -83,13 +83,16 @@ ENV HOME=/home/nonroot \
|
|||
PATH="/app/.venv/bin:${PATH}" \
|
||||
PYTHONPATH="/app" \
|
||||
PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
|
||||
|
||||
COPY --from=builder --chown=nonroot:nonroot /app /app
|
||||
COPY --from=builder --chown=nonroot:nonroot /home/nonroot/.cache /home/nonroot/.cache
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
|
||||
|
||||
USER nonroot
|
||||
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/router/",
|
||||
"/router_settings",
|
||||
"/adaptive_router/",
|
||||
"/auto_router/",
|
||||
"/fallback",
|
||||
"/fallbacks",
|
||||
"/cache_settings",
|
||||
|
|
@ -81,6 +82,9 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/user_agent",
|
||||
"/usage/",
|
||||
"/daily/",
|
||||
# Deployment-wide gateway request counts. Scoped to the analytics read rather
|
||||
# than all of /gateway/, which stays free for data-plane routes.
|
||||
"/gateway/daily/",
|
||||
# CloudZero cost-export admin (init / settings / export / dry-run / delete)
|
||||
"/cloudzero/",
|
||||
# Caching admin
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 29813
|
||||
"limit": 29204
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2645
|
||||
"limit": 2635
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 329
|
||||
|
|
@ -15,52 +15,52 @@
|
|||
"limit": 123
|
||||
},
|
||||
"reportConstantRedefinition": {
|
||||
"limit": 59
|
||||
"limit": 40
|
||||
},
|
||||
"reportDeprecated": {
|
||||
"limit": 325
|
||||
"limit": 215
|
||||
},
|
||||
"reportDuplicateImport": {
|
||||
"limit": 42
|
||||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 9473
|
||||
"limit": 9227
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 11
|
||||
"limit": 7
|
||||
},
|
||||
"reportGeneralTypeIssues": {
|
||||
"limit": 227
|
||||
"limit": 157
|
||||
},
|
||||
"reportIncompatibleMethodOverride": {
|
||||
"limit": 77
|
||||
"limit": 56
|
||||
},
|
||||
"reportIncompatibleVariableOverride": {
|
||||
"limit": 12
|
||||
"limit": 8
|
||||
},
|
||||
"reportInconsistentOverload": {
|
||||
"limit": 18
|
||||
"limit": 12
|
||||
},
|
||||
"reportIndexIssue": {
|
||||
"limit": 37
|
||||
"limit": 35
|
||||
},
|
||||
"reportInvalidTypeForm": {
|
||||
"limit": 35
|
||||
},
|
||||
"reportInvalidTypeVarUse": {
|
||||
"limit": 5
|
||||
"limit": 2
|
||||
},
|
||||
"reportMatchNotExhaustive": {
|
||||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5855
|
||||
"limit": 5850
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15852
|
||||
"limit": 15833
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 41
|
||||
"limit": 40
|
||||
},
|
||||
"reportOperatorIssue": {
|
||||
"limit": 0
|
||||
|
|
@ -72,7 +72,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"limit": 1079
|
||||
"limit": 1078
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"limit": 0
|
||||
|
|
@ -81,16 +81,16 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportPossiblyUnboundVariable": {
|
||||
"limit": 77
|
||||
"limit": 56
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 2437
|
||||
"limit": 1825
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 12
|
||||
"limit": 8
|
||||
},
|
||||
"reportReturnType": {
|
||||
"limit": 219
|
||||
"limit": 218
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"limit": 27
|
||||
|
|
@ -99,48 +99,48 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 45324
|
||||
"limit": 45242
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 113
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 40452
|
||||
"limit": 40340
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 20309
|
||||
"limit": 20293
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 31978
|
||||
"limit": 31796
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 177
|
||||
"limit": 122
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 1021
|
||||
"limit": 703
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 7
|
||||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 1204
|
||||
"limit": 865
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 165
|
||||
"limit": 72
|
||||
},
|
||||
"reportUntypedFunctionDecorator": {
|
||||
"limit": 33
|
||||
},
|
||||
"reportUnusedClass": {
|
||||
"limit": 33
|
||||
"limit": 23
|
||||
},
|
||||
"reportUnusedFunction": {
|
||||
"limit": 204
|
||||
"limit": 139
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"limit": 1003
|
||||
"limit": 555
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"limit": 1297
|
||||
"limit": 146
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -133,7 +133,8 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
|||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
|
||||
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js
|
||||
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \
|
||||
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
|
||||
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -185,7 +185,8 @@ RUN mkdir -p /nonexistent /app/.cache /var/lib/litellm/assets /var/lib/litellm/u
|
|||
chmod -R a+rX /opt/prisma && \
|
||||
test -x /opt/prisma/binaries/node_modules/.bin/prisma && \
|
||||
test -f /opt/prisma/binaries/node_modules/prisma/build/index.js && \
|
||||
ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null 2>&1
|
||||
ls /opt/prisma/binaries/node_modules/@prisma/engines/query-engine-* >/dev/null 2>&1 && \
|
||||
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
|
||||
|
||||
USER 65534
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -47,7 +47,13 @@ RUN uv venv --python python && \
|
|||
"prisma==0.11.0" \
|
||||
"openai==2.24.0"
|
||||
|
||||
RUN prisma generate --schema=./schema.prisma
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
python -c "import sys; from prisma.client import BINARY_PATHS; bad = sorted(p for group in BINARY_PATHS.model_dump().values() for p in group.values() if not p.startswith('/opt/prisma/')); sys.exit('prisma engines baked outside /opt/prisma: %r' % bad) if bad else None"
|
||||
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
|
||||
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
|
|
|
|||
0
enterprise/litellm_enterprise/integrations/__init__.py
Normal file
0
enterprise/litellm_enterprise/integrations/__init__.py
Normal file
|
|
@ -296,17 +296,13 @@ class CheckBatchCost:
|
|||
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,
|
||||
resolve_managed_output_file_model_name,
|
||||
)
|
||||
|
||||
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 []
|
||||
return resolve_managed_output_file_model_name(
|
||||
unified_input_file_id=cls._get_input_file_id(job),
|
||||
fallback_model_name=deployment_info.model_name or None,
|
||||
)
|
||||
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]:
|
||||
|
|
|
|||
0
enterprise/litellm_enterprise/proxy/hooks/__init__.py
Normal file
0
enterprise/litellm_enterprise/proxy/hooks/__init__.py
Normal file
|
|
@ -4,7 +4,8 @@
|
|||
import base64
|
||||
import json
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast
|
||||
from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -33,8 +34,8 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
get_batch_id_from_unified_batch_id,
|
||||
get_content_type_from_file_object,
|
||||
get_model_id_from_unified_batch_id,
|
||||
get_models_from_unified_file_id,
|
||||
normalize_mime_type_for_provider,
|
||||
resolve_managed_output_file_model_name,
|
||||
)
|
||||
from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccessIssue]
|
||||
AllMessageValues,
|
||||
|
|
@ -382,7 +383,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
"flat_model_file_ids": {"hasSome": model_object_ids},
|
||||
}
|
||||
)
|
||||
return [OpenAIFileObject.model_validate(file_object.file_object) for file_object in file_ids]
|
||||
return [
|
||||
OpenAIFileObject.model_validate(row.file_object).model_copy(
|
||||
update={"id": row.unified_file_id}
|
||||
)
|
||||
for row in file_ids
|
||||
if row.file_object is not None
|
||||
]
|
||||
|
||||
async def check_managed_file_id_access(
|
||||
self, data: Dict, user_api_key_dict: UserAPIKeyAuth
|
||||
|
|
@ -1055,10 +1062,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
def get_unified_output_file_id(
|
||||
self, output_file_id: str, model_id: str, model_name: Optional[str]
|
||||
) -> str:
|
||||
deterministic_uuid: Final = uuid5(
|
||||
uuid5(NAMESPACE_URL, model_id), output_file_id
|
||||
)
|
||||
unified_output_file_id = (
|
||||
SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
|
||||
"application/json",
|
||||
str(uuid.uuid4()),
|
||||
str(deterministic_uuid),
|
||||
model_name or "",
|
||||
output_file_id,
|
||||
model_id,
|
||||
|
|
@ -1094,21 +1104,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
) # managed batch id
|
||||
model_id = cast(Optional[str], response._hidden_params.get("model_id"))
|
||||
model_name = cast(Optional[str], response._hidden_params.get("model_name"))
|
||||
resolved_model_name = model_name
|
||||
|
||||
# Some providers (e.g. Vertex batch retrieve) do not set model_name on
|
||||
# the response. In that case, recover target_model_names from the input
|
||||
# managed file metadata so unified output IDs preserve routing metadata.
|
||||
if not resolved_model_name and isinstance(unified_file_id, str):
|
||||
decoded_unified_file_id = (
|
||||
_is_base64_encoded_unified_file_id(unified_file_id)
|
||||
or unified_file_id
|
||||
)
|
||||
target_model_names = get_models_from_unified_file_id(
|
||||
decoded_unified_file_id
|
||||
)
|
||||
if target_model_names:
|
||||
resolved_model_name = ",".join(target_model_names)
|
||||
resolved_model_name = resolve_managed_output_file_model_name(
|
||||
unified_input_file_id=unified_file_id
|
||||
if isinstance(unified_file_id, str)
|
||||
else response.input_file_id,
|
||||
fallback_model_name=model_name,
|
||||
)
|
||||
original_response_id = response.id
|
||||
|
||||
if (unified_batch_id or unified_file_id) and model_id:
|
||||
|
|
|
|||
|
|
@ -831,7 +831,7 @@ async def project_info(
|
|||
)
|
||||
|
||||
# Check if user has access to this project (admin or team member)
|
||||
is_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
is_admin = user_api_key_has_admin_view(user_api_key_dict)
|
||||
is_team_member = False
|
||||
|
||||
if project.team_id and user_api_key_dict.user_id:
|
||||
|
|
@ -886,7 +886,7 @@ async def list_projects(
|
|||
)
|
||||
|
||||
# If proxy admin, get all projects
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
if user_api_key_has_admin_view(user_api_key_dict):
|
||||
projects: Sequence[
|
||||
prisma_models.LiteLLM_ProjectTable
|
||||
] = await prisma_client.db.litellm_projecttable.find_many(
|
||||
|
|
|
|||
0
enterprise/litellm_enterprise/py.typed
Normal file
0
enterprise/litellm_enterprise/py.typed
Normal file
0
enterprise/litellm_enterprise/types/__init__.py
Normal file
0
enterprise/litellm_enterprise/types/__init__.py
Normal file
0
enterprise/litellm_enterprise/types/proxy/__init__.py
Normal file
0
enterprise/litellm_enterprise/types/proxy/__init__.py
Normal file
|
|
@ -61,9 +61,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra bedrock-realtime \
|
||||
--python python3
|
||||
|
||||
RUN mkdir -p /home/nonroot && \
|
||||
HOME=/home/nonroot prisma generate --schema=./schema.prisma && \
|
||||
chown -R nonroot:nonroot /home/nonroot/.cache
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
prisma generate --schema=./schema.prisma
|
||||
|
||||
RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh
|
||||
|
||||
|
|
@ -85,13 +85,16 @@ ENV HOME=/home/nonroot \
|
|||
PATH="/app/.venv/bin:${PATH}" \
|
||||
PYTHONPATH="/app" \
|
||||
PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries
|
||||
|
||||
COPY --from=builder --chown=nonroot:nonroot /app /app
|
||||
COPY --from=builder --chown=nonroot:nonroot /home/nonroot/.cache /home/nonroot/.cache
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
chmod -R a+rX /opt/prisma && \
|
||||
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
|
||||
|
||||
USER nonroot
|
||||
|
||||
|
|
|
|||
|
|
@ -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 |
|
||||
|
|
|
|||
|
|
@ -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: []
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -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;
|
||||
|
|
@ -0,0 +1,15 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGatewayRequests" (
|
||||
"date" TEXT NOT NULL,
|
||||
"category" TEXT NOT NULL,
|
||||
"route" TEXT NOT NULL,
|
||||
"successful_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"failed_requests" BIGINT NOT NULL DEFAULT 0,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL,
|
||||
|
||||
CONSTRAINT "LiteLLM_DailyGatewayRequests_pkey" PRIMARY KEY ("date","category","route")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGatewayRequests_date_idx" ON "LiteLLM_DailyGatewayRequests"("date");
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterSession" (
|
||||
"api_key" TEXT NOT NULL,
|
||||
"session_id" TEXT NOT NULL,
|
||||
"router_name" TEXT NOT NULL,
|
||||
"router_type" TEXT NOT NULL,
|
||||
"first_turn_at" TIMESTAMP(3) NOT NULL,
|
||||
"last_turn_at" TIMESTAMP(3) NOT NULL,
|
||||
"last_model" TEXT NOT NULL,
|
||||
"models" JSONB NOT NULL DEFAULT '{}',
|
||||
"turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"unordered_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"covered_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"cache_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"same_model_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"same_model_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"first_visit_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"first_visit_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_hits" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_expired_misses" INTEGER NOT NULL DEFAULT 0,
|
||||
"return_within_ttl_misses" INTEGER NOT NULL DEFAULT 0,
|
||||
"ttl_5m_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"ttl_1h_turns" INTEGER NOT NULL DEFAULT 0,
|
||||
"total_tokens" BIGINT NOT NULL DEFAULT 0,
|
||||
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
"saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
|
||||
CONSTRAINT "LiteLLM_AutoRouterSession_pkey" PRIMARY KEY ("api_key", "session_id", "router_name")
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "idx_autorouter_session_last_turn" ON "LiteLLM_AutoRouterSession"("last_turn_at");
|
||||
181
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
Normal file
181
litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py
Normal file
|
|
@ -0,0 +1,181 @@
|
|||
"""Prepare the Node toolchain the Prisma CLI needs, separately from migrations.
|
||||
|
||||
The Prisma CLI is a Node program. The first invocation inside a fresh
|
||||
container installs a private Node runtime and npm-installs the CLI itself,
|
||||
which can take minutes on a cold or slow machine. Sharing one timeout between
|
||||
that one-time bootstrap and the migration commands makes a slow bootstrap
|
||||
indistinguishable from a slow migration, so the bootstrap gets killed long
|
||||
before it can finish.
|
||||
|
||||
A killed bootstrap does not correct itself. The installer leaves its cache
|
||||
directory behind, and Prisma decides whether to install by testing that
|
||||
directory for existence alone, so every later attempt skips the install and
|
||||
then fails on a Node binary that was never written. Deleting a cache directory
|
||||
that exists without a Node binary is what turns a killed bootstrap back into a
|
||||
recoverable one.
|
||||
|
||||
Both budgets are overridable so an operator can widen them without a release:
|
||||
``LITELLM_PRISMA_BOOTSTRAP_TIMEOUT`` for the toolchain install and
|
||||
``LITELLM_PRISMA_COMMAND_TIMEOUT`` for every individual Prisma command.
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
|
||||
try:
|
||||
from prisma import config as prisma_config
|
||||
except ImportError:
|
||||
prisma_config = None
|
||||
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_COMMAND_TIMEOUT"
|
||||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR = "LITELLM_PRISMA_BOOTSTRAP_TIMEOUT"
|
||||
NODEENV_CACHE_DIR_ENV_VAR = "PRISMA_NODEENV_CACHE_DIR"
|
||||
|
||||
DEFAULT_PRISMA_COMMAND_TIMEOUT = 60.0
|
||||
DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0
|
||||
|
||||
BOOTSTRAP_ARG = "--version"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolchainBootstrap:
|
||||
"""Outcome of preparing the Prisma toolchain."""
|
||||
|
||||
healed_incomplete_cache: bool
|
||||
ready: bool
|
||||
|
||||
|
||||
def _timeout_from_env(env_var: str, default: float) -> float:
|
||||
raw = os.getenv(env_var)
|
||||
if raw is None:
|
||||
return default
|
||||
try:
|
||||
seconds = float(raw)
|
||||
except ValueError:
|
||||
logger.warning(
|
||||
"%s=%r is not a number, falling back to %ss", env_var, raw, default
|
||||
)
|
||||
return default
|
||||
if not math.isfinite(seconds) or seconds <= 0:
|
||||
logger.warning(
|
||||
"%s=%r is not a finite positive number, falling back to %ss",
|
||||
env_var,
|
||||
raw,
|
||||
default,
|
||||
)
|
||||
return default
|
||||
return seconds
|
||||
|
||||
|
||||
def prisma_command_timeout() -> float:
|
||||
"""Seconds any single Prisma command may run for."""
|
||||
return _timeout_from_env(
|
||||
PRISMA_COMMAND_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_COMMAND_TIMEOUT
|
||||
)
|
||||
|
||||
|
||||
def prisma_bootstrap_timeout() -> float:
|
||||
"""Seconds the one-time Node toolchain install may run for."""
|
||||
return _timeout_from_env(
|
||||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR, DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT
|
||||
)
|
||||
|
||||
|
||||
def nodeenv_cache_dir() -> Optional[Path]:
|
||||
"""Where Prisma installs its private Node runtime, or None if unknowable."""
|
||||
override = os.getenv(NODEENV_CACHE_DIR_ENV_VAR)
|
||||
if override:
|
||||
return Path(override).absolute()
|
||||
if prisma_config is not None:
|
||||
try:
|
||||
return Path(prisma_config.nodeenv_cache_dir).absolute()
|
||||
except (OSError, ValueError) as e:
|
||||
logger.warning("Could not read the Prisma nodeenv cache dir: %s", e)
|
||||
try:
|
||||
return Path.home() / ".cache" / "prisma-python" / "nodeenv"
|
||||
except RuntimeError:
|
||||
logger.warning(
|
||||
"No resolvable home directory, cannot locate the Prisma nodeenv cache"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def node_binary_path(cache_dir: Path) -> Path:
|
||||
"""Path the Node binary occupies once the toolchain is fully installed."""
|
||||
if os.name == "nt":
|
||||
return cache_dir / "Scripts" / "node.exe"
|
||||
return cache_dir / "bin" / "node"
|
||||
|
||||
|
||||
def heal_incomplete_nodeenv_cache() -> bool:
|
||||
"""Delete a nodeenv cache directory left without a Node binary.
|
||||
|
||||
Returns True when a half-installed toolchain was removed, so the next
|
||||
Prisma invocation reinstalls it instead of failing on a missing binary.
|
||||
"""
|
||||
cache_dir = nodeenv_cache_dir()
|
||||
if cache_dir is None:
|
||||
return False
|
||||
try:
|
||||
if not cache_dir.is_dir() or node_binary_path(cache_dir).exists():
|
||||
return False
|
||||
except OSError as e:
|
||||
logger.warning("Could not inspect the Node toolchain at %s: %s", cache_dir, e)
|
||||
return False
|
||||
logger.warning(
|
||||
"Node toolchain at %s has no %s, so a previous install was interrupted. "
|
||||
"Removing it so it can be reinstalled.",
|
||||
cache_dir,
|
||||
node_binary_path(cache_dir).name,
|
||||
)
|
||||
try:
|
||||
shutil.rmtree(cache_dir)
|
||||
except OSError as e:
|
||||
logger.warning("Could not remove %s: %s", cache_dir, e)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def ensure_prisma_toolchain(
|
||||
prisma_command: str, prisma_env: dict[str, str]
|
||||
) -> ToolchainBootstrap:
|
||||
"""Install whatever the Prisma CLI needs to run, under its own timeout.
|
||||
|
||||
Never raises. A toolchain that cannot be prepared is reported so the
|
||||
caller can go on and let the real Prisma command produce the real error.
|
||||
"""
|
||||
healed = heal_incomplete_nodeenv_cache()
|
||||
timeout = prisma_bootstrap_timeout()
|
||||
logger.info("Preparing the Prisma CLI toolchain (timeout %ss)", timeout)
|
||||
try:
|
||||
subprocess.run(
|
||||
[prisma_command, BOOTSTRAP_ARG],
|
||||
timeout=timeout,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=prisma_env,
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
"Preparing the Prisma CLI toolchain timed out after %ss. Raise %s "
|
||||
"if this machine needs longer to install it.",
|
||||
timeout,
|
||||
PRISMA_BOOTSTRAP_TIMEOUT_ENV_VAR,
|
||||
)
|
||||
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning("Preparing the Prisma CLI toolchain failed: %s", e.stderr)
|
||||
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
|
||||
except OSError as e:
|
||||
logger.warning("Could not run the Prisma CLI: %s", e)
|
||||
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=False)
|
||||
logger.info("Prisma CLI toolchain ready")
|
||||
return ToolchainBootstrap(healed_incomplete_cache=healed, ready=True)
|
||||
0
litellm-proxy-extras/litellm_proxy_extras/py.typed
Normal file
0
litellm-proxy-extras/litellm_proxy_extras/py.typed
Normal file
|
|
@ -16,6 +16,7 @@ import tempfile
|
|||
from pathlib import Path
|
||||
|
||||
from litellm_proxy_extras._logging import logger
|
||||
from litellm_proxy_extras.prisma_toolchain import prisma_command_timeout
|
||||
|
||||
REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"
|
||||
|
||||
|
|
@ -75,7 +76,7 @@ def apply_replica_identity_full(
|
|||
"--schema",
|
||||
schema_path,
|
||||
],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -1110,6 +1118,26 @@ model LiteLLM_DailyToolSpend {
|
|||
@@id([date, tool_name])
|
||||
}
|
||||
|
||||
// Gateway request counts recorded at the ASGI edge by
|
||||
// BillableRequestMetricsMiddleware. This is the source of truth for SGR
|
||||
// (successful gateway requests): it counts what the proxy actually answered,
|
||||
// independent of whether the request reached litellm's logging callbacks.
|
||||
// The key carries no deployment or caller dimension. Every part of it is
|
||||
// chosen by the proxy and drawn from a closed set, so the table is bounded by
|
||||
// (days x categories x routes) rather than by anything a caller can vary.
|
||||
model LiteLLM_DailyGatewayRequests {
|
||||
date String
|
||||
category String
|
||||
route String
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@id([date, category, route])
|
||||
@@index([date])
|
||||
}
|
||||
|
||||
// Prompt table for storing prompt configurations
|
||||
model LiteLLM_PromptTable {
|
||||
id String @id @default(uuid())
|
||||
|
|
@ -1385,6 +1413,37 @@ model LiteLLM_AdaptiveRouterSession {
|
|||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
||||
model LiteLLM_AutoRouterSession {
|
||||
api_key String
|
||||
session_id String
|
||||
router_name String
|
||||
router_type String
|
||||
first_turn_at DateTime
|
||||
last_turn_at DateTime
|
||||
last_model String
|
||||
models Json @default("{}")
|
||||
turns Int @default(0)
|
||||
unordered_turns Int @default(0)
|
||||
covered_turns Int @default(0)
|
||||
cache_hits Int @default(0)
|
||||
same_model_turns Int @default(0)
|
||||
same_model_hits Int @default(0)
|
||||
first_visit_turns Int @default(0)
|
||||
first_visit_hits Int @default(0)
|
||||
return_turns Int @default(0)
|
||||
return_hits Int @default(0)
|
||||
return_expired_misses Int @default(0)
|
||||
return_within_ttl_misses Int @default(0)
|
||||
ttl_5m_turns Int @default(0)
|
||||
ttl_1h_turns Int @default(0)
|
||||
total_tokens BigInt @default(0)
|
||||
spend Float @default(0)
|
||||
saved_spend Float @default(0)
|
||||
|
||||
@@id([api_key, session_id, router_name])
|
||||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Workflow Run Tracking
|
||||
//
|
||||
|
|
|
|||
|
|
@ -14,6 +14,10 @@ from litellm_proxy_extras.replica_identity import (
|
|||
REPLICA_IDENTITY_FULL_ENV_VAR,
|
||||
apply_replica_identity_full,
|
||||
)
|
||||
from litellm_proxy_extras.prisma_toolchain import (
|
||||
ensure_prisma_toolchain,
|
||||
prisma_command_timeout,
|
||||
)
|
||||
|
||||
|
||||
def str_to_bool(value: Optional[str]) -> bool:
|
||||
|
|
@ -142,7 +146,7 @@ class ProxyExtrasDBManager:
|
|||
],
|
||||
stdout=open(migration_file, "w"),
|
||||
check=True,
|
||||
timeout=30,
|
||||
timeout=prisma_command_timeout(),
|
||||
env=prisma_env,
|
||||
)
|
||||
|
||||
|
|
@ -157,7 +161,7 @@ class ProxyExtrasDBManager:
|
|||
"0_init",
|
||||
],
|
||||
check=True,
|
||||
timeout=30,
|
||||
timeout=prisma_command_timeout(),
|
||||
env=prisma_env,
|
||||
)
|
||||
|
||||
|
|
@ -193,7 +197,7 @@ class ProxyExtrasDBManager:
|
|||
"--rolled-back",
|
||||
migration_name,
|
||||
],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
env=prisma_env,
|
||||
|
|
@ -205,7 +209,7 @@ class ProxyExtrasDBManager:
|
|||
prisma_env = _get_prisma_env()
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "resolve", "--applied", migration_name],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
env=prisma_env,
|
||||
|
|
@ -303,7 +307,7 @@ class ProxyExtrasDBManager:
|
|||
"--script",
|
||||
],
|
||||
check=True,
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
stdout=f,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
|
|
@ -335,7 +339,7 @@ class ProxyExtrasDBManager:
|
|||
"--schema",
|
||||
schema_path,
|
||||
],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -364,7 +368,7 @@ class ProxyExtrasDBManager:
|
|||
"--schema",
|
||||
schema_path,
|
||||
],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -393,7 +397,7 @@ class ProxyExtrasDBManager:
|
|||
"--applied",
|
||||
migration_name,
|
||||
],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -530,7 +534,7 @@ class ProxyExtrasDBManager:
|
|||
try:
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
|
|
@ -555,7 +559,7 @@ class ProxyExtrasDBManager:
|
|||
try:
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -731,6 +735,9 @@ class ProxyExtrasDBManager:
|
|||
Returns:
|
||||
bool: True if setup was successful, False otherwise
|
||||
"""
|
||||
ensure_prisma_toolchain(
|
||||
prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env()
|
||||
)
|
||||
migrated = ProxyExtrasDBManager._run_migrations(
|
||||
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
|
||||
)
|
||||
|
|
@ -757,7 +764,7 @@ class ProxyExtrasDBManager:
|
|||
# Set migrations directory for Prisma
|
||||
result = subprocess.run(
|
||||
[_get_prisma_command(), "migrate", "deploy"],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -840,7 +847,7 @@ class ProxyExtrasDBManager:
|
|||
"--rolled-back",
|
||||
failed_migration,
|
||||
],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
|
|
@ -968,7 +975,7 @@ class ProxyExtrasDBManager:
|
|||
# Use prisma db push with increased timeout
|
||||
subprocess.run(
|
||||
[_get_prisma_command(), "db", "push", "--accept-data-loss"],
|
||||
timeout=60,
|
||||
timeout=prisma_command_timeout(),
|
||||
check=True,
|
||||
)
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -243,6 +244,7 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
|
|||
# Or via `litellm_settings.strip_anthropic_total_tokens: true` in
|
||||
# config.yaml.
|
||||
strip_anthropic_total_tokens: bool = False
|
||||
anthropic_sse_ping_interval_seconds: float = 15.0
|
||||
route_all_chat_openai_to_responses: bool = (
|
||||
os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true"
|
||||
) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge
|
||||
|
|
@ -264,6 +266,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
|
||||
|
|
@ -680,12 +683,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
|
||||
|
||||
|
||||
|
|
@ -702,9 +705,8 @@ def is_openai_finetune_model(key: str) -> bool:
|
|||
return key.startswith("ft:") and not key.count(":") > 1
|
||||
|
||||
|
||||
def add_known_models(model_cost_map: Optional[Dict] = None):
|
||||
_map = model_cost_map if model_cost_map is not None else model_cost
|
||||
for key, value in _map.items():
|
||||
def _populate_provider_model_sets(model_cost_map: Dict) -> None:
|
||||
for key, value in model_cost_map.items():
|
||||
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
|
||||
open_ai_chat_completion_models.add(key)
|
||||
elif value.get("litellm_provider") == "text-completion-openai":
|
||||
|
|
@ -947,7 +949,16 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
|||
bedrock_mantle_models.add(key)
|
||||
|
||||
|
||||
add_known_models()
|
||||
def add_known_models(model_cost_map: Optional[Dict] = None):
|
||||
"""Fold `model_cost_map` (defaults to `litellm.model_cost`) into the per-provider model sets,
|
||||
then refresh `models_by_provider` from those sets so the additions reach wildcard expansion.
|
||||
The refresh updates the dict in place, so references captured before a reload stay live.
|
||||
"""
|
||||
_populate_provider_model_sets(model_cost_map if model_cost_map is not None else model_cost)
|
||||
models_by_provider.update(_build_models_by_provider())
|
||||
|
||||
|
||||
_populate_provider_model_sets(model_cost)
|
||||
# known openai compatible endpoints - we'll eventually move this list to the model_prices_and_context_window.json dictionary
|
||||
|
||||
# this is maintained for Exception Mapping
|
||||
|
|
@ -1069,112 +1080,116 @@ model_list_set = set(model_list)
|
|||
# provider_list is lazy-loaded via __getattr__ to avoid importing LlmProviders at import time
|
||||
|
||||
|
||||
models_by_provider: dict = {
|
||||
"openai": open_ai_chat_completion_models | open_ai_text_completion_models,
|
||||
"text-completion-openai": open_ai_text_completion_models,
|
||||
"cohere": cohere_models | cohere_chat_models,
|
||||
"cohere_chat": cohere_chat_models,
|
||||
"anthropic": anthropic_models,
|
||||
"replicate": replicate_models,
|
||||
"huggingface": huggingface_models,
|
||||
"together_ai": together_ai_models,
|
||||
"baseten": baseten_models,
|
||||
"openrouter": openrouter_models,
|
||||
"vercel_ai_gateway": vercel_ai_gateway_models,
|
||||
"datarobot": datarobot_models,
|
||||
"vertex_ai": vertex_chat_models
|
||||
| vertex_text_models
|
||||
| vertex_anthropic_models
|
||||
| vertex_vision_models
|
||||
| vertex_language_models
|
||||
| vertex_deepseek_models
|
||||
| vertex_minimax_models
|
||||
| vertex_moonshot_models
|
||||
| vertex_zai_models,
|
||||
"ai21": ai21_models,
|
||||
"bedrock": bedrock_models | bedrock_converse_models,
|
||||
"petals": petals_models,
|
||||
"ollama": ollama_models,
|
||||
"ollama_chat": ollama_models,
|
||||
"deepinfra": deepinfra_models,
|
||||
"perplexity": perplexity_models,
|
||||
"maritalk": maritalk_models,
|
||||
"watsonx": watsonx_models,
|
||||
"gemini": gemini_models,
|
||||
"fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models,
|
||||
"aleph_alpha": aleph_alpha_models,
|
||||
"text-completion-codestral": text_completion_codestral_models,
|
||||
"text-completion-inception": text_completion_inception_models,
|
||||
"xai": xai_models,
|
||||
"zai": zai_models,
|
||||
"fal_ai": fal_ai_models,
|
||||
"deepseek": deepseek_models,
|
||||
"tencent": tencent_models,
|
||||
"runwayml": runwayml_models,
|
||||
"mistral": mistral_chat_models,
|
||||
"azure_ai": azure_ai_models,
|
||||
"voyage": voyage_models,
|
||||
"infinity": infinity_models,
|
||||
"databricks": databricks_models,
|
||||
"cloudflare": cloudflare_models,
|
||||
"codestral": codestral_models,
|
||||
"nlp_cloud": nlp_cloud_models,
|
||||
"friendliai": friendliai_models,
|
||||
"palm": palm_models,
|
||||
"groq": groq_models,
|
||||
"azure": azure_models | azure_text_models,
|
||||
"azure_anthropic": azure_anthropic_models,
|
||||
"azure_text": azure_text_models,
|
||||
"anyscale": anyscale_models,
|
||||
"cerebras": cerebras_models,
|
||||
"galadriel": galadriel_models,
|
||||
"nvidia_nim": nvidia_nim_models,
|
||||
"nvidia_riva": nvidia_riva_models,
|
||||
"soniox": soniox_models,
|
||||
"sambanova": sambanova_models | sambanova_embedding_models,
|
||||
"novita": novita_models,
|
||||
"nebius": nebius_models | nebius_embedding_models,
|
||||
"aiml": aiml_models,
|
||||
"assemblyai": assemblyai_models,
|
||||
"jina_ai": jina_ai_models,
|
||||
"snowflake": snowflake_models,
|
||||
"gradient_ai": gradient_ai_models,
|
||||
"meta_llama": llama_models,
|
||||
"nscale": nscale_models,
|
||||
"featherless_ai": featherless_ai_models,
|
||||
"deepgram": deepgram_models,
|
||||
"elevenlabs": elevenlabs_models,
|
||||
"heroku": heroku_models,
|
||||
"dashscope": dashscope_models,
|
||||
"modelscope": modelscope_models,
|
||||
"moonshot": moonshot_models,
|
||||
"publicai": publicai_models,
|
||||
"darkbloom": darkbloom_models,
|
||||
"v0": v0_models,
|
||||
"morph": morph_models,
|
||||
"lambda_ai": lambda_ai_models,
|
||||
"inception": inception_models,
|
||||
"hyperbolic": hyperbolic_models,
|
||||
"black_forest_labs": black_forest_labs_models,
|
||||
"recraft": recraft_models,
|
||||
"cometapi": cometapi_models,
|
||||
"oci": oci_models,
|
||||
"volcengine": volcengine_models,
|
||||
"wandb": wandb_models,
|
||||
"ovhcloud": ovhcloud_models | ovhcloud_embedding_models,
|
||||
"lemonade": lemonade_models,
|
||||
"clarifai": clarifai_models,
|
||||
"amazon_nova": amazon_nova_models,
|
||||
"stability": stability_models,
|
||||
"github_copilot": github_copilot_models,
|
||||
"chatgpt": chatgpt_models,
|
||||
"minimax": minimax_models,
|
||||
"aws_polly": aws_polly_models,
|
||||
"gigachat": gigachat_models,
|
||||
"llamagate": llamagate_models,
|
||||
"reducto": reducto_models,
|
||||
"bedrock_mantle": bedrock_mantle_models,
|
||||
}
|
||||
def _build_models_by_provider() -> dict:
|
||||
return {
|
||||
"openai": open_ai_chat_completion_models | open_ai_text_completion_models,
|
||||
"text-completion-openai": open_ai_text_completion_models,
|
||||
"cohere": cohere_models | cohere_chat_models,
|
||||
"cohere_chat": cohere_chat_models,
|
||||
"anthropic": anthropic_models,
|
||||
"replicate": replicate_models,
|
||||
"huggingface": huggingface_models,
|
||||
"together_ai": together_ai_models,
|
||||
"baseten": baseten_models,
|
||||
"openrouter": openrouter_models,
|
||||
"vercel_ai_gateway": vercel_ai_gateway_models,
|
||||
"datarobot": datarobot_models,
|
||||
"vertex_ai": vertex_chat_models
|
||||
| vertex_text_models
|
||||
| vertex_anthropic_models
|
||||
| vertex_vision_models
|
||||
| vertex_language_models
|
||||
| vertex_deepseek_models
|
||||
| vertex_minimax_models
|
||||
| vertex_moonshot_models
|
||||
| vertex_zai_models,
|
||||
"ai21": ai21_models,
|
||||
"bedrock": bedrock_models | bedrock_converse_models,
|
||||
"petals": petals_models,
|
||||
"ollama": ollama_models,
|
||||
"ollama_chat": ollama_models,
|
||||
"deepinfra": deepinfra_models,
|
||||
"perplexity": perplexity_models,
|
||||
"maritalk": maritalk_models,
|
||||
"watsonx": watsonx_models,
|
||||
"gemini": gemini_models,
|
||||
"fireworks_ai": fireworks_ai_models | fireworks_ai_embedding_models,
|
||||
"aleph_alpha": aleph_alpha_models,
|
||||
"text-completion-codestral": text_completion_codestral_models,
|
||||
"text-completion-inception": text_completion_inception_models,
|
||||
"xai": xai_models,
|
||||
"zai": zai_models,
|
||||
"fal_ai": fal_ai_models,
|
||||
"deepseek": deepseek_models,
|
||||
"tencent": tencent_models,
|
||||
"runwayml": runwayml_models,
|
||||
"mistral": mistral_chat_models,
|
||||
"azure_ai": azure_ai_models,
|
||||
"voyage": voyage_models,
|
||||
"infinity": infinity_models,
|
||||
"databricks": databricks_models,
|
||||
"cloudflare": cloudflare_models,
|
||||
"codestral": codestral_models,
|
||||
"nlp_cloud": nlp_cloud_models,
|
||||
"friendliai": friendliai_models,
|
||||
"palm": palm_models,
|
||||
"groq": groq_models,
|
||||
"azure": azure_models | azure_text_models,
|
||||
"azure_anthropic": azure_anthropic_models,
|
||||
"azure_text": azure_text_models,
|
||||
"anyscale": anyscale_models,
|
||||
"cerebras": cerebras_models,
|
||||
"galadriel": galadriel_models,
|
||||
"nvidia_nim": nvidia_nim_models,
|
||||
"nvidia_riva": nvidia_riva_models,
|
||||
"soniox": soniox_models,
|
||||
"sambanova": sambanova_models | sambanova_embedding_models,
|
||||
"novita": novita_models,
|
||||
"nebius": nebius_models | nebius_embedding_models,
|
||||
"aiml": aiml_models,
|
||||
"assemblyai": assemblyai_models,
|
||||
"jina_ai": jina_ai_models,
|
||||
"snowflake": snowflake_models,
|
||||
"gradient_ai": gradient_ai_models,
|
||||
"meta_llama": llama_models,
|
||||
"nscale": nscale_models,
|
||||
"featherless_ai": featherless_ai_models,
|
||||
"deepgram": deepgram_models,
|
||||
"elevenlabs": elevenlabs_models,
|
||||
"heroku": heroku_models,
|
||||
"dashscope": dashscope_models,
|
||||
"modelscope": modelscope_models,
|
||||
"moonshot": moonshot_models,
|
||||
"publicai": publicai_models,
|
||||
"darkbloom": darkbloom_models,
|
||||
"v0": v0_models,
|
||||
"morph": morph_models,
|
||||
"lambda_ai": lambda_ai_models,
|
||||
"inception": inception_models,
|
||||
"hyperbolic": hyperbolic_models,
|
||||
"black_forest_labs": black_forest_labs_models,
|
||||
"recraft": recraft_models,
|
||||
"cometapi": cometapi_models,
|
||||
"oci": oci_models,
|
||||
"volcengine": volcengine_models,
|
||||
"wandb": wandb_models,
|
||||
"ovhcloud": ovhcloud_models | ovhcloud_embedding_models,
|
||||
"lemonade": lemonade_models,
|
||||
"clarifai": clarifai_models,
|
||||
"amazon_nova": amazon_nova_models,
|
||||
"stability": stability_models,
|
||||
"github_copilot": github_copilot_models,
|
||||
"chatgpt": chatgpt_models,
|
||||
"minimax": minimax_models,
|
||||
"aws_polly": aws_polly_models,
|
||||
"gigachat": gigachat_models,
|
||||
"llamagate": llamagate_models,
|
||||
"reducto": reducto_models,
|
||||
"bedrock_mantle": bedrock_mantle_models,
|
||||
}
|
||||
|
||||
|
||||
models_by_provider: dict = _build_models_by_provider()
|
||||
|
||||
# mapping for those models which have larger equivalents
|
||||
longer_context_model_fallback_dict: dict = {
|
||||
|
|
@ -1267,8 +1282,8 @@ from .llms.xai.common_utils import XAIModelInfo
|
|||
from litellm.types.utils import LlmProviders
|
||||
|
||||
## Lazy loading this is not straightforward, will leave it here for now.
|
||||
from .main import * # type: ignore
|
||||
from .compression import compress # type: ignore[no-redef]
|
||||
from .main import *
|
||||
from .compression import compress
|
||||
|
||||
# Skills API
|
||||
from .skills.main import (
|
||||
|
|
@ -1339,7 +1354,7 @@ from .assistants.main import *
|
|||
from .batches.main import *
|
||||
from .images.main import *
|
||||
from .videos.main import *
|
||||
from .batch_completion.main import * # type: ignore
|
||||
from .batch_completion.main import *
|
||||
from .rerank_api.main import *
|
||||
from .llms.anthropic.experimental_pass_through.messages.handler import *
|
||||
from .responses.main import *
|
||||
|
|
@ -2052,7 +2067,7 @@ if TYPE_CHECKING:
|
|||
supports_reasoning: Callable[..., bool]
|
||||
acreate: Callable[..., Any]
|
||||
get_max_tokens: Callable[..., int]
|
||||
get_model_info: Callable[..., _ModelInfoType] # type: ignore[no-redef]
|
||||
get_model_info: Callable[..., _ModelInfoType]
|
||||
register_prompt_template: Callable[..., None]
|
||||
validate_environment: Callable[..., dict]
|
||||
check_valid_key: Callable[..., bool]
|
||||
|
|
@ -2139,18 +2154,18 @@ 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
|
||||
if name == "encoding":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
# Check if already cached
|
||||
if "encoding" not in _globals:
|
||||
from .main import encoding as _encoding
|
||||
|
|
@ -2160,9 +2175,9 @@ def __getattr__(name: str) -> Any:
|
|||
|
||||
# Lazy load bedrock_tool_name_mappings instance
|
||||
if name == "bedrock_tool_name_mappings":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
# Check if already cached
|
||||
if "bedrock_tool_name_mappings" not in _globals:
|
||||
from .llms.bedrock.chat.invoke_handler import (
|
||||
|
|
@ -2174,9 +2189,9 @@ def __getattr__(name: str) -> Any:
|
|||
|
||||
# Lazy load AzureOpenAIError exception class
|
||||
if name == "AzureOpenAIError":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
# Check if already cached
|
||||
if "AzureOpenAIError" not in _globals:
|
||||
from .llms.azure.common_utils import AzureOpenAIError as _AzureOpenAIError
|
||||
|
|
@ -2186,9 +2201,9 @@ def __getattr__(name: str) -> Any:
|
|||
|
||||
# Lazy load openaiOSeriesConfig instance
|
||||
if name == "openaiOSeriesConfig":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
if "openaiOSeriesConfig" not in _globals:
|
||||
# Import the config class and instantiate it
|
||||
config_class = __getattr__("OpenAIOSeriesConfig")
|
||||
|
|
@ -2196,7 +2211,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",
|
||||
|
|
@ -2204,9 +2219,9 @@ def __getattr__(name: str) -> Any:
|
|||
"nvidiaNimEmbeddingConfig": "NvidiaNimEmbeddingConfig",
|
||||
}
|
||||
if name in _config_instances:
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
if name not in _globals:
|
||||
# Import the config class and instantiate it
|
||||
config_class = __getattr__(_config_instances[name])
|
||||
|
|
@ -2219,9 +2234,9 @@ def __getattr__(name: str) -> Any:
|
|||
|
||||
# Lazy load provider_list
|
||||
if name == "provider_list":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
# Check if already cached
|
||||
if "provider_list" not in _globals:
|
||||
# LlmProviders is eagerly imported above, so we can import it directly
|
||||
|
|
@ -2232,33 +2247,33 @@ def __getattr__(name: str) -> Any:
|
|||
|
||||
# Lazy load priority_reservation_settings instance
|
||||
if name == "priority_reservation_settings":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
# 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"]
|
||||
|
||||
# Lazy load logging_callback_manager instance
|
||||
if name == "logging_callback_manager":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
# 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"]
|
||||
|
||||
# Lazy load _service_logger module
|
||||
if name == "_service_logger":
|
||||
from ._lazy_imports import _get_litellm_globals
|
||||
from ._lazy_imports import get_litellm_globals
|
||||
|
||||
_globals = _get_litellm_globals()
|
||||
_globals = get_litellm_globals()
|
||||
# Check if already cached
|
||||
if "_service_logger" not in _globals:
|
||||
# Import the module lazily
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -18,11 +18,12 @@ until they're actually needed.
|
|||
import importlib
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
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 (
|
||||
# Import maps
|
||||
_BEDROCK_TYPES_IMPORT_MAP,
|
||||
_CACHING_IMPORT_MAP,
|
||||
_COST_CALCULATOR_IMPORT_MAP,
|
||||
|
|
@ -33,12 +34,11 @@ from ._lazy_imports_registry import (
|
|||
_TOKEN_COUNTER_IMPORT_MAP,
|
||||
_TYPES_IMPORT_MAP,
|
||||
_TYPES_UTILS_IMPORT_MAP,
|
||||
# Import maps
|
||||
_UTILS_IMPORT_MAP,
|
||||
_UTILS_MODULE_IMPORT_MAP,
|
||||
# Name tuples
|
||||
BEDROCK_TYPES_NAMES,
|
||||
CACHING_NAMES,
|
||||
# Name tuples
|
||||
COST_CALCULATOR_NAMES,
|
||||
DOTPROMPT_NAMES,
|
||||
HTTP_HANDLER_NAMES,
|
||||
|
|
@ -54,7 +54,7 @@ from ._lazy_imports_registry import (
|
|||
)
|
||||
|
||||
|
||||
def _get_litellm_globals() -> dict:
|
||||
def get_litellm_globals() -> dict:
|
||||
"""
|
||||
Get the globals dictionary of the litellm module.
|
||||
|
||||
|
|
@ -233,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:
|
||||
|
|
@ -255,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
|
||||
|
|
@ -332,14 +332,14 @@ def _lazy_import_utils_module(name: str) -> Any:
|
|||
Handler for utils module lazy imports.
|
||||
|
||||
This uses a custom implementation because utils module needs to use
|
||||
_get_utils_globals() instead of _get_litellm_globals() for caching.
|
||||
_get_utils_globals() instead of get_litellm_globals() for caching.
|
||||
"""
|
||||
# Check if this attribute exists in our map
|
||||
if name not in _UTILS_MODULE_IMPORT_MAP:
|
||||
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:
|
||||
|
|
@ -355,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
|
||||
|
|
@ -379,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":
|
||||
|
|
@ -396,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
|
||||
|
||||
|
|
@ -412,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
|
||||
|
|
@ -420,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,
|
||||
)
|
||||
|
|
@ -438,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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
@ -1459,32 +1461,30 @@ _UTILS_MODULE_IMPORT_MAP = {
|
|||
|
||||
# Export all name tuples and import maps for use in _lazy_imports.py
|
||||
__all__ = [
|
||||
# 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",
|
||||
"COST_CALCULATOR_NAMES",
|
||||
"DOTPROMPT_NAMES",
|
||||
"HTTP_HANDLER_NAMES",
|
||||
"LITELLM_LOGGING_NAMES",
|
||||
"LLM_CLIENT_CACHE_NAMES",
|
||||
"LLM_CONFIG_NAMES",
|
||||
"TYPES_NAMES",
|
||||
"LLM_PROVIDER_LOGIC_NAMES",
|
||||
"TOKEN_COUNTER_NAMES",
|
||||
"TYPES_NAMES",
|
||||
"TYPES_UTILS_NAMES",
|
||||
"UTILS_MODULE_NAMES",
|
||||
# Import maps
|
||||
"_UTILS_IMPORT_MAP",
|
||||
"_COST_CALCULATOR_IMPORT_MAP",
|
||||
"_TYPES_UTILS_IMPORT_MAP",
|
||||
"_TOKEN_COUNTER_IMPORT_MAP",
|
||||
"UTILS_NAMES",
|
||||
"_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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import os
|
|||
import sys
|
||||
from datetime import datetime
|
||||
from logging import Formatter
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
|
|
@ -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,14 +74,14 @@ 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)
|
||||
|
||||
|
|
@ -94,10 +94,10 @@ def _try_parse_json_message(message: str) -> dict[str, Any] | None:
|
|||
"""
|
||||
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
|
||||
|
|
@ -144,7 +144,7 @@ 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):
|
||||
|
|
@ -153,12 +153,12 @@ class JsonFormatter(Formatter):
|
|||
|
||||
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,54 @@ 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)
|
||||
|
||||
|
||||
_REDACTED_THIRD_PARTY_LOGGERS: Final[tuple[str, ...]] = (
|
||||
"apscheduler.executors.default",
|
||||
"apscheduler.scheduler",
|
||||
"asyncio",
|
||||
"backoff",
|
||||
"httpx",
|
||||
"uvicorn.error",
|
||||
)
|
||||
|
||||
|
||||
def _redact_third_party_loggers() -> None:
|
||||
"""Extend secret redaction to records litellm does not emit directly.
|
||||
|
||||
litellm's own loggers are covered by the filter on their shared handler, but a
|
||||
litellm value can also reach a log record through a dependency that logs on its
|
||||
own logger. Those records never pass through a litellm handler.
|
||||
|
||||
The filter is attached to each emitting logger rather than to the root logger or
|
||||
to root's handlers. `Logger.handle` applies the emitting logger's filters before
|
||||
any handler runs, so redaction happens once, at the earliest point in the
|
||||
record's life, and covers every downstream handler regardless of who owns it.
|
||||
The alternatives do not hold: `callHandlers` consults ancestors for handlers but
|
||||
never for filters, so a filter on the root logger never sees these records at
|
||||
all, and a filter on a root handler only covers that one handler, leaving
|
||||
handlers registered earlier or on the emitting logger itself untouched.
|
||||
|
||||
Each name is the exact logger a dependency emits on; a parent name would not
|
||||
cover its children, for the same reason the root logger does not.
|
||||
"""
|
||||
for name in _REDACTED_THIRD_PARTY_LOGGERS:
|
||||
logging.getLogger(name).addFilter(_secret_filter)
|
||||
|
||||
|
||||
# Call the suppression function
|
||||
_suppress_loggers()
|
||||
_redact_third_party_loggers()
|
||||
|
||||
ALL_LOGGERS = [
|
||||
ALL_LOGGERS: Final = [
|
||||
logging.getLogger(),
|
||||
verbose_logger,
|
||||
verbose_router_logger,
|
||||
|
|
@ -293,11 +327,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 +359,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 +418,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
|
||||
|
|
|
|||
|
|
@ -13,9 +13,10 @@ import json
|
|||
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import redis # type: ignore
|
||||
import redis.asyncio as async_redis # type: ignore
|
||||
import redis
|
||||
import redis.asyncio as async_redis
|
||||
|
||||
from litellm import get_secret, get_secret_str
|
||||
from litellm._redis_credential_provider import (
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -92,9 +93,9 @@ def _get_redis_url_kwargs(client: type | None = 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: type | None = 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: type | None = 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,17 +143,17 @@ 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
|
||||
value = get_secret(k, default_value=None)
|
||||
if value is not None:
|
||||
return_dict[v] = value
|
||||
return return_dict
|
||||
|
|
@ -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:
|
||||
|
|
@ -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(
|
||||
|
|
@ -253,12 +254,12 @@ 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
|
||||
|
||||
|
||||
|
|
@ -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:
|
||||
|
|
@ -316,7 +317,7 @@ def create_azure_ad_redis_connect_func(
|
|||
# AzureADCredentialProvider for refresh-aware token retrieval. The raw
|
||||
# client_id/tenant_id/secret are intentionally NOT exposed here — the
|
||||
# credential closure already holds them.
|
||||
ad_connect._azure_credential = credential # type: ignore[attr-defined]
|
||||
ad_connect._azure_credential = credential
|
||||
return ad_connect
|
||||
|
||||
|
||||
|
|
@ -350,26 +351,26 @@ def _get_redis_client_logic(**env_overrides):
|
|||
for k, v in env_overrides.items():
|
||||
if isinstance(v, str) and v.startswith("os.environ/"):
|
||||
v = v.replace("os.environ/", "")
|
||||
value = get_secret(v) # type: ignore
|
||||
value = get_secret(v)
|
||||
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: str | list | None = 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(
|
||||
"REDIS_CLUSTER_NODES"
|
||||
)
|
||||
|
||||
|
|
@ -380,30 +381,28 @@ def _get_redis_client_logic(**env_overrides):
|
|||
elif _startup_nodes is None:
|
||||
redis_kwargs.pop("startup_nodes", None)
|
||||
|
||||
_sentinel_nodes: str | list | None = 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(
|
||||
"REDIS_SENTINEL_NODES"
|
||||
)
|
||||
|
||||
if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str):
|
||||
redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes)
|
||||
|
||||
_sentinel_password: str | None = 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: str | None = redis_kwargs.get("service_name", None) or get_secret( # type: ignore
|
||||
"REDIS_SERVICE_NAME"
|
||||
)
|
||||
_service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret("REDIS_SERVICE_NAME")
|
||||
|
||||
if _service_name is not None:
|
||||
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.")
|
||||
|
|
@ -411,7 +410,7 @@ def _get_redis_client_logic(**env_overrides):
|
|||
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
|
||||
)
|
||||
# Store GCP service account in redis_connect_func for async cluster access
|
||||
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account # type: ignore[attr-defined]
|
||||
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
|
||||
|
||||
# Remove GCP-specific kwargs that shouldn't be passed to Redis client
|
||||
redis_kwargs.pop("gcp_service_account", None)
|
||||
|
|
@ -422,9 +421,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 +432,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(
|
||||
|
|
@ -448,7 +447,7 @@ def _get_redis_client_logic(**env_overrides):
|
|||
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
|
||||
# are intentionally NOT exposed on the function to avoid leaking
|
||||
# credentials via inspection or logging.
|
||||
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True # type: ignore[attr-defined]
|
||||
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
|
||||
|
||||
# Always remove Azure-specific kwargs that shouldn't be passed to Redis client
|
||||
redis_kwargs.pop("azure_redis_ad_token", None)
|
||||
|
|
@ -480,7 +479,7 @@ def _get_redis_client_logic(**env_overrides):
|
|||
|
||||
|
||||
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
|
||||
_redis_cluster_nodes_in_env: str | None = get_secret("REDIS_CLUSTER_NODES") # type: ignore
|
||||
_redis_cluster_nodes_in_env: Final[str | None] = get_secret("REDIS_CLUSTER_NODES")
|
||||
if _redis_cluster_nodes_in_env is not None:
|
||||
try:
|
||||
redis_kwargs["startup_nodes"] = json.loads(_redis_cluster_nodes_in_env)
|
||||
|
|
@ -492,24 +491,24 @@ 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))
|
||||
|
||||
cluster_kwargs.pop("startup_nodes", None)
|
||||
return redis.RedisCluster(startup_nodes=new_startup_nodes, **cluster_kwargs) # type: ignore
|
||||
return redis.RedisCluster(startup_nodes=new_startup_nodes, **cluster_kwargs)
|
||||
|
||||
|
||||
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]
|
||||
|
|
@ -518,12 +517,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:
|
||||
|
|
@ -532,7 +531,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,
|
||||
)
|
||||
|
|
@ -543,12 +542,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:
|
||||
|
|
@ -557,7 +556,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,
|
||||
)
|
||||
|
|
@ -568,14 +567,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]
|
||||
|
|
@ -593,13 +592,13 @@ def get_redis_async_client(
|
|||
connection_pool: async_redis.BlockingConnectionPool | None = None,
|
||||
**env_overrides,
|
||||
) -> async_redis.Redis | async_redis.RedisCluster:
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
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]
|
||||
|
|
@ -621,7 +620,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))
|
||||
|
|
@ -635,9 +634,9 @@ 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
|
||||
**cluster_kwargs,
|
||||
)
|
||||
|
||||
return cluster_client
|
||||
|
|
@ -646,12 +645,14 @@ 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(f"REDIS: ignoring argument: {arg}. Not an allowed async_redis.Redis.from_url arg.")
|
||||
verbose_logger.debug(
|
||||
"REDIS: ignoring argument: %s. Not an allowed async_redis.Redis.from_url arg.", arg
|
||||
)
|
||||
return async_redis.Redis.from_url(**url_kwargs)
|
||||
|
||||
# Check for Redis Sentinel
|
||||
|
|
@ -684,15 +685,15 @@ def get_redis_async_client(
|
|||
def get_redis_connection_pool(
|
||||
**env_overrides,
|
||||
) -> async_redis.BlockingConnectionPool | None:
|
||||
redis_kwargs = _get_redis_client_logic(**env_overrides)
|
||||
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:
|
||||
|
|
@ -708,7 +709,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,
|
||||
|
|
@ -735,7 +736,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()
|
||||
|
|
@ -744,10 +745,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",
|
||||
|
|
@ -784,7 +785,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",
|
||||
|
|
@ -805,6 +806,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)
|
||||
|
|
|
|||
|
|
@ -1,21 +1,21 @@
|
|||
import asyncio
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
from redis.credentials import CredentialProvider # type: ignore[attr-defined]
|
||||
from redis.credentials import CredentialProvider
|
||||
|
||||
# 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)
|
||||
|
||||
|
||||
|
|
@ -96,11 +96,11 @@ class GCPIAMCredentialProvider(CredentialProvider):
|
|||
self._gcp_service_account = gcp_service_account
|
||||
|
||||
def get_credentials(self) -> tuple[str]:
|
||||
token = _get_cached_gcp_iam_token(self._gcp_service_account)
|
||||
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)
|
||||
token: Final = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account)
|
||||
return (token,)
|
||||
|
||||
|
||||
|
|
@ -120,13 +120,13 @@ class AzureADCredentialProvider(CredentialProvider):
|
|||
self._username = username
|
||||
|
||||
def get_credentials(self) -> tuple[str] | tuple[str, str]:
|
||||
token = self._credential.get_token(AZURE_REDIS_SCOPE).token
|
||||
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) -> tuple[str] | tuple[str, str]:
|
||||
token_obj = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE)
|
||||
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,)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -16,7 +16,7 @@ if TYPE_CHECKING:
|
|||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
Span = Union[_Span, Any]
|
||||
Span = _Span | Any
|
||||
OTELClass = OpenTelemetry
|
||||
else:
|
||||
Span = Any
|
||||
|
|
@ -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):
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -267,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,
|
||||
|
|
@ -278,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()
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ Internal unified UUID helper.
|
|||
Always uses fastuuid for performance.
|
||||
"""
|
||||
|
||||
import fastuuid as _uuid # type: ignore
|
||||
import fastuuid as _uuid
|
||||
|
||||
# Expose a module-like alias so callers can use: uuid.uuid4()
|
||||
uuid = _uuid
|
||||
|
|
|
|||
|
|
@ -55,19 +55,15 @@ from litellm.a2a_protocol.main import (
|
|||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
|
||||
__all__ = [
|
||||
# Client
|
||||
"A2AClient",
|
||||
# Functions
|
||||
"asend_message",
|
||||
"send_message",
|
||||
"asend_message_streaming",
|
||||
"aget_agent_card",
|
||||
"create_a2a_client",
|
||||
# Response types
|
||||
"LiteLLMSendMessageResponse",
|
||||
# Exceptions
|
||||
"A2AError",
|
||||
"A2AConnectionError",
|
||||
"A2AAgentCardError",
|
||||
"A2AClient",
|
||||
"A2AConnectionError",
|
||||
"A2AError",
|
||||
"A2ALocalhostURLError",
|
||||
"LiteLLMSendMessageResponse",
|
||||
"aget_agent_card",
|
||||
"asend_message",
|
||||
"asend_message_streaming",
|
||||
"create_a2a_client",
|
||||
"send_message",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import LOCALHOST_URL_PATTERNS
|
||||
|
|
@ -18,8 +18,8 @@ AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent-card.json"
|
|||
PREV_AGENT_CARD_WELL_KNOWN_PATH: str = "/.well-known/agent.json"
|
||||
|
||||
try:
|
||||
from a2a.client import A2ACardResolver as _A2ACardResolver # type: ignore[no-redef]
|
||||
from a2a.utils.constants import ( # type: ignore[no-redef]
|
||||
from a2a.client import A2ACardResolver as _A2ACardResolver
|
||||
from a2a.utils.constants import (
|
||||
AGENT_CARD_WELL_KNOWN_PATH,
|
||||
PREV_AGENT_CARD_WELL_KNOWN_PATH,
|
||||
)
|
||||
|
|
@ -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,23 +86,23 @@ 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("/") + "/"
|
||||
|
||||
return agent_card
|
||||
|
||||
|
||||
class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc]
|
||||
class LiteLLMA2ACardResolver(_A2ACardResolver):
|
||||
"""
|
||||
Custom A2A card resolver that supports 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
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Provides a class-based interface for A2A agent invocation.
|
|||
"""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.types.agents import LiteLLMSendMessageResponse
|
||||
|
||||
|
|
@ -92,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(
|
||||
|
|
@ -101,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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -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(
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils.
|
|||
Maps A2A SDK exceptions to LiteLLM A2A exception types.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.card_resolver import (
|
||||
|
|
@ -29,9 +29,9 @@ try:
|
|||
A2A_SDK_AVAILABLE = True
|
||||
except ImportError:
|
||||
A2A_SDK_AVAILABLE = False
|
||||
Client = None # type: ignore[misc, assignment]
|
||||
ClientConfig = None # type: ignore[misc, assignment]
|
||||
create_client = None # type: ignore[misc, assignment]
|
||||
Client = None
|
||||
ClientConfig = None
|
||||
create_client = None
|
||||
|
||||
|
||||
class A2AExceptionCheckers:
|
||||
|
|
@ -53,7 +53,7 @@ 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
|
||||
|
|
@ -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",
|
||||
|
|
@ -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,20 +205,20 @@ 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,
|
||||
streaming=is_streaming,
|
||||
),
|
||||
)
|
||||
new_client._litellm_httpx_client = httpx_client # type: ignore[attr-defined]
|
||||
new_client._litellm_agent_card = agent_card # type: ignore[attr-defined]
|
||||
new_client._litellm_httpx_client = httpx_client
|
||||
new_client._litellm_agent_card = agent_card
|
||||
return new_client
|
||||
|
|
|
|||
|
|
@ -10,8 +10,8 @@ A2A Streaming Events (in order):
|
|||
4. Status update (kind: "status-update") - Final status "completed" with final=true
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -21,14 +21,16 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
|
|||
)
|
||||
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
|
||||
from litellm.interactions.agents.utils import merge_agent_headers
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
# 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",
|
||||
|
|
@ -45,56 +47,23 @@ class A2ACompletionBridgeHandler:
|
|||
"""
|
||||
|
||||
@staticmethod
|
||||
async def handle_non_streaming(
|
||||
request_id: str,
|
||||
def _build_completion_params(
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
litellm_params: Mapping[str, Any],
|
||||
api_base: str | None,
|
||||
agent_extra_headers: Mapping[str, str] | None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Handle non-streaming A2A request via litellm.acompletion.
|
||||
|
||||
Args:
|
||||
request_id: A2A JSON-RPC request ID
|
||||
params: A2A MessageSendParams containing the message
|
||||
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
|
||||
api_base: API base URL from agent_card_params
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Returns:
|
||||
A2A SendMessageResponse dict
|
||||
"""
|
||||
custom_llm_provider = litellm_params.get("custom_llm_provider")
|
||||
if not _skip_a2a_provider_routing:
|
||||
a2a_provider_config = 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}")
|
||||
|
||||
return await a2a_provider_config.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
stream: bool,
|
||||
) -> Mapping[str, Any]:
|
||||
# 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")
|
||||
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
|
||||
model: Final = litellm_params.get("model", "agent")
|
||||
|
||||
# Build full model string if provider specified
|
||||
# Skip prepending if model already starts with the provider prefix
|
||||
|
|
@ -103,17 +72,20 @@ class A2ACompletionBridgeHandler:
|
|||
else:
|
||||
full_model = model
|
||||
|
||||
verbose_logger.info(f"A2A completion bridge: model={full_model}, api_base={api_base}")
|
||||
if stream:
|
||||
verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base)
|
||||
else:
|
||||
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,
|
||||
"stream": stream,
|
||||
}
|
||||
# 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,16 +106,72 @@ class A2ACompletionBridgeHandler:
|
|||
static_headers=completion_params.get("extra_headers"),
|
||||
)
|
||||
|
||||
return completion_params
|
||||
|
||||
@staticmethod
|
||||
async def _acompletion(completion_params: Mapping[str, Any]) -> ModelResponse | CustomStreamWrapper:
|
||||
return await litellm.acompletion(**completion_params)
|
||||
|
||||
@staticmethod
|
||||
async def handle_non_streaming(
|
||||
request_id: str,
|
||||
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, object]:
|
||||
"""
|
||||
Handle non-streaming A2A request via litellm.acompletion.
|
||||
|
||||
Args:
|
||||
request_id: A2A JSON-RPC request ID
|
||||
params: A2A MessageSendParams containing the message
|
||||
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
|
||||
api_base: API base URL from agent_card_params
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Returns:
|
||||
A2A SendMessageResponse dict
|
||||
"""
|
||||
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
|
||||
if not _skip_a2a_provider_routing:
|
||||
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("A2A: Using provider config for %s", custom_llm_provider)
|
||||
|
||||
return await a2a_provider_config.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
completion_params: Final = A2ACompletionBridgeHandler._build_completion_params(
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
# Call litellm.acompletion
|
||||
response = await litellm.acompletion(**completion_params)
|
||||
response: Final = await A2ACompletionBridgeHandler._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
|
||||
|
||||
|
|
@ -156,7 +184,7 @@ class A2ACompletionBridgeHandler:
|
|||
agent_extra_headers: dict[str, str] | None = None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
) -> AsyncIterator[dict[str, object]]:
|
||||
"""
|
||||
Handle streaming A2A request via litellm.acompletion with stream=True.
|
||||
|
||||
|
|
@ -177,15 +205,15 @@ class A2ACompletionBridgeHandler:
|
|||
Yields:
|
||||
A2A streaming response events
|
||||
"""
|
||||
custom_llm_provider = litellm_params.get("custom_llm_provider")
|
||||
custom_llm_provider: Final = 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,66 +226,26 @@ class A2ACompletionBridgeHandler:
|
|||
|
||||
return
|
||||
|
||||
# Extract message from params
|
||||
message = params.get("message", {})
|
||||
|
||||
# Create streaming context
|
||||
ctx = A2AStreamingContext(
|
||||
ctx: Final = A2AStreamingContext(
|
||||
request_id=request_id,
|
||||
input_message=message,
|
||||
input_message=params.get("message", {}),
|
||||
)
|
||||
|
||||
# Transform A2A message to OpenAI format
|
||||
openai_messages = 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")
|
||||
|
||||
# Build full model string if provider specified
|
||||
# Skip prepending if model already starts with the provider prefix
|
||||
if custom_llm_provider and not model.startswith(f"{custom_llm_provider}/"):
|
||||
full_model = f"{custom_llm_provider}/{model}"
|
||||
else:
|
||||
full_model = model
|
||||
|
||||
verbose_logger.info(f"A2A completion bridge streaming: model={full_model}, api_base={api_base}")
|
||||
|
||||
# Build completion params dict
|
||||
completion_params: 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 = {
|
||||
k: v
|
||||
for k, v in litellm_params.items()
|
||||
if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS
|
||||
}
|
||||
completion_params.update(litellm_params_to_add)
|
||||
# Apply forward metadata AFTER the litellm_params merge so the helper
|
||||
# sees any agent-owner-configured ``extra_body.metadata`` and can keep
|
||||
# those keys authoritative over the client-supplied A2A metadata.
|
||||
A2ACompletionBridgeTransformation.apply_forward_metadata_to_completion_params(
|
||||
completion_params=completion_params,
|
||||
a2a_message=message,
|
||||
completion_params: Final = A2ACompletionBridgeHandler._build_completion_params(
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
if agent_extra_headers:
|
||||
completion_params["extra_headers"] = merge_agent_headers(
|
||||
dynamic_headers=agent_extra_headers,
|
||||
static_headers=completion_params.get("extra_headers"),
|
||||
)
|
||||
|
||||
# 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,
|
||||
|
|
@ -266,12 +254,12 @@ class A2ACompletionBridgeHandler:
|
|||
yield working_event
|
||||
|
||||
# Call litellm.acompletion with streaming
|
||||
response = await litellm.acompletion(**completion_params)
|
||||
response: Final = await A2ACompletionBridgeHandler._acompletion(completion_params)
|
||||
|
||||
# 3. Accumulate content and emit artifact update
|
||||
accumulated_text = ""
|
||||
chunk_count = 0
|
||||
async for chunk in response: # type: ignore[union-attr]
|
||||
async for chunk in response:
|
||||
chunk_count += 1
|
||||
|
||||
# Extract delta content
|
||||
|
|
@ -286,21 +274,23 @@ 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
|
||||
|
|
@ -310,7 +300,7 @@ async def handle_a2a_completion(
|
|||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""Convenience function for non-streaming A2A completion."""
|
||||
return await A2ACompletionBridgeHandler.handle_non_streaming(
|
||||
request_id=request_id,
|
||||
|
|
@ -327,7 +317,7 @@ async def handle_a2a_completion_streaming(
|
|||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
) -> AsyncIterator[dict[str, object]]:
|
||||
"""Convenience function for streaming A2A completion."""
|
||||
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
|
||||
request_id=request_id,
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ A2A Streaming Events:
|
|||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -48,7 +48,7 @@ class A2ACompletionBridgeTransformation:
|
|||
@staticmethod
|
||||
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
|
||||
|
|
@ -71,10 +71,10 @@ class A2ACompletionBridgeTransformation:
|
|||
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
|
||||
|
|
@ -90,7 +90,7 @@ class A2ACompletionBridgeTransformation:
|
|||
|
||||
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,13 +103,13 @@ 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(
|
||||
|
|
@ -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,13 +139,15 @@ 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]
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -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(),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -13,12 +13,7 @@ import asyncio
|
|||
import datetime
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Coroutine
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Optional,
|
||||
cast,
|
||||
)
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
|
|
@ -64,9 +59,9 @@ try:
|
|||
|
||||
A2A_SDK_AVAILABLE = True
|
||||
except ImportError:
|
||||
Client = None # type: ignore[misc, assignment]
|
||||
ClientConfig = None # type: ignore[misc, assignment]
|
||||
create_client = None # type: ignore[misc, assignment]
|
||||
Client = None
|
||||
ClientConfig = None
|
||||
create_client = None
|
||||
|
||||
# Import our custom card resolver that supports multiple well-known paths
|
||||
from litellm.a2a_protocol.card_resolver import (
|
||||
|
|
@ -80,7 +75,7 @@ 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(
|
||||
|
|
@ -96,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,
|
||||
|
|
@ -120,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(
|
||||
|
|
@ -141,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
|
||||
|
||||
|
|
@ -149,7 +144,7 @@ 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}
|
||||
|
||||
|
||||
|
|
@ -162,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
|
||||
|
|
@ -204,7 +199,7 @@ async def _send_message_via_completion_bridge(
|
|||
|
||||
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,
|
||||
|
|
@ -212,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,
|
||||
|
|
@ -230,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 "
|
||||
|
|
@ -305,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,
|
||||
|
|
@ -425,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:
|
||||
|
|
@ -450,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)
|
||||
|
|
@ -461,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,
|
||||
|
|
@ -478,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,
|
||||
|
|
@ -549,10 +544,10 @@ def _build_streaming_logging_obj(
|
|||
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,
|
||||
|
|
@ -569,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:
|
||||
|
|
@ -632,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:
|
||||
|
|
@ -640,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)
|
||||
)
|
||||
|
||||
|
|
@ -664,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:
|
||||
|
|
@ -685,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(
|
||||
|
|
@ -697,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,
|
||||
|
|
@ -759,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
|
||||
|
|
@ -769,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,
|
||||
|
|
@ -793,12 +788,12 @@ async def create_a2a_client(
|
|||
# Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse
|
||||
# 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)
|
||||
a2a_client._litellm_httpx_client = httpx_client
|
||||
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]
|
||||
a2a_client._litellm_agent_card = agent_card
|
||||
|
||||
verbose_logger.info(f"A2A client created for {base_url}")
|
||||
verbose_logger.info("A2A client created for %s", base_url)
|
||||
|
||||
return a2a_client
|
||||
|
||||
|
|
@ -824,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
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Bedrock AgentCore A2A provider configuration.
|
|||
"""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.handler import (
|
||||
|
|
@ -28,7 +28,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
|
|||
**kwargs,
|
||||
) -> 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)"
|
||||
|
|
@ -48,7 +48,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
|
|||
**kwargs,
|
||||
) -> 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)"
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ completion bridge that would otherwise strip the envelope.
|
|||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, cast
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
|
||||
|
|
@ -53,21 +53,21 @@ 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
|
||||
|
||||
|
|
@ -100,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,
|
||||
|
|
@ -114,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
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT).
|
|||
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
|
||||
|
|
@ -23,13 +23,13 @@ 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-",
|
||||
)
|
||||
|
|
@ -47,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):
|
||||
|
|
@ -107,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,
|
||||
|
|
@ -129,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,
|
||||
|
|
@ -138,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)
|
||||
|
||||
|
|
@ -195,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
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ This handler provides fake streaming by converting non-streaming responses into
|
|||
"""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
|
||||
|
|
@ -47,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,
|
||||
|
|
@ -92,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,
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ This module provides fake streaming by converting non-streaming responses into s
|
|||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, cast
|
||||
from typing import Any, Final, cast
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -118,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
|
||||
|
|
@ -163,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",
|
||||
|
|
@ -171,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={
|
||||
|
|
@ -192,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,
|
||||
|
|
@ -209,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
|
||||
|
||||
|
|
@ -235,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,
|
||||
|
|
@ -313,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}],
|
||||
|
|
@ -342,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", [])
|
||||
|
|
@ -356,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", [])
|
||||
|
|
@ -369,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()))
|
||||
|
|
@ -410,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":
|
||||
|
|
@ -419,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": {
|
||||
|
|
@ -452,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": {
|
||||
|
|
@ -503,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": {
|
||||
|
|
@ -518,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)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ A2A provider configuration for IBM watsonx Orchestrate (WXO).
|
|||
"""
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
|
||||
from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import (
|
||||
|
|
@ -22,7 +22,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
|
|||
**kwargs: 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 "
|
||||
|
|
@ -42,7 +42,7 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig):
|
|||
**kwargs: 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 "
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import hashlib
|
|||
import json
|
||||
import time
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, NamedTuple, cast
|
||||
from typing import Any, Final, NamedTuple, cast
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -21,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):
|
||||
|
|
@ -53,14 +53,14 @@ class WatsonxOrchestrateHandler:
|
|||
api_key: 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: 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
|
||||
|
|
@ -71,9 +71,9 @@ class WatsonxOrchestrateHandler:
|
|||
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]
|
||||
|
||||
|
|
@ -96,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},
|
||||
|
|
@ -105,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:
|
||||
|
|
@ -127,7 +127,7 @@ class WatsonxOrchestrateHandler:
|
|||
max_attempts: int = _MAX_POLL_ATTEMPTS,
|
||||
interval_s: float = _POLL_INTERVAL_S,
|
||||
) -> dict[str, Any]:
|
||||
url = f"{base_url}/v1/orchestrate/runs/{run_id}"
|
||||
url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}"
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
await asyncio.sleep(interval_s)
|
||||
|
|
@ -135,7 +135,7 @@ class WatsonxOrchestrateHandler:
|
|||
response.raise_for_status()
|
||||
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
|
||||
|
||||
|
|
@ -152,7 +152,7 @@ class WatsonxOrchestrateHandler:
|
|||
) -> 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(
|
||||
|
|
@ -188,10 +188,10 @@ class WatsonxOrchestrateHandler:
|
|||
|
||||
@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 ""
|
||||
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")
|
||||
|
|
@ -218,29 +218,29 @@ class WatsonxOrchestrateHandler:
|
|||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
|
||||
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,
|
||||
|
|
@ -255,7 +255,7 @@ 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
|
||||
|
|
@ -266,29 +266,29 @@ class WatsonxOrchestrateHandler:
|
|||
chunk_size: int = 50,
|
||||
delay_ms: int = 10,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
|
||||
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,
|
||||
|
|
@ -297,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(
|
||||
|
|
@ -306,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,
|
||||
|
|
@ -316,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,
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model:
|
|||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -35,9 +35,9 @@ class WatsonxOrchestrateTransformation:
|
|||
|
||||
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
|
||||
|
|
@ -53,7 +53,7 @@ class WatsonxOrchestrateTransformation:
|
|||
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",
|
||||
|
|
@ -96,7 +96,7 @@ 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
|
||||
|
||||
|
|
@ -104,11 +104,11 @@ class WatsonxOrchestrateTransformation:
|
|||
|
||||
@staticmethod
|
||||
def extract_text_from_a2a_message_response(a2a_response: dict[str, Any]) -> str:
|
||||
result = a2a_response.get("result")
|
||||
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 ""
|
||||
|
|
@ -150,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 {
|
||||
|
|
@ -181,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)
|
||||
|
|
@ -214,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)
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ 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
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -47,7 +47,7 @@ class A2AStreamingIterator:
|
|||
|
||||
async def __anext__(self) -> "SendStreamingMessageResponse":
|
||||
try:
|
||||
chunk = await self.stream.__anext__()
|
||||
chunk: Final = await self.stream.__anext__()
|
||||
|
||||
# Store chunk
|
||||
self.chunks.append(chunk)
|
||||
|
|
@ -71,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:
|
||||
|
|
@ -81,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:
|
||||
|
|
@ -94,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,
|
||||
|
|
@ -120,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(
|
||||
|
|
@ -138,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]:
|
||||
"""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)),
|
||||
|
|
@ -157,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
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
Utility functions for A2A protocol.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
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":
|
||||
|
|
@ -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,7 +66,7 @@ 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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ Environment Variables:
|
|||
import json
|
||||
import os
|
||||
from importlib.resources import files
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -46,12 +47,12 @@ class GetAnthropicBetaHeadersConfig:
|
|||
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": {},
|
||||
|
|
@ -79,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(
|
||||
|
|
@ -113,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()
|
||||
|
||||
|
|
@ -138,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.",
|
||||
|
|
@ -206,8 +207,8 @@ 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)
|
||||
|
||||
|
||||
|
|
@ -233,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
|
||||
|
|
@ -254,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
|
||||
|
|
@ -277,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
|
||||
|
|
@ -301,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:
|
||||
|
|
@ -330,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,
|
||||
)
|
||||
|
|
@ -372,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,
|
||||
)
|
||||
|
|
@ -399,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,13 +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 .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",
|
||||
|
|
@ -50,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,
|
||||
|
|
@ -76,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):
|
||||
|
|
@ -151,7 +153,7 @@ class AnthropicExceptionMapping:
|
|||
# Optionally add request_id if provided and not present
|
||||
if request_id and "request_id" not in parsed:
|
||||
parsed["request_id"] = request_id
|
||||
return parsed # type: ignore
|
||||
return parsed
|
||||
|
||||
# Extract message - use parsed dict if available, otherwise raw string
|
||||
if parsed is not None:
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import contextvars
|
|||
import os
|
||||
from collections.abc import Coroutine, Iterable
|
||||
from functools import partial
|
||||
from typing import Any, Literal
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
import httpx
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
|
|
@ -29,8 +29,8 @@ 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 ###
|
||||
|
||||
|
|
@ -40,28 +40,26 @@ async def aget_assistants(
|
|||
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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model="",
|
||||
|
|
@ -80,11 +78,11 @@ def get_assistants(
|
|||
api_version: str | None = None,
|
||||
**kwargs,
|
||||
) -> SyncCursorPage[Assistant]:
|
||||
aget_assistants: bool | None = 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
|
||||
|
|
@ -95,10 +93,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
|
||||
|
|
@ -111,7 +109,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)
|
||||
|
|
@ -132,12 +130,12 @@ def get_assistants(
|
|||
max_retries=optional_params.max_retries,
|
||||
organization=organization,
|
||||
client=client,
|
||||
aget_assistants=aget_assistants, # type: ignore
|
||||
) # type: ignore
|
||||
aget_assistants=aget_assistants,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -145,14 +143,14 @@ def get_assistants(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
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: str | None = None
|
||||
if extra_body is not None:
|
||||
azure_ad_token = extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_assistants_api.get_assistants(
|
||||
api_base=api_base,
|
||||
|
|
@ -162,7 +160,7 @@ def get_assistants(
|
|||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
client=client,
|
||||
aget_assistants=aget_assistants, # type: ignore
|
||||
aget_assistants=aget_assistants,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
else:
|
||||
|
|
@ -173,7 +171,7 @@ def get_assistants(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -185,7 +183,7 @@ def get_assistants(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -197,30 +195,28 @@ async def acreate_assistants(
|
|||
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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model=model,
|
||||
|
|
@ -249,11 +245,11 @@ def create_assistants(
|
|||
api_version: str | None = None,
|
||||
**kwargs,
|
||||
) -> Assistant | Coroutine[Any, Any, Assistant]:
|
||||
async_create_assistants: bool | None = kwargs.pop("async_create_assistants", None)
|
||||
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
|
||||
|
|
@ -264,10 +260,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
|
||||
|
|
@ -296,7 +292,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)
|
||||
|
|
@ -318,12 +314,12 @@ def create_assistants(
|
|||
organization=organization,
|
||||
create_assistant_data=create_assistant_data,
|
||||
client=client,
|
||||
async_create_assistants=async_create_assistants, # type: ignore
|
||||
) # type: ignore
|
||||
async_create_assistants=async_create_assistants,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -331,14 +327,14 @@ def create_assistants(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
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: str | None = None
|
||||
if extra_body is not None:
|
||||
azure_ad_token = extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
if isinstance(client, OpenAI):
|
||||
client = None # only pass client if it's AzureOpenAI
|
||||
|
|
@ -363,7 +359,7 @@ def create_assistants(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
if response is None:
|
||||
|
|
@ -380,29 +376,27 @@ async def adelete_assistant(
|
|||
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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model="",
|
||||
|
|
@ -422,11 +416,11 @@ def delete_assistant(
|
|||
api_version: str | None = None,
|
||||
**kwargs,
|
||||
) -> AssistantDeleted | Coroutine[Any, Any, AssistantDeleted]:
|
||||
optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs)
|
||||
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: bool | None = 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")
|
||||
|
||||
|
|
@ -439,10 +433,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
|
||||
|
|
@ -455,7 +449,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
|
||||
|
|
@ -472,9 +466,9 @@ def delete_assistant(
|
|||
async_delete_assistants=async_delete_assistants,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -482,14 +476,14 @@ def delete_assistant(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
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: str | None = None
|
||||
if extra_body is not None:
|
||||
azure_ad_token = extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
if isinstance(client, OpenAI):
|
||||
client = None # only pass client if it's AzureOpenAI
|
||||
|
|
@ -530,28 +524,26 @@ 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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model="",
|
||||
|
|
@ -592,9 +584,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
|
||||
|
|
@ -605,10 +597,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
|
||||
|
|
@ -624,7 +616,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)
|
||||
|
|
@ -649,7 +641,7 @@ def create_thread(
|
|||
acreate_thread=acreate_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_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -657,16 +649,16 @@ def create_thread(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
or get_secret("AZURE_API_KEY")
|
||||
) # type: ignore
|
||||
)
|
||||
|
||||
api_version: str | None = 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")
|
||||
|
||||
extra_body = optional_params.get("extra_body", {})
|
||||
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:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
if isinstance(client, OpenAI):
|
||||
client = None # only pass client if it's AzureOpenAI
|
||||
|
|
@ -692,10 +684,10 @@ def create_thread(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
return response # type: ignore
|
||||
return response
|
||||
|
||||
|
||||
async def aget_thread(
|
||||
|
|
@ -704,28 +696,26 @@ async def aget_thread(
|
|||
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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model="",
|
||||
|
|
@ -743,9 +733,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
|
||||
|
|
@ -755,10 +745,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
api_base: str | None = None
|
||||
|
|
@ -772,7 +762,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)
|
||||
|
|
@ -797,9 +787,9 @@ def get_thread(
|
|||
aget_thread=aget_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_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_version: str | None = 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")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -807,14 +797,14 @@ def get_thread(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
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: str | None = None
|
||||
if extra_body is not None:
|
||||
azure_ad_token = extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
if isinstance(client, OpenAI):
|
||||
client = None # only pass client if it's AzureOpenAI
|
||||
|
|
@ -839,10 +829,10 @@ def get_thread(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
return response # type: ignore
|
||||
return response
|
||||
|
||||
|
||||
### MESSAGES ###
|
||||
|
|
@ -858,12 +848,12 @@ async def a_add_message(
|
|||
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,
|
||||
|
|
@ -876,21 +866,19 @@ 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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
# Call the synchronous function using run_in_executor
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model="",
|
||||
|
|
@ -912,12 +900,12 @@ def add_message(
|
|||
**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"],
|
||||
|
|
@ -934,10 +922,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
api_key: str | None = None
|
||||
|
|
@ -951,7 +939,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)
|
||||
|
|
@ -976,9 +964,9 @@ def add_message(
|
|||
a_add_message=a_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_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_version: str | None = 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")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -986,14 +974,14 @@ def add_message(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
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: str | None = None
|
||||
if extra_body is not None:
|
||||
azure_ad_token = extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_assistants_api.add_message(
|
||||
thread_id=thread_id,
|
||||
|
|
@ -1016,11 +1004,11 @@ def add_message(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
|
||||
return response # type: ignore
|
||||
return response
|
||||
|
||||
|
||||
async def aget_messages(
|
||||
|
|
@ -1029,12 +1017,12 @@ async def aget_messages(
|
|||
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,
|
||||
|
|
@ -1043,21 +1031,19 @@ 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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
# Call the synchronous function using run_in_executor
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model="",
|
||||
|
|
@ -1074,9 +1060,9 @@ def get_messages(
|
|||
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
|
||||
|
|
@ -1087,10 +1073,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
|
||||
|
|
@ -1105,7 +1091,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)
|
||||
|
|
@ -1129,9 +1115,9 @@ def get_messages(
|
|||
aget_messages=aget_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_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_version: str | None = 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")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -1139,14 +1125,14 @@ def get_messages(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
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: str | None = None
|
||||
if extra_body is not None:
|
||||
azure_ad_token = extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_assistants_api.get_messages(
|
||||
thread_id=thread_id,
|
||||
|
|
@ -1168,11 +1154,11 @@ def get_messages(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
|
||||
return response # type: ignore
|
||||
return response
|
||||
|
||||
|
||||
### RUNS ###
|
||||
|
|
@ -1182,7 +1168,7 @@ def arun_thread_stream(
|
|||
**kwargs,
|
||||
) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]:
|
||||
kwargs["arun_thread"] = True
|
||||
return run_thread(stream=True, event_handler=event_handler, **kwargs) # type: ignore
|
||||
return run_thread(stream=True, event_handler=event_handler, **kwargs)
|
||||
|
||||
|
||||
async def arun_thread(
|
||||
|
|
@ -1198,12 +1184,12 @@ async def arun_thread(
|
|||
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,
|
||||
|
|
@ -1219,21 +1205,19 @@ 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
|
||||
_, custom_llm_provider, _, _ = get_llm_provider(model="", custom_llm_provider=custom_llm_provider)
|
||||
|
||||
# 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:
|
||||
# Call the synchronous function using run_in_executor
|
||||
response = init_response
|
||||
return response # type: ignore
|
||||
return response
|
||||
except Exception as e:
|
||||
raise exception_type(
|
||||
model="",
|
||||
|
|
@ -1249,7 +1233,7 @@ def run_thread_stream(
|
|||
event_handler: AssistantEventHandler | None = None,
|
||||
**kwargs,
|
||||
) -> AssistantStreamManager[AssistantEventHandler]:
|
||||
return run_thread(stream=True, event_handler=event_handler, **kwargs) # type: ignore
|
||||
return run_thread(stream=True, event_handler=event_handler, **kwargs)
|
||||
|
||||
|
||||
def run_thread(
|
||||
|
|
@ -1267,9 +1251,9 @@ def run_thread(
|
|||
**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
|
||||
|
|
@ -1280,10 +1264,10 @@ 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
|
||||
timeout = float(timeout)
|
||||
elif timeout is None:
|
||||
timeout = 600.0
|
||||
|
||||
|
|
@ -1296,7 +1280,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)
|
||||
|
|
@ -1329,9 +1313,9 @@ def run_thread(
|
|||
event_handler=event_handler,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE")
|
||||
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
|
||||
api_version = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION")
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
|
|
@ -1339,14 +1323,14 @@ def run_thread(
|
|||
or litellm.azure_key
|
||||
or get_secret("AZURE_OPENAI_API_KEY")
|
||||
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)
|
||||
else:
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN") # type: ignore
|
||||
azure_ad_token = get_secret("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_assistants_api.run_thread(
|
||||
thread_id=thread_id,
|
||||
|
|
@ -1366,7 +1350,7 @@ def run_thread(
|
|||
client=client,
|
||||
arun_thread=arun_thread,
|
||||
litellm_params=litellm_params_dict,
|
||||
) # type: ignore
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message=f"LiteLLM doesn't support {custom_llm_provider} for 'run_thread'. Only 'openai' is supported.",
|
||||
|
|
@ -1375,7 +1359,7 @@ def run_thread(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
return response # type: ignore
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
from typing import Final
|
||||
|
||||
import litellm
|
||||
|
||||
from ..exceptions import UnsupportedParamsError
|
||||
|
|
@ -6,8 +8,8 @@ from ..types.llms.openai import *
|
|||
|
||||
def get_optional_params_add_message(
|
||||
role: str | None,
|
||||
content: str | List[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None,
|
||||
attachments: List[Attachment] | None,
|
||||
content: str | list[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None,
|
||||
attachments: list[Attachment] | None,
|
||||
metadata: dict | None,
|
||||
custom_llm_provider: str,
|
||||
**kwargs,
|
||||
|
|
@ -17,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,
|
||||
|
|
@ -36,7 +38,7 @@ 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)
|
||||
|
|
@ -50,12 +52,12 @@ def get_optional_params_add_message(
|
|||
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():
|
||||
for k in passed_params:
|
||||
if k not in default_params:
|
||||
optional_params[k] = passed_params[k]
|
||||
return optional_params
|
||||
|
|
@ -72,13 +74,13 @@ def get_optional_params_image_gen(
|
|||
**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,
|
||||
|
|
@ -93,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)
|
||||
|
|
@ -126,7 +128,7 @@ def get_optional_params_image_gen(
|
|||
if n is not None:
|
||||
optional_params["sampleCount"] = int(n)
|
||||
|
||||
for k in passed_params.keys():
|
||||
for k in passed_params:
|
||||
if k not in default_params:
|
||||
optional_params[k] = passed_params[k]
|
||||
return optional_params
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose
|
||||
|
|
@ -55,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,
|
||||
|
|
@ -145,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:
|
||||
|
|
@ -156,10 +157,10 @@ 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:
|
||||
|
|
@ -238,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:
|
||||
|
|
@ -249,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: {e!s}")
|
||||
print_verbose(f"batch_completion_models_all_responses: model request failed: {e}")
|
||||
continue
|
||||
|
||||
return responses
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -141,9 +141,9 @@ def _aggregate_batch_cost_usage_models(
|
|||
) -> 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)),
|
||||
|
|
@ -151,14 +151,14 @@ 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
|
||||
|
||||
|
|
@ -184,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")
|
||||
|
|
@ -254,27 +254,27 @@ 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)
|
||||
return _file_content.content
|
||||
|
||||
|
||||
|
|
@ -291,11 +291,11 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> 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",
|
||||
|
|
@ -355,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:
|
||||
|
|
@ -370,18 +370,18 @@ def _count_entry_tokens(
|
|||
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)
|
||||
|
||||
|
|
@ -432,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
|
||||
|
||||
|
||||
|
|
@ -455,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
|
||||
|
||||
|
||||
|
|
@ -472,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
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import contextvars
|
|||
import os
|
||||
from collections.abc import Coroutine
|
||||
from functools import partial
|
||||
from typing import Any, Literal, cast
|
||||
from typing import Any, Final, Literal, cast
|
||||
|
||||
import httpx
|
||||
from openai.types.batch import BatchRequestCounts
|
||||
|
|
@ -54,10 +54,10 @@ 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()
|
||||
#################################################
|
||||
|
||||
|
|
@ -80,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
|
||||
|
|
@ -104,7 +104,7 @@ 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: dict[str, str] | None = None,
|
||||
|
|
@ -119,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,
|
||||
|
|
@ -137,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
|
||||
|
|
@ -154,7 +154,7 @@ 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: dict[str, str] | None = None,
|
||||
|
|
@ -169,10 +169,10 @@ def create_batch(
|
|||
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)
|
||||
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:
|
||||
|
|
@ -182,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 - {e!s}"
|
||||
"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,
|
||||
|
|
@ -206,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,
|
||||
|
|
@ -248,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)
|
||||
|
|
@ -287,7 +287,7 @@ def create_batch(
|
|||
if extra_body is not None:
|
||||
extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
get_secret_str("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_batches_instance.create_batch(
|
||||
_is_async=_is_async,
|
||||
|
|
@ -301,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,
|
||||
|
|
@ -327,7 +327,7 @@ def create_batch(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
return response
|
||||
|
|
@ -350,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,
|
||||
|
|
@ -364,13 +364,13 @@ 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:
|
||||
response = init_response # type: ignore
|
||||
response = init_response
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
|
|
@ -397,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)
|
||||
|
|
@ -422,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
|
||||
|
|
@ -432,11 +432,11 @@ 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:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
get_secret_str("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
|
|
@ -450,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,
|
||||
|
|
@ -498,7 +498,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="retrieve_batch", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
return response
|
||||
|
|
@ -519,11 +519,11 @@ def retrieve_batch(
|
|||
LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id}
|
||||
"""
|
||||
try:
|
||||
optional_params = GenericLiteLLMParams(**kwargs)
|
||||
litellm_logging_obj: LiteLLMLoggingObj | None = 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,
|
||||
)
|
||||
|
|
@ -542,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
|
||||
timeout = float(timeout)
|
||||
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
|
||||
|
|
@ -568,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(
|
||||
|
|
@ -578,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(
|
||||
|
|
@ -589,7 +589,7 @@ def retrieve_batch(
|
|||
)
|
||||
|
||||
# Try to use provider config first (for providers like bedrock)
|
||||
model: str | None = 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,
|
||||
|
|
@ -599,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,
|
||||
|
|
@ -656,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,
|
||||
|
|
@ -671,13 +671,13 @@ 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:
|
||||
response = init_response # type: ignore
|
||||
response = init_response
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
|
|
@ -700,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,
|
||||
)
|
||||
|
|
@ -720,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
|
||||
timeout = float(timeout)
|
||||
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 = (
|
||||
|
|
@ -737,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)
|
||||
|
|
@ -755,7 +755,7 @@ def list_batches(
|
|||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") # type: ignore
|
||||
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_key = (
|
||||
|
|
@ -770,7 +770,7 @@ def list_batches(
|
|||
if extra_body is not None:
|
||||
extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
get_secret_str("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_batches_instance.list_batches(
|
||||
_is_async=_is_async,
|
||||
|
|
@ -783,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,
|
||||
|
|
@ -813,7 +813,7 @@ def list_batches(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
return response
|
||||
|
|
@ -836,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,
|
||||
|
|
@ -854,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:
|
||||
|
|
@ -890,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 - {e!s}"
|
||||
"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,
|
||||
)
|
||||
|
|
@ -906,20 +906,20 @@ 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
|
||||
timeout = float(timeout)
|
||||
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
|
||||
_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 = (
|
||||
|
|
@ -929,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")
|
||||
|
|
@ -959,7 +959,7 @@ def cancel_batch(
|
|||
if extra_body is not None:
|
||||
extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
get_secret_str("AZURE_AD_TOKEN")
|
||||
|
||||
response = azure_batches_instance.cancel_batch(
|
||||
_is_async=_is_async,
|
||||
|
|
@ -973,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,
|
||||
|
|
@ -999,7 +999,7 @@ def cancel_batch(
|
|||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="cancel_batch", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
request=httpx.Request(method="cancel_batch", url="https://github.com/BerriAI/litellm"),
|
||||
),
|
||||
)
|
||||
return response
|
||||
|
|
@ -1025,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,
|
||||
|
|
@ -1040,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
|
||||
|
||||
|
|
@ -1073,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,
|
||||
|
|
@ -1105,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())
|
||||
|
|
@ -1113,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()
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import json
|
|||
import os
|
||||
import threading
|
||||
import time
|
||||
from typing import Literal
|
||||
from typing import Final, Literal
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
|
|
@ -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":
|
||||
|
|
@ -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):
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from collections.abc import Callable
|
||||
from functools import lru_cache
|
||||
from typing import TypeVar
|
||||
from typing import Final, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
|
@ -21,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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -9,12 +9,12 @@ Has 4 methods:
|
|||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
Span = Union[_Span, Any]
|
||||
Span = _Span | Any
|
||||
else:
|
||||
Span = Any
|
||||
|
||||
|
|
@ -24,7 +24,7 @@ class BaseCache(ABC):
|
|||
self.default_ttl = default_ttl
|
||||
|
||||
def get_ttl(self, **kwargs) -> int | None:
|
||||
kwargs_ttl: int | None = kwargs.get("ttl")
|
||||
kwargs_ttl: Final[int | None] = kwargs.get("ttl")
|
||||
if kwargs_ttl is not None:
|
||||
try:
|
||||
return int(kwargs_ttl)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import json
|
|||
import time
|
||||
import traceback
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -169,13 +169,13 @@ class Cache:
|
|||
if type == LiteLLMCacheType.REDIS:
|
||||
# Check REDIS_CLUSTER_NODES env var if no explicit startup nodes
|
||||
if not redis_startup_nodes:
|
||||
_env_cluster_nodes = litellm.get_secret("REDIS_CLUSTER_NODES")
|
||||
_env_cluster_nodes: Final = litellm.get_secret("REDIS_CLUSTER_NODES")
|
||||
if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str):
|
||||
redis_startup_nodes = json.loads(_env_cluster_nodes)
|
||||
|
||||
if redis_startup_nodes:
|
||||
# Only pass GCP parameters if they are provided
|
||||
cluster_kwargs = {
|
||||
cluster_kwargs: Final = {
|
||||
"host": host,
|
||||
"port": port,
|
||||
"password": password,
|
||||
|
|
@ -312,9 +312,9 @@ class Cache:
|
|||
)
|
||||
|
||||
def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str:
|
||||
metadata: dict = kwargs.get("metadata") or {}
|
||||
litellm_params: dict = kwargs.get("litellm_params") or {}
|
||||
metadata_in_litellm_params: dict = litellm_params.get("metadata") or {}
|
||||
metadata: Final[dict] = kwargs.get("metadata") or {}
|
||||
litellm_params: Final[dict] = kwargs.get("litellm_params") or {}
|
||||
metadata_in_litellm_params: Final[dict] = litellm_params.get("metadata") or {}
|
||||
|
||||
scope = ""
|
||||
for field in self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS:
|
||||
|
|
@ -338,28 +338,28 @@ class Cache:
|
|||
cache_key = ""
|
||||
# verbose_logger.debug("\nGetting Cache key. Kwargs: %s", kwargs)
|
||||
|
||||
preset_cache_key = self._get_preset_cache_key_from_kwargs(**kwargs)
|
||||
preset_cache_key: Final = self._get_preset_cache_key_from_kwargs(**kwargs)
|
||||
if preset_cache_key is not None:
|
||||
verbose_logger.debug("\nReturning preset cache key: %s", preset_cache_key)
|
||||
return preset_cache_key
|
||||
|
||||
combined_kwargs = ModelParamHelper._get_all_llm_api_params()
|
||||
litellm_param_kwargs = all_litellm_params
|
||||
is_semantic_cache = self._is_semantic_cache()
|
||||
scope_excluded_params = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset()
|
||||
combined_kwargs: Final = ModelParamHelper._get_all_llm_api_params()
|
||||
litellm_param_kwargs: Final = all_litellm_params
|
||||
is_semantic_cache: Final = self._is_semantic_cache()
|
||||
scope_excluded_params: Final = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset()
|
||||
for param in kwargs:
|
||||
if param in scope_excluded_params:
|
||||
continue
|
||||
if param in combined_kwargs:
|
||||
param_value: str | None = self._get_param_value(param, kwargs)
|
||||
if param_value is not None:
|
||||
cache_key += f"{param!s}: {param_value!s}"
|
||||
cache_key += f"{param}: {param_value}"
|
||||
elif param not in litellm_param_kwargs: # check if user passed in optional param - e.g. top_k
|
||||
if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now
|
||||
if kwargs[param] is None:
|
||||
continue # ignore None params
|
||||
param_value = kwargs[param]
|
||||
cache_key += f"{param!s}: {param_value!s}"
|
||||
cache_key += f"{param}: {param_value}"
|
||||
|
||||
if is_semantic_cache:
|
||||
cache_key += self._get_semantic_cache_tenant_scope(kwargs)
|
||||
|
|
@ -373,7 +373,7 @@ class Cache:
|
|||
)
|
||||
# Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError
|
||||
# when kwargs already contains preset_cache_key from upstream callers
|
||||
kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"}
|
||||
kwargs_for_preset: Final = {k: v for k, v in kwargs.items() if k != "preset_cache_key"}
|
||||
self._set_preset_cache_key_in_kwargs(preset_cache_key=hashed_cache_key, **kwargs_for_preset)
|
||||
return hashed_cache_key
|
||||
|
||||
|
|
@ -399,15 +399,15 @@ class Cache:
|
|||
2. Else if a model_group is set, then return the model_group as the model. This is used for all requests sent through the litellm.Router()
|
||||
3. Else use the `model` passed in kwargs
|
||||
"""
|
||||
metadata: dict = kwargs.get("metadata", {}) or {}
|
||||
litellm_params: dict = kwargs.get("litellm_params", {}) or {}
|
||||
metadata_in_litellm_params: dict = litellm_params.get("metadata", {}) or {}
|
||||
model_group: str | None = metadata.get("model_group") or metadata_in_litellm_params.get("model_group")
|
||||
caching_group = self._get_caching_group(metadata, model_group)
|
||||
metadata: Final[dict] = kwargs.get("metadata", {}) or {}
|
||||
litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {}
|
||||
metadata_in_litellm_params: Final[dict] = litellm_params.get("metadata", {}) or {}
|
||||
model_group: Final[str | None] = metadata.get("model_group") or metadata_in_litellm_params.get("model_group")
|
||||
caching_group: Final = self._get_caching_group(metadata, model_group)
|
||||
return caching_group or model_group or kwargs["model"]
|
||||
|
||||
def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None:
|
||||
caching_groups: list | None = metadata.get("caching_groups", [])
|
||||
caching_groups: Final[list | None] = metadata.get("caching_groups", [])
|
||||
if caching_groups:
|
||||
for group in caching_groups:
|
||||
if model_group in group:
|
||||
|
|
@ -418,9 +418,9 @@ class Cache:
|
|||
"""
|
||||
Handles getting the value for the 'file' param from kwargs. Used for `transcription` requests
|
||||
"""
|
||||
file = kwargs.get("file")
|
||||
metadata = kwargs.get("metadata", {})
|
||||
litellm_params = kwargs.get("litellm_params", {})
|
||||
file: Final = kwargs.get("file")
|
||||
metadata: Final = kwargs.get("metadata", {})
|
||||
litellm_params: Final = kwargs.get("litellm_params", {})
|
||||
return (
|
||||
metadata.get("file_checksum")
|
||||
or getattr(file, "name", None)
|
||||
|
|
@ -467,9 +467,9 @@ class Cache:
|
|||
Returns:
|
||||
str: The hashed cache key.
|
||||
"""
|
||||
hash_object = hashlib.sha256(cache_key.encode())
|
||||
hash_object: Final = hashlib.sha256(cache_key.encode())
|
||||
# Hexadecimal representation of the hash
|
||||
hash_hex = hash_object.hexdigest()
|
||||
hash_hex: Final = hash_object.hexdigest()
|
||||
verbose_logger.debug("Hashed cache key (SHA-256): %s", hash_hex)
|
||||
return hash_hex
|
||||
|
||||
|
|
@ -484,16 +484,16 @@ class Cache:
|
|||
Returns:
|
||||
str: The final hashed cache key with the redis namespace.
|
||||
"""
|
||||
dynamic_cache_control: DynamicCacheControl = kwargs.get("cache", {})
|
||||
metadata = kwargs.get("metadata") or {}
|
||||
namespace = dynamic_cache_control.get("namespace") or metadata.get("redis_namespace") or self.namespace
|
||||
dynamic_cache_control: Final[DynamicCacheControl] = kwargs.get("cache", {})
|
||||
metadata: Final = kwargs.get("metadata") or {}
|
||||
namespace: Final = dynamic_cache_control.get("namespace") or metadata.get("redis_namespace") or self.namespace
|
||||
if namespace:
|
||||
hash_hex = f"{namespace}:{hash_hex}"
|
||||
verbose_logger.debug("Final hashed key: %s", hash_hex)
|
||||
return hash_hex
|
||||
|
||||
def generate_streaming_content(self, content):
|
||||
chunk_size = 5 # Adjust the chunk size as needed
|
||||
chunk_size: Final = 5 # Adjust the chunk size as needed
|
||||
for i in range(0, len(content), chunk_size):
|
||||
yield {
|
||||
"choices": [
|
||||
|
|
@ -517,11 +517,11 @@ class Cache:
|
|||
"""
|
||||
# Check if a timestamp was stored with the cached response
|
||||
if cached_result is not None and isinstance(cached_result, dict) and "timestamp" in cached_result:
|
||||
timestamp = cached_result["timestamp"]
|
||||
current_time = time.time()
|
||||
timestamp: Final = cached_result["timestamp"]
|
||||
current_time: Final = time.time()
|
||||
|
||||
# Calculate age of the cached response
|
||||
response_age = current_time - timestamp
|
||||
response_age: Final = current_time - timestamp
|
||||
|
||||
# Check if the cached response is older than the max-age
|
||||
if max_age is not None and response_age > max_age:
|
||||
|
|
@ -534,22 +534,20 @@ class Cache:
|
|||
if isinstance(cached_response, dict):
|
||||
pass
|
||||
else:
|
||||
cached_response = json.loads(
|
||||
cached_response # type: ignore
|
||||
) # Convert string to dictionary
|
||||
cached_response = json.loads(cached_response) # Convert string to dictionary
|
||||
except Exception:
|
||||
cached_response = ast.literal_eval(cached_response) # type: ignore
|
||||
cached_response = ast.literal_eval(cached_response)
|
||||
return cached_response
|
||||
return cached_result
|
||||
|
||||
@staticmethod
|
||||
def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
cache_lookup_kwargs: dict[str, Any] = {}
|
||||
cache_lookup_kwargs: Final[dict[str, Any]] = {}
|
||||
for prompt_kwarg in ("messages", "input"):
|
||||
if prompt_kwarg in kwargs:
|
||||
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
|
||||
|
||||
metadata = kwargs.get("metadata")
|
||||
metadata: Final = kwargs.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
cache_lookup_kwargs["metadata"] = dict(metadata)
|
||||
|
||||
|
|
@ -559,8 +557,8 @@ class Cache:
|
|||
def _update_metadata_from_cache_lookup_kwargs(
|
||||
original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any]
|
||||
) -> None:
|
||||
original_metadata = original_kwargs.get("metadata")
|
||||
cache_lookup_metadata = cache_lookup_kwargs.get("metadata")
|
||||
original_metadata: Final = original_kwargs.get("metadata")
|
||||
cache_lookup_metadata: Final = cache_lookup_kwargs.get("metadata")
|
||||
if not isinstance(original_metadata, dict) or not isinstance(cache_lookup_metadata, dict):
|
||||
return
|
||||
|
||||
|
|
@ -586,9 +584,9 @@ class Cache:
|
|||
else:
|
||||
cache_key = self.get_cache_key(**kwargs)
|
||||
if cache_key is not None:
|
||||
cache_control_args: DynamicCacheControl = kwargs.get("cache", {})
|
||||
cache_control_args: Final[DynamicCacheControl] = kwargs.get("cache", {})
|
||||
max_age = cache_control_args.get("s-maxage") or cache_control_args.get("s-max-age") or float("inf")
|
||||
cache_lookup_kwargs = self._get_safe_cache_lookup_kwargs(kwargs)
|
||||
cache_lookup_kwargs: Final = self._get_safe_cache_lookup_kwargs(kwargs)
|
||||
if dynamic_cache_object is not None:
|
||||
cached_result = dynamic_cache_object.get_cache(cache_key, **cache_lookup_kwargs)
|
||||
else:
|
||||
|
|
@ -618,8 +616,8 @@ class Cache:
|
|||
else:
|
||||
cache_key = self.get_cache_key(**kwargs)
|
||||
if cache_key is not None:
|
||||
cache_control_args = kwargs.get("cache", {})
|
||||
max_age = cache_control_args.get("s-max-age", cache_control_args.get("s-maxage", float("inf")))
|
||||
cache_control_args: Final = kwargs.get("cache", {})
|
||||
max_age: Final = cache_control_args.get("s-max-age", cache_control_args.get("s-maxage", float("inf")))
|
||||
if dynamic_cache_object is not None:
|
||||
cached_result = await dynamic_cache_object.async_get_cache(cache_key, **kwargs)
|
||||
else:
|
||||
|
|
@ -646,13 +644,13 @@ class Cache:
|
|||
if self.ttl is not None:
|
||||
kwargs["ttl"] = self.ttl
|
||||
## Get Cache-Controls ##
|
||||
_cache_kwargs = kwargs.get("cache", None)
|
||||
_cache_kwargs: Final = kwargs.get("cache", None)
|
||||
if isinstance(_cache_kwargs, dict):
|
||||
for k, v in _cache_kwargs.items():
|
||||
if k == "ttl":
|
||||
kwargs["ttl"] = v
|
||||
|
||||
cached_data = {"timestamp": time.time(), "response": result}
|
||||
cached_data: Final = {"timestamp": time.time(), "response": result}
|
||||
return cache_key, cached_data, kwargs
|
||||
else:
|
||||
raise Exception("cache key is None")
|
||||
|
|
@ -676,7 +674,7 @@ class Cache:
|
|||
cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs)
|
||||
self.cache.set_cache(cache_key, cached_data, **kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e!s}")
|
||||
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
|
||||
|
||||
async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs):
|
||||
"""
|
||||
|
|
@ -695,7 +693,7 @@ class Cache:
|
|||
else:
|
||||
await self.cache.async_set_cache(cache_key, cached_data, **kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e!s}")
|
||||
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
|
||||
|
||||
def _convert_to_cached_embedding(
|
||||
self,
|
||||
|
|
@ -756,7 +754,7 @@ class Cache:
|
|||
if result.usage is None or result.usage.prompt_tokens_details is None:
|
||||
return None
|
||||
|
||||
details = result.usage.prompt_tokens_details
|
||||
details: Final = result.usage.prompt_tokens_details
|
||||
if hasattr(details, "model_dump"):
|
||||
details_dict = details.model_dump(exclude_none=True)
|
||||
elif isinstance(details, dict):
|
||||
|
|
@ -767,12 +765,12 @@ class Cache:
|
|||
if not details_dict:
|
||||
return None
|
||||
|
||||
num_items = len(result.data)
|
||||
num_items: Final = len(result.data)
|
||||
if num_items <= 1:
|
||||
return details_dict
|
||||
|
||||
# Distribute integer/float fields evenly across items
|
||||
per_item: dict = {}
|
||||
per_item: Final[dict] = {}
|
||||
for key, value in details_dict.items():
|
||||
if isinstance(value, int):
|
||||
quotient, remainder = divmod(value, num_items)
|
||||
|
|
@ -798,8 +796,8 @@ class Cache:
|
|||
if result.usage is None or result.usage.prompt_tokens is None:
|
||||
return None
|
||||
|
||||
total = result.usage.prompt_tokens
|
||||
num_items = len(result.data)
|
||||
total: Final = result.usage.prompt_tokens
|
||||
num_items: Final = len(result.data)
|
||||
if num_items <= 1:
|
||||
return total
|
||||
|
||||
|
|
@ -813,23 +811,23 @@ class Cache:
|
|||
kwargs: dict,
|
||||
idx_in_result_data: int = 0,
|
||||
) -> tuple[str, dict, dict]:
|
||||
preset_cache_key = self.get_cache_key(**{**kwargs, "input": input})
|
||||
preset_cache_key: Final = self.get_cache_key(**{**kwargs, "input": input})
|
||||
kwargs["cache_key"] = preset_cache_key
|
||||
embedding_response = result.data[idx_in_result_data]
|
||||
embedding_response: Final = result.data[idx_in_result_data]
|
||||
|
||||
# Extract per-item prompt_tokens + details from response usage
|
||||
prompt_tokens = self._get_per_item_prompt_tokens(
|
||||
prompt_tokens: Final = self._get_per_item_prompt_tokens(
|
||||
result=result,
|
||||
idx_in_result_data=idx_in_result_data,
|
||||
)
|
||||
prompt_tokens_details = self._get_per_item_prompt_tokens_details(
|
||||
prompt_tokens_details: Final = self._get_per_item_prompt_tokens_details(
|
||||
result=result,
|
||||
idx_in_result_data=idx_in_result_data,
|
||||
)
|
||||
|
||||
# Always convert to properly typed CachedEmbedding
|
||||
model_name = result.model
|
||||
embedding_dict: CachedEmbedding = self._convert_to_cached_embedding(
|
||||
model_name: Final = result.model
|
||||
embedding_dict: Final[CachedEmbedding] = self._convert_to_cached_embedding(
|
||||
embedding_response,
|
||||
model_name,
|
||||
prompt_tokens=prompt_tokens,
|
||||
|
|
@ -856,7 +854,7 @@ class Cache:
|
|||
if self.ttl is not None:
|
||||
kwargs["ttl"] = self.ttl
|
||||
|
||||
cache_list = []
|
||||
cache_list: Final = []
|
||||
if isinstance(kwargs["input"], list):
|
||||
for idx, i in enumerate(kwargs["input"]):
|
||||
(
|
||||
|
|
@ -874,7 +872,7 @@ class Cache:
|
|||
else:
|
||||
await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {e!s}")
|
||||
verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e)
|
||||
|
||||
def should_use_cache(self, **kwargs):
|
||||
"""
|
||||
|
|
@ -887,7 +885,7 @@ class Cache:
|
|||
return True
|
||||
|
||||
# when mode == default_off -> Cache is opt in only
|
||||
_cache = kwargs.get("cache", None)
|
||||
_cache: Final = kwargs.get("cache", None)
|
||||
verbose_logger.debug("should_use_cache: kwargs: %s; _cache: %s", kwargs, _cache)
|
||||
if _cache and isinstance(_cache, dict):
|
||||
if _cache.get("use-cache", False) is True:
|
||||
|
|
@ -899,13 +897,13 @@ class Cache:
|
|||
await self.cache.batch_cache_write(cache_key, cached_data, **kwargs)
|
||||
|
||||
async def ping(self):
|
||||
cache_ping = getattr(self.cache, "ping")
|
||||
cache_ping: Final = getattr(self.cache, "ping")
|
||||
if cache_ping:
|
||||
return await cache_ping()
|
||||
return None
|
||||
|
||||
async def delete_cache_keys(self, keys):
|
||||
cache_delete_cache_keys = getattr(self.cache, "delete_cache_keys")
|
||||
cache_delete_cache_keys: Final = getattr(self.cache, "delete_cache_keys")
|
||||
if cache_delete_cache_keys:
|
||||
return await cache_delete_cache_keys(keys)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -19,11 +19,7 @@ import datetime
|
|||
import inspect
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Callable, Generator
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Optional,
|
||||
)
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -76,7 +72,7 @@ class CachingHandlerResponse(BaseModel):
|
|||
embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
|
||||
|
||||
|
||||
in_memory_cache_obj = InMemoryCache()
|
||||
in_memory_cache_obj: Final = InMemoryCache()
|
||||
|
||||
|
||||
def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str, object]:
|
||||
|
|
@ -96,10 +92,10 @@ def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str
|
|||
|
||||
|
||||
def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
|
||||
cached_id = cached_result.get("id")
|
||||
cached_id: Final = cached_result.get("id")
|
||||
if isinstance(cached_id, str) and cached_id.startswith("chatcmpl"):
|
||||
return True
|
||||
obj = cached_result.get("object")
|
||||
obj: Final = cached_result.get("object")
|
||||
if isinstance(obj, str):
|
||||
return obj.startswith("chat.completion")
|
||||
return "choices" in cached_result
|
||||
|
|
@ -184,10 +180,10 @@ class LLMCachingHandler:
|
|||
#########################################################
|
||||
# Init cache timing metrics
|
||||
#########################################################
|
||||
cache_check_start_time = time.perf_counter()
|
||||
cache_check_start_time: Final = time.perf_counter()
|
||||
cache_check_end_time: float | None = None
|
||||
#########################################################
|
||||
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
|
||||
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
|
||||
kwargs["parent_otel_span"] = parent_otel_span
|
||||
|
||||
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
|
||||
|
|
@ -201,15 +197,15 @@ class LLMCachingHandler:
|
|||
|
||||
if cached_result is not None and not isinstance(cached_result, list):
|
||||
verbose_logger.debug("Cache Hit!")
|
||||
cache_hit = True
|
||||
end_time = datetime.datetime.now()
|
||||
cache_hit: Final = True
|
||||
end_time: Final = datetime.datetime.now()
|
||||
model, custom_llm_provider, _, _ = litellm.get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=kwargs.get("custom_llm_provider", None),
|
||||
api_base=kwargs.get("api_base", None),
|
||||
api_key=kwargs.get("api_key", None),
|
||||
)
|
||||
cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000
|
||||
cache_duration_ms: Final = (cache_check_end_time - cache_check_start_time) * 1000
|
||||
self._update_litellm_logging_obj_environment(
|
||||
logging_obj=logging_obj,
|
||||
model=model,
|
||||
|
|
@ -240,13 +236,13 @@ class LLMCachingHandler:
|
|||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
)
|
||||
cache_key = (
|
||||
cache_key: Final = (
|
||||
self.preset_cache_key
|
||||
or self.request_kwargs.get("cache_key")
|
||||
or litellm.cache.get_cache_key(**self.request_kwargs)
|
||||
)
|
||||
if hasattr(cached_result, "_hidden_params"):
|
||||
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
|
||||
cached_result._hidden_params["cache_key"] = cache_key
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
elif (
|
||||
call_type == CallTypes.aembedding.value
|
||||
|
|
@ -271,7 +267,7 @@ class LLMCachingHandler:
|
|||
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
|
||||
)
|
||||
|
||||
verbose_logger.debug(f"CACHE RESULT: {cached_result}")
|
||||
verbose_logger.debug("CACHE RESULT: %s", cached_result)
|
||||
return CachingHandlerResponse(
|
||||
cached_result=cached_result,
|
||||
final_embedding_cached_response=final_embedding_cached_response,
|
||||
|
|
@ -295,7 +291,7 @@ class LLMCachingHandler:
|
|||
if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function):
|
||||
args = args or ()
|
||||
# Now that we confirmed caching will happen, prepare kwargs
|
||||
new_kwargs = kwargs.copy()
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs.update(
|
||||
convert_args_to_kwargs(
|
||||
self.original_function,
|
||||
|
|
@ -326,8 +322,8 @@ class LLMCachingHandler:
|
|||
)
|
||||
|
||||
# LOG SUCCESS
|
||||
cache_hit = True
|
||||
end_time = datetime.datetime.now()
|
||||
cache_hit: Final = True
|
||||
end_time: Final = datetime.datetime.now()
|
||||
(
|
||||
model,
|
||||
custom_llm_provider,
|
||||
|
|
@ -354,13 +350,13 @@ class LLMCachingHandler:
|
|||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
)
|
||||
cache_key = (
|
||||
cache_key: Final = (
|
||||
self.preset_cache_key
|
||||
or self.request_kwargs.get("cache_key")
|
||||
or litellm.cache.get_cache_key(**self.request_kwargs)
|
||||
)
|
||||
if hasattr(cached_result, "_hidden_params"):
|
||||
cached_result._hidden_params["cache_key"] = cache_key # type: ignore
|
||||
cached_result._hidden_params["cache_key"] = cache_key
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
return CachingHandlerResponse(cached_result=cached_result)
|
||||
|
||||
|
|
@ -420,9 +416,9 @@ class LLMCachingHandler:
|
|||
|
||||
"""
|
||||
embedding_all_elements_cache_hit: bool = False
|
||||
remaining_list = []
|
||||
non_null_list = []
|
||||
kwargs_input_as_list = self.handle_kwargs_input_list_or_str(kwargs)
|
||||
remaining_list: Final = []
|
||||
non_null_list: Final = []
|
||||
kwargs_input_as_list: Final = self.handle_kwargs_input_list_or_str(kwargs)
|
||||
for idx, cr in enumerate(cached_result):
|
||||
if cr is None:
|
||||
remaining_list.append(kwargs_input_as_list[idx])
|
||||
|
|
@ -479,7 +475,7 @@ class LLMCachingHandler:
|
|||
prompt_tokens_details = PromptTokensDetailsWrapper(**aggregated_details)
|
||||
except Exception:
|
||||
prompt_tokens_details = None
|
||||
usage = Usage(
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=0,
|
||||
total_tokens=prompt_tokens,
|
||||
|
|
@ -488,9 +484,9 @@ class LLMCachingHandler:
|
|||
final_embedding_cached_response.usage = usage
|
||||
if len(remaining_list) == 0:
|
||||
# LOG SUCCESS
|
||||
cache_hit = True
|
||||
cache_hit: Final = True
|
||||
embedding_all_elements_cache_hit = True
|
||||
end_time = datetime.datetime.now()
|
||||
end_time: Final = datetime.datetime.now()
|
||||
(
|
||||
model,
|
||||
custom_llm_provider,
|
||||
|
|
@ -546,10 +542,10 @@ class LLMCachingHandler:
|
|||
if details2 is None:
|
||||
return details1
|
||||
|
||||
dict1 = details1.model_dump(exclude_none=True) if hasattr(details1, "model_dump") else {}
|
||||
dict2 = details2.model_dump(exclude_none=True) if hasattr(details2, "model_dump") else {}
|
||||
dict1: Final = details1.model_dump(exclude_none=True) if hasattr(details1, "model_dump") else {}
|
||||
dict2: Final = details2.model_dump(exclude_none=True) if hasattr(details2, "model_dump") else {}
|
||||
|
||||
merged: dict = {}
|
||||
merged: Final[dict] = {}
|
||||
for key in set(dict1.keys()) | set(dict2.keys()):
|
||||
v1 = dict1.get(key, 0)
|
||||
v2 = dict2.get(key, 0)
|
||||
|
|
@ -607,7 +603,7 @@ class LLMCachingHandler:
|
|||
return embedding_response
|
||||
|
||||
idx = 0
|
||||
final_data_list = []
|
||||
final_data_list: Final = []
|
||||
for item in _caching_handler_response.final_embedding_cached_response.data:
|
||||
if item is None and embedding_response.data is not None:
|
||||
final_data_list.append(embedding_response.data[idx])
|
||||
|
|
@ -690,7 +686,7 @@ class LLMCachingHandler:
|
|||
if litellm.cache is None:
|
||||
return None
|
||||
|
||||
new_kwargs = kwargs.copy()
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs.update(
|
||||
convert_args_to_kwargs(
|
||||
self.original_function,
|
||||
|
|
@ -708,7 +704,7 @@ class LLMCachingHandler:
|
|||
new_kwargs["input"] = [new_kwargs["input"]]
|
||||
elif not isinstance(new_kwargs["input"], list):
|
||||
raise ValueError("input must be a string or a list")
|
||||
tasks = []
|
||||
tasks: Final = []
|
||||
for idx, i in enumerate(new_kwargs["input"]):
|
||||
preset_cache_key = litellm.cache.get_cache_key(**{**new_kwargs, "input": i})
|
||||
tasks.append(
|
||||
|
|
@ -724,8 +720,8 @@ class LLMCachingHandler:
|
|||
if all(result is None for result in cached_result):
|
||||
cached_result = None
|
||||
else:
|
||||
request_kwargs = new_kwargs.copy()
|
||||
request_cache_key = request_kwargs.pop("cache_key", None)
|
||||
request_kwargs: Final = new_kwargs.copy()
|
||||
request_cache_key: Final = request_kwargs.pop("cache_key", None)
|
||||
if litellm.cache._supports_async() is True:
|
||||
## check if dual cache is supported ##
|
||||
self.preset_cache_key = request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
|
||||
|
|
@ -828,7 +824,7 @@ class LLMCachingHandler:
|
|||
elif (call_type == CallTypes.atranscription.value or call_type == CallTypes.transcription.value) and isinstance(
|
||||
cached_result, dict
|
||||
):
|
||||
hidden_params = {
|
||||
hidden_params: Final = {
|
||||
"model": "whisper-1",
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
"cache_hit": True,
|
||||
|
|
@ -840,10 +836,10 @@ class LLMCachingHandler:
|
|||
hidden_params=hidden_params,
|
||||
)
|
||||
elif (call_type == "aresponses" or call_type == "responses") and isinstance(cached_result, dict):
|
||||
use_chat_completion_cache = _is_chat_completion_cached_dict(cached_result)
|
||||
use_chat_completion_cache: Final = _is_chat_completion_cached_dict(cached_result)
|
||||
if use_chat_completion_cache:
|
||||
if kwargs.get("stream", False) is True:
|
||||
bridge_call_type = (
|
||||
bridge_call_type: Final = (
|
||||
CallTypes.acompletion.value if call_type == "aresponses" else CallTypes.completion.value
|
||||
)
|
||||
cached_result = self._convert_cached_stream_response(
|
||||
|
|
@ -862,7 +858,7 @@ class LLMCachingHandler:
|
|||
CachedResponsesAPIStreamingIterator,
|
||||
)
|
||||
|
||||
response_obj = ResponsesAPIResponse(**cached_result)
|
||||
response_obj: Final = ResponsesAPIResponse(**cached_result)
|
||||
if (
|
||||
hasattr(response_obj, "_hidden_params")
|
||||
and response_obj._hidden_params is not None
|
||||
|
|
@ -957,14 +953,14 @@ class LLMCachingHandler:
|
|||
if litellm.cache is None:
|
||||
return
|
||||
|
||||
new_kwargs = kwargs.copy()
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs.update(
|
||||
convert_args_to_kwargs(
|
||||
original_function,
|
||||
args,
|
||||
)
|
||||
)
|
||||
parent_otel_span = _get_parent_otel_span_from_kwargs(new_kwargs)
|
||||
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(new_kwargs)
|
||||
new_kwargs["parent_otel_span"] = parent_otel_span
|
||||
# [OPTIONAL] ADD TO CACHE
|
||||
if self._should_store_result_in_cache(original_function=original_function, kwargs=new_kwargs):
|
||||
|
|
@ -1006,7 +1002,7 @@ class LLMCachingHandler:
|
|||
Sync internal method to add the result to the cache
|
||||
"""
|
||||
|
||||
new_kwargs = kwargs.copy()
|
||||
new_kwargs: Final = kwargs.copy()
|
||||
new_kwargs.update(
|
||||
convert_args_to_kwargs(
|
||||
self.original_function,
|
||||
|
|
@ -1067,7 +1063,7 @@ class LLMCachingHandler:
|
|||
|
||||
"""
|
||||
|
||||
complete_streaming_response: ModelResponse | TextCompletionResponse | None = (
|
||||
complete_streaming_response: Final[ModelResponse | TextCompletionResponse | None] = (
|
||||
_assemble_complete_response_from_streaming_chunks(
|
||||
result=processed_chunk,
|
||||
start_time=self.start_time,
|
||||
|
|
@ -1089,7 +1085,7 @@ class LLMCachingHandler:
|
|||
"""
|
||||
Sync internal method to add the streaming response to the cache
|
||||
"""
|
||||
complete_streaming_response: ModelResponse | TextCompletionResponse | None = (
|
||||
complete_streaming_response: Final[ModelResponse | TextCompletionResponse | None] = (
|
||||
_assemble_complete_response_from_streaming_chunks(
|
||||
result=processed_chunk,
|
||||
start_time=self.start_time,
|
||||
|
|
@ -1133,7 +1129,7 @@ class LLMCachingHandler:
|
|||
Returns:
|
||||
None
|
||||
"""
|
||||
litellm_params = {
|
||||
litellm_params: Final = {
|
||||
"logger_fn": kwargs.get("logger_fn", None),
|
||||
"acompletion": is_async,
|
||||
"api_base": kwargs.get("api_base", ""),
|
||||
|
|
@ -1173,13 +1169,13 @@ def convert_args_to_kwargs(
|
|||
args: tuple[Any, ...] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
# Get the signature of the original function
|
||||
signature = inspect.signature(original_function)
|
||||
signature: Final = inspect.signature(original_function)
|
||||
|
||||
# Get parameter names in the order they appear in the original function
|
||||
param_names = list(signature.parameters.keys())
|
||||
param_names: Final = list(signature.parameters.keys())
|
||||
|
||||
# Create a mapping of positional arguments to parameter names
|
||||
args_to_kwargs = {}
|
||||
args_to_kwargs: Final = {}
|
||||
if args:
|
||||
for index, arg in enumerate(args):
|
||||
if index < len(param_names):
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
import json
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from .base_cache import BaseCache
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
Span = Union[_Span, Any]
|
||||
Span = _Span | Any
|
||||
else:
|
||||
Span = Any
|
||||
|
||||
|
|
@ -41,17 +41,17 @@ class DiskCache(BaseCache):
|
|||
self.set_cache(key=cache_key, value=cache_value)
|
||||
|
||||
def get_cache(self, key, **kwargs):
|
||||
original_cached_response = self.disk_cache.get(key)
|
||||
original_cached_response: Final = self.disk_cache.get(key)
|
||||
if original_cached_response:
|
||||
try:
|
||||
cached_response = json.loads(original_cached_response) # type: ignore
|
||||
cached_response = json.loads(original_cached_response)
|
||||
except Exception:
|
||||
cached_response = original_cached_response
|
||||
return cached_response
|
||||
return None
|
||||
|
||||
def batch_get_cache(self, keys: list, **kwargs):
|
||||
return_val = []
|
||||
return_val: Final = []
|
||||
for k in keys:
|
||||
val = self.get_cache(key=k, **kwargs)
|
||||
return_val.append(val)
|
||||
|
|
@ -59,9 +59,9 @@ class DiskCache(BaseCache):
|
|||
|
||||
def increment_cache(self, key, value: int, **kwargs) -> int:
|
||||
with self.disk_cache.transact():
|
||||
cached_value = self.get_cache(key=key)
|
||||
init_value = cached_value if isinstance(cached_value, int) else 0
|
||||
new_value = init_value + value
|
||||
cached_value: Final = self.get_cache(key=key)
|
||||
init_value: Final = cached_value if isinstance(cached_value, int) else 0
|
||||
new_value: Final = init_value + value
|
||||
self.set_cache(key, new_value, **kwargs)
|
||||
return new_value
|
||||
|
||||
|
|
@ -69,7 +69,7 @@ class DiskCache(BaseCache):
|
|||
return self.get_cache(key=key, **kwargs)
|
||||
|
||||
async def async_batch_get_cache(self, keys: list, **kwargs):
|
||||
return_val = []
|
||||
return_val: Final = []
|
||||
for k in keys:
|
||||
val = self.get_cache(key=k, **kwargs)
|
||||
return_val.append(val)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import time
|
|||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from threading import Lock
|
||||
from typing import TYPE_CHECKING, Any, Union
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
|
|
@ -29,7 +29,7 @@ from .redis_cache import RedisCache
|
|||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
|
||||
Span = Union[_Span, Any]
|
||||
Span = _Span | Any
|
||||
else:
|
||||
Span = Any
|
||||
|
||||
|
|
@ -147,7 +147,7 @@ class DualCache(BaseCache):
|
|||
|
||||
return result
|
||||
except Exception as e:
|
||||
verbose_logger.error(f"LiteLLM Cache: Excepton async add_cache: {e!s}")
|
||||
verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e)
|
||||
raise e
|
||||
|
||||
def get_cache(
|
||||
|
|
@ -161,14 +161,14 @@ class DualCache(BaseCache):
|
|||
try:
|
||||
result = None
|
||||
if self.in_memory_cache is not None:
|
||||
in_memory_result = self.in_memory_cache.get_cache(key, **kwargs)
|
||||
in_memory_result: Final = self.in_memory_cache.get_cache(key, **kwargs)
|
||||
|
||||
if in_memory_result is not None:
|
||||
result = in_memory_result
|
||||
|
||||
if result is None and self.redis_cache is not None and local_only is False:
|
||||
# If not found in in-memory cache, try fetching from Redis
|
||||
redis_result = self.redis_cache.get_cache(key, parent_otel_span=parent_otel_span)
|
||||
redis_result: Final = self.redis_cache.get_cache(key, parent_otel_span=parent_otel_span)
|
||||
|
||||
if redis_result is not None:
|
||||
# Update in-memory cache with the value from Redis
|
||||
|
|
@ -188,12 +188,12 @@ class DualCache(BaseCache):
|
|||
local_only: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
received_args = locals()
|
||||
received_args: Final = locals()
|
||||
received_args.pop("self")
|
||||
|
||||
def run_in_new_loop():
|
||||
"""Run the coroutine in a new event loop within this thread."""
|
||||
new_loop = asyncio.new_event_loop()
|
||||
new_loop: Final = asyncio.new_event_loop()
|
||||
try:
|
||||
asyncio.set_event_loop(new_loop)
|
||||
return new_loop.run_until_complete(self.async_batch_get_cache(**received_args))
|
||||
|
|
@ -207,7 +207,7 @@ class DualCache(BaseCache):
|
|||
# If we're already in an event loop, run in a separate thread
|
||||
# to avoid nested event loop issues
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future = executor.submit(run_in_new_loop)
|
||||
future: Final = executor.submit(run_in_new_loop)
|
||||
return future.result()
|
||||
|
||||
except RuntimeError:
|
||||
|
|
@ -226,7 +226,7 @@ class DualCache(BaseCache):
|
|||
print_verbose(f"async get cache: cache key: {key}; local_only: {local_only}")
|
||||
result = None
|
||||
if self.in_memory_cache is not None:
|
||||
in_memory_result = await self.in_memory_cache.async_get_cache(key, **kwargs)
|
||||
in_memory_result: Final = await self.in_memory_cache.async_get_cache(key, **kwargs)
|
||||
|
||||
print_verbose(f"in_memory_result: {in_memory_result}")
|
||||
if in_memory_result is not None:
|
||||
|
|
@ -234,7 +234,7 @@ class DualCache(BaseCache):
|
|||
|
||||
if result is None and self.redis_cache is not None and local_only is False:
|
||||
# If not found in in-memory cache, try fetching from Redis
|
||||
redis_result = await self.redis_cache.async_get_cache(key, parent_otel_span=parent_otel_span)
|
||||
redis_result: Final = await self.redis_cache.async_get_cache(key, parent_otel_span=parent_otel_span)
|
||||
|
||||
if redis_result is not None:
|
||||
# Update in-memory cache with the value from Redis
|
||||
|
|
@ -257,8 +257,8 @@ class DualCache(BaseCache):
|
|||
Atomically choose keys to fetch from Redis and reserve their access time.
|
||||
This prevents check-then-act races under concurrent async callers.
|
||||
"""
|
||||
sublist_keys: list[str] = []
|
||||
previous_access_times: dict[str, float | None] = {}
|
||||
sublist_keys: Final[list[str]] = []
|
||||
previous_access_times: Final[dict[str, float | None]] = {}
|
||||
|
||||
with self._last_redis_batch_access_time_lock:
|
||||
for key, value in zip(keys, result):
|
||||
|
|
@ -293,7 +293,7 @@ class DualCache(BaseCache):
|
|||
try:
|
||||
result = [None] * len(keys)
|
||||
if self.in_memory_cache is not None:
|
||||
in_memory_result = await self.in_memory_cache.async_batch_get_cache(keys, **kwargs)
|
||||
in_memory_result: Final = await self.in_memory_cache.async_batch_get_cache(keys, **kwargs)
|
||||
|
||||
if in_memory_result is not None:
|
||||
result = in_memory_result
|
||||
|
|
@ -303,14 +303,14 @@ class DualCache(BaseCache):
|
|||
- for the none values in the result
|
||||
- check the redis cache
|
||||
"""
|
||||
current_time = time.time()
|
||||
current_time: Final = time.time()
|
||||
sublist_keys, previous_access_times = self._reserve_redis_batch_keys(current_time, keys, result)
|
||||
|
||||
# Only hit Redis if enough time has passed since last access.
|
||||
if len(sublist_keys) > 0:
|
||||
try:
|
||||
# If not found in in-memory cache, try fetching from Redis
|
||||
redis_result = await self.redis_cache.async_batch_get_cache(
|
||||
redis_result: Final = await self.redis_cache.async_batch_get_cache(
|
||||
sublist_keys, parent_otel_span=parent_otel_span
|
||||
)
|
||||
except Exception:
|
||||
|
|
@ -323,7 +323,7 @@ class DualCache(BaseCache):
|
|||
return result
|
||||
|
||||
# Pre-compute key-to-index mapping for O(1) lookup
|
||||
key_to_index = {key: i for i, key in enumerate(keys)}
|
||||
key_to_index: Final = {key: i for i, key in enumerate(keys)}
|
||||
|
||||
# Update both result and in-memory cache in a single loop
|
||||
for key, value in redis_result.items():
|
||||
|
|
@ -347,7 +347,7 @@ class DualCache(BaseCache):
|
|||
if self.redis_cache is not None and local_only is False:
|
||||
await self.redis_cache.async_set_cache(key, value, **kwargs)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {e!s}")
|
||||
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
|
||||
|
||||
# async_batch_set_cache
|
||||
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs):
|
||||
|
|
@ -366,7 +366,7 @@ class DualCache(BaseCache):
|
|||
cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {e!s}")
|
||||
verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e)
|
||||
|
||||
async def async_increment_cache(
|
||||
self,
|
||||
|
|
|
|||
276
litellm/caching/evicted_client_closer.py
Normal file
276
litellm/caching/evicted_client_closer.py
Normal file
|
|
@ -0,0 +1,276 @@
|
|||
"""
|
||||
Deferred close of HTTP/SDK clients that the LLM client cache has evicted.
|
||||
|
||||
Eviction only drops the cache's reference to a client. Every OpenAI/Azure SDK
|
||||
client is a reference cycle (each resource namespace holds the client back), so
|
||||
an evicted client and its pooled TCP connections survive until a generational
|
||||
collection runs, which under load is thousands of requests later.
|
||||
|
||||
Closing at eviction time is not an option: a request that was handed the client
|
||||
just before it was evicted is still using it, and closing it underneath that
|
||||
request raises ``RuntimeError: Cannot send a request, as the client has been
|
||||
closed.``
|
||||
|
||||
So an evicted client is closed once two conditions hold. A grace window must
|
||||
have passed since its eviction, which covers a request that holds the client
|
||||
but is momentarily not on the wire, and the client must report no connection in
|
||||
flight. The second condition is what keeps the first honest: a request may run
|
||||
for ``litellm.request_timeout`` seconds, 6000 by default, and a streaming
|
||||
response is bounded only by how long the upstream keeps sending, so no deadline
|
||||
on its own can promise that a request has finished.
|
||||
|
||||
Only clients litellm itself created are closed; a client the caller supplied is
|
||||
left alone because litellm does not own its lifecycle.
|
||||
|
||||
A client that closes synchronously is closed from wherever the cache is next
|
||||
used. One whose close is a coroutine needs the event loop it was evicted on, so
|
||||
it waits for a call from that loop rather than having work scheduled onto a loop
|
||||
it does not belong to. Queued clients are therefore bucketed by what it takes to
|
||||
close them, and each bucket is ordered by deadline, so a reap walks the entries
|
||||
that are due rather than the whole queue.
|
||||
|
||||
The queue holds its clients weakly, so waiting out a grace window never keeps
|
||||
alive anything the collector would have reclaimed first.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import inspect
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from collections import deque
|
||||
from collections.abc import Awaitable, Callable, Iterator
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import (
|
||||
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
|
||||
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
|
||||
)
|
||||
|
||||
_CLOSABLE_ANYWHERE: Final = "closable-anywhere"
|
||||
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
|
||||
|
||||
_BucketKey = str | int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PendingClose:
|
||||
"""A queued close.
|
||||
|
||||
The client is held weakly, so queueing one never keeps alive anything the
|
||||
collector would otherwise have reclaimed first.
|
||||
|
||||
``needs_loop`` is set for a client whose close is a coroutine; those can only
|
||||
be closed from the event loop they were evicted on, recorded in ``loop_id``.
|
||||
A client that closes synchronously carries neither constraint.
|
||||
"""
|
||||
|
||||
client_ref: "weakref.ref[object]"
|
||||
loop_id: int | None
|
||||
needs_loop: bool
|
||||
close_after: float
|
||||
|
||||
|
||||
def _bucket_key(pending: _PendingClose) -> _BucketKey:
|
||||
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
|
||||
if not pending.needs_loop:
|
||||
return _CLOSABLE_ANYWHERE
|
||||
if pending.loop_id is None:
|
||||
return _CLOSABLE_ON_ANY_LOOP
|
||||
return pending.loop_id
|
||||
|
||||
|
||||
def _running_loop_id() -> int | None:
|
||||
try:
|
||||
return id(asyncio.get_running_loop())
|
||||
except RuntimeError:
|
||||
return None
|
||||
|
||||
|
||||
def _close_function(client: object) -> Callable[[], object] | None:
|
||||
close_fn: Final[Callable[[], object] | None] = getattr(client, "aclose", None) or getattr(client, "close", None)
|
||||
return close_fn
|
||||
|
||||
|
||||
def _transport_of(client: object) -> object:
|
||||
"""The httpx transport behind an SDK wrapper, a litellm handler, or a bare client."""
|
||||
for holder in (getattr(client, "_client", None), getattr(client, "client", None), client):
|
||||
transport: object = getattr(holder, "_transport", None)
|
||||
if transport is not None:
|
||||
return transport
|
||||
return None
|
||||
|
||||
|
||||
def _connection_is_idle(connection: object) -> bool:
|
||||
"""A pooled connection is idle unless it is servicing a request."""
|
||||
is_idle: Final[object] = getattr(connection, "is_idle", None)
|
||||
return bool(is_idle()) if callable(is_idle) else True
|
||||
|
||||
|
||||
def _pool_has_busy_connection(transport: object) -> bool | None:
|
||||
"""Whether the httpcore pool behind the transport is servicing a request.
|
||||
|
||||
``None`` when there is no such pool, so the caller can ask the other backend.
|
||||
"""
|
||||
pooled: Final[object] = getattr(getattr(transport, "_pool", None), "connections", None)
|
||||
if not isinstance(pooled, (list, tuple)):
|
||||
return None
|
||||
return any(
|
||||
not _connection_is_idle(connection) # pyright: ignore[reportUnknownArgumentType] # untyped pool list
|
||||
for connection in pooled # pyright: ignore[reportUnknownVariableType] # untyped pool list
|
||||
)
|
||||
|
||||
|
||||
def _has_connection_in_flight(client: object) -> bool:
|
||||
"""Whether the client is servicing a request right now.
|
||||
|
||||
Both connection backends litellm uses already account for the connections
|
||||
they have handed out, so this reads the client's own lease accounting rather
|
||||
than inferring it from elapsed time: httpcore reports a non-idle connection
|
||||
for the whole of a response including a stream, and aiohttp holds the
|
||||
connection in ``_acquired`` over the same span.
|
||||
|
||||
A client that cannot answer is reported as idle, which leaves the grace
|
||||
window as the only guard, exactly as it was before this check existed.
|
||||
"""
|
||||
try:
|
||||
transport: Final = _transport_of(client)
|
||||
pooled_busy: Final = _pool_has_busy_connection(transport)
|
||||
if pooled_busy is not None:
|
||||
return pooled_busy
|
||||
session: Final[object] = getattr(transport, "client", None)
|
||||
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
|
||||
except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle
|
||||
return False
|
||||
|
||||
|
||||
async def _close_quietly(closing: Awaitable[object]) -> None:
|
||||
with contextlib.suppress(Exception):
|
||||
await closing
|
||||
|
||||
|
||||
class EvictedClientCloser:
|
||||
"""Closes evicted, litellm-owned clients once they are idle and out of grace."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
grace_seconds: float = EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
|
||||
max_pending: int = EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._grace_seconds = grace_seconds
|
||||
self._max_pending = max_pending
|
||||
self._clock = clock
|
||||
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
|
||||
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
|
||||
self._pending_count = 0
|
||||
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
|
||||
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
|
||||
|
||||
def mark_owned(self, client: object) -> None:
|
||||
"""Record that litellm created this client, so it may be closed on eviction."""
|
||||
try:
|
||||
self._owned.add(client)
|
||||
except TypeError:
|
||||
pass # values that cannot be weak-referenced are never litellm clients
|
||||
|
||||
def _is_owned(self, client: object) -> bool:
|
||||
try:
|
||||
return client in self._owned
|
||||
except TypeError:
|
||||
return False # unhashable values are never litellm clients
|
||||
|
||||
def schedule(self, client: object) -> None:
|
||||
"""Queue an evicted client for closing once it is idle and out of grace.
|
||||
|
||||
Past ``max_pending`` the client is left to the collector instead, so a
|
||||
workload that churns the cache cannot grow this queue without bound.
|
||||
Every queued entry comes due within one grace window, so the capacity it
|
||||
occupies is returned within that window rather than held.
|
||||
"""
|
||||
if client is None or not self._is_owned(client):
|
||||
return
|
||||
close_fn: Final = _close_function(client)
|
||||
if close_fn is None:
|
||||
return
|
||||
if self._pending_count >= self._max_pending:
|
||||
return
|
||||
self._enqueue(
|
||||
_PendingClose(
|
||||
client_ref=weakref.ref(client),
|
||||
loop_id=_running_loop_id(),
|
||||
needs_loop=inspect.iscoroutinefunction(close_fn),
|
||||
close_after=self._clock() + self._grace_seconds,
|
||||
)
|
||||
)
|
||||
|
||||
def reap(self) -> None:
|
||||
"""Close every queued client that is due, idle, and closable from here.
|
||||
|
||||
Called from the cache's read path, so the empty-queue exit comes first and
|
||||
the work done past it is proportional to what is due, not to the queue.
|
||||
"""
|
||||
if not self._pending_count:
|
||||
return
|
||||
now: Final = self._clock()
|
||||
for pending in self._take_due(_running_loop_id(), now):
|
||||
client = pending.client_ref()
|
||||
if client is None:
|
||||
continue
|
||||
if _has_connection_in_flight(client):
|
||||
self._enqueue(replace(pending, close_after=now + self._grace_seconds))
|
||||
continue
|
||||
self._close(client)
|
||||
|
||||
@property
|
||||
def pending_count(self) -> int:
|
||||
return self._pending_count
|
||||
|
||||
def _enqueue(self, pending: _PendingClose) -> None:
|
||||
"""Append to the entry's bucket, dropping any dead entries it queues behind.
|
||||
|
||||
Deadlines only ever move forward, so appending keeps each bucket ordered
|
||||
by deadline, and entries whose client the collector already took sit at
|
||||
the front rather than having to be searched for.
|
||||
"""
|
||||
with self._queue_lock:
|
||||
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque()) # mutable-ok: FIFO by design
|
||||
while bucket and bucket[0].client_ref() is None:
|
||||
bucket.popleft()
|
||||
self._pending_count -= 1
|
||||
bucket.append(pending)
|
||||
self._pending_count += 1
|
||||
|
||||
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
|
||||
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
|
||||
with self._queue_lock:
|
||||
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
|
||||
|
||||
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:
|
||||
bucket: Final = self._buckets.get(key)
|
||||
if bucket is None:
|
||||
return
|
||||
while bucket and bucket[0].close_after <= now:
|
||||
self._pending_count -= 1
|
||||
yield bucket.popleft()
|
||||
if not bucket:
|
||||
del self._buckets[key]
|
||||
|
||||
def _close(self, client: object) -> None:
|
||||
close_fn: Final = _close_function(client)
|
||||
if close_fn is None:
|
||||
return
|
||||
try:
|
||||
closing: Final = close_fn()
|
||||
except Exception: # noqa: BLE001 - a discarded client's close must never surface to callers
|
||||
return
|
||||
if not inspect.isawaitable(closing):
|
||||
return
|
||||
task: Final = asyncio.get_running_loop().create_task(_close_quietly(closing))
|
||||
self._close_tasks.add(task)
|
||||
task.add_done_callback(self._close_tasks.discard)
|
||||
|
||||
|
||||
default_evicted_client_closer: Final = EvictedClientCloser()
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue