diff --git a/.circleci/config.yml b/.circleci/config.yml index 2f01b6de4f3..cc485aa0595 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -88,6 +88,59 @@ commands: rm -f /tmp/uv-install.sh echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$BASH_ENV" export PATH="$HOME/.local/bin:$PATH" + install_node: + description: "Install the Node.js version pinned in ui/litellm-dashboard/.nvmrc (24.19.0, which bundles npm 11.17.0) with checksum verification, and prepend it to PATH. Run this on any executor whose image does not already ship that version, or `npm ci` in ui/litellm-dashboard fails EBADENGINE against the engines floor. Installs into /opt/node rather than over /usr/local on purpose: cimg/python:*-browsers ships its own node there, and unpacking the tarball on top of it leaves npm 11.17 files merged with the image's npm 11.9 tree, which reports the new version and then exits 1 on `npm ci` with no error text at all. Requires checkout, which the .nvmrc drift check reads." + steps: + - run: + name: Install Node.js 24.19.0 + command: | + NODE_VERSION="24.19.0" + NODE_TARBALL="node-v${NODE_VERSION}-linux-x64.tar.xz" + NODE_EXPECTED_SHA="14b342e71204f811bde6153be8e04b62aef63c236fef92b55f9c83154b409647" + NVMRC_VERSION="$(tr -d '[:space:]' < ui/litellm-dashboard/.nvmrc)" + if [ "$NVMRC_VERSION" != "$NODE_VERSION" ]; then + echo "install_node: ui/litellm-dashboard/.nvmrc pins ${NVMRC_VERSION} but this command pins ${NODE_VERSION}; update NODE_VERSION and NODE_EXPECTED_SHA together" >&2 + exit 1 + fi + curl -sSLf -o "/tmp/${NODE_TARBALL}" "https://nodejs.org/dist/v${NODE_VERSION}/${NODE_TARBALL}" + echo "${NODE_EXPECTED_SHA} /tmp/${NODE_TARBALL}" | sha256sum -c - + sudo mkdir -p /opt/node + sudo tar -xJf "/tmp/${NODE_TARBALL}" -C /opt/node --strip-components=1 + rm -f "/tmp/${NODE_TARBALL}" + echo 'export PATH="/opt/node/bin:$PATH"' >> "$BASH_ENV" + export PATH="/opt/node/bin:$PATH" + node --version + npm --version + install_rust: + description: "Install pinned rustup (1.28.2) and Rust toolchain (1.97.1) with checksum verification. Adds ~/.cargo/bin to PATH. Run this before any `uv sync` or `uv build` of the workspace: the root package builds litellm-rust through maturin, and on an image without cargo maturin fetches an unpinned rustup and a floating toolchain by itself." + steps: + - run: + name: Install Rust (rustup 1.28.2, toolchain 1.97.1) + command: | + case "$(uname -m)" in + x86_64) + RUSTUP_TRIPLE=x86_64-unknown-linux-gnu + RUSTUP_SHA256=20a06e644b0d9bd2fbdbfd52d42540bdde820ea7df86e92e533c073da0cdd43c + ;; + aarch64) + RUSTUP_TRIPLE=aarch64-unknown-linux-gnu + RUSTUP_SHA256=e3853c5a252fca15252d07cb23a1bdd9377a8c6f3efa01531109281ae47f841c + ;; + *) + echo "install_rust: unsupported architecture $(uname -m)" >&2 + exit 1 + ;; + esac + curl -sSLf -o /tmp/rustup-init \ + "https://static.rust-lang.org/rustup/archive/1.28.2/${RUSTUP_TRIPLE}/rustup-init" + echo "${RUSTUP_SHA256} /tmp/rustup-init" | sha256sum -c - + chmod +x /tmp/rustup-init + /tmp/rustup-init -y --no-modify-path --profile minimal --default-toolchain 1.97.1 + rm -f /tmp/rustup-init + echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> "$BASH_ENV" + export PATH="$HOME/.cargo/bin:$PATH" + rustc --version + cargo --version start_postgres: description: "Start a postgres-db container on port 5432 and wait until it accepts connections." parameters: @@ -163,6 +216,26 @@ commands: done echo "fake OpenAI endpoint did not become ready" >&2 exit 1 + start_cost_center_service: + description: "Start the stand-in cost center validation service (tests/store_model_in_db_tests/cost_center_service.py) on host port 9414 and wait until healthy. The proxy's team-metadata validator (team_metadata_validator_e2e.py, impl 'http') reaches it via TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate. Run after uv deps are synced." + steps: + - run: + name: Start cost center validation service + background: true + command: | + uv run --no-sync python tests/store_model_in_db_tests/cost_center_service.py --host 0.0.0.0 --port 9414 + - run: + name: Wait for cost center validation service + command: | + for i in $(seq 1 30); do + if curl -sf http://localhost:9414/health >/dev/null 2>&1; then + echo "cost center validation service is up" + exit 0 + fi + sleep 1 + done + echo "cost center validation service did not become ready" >&2 + exit 1 setup_litellm_enterprise_pip: steps: - run: @@ -178,6 +251,7 @@ commands: - checkout - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -281,6 +355,33 @@ jobs: uv build --wheel --out-dir dist uv run --no-sync python tests/windows_tests/check_windows_wheel_install.py + base_sdk_install: + docker: + - image: cimg/python:3.12@sha256:9c796c23c84e84a66a964acb508d39dc5433c81a47e07efd56dccbbc2427e07c + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + steps: + - checkout + - setup_google_dns + - install_uv + - install_rust + - run: + name: Build the wheel + environment: + UV_HTTP_TIMEOUT: "300" + command: | + uv build --wheel --out-dir dist + - run: + name: Install the wheel with no extras and smoke-check it + environment: + UV_HTTP_TIMEOUT: "300" + command: | + uv venv /tmp/base-sdk --python 3.12 + VIRTUAL_ENV=/tmp/base-sdk uv pip install dist/*.whl + /tmp/base-sdk/bin/python tests/base_sdk_tests/check_base_sdk_install.py + local_testing_part1: docker: - &python312_image @@ -298,6 +399,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -371,6 +473,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -445,6 +548,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -496,6 +600,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -562,6 +667,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -602,6 +708,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -643,6 +750,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -676,6 +784,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -726,6 +835,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -777,6 +887,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -810,6 +921,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -856,6 +968,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -902,6 +1015,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -944,6 +1058,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -990,6 +1105,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1037,6 +1153,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -1077,6 +1194,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1122,6 +1240,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1166,6 +1285,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1198,6 +1318,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1241,6 +1362,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1285,6 +1407,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1329,6 +1452,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1360,6 +1484,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1406,6 +1531,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1451,6 +1577,7 @@ jobs: keys: - v1-uv-cache-{{ checksum "uv.lock" }} - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1501,6 +1628,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1525,6 +1653,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1551,6 +1680,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1652,6 +1782,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1747,6 +1878,7 @@ jobs: at: ~/project - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1835,6 +1967,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -1918,6 +2051,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2050,6 +2184,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2136,6 +2271,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2232,12 +2368,14 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | uv sync --frozen --all-groups --all-extras --python 3.12 - start_postgres - start_fake_openai_endpoint + - start_cost_center_service - attach_workspace: at: ~/project - run: @@ -2257,11 +2395,13 @@ jobs: -e STORE_MODEL_IN_DB="True" \ -e LITELLM_MASTER_KEY="sk-1234" \ -e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \ + -e TEAM_METADATA_VALIDATION_SERVICE_URL=http://host.docker.internal:9414/validate \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ -e LITELLM_LOG=ERROR \ --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ + -v $(pwd)/litellm/proxy/example_config_yaml/team_metadata_validator_e2e.py:/app/team_metadata_validator_e2e.py \ litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 @@ -2307,6 +2447,7 @@ jobs: - setup_google_dns # Remove Docker CLI installation since it's already available in machine executor - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2388,6 +2529,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2473,20 +2615,9 @@ jobs: bundle exec rspec no_output_timeout: 30m # Install Node.js directly from nodejs.org with SHA256 verification, - # instead of piping NodeSource's setup_18.x apt-repo installer into + # instead of piping NodeSource's setup_24.x apt-repo installer into # sudo bash (which runs a mutable upstream script unattended). - - run: - name: Install Node.js 18.20.8 - command: | - NODE_VERSION="18.20.8" - NODE_TARBALL="node-v${NODE_VERSION}-linux-x64.tar.xz" - NODE_EXPECTED_SHA="5467ee62d6af1411d46b6a10e3fb5cacc92734dbcef465fea14e7b90993001c9" - curl -sSLf -o "/tmp/${NODE_TARBALL}" "https://nodejs.org/dist/v${NODE_VERSION}/${NODE_TARBALL}" - echo "${NODE_EXPECTED_SHA} /tmp/${NODE_TARBALL}" | sha256sum -c - - sudo tar -xJf "/tmp/${NODE_TARBALL}" -C /usr/local --strip-components=1 - rm -f "/tmp/${NODE_TARBALL}" - node --version - npm --version + - install_node - run: name: Install Node.js test dependencies @@ -2527,6 +2658,7 @@ jobs: - skip_if_unrelated_changes - setup_google_dns - install_uv + - install_rust - run: name: Install Dependencies command: | @@ -2614,7 +2746,7 @@ jobs: ui_build: docker: - - image: cimg/node:20.19@sha256:35e64883e8d21bc345b0a7b04c35ee46442c127607ed1d8d7d37d8a1ed76db81 + - image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad auth: username: ${DOCKERHUB_USERNAME} password: ${DOCKERHUB_PASSWORD} @@ -2658,7 +2790,7 @@ jobs: ui_unit_tests: docker: - - image: cimg/node:20.19@sha256:35e64883e8d21bc345b0a7b04c35ee46442c127607ed1d8d7d37d8a1ed76db81 + - image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad auth: username: ${DOCKERHUB_USERNAME} password: ${DOCKERHUB_PASSWORD} @@ -2716,7 +2848,9 @@ jobs: - skip_if_unrelated_changes: category: client - setup_google_dns + - install_node - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -2731,7 +2865,7 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} + - ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright # The cimg/python:3.12-browsers image already ships the Chromium system @@ -2746,7 +2880,7 @@ jobs: npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} + key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules - tests/e2e/ui/node_modules @@ -2858,7 +2992,9 @@ jobs: - skip_if_unrelated_changes: category: client - setup_google_dns + - install_node - install_uv + - install_rust - restore_cache: keys: - v1-uv-cache-{{ checksum "uv.lock" }} @@ -2873,7 +3009,7 @@ jobs: - ~/.cache/uv - restore_cache: keys: - - ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} + - ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} - run: name: Install Node dependencies and Playwright command: | @@ -2883,7 +3019,7 @@ jobs: npm ci npx playwright install chromium - save_cache: - key: ui-e2e-node-deps-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} + key: ui-e2e-node-deps-v4-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }} paths: - ui/litellm-dashboard/node_modules - tests/e2e/ui/node_modules @@ -3031,6 +3167,8 @@ workflows: only: - main - /litellm_.*/ + - base_sdk_install: + filters: *main_branches - local_testing_part1: filters: *main_branches - local_testing_part2: diff --git a/.flake8 b/.flake8 deleted file mode 100644 index afd4596076b..00000000000 --- a/.flake8 +++ /dev/null @@ -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 diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 9c24bad00f1..597daebd720 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -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 diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index 4d4a3242399..23aa6114f08 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -9,6 +9,8 @@ on: - "litellm_**" paths: - docker/Dockerfile.non_root + - migrations/Dockerfile + - migrations/run.py - tests/proxy_migration_tests/test_offline_image_migration.py - uv.lock - ui/litellm-dashboard/package-lock.json @@ -83,3 +85,34 @@ jobs: --only-fixed \ --fail-on high \ --output table + + migrations-image: + name: migrations-image + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build migrations image + run: docker build -f migrations/Dockerfile -t litellm-migrations-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify offline migration as a non-root uid + env: + LITELLM_IMAGE: litellm-migrations-scan:${{ github.sha }} + LITELLM_MIGRATION_INTERPRETER: python3 + LITELLM_MIGRATION_SCRIPT: /app/run.py + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index ae31395521a..fab05fc2bbb 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -7,13 +7,17 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: code-quality: diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 8d2b2c2f972..b539ec4be88 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -104,9 +104,8 @@ jobs: - name: Check basedpyright budget (delta vs base) env: BASE_SHA: ${{ github.event.pull_request.base.sha }} - NODE_OPTIONS: --max-old-space-size=12288 run: | - (uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA" + uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA" - name: Check tests/e2e basedpyright (zero errors) env: diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml index 525e2c5b949..39f4bc1428a 100644 --- a/.github/workflows/test-litellm-ui-build.yml +++ b/.github/workflows/test-litellm-ui-build.yml @@ -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 diff --git a/.github/workflows/test-litellm-ui-lint.yml b/.github/workflows/test-litellm-ui-lint.yml index 5173eb6da35..2c31abc6609 100644 --- a/.github/workflows/test-litellm-ui-lint.yml +++ b/.github/workflows/test-litellm-ui-lint.yml @@ -61,7 +61,7 @@ jobs: if: steps.changed.outputs.has_files == 'true' uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0 with: - node-version: "20" + node-version-file: ui/litellm-dashboard/.nvmrc cache: "npm" cache-dependency-path: ui/litellm-dashboard/package-lock.json diff --git a/.github/workflows/test-litellm-ui-unit.yml b/.github/workflows/test-litellm-ui-unit.yml index 5374a0059de..20c3611eb38 100644 --- a/.github/workflows/test-litellm-ui-unit.yml +++ b/.github/workflows/test-litellm-ui-unit.yml @@ -35,7 +35,7 @@ jobs: - name: Setup Node.js uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0 with: - node-version: "20" + node-version-file: ui/litellm-dashboard/.nvmrc cache: "npm" cache-dependency-path: ui/litellm-dashboard/package-lock.json diff --git a/.github/workflows/test-unit-core-utils.yml b/.github/workflows/test-unit-core-utils.yml index d6d6353238f..a01f09559c6 100644 --- a/.github/workflows/test-unit-core-utils.yml +++ b/.github/workflows/test-unit-core-utils.yml @@ -7,6 +7,10 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -14,8 +18,8 @@ permissions: pull-requests: write concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: core-utils: diff --git a/.github/workflows/test-unit-documentation.yml b/.github/workflows/test-unit-documentation.yml index c12a289ce9f..50589cb5926 100644 --- a/.github/workflows/test-unit-documentation.yml +++ b/.github/workflows/test-unit-documentation.yml @@ -7,13 +7,17 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: documentation: diff --git a/.github/workflows/test-unit-enterprise-routing.yml b/.github/workflows/test-unit-enterprise-routing.yml index 13136c968d1..a64f00f4744 100644 --- a/.github/workflows/test-unit-enterprise-routing.yml +++ b/.github/workflows/test-unit-enterprise-routing.yml @@ -7,6 +7,10 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -14,8 +18,8 @@ permissions: pull-requests: write concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: enterprise-routing: diff --git a/.github/workflows/test-unit-integrations.yml b/.github/workflows/test-unit-integrations.yml index c95ed4e7c24..39752cf8e5d 100644 --- a/.github/workflows/test-unit-integrations.yml +++ b/.github/workflows/test-unit-integrations.yml @@ -7,6 +7,10 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -14,8 +18,8 @@ permissions: pull-requests: write concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: integrations: diff --git a/.github/workflows/test-unit-llm-providers.yml b/.github/workflows/test-unit-llm-providers.yml index df78564ab0c..4d1c921f723 100644 --- a/.github/workflows/test-unit-llm-providers.yml +++ b/.github/workflows/test-unit-llm-providers.yml @@ -7,13 +7,17 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: vertex-ai: diff --git a/.github/workflows/test-unit-misc.yml b/.github/workflows/test-unit-misc.yml index 9afaaaead93..505e22cfed4 100644 --- a/.github/workflows/test-unit-misc.yml +++ b/.github/workflows/test-unit-misc.yml @@ -7,6 +7,10 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -14,8 +18,8 @@ permissions: pull-requests: write concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: misc: diff --git a/.github/workflows/test-unit-proxy-auth.yml b/.github/workflows/test-unit-proxy-auth.yml index 97dfaed6e81..c27fe16d611 100644 --- a/.github/workflows/test-unit-proxy-auth.yml +++ b/.github/workflows/test-unit-proxy-auth.yml @@ -7,6 +7,10 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -14,8 +18,8 @@ permissions: pull-requests: write concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: proxy-auth: diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index b0ee56f5a5c..60d2e471862 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -7,13 +7,17 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} # Semantic matrix: each shard groups tests by concern (auth, server, logging, …) # rather than alphabetical letter ranges. Adding a new test file means adding it diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index b3eb8f79a43..6a51d2a8578 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -7,14 +7,18 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging workflow_dispatch: permissions: contents: read concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: proxy-endpoints: diff --git a/.github/workflows/test-unit-proxy-infra.yml b/.github/workflows/test-unit-proxy-infra.yml index 884d62289b9..913653a1711 100644 --- a/.github/workflows/test-unit-proxy-infra.yml +++ b/.github/workflows/test-unit-proxy-infra.yml @@ -7,6 +7,10 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -14,8 +18,8 @@ permissions: pull-requests: write concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: proxy-infra: diff --git a/.github/workflows/test-unit-proxy-legacy.yml b/.github/workflows/test-unit-proxy-legacy.yml index bcbf365babf..49aa5f9f51d 100644 --- a/.github/workflows/test-unit-proxy-legacy.yml +++ b/.github/workflows/test-unit-proxy-legacy.yml @@ -7,13 +7,17 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: test: diff --git a/.github/workflows/test-unit-responses-caching-types.yml b/.github/workflows/test-unit-responses-caching-types.yml index 2f177587997..5b336452069 100644 --- a/.github/workflows/test-unit-responses-caching-types.yml +++ b/.github/workflows/test-unit-responses-caching-types.yml @@ -7,6 +7,10 @@ on: - litellm_internal_staging - litellm_oss_staging - "litellm_**" + push: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -14,8 +18,8 @@ permissions: pull-requests: write concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} jobs: responses-caching-types: diff --git a/CLAUDE.md b/CLAUDE.md index c3c8138d2ac..209d9aaf326 100644 --- a/CLAUDE.md +++ b/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,11 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a Python max line length is 120, not 88 -On a fresh worktree or clone, run `make bootstrap` before anything else. It provisions everything tests, `make pre-commit`, and a local proxy need - -Run tests before you commit. Also, run `make pre-commit` right before each commit, which generates types (as needed) and formats/lints your code. Any errors found must be fixed. It only runs when there are staged frontend and/or backend changes and calculates violations, generates types, etc. based on the worktree, so stage what you need or stash/delete unwanted files in litellm/ or ui/ (where backend and frontend lint run, respectively) before running it. If it fails because dashboard api types are stale, it already regenerated them for you. You just need to stage the schema.d.ts, re-run `make pre-commit` to confirm it passes, and commit - When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in -If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason +If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # `. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing @@ -75,6 +71,7 @@ Follow these coding conventions for new/updated code (a three-line fix in a lega - Never-nester: early returns over deep nesting - Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never) - No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), etc. + - Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: ` 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 diff --git a/Dockerfile b/Dockerfile index a127cdabd59..1fb34f6ebf9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 diff --git a/Makefile b/Makefile index e9b2fb9d8f1..68753605e27 100644 --- a/Makefile +++ b/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 ci --no-audit --no-fund + cd ui/litellm-dashboard && ../../scripts/with_dashboard_node.sh npm install --no-audit --no-fund @main_root=$$(git worktree list --porcelain | head -1 | sed 's/^worktree //'); \ if [ "$$main_root" != "$$(git rev-parse --show-toplevel)" ] && [ -f "$$main_root/.env" ] && [ ! -f .env ]; then \ cp "$$main_root/.env" .env && echo "bootstrap: copied .env from $$main_root"; \ @@ -176,10 +176,8 @@ lint-ruff-FULL-dev: install-dev if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \ else echo "No changed .py files to check."; fi -lint-basedpyright lint-basedpyright-budget-update: export NODE_OPTIONS := --max-old-space-size=12288 - lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL) $(UV_RUN) basedpyright tests/e2e @@ -192,7 +190,7 @@ lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) # --update lowers each limit by what this branch fixed since its branch point, so # it needs the base ref fetched to resolve the merge-base. lint-basedpyright-budget-update: install-dev lint-fetch-base - ($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --update + $(UV_RUN) python scripts/type_check_gate.py --update lint-format: format-check @@ -239,7 +237,7 @@ lint-dev: lint-format-changed check-circular-imports check-import-safety # test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and # check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage. # Not auto-installed as a git hook so it never slows an unrelated human commit. -pre-commit: +pre-commit: bootstrap ./scripts/pre_commit_lint.sh # Testing targets diff --git a/backend/Dockerfile b/backend/Dockerfile index 62bd8b56483..9c93259adc3 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -63,6 +63,8 @@ RUN mkdir -p /home/nonroot && \ HOME=/home/nonroot prisma generate --schema=./schema.prisma && \ chown -R nonroot:nonroot /home/nonroot/.cache +RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh + # ---------- Runtime ---------- FROM $LITELLM_RUNTIME_IMAGE AS runtime @@ -93,5 +95,5 @@ USER nonroot EXPOSE 4001/tcp -ENTRYPOINT ["uvicorn", "backend.main:app"] +ENTRYPOINT ["/app/docker/component_entrypoint.sh", "uvicorn", "backend.main:app"] CMD ["--host", "0.0.0.0", "--port", "4001"] diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index a0efa19f320..96e224a7dc6 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -44,6 +44,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/router/", "/router_settings", "/adaptive_router/", + "/auto_router/", "/fallback", "/fallbacks", "/cache_settings", diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 65142091712..614a8e5d2c0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 31903 + "limit": 29809 }, "reportArgumentType": { "limit": 2645 @@ -15,25 +15,25 @@ "limit": 123 }, "reportConstantRedefinition": { - "limit": 59 + "limit": 40 }, "reportDeprecated": { "limit": 325 }, "reportDuplicateImport": { - "limit": 42 + "limit": 24 }, "reportExplicitAny": { - "limit": 10214 + "limit": 9473 }, "reportFunctionMemberAccess": { "limit": 11 }, "reportGeneralTypeIssues": { - "limit": 227 + "limit": 157 }, "reportIncompatibleMethodOverride": { - "limit": 78 + "limit": 77 }, "reportIncompatibleVariableOverride": { "limit": 12 @@ -42,7 +42,7 @@ "limit": 18 }, "reportIndexIssue": { - "limit": 37 + "limit": 35 }, "reportInvalidTypeForm": { "limit": 35 @@ -54,13 +54,13 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5869 + "limit": 5855 }, "reportMissingTypeArgument": { - "limit": 15861 + "limit": 15849 }, "reportMissingTypeStubs": { - "limit": 41 + "limit": 40 }, "reportOperatorIssue": { "limit": 0 @@ -81,13 +81,13 @@ "limit": 0 }, "reportPossiblyUnboundVariable": { - "limit": 77 + "limit": 56 }, "reportPrivateUsage": { - "limit": 2437 + "limit": 2436 }, "reportRedeclaration": { - "limit": 12 + "limit": 8 }, "reportReturnType": { "limit": 219 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45366 + "limit": 45262 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 40477 + "limit": 40452 }, "reportUnknownParameterType": { - "limit": 20338 + "limit": 20309 }, "reportUnknownVariableType": { - "limit": 32047 + "limit": 31978 }, "reportUnnecessaryCast": { - "limit": 177 + "limit": 124 }, "reportUnnecessaryComparison": { - "limit": 1021 + "limit": 703 }, "reportUnnecessaryContains": { - "limit": 7 + "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 1205 + "limit": 866 }, "reportUntypedBaseClass": { "limit": 165 @@ -132,15 +132,15 @@ "limit": 33 }, "reportUnusedClass": { - "limit": 33 + "limit": 23 }, "reportUnusedFunction": { - "limit": 204 + "limit": 139 }, "reportUnusedImport": { - "limit": 1003 + "limit": 588 }, "reportUnusedVariable": { - "limit": 1297 + "limit": 147 } } diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 9ee076ce825..c93a08409a2 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -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 diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 946b4de6f5e..1545a84d379 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -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 diff --git a/docker/build_admin_ui.sh b/docker/build_admin_ui.sh index 68acdd78e3e..f84aba7c053 100755 --- a/docker/build_admin_ui.sh +++ b/docker/build_admin_ui.sh @@ -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 diff --git a/docker/component_entrypoint.sh b/docker/component_entrypoint.sh new file mode 100755 index 00000000000..1748f1e13a7 --- /dev/null +++ b/docker/component_entrypoint.sh @@ -0,0 +1,8 @@ +#!/bin/sh + +if [ "$USE_DDTRACE" = "true" ]; then + export DD_TRACE_OPENAI_ENABLED="False" + exec ddtrace-run "$@" +fi + +exec "$@" diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index f209ab54f64..22f9f40ecd8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: from litellm.proxy._types import LiteLLM_ManagedObjectTable from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router + from litellm.types.router import Deployment from litellm.types.utils import LiteLLMBatch @@ -281,6 +282,32 @@ class CheckBatchCost: return deployment_id return None + @classmethod + def _get_managed_file_model_name( + cls, + job: "LiteLLM_ManagedObjectTable", + deployment_info: "Deployment", + ) -> Optional[str]: + """ + Public model group name to encode as ``target_model_names`` on unified output file ids. + + Key model-access checks resolve a managed file id back to a model via its + ``target_model_names``, so this must be the model group the caller requested, never the + underlying provider model (e.g. ``gpt-5.5``), which no key is allowed to call. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + convert_b64_uid_to_unified_uid, + get_models_from_unified_file_id, + ) + + input_file_id = cls._get_input_file_id(job) + target_model_names = ( + get_models_from_unified_file_id(convert_b64_uid_to_unified_uid(input_file_id)) if input_file_id else [] + ) + if target_model_names: + return ",".join(target_model_names) + return deployment_info.model_name or None + @staticmethod def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]: import json @@ -406,6 +433,10 @@ class CheckBatchCost: managed_files_hook = self.proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_hook is not None: from litellm.proxy._types import UserAPIKeyAuth + + managed_file_model_name = self._get_managed_file_model_name( + job=job, deployment_info=deployment_info + ) _minimal_auth = UserAPIKeyAuth( user_id=job.created_by or "default-user-id", team_id=getattr(job, "team_id", None), @@ -417,7 +448,7 @@ class CheckBatchCost: _unified_file_id = managed_files_hook.get_unified_output_file_id( output_file_id=_raw_file_id, model_id=model_id, - model_name=str(model_name) if model_name else deployment_info.model_name or None, + model_name=managed_file_model_name, ) await managed_files_hook.store_unified_file_id( file_id=_unified_file_id, diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index b8c97dd2bb5..ec47b6ac0e6 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -215,7 +215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) if result: - return LiteLLM_ManagedFileTable(**result) + return LiteLLM_ManagedFileTable.model_validate(result) ## CHECK DB db_object = await self.prisma_client.db.litellm_managedfiletable.find_first( @@ -223,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) if db_object: - return LiteLLM_ManagedFileTable(**db_object.model_dump()) + return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump()) return None async def delete_unified_file_id( @@ -349,7 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if isinstance(batch.file_object, str) else batch.file_object ) - batch_obj = LiteLLMBatch(**batch_data) + batch_obj = LiteLLMBatch.model_validate(batch_data) batch_obj.id = batch.unified_object_id batch_objects.append(batch_obj) @@ -383,7 +383,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) return [ - OpenAIFileObject(**file_object.file_object) + OpenAIFileObject.model_validate(file_object.file_object) for file_object in file_ids if file_object.file_object is not None ] diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index fa209e55eb8..5489eba1494 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.52" +version = "0.1.53" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.52" +version = "0.1.53" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 4b000912393..3b4f94d5dc9 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -65,6 +65,8 @@ RUN mkdir -p /home/nonroot && \ HOME=/home/nonroot prisma generate --schema=./schema.prisma && \ chown -R nonroot:nonroot /home/nonroot/.cache +RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component_entrypoint.sh + # ---------- Runtime ---------- FROM $LITELLM_RUNTIME_IMAGE AS runtime @@ -95,5 +97,5 @@ USER nonroot EXPOSE 4000/tcp -ENTRYPOINT ["sh", "-c", "exec uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"] +ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"] CMD ["--host", "0.0.0.0", "--port", "4000"] diff --git a/helm/litellm-helm/README.md b/helm/litellm-helm/README.md index 4e0884dd08c..4c8712ea7b9 100644 --- a/helm/litellm-helm/README.md +++ b/helm/litellm-helm/README.md @@ -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 | diff --git a/helm/litellm-helm/templates/migrations-job.yaml b/helm/litellm-helm/templates/migrations-job.yaml index 5ec7f5b7f3e..7bc1a133883 100644 --- a/helm/litellm-helm/templates/migrations-job.yaml +++ b/helm/litellm-helm/templates/migrations-job.yaml @@ -35,6 +35,8 @@ spec: {{- toYaml . | nindent 8 }} {{- end }} serviceAccountName: {{ include "litellm.migrationServiceAccountName" . }} + securityContext: + {{- toYaml .Values.podSecurityContext | nindent 8 }} {{- with .Values.migrationJob.extraInitContainers }} initContainers: {{- tpl (toYaml .) $ | nindent 8 }} diff --git a/helm/litellm-helm/tests/migrations-job_tests.yaml b/helm/litellm-helm/tests/migrations-job_tests.yaml index 05dd37b4857..6bfc1f38adc 100644 --- a/helm/litellm-helm/tests/migrations-job_tests.yaml +++ b/helm/litellm-helm/tests/migrations-job_tests.yaml @@ -254,3 +254,39 @@ tests: content: name: sidecar-tpl image: "ghcr.io/berriai/litellm-database:test" + - it: should render the pod-level securityContext from podSecurityContext + template: migrations-job.yaml + set: + migrationJob: + enabled: true + podSecurityContext: + fsGroup: 10000 + runAsUser: 10000 + runAsNonRoot: true + asserts: + - equal: + path: spec.template.spec.securityContext + value: + fsGroup: 10000 + runAsUser: 10000 + runAsNonRoot: true + - it: should keep the pod-level and container-level securityContext separate + template: migrations-job.yaml + set: + migrationJob: + enabled: true + podSecurityContext: + fsGroup: 10000 + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + asserts: + - equal: + path: spec.template.spec.securityContext + value: + fsGroup: 10000 + - equal: + path: spec.template.spec.containers[0].securityContext + value: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 7235bb0bd78..df2b55723fe 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -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: [] diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index a0205c0a3a2..bffd627393a 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -138,6 +138,59 @@ is false the chart uses the provided name, or the namespace `default` SA. {{- end -}} {{- end -}} +{{/* +ServiceAccount name for the migrations Job. + +The Job is a pre-install / pre-upgrade hook, so it is created before the +chart's ordinary resources. A ServiceAccount the chart creates is one of +those ordinary resources, which makes borrowing the backend name a cycle: +the hook pod is rejected because the account does not exist yet. So when +`serviceAccounts.backend.create` is true the Job falls back to the namespace +`default` account unless the operator names one that already exists. With +`create` false the backend name is either an operator-supplied existing +account or `default`, both of which are safe for the hook, so the Job keeps +sharing it. + +`migrationJob.serviceAccountName` always wins when set, which is how a Job +that needs credentials of its own (IRSA / Workload Identity for IAM database +auth) gets them. +*/}} +{{- define "litellm.migrations.serviceAccountName" -}} +{{- if .Values.migrationJob.serviceAccountName -}} +{{ .Values.migrationJob.serviceAccountName }} +{{- else if .Values.serviceAccounts.backend.create -}} +default +{{- else -}} +{{ include "litellm.backend.serviceAccountName" . }} +{{- end -}} +{{- end -}} + +{{/* +Extra pod labels for a component's Deployment, validated against its selector. + +Invoke with a dict: + (dict "podLabels" .Values.gateway.podLabels "componentName" "gateway") + +The three selector keys are also emitted on the pod template, so a podLabels +entry reusing one renders a duplicate YAML key whose later value wins. That +leaves the pod template no longer matching the (immutable) selector and the +apiserver rejects the Deployment. Fail at template time naming the key +instead, so the operator gets the reason here rather than an opaque +`selector does not match template labels` from the apiserver. + +The migrations Job takes podLabels unvalidated: a Job's selector is generated +by the controller rather than declared, so nothing there can collide. +*/}} +{{- define "litellm.podLabels" -}} +{{- $componentName := .componentName -}} +{{- range $key, $value := .podLabels }} +{{- if has $key (list "app.kubernetes.io/name" "app.kubernetes.io/instance" "app.kubernetes.io/component") }} +{{- fail (printf "%s.podLabels cannot set %s: it is part of the Deployment's immutable selector" $componentName $key) }} +{{- end }} +{{- end }} +{{- toYaml .podLabels }} +{{- end -}} + {{/* Master-key + database + redis env block — shared by gateway, backend, and the migrations Job. diff --git a/helm/litellm/templates/backend/deployment.yaml b/helm/litellm/templates/backend/deployment.yaml index 892b84ff7d5..c5d799a0faf 100644 --- a/helm/litellm/templates/backend/deployment.yaml +++ b/helm/litellm/templates/backend/deployment.yaml @@ -23,9 +23,16 @@ spec: {{- end }} labels: {{- include "litellm.backend.selectorLabels" . | nindent 8 }} + {{- with .Values.backend.podLabels }} + {{- include "litellm.podLabels" (dict "podLabels" . "componentName" "backend") | nindent 8 }} + {{- end }} spec: serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }} automountServiceAccountToken: {{ .Values.serviceAccounts.backend.automount }} + {{- with .Values.backend.podSecurityContext }} + securityContext: + {{- toYaml . | nindent 8 }} + {{- end }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} @@ -34,6 +41,10 @@ spec: - name: backend image: "{{ .Values.backend.image.repository }}:{{ .Values.backend.image.tag | default .Chart.AppVersion }}" imagePullPolicy: {{ .Values.backend.image.pullPolicy }} + {{- with .Values.backend.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} ports: - name: http containerPort: 4001 @@ -70,8 +81,15 @@ spec: readinessProbe: {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.backend.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} resources: {{- toYaml .Values.backend.resources | nindent 12 }} + {{- with .Values.backend.extraContainers }} + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} {{- if or .Values.gateway.config.create .Values.backend.volumes .Values.billingMetrics.enabled }} volumes: {{- if .Values.gateway.config.create }} @@ -102,4 +120,8 @@ spec: topologySpreadConstraints: {{- toYaml . | nindent 8 }} {{- end }} + {{- $gracePeriod := .Values.backend.terminationGracePeriodSeconds }} + {{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }} + terminationGracePeriodSeconds: {{ $gracePeriod }} + {{- end }} {{- end }} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index b2e22612905..7d16134a53d 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -21,9 +21,16 @@ spec: {{- end }} labels: {{- include "litellm.gateway.selectorLabels" . | nindent 8 }} + {{- with .Values.gateway.podLabels }} + {{- include "litellm.podLabels" (dict "podLabels" . "componentName" "gateway") | nindent 8 }} + {{- end }} spec: serviceAccountName: {{ include "litellm.gateway.serviceAccountName" . }} automountServiceAccountToken: {{ .Values.serviceAccounts.gateway.automount }} + {{- with .Values.gateway.podSecurityContext }} + securityContext: + {{- toYaml . | nindent 8 }} + {{- end }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} @@ -32,6 +39,10 @@ spec: - name: gateway image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}" imagePullPolicy: {{ .Values.gateway.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} ports: - name: http containerPort: 4000 @@ -72,8 +83,15 @@ spec: readinessProbe: {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.gateway.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} resources: {{- toYaml .Values.gateway.resources | nindent 12 }} + {{- with .Values.gateway.extraContainers }} + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }} volumes: {{- if .Values.gateway.config.create }} @@ -104,4 +122,8 @@ spec: topologySpreadConstraints: {{- toYaml . | nindent 8 }} {{- end }} + {{- $gracePeriod := .Values.gateway.terminationGracePeriodSeconds }} + {{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }} + terminationGracePeriodSeconds: {{ $gracePeriod }} + {{- end }} {{- end }} diff --git a/helm/litellm/templates/migrations-job.yaml b/helm/litellm/templates/migrations-job.yaml index 92671388546..2debe8a1e10 100644 --- a/helm/litellm/templates/migrations-job.yaml +++ b/helm/litellm/templates/migrations-job.yaml @@ -23,12 +23,21 @@ spec: ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }} template: metadata: + {{- /* The Job's selector is generated by the controller rather than + declared, so podLabels may override a chart label here. Merge + instead of appending so an override replaces the key rather than + rendering it twice. */}} + {{- $chartLabels := merge (dict "app.kubernetes.io/component" "migrations") (fromYaml (include "litellm.commonLabels" .)) }} labels: - {{- include "litellm.commonLabels" . | nindent 8 }} - app.kubernetes.io/component: migrations + {{- toYaml (merge (deepCopy .Values.migrationJob.podLabels) $chartLabels) | nindent 8 }} spec: restartPolicy: Never - serviceAccountName: {{ include "litellm.backend.serviceAccountName" . }} + serviceAccountName: {{ include "litellm.migrations.serviceAccountName" . }} + automountServiceAccountToken: {{ .Values.migrationJob.automountServiceAccountToken }} + {{- with .Values.migrationJob.podSecurityContext }} + securityContext: + {{- toYaml . | nindent 8 }} + {{- end }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} @@ -37,10 +46,22 @@ spec: - name: prisma-migrations image: "{{ .Values.migrationJob.image.repository }}:{{ .Values.migrationJob.image.tag | default .Chart.AppVersion }}" imagePullPolicy: {{ .Values.migrationJob.image.pullPolicy }} + {{- with .Values.migrationJob.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} env: {{- include "litellm.serverEnv" (dict "root" $ "component" .Values.migrationJob) | nindent 12 }} + {{- with .Values.migrationJob.volumeMounts }} + volumeMounts: + {{- toYaml . | nindent 12 }} + {{- end }} {{- with .Values.migrationJob.resources }} resources: {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.migrationJob.volumes }} + volumes: + {{- toYaml . | nindent 8 }} + {{- end }} {{- end }} diff --git a/helm/litellm/templates/ui/deployment.yaml b/helm/litellm/templates/ui/deployment.yaml index cd1f8c08fd4..b4129dbc8ac 100644 --- a/helm/litellm/templates/ui/deployment.yaml +++ b/helm/litellm/templates/ui/deployment.yaml @@ -18,9 +18,16 @@ spec: {{- end }} labels: {{- include "litellm.ui.selectorLabels" . | nindent 8 }} + {{- with .Values.ui.podLabels }} + {{- include "litellm.podLabels" (dict "podLabels" . "componentName" "ui") | nindent 8 }} + {{- end }} spec: serviceAccountName: {{ include "litellm.ui.serviceAccountName" . }} automountServiceAccountToken: {{ .Values.serviceAccounts.ui.automount }} + {{- with .Values.ui.podSecurityContext }} + securityContext: + {{- toYaml . | nindent 8 }} + {{- end }} {{- with .Values.imagePullSecrets }} imagePullSecrets: {{- toYaml . | nindent 8 }} @@ -29,6 +36,10 @@ spec: - name: ui image: "{{ .Values.ui.image.repository }}:{{ .Values.ui.image.tag | default .Chart.AppVersion }}" imagePullPolicy: {{ .Values.ui.image.pullPolicy }} + {{- with .Values.ui.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} ports: - name: http containerPort: 3000 @@ -58,8 +69,15 @@ spec: readinessProbe: {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.ui.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} resources: {{- toYaml .Values.ui.resources | nindent 12 }} + {{- with .Values.ui.extraContainers }} + {{- tpl (toYaml .) $ | nindent 8 }} + {{- end }} {{- with .Values.ui.volumes }} volumes: {{- toYaml . | nindent 8 }} @@ -80,4 +98,8 @@ spec: topologySpreadConstraints: {{- toYaml . | nindent 8 }} {{- end }} + {{- $gracePeriod := .Values.ui.terminationGracePeriodSeconds }} + {{- if not (or (kindIs "invalid" $gracePeriod) (eq (printf "%v" $gracePeriod) "")) }} + terminationGracePeriodSeconds: {{ $gracePeriod }} + {{- end }} {{- end }} diff --git a/helm/litellm/tests/migration_job_tests.yaml b/helm/litellm/tests/migration_job_tests.yaml new file mode 100644 index 00000000000..12e525c5a8c --- /dev/null +++ b/helm/litellm/tests/migration_job_tests.yaml @@ -0,0 +1,169 @@ +suite: test migrations Job ServiceAccount resolution and pod hardening +templates: + - migrations-job.yaml +values: + - ./values/required.yaml +tests: + - it: borrows the namespace default account when no ServiceAccount is configured + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: default + + - it: falls back to the namespace default account when the chart creates the backend ServiceAccount + set: + serviceAccounts.backend.create: true + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: default + - notEqual: + path: spec.template.spec.serviceAccountName + value: RELEASE-NAME-litellm-backend + + - it: keeps sharing an existing backend ServiceAccount the chart does not create + set: + serviceAccounts.backend.create: false + serviceAccounts.backend.name: existing-backend-sa + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: existing-backend-sa + + - it: prefers an explicit migration ServiceAccount over the created backend one + set: + serviceAccounts.backend.create: true + migrationJob.serviceAccountName: migrations-sa + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: migrations-sa + + - it: prefers an explicit migration ServiceAccount over an existing backend one + set: + serviceAccounts.backend.create: false + serviceAccounts.backend.name: existing-backend-sa + migrationJob.serviceAccountName: migrations-sa + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: migrations-sa + + - it: mounts no ServiceAccount token by default + asserts: + - equal: + path: spec.template.spec.automountServiceAccountToken + value: false + + - it: mounts a ServiceAccount token when the operator asks for one + set: + migrationJob.automountServiceAccountToken: true + asserts: + - equal: + path: spec.template.spec.automountServiceAccountToken + value: true + + - it: keeps the token off the Job when the backend disables automounting + set: + serviceAccounts.backend.create: true + serviceAccounts.backend.automount: false + asserts: + - equal: + path: spec.template.spec.serviceAccountName + value: default + - equal: + path: spec.template.spec.automountServiceAccountToken + value: false + + - it: renders no hardening fields by default + asserts: + - isNull: + path: spec.template.spec.securityContext + - isNull: + path: spec.template.spec.containers[0].securityContext + - isNull: + path: spec.template.spec.volumes + - isNull: + path: spec.template.spec.containers[0].volumeMounts + - equal: + path: spec.template.metadata.labels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/managed-by: Helm + helm.sh/chart: litellm-0.1.0 + app.kubernetes.io/component: migrations + + - it: renders pod-level and container-level securityContext in their own scopes + set: + migrationJob.podSecurityContext: + runAsNonRoot: true + runAsUser: 65532 + seccompProfile: + type: RuntimeDefault + migrationJob.securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: + - ALL + asserts: + - equal: + path: spec.template.spec.securityContext + value: + runAsNonRoot: true + runAsUser: 65532 + seccompProfile: + type: RuntimeDefault + - equal: + path: spec.template.spec.containers[0].securityContext + value: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: + - ALL + + - it: renders volumes on the pod and volumeMounts on the migration container + set: + migrationJob.volumes: + - name: tmp + emptyDir: + sizeLimit: 64Mi + migrationJob.volumeMounts: + - name: tmp + mountPath: /tmp + asserts: + - equal: + path: spec.template.spec.volumes + value: + - name: tmp + emptyDir: + sizeLimit: 64Mi + - equal: + path: spec.template.spec.containers[0].volumeMounts + value: + - name: tmp + mountPath: /tmp + + - it: merges podLabels with the chart labels on the Job pod + set: + migrationJob.podLabels: + egress-policy: restricted + asserts: + - equal: + path: spec.template.metadata.labels['egress-policy'] + value: restricted + - equal: + path: spec.template.metadata.labels['app.kubernetes.io/component'] + value: migrations + + - it: accepts a podLabel that reuses a chart label, since the Job selector is controller-generated + set: + migrationJob.podLabels: + app.kubernetes.io/component: batch-migrations + asserts: + - notFailedTemplate: {} + - equal: + path: spec.template.metadata.labels['app.kubernetes.io/component'] + value: batch-migrations diff --git a/helm/litellm/tests/pod_hardening_tests.yaml b/helm/litellm/tests/pod_hardening_tests.yaml new file mode 100644 index 00000000000..18e836670c0 --- /dev/null +++ b/helm/litellm/tests/pod_hardening_tests.yaml @@ -0,0 +1,298 @@ +suite: test pod hardening knobs on the component deployments +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml + - ui/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: gateway renders no hardening fields by default + template: gateway/deployment.yaml + asserts: + - isNull: + path: spec.template.spec.securityContext + - isNull: + path: spec.template.spec.containers[0].securityContext + - isNull: + path: spec.template.spec.containers[0].lifecycle + - isNull: + path: spec.template.spec.terminationGracePeriodSeconds + - lengthEqual: + path: spec.template.spec.containers + count: 1 + - equal: + path: spec.template.metadata.labels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: gateway + + - it: gateway renders pod-level and container-level securityContext in their own scopes + template: gateway/deployment.yaml + set: + gateway.podSecurityContext: + runAsNonRoot: true + runAsUser: 65532 + fsGroup: 65532 + seccompProfile: + type: RuntimeDefault + gateway.securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: + - ALL + asserts: + - equal: + path: spec.template.spec.securityContext + value: + runAsNonRoot: true + runAsUser: 65532 + fsGroup: 65532 + seccompProfile: + type: RuntimeDefault + - equal: + path: spec.template.spec.containers[0].securityContext + value: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: + - ALL + + - it: gateway merges podLabels with the selector labels + template: gateway/deployment.yaml + set: + gateway.podLabels: + egress-policy: restricted + team: platform + asserts: + - equal: + path: spec.template.metadata.labels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: gateway + egress-policy: restricted + team: platform + - equal: + path: spec.selector.matchLabels + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: gateway + + - it: gateway rejects a podLabel that collides with the selector + template: gateway/deployment.yaml + set: + gateway.podLabels: + app.kubernetes.io/component: not-gateway + asserts: + - failedTemplate: + errorMessage: "gateway.podLabels cannot set app.kubernetes.io/component: it is part of the Deployment's immutable selector" + + - it: backend rejects a podLabel that collides with the selector + template: backend/deployment.yaml + set: + backend.podLabels: + app.kubernetes.io/name: not-litellm + asserts: + - failedTemplate: + errorMessage: "backend.podLabels cannot set app.kubernetes.io/name: it is part of the Deployment's immutable selector" + + - it: ui rejects a podLabel that collides with the selector + template: ui/deployment.yaml + set: + ui.podLabels: + app.kubernetes.io/instance: not-the-release + asserts: + - failedTemplate: + errorMessage: "ui.podLabels cannot set app.kubernetes.io/instance: it is part of the Deployment's immutable selector" + + - it: gateway renders lifecycle hooks on the container + template: gateway/deployment.yaml + set: + gateway.lifecycle: + preStop: + httpGet: + path: /health/drain + port: 4000 + asserts: + - equal: + path: spec.template.spec.containers[0].lifecycle + value: + preStop: + httpGet: + path: /health/drain + port: 4000 + + - it: gateway renders terminationGracePeriodSeconds on the pod spec + template: gateway/deployment.yaml + set: + gateway.terminationGracePeriodSeconds: 90 + asserts: + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 90 + + - it: gateway honors an explicit terminationGracePeriodSeconds of zero + template: gateway/deployment.yaml + set: + gateway.terminationGracePeriodSeconds: 0 + asserts: + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 0 + + - it: gateway appends extraContainers after the gateway container + template: gateway/deployment.yaml + set: + gateway.extraContainers: + - name: auth-sidecar + image: registry.example.com/auth-proxy:1.2.3 + args: + - --upstream + - http://127.0.0.1:4000 + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 2 + - equal: + path: spec.template.spec.containers[0].name + value: gateway + - equal: + path: spec.template.spec.containers[1] + value: + name: auth-sidecar + image: registry.example.com/auth-proxy:1.2.3 + args: + - --upstream + - http://127.0.0.1:4000 + + - it: gateway templates chart context inside extraContainers + template: gateway/deployment.yaml + set: + gateway.extraContainers: + - name: auth-sidecar + image: registry.example.com/auth-proxy:1.2.3 + env: + - name: RELEASE + value: "{{ .Release.Name }}" + asserts: + - equal: + path: spec.template.spec.containers[1].env[0].value + value: RELEASE-NAME + + - it: backend renders every hardening knob in the right scope + template: backend/deployment.yaml + set: + backend.podLabels: + egress-policy: restricted + backend.podSecurityContext: + runAsNonRoot: true + backend.securityContext: + readOnlyRootFilesystem: true + backend.lifecycle: + preStop: + exec: + command: + - sleep + - "5" + backend.terminationGracePeriodSeconds: 60 + backend.extraContainers: + - name: auth-sidecar + image: registry.example.com/auth-proxy:1.2.3 + asserts: + - equal: + path: spec.template.metadata.labels['egress-policy'] + value: restricted + - equal: + path: spec.template.spec.securityContext + value: + runAsNonRoot: true + - equal: + path: spec.template.spec.containers[0].securityContext + value: + readOnlyRootFilesystem: true + - equal: + path: spec.template.spec.containers[0].lifecycle + value: + preStop: + exec: + command: + - sleep + - "5" + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 60 + - equal: + path: spec.template.spec.containers[1].name + value: auth-sidecar + + - it: ui renders every hardening knob in the right scope + template: ui/deployment.yaml + set: + ui.podLabels: + egress-policy: restricted + ui.podSecurityContext: + runAsNonRoot: true + fsGroup: 101 + ui.securityContext: + readOnlyRootFilesystem: true + ui.lifecycle: + preStop: + exec: + command: + - /bin/sh + - -c + - nginx -s quit + ui.terminationGracePeriodSeconds: 30 + ui.extraContainers: + - name: auth-sidecar + image: registry.example.com/auth-proxy:1.2.3 + asserts: + - equal: + path: spec.template.metadata.labels['egress-policy'] + value: restricted + - equal: + path: spec.template.spec.securityContext + value: + runAsNonRoot: true + fsGroup: 101 + - equal: + path: spec.template.spec.containers[0].securityContext + value: + readOnlyRootFilesystem: true + - equal: + path: spec.template.spec.containers[0].lifecycle + value: + preStop: + exec: + command: + - /bin/sh + - -c + - nginx -s quit + - equal: + path: spec.template.spec.terminationGracePeriodSeconds + value: 30 + - equal: + path: spec.template.spec.containers[1].name + value: auth-sidecar + + - it: backend and ui render no hardening fields by default + templates: + - backend/deployment.yaml + - ui/deployment.yaml + asserts: + - isNull: + path: spec.template.spec.securityContext + - isNull: + path: spec.template.spec.containers[0].securityContext + - isNull: + path: spec.template.spec.containers[0].lifecycle + - isNull: + path: spec.template.spec.terminationGracePeriodSeconds + - lengthEqual: + path: spec.template.spec.containers + count: 1 diff --git a/helm/litellm/tests/probe_tests.yaml b/helm/litellm/tests/probe_tests.yaml new file mode 100644 index 00000000000..a04709db2f5 --- /dev/null +++ b/helm/litellm/tests/probe_tests.yaml @@ -0,0 +1,106 @@ +suite: test liveness and readiness probe timeouts +templates: + - gateway/deployment.yaml + - gateway/configmap.yaml + - backend/deployment.yaml +values: + - ./values/required.yaml +tests: + - it: gateway probes set an explicit timeout that outlasts a saturated event loop + template: gateway/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].livenessProbe + value: + httpGet: + path: /health/liveliness + port: http + initialDelaySeconds: 10 + periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 6 + - equal: + path: spec.template.spec.containers[0].readinessProbe + value: + httpGet: + path: /health/readiness + port: http + initialDelaySeconds: 5 + periodSeconds: 10 + timeoutSeconds: 10 + + - it: backend probes set an explicit timeout that outlasts a saturated event loop + template: backend/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].livenessProbe + value: + httpGet: + path: /health/liveliness + port: http + initialDelaySeconds: 10 + periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 6 + - equal: + path: spec.template.spec.containers[0].readinessProbe + value: + httpGet: + path: /health/readiness + port: http + initialDelaySeconds: 5 + periodSeconds: 10 + timeoutSeconds: 10 + + - it: no single-event-loop component is left on the kubernetes default 1s probe timeout + templates: + - gateway/deployment.yaml + - backend/deployment.yaml + asserts: + - isNotNullOrEmpty: + path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds + - isNotNullOrEmpty: + path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds + - equal: + path: spec.template.spec.containers[0].livenessProbe.timeoutSeconds + value: 10 + - equal: + path: spec.template.spec.containers[0].readinessProbe.timeoutSeconds + value: 10 + + - it: gateway liveness tolerates a longer outage than readiness before acting + template: gateway/deployment.yaml + asserts: + - equal: + path: spec.template.spec.containers[0].livenessProbe.failureThreshold + value: 6 + - notExists: + path: spec.template.spec.containers[0].readinessProbe.failureThreshold + + - it: probe timeouts and thresholds stay overridable per component + template: gateway/deployment.yaml + set: + gateway.readinessProbe.timeoutSeconds: 3 + gateway.readinessProbe.periodSeconds: 20 + gateway.livenessProbe.timeoutSeconds: 4 + gateway.livenessProbe.failureThreshold: 3 + asserts: + - equal: + path: spec.template.spec.containers[0].readinessProbe + value: + httpGet: + path: /health/readiness + port: http + initialDelaySeconds: 5 + periodSeconds: 20 + timeoutSeconds: 3 + - equal: + path: spec.template.spec.containers[0].livenessProbe + value: + httpGet: + path: /health/liveliness + port: http + initialDelaySeconds: 10 + periodSeconds: 15 + timeoutSeconds: 4 + failureThreshold: 3 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 461935b2f50..cd377667602 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -57,6 +57,42 @@ migrationJob: backoffLimit: 4 ttlSecondsAfterFinished: 120 resources: {} + # ServiceAccount for the Job pod only. + # + # The Job is a pre-install / pre-upgrade hook, so it runs before the chart's + # ordinary resources exist. With `serviceAccounts.backend.create: true` the + # backend ServiceAccount is one of those ordinary resources, so a Job that + # borrowed its name would reference an account that does not exist yet and + # the first install would fail with a forbidden pod creation. The name set + # here always wins; when it is empty the Job falls back to `default` if the + # chart creates the backend ServiceAccount, and to the backend + # ServiceAccount name otherwise (that name is either an existing account you + # supplied or `default`). + # + # Point this at a pre-existing ServiceAccount when the Job needs credentials + # of its own, e.g. the IRSA / Workload Identity annotations that + # `database.writer.useIAMAuth` relies on. That is also the upgrade path to + # watch: a release already running with `serviceAccounts.backend.create: + # true` used to hand the Job the created backend account on every upgrade, + # and now hands it `default` unless you name an account here. + serviceAccountName: "" + # The Job runs `prisma migrate deploy` against Postgres and never calls the + # K8s API, so it defaults to no projected ServiceAccount token, the same + # reasoning the ui SA above uses. Flip to true if your Job genuinely needs + # one; IAM database auth does not, since EKS Pod Identity injects its own + # projected token volume and GKE Workload Identity goes through the + # metadata server, neither of which is the default token mount. + automountServiceAccountToken: false + # Standard k8s pod-level and container-level securityContext for the Job + # pod. Same shape as gateway.podSecurityContext / gateway.securityContext. + podSecurityContext: {} + securityContext: {} + # Extra pod labels on the Job pod, merged into the chart's common labels. + podLabels: {} + # Additional volumes on the Job pod and volumeMounts on its container, e.g. + # the writable scratch space a read-only root filesystem needs. + volumes: [] + volumeMounts: [] image: repository: ghcr.io/berriai/litellm-migrations tag: "" # defaults to .Chart.AppVersion @@ -180,10 +216,13 @@ gateway: httpGet: { path: /health/liveliness, port: http } initialDelaySeconds: 10 periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 6 readinessProbe: httpGet: { path: /health/readiness, port: http } initialDelaySeconds: 5 periodSeconds: 10 + timeoutSeconds: 10 hpa: enabled: true minReplicas: 1 @@ -200,6 +239,37 @@ gateway: minAvailable: "" maxUnavailable: "" podAnnotations: {} + # Extra pod labels, merged into the chart's selector labels. Do not + # re-declare `app.kubernetes.io/name` / `instance` / `component` here: they + # form the Deployment's immutable selector. + podLabels: {} + # Pod-level securityContext, applied to every container in the pod + # (runAsNonRoot, runAsUser, fsGroup, seccompProfile, ...). Empty by default + # so the cluster's own defaults keep applying to existing installs; clusters + # enforcing a restricted Pod Security Standard usually want at least + # `runAsNonRoot: true` and `seccompProfile.type: RuntimeDefault`. + podSecurityContext: {} + # Container-level securityContext for the gateway container. Empty by + # default for the same reason. Example: + # allowPrivilegeEscalation: false + # readOnlyRootFilesystem: true + # capabilities: + # drop: + # - ALL + # `readOnlyRootFilesystem: true` needs writable scratch space; supply it + # through `volumes` / `volumeMounts` above rather than expecting the chart + # to guess the paths your workload writes to. + securityContext: {} + # Extra sidecar containers appended to the gateway pod, e.g. an auth or + # egress proxy. Rendered through `tpl`, so entries may reference chart + # values and release metadata. + extraContainers: [] + # Container lifecycle hooks (postStart / preStop) for the gateway container. + lifecycle: {} + # Grace period the kubelet allows between SIGTERM and SIGKILL. Leave empty + # to inherit the Kubernetes default of 30s. Set it a few seconds above the + # proxy's GRACEFUL_SHUTDOWN_TIMEOUT when you use a draining preStop hook. + terminationGracePeriodSeconds: "" nodeSelector: {} tolerations: [] affinity: {} @@ -242,10 +312,13 @@ backend: httpGet: { path: /health/liveliness, port: http } initialDelaySeconds: 10 periodSeconds: 15 + timeoutSeconds: 10 + failureThreshold: 6 readinessProbe: httpGet: { path: /health/readiness, port: http } initialDelaySeconds: 5 periodSeconds: 10 + timeoutSeconds: 10 hpa: enabled: true minReplicas: 1 @@ -257,6 +330,13 @@ backend: minAvailable: "" maxUnavailable: "" podAnnotations: {} + # Same shape as the gateway blocks of the same name. + podLabels: {} + podSecurityContext: {} + securityContext: {} + extraContainers: [] + lifecycle: {} + terminationGracePeriodSeconds: "" nodeSelector: {} tolerations: [] affinity: {} @@ -310,6 +390,16 @@ ui: minAvailable: "" maxUnavailable: "" podAnnotations: {} + # Same shape as the gateway blocks of the same name. The nginx runtime + # writes its pid, cache, and proxy temp files under the image's root + # filesystem, so `securityContext.readOnlyRootFilesystem: true` here needs + # emptyDir volumes mounted over those paths. + podLabels: {} + podSecurityContext: {} + securityContext: {} + extraContainers: [] + lifecycle: {} + terminationGracePeriodSeconds: "" nodeSelector: {} tolerations: [] affinity: {} diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql new file mode 100644 index 00000000000..2a8460a9b60 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260729000000_add_reload_tracking_to_litellm_config/migration.sql @@ -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; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260731000000_add_autorouter_savings_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260731000000_add_autorouter_savings_spend/migration.sql new file mode 100644 index 00000000000..a7efd444fdb --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260731000000_add_autorouter_savings_spend/migration.sql @@ -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; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 37ea55f8c13..17339541fd9 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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) diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 79984dcab68..beddd899472 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -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==", diff --git a/litellm/__init__.py b/litellm/__init__.py index 3f8c742c5a2..c2dbc9687f9 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -27,18 +27,19 @@ if os.getenv("LITELLM_MODE", "DEV") == "DEV": _dotenv.load_dotenv(override=_dev_env_hot_reload_enabled()) from typing import ( - Callable, - List, - Optional, - Dict, - Union, Any, - Literal, + Callable, + Dict, + Final, get_args, - TYPE_CHECKING, - Tuple, + List, + Literal, + Optional, overload, + Tuple, Type, + TYPE_CHECKING, + Union, ) from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams @@ -264,6 +265,7 @@ databricks_key: Optional[str] = None openai_like_key: Optional[str] = None azure_key: Optional[str] = None anthropic_key: Optional[str] = None +autorouter_savings_baseline_model: Optional[str] = None replicate_key: Optional[str] = None bytez_key: Optional[str] = None gdc_key: Optional[str] = None @@ -449,6 +451,8 @@ enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None custom_prometheus_metadata_labels: List[str] = [] custom_prometheus_tags: List[str] = [] prometheus_metrics_config: Optional[List] = None +prometheus_exclude_metrics: Optional[List[str]] = None +prometheus_exclude_labels: Optional[List[str]] = None prometheus_emit_stream_label: bool = False # Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on # `litellm_proxy_failed_requests_metric`. Off by default to preserve the @@ -678,12 +682,12 @@ def is_bedrock_pricing_only_model(key: str) -> bool: bool: True if the key matches the Bedrock pattern, False otherwise. """ # Regex to match 'bedrock//' - bedrock_pattern = re.compile(r"^bedrock/[a-zA-Z0-9_-]+/.+$") + bedrock_pattern: Final = re.compile(r"^bedrock/[a-zA-Z0-9_-]+/.+$") if "month-commitment" in key: return True - is_match = bedrock_pattern.match(key) + is_match: Final = bedrock_pattern.match(key) return is_match is not None @@ -701,7 +705,7 @@ def is_openai_finetune_model(key: str) -> bool: def add_known_models(model_cost_map: Optional[Dict] = None): - _map = model_cost_map if model_cost_map is not None else model_cost + _map: Final = model_cost_map if model_cost_map is not None else model_cost for key, value in _map.items(): if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): open_ai_chat_completion_models.add(key) @@ -2137,11 +2141,11 @@ def __getattr__(name: str) -> Any: # Use cached registry from _lazy_imports instead of importing tuples every time from ._lazy_imports import _get_lazy_import_registry - registry = _get_lazy_import_registry() + registry: Final = _get_lazy_import_registry() # Check if name is in registry and call the cached handler function if name in registry: - handler_func = registry[name] + handler_func: Final = registry[name] return handler_func(name) # Lazy load encoding from main.py to avoid heavy tiktoken import @@ -2194,7 +2198,7 @@ def __getattr__(name: str) -> Any: return _globals["openaiOSeriesConfig"] # Lazy load other config instances - _config_instances = { + _config_instances: Final = { "openAIGPTConfig": "OpenAIGPTConfig", "openAIGPTAudioConfig": "OpenAIGPTAudioConfig", "openAIGPT5Config": "OpenAIGPT5Config", @@ -2236,7 +2240,7 @@ def __getattr__(name: str) -> Any: # Check if already cached if "priority_reservation_settings" not in _globals: # Import the class and instantiate it - PriorityReservationSettings = __getattr__("PriorityReservationSettings") + PriorityReservationSettings: Final = __getattr__("PriorityReservationSettings") _globals["priority_reservation_settings"] = PriorityReservationSettings() return _globals["priority_reservation_settings"] @@ -2248,7 +2252,7 @@ def __getattr__(name: str) -> Any: # Check if already cached if "logging_callback_manager" not in _globals: # Import the class and instantiate it - LoggingCallbackManager = __getattr__("LoggingCallbackManager") + LoggingCallbackManager: Final = __getattr__("LoggingCallbackManager") _globals["logging_callback_manager"] = LoggingCallbackManager() return _globals["logging_callback_manager"] diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 727ca87d3e8..f856fe0f2b3 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -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) diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index b04fae86e47..63142ee4f2f 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -17,39 +17,40 @@ until they're actually needed. import importlib import sys -from typing import Any, Optional, cast, Callable +from collections.abc import Callable +from typing import Any, Final, cast # Import all the data structures that define what can be lazy-loaded # These are just lists of names and maps of where to find them from ._lazy_imports_registry import ( - # Name tuples - COST_CALCULATOR_NAMES, - LITELLM_LOGGING_NAMES, - UTILS_NAMES, - TOKEN_COUNTER_NAMES, - LLM_CLIENT_CACHE_NAMES, - BEDROCK_TYPES_NAMES, - TYPES_UTILS_NAMES, - CACHING_NAMES, - HTTP_HANDLER_NAMES, - DOTPROMPT_NAMES, - LLM_CONFIG_NAMES, - TYPES_NAMES, - LLM_PROVIDER_LOGIC_NAMES, - UTILS_MODULE_NAMES, # Import maps - _UTILS_IMPORT_MAP, - _COST_CALCULATOR_IMPORT_MAP, - _TYPES_UTILS_IMPORT_MAP, - _TOKEN_COUNTER_IMPORT_MAP, _BEDROCK_TYPES_IMPORT_MAP, _CACHING_IMPORT_MAP, - _LITELLM_LOGGING_IMPORT_MAP, + _COST_CALCULATOR_IMPORT_MAP, _DOTPROMPT_IMPORT_MAP, - _TYPES_IMPORT_MAP, + _LITELLM_LOGGING_IMPORT_MAP, _LLM_CONFIGS_IMPORT_MAP, _LLM_PROVIDER_LOGIC_IMPORT_MAP, + _TOKEN_COUNTER_IMPORT_MAP, + _TYPES_IMPORT_MAP, + _TYPES_UTILS_IMPORT_MAP, + _UTILS_IMPORT_MAP, _UTILS_MODULE_IMPORT_MAP, + # Name tuples + BEDROCK_TYPES_NAMES, + CACHING_NAMES, + COST_CALCULATOR_NAMES, + DOTPROMPT_NAMES, + HTTP_HANDLER_NAMES, + LITELLM_LOGGING_NAMES, + LLM_CLIENT_CACHE_NAMES, + LLM_CONFIG_NAMES, + LLM_PROVIDER_LOGIC_NAMES, + TOKEN_COUNTER_NAMES, + TYPES_NAMES, + TYPES_UTILS_NAMES, + UTILS_MODULE_NAMES, + UTILS_NAMES, ) @@ -77,7 +78,7 @@ def _get_utils_globals() -> dict: # They're separate from the main lazy import system because they have specific use cases # Lazy loader for default encoding - avoids importing heavy tiktoken library at startup -_default_encoding: Optional[Any] = None +_default_encoding: Any | None = None def _get_default_encoding() -> Any: @@ -99,7 +100,7 @@ def _get_default_encoding() -> Any: # Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time -_get_modified_max_tokens_func: Optional[Any] = None +_get_modified_max_tokens_func: Any | None = None def _get_modified_max_tokens() -> Any: @@ -123,7 +124,7 @@ def _get_modified_max_tokens() -> Any: # Lazy loader for token_counter to avoid importing token_counter module at module import time -_token_counter_new_func: Optional[Any] = None +_token_counter_new_func: Any | None = None def _get_token_counter_new() -> Any: @@ -153,7 +154,7 @@ def _get_token_counter_new() -> Any: # This registry maps attribute names (like "ModelResponse") to handler functions # It's built once the first time someone accesses a lazy-loaded attribute # Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...} -_LAZY_IMPORT_REGISTRY: Optional[dict[str, Callable[[str], Any]]] = None +_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: @@ -232,7 +233,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate raise AttributeError(f"{category} lazy import: unknown attribute {name!r}") # Step 2: Get the cache (where we store imported things) - _globals = _get_litellm_globals() + _globals: Final = _get_litellm_globals() # Step 3: If we've already imported it, just return the cached version if name in _globals: @@ -254,7 +255,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # Step 6: Get the actual attribute from the module # Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class - value = getattr(module, attr_name) + value: Final = getattr(module, attr_name) # Step 7: Cache it so we don't have to import again next time _globals[name] = value @@ -338,7 +339,7 @@ def _lazy_import_utils_module(name: str) -> Any: raise AttributeError(f"Utils module lazy import: unknown attribute {name!r}") # Get the cache (where we store imported things) - use utils globals - _globals = _get_utils_globals() + _globals: Final = _get_utils_globals() # If we've already imported it, just return the cached version if name in _globals: @@ -354,7 +355,7 @@ def _lazy_import_utils_module(name: str) -> Any: module = importlib.import_module(module_path) # Get the actual attribute from the module - value = getattr(module, attr_name) + value: Final = getattr(module, attr_name) # Cache it so we don't have to import again next time _globals[name] = value @@ -378,15 +379,15 @@ def _lazy_import_llm_client_cache(name: str) -> Any: - "in_memory_llm_clients_cache" is a singleton instance of that class So we need custom logic to handle both cases. """ - _globals = _get_litellm_globals() + _globals: Final = _get_litellm_globals() # If already cached, return it if name in _globals: return _globals[name] # Import the class - module = importlib.import_module("litellm.caching.llm_caching_handler") - LLMClientCache = getattr(module, "LLMClientCache") + module: Final = importlib.import_module("litellm.caching.llm_caching_handler") + LLMClientCache: Final = getattr(module, "LLMClientCache") # If they want the class itself, return it if name == "LLMClientCache": @@ -395,7 +396,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: # If they want the singleton instance, create it (only once) if name == "in_memory_llm_clients_cache": - instance = LLMClientCache() + instance: Final = LLMClientCache() _globals["in_memory_llm_clients_cache"] = instance return instance @@ -411,7 +412,7 @@ def _lazy_import_http_handlers(name: str) -> Any: - They need configuration (timeout, etc.) from the module globals - They use factory functions instead of direct instantiation """ - _globals = _get_litellm_globals() + _globals: Final = _get_litellm_globals() if name == "module_level_aclient": # Create an async HTTP client using the factory function @@ -419,11 +420,11 @@ def _lazy_import_http_handlers(name: str) -> Any: # Get timeout from module config (if set) timeout = _globals.get("request_timeout") - params = {"timeout": timeout, "client_alias": "module level aclient"} + params: Final = {"timeout": timeout, "client_alias": "module level aclient"} # Create the client instance - provider_id = cast(Any, "litellm_module_level_client") - async_client = get_async_httpx_client( + provider_id: Final = cast(Any, "litellm_module_level_client") + async_client: Final = get_async_httpx_client( llm_provider=provider_id, params=params, ) @@ -437,7 +438,7 @@ def _lazy_import_http_handlers(name: str) -> Any: from litellm.llms.custom_httpx.http_handler import HTTPHandler timeout = _globals.get("request_timeout") - sync_client = HTTPHandler(timeout=timeout) + sync_client: Final = HTTPHandler(timeout=timeout) # Cache it _globals["module_level_client"] = sync_client diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 488331e3895..37f111c2324 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -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", diff --git a/litellm/_logging.py b/litellm/_logging.py index 5f3c483869d..c5a9fd0f8c7 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -4,11 +4,11 @@ import os import sys from datetime import datetime from logging import Formatter -from typing import Any, Dict, Optional +from typing import Any, Final -from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads +from litellm.litellm_core_utils.secret_redaction import redact_string set_verbose = False @@ -17,7 +17,7 @@ if set_verbose is True: "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs." ) -_ENABLE_SECRET_REDACTION = os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true" +_ENABLE_SECRET_REDACTION: Final = os.getenv("LITELLM_DISABLE_REDACT_SECRETS", "").lower() != "true" def _redact_string(value: str) -> str: @@ -74,19 +74,19 @@ class SecretRedactionFilter(logging.Filter): return True -_secret_filter = SecretRedactionFilter() +_secret_filter: Final = SecretRedactionFilter() json_logs = bool(os.getenv("JSON_LOGS", False)) # Create a handler for the logger (you may need to adapt this based on your needs) -log_level = os.getenv("LITELLM_LOG", "DEBUG") -numeric_level: str = getattr(logging, log_level.upper()) -handler = logging.StreamHandler() +log_level: Final = os.getenv("LITELLM_LOG", "DEBUG") +numeric_level: Final[str] = getattr(logging, log_level.upper()) +handler: Final = logging.StreamHandler() handler.setLevel(numeric_level) handler.addFilter(_secret_filter) -def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]: +def _try_parse_json_message(message: str) -> dict[str, Any] | None: """ Try to parse a log message as JSON. Returns parsed dict if valid, else None. Handles messages that are entirely valid JSON (e.g. json.dumps output). @@ -94,16 +94,16 @@ def _try_parse_json_message(message: str) -> Optional[Dict[str, Any]]: """ if not message or not isinstance(message, str): return None - msg_stripped = message.strip() + msg_stripped: Final = message.strip() if not (msg_stripped.startswith("{") or msg_stripped.startswith("[")): return None - parsed = safe_json_loads(message, default=None) + parsed: Final = safe_json_loads(message, default=None) if parsed is None or not isinstance(parsed, dict): return None return parsed -def _try_parse_embedded_python_dict(message: str) -> Optional[Dict[str, Any]]: +def _try_parse_embedded_python_dict(message: str) -> dict[str, Any] | None: """ Try to find and parse a Python dict repr (e.g. str(d) or repr(d)) embedded in the message. Handles patterns like: @@ -144,21 +144,21 @@ def _get_standard_record_attrs() -> frozenset: return frozenset(logging.LogRecord("", 0, "", 0, "", (), None).__dict__.keys()) -_STANDARD_RECORD_ATTRS = _get_standard_record_attrs() +_STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs() class JsonFormatter(Formatter): def __init__(self): - super(JsonFormatter, self).__init__() + super().__init__() def formatTime(self, record, datefmt=None): # Use datetime to format the timestamp in ISO 8601 format - dt = datetime.fromtimestamp(record.created) + dt: Final = datetime.fromtimestamp(record.created) return dt.isoformat() def format(self, record): - message_str = record.getMessage() - json_record: Dict[str, Any] = { + message_str: Final = record.getMessage() + json_record: Final[dict[str, Any]] = { "message": message_str, "level": record.levelname, "timestamp": self.formatTime(record), @@ -193,13 +193,13 @@ class JsonFormatter(Formatter): # Function to set up exception handlers for JSON logging def _setup_json_exception_handlers(formatter): # Create a handler with JSON formatting for exceptions - error_handler = logging.StreamHandler() + error_handler: Final = logging.StreamHandler() error_handler.setFormatter(formatter) error_handler.addFilter(_secret_filter) # Setup excepthook for uncaught exceptions def json_excepthook(exc_type, exc_value, exc_traceback): - record = logging.LogRecord( + record: Final = logging.LogRecord( name="LiteLLM", level=logging.ERROR, pathname="", @@ -217,10 +217,10 @@ def _setup_json_exception_handlers(formatter): import asyncio def async_json_exception_handler(loop, context): - exception = context.get("exception") + exception: Final = context.get("exception") if exception: - exc_type = type(exception) - record = logging.LogRecord( + exc_type: Final = type(exception) + record: Final = logging.LogRecord( name="LiteLLM", level=logging.ERROR, pathname="", @@ -243,7 +243,7 @@ if json_logs: handler.setFormatter(JsonFormatter()) _setup_json_exception_handlers(JsonFormatter()) else: - formatter = logging.Formatter( + formatter: Final = logging.Formatter( "\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s", datefmt="%H:%M:%S", ) @@ -263,20 +263,20 @@ verbose_logger.addHandler(handler) def _suppress_loggers(): """Suppress noisy loggers at INFO level""" # Suppress httpx request logging at INFO level - httpx_logger = logging.getLogger("httpx") + httpx_logger: Final = logging.getLogger("httpx") httpx_logger.setLevel(logging.WARNING) # Suppress APScheduler logging at INFO level - apscheduler_executors_logger = logging.getLogger("apscheduler.executors.default") + apscheduler_executors_logger: Final = logging.getLogger("apscheduler.executors.default") apscheduler_executors_logger.setLevel(logging.WARNING) - apscheduler_scheduler_logger = logging.getLogger("apscheduler.scheduler") + apscheduler_scheduler_logger: Final = logging.getLogger("apscheduler.scheduler") apscheduler_scheduler_logger.setLevel(logging.WARNING) # Call the suppression function _suppress_loggers() -ALL_LOGGERS = [ +ALL_LOGGERS: Final = [ logging.getLogger(), verbose_logger, verbose_router_logger, @@ -293,11 +293,11 @@ def _get_loggers_to_initialize(): """ import litellm - loggers = list(ALL_LOGGERS) + loggers: Final = list(ALL_LOGGERS) # Add langfuse logger if langfuse is being used as a callback - langfuse_callbacks = {"langfuse", "langfuse_otel"} - all_callbacks = set(litellm.success_callback + litellm.failure_callback) + langfuse_callbacks: Final = {"langfuse", "langfuse_otel"} + all_callbacks: Final = set(litellm.success_callback + litellm.failure_callback) if langfuse_callbacks & all_callbacks: loggers.append(logging.getLogger("langfuse")) @@ -325,12 +325,12 @@ def _get_uvicorn_json_log_config(): This ensures that uvicorn's access logs, error logs, and all application logs are formatted as JSON when json_logs is enabled. """ - json_formatter_class = "litellm._logging.JsonFormatter" + json_formatter_class: Final = "litellm._logging.JsonFormatter" # Use the module-level log_level variable for consistency - uvicorn_log_level = log_level.upper() + uvicorn_log_level: Final = log_level.upper() - log_config = { + log_config: Final = { "version": 1, "disable_existing_loggers": False, "formatters": { @@ -384,7 +384,7 @@ def _turn_on_json(): - Adds a JSON formatter to all loggers """ - handler = logging.StreamHandler() + handler: Final = logging.StreamHandler() handler.setFormatter(JsonFormatter()) _initialize_loggers_with_handler(handler) # Set up exception handlers diff --git a/litellm/_redis.py b/litellm/_redis.py index 9e3b247f577..ed014a83c25 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -12,7 +12,8 @@ import json # s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation import os -from typing import Callable, List, Optional, Union +from collections.abc import Callable +from typing import Final import redis # type: ignore import redis.asyncio as async_redis # type: ignore @@ -32,20 +33,20 @@ from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from ._logging import verbose_logger -AZURE_REDIS_SCOPE = "https://redis.azure.com/.default" +AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" def _get_redis_kwargs(): - arg_spec = inspect.getfullargspec(redis.Redis) + arg_spec: Final = inspect.getfullargspec(redis.Redis) # Only allow primitive arguments - exclude_args = { + exclude_args: Final = { "self", "connection_pool", "retry", } - include_args = { + include_args: Final = { "url", "redis_connect_func", "gcp_service_account", @@ -56,7 +57,7 @@ def _get_redis_kwargs(): "azure_client_secret", } - available_args = {x for x in arg_spec.args if x not in exclude_args} | include_args + available_args: Final = {x for x in arg_spec.args if x not in exclude_args} | include_args return available_args @@ -76,7 +77,7 @@ def _init_arg_names(cls: type) -> frozenset[str]: ) -def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]: +def _get_redis_url_kwargs(client: type | None = None) -> tuple[str, ...]: """Connection kwargs that redis-py forwards from ``from_url`` down to the connection. ``from_url`` is declared as ``(cls, url, **kwargs)``, so introspecting it yields no @@ -92,9 +93,9 @@ def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]: """ if client is None: client = redis.Redis - connection_cls = async_redis.Connection if client is async_redis.Redis else redis.Connection + connection_cls: Final = async_redis.Connection if client is async_redis.Redis else redis.Connection - exclude_args = frozenset( + exclude_args: Final = frozenset( { "self", "connection_pool", @@ -103,7 +104,7 @@ def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]: ) # Only allow primitive arguments - include_args = ("url", "max_connections") + include_args: Final = ("url", "max_connections") return tuple(x for x in _init_arg_names(connection_cls) if x not in exclude_args) + include_args @@ -111,10 +112,10 @@ def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]: def _get_redis_cluster_kwargs(client=None): if client is None: client = redis.Redis.from_url - arg_spec = inspect.getfullargspec(redis.RedisCluster) + arg_spec: Final = inspect.getfullargspec(redis.RedisCluster) # Only allow primitive arguments - exclude_args = {"self", "connection_pool", "retry", "host", "port", "startup_nodes"} + exclude_args: Final = {"self", "connection_pool", "retry", "host", "port", "startup_nodes"} available_args = {x for x in arg_spec.args if x not in exclude_args} available_args |= { @@ -142,15 +143,15 @@ def _get_redis_cluster_kwargs(client=None): def _get_redis_env_kwarg_mapping(): - PREFIX = "REDIS_" + PREFIX: Final = "REDIS_" return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs()} def _redis_kwargs_from_environment(): - mapping = _get_redis_env_kwarg_mapping() + mapping: Final = _get_redis_env_kwarg_mapping() - return_dict = {} + return_dict: Final = {} for k, v in mapping.items(): value = get_secret(k, default_value=None) # type: ignore if value is not None: @@ -160,7 +161,7 @@ def _redis_kwargs_from_environment(): def create_gcp_iam_redis_connect_func( service_account: str, - ssl_ca_certs: Optional[str] = None, + ssl_ca_certs: str | None = None, ) -> Callable: """ Creates a custom Redis connection function for GCP IAM authentication. @@ -183,7 +184,7 @@ def create_gcp_iam_redis_connect_func( self._parser.on_connect(self) - auth_args = (_generate_gcp_iam_access_token(service_account),) + auth_args: Final = (_generate_gcp_iam_access_token(service_account),) self.send_command("AUTH", *auth_args, check_health=False) try: @@ -203,9 +204,9 @@ def create_gcp_iam_redis_connect_func( def _build_azure_credential( - azure_client_id: Optional[str] = None, - azure_tenant_id: Optional[str] = None, - azure_client_secret: Optional[str] = None, + azure_client_id: str | None = None, + azure_tenant_id: str | None = None, + azure_client_secret: str | None = None, ): """ Build a long-lived Azure credential object. @@ -224,9 +225,9 @@ def _build_azure_credential( "azure-identity is required for Azure AD Redis authentication. Install it with: pip install azure-identity" ) - _client_id = azure_client_id or os.environ.get("AZURE_CLIENT_ID") - _tenant_id = azure_tenant_id or os.environ.get("AZURE_TENANT_ID") - _client_secret = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET") + _client_id: Final = azure_client_id or os.environ.get("AZURE_CLIENT_ID") + _tenant_id: Final = azure_tenant_id or os.environ.get("AZURE_TENANT_ID") + _client_secret: Final = azure_client_secret or os.environ.get("AZURE_CLIENT_SECRET") if _client_id and _tenant_id and _client_secret: return ClientSecretCredential( @@ -241,9 +242,9 @@ def _build_azure_credential( def _generate_azure_ad_redis_token( - azure_client_id: Optional[str] = None, - azure_tenant_id: Optional[str] = None, - azure_client_secret: Optional[str] = None, + azure_client_id: str | None = None, + azure_tenant_id: str | None = None, + azure_client_secret: str | None = None, ) -> str: """ One-shot helper that builds a credential and fetches a single Azure AD @@ -253,19 +254,19 @@ def _generate_azure_ad_redis_token( (``AzureADCredentialProvider``) keep the credential alive across connections so the Azure SDK's internal cache + silent refresh apply. """ - credential = _build_azure_credential( + credential: Final = _build_azure_credential( azure_client_id=azure_client_id, azure_tenant_id=azure_tenant_id, azure_client_secret=azure_client_secret, ) - token = credential.get_token(AZURE_REDIS_SCOPE) + token: Final = credential.get_token(AZURE_REDIS_SCOPE) return token.token def create_azure_ad_redis_connect_func( - azure_client_id: Optional[str] = None, - azure_tenant_id: Optional[str] = None, - azure_client_secret: Optional[str] = None, + azure_client_id: str | None = None, + azure_tenant_id: str | None = None, + azure_client_secret: str | None = None, ) -> Callable: """ Creates a custom Redis connection function for Azure AD authentication. @@ -274,7 +275,7 @@ def create_azure_ad_redis_connect_func( closure) and reused across connections — the Azure SDK handles token caching and silent renewal internally. Only ``get_token`` is called per connection. """ - credential = _build_azure_credential( + credential: Final = _build_azure_credential( azure_client_id=azure_client_id, azure_tenant_id=azure_tenant_id, azure_client_secret=azure_client_secret, @@ -290,11 +291,11 @@ def create_azure_ad_redis_connect_func( self._parser.on_connect(self) - access_token = credential.get_token(AZURE_REDIS_SCOPE).token + access_token: Final = credential.get_token(AZURE_REDIS_SCOPE).token # Only include username when explicitly set — sending AUTH "" # is invalid for most ACL-configured Azure Redis instances. - username = os.environ.get("REDIS_USERNAME", "") + username: Final = os.environ.get("REDIS_USERNAME", "") if username: auth_args = (username, access_token) else: @@ -353,23 +354,23 @@ def _get_redis_client_logic(**env_overrides): value = get_secret(v) # type: ignore env_overrides[k] = value - environment_kwargs = _redis_kwargs_from_environment() + environment_kwargs: Final = _redis_kwargs_from_environment() # An explicitly configured connection target outranks REDIS_URL from the # environment. Without this, the url branch below strips the caller's # host/port/password and silently connects to whatever REDIS_URL names. - caller_named_a_target = any( + caller_named_a_target: Final = any( env_overrides.get(key) is not None for key in ("host", "startup_nodes", "sentinel_nodes") ) if caller_named_a_target and env_overrides.get("url") is None: environment_kwargs.pop("url", None) - redis_kwargs = { + redis_kwargs: Final = { **environment_kwargs, **env_overrides, } - _startup_nodes: Optional[Union[str, list]] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore + _startup_nodes: Final[str | list | None] = redis_kwargs.get("startup_nodes", None) or get_secret( # type: ignore "REDIS_CLUSTER_NODES" ) @@ -380,21 +381,21 @@ def _get_redis_client_logic(**env_overrides): elif _startup_nodes is None: redis_kwargs.pop("startup_nodes", None) - _sentinel_nodes: Optional[Union[str, list]] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore + _sentinel_nodes: Final[str | list | None] = redis_kwargs.get("sentinel_nodes", None) or get_secret( # type: ignore "REDIS_SENTINEL_NODES" ) if _sentinel_nodes is not None and isinstance(_sentinel_nodes, str): redis_kwargs["sentinel_nodes"] = json.loads(_sentinel_nodes) - _sentinel_password: Optional[str] = redis_kwargs.get("sentinel_password", None) or get_secret_str( + _sentinel_password: Final[str | None] = redis_kwargs.get("sentinel_password", None) or get_secret_str( "REDIS_SENTINEL_PASSWORD" ) if _sentinel_password is not None: redis_kwargs["sentinel_password"] = _sentinel_password - _service_name: Optional[str] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore + _service_name: Final[str | None] = redis_kwargs.get("service_name", None) or get_secret( # type: ignore "REDIS_SERVICE_NAME" ) @@ -402,8 +403,8 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs["service_name"] = _service_name # Handle GCP IAM authentication - _gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT") - _gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS") + _gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT") + _gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS") if _gcp_service_account is not None: verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.") @@ -422,9 +423,9 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs # Handle Azure AD authentication (after GCP IAM block) - _azure_redis_ad_token = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN") + _azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN") - _azure_ad_enabled = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true" + _azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true" if _azure_ad_enabled and _gcp_service_account is not None: verbose_logger.warning( @@ -433,9 +434,9 @@ def _get_redis_client_logic(**env_overrides): ) if _azure_ad_enabled and _gcp_service_account is None: - _azure_client_id = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID") - _azure_tenant_id = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID") - _azure_client_secret = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET") + _azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID") + _azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID") + _azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET") verbose_logger.debug("Setting up Azure AD authentication for Redis.") redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func( @@ -465,9 +466,12 @@ def _get_redis_client_logic(**env_overrides): redis_kwargs.pop("port", None) redis_kwargs.pop("db", None) redis_kwargs.pop("password", None) - elif "startup_nodes" in redis_kwargs and redis_kwargs["startup_nodes"] is not None: - pass - elif "sentinel_nodes" in redis_kwargs and redis_kwargs["sentinel_nodes"] is not None: + elif ( + "startup_nodes" in redis_kwargs + and redis_kwargs["startup_nodes"] is not None + or "sentinel_nodes" in redis_kwargs + and redis_kwargs["sentinel_nodes"] is not None + ): pass elif "host" not in redis_kwargs or redis_kwargs["host"] is None: raise ValueError("Either 'host' or 'url' must be specified for redis.") @@ -477,7 +481,7 @@ def _get_redis_client_logic(**env_overrides): def init_redis_cluster(redis_kwargs) -> redis.RedisCluster: - _redis_cluster_nodes_in_env: Optional[str] = get_secret("REDIS_CLUSTER_NODES") # type: ignore + _redis_cluster_nodes_in_env: Final[str | None] = get_secret("REDIS_CLUSTER_NODES") # type: ignore if _redis_cluster_nodes_in_env is not None: try: redis_kwargs["startup_nodes"] = json.loads(_redis_cluster_nodes_in_env) @@ -489,13 +493,13 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster: verbose_logger.debug("init_redis_cluster: startup nodes are being initialized.") from redis.cluster import ClusterNode - args = _get_redis_cluster_kwargs() - cluster_kwargs = {} + args: Final = _get_redis_cluster_kwargs() + cluster_kwargs: Final = {} for arg in redis_kwargs: if arg in args: cluster_kwargs[arg] = redis_kwargs[arg] - new_startup_nodes: List[ClusterNode] = [] + new_startup_nodes: Final[list[ClusterNode]] = [] for item in redis_kwargs["startup_nodes"]: new_startup_nodes.append(ClusterNode(**item)) @@ -505,8 +509,8 @@ def init_redis_cluster(redis_kwargs) -> redis.RedisCluster: def _get_redis_sentinel_connection_kwargs(redis_kwargs: dict) -> dict: - connection_kwargs = {} - args = _get_redis_kwargs() + connection_kwargs: Final = {} + args: Final = _get_redis_kwargs() for arg in redis_kwargs: if arg in args: connection_kwargs[arg] = redis_kwargs[arg] @@ -515,12 +519,12 @@ def _get_redis_sentinel_connection_kwargs(redis_kwargs: dict) -> dict: def _init_redis_sentinel(redis_kwargs) -> redis.Redis: - sentinel_nodes = redis_kwargs.get("sentinel_nodes") - sentinel_password = redis_kwargs.get("sentinel_password") - service_name = redis_kwargs.get("service_name") - connection_kwargs = _get_redis_sentinel_connection_kwargs(redis_kwargs) + sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes") + sentinel_password: Final = redis_kwargs.get("sentinel_password") + service_name: Final = redis_kwargs.get("service_name") + connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs) connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT) - sentinel_kwargs = dict(connection_kwargs) + sentinel_kwargs: Final = dict(connection_kwargs) sentinel_kwargs["password"] = sentinel_password if not sentinel_nodes or not service_name: @@ -529,7 +533,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis: verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.") # Set up the Sentinel client - sentinel = redis.Sentinel( + sentinel: Final = redis.Sentinel( sentinel_nodes, sentinel_kwargs=sentinel_kwargs, ) @@ -540,12 +544,12 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis: def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: - sentinel_nodes = redis_kwargs.get("sentinel_nodes") - sentinel_password = redis_kwargs.get("sentinel_password") - service_name = redis_kwargs.get("service_name") - connection_kwargs = _get_redis_sentinel_connection_kwargs(redis_kwargs) + sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes") + sentinel_password: Final = redis_kwargs.get("sentinel_password") + service_name: Final = redis_kwargs.get("service_name") + connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs) connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT) - sentinel_kwargs = dict(connection_kwargs) + sentinel_kwargs: Final = dict(connection_kwargs) sentinel_kwargs["password"] = sentinel_password if not sentinel_nodes or not service_name: @@ -554,7 +558,7 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: verbose_logger.debug("init_redis_sentinel: sentinel nodes are being initialized.") # Set up the Sentinel client - sentinel = async_redis.Sentinel( + sentinel: Final = async_redis.Sentinel( sentinel_nodes, sentinel_kwargs=sentinel_kwargs, ) @@ -565,14 +569,14 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis: def get_redis_client(**env_overrides): - redis_kwargs = _get_redis_client_logic(**env_overrides) + redis_kwargs: Final = _get_redis_client_logic(**env_overrides) if "startup_nodes" in redis_kwargs: return init_redis_cluster(redis_kwargs) if "url" in redis_kwargs and redis_kwargs["url"] is not None: - args = _get_redis_url_kwargs() - url_kwargs = {} + args: Final = _get_redis_url_kwargs() + url_kwargs: Final = {} for arg in redis_kwargs: if arg in args: url_kwargs[arg] = redis_kwargs[arg] @@ -587,16 +591,16 @@ def get_redis_client(**env_overrides): def get_redis_async_client( - connection_pool: Optional[async_redis.BlockingConnectionPool] = None, + connection_pool: async_redis.BlockingConnectionPool | None = None, **env_overrides, -) -> Union[async_redis.Redis, async_redis.RedisCluster]: - redis_kwargs = _get_redis_client_logic(**env_overrides) +) -> async_redis.Redis | async_redis.RedisCluster: + redis_kwargs: Final = _get_redis_client_logic(**env_overrides) if "startup_nodes" in redis_kwargs: from redis.cluster import ClusterNode args = _get_redis_cluster_kwargs() - cluster_kwargs = {} + cluster_kwargs: Final = {} for arg in redis_kwargs: if arg in args: cluster_kwargs[arg] = redis_kwargs[arg] @@ -618,7 +622,7 @@ def get_redis_async_client( username=os.environ.get("REDIS_USERNAME") or None, ) - new_startup_nodes: List[ClusterNode] = [] + new_startup_nodes: Final[list[ClusterNode]] = [] for item in redis_kwargs["startup_nodes"]: new_startup_nodes.append(ClusterNode(**item)) @@ -632,7 +636,7 @@ def get_redis_async_client( cluster_kwargs.setdefault("socket_keepalive", True) # Create async RedisCluster with IAM token as password if available - cluster_client = async_redis.RedisCluster( + cluster_client: Final = async_redis.RedisCluster( startup_nodes=new_startup_nodes, **cluster_kwargs, # type: ignore ) @@ -643,13 +647,13 @@ def get_redis_async_client( if connection_pool is not None: return async_redis.Redis(connection_pool=connection_pool) args = _get_redis_url_kwargs(client=async_redis.Redis) - url_kwargs = {} + url_kwargs: Final = {} for arg in redis_kwargs: if arg in args: url_kwargs[arg] = redis_kwargs[arg] else: verbose_logger.debug( - "REDIS: ignoring argument: {}. Not an allowed async_redis.Redis.from_url arg.".format(arg) + "REDIS: ignoring argument: %s. Not an allowed async_redis.Redis.from_url arg.", arg ) return async_redis.Redis.from_url(**url_kwargs) @@ -682,16 +686,16 @@ def get_redis_async_client( def get_redis_connection_pool( **env_overrides, -) -> Optional[async_redis.BlockingConnectionPool]: - redis_kwargs = _get_redis_client_logic(**env_overrides) +) -> async_redis.BlockingConnectionPool | None: + redis_kwargs: Final = _get_redis_client_logic(**env_overrides) verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs) if "startup_nodes" in redis_kwargs: return None if "url" in redis_kwargs and redis_kwargs["url"] is not None: - allowed_args = _get_redis_url_kwargs(client=async_redis.Redis) - pool_kwargs = {k: v for k, v in redis_kwargs.items() if k in allowed_args and k != "max_connections"} + allowed_args: Final = _get_redis_url_kwargs(client=async_redis.Redis) + pool_kwargs: Final = {k: v for k, v in redis_kwargs.items() if k in allowed_args and k != "max_connections"} pool_kwargs["timeout"] = REDIS_CONNECTION_POOL_TIMEOUT pool_kwargs["url"] = redis_kwargs["url"] if "max_connections" in redis_kwargs: @@ -707,7 +711,7 @@ def get_redis_connection_pool( # Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed # connections re-fetch tokens via the SDK's internal cache + silent refresh # rather than reusing a single token captured at pool creation. - redis_connect_func = redis_kwargs.pop("redis_connect_func", None) + redis_connect_func: Final = redis_kwargs.pop("redis_connect_func", None) if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"): redis_kwargs["credential_provider"] = AzureADCredentialProvider( redis_connect_func._azure_credential, @@ -734,7 +738,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: if not verbose_logger.isEnabledFor(logging.DEBUG): return - console = Console() + console: Final = Console() # Initialize the sensitive data masker masker = SensitiveDataMasker() @@ -743,10 +747,10 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: masked_redis_kwargs = masker.mask_dict(redis_kwargs) # Create main panel title - title = Text("Redis Configuration", style="bold blue") + title: Final = Text("Redis Configuration", style="bold blue") # Create configuration table - config_table = Table( + config_table: Final = Table( title="🔧 Redis Connection Parameters", show_header=True, header_style="bold magenta", @@ -783,7 +787,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: connection_type = "Redis (URL-based)" # Create connection type info - info_table = Table( + info_table: Final = Table( title="📊 Connection Info", show_header=True, header_style="bold green", @@ -804,6 +808,6 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None: # Fallback to simple logging if rich is not available masker = SensitiveDataMasker() masked_redis_kwargs = masker.mask_dict(redis_kwargs) - verbose_logger.info(f"Redis configuration: {masked_redis_kwargs}") + verbose_logger.info("Redis configuration: %s", masked_redis_kwargs) except Exception as e: - verbose_logger.error(f"Error pretty printing Redis configuration: {e}") + verbose_logger.error("Error pretty printing Redis configuration: %s", e) diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py index b973e292a17..8b8bbf9366f 100644 --- a/litellm/_redis_credential_provider.py +++ b/litellm/_redis_credential_provider.py @@ -1,21 +1,21 @@ import asyncio import threading import time -from typing import Any, Dict, Optional, Tuple, Union +from typing import Any, Final from redis.credentials import CredentialProvider # type: ignore[attr-defined] # Azure AD scope for Redis Cache for Azure. -AZURE_REDIS_SCOPE = "https://redis.azure.com/.default" +AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default" # GCP IAM tokens are valid for 1 hour. Cache for 55 minutes to refresh before expiry. -_GCP_IAM_TOKEN_TTL_SECONDS = 3300 +_GCP_IAM_TOKEN_TTL_SECONDS: Final = 3300 # Module-level cache shared across all GCPIAMCredentialProvider instances for the # same service account, so multiple Redis connections on the same pod share one token. # Keyed by service_account → (token, expiry_monotonic_timestamp). -_token_cache: Dict[str, Tuple[str, float]] = {} -_token_cache_lock = threading.Lock() +_token_cache: Final[dict[str, tuple[str, float]]] = {} +_token_cache_lock: Final = threading.Lock() def _generate_gcp_iam_access_token(service_account: str) -> str: @@ -36,12 +36,12 @@ def _generate_gcp_iam_access_token(service_account: str) -> str: "Install it with: pip install google-cloud-iam" ) - client = iam_credentials_v1.IAMCredentialsClient() - request = iam_credentials_v1.GenerateAccessTokenRequest( + client: Final = iam_credentials_v1.IAMCredentialsClient() + request: Final = iam_credentials_v1.GenerateAccessTokenRequest( name=service_account, scope=["https://www.googleapis.com/auth/cloud-platform"], ) - response = client.generate_access_token(request=request) + response: Final = client.generate_access_token(request=request) return str(response.access_token) @@ -95,12 +95,12 @@ class GCPIAMCredentialProvider(CredentialProvider): def __init__(self, gcp_service_account: str) -> None: self._gcp_service_account = gcp_service_account - def get_credentials(self) -> Tuple[str]: - token = _get_cached_gcp_iam_token(self._gcp_service_account) + def get_credentials(self) -> tuple[str]: + token: Final = _get_cached_gcp_iam_token(self._gcp_service_account) return (token,) - async def get_credentials_async(self) -> Tuple[str]: - token = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account) + async def get_credentials_async(self) -> tuple[str]: + token: Final = await asyncio.to_thread(_get_cached_gcp_iam_token, self._gcp_service_account) return (token,) @@ -115,18 +115,18 @@ class AzureADCredentialProvider(CredentialProvider): fail authentication after the initial token expired (~1 hour TTL). """ - def __init__(self, credential: Any, username: Optional[str] = None) -> None: + def __init__(self, credential: Any, username: str | None = None) -> None: self._credential = credential self._username = username - def get_credentials(self) -> Union[Tuple[str], Tuple[str, str]]: - token = self._credential.get_token(AZURE_REDIS_SCOPE).token + def get_credentials(self) -> tuple[str] | tuple[str, str]: + token: Final = self._credential.get_token(AZURE_REDIS_SCOPE).token if self._username: return (self._username, token) return (token,) - async def get_credentials_async(self) -> Union[Tuple[str], Tuple[str, str]]: - token_obj = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE) + async def get_credentials_async(self) -> tuple[str] | tuple[str, str]: + token_obj: Final = await asyncio.to_thread(self._credential.get_token, AZURE_REDIS_SCOPE) if self._username: return (self._username, token_obj.token) return (token_obj.token,) diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index b1bd0a3bba2..06ae3f41c19 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime, timedelta -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Final, Union import litellm from litellm._logging import verbose_logger @@ -24,7 +24,7 @@ else: UserAPIKeyAuth = Any -def _get_otel_v2_class() -> Optional[type]: +def _get_otel_v2_class() -> type | None: """Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent. Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry @@ -54,7 +54,7 @@ class ServiceLogging(CustomLogger): if "prometheus_system" in litellm.service_callback: self.prometheusServicesLogger = PrometheusServicesLogger() - def _resolve_otel_service_logger(self, callback: Any) -> Optional[Any]: + def _resolve_otel_service_logger(self, callback: Any) -> Any | None: """Resolve the OTel logger (legacy or V2) to emit a service span on. Returns the logger instance whose ``async_service_*_hook`` should fire for @@ -67,7 +67,7 @@ class ServiceLogging(CustomLogger): whether the callback is the logger instance itself or the ``"otel"`` string (which routes to the proxy's registered ``open_telemetry_logger``). """ - otel_v2_cls = _get_otel_v2_class() + otel_v2_cls: Final = _get_otel_v2_class() def _is_otel_logger(obj: Any) -> bool: if isinstance(obj, OpenTelemetry): @@ -88,9 +88,9 @@ class ServiceLogging(CustomLogger): service: ServiceTypes, duration: float, call_type: str, - parent_otel_span: Optional[Span] = None, - start_time: Optional[Union[datetime, float]] = None, - end_time: Optional[Union[float, datetime]] = None, + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: float | datetime | None = None, ): """ Handles both sync and async monitoring by checking for existing event loop. @@ -101,7 +101,7 @@ class ServiceLogging(CustomLogger): try: # Try to get the current event loop - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() # Check if the loop is running if loop.is_running(): # If we're in a running loop, create a task @@ -152,10 +152,10 @@ class ServiceLogging(CustomLogger): service: ServiceTypes, call_type: str, duration: float, - parent_otel_span: Optional[Span] = None, - start_time: Optional[Union[datetime, float]] = None, - end_time: Optional[Union[datetime, float]] = None, - event_metadata: Optional[dict] = None, + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: datetime | float | None = None, + event_metadata: dict | None = None, ): """ - For counting if the redis, postgres call is successful @@ -163,7 +163,7 @@ class ServiceLogging(CustomLogger): if self.mock_testing: self.mock_testing_async_success_hook += 1 - payload = ServiceLoggerPayload( + payload: Final = ServiceLoggerPayload( is_error=False, error=None, service=service, @@ -178,7 +178,7 @@ class ServiceLogging(CustomLogger): # (the V2 logger self-registers its instance even when the string is # present, unlike V1). Without this guard each such reference emits its own # span, so a single DB call shows up as duplicate ``postgres ...`` spans. - emitted_otel_logger_ids: set = set() + emitted_otel_logger_ids: Final[set] = set() for callback in litellm.service_callback: if callback == "prometheus_system": await self.init_prometheus_services_logger_if_none() @@ -218,7 +218,6 @@ class ServiceLogging(CustomLogger): self.prometheusServicesLogger = PrometheusServicesLogger() elif self.prometheusServicesLogger is None: self.prometheusServicesLogger = self.prometheusServicesLogger() - return async def init_datadog_logger_if_none(self): """ @@ -230,8 +229,6 @@ class ServiceLogging(CustomLogger): if not hasattr(self, "dd_logger"): self.dd_logger: DataDogLogger = DataDogLogger() - return - async def init_otel_logger_if_none(self): """ initializes otel_logger if it is None or no attribute exists on ServiceLogging Object @@ -246,18 +243,17 @@ class ServiceLogging(CustomLogger): verbose_logger.warning( "ServiceLogger: open_telemetry_logger is None or not an instance of OpenTelemetry" ) - return async def async_service_failure_hook( self, service: ServiceTypes, duration: float, - error: Union[str, Exception], + error: str | Exception, call_type: str, - parent_otel_span: Optional[Span] = None, - start_time: Optional[Union[datetime, float]] = None, - end_time: Optional[Union[float, datetime]] = None, - event_metadata: Optional[dict] = None, + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: float | datetime | None = None, + event_metadata: dict | None = None, ): """ - For counting if the redis, postgres call is unsuccessful @@ -271,7 +267,7 @@ class ServiceLogging(CustomLogger): elif isinstance(error, str): error_message = error - payload = ServiceLoggerPayload( + payload: Final = ServiceLoggerPayload( is_error=True, error=error_message, service=service, @@ -282,7 +278,7 @@ class ServiceLogging(CustomLogger): # Dedupe OTel loggers per event — see ``async_service_success_hook`` for why # the same logger can be referenced twice in ``service_callback``. - emitted_otel_logger_ids: set = set() + emitted_otel_logger_ids: Final[set] = set() for callback in litellm.service_callback: if callback == "prometheus_system": await self.init_prometheus_services_logger_if_none() @@ -324,7 +320,7 @@ class ServiceLogging(CustomLogger): request_data: dict, original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, - traceback_str: Optional[str] = None, + traceback_str: str | None = None, ): """ Hook to track failed litellm-service calls @@ -347,7 +343,7 @@ class ServiceLogging(CustomLogger): pass else: raise Exception( - "Duration={} is not a float or timedelta object. type={}".format(_duration, type(_duration)) + f"Duration={_duration} is not a float or timedelta object. type={type(_duration)}" ) # invalid _duration value # Batch polling callbacks (check_batch_cost) don't include call_type in kwargs. # Use .get() to avoid KeyError. diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index 412c7a0897d..f6b74bcbb42 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -4,7 +4,7 @@ Custom A2A Card Resolver for LiteLLM. Extends the A2A SDK's card resolver to support multiple well-known paths. """ -from typing import TYPE_CHECKING, Any, Dict +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger from litellm.constants import LOCALHOST_URL_PATTERNS @@ -43,18 +43,18 @@ def is_localhost_or_internal_url(url: str | None) -> bool: if not url: return False - url_lower = url.lower() + url_lower: Final = url.lower() return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS) def get_agent_card_url(agent_card: "AgentCard") -> str | None: """Return the agent endpoint URL from the resolved SDK card.""" - url = getattr(agent_card, "url", None) + url: Final = getattr(agent_card, "url", None) if url: return url - interfaces = getattr(agent_card, "supported_interfaces", None) + interfaces: Final = getattr(agent_card, "supported_interfaces", None) if interfaces: return getattr(interfaces[0], "url", None) return None @@ -62,11 +62,11 @@ def get_agent_card_url(agent_card: "AgentCard") -> str | None: def set_agent_card_url(agent_card: "AgentCard", url: str) -> None: """Set the agent endpoint URL on the resolved SDK card.""" - normalized = url.rstrip("/") + "/" + normalized: Final = url.rstrip("/") + "/" if hasattr(agent_card, "url"): agent_card.url = normalized - interfaces = getattr(agent_card, "supported_interfaces", None) + interfaces: Final = getattr(agent_card, "supported_interfaces", None) if interfaces: interfaces[0].url = normalized @@ -86,16 +86,16 @@ def fix_agent_card_url(agent_card: "AgentCard", base_url: str) -> "AgentCard": Returns: The agent card with the URL fixed if necessary """ - card_url = getattr(agent_card, "url", None) + card_url: Final = getattr(agent_card, "url", None) if card_url and is_localhost_or_internal_url(card_url): # Normalize base_url to ensure it ends with / - fixed_url = base_url.rstrip("/") + "/" + fixed_url: Final = base_url.rstrip("/") + "/" agent_card.url = fixed_url - interfaces = getattr(agent_card, "supported_interfaces", None) + interfaces: Final = getattr(agent_card, "supported_interfaces", None) if interfaces: - interface_url = getattr(interfaces[0], "url", None) + interface_url: Final = getattr(interfaces[0], "url", None) if interface_url and is_localhost_or_internal_url(interface_url): interfaces[0].url = base_url.rstrip("/") + "/" @@ -114,7 +114,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] async def get_agent_card( self, relative_card_path: str | None = None, - http_kwargs: Dict[str, Any] | None = None, + http_kwargs: dict[str, Any] | None = None, ) -> "AgentCard": """ Fetch the agent card, trying multiple well-known paths. @@ -140,7 +140,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] ) # Try both well-known paths - paths = [ + paths: Final = [ AGENT_CARD_WELL_KNOWN_PATH, PREV_AGENT_CARD_WELL_KNOWN_PATH, ] @@ -148,13 +148,13 @@ class LiteLLMA2ACardResolver(_A2ACardResolver): # type: ignore[misc] last_error = None for path in paths: try: - verbose_logger.debug(f"Attempting to fetch agent card from {self.base_url}{path}") + verbose_logger.debug("Attempting to fetch agent card from %s%s", self.base_url, path) return await super().get_agent_card( relative_card_path=path, http_kwargs=http_kwargs, ) except Exception as e: - verbose_logger.debug(f"Failed to fetch agent card from {self.base_url}{path}: {e}") + verbose_logger.debug("Failed to fetch agent card from %s%s: %s", self.base_url, path, e) last_error = e continue diff --git a/litellm/a2a_protocol/client.py b/litellm/a2a_protocol/client.py index a05f8dc390c..0ded50d25b7 100644 --- a/litellm/a2a_protocol/client.py +++ b/litellm/a2a_protocol/client.py @@ -4,7 +4,8 @@ LiteLLM A2A Client class. Provides a class-based interface for A2A agent invocation. """ -from typing import TYPE_CHECKING, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Final from litellm.types.agents import LiteLLMSendMessageResponse @@ -50,7 +51,7 @@ class A2AClient: self, base_url: str, timeout: float = 60.0, - extra_headers: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, ): """ Initialize the A2A client wrapper. @@ -63,7 +64,7 @@ class A2AClient: self.base_url = base_url self.timeout = timeout self.extra_headers = extra_headers - self._a2a_client: Optional["A2AClientType"] = None + self._a2a_client: A2AClientType | None = None async def _get_client(self) -> "A2AClientType": """Get or create the underlying A2A client.""" @@ -91,7 +92,7 @@ class A2AClient: """Send a message to the A2A agent.""" from litellm.a2a_protocol.main import asend_message - a2a_client = await self._get_client() + a2a_client: Final = await self._get_client() return await asend_message(a2a_client=a2a_client, request=request) async def send_message_streaming( @@ -100,6 +101,6 @@ class A2AClient: """Send a streaming message to the A2A agent.""" from litellm.a2a_protocol.main import asend_message_streaming - a2a_client = await self._get_client() + a2a_client: Final = await self._get_client() async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request): yield chunk diff --git a/litellm/a2a_protocol/cost_calculator.py b/litellm/a2a_protocol/cost_calculator.py index f3e84c5b84d..31ca81c44dd 100644 --- a/litellm/a2a_protocol/cost_calculator.py +++ b/litellm/a2a_protocol/cost_calculator.py @@ -5,7 +5,7 @@ Supports dynamic cost parameters that allow platform owners to define custom costs per agent query or per token. """ -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import ( @@ -18,7 +18,7 @@ else: class A2ACostCalculator: @staticmethod def calculate_a2a_cost( - litellm_logging_obj: Optional[LitellmLoggingObject], + litellm_logging_obj: LitellmLoggingObject | None, ) -> float: """ Calculate the cost of an A2A send_message call. @@ -42,23 +42,23 @@ class A2ACostCalculator: if litellm_logging_obj is None: return 0.0 - model_call_details = litellm_logging_obj.model_call_details + model_call_details: Final = litellm_logging_obj.model_call_details # Check if user set a custom response cost (backward compatibility) - response_cost = model_call_details.get("response_cost", None) + response_cost: Final = model_call_details.get("response_cost", None) if response_cost is not None: return float(response_cost) # Get litellm_params for cost parameters - litellm_params = model_call_details.get("litellm_params", {}) or {} + litellm_params: Final = model_call_details.get("litellm_params", {}) or {} # Check for cost_per_query (fixed cost per query) if litellm_params.get("cost_per_query") is not None: return float(litellm_params["cost_per_query"]) # Check for token-based pricing - input_cost_per_token = litellm_params.get("input_cost_per_token") - output_cost_per_token = litellm_params.get("output_cost_per_token") + input_cost_per_token: Final = litellm_params.get("input_cost_per_token") + output_cost_per_token: Final = litellm_params.get("output_cost_per_token") if input_cost_per_token is not None or output_cost_per_token is not None: return A2ACostCalculator._calculate_token_based_cost( @@ -73,8 +73,8 @@ class A2ACostCalculator: @staticmethod def _calculate_token_based_cost( model_call_details: dict, - input_cost_per_token: Optional[float], - output_cost_per_token: Optional[float], + input_cost_per_token: float | None, + output_cost_per_token: float | None, ) -> float: """ Calculate cost based on token usage and per-token pricing. @@ -88,16 +88,16 @@ class A2ACostCalculator: float: The calculated cost """ # Get usage from model_call_details - usage = model_call_details.get("usage") + usage: Final = model_call_details.get("usage") if usage is None: return 0.0 # Get token counts - prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0 - completion_tokens = getattr(usage, "completion_tokens", 0) or 0 + prompt_tokens: Final = getattr(usage, "prompt_tokens", 0) or 0 + completion_tokens: Final = getattr(usage, "completion_tokens", 0) or 0 # Calculate costs - input_cost = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0) - output_cost = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0) + input_cost: Final = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0) + output_cost: Final = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0) return input_cost + output_cost diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py index 89b831351ab..d2c4cdf7a65 100644 --- a/litellm/a2a_protocol/exception_mapping_utils.py +++ b/litellm/a2a_protocol/exception_mapping_utils.py @@ -4,7 +4,7 @@ A2A Protocol Exception Mapping Utils. Maps A2A SDK exceptions to LiteLLM A2A exception types. """ -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger from litellm.a2a_protocol.card_resolver import ( @@ -53,11 +53,11 @@ class A2AExceptionCheckers: if not isinstance(error_str, str): return False - error_str_lower = error_str.lower() + error_str_lower: Final = error_str.lower() return any(pattern in error_str_lower for pattern in CONNECTION_ERROR_PATTERNS) @staticmethod - def is_localhost_url(url: Optional[str]) -> bool: + def is_localhost_url(url: str | None) -> bool: """ Check if a URL is a localhost/internal URL. @@ -83,8 +83,8 @@ class A2AExceptionCheckers: if not isinstance(error_str, str): return False - error_str_lower = error_str.lower() - agent_card_patterns = [ + error_str_lower: Final = error_str.lower() + agent_card_patterns: Final = [ "agent card", "agent-card", ".well-known", @@ -96,9 +96,9 @@ class A2AExceptionCheckers: def map_a2a_exception( original_exception: Exception, - card_url: Optional[str] = None, - api_base: Optional[str] = None, - model: Optional[str] = None, + card_url: str | None = None, + api_base: str | None = None, + model: str | None = None, ) -> Exception: """ Map an A2A SDK exception to a LiteLLM A2A exception type. @@ -118,7 +118,7 @@ def map_a2a_exception( A2AAgentCardError: If the error is related to agent card issues A2AError: For other A2A-related errors """ - error_str = str(original_exception) + error_str: Final = str(original_exception) # Check for localhost URL connection error (special case - retryable) if ( @@ -190,11 +190,13 @@ async def handle_a2a_localhost_retry( "rewrite, so the upstream URL cannot be corrected." ) - request_type = "streaming " if is_streaming else "" + request_type: Final = "streaming " if is_streaming else "" verbose_logger.warning( - f"A2A {request_type}request to '{error.localhost_url}' failed: {error.original_error}. " - f"Agent card contains localhost/internal URL. " - f"Retrying with base_url '{error.base_url}'." + "A2A %srequest to '%s' failed: %s. Agent card contains localhost/internal URL. Retrying with base_url '%s'.", + request_type, + error.localhost_url, + error.original_error, + error.base_url, ) # Fix the agent card URL @@ -203,14 +205,14 @@ async def handle_a2a_localhost_retry( # Reuse the httpx client LiteLLM attached at creation. It carries this agent's # trace-id and auth headers, so a fresh client would drop them. Only clients built # by ``create_a2a_client`` have it; an externally-supplied client cannot be retried. - httpx_client = getattr(a2a_client, "_litellm_httpx_client", None) + httpx_client: Final = getattr(a2a_client, "_litellm_httpx_client", None) if httpx_client is None: raise RuntimeError( "Cannot retry A2A localhost URL fix: the client was not created by " "create_a2a_client, so no LiteLLM httpx client is attached." ) - new_client = await create_client( # pyright: ignore[reportOptionalCall] + new_client: Final = await create_client( # pyright: ignore[reportOptionalCall] agent_card, client_config=ClientConfig( # pyright: ignore[reportOptionalCall] httpx_client=httpx_client, diff --git a/litellm/a2a_protocol/exceptions.py b/litellm/a2a_protocol/exceptions.py index b672971e727..2542cbc67b0 100644 --- a/litellm/a2a_protocol/exceptions.py +++ b/litellm/a2a_protocol/exceptions.py @@ -4,8 +4,6 @@ A2A Protocol Exceptions. Custom exception types for A2A protocol operations, following LiteLLM's exception pattern. """ -from typing import Optional - import httpx @@ -21,11 +19,11 @@ class A2AError(Exception): message: str, status_code: int = 500, llm_provider: str = "a2a_agent", - model: Optional[str] = None, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + model: str | None = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = status_code self.message = f"litellm.A2AError: {message}" @@ -65,12 +63,12 @@ class A2AConnectionError(A2AError): def __init__( self, message: str, - url: Optional[str] = None, - model: Optional[str] = None, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + url: str | None = None, + model: str | None = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.url = url super().__init__( @@ -98,10 +96,10 @@ class A2AAgentCardError(A2AError): def __init__( self, message: str, - url: Optional[str] = None, - model: Optional[str] = None, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, + url: str | None = None, + model: str | None = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, ): self.url = url super().__init__( @@ -132,8 +130,8 @@ class A2ALocalhostURLError(A2AConnectionError): self, localhost_url: str, base_url: str, - original_error: Optional[Exception] = None, - model: Optional[str] = None, + original_error: Exception | None = None, + model: str | None = None, ): self.localhost_url = localhost_url self.base_url = base_url diff --git a/litellm/a2a_protocol/litellm_completion_bridge/__init__.py b/litellm/a2a_protocol/litellm_completion_bridge/__init__.py index 6c9df0ee285..a81f5304f7c 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/__init__.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/__init__.py @@ -16,8 +16,8 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import ( ) __all__ = [ - "A2ACompletionBridgeTransformation", "A2ACompletionBridgeHandler", + "A2ACompletionBridgeTransformation", "handle_a2a_completion", "handle_a2a_completion_streaming", ] diff --git a/litellm/a2a_protocol/litellm_completion_bridge/handler.py b/litellm/a2a_protocol/litellm_completion_bridge/handler.py index a84b23a2170..9c0564ca594 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/handler.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/handler.py @@ -10,7 +10,8 @@ A2A Streaming Events (in order): 4. Status update (kind: "status-update") - Final status "completed" with final=true """ -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any, Final import litellm from litellm._logging import verbose_logger @@ -24,10 +25,10 @@ from litellm.interactions.agents.utils import merge_agent_headers # litellm_params key carrying the authenticated principal (hashed virtual key) so # A2A provider configs can scope provider-side state (e.g. LangFlow session memory) # per key instead of trusting the client-supplied A2A contextId. -A2A_USER_API_KEY_HASH_PARAM = "litellm_a2a_user_api_key_hash" +A2A_USER_API_KEY_HASH_PARAM: Final = "litellm_a2a_user_api_key_hash" # Agent metadata fields stored in litellm_params that are not valid litellm.acompletion() kwargs -_AGENT_ONLY_PARAMS = frozenset( +_AGENT_ONLY_PARAMS: Final = frozenset( { "is_public", "agent_name", @@ -46,13 +47,13 @@ class A2ACompletionBridgeHandler: @staticmethod async def handle_non_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], - api_base: Optional[str] = None, - agent_extra_headers: Optional[Dict[str, str]] = None, + params: dict[str, Any], + litellm_params: dict[str, Any], + api_base: str | None = None, + agent_extra_headers: dict[str, str] | None = None, *, _skip_a2a_provider_routing: bool = False, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Handle non-streaming A2A request via litellm.acompletion. @@ -69,13 +70,13 @@ class A2ACompletionBridgeHandler: """ custom_llm_provider = litellm_params.get("custom_llm_provider") if not _skip_a2a_provider_routing: - a2a_provider_config = A2AProviderConfigManager.get_provider_config( + a2a_provider_config: Final = A2AProviderConfigManager.get_provider_config( custom_llm_provider=custom_llm_provider, model=litellm_params.get("model"), ) if a2a_provider_config is not None: - verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}") + verbose_logger.info("A2A: Using provider config for %s", custom_llm_provider) return await a2a_provider_config.handle_non_streaming( request_id=request_id, @@ -86,14 +87,14 @@ class A2ACompletionBridgeHandler: ) # Extract message from params - message = params.get("message", {}) + message: Final = params.get("message", {}) # Transform A2A message to OpenAI format - openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) + openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) # Get completion params custom_llm_provider = litellm_params.get("custom_llm_provider") - model = litellm_params.get("model", "agent") + model: Final = litellm_params.get("model", "agent") # Build full model string if provider specified # Skip prepending if model already starts with the provider prefix @@ -102,17 +103,17 @@ class A2ACompletionBridgeHandler: else: full_model = model - verbose_logger.info(f"A2A completion bridge: model={full_model}, api_base={api_base}") + verbose_logger.info("A2A completion bridge: model=%s, api_base=%s", full_model, api_base) # Build completion params dict - completion_params: Dict[str, Any] = { + completion_params: Final[dict[str, Any]] = { "model": full_model, "messages": openai_messages, "api_base": api_base, "stream": False, } # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) - litellm_params_to_add = { + litellm_params_to_add: Final = { k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS @@ -134,28 +135,28 @@ class A2ACompletionBridgeHandler: ) # Call litellm.acompletion - response = await litellm.acompletion(**completion_params) + response: Final = await litellm.acompletion(**completion_params) # Transform response to A2A format - a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response( + a2a_response: Final = A2ACompletionBridgeTransformation.openai_response_to_a2a_response( response=response, request_id=request_id, ) - verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}") + verbose_logger.info("A2A completion bridge completed: request_id=%s", request_id) return a2a_response @staticmethod async def handle_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], - api_base: Optional[str] = None, - agent_extra_headers: Optional[Dict[str, str]] = None, + params: dict[str, Any], + litellm_params: dict[str, Any], + api_base: str | None = None, + agent_extra_headers: dict[str, str] | None = None, *, _skip_a2a_provider_routing: bool = False, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: """ Handle streaming A2A request via litellm.acompletion with stream=True. @@ -178,13 +179,13 @@ class A2ACompletionBridgeHandler: """ custom_llm_provider = litellm_params.get("custom_llm_provider") if not _skip_a2a_provider_routing: - a2a_provider_config = A2AProviderConfigManager.get_provider_config( + a2a_provider_config: Final = A2AProviderConfigManager.get_provider_config( custom_llm_provider=custom_llm_provider, model=litellm_params.get("model"), ) if a2a_provider_config is not None: - verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider} (streaming)") + verbose_logger.info("A2A: Using provider config for %s (streaming)", custom_llm_provider) async for chunk in a2a_provider_config.handle_streaming( request_id=request_id, @@ -198,20 +199,20 @@ class A2ACompletionBridgeHandler: return # Extract message from params - message = params.get("message", {}) + message: Final = params.get("message", {}) # Create streaming context - ctx = A2AStreamingContext( + ctx: Final = A2AStreamingContext( request_id=request_id, input_message=message, ) # Transform A2A message to OpenAI format - openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) + openai_messages: Final = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message) # Get completion params custom_llm_provider = litellm_params.get("custom_llm_provider") - model = litellm_params.get("model", "agent") + model: Final = litellm_params.get("model", "agent") # Build full model string if provider specified # Skip prepending if model already starts with the provider prefix @@ -220,17 +221,17 @@ class A2ACompletionBridgeHandler: else: full_model = model - verbose_logger.info(f"A2A completion bridge streaming: model={full_model}, api_base={api_base}") + verbose_logger.info("A2A completion bridge streaming: model=%s, api_base=%s", full_model, api_base) # Build completion params dict - completion_params: Dict[str, Any] = { + completion_params: Final[dict[str, Any]] = { "model": full_model, "messages": openai_messages, "api_base": api_base, "stream": True, } # Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.) - litellm_params_to_add = { + litellm_params_to_add: Final = { k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider") and k not in _AGENT_ONLY_PARAMS @@ -252,11 +253,11 @@ class A2ACompletionBridgeHandler: ) # 1. Emit initial task event (kind: "task", status: "submitted") - task_event = A2ACompletionBridgeTransformation.create_task_event(ctx) + task_event: Final = A2ACompletionBridgeTransformation.create_task_event(ctx) yield task_event # 2. Emit status update (kind: "status-update", status: "working") - working_event = A2ACompletionBridgeTransformation.create_status_update_event( + working_event: Final = A2ACompletionBridgeTransformation.create_status_update_event( ctx=ctx, state="working", final=False, @@ -265,7 +266,7 @@ class A2ACompletionBridgeHandler: yield working_event # Call litellm.acompletion with streaming - response = await litellm.acompletion(**completion_params) + response: Final = await litellm.acompletion(**completion_params) # 3. Accumulate content and emit artifact update accumulated_text = "" @@ -285,31 +286,33 @@ class A2ACompletionBridgeHandler: # Emit artifact update with accumulated content if accumulated_text: - artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event( + artifact_event: Final = A2ACompletionBridgeTransformation.create_artifact_update_event( ctx=ctx, text=accumulated_text, ) yield artifact_event # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event = A2ACompletionBridgeTransformation.create_status_update_event( + completed_event: Final = A2ACompletionBridgeTransformation.create_status_update_event( ctx=ctx, state="completed", final=True, ) yield completed_event - verbose_logger.info(f"A2A completion bridge streaming completed: request_id={request_id}, chunks={chunk_count}") + verbose_logger.info( + "A2A completion bridge streaming completed: request_id=%s, chunks=%s", request_id, chunk_count + ) # Convenience functions that delegate to the class methods async def handle_a2a_completion( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], - api_base: Optional[str] = None, - agent_extra_headers: Optional[Dict[str, str]] = None, -) -> Dict[str, Any]: + params: dict[str, Any], + litellm_params: dict[str, Any], + api_base: str | None = None, + agent_extra_headers: dict[str, str] | None = None, +) -> dict[str, Any]: """Convenience function for non-streaming A2A completion.""" return await A2ACompletionBridgeHandler.handle_non_streaming( request_id=request_id, @@ -322,11 +325,11 @@ async def handle_a2a_completion( async def handle_a2a_completion_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], - api_base: Optional[str] = None, - agent_extra_headers: Optional[Dict[str, str]] = None, -) -> AsyncIterator[Dict[str, Any]]: + params: dict[str, Any], + litellm_params: dict[str, Any], + api_base: str | None = None, + agent_extra_headers: dict[str, str] | None = None, +) -> AsyncIterator[dict[str, Any]]: """Convenience function for streaming A2A completion.""" async for chunk in A2ACompletionBridgeHandler.handle_streaming( request_id=request_id, diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index b32963dd6fb..15cf77708f9 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -18,7 +18,7 @@ A2A Streaming Events: """ from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any, Final from uuid import uuid4 from litellm._logging import verbose_logger @@ -30,7 +30,7 @@ class A2AStreamingContext: Tracks task_id, context_id, and message accumulation. """ - def __init__(self, request_id: str, input_message: Dict[str, Any]): + def __init__(self, request_id: str, input_message: dict[str, Any]): self.request_id = request_id self.task_id = str(uuid4()) self.context_id = str(uuid4()) @@ -46,9 +46,9 @@ class A2ACompletionBridgeTransformation: """ @staticmethod - def _extract_text_from_a2a_parts(parts: List[Dict[str, Any]]) -> str: + def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str: """Extract text from A2A parts (with or without explicit ``kind``).""" - content_parts: List[str] = [] + content_parts: Final[list[str]] = [] for part in parts: if not isinstance(part, dict): continue @@ -62,35 +62,35 @@ class A2ACompletionBridgeTransformation: @staticmethod def get_forward_metadata( - a2a_message: Dict[str, Any], - params: Optional[Dict[str, Any]] = None, - ) -> Optional[Dict[str, Any]]: + a2a_message: dict[str, Any], + params: dict[str, Any] | None = None, + ) -> dict[str, Any] | None: """ Merge A2A metadata from MessageSendParams and the message for downstream providers. Forwarded once on the LangGraph run payload (``metadata``), not duplicated on each input message — see ``apply_forward_metadata_to_completion_params``. """ - merged: Dict[str, Any] = {} + merged: Final[dict[str, Any]] = {} if params and isinstance(params.get("metadata"), dict): merged.update(params["metadata"]) - message_metadata = a2a_message.get("metadata") + message_metadata: Final = a2a_message.get("metadata") if isinstance(message_metadata, dict): merged.update(message_metadata) return merged or None @staticmethod def apply_forward_metadata_to_completion_params( - completion_params: Dict[str, Any], - a2a_message: Dict[str, Any], - params: Optional[Dict[str, Any]] = None, + completion_params: dict[str, Any], + a2a_message: dict[str, Any], + params: dict[str, Any] | None = None, ) -> None: """ Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph). Uses ``extra_body`` so we do not collide with LiteLLM's spend-log ``metadata`` kwarg. """ - forward_metadata = A2ACompletionBridgeTransformation.get_forward_metadata( + forward_metadata: Final = A2ACompletionBridgeTransformation.get_forward_metadata( a2a_message=a2a_message, params=params, ) @@ -103,18 +103,18 @@ class A2ACompletionBridgeTransformation: # Layer client-supplied A2A metadata under any agent-owner-configured # ``extra_body.metadata`` so the configured keys remain authoritative # and an A2A caller cannot overwrite server-set run metadata. - existing_metadata = extra_body.get("metadata") - existing_dict: Dict[str, Any] = existing_metadata if isinstance(existing_metadata, dict) else {} - merged_metadata: Dict[str, Any] = {**forward_metadata, **existing_dict} + existing_metadata: Final = extra_body.get("metadata") + existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {} + merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict} extra_body = {**extra_body, "metadata": merged_metadata} completion_params["extra_body"] = extra_body - verbose_logger.debug(f"A2A -> completion forward metadata keys={list(forward_metadata.keys())}") + verbose_logger.debug("A2A -> completion forward metadata keys=%s", list(forward_metadata.keys())) @staticmethod def a2a_message_to_openai_messages( - a2a_message: Dict[str, Any], - ) -> List[Dict[str, Any]]: + a2a_message: dict[str, Any], + ) -> list[dict[str, Any]]: """ Transform an A2A message to OpenAI message format. @@ -124,7 +124,7 @@ class A2ACompletionBridgeTransformation: Returns: List of OpenAI-format messages """ - role = a2a_message.get("role", "user") + role: Final = a2a_message.get("role", "user") parts = a2a_message.get("parts", []) # Map A2A roles to OpenAI roles @@ -139,21 +139,23 @@ class A2ACompletionBridgeTransformation: if not isinstance(parts, list): parts = [] - content = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts) + content: Final = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts) # Do not attach A2A message.metadata here — the completion bridge forwards it # once at run level via extra_body.metadata (LangGraph POST /runs/wait shape). - openai_message: Dict[str, Any] = {"role": openai_role, "content": content} + openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content} - verbose_logger.debug(f"A2A -> OpenAI transform: role={role} -> {openai_role}, content_length={len(content)}") + verbose_logger.debug( + "A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content) + ) return [openai_message] @staticmethod def openai_response_to_a2a_response( response: Any, - request_id: Optional[str] = None, - ) -> Dict[str, Any]: + request_id: str | None = None, + ) -> dict[str, Any]: """ Transform a LiteLLM ModelResponse to A2A SendMessageResponse format. @@ -167,12 +169,12 @@ class A2ACompletionBridgeTransformation: # Extract content from response content = "" if hasattr(response, "choices") and response.choices: - choice = response.choices[0] + choice: Final = response.choices[0] if hasattr(choice, "message") and choice.message: content = choice.message.content or "" # Build A2A message - a2a_message = { + a2a_message: Final = { "kind": "message", "role": "agent", "parts": [{"kind": "text", "text": content}], @@ -180,13 +182,13 @@ class A2ACompletionBridgeTransformation: } # Build A2A response - a2a_response = { + a2a_response: Final = { "jsonrpc": "2.0", "id": request_id, "result": a2a_message, } - verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}") + verbose_logger.debug("OpenAI -> A2A transform: content_length=%s", len(content)) return a2a_response @@ -198,7 +200,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def create_task_event( ctx: A2AStreamingContext, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Create the initial task event with status 'submitted'. @@ -232,8 +234,8 @@ class A2ACompletionBridgeTransformation: ctx: A2AStreamingContext, state: str, final: bool = False, - message_text: Optional[str] = None, - ) -> Dict[str, Any]: + message_text: str | None = None, + ) -> dict[str, Any]: """ Create a status update event. @@ -243,7 +245,7 @@ class A2ACompletionBridgeTransformation: final: Whether this is the final event message_text: Optional message text for 'working' status """ - status: Dict[str, Any] = { + status: Final[dict[str, Any]] = { "state": state, "timestamp": A2ACompletionBridgeTransformation._get_timestamp(), } @@ -275,7 +277,7 @@ class A2ACompletionBridgeTransformation: def create_artifact_update_event( ctx: A2AStreamingContext, text: str, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Create an artifact update event with content. diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index f04edf2579b..6b2541bc8a9 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -12,16 +12,8 @@ Provides standalone functions with @client decorator for LiteLLM logging integra import asyncio import datetime import uuid -from typing import ( - TYPE_CHECKING, - Any, - AsyncIterator, - Coroutine, - Dict, - Optional, - Union, - cast, -) +from collections.abc import AsyncIterator, Coroutine +from typing import TYPE_CHECKING, Any, Final, Optional, cast import litellm from litellm._logging import verbose_logger, verbose_proxy_logger @@ -83,11 +75,11 @@ from litellm.a2a_protocol.exception_mapping_utils import ( from litellm.a2a_protocol.exceptions import A2ALocalhostURLError # Use our custom resolver instead of the default A2A SDK resolver -A2ACardResolver = LiteLLMA2ACardResolver +A2ACardResolver: Final = LiteLLMA2ACardResolver def _set_usage_on_logging_obj( - kwargs: Dict[str, Any], + kwargs: dict[str, Any], prompt_tokens: int, completion_tokens: int, ) -> None: @@ -99,9 +91,9 @@ def _set_usage_on_logging_obj( prompt_tokens: Number of input tokens completion_tokens: Number of output tokens """ - litellm_logging_obj = kwargs.get("litellm_logging_obj") + litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") if litellm_logging_obj is not None: - usage = litellm.Usage( + usage: Final = litellm.Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, @@ -110,7 +102,7 @@ def _set_usage_on_logging_obj( def _set_agent_id_on_logging_obj( - kwargs: Dict[str, Any], + kwargs: dict[str, Any], agent_id: str | None, ) -> None: """ @@ -123,13 +115,13 @@ def _set_agent_id_on_logging_obj( if agent_id is None: return - litellm_logging_obj = kwargs.get("litellm_logging_obj") + litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") if litellm_logging_obj is not None: # Set agent_id directly on model_call_details (same pattern as custom_llm_provider) litellm_logging_obj.model_call_details["agent_id"] = agent_id -_A2A_COST_PARAM_KEYS = ("cost_per_query", "input_cost_per_token", "output_cost_per_token") +_A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output_cost_per_token") def _set_litellm_params_on_logging_obj( @@ -144,7 +136,7 @@ def _set_litellm_params_on_logging_obj( litellm_params already carries metadata / proxy_server_request / user-key context, so merge the pricing keys in rather than replacing the dict. """ - logging_obj = kwargs.get("litellm_logging_obj") + logging_obj: Final = kwargs.get("litellm_logging_obj") if logging_obj is None: return @@ -152,11 +144,11 @@ def _set_litellm_params_on_logging_obj( if not cost_params: return - existing = logging_obj.model_call_details.get("litellm_params") or {} + existing: Final = logging_obj.model_call_details.get("litellm_params") or {} logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params} -def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: +def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str: """ Extract agent info and set model/custom_llm_provider for cost tracking. @@ -165,17 +157,17 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: """ agent_name = "unknown" - agent_card = _get_a2a_client_agent_card(a2a_client) + agent_card: Final = _get_a2a_client_agent_card(a2a_client) if agent_card is not None: agent_name = getattr(agent_card, "name", "unknown") or "unknown" # Build model string - model = f"a2a_agent/{agent_name}" - custom_llm_provider = "a2a_agent" + model: Final = f"a2a_agent/{agent_name}" + custom_llm_provider: Final = "a2a_agent" # Set on litellm_logging_obj if available (for standard logging payload) - litellm_logging_obj = kwargs.get("litellm_logging_obj") + litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") if litellm_logging_obj is not None: litellm_logging_obj.model = model litellm_logging_obj.custom_llm_provider = custom_llm_provider @@ -199,15 +191,15 @@ async def _send_message_via_completion_bridge( request: "SendMessageRequest", custom_llm_provider: str, api_base: str | None, - litellm_params: Dict[str, Any], - agent_extra_headers: Dict[str, str] | None = None, + litellm_params: dict[str, Any], + agent_extra_headers: dict[str, str] | None = None, ) -> LiteLLMSendMessageResponse: """ Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore). Requires request; api_base is optional for providers that derive endpoint from model. """ - verbose_logger.info(f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}") + verbose_logger.info("A2A using completion bridge: provider=%s, api_base=%s", custom_llm_provider, api_base) from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, @@ -215,7 +207,7 @@ async def _send_message_via_completion_bridge( params = request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) - response_dict = await A2ACompletionBridgeHandler.handle_non_streaming( + response_dict: Final = await A2ACompletionBridgeHandler.handle_non_streaming( request_id=str(request.id), params=params, litellm_params=litellm_params, @@ -233,18 +225,18 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - pb_request = _a2a_conversions.to_core_send_message_request(request) + pb_request: Final = _a2a_conversions.to_core_send_message_request(request) last_event = None async for event in a2a_client.send_message(pb_request): last_event = event if last_event is None: raise RuntimeError("A2A send_message failed: no response received from agent.") - stream_compat = _a2a_conversions.to_compat_stream_response( + stream_compat: Final = _a2a_conversions.to_compat_stream_response( last_event, request_id=request.id, ) - result = stream_compat.result + result: Final = stream_compat.result if not isinstance(result, (Message, Task)): raise RuntimeError( "A2A send_message failed: non-streaming message/send expects the " @@ -308,7 +300,7 @@ async def _stream_messages( "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - pb_request = _a2a_conversions.to_core_send_message_request(request) + pb_request: Final = _a2a_conversions.to_core_send_message_request(request) async for event in a2a_client.send_message(pb_request): compat_chunk = _a2a_conversions.to_compat_stream_response( event, @@ -370,9 +362,9 @@ async def asend_message( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendMessageRequest"] = None, api_base: str | None = None, - litellm_params: Dict[str, Any] | None = None, + litellm_params: dict[str, Any] | None = None, agent_id: str | None = None, - agent_extra_headers: Dict[str, str] | None = None, + agent_extra_headers: dict[str, str] | None = None, **kwargs: Any, ) -> LiteLLMSendMessageResponse: """ @@ -428,9 +420,9 @@ async def asend_message( ``` """ litellm_params = litellm_params or {} - logging_obj = kwargs.get("litellm_logging_obj") + logging_obj: Final = kwargs.get("litellm_logging_obj") trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None - custom_llm_provider = litellm_params.get("custom_llm_provider") + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") # Route through completion bridge if custom_llm_provider is set if custom_llm_provider: @@ -453,7 +445,7 @@ async def asend_message( if api_base is None: raise ValueError("Either a2a_client or api_base is required for standard A2A flow") trace_id = trace_id or str(uuid.uuid4()) - extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} + extra_headers: Final[dict[str, str]] = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: extra_headers["X-LiteLLM-Agent-Id"] = agent_id # Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones) @@ -464,15 +456,15 @@ async def asend_message( # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None - agent_name = _get_a2a_model_info(a2a_client, kwargs) + agent_name: Final = _get_a2a_model_info(a2a_client, kwargs) - verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}") + verbose_logger.info("A2A send_message request_id=%s, agent=%s", request.id, agent_name) # Get agent card URL for localhost retry logic - agent_card = _get_a2a_client_agent_card(a2a_client) - card_url = get_agent_card_url(agent_card) if agent_card else None + agent_card: Final = _get_a2a_client_agent_card(a2a_client) + card_url: Final = get_agent_card_url(agent_card) if agent_card else None - a2a_response = await _execute_a2a_send_with_retry( + a2a_response: Final = await _execute_a2a_send_with_retry( a2a_client=a2a_client, request=request, agent_card=agent_card, @@ -481,13 +473,13 @@ async def asend_message( agent_name=agent_name, ) - verbose_logger.info(f"A2A send_message completed, request_id={request.id}") + verbose_logger.info("A2A send_message completed, request_id=%s", request.id) # Wrap in LiteLLM response type for _hidden_params support - response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id)) + response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id)) # Calculate token usage from request and response - response_dict = a2a_response.model_dump(mode="json", exclude_none=True) + response_dict: Final = a2a_response.model_dump(mode="json", exclude_none=True) ( prompt_tokens, completion_tokens, @@ -518,7 +510,7 @@ def send_message( a2a_client: "A2AClientType", request: "SendMessageRequest", **kwargs: Any, -) -> Union[LiteLLMSendMessageResponse, Coroutine[Any, Any, LiteLLMSendMessageResponse]]: +) -> LiteLLMSendMessageResponse | Coroutine[Any, Any, LiteLLMSendMessageResponse]: """ Sync: Send a message to an A2A agent. @@ -547,15 +539,15 @@ def _build_streaming_logging_obj( request: "SendStreamingMessageRequest", agent_name: str, agent_id: str | None, - litellm_params: Dict[str, Any] | None, - metadata: Dict[str, Any] | None, - proxy_server_request: Dict[str, Any] | None, + litellm_params: dict[str, Any] | None, + metadata: dict[str, Any] | None, + proxy_server_request: dict[str, Any] | None, ) -> Logging: """Build logging object for streaming A2A requests.""" - start_time = datetime.datetime.now() - model = f"a2a_agent/{agent_name}" + start_time: Final = datetime.datetime.now() + model: Final = f"a2a_agent/{agent_name}" - logging_obj = Logging( + logging_obj: Final = Logging( model=model, messages=[{"role": "user", "content": "streaming-request"}], stream=False, @@ -572,7 +564,7 @@ def _build_streaming_logging_obj( if agent_id: logging_obj.model_call_details["agent_id"] = agent_id - _litellm_params = litellm_params.copy() if litellm_params else {} + _litellm_params: Final = litellm_params.copy() if litellm_params else {} if metadata: _litellm_params["metadata"] = metadata if proxy_server_request: @@ -590,11 +582,11 @@ async def asend_message_streaming( a2a_client: Optional["A2AClientType"] = None, request: Optional["SendStreamingMessageRequest"] = None, api_base: str | None = None, - litellm_params: Dict[str, Any] | None = None, + litellm_params: dict[str, Any] | None = None, agent_id: str | None = None, - metadata: Dict[str, Any] | None = None, - proxy_server_request: Dict[str, Any] | None = None, - agent_extra_headers: Dict[str, str] | None = None, + metadata: dict[str, Any] | None = None, + proxy_server_request: dict[str, Any] | None = None, + agent_extra_headers: dict[str, str] | None = None, **kwargs: object, ) -> AsyncIterator[Any]: """ @@ -635,7 +627,7 @@ async def asend_message_streaming( ``` """ litellm_params = litellm_params or {} - custom_llm_provider = litellm_params.get("custom_llm_provider") + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") # Route through completion bridge if custom_llm_provider is set if custom_llm_provider: @@ -643,14 +635,14 @@ async def asend_message_streaming( raise ValueError("request is required for completion bridge") # api_base is optional for providers that derive endpoint from model (e.g., bedrock/agentcore) - verbose_logger.info(f"A2A streaming using completion bridge: provider={custom_llm_provider}") + verbose_logger.info("A2A streaming using completion bridge: provider=%s", custom_llm_provider) from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2ACompletionBridgeHandler, ) # Extract params from request - params = ( + params: Final = ( request.params.model_dump(mode="json") if hasattr(request.params, "model_dump") else dict(request.params) ) @@ -667,15 +659,15 @@ async def asend_message_streaming( if request is None: raise ValueError("request is required") - _raw_logging_obj = kwargs.get("litellm_logging_obj") + _raw_logging_obj: Final = kwargs.get("litellm_logging_obj") logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None if a2a_client is None: if api_base is None: raise ValueError("Either a2a_client or api_base is required for standard A2A flow") - logging_trace_id = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None - trace_id = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4())) - extra_headers: dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} + logging_trace_id: Final = getattr(logging_obj, "litellm_trace_id", None) if logging_obj else None + trace_id: Final = logging_trace_id or (str(request.id) if request.id else str(uuid.uuid4())) + extra_headers: Final[dict[str, str]] = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: extra_headers["X-LiteLLM-Agent-Id"] = agent_id if agent_extra_headers: @@ -688,7 +680,7 @@ async def asend_message_streaming( assert a2a_client is not None - agent_name = _get_a2a_model_info(a2a_client, kwargs) + agent_name: Final = _get_a2a_model_info(a2a_client, kwargs) if logging_obj is None: logging_obj = _build_streaming_logging_obj( @@ -700,12 +692,12 @@ async def asend_message_streaming( proxy_server_request=proxy_server_request, ) - verbose_logger.info(f"A2A send_message_streaming request_id={request.id}, agent={agent_name}") + verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name) - agent_card = _get_a2a_client_agent_card(a2a_client) - card_url = get_agent_card_url(agent_card) if agent_card else None + agent_card: Final = _get_a2a_client_agent_card(a2a_client) + card_url: Final = get_agent_card_url(agent_card) if agent_card else None - stream = _execute_a2a_stream_with_retry( + stream: Final = _execute_a2a_stream_with_retry( a2a_client=a2a_client, request=request, agent_card=agent_card, @@ -728,7 +720,7 @@ async def asend_message_streaming( async def create_a2a_client( base_url: str, timeout: float = DEFAULT_A2A_AGENT_TIMEOUT, - extra_headers: Dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, streaming: bool = False, ) -> "A2AClientType": """ @@ -762,7 +754,7 @@ async def create_a2a_client( "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - verbose_logger.info(f"Creating A2A client for {base_url}") + verbose_logger.info("Creating A2A client for %s", base_url) # Use get_async_httpx_client with per-agent params so that different agents # (with different extra_headers) get separate cached clients. The params @@ -772,21 +764,21 @@ async def create_a2a_client( # Only pass params that AsyncHTTPHandler.__init__ accepts (e.g. timeout). # Use "disable_aiohttp_transport" key for cache-key-only data (it's # filtered out before reaching the constructor). - _client_params: dict = {"timeout": timeout} + _client_params: Final[dict] = {"timeout": timeout} if extra_headers: # Encode headers into a cache-key-only param so each unique header # set produces a distinct cache key. _client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items())) - _async_handler = get_async_httpx_client( + _async_handler: Final = get_async_httpx_client( llm_provider=httpxSpecialProvider.A2AProvider, params=_client_params, ) - httpx_client = _async_handler.client + httpx_client: Final = _async_handler.client if extra_headers: httpx_client.headers.update(extra_headers) - verbose_proxy_logger.debug(f"A2A client created with extra_headers={list(extra_headers.keys())}") + verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys())) - a2a_client = await create_client( # pyright: ignore[reportOptionalCall] + a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall] base_url, client_config=ClientConfig( # pyright: ignore[reportOptionalCall] httpx_client=httpx_client, @@ -797,11 +789,11 @@ async def create_a2a_client( # the configured httpx client (with this agent's trace-id/auth headers) without # excavating a2a-sdk private internals. a2a_client._litellm_httpx_client = httpx_client # type: ignore[attr-defined] - agent_card = getattr(a2a_client, "_card", None) + agent_card: Final = getattr(a2a_client, "_card", None) if agent_card is not None: a2a_client._litellm_agent_card = agent_card # type: ignore[attr-defined] - verbose_logger.info(f"A2A client created for {base_url}") + verbose_logger.info("A2A client created for %s", base_url) return a2a_client @@ -809,7 +801,7 @@ async def create_a2a_client( async def aget_agent_card( base_url: str, timeout: float = DEFAULT_A2A_AGENT_TIMEOUT, - extra_headers: Dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, ) -> "AgentCard": """ Fetch the agent card from an A2A agent. @@ -827,20 +819,20 @@ async def aget_agent_card( "The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk" ) - verbose_logger.info(f"Fetching agent card from {base_url}") + verbose_logger.info("Fetching agent card from %s", base_url) # Use LiteLLM's cached httpx client - http_handler = get_async_httpx_client( + http_handler: Final = get_async_httpx_client( llm_provider=httpxSpecialProvider.A2A, params={"timeout": timeout}, ) - httpx_client = http_handler.client + httpx_client: Final = http_handler.client - resolver = A2ACardResolver( + resolver: Final = A2ACardResolver( httpx_client=httpx_client, base_url=base_url, ) - agent_card = await resolver.get_agent_card() + agent_card: Final = await resolver.get_agent_card() - verbose_logger.info(f"Fetched agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}") + verbose_logger.info("Fetched agent card: %s", agent_card.name if hasattr(agent_card, "name") else "unknown") return agent_card diff --git a/litellm/a2a_protocol/providers/__init__.py b/litellm/a2a_protocol/providers/__init__.py index a21fa5f8f5e..8f16fcf15c8 100644 --- a/litellm/a2a_protocol/providers/__init__.py +++ b/litellm/a2a_protocol/providers/__init__.py @@ -7,4 +7,4 @@ This module contains provider-specific implementations for the A2A protocol. from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager -__all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"] +__all__ = ["A2AProviderConfigManager", "BaseA2AProviderConfig"] diff --git a/litellm/a2a_protocol/providers/base.py b/litellm/a2a_protocol/providers/base.py index 3ac1cb47fc8..5a5eff8cf35 100644 --- a/litellm/a2a_protocol/providers/base.py +++ b/litellm/a2a_protocol/providers/base.py @@ -3,7 +3,8 @@ Base configuration for A2A protocol providers. """ from abc import ABC, abstractmethod -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any class BaseA2AProviderConfig(ABC): @@ -18,10 +19,10 @@ class BaseA2AProviderConfig(ABC): async def handle_non_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Handle non-streaming A2A request. @@ -34,16 +35,15 @@ class BaseA2AProviderConfig(ABC): Returns: A2A SendMessageResponse dict """ - pass @abstractmethod async def handle_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: """ Handle streaming A2A request. diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py index f624aa393ed..2b37c0c4906 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/config.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/config.py @@ -2,7 +2,8 @@ Bedrock AgentCore A2A provider configuration. """ -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any, Final from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.bedrock_agentcore.handler import ( @@ -22,12 +23,12 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Handle non-streaming request to AgentCore A2A agent.""" - litellm_params = kwargs.get("litellm_params") + litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: raise ValueError( "litellm_params is required for BedrockAgentCoreA2AConfig (must contain model with AgentCore ARN)" @@ -42,12 +43,12 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig): async def handle_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: """Handle streaming request to AgentCore A2A agent.""" - litellm_params = kwargs.get("litellm_params") + litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: raise ValueError( "litellm_params is required for BedrockAgentCoreA2AConfig (must contain model with AgentCore ARN)" diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index c613b68668f..db57072ca38 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -6,7 +6,8 @@ completion bridge that would otherwise strip the envelope. """ import json -from typing import Any, AsyncIterator, Dict, Optional, cast +from collections.abc import AsyncIterator +from typing import Any, Final, cast from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import ( @@ -27,10 +28,10 @@ class BedrockAgentCoreA2AHandler: @staticmethod async def handle_non_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: + params: dict[str, Any], + litellm_params: dict[str, Any], + agent_extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: """ Handle non-streaming A2A request to AgentCore. @@ -52,31 +53,31 @@ class BedrockAgentCoreA2AHandler: agent_extra_headers=agent_extra_headers, ) - verbose_logger.info(f"BedrockAgentCore A2A: Sending non-streaming request to {url}") + verbose_logger.info("BedrockAgentCore A2A: Sending non-streaming request to %s", url) - client = get_async_httpx_client( + client: Final = get_async_httpx_client( llm_provider=cast(Any, httpxSpecialProvider.A2AProvider), ) - response = await client.post( + response: Final = await client.post( url, headers=headers, data=body, ) response.raise_for_status() - response_data = response.json() + response_data: Final = response.json() if "error" in response_data: - verbose_logger.warning(f"BedrockAgentCore A2A: Agent returned error: {response_data['error']}") + verbose_logger.warning("BedrockAgentCore A2A: Agent returned error: %s", response_data["error"]) return response_data @staticmethod async def handle_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> AsyncIterator[Dict[str, Any]]: + params: dict[str, Any], + litellm_params: dict[str, Any], + agent_extra_headers: dict[str, str] | None = None, + ) -> AsyncIterator[dict[str, Any]]: """ Handle streaming A2A request to AgentCore. @@ -99,12 +100,12 @@ class BedrockAgentCoreA2AHandler: agent_extra_headers=agent_extra_headers, ) - verbose_logger.info(f"BedrockAgentCore A2A: Sending streaming request to {url}") + verbose_logger.info("BedrockAgentCore A2A: Sending streaming request to %s", url) - client = get_async_httpx_client( + client: Final = get_async_httpx_client( llm_provider=cast(Any, httpxSpecialProvider.A2AProvider), ) - response = await client.post( + response: Final = await client.post( url, headers=headers, data=body, @@ -113,15 +114,15 @@ class BedrockAgentCoreA2AHandler: response.raise_for_status() # Check content type — AgentCore may return JSON instead of SSE - content_type = response.headers.get("content-type", "").lower() + content_type: Final = response.headers.get("content-type", "").lower() if "application/json" in content_type: # Single JSON response fallback (not SSE) verbose_logger.debug( "BedrockAgentCore A2A streaming: received JSON instead of SSE, yielding as single event" ) - response_body = await response.aread() - response_data = json.loads(response_body) + response_body: Final = await response.aread() + response_data: Final = json.loads(response_body) yield response_data else: # SSE stream — parse data: lines diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 091a13ccea5..32252711997 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -6,7 +6,8 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT). """ import json -from typing import Any, AsyncIterator, Dict, Mapping, Optional, Tuple +from collections.abc import AsyncIterator, Mapping +from typing import Any, Final from litellm._logging import verbose_logger from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig @@ -22,21 +23,21 @@ from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreCo # ``runtimeSessionId`` / ``runtimeUserId`` in the agent's ``litellm_params``; # ``authorization`` is set by the AgentCore signer (JWT or SigV4); ``host`` and # the ``x-amz-*`` family are owned by SigV4 itself. -_RESERVED_EXACT_HEADERS = frozenset( +_RESERVED_EXACT_HEADERS: Final = frozenset( { "authorization", "host", } ) -_RESERVED_PREFIX_HEADERS: Tuple[str, ...] = ( +_RESERVED_PREFIX_HEADERS: Final[tuple[str, ...]] = ( "x-amzn-bedrock-agentcore-runtime-", "x-amz-", ) def _filter_reserved_headers( - agent_extra_headers: Optional[Mapping[str, str]], -) -> Optional[Dict[str, str]]: + agent_extra_headers: Mapping[str, str] | None, +) -> dict[str, str] | None: """ Strip reserved AWS / AgentCore headers from caller-supplied ``agent_extra_headers`` before they are merged into the signed request. @@ -46,8 +47,8 @@ def _filter_reserved_headers( if not agent_extra_headers: return None - filtered: Dict[str, str] = {} - dropped: list = [] + filtered: Final[dict[str, str]] = {} + dropped: Final[list] = [] for k, v in agent_extra_headers.items(): k_lower = k.lower() if k_lower in _RESERVED_EXACT_HEADERS or any(k_lower.startswith(prefix) for prefix in _RESERVED_PREFIX_HEADERS): @@ -76,12 +77,12 @@ class BedrockAgentCoreA2ATransformation: @staticmethod def get_url_and_signed_request( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], + params: dict[str, Any], + litellm_params: dict[str, Any], method: str = "message/send", stream: bool = False, - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> Tuple[str, dict, bytes]: + agent_extra_headers: dict[str, str] | None = None, + ) -> tuple[str, dict, bytes]: """ Build the AgentCore URL, construct a JSON-RPC envelope, and sign the request. @@ -106,19 +107,19 @@ class BedrockAgentCoreA2ATransformation: """ # Extract model and strip the "bedrock/" prefix # "bedrock/agentcore/arn:aws:..." → "agentcore/arn:aws:..." - model = litellm_params.get("model", "") + model: Final = litellm_params.get("model", "") if model.startswith("bedrock/"): agentcore_model = model[len("bedrock/") :] else: agentcore_model = model # Build optional_params from litellm_params (everything except model and custom_llm_provider) - optional_params = {k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider")} + optional_params: Final = {k: v for k, v in litellm_params.items() if k not in ("model", "custom_llm_provider")} - agentcore_config = AmazonAgentCoreConfig() + agentcore_config: Final = AmazonAgentCoreConfig() # Derive URL from ARN - url = agentcore_config.get_complete_url( + url: Final = agentcore_config.get_complete_url( api_base=optional_params.get("api_base"), api_key=optional_params.get("api_key"), model=agentcore_model, @@ -128,7 +129,7 @@ class BedrockAgentCoreA2ATransformation: ) # Construct JSON-RPC 2.0 envelope - json_rpc_body = { + json_rpc_body: Final = { "jsonrpc": "2.0", "method": method, "id": request_id, @@ -137,17 +138,17 @@ class BedrockAgentCoreA2ATransformation: # Set required AgentCore session headers (normally set by transform_request, # which we skip because it also builds {"prompt": "..."}) - headers: dict = {} - session_id = agentcore_config._get_runtime_session_id(optional_params) + headers: Final[dict] = {} + session_id: Final = agentcore_config._get_runtime_session_id(optional_params) headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = session_id - runtime_user_id = agentcore_config._get_runtime_user_id(optional_params) + runtime_user_id: Final = agentcore_config._get_runtime_user_id(optional_params) if runtime_user_id: headers["X-Amzn-Bedrock-AgentCore-Runtime-User-Id"] = runtime_user_id # Merge per-request agent headers before signing so SigV4 covers them. # Reserved headers are stripped first to prevent client-controlled values # from spoofing the AgentCore runtime identity / SigV4 metadata. - safe_extra_headers = _filter_reserved_headers(agent_extra_headers) + safe_extra_headers: Final = _filter_reserved_headers(agent_extra_headers) if safe_extra_headers: headers.update(safe_extra_headers) @@ -169,7 +170,7 @@ class BedrockAgentCoreA2ATransformation: return url, signed_headers, signed_body @staticmethod - async def parse_sse_events(response: Any) -> AsyncIterator[Dict[str, Any]]: + async def parse_sse_events(response: Any) -> AsyncIterator[dict[str, Any]]: """ Parse SSE events from an httpx streaming response. @@ -194,5 +195,5 @@ class BedrockAgentCoreA2ATransformation: event = json.loads(data_str) yield event except json.JSONDecodeError: - verbose_logger.debug(f"BedrockAgentCore A2A: Skipping non-JSON SSE line: {data_str[:100]}") + verbose_logger.debug("BedrockAgentCore A2A: Skipping non-JSON SSE line: %s", data_str[:100]) continue diff --git a/litellm/a2a_protocol/providers/config_manager.py b/litellm/a2a_protocol/providers/config_manager.py index a421afec184..2eab2adb1ba 100644 --- a/litellm/a2a_protocol/providers/config_manager.py +++ b/litellm/a2a_protocol/providers/config_manager.py @@ -4,8 +4,6 @@ A2A Provider Config Manager. Manages provider-specific configurations for A2A protocol. """ -from typing import Optional - from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig @@ -18,9 +16,9 @@ class A2AProviderConfigManager: @staticmethod def get_provider_config( - custom_llm_provider: Optional[str], - model: Optional[str] = None, - ) -> Optional[BaseA2AProviderConfig]: + custom_llm_provider: str | None, + model: str | None = None, + ) -> BaseA2AProviderConfig | None: """ Get the provider configuration for a given custom_llm_provider. diff --git a/litellm/a2a_protocol/providers/langflow/config.py b/litellm/a2a_protocol/providers/langflow/config.py index 9edaf151c71..54d403f88c0 100644 --- a/litellm/a2a_protocol/providers/langflow/config.py +++ b/litellm/a2a_protocol/providers/langflow/config.py @@ -1,4 +1,5 @@ -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any from litellm.a2a_protocol.litellm_completion_bridge.handler import ( A2A_USER_API_KEY_HASH_PARAM, @@ -15,10 +16,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( @@ -38,10 +39,10 @@ class LangFlowA2AConfig(BaseA2AProviderConfig): async def handle_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: litellm_params = kwargs.get("litellm_params") if not litellm_params: raise ValueError( diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py index 8e9cd6fc87e..078e0633e04 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/__init__.py @@ -13,4 +13,4 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( PydanticAITransformation, ) -__all__ = ["PydanticAIHandler", "PydanticAITransformation", "PydanticAIProviderConfig"] +__all__ = ["PydanticAIHandler", "PydanticAIProviderConfig", "PydanticAITransformation"] diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py index 6f067aecd2b..b7546e1a2a1 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/config.py @@ -2,7 +2,8 @@ Pydantic AI provider configuration. """ -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAIHandler @@ -19,10 +20,10 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs: Any, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Handle non-streaming request to Pydantic AI agent.""" if api_base is None: raise ValueError("api_base is required for PydanticAIProviderConfig") @@ -37,10 +38,10 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig): async def handle_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: """Handle streaming request with fake streaming.""" if not api_base: raise ValueError("api_base is required for Pydantic AI agents") diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py index 352005ff549..c083c0267f7 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/handler.py @@ -5,7 +5,8 @@ Pydantic AI agents follow A2A protocol but don't support streaming natively. This handler provides fake streaming by converting non-streaming responses into streaming chunks. """ -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any, Final from litellm._logging import verbose_logger from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import ( @@ -25,11 +26,11 @@ class PydanticAIHandler: @staticmethod async def handle_non_streaming( request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, timeout: float = 60.0, - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: + agent_extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: """ Handle non-streaming request to Pydantic AI agent. @@ -46,10 +47,10 @@ class PydanticAIHandler: """ if api_base is None: raise ValueError("api_base is required for Pydantic AI agents") - verbose_logger.info(f"Pydantic AI: Routing to Pydantic AI agent at {api_base}") + verbose_logger.info("Pydantic AI: Routing to Pydantic AI agent at %s", api_base) # Send request directly to Pydantic AI agent - response_data = await PydanticAITransformation.send_non_streaming_request( + response_data: Final = await PydanticAITransformation.send_non_streaming_request( api_base=api_base, request_id=request_id, params=params, @@ -62,13 +63,13 @@ class PydanticAIHandler: @staticmethod async def handle_streaming( request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, timeout: float = 60.0, chunk_size: int = 50, delay_ms: int = 10, - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> AsyncIterator[Dict[str, Any]]: + agent_extra_headers: dict[str, str] | None = None, + ) -> AsyncIterator[dict[str, Any]]: """ Handle streaming request to Pydantic AI agent with fake streaming. @@ -91,10 +92,10 @@ class PydanticAIHandler: """ if api_base is None: raise ValueError("api_base is required for Pydantic AI agents") - verbose_logger.info(f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}") + verbose_logger.info("Pydantic AI: Faking streaming for Pydantic AI agent at %s", api_base) # Get raw task response first (not the transformed A2A format) - raw_response = await PydanticAITransformation.send_and_get_raw_response( + raw_response: Final = await PydanticAITransformation.send_and_get_raw_response( api_base=api_base, request_id=request_id, params=params, diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index b9943d83c8a..339da998d56 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,7 +6,8 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from typing import Any, AsyncIterator, Dict, Optional, cast +from collections.abc import AsyncIterator +from typing import Any, Final, cast from uuid import uuid4 from litellm._logging import verbose_logger @@ -48,7 +49,7 @@ class PydanticAITransformation: return obj @staticmethod - def _params_to_dict(params: Any) -> Dict[str, Any]: + def _params_to_dict(params: Any) -> dict[str, Any]: """ Convert params to a dict, handling Pydantic models. @@ -78,8 +79,8 @@ class PydanticAITransformation: request_id: str, max_attempts: int = 30, poll_interval: float = 0.5, - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: + agent_extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: """ Poll for task completion using tasks/get method. @@ -117,7 +118,7 @@ class PydanticAITransformation: status = result.get("status", {}) state = status.get("state", "") - verbose_logger.debug(f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}") + verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state) if state == "completed": return poll_data @@ -134,8 +135,8 @@ class PydanticAITransformation: request_id: str, params: Any, timeout: float = 60.0, - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: + agent_extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -162,7 +163,7 @@ class PydanticAITransformation: params_dict["message"]["kind"] = "message" # Build A2A JSON-RPC request using message/send method for FastA2A compatibility - a2a_request = { + a2a_request: Final = { "jsonrpc": "2.0", "id": request_id, "method": "message/send", @@ -170,16 +171,16 @@ class PydanticAITransformation: } # FastA2A uses root endpoint (/) not /messages - endpoint = api_base.rstrip("/") + endpoint: Final = api_base.rstrip("/") - verbose_logger.info(f"Pydantic AI: Sending non-streaming request to {endpoint}") + verbose_logger.info("Pydantic AI: Sending non-streaming request to %s", endpoint) # Send request to Pydantic AI agent using shared async HTTP client - client = get_async_httpx_client( + client: Final = get_async_httpx_client( llm_provider=cast(Any, "pydantic_ai_agent"), params={"timeout": timeout}, ) - response = await client.post( + response: Final = await client.post( endpoint, json=a2a_request, headers={ @@ -191,15 +192,15 @@ class PydanticAITransformation: response_data = response.json() # Check if task is already completed - result = response_data.get("result", {}) - status = result.get("status", {}) - state = status.get("state", "") + result: Final = response_data.get("result", {}) + status: Final = result.get("status", {}) + state: Final = status.get("state", "") if state != "completed": # Need to poll for completion - task_id = result.get("id") + task_id: Final = result.get("id") if task_id: - verbose_logger.info(f"Pydantic AI: Task {task_id} submitted, polling for completion...") + verbose_logger.info("Pydantic AI: Task %s submitted, polling for completion...", task_id) response_data = await PydanticAITransformation._poll_for_completion( client=client, endpoint=endpoint, @@ -208,7 +209,7 @@ class PydanticAITransformation: agent_extra_headers=agent_extra_headers, ) - verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}") + verbose_logger.info("Pydantic AI: Received completed response for request_id=%s", request_id) return response_data @@ -218,8 +219,8 @@ class PydanticAITransformation: request_id: str, params: Any, timeout: float = 60.0, - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: + agent_extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -234,7 +235,7 @@ class PydanticAITransformation: Standard A2A non-streaming response format with message """ # Get raw task response - raw_response = await PydanticAITransformation._send_and_poll_raw( + raw_response: Final = await PydanticAITransformation._send_and_poll_raw( api_base=api_base, request_id=request_id, params=params, @@ -254,8 +255,8 @@ class PydanticAITransformation: request_id: str, params: Any, timeout: float = 60.0, - agent_extra_headers: Optional[Dict[str, str]] = None, - ) -> Dict[str, Any]: + agent_extra_headers: dict[str, str] | None = None, + ) -> dict[str, Any]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -281,9 +282,9 @@ class PydanticAITransformation: @staticmethod def _transform_to_a2a_response( - response_data: Dict[str, Any], + response_data: dict[str, Any], request_id: str, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Transform Pydantic AI task response to standard A2A non-streaming format. @@ -312,7 +313,7 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Build standard A2A message - a2a_message = { + a2a_message: Final = { "kind": "message", "role": "agent", "parts": parts if parts else [{"kind": "text", "text": full_text}], @@ -327,7 +328,7 @@ class PydanticAITransformation: } @staticmethod - def _extract_response_text(response_data: Dict[str, Any]) -> tuple[str, str, list]: + def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]: """ Extract response text from completed task response. @@ -341,10 +342,10 @@ class PydanticAITransformation: Returns: Tuple of (full_text, message_id, parts) """ - result = response_data.get("result", {}) + result: Final = response_data.get("result", {}) # Try to extract from artifacts first (preferred for results) - artifacts = result.get("artifacts", []) + artifacts: Final = result.get("artifacts", []) if artifacts: for artifact in artifacts: parts = artifact.get("parts", []) @@ -355,7 +356,7 @@ class PydanticAITransformation: return text, str(uuid4()), parts # Fall back to history - get the last agent message - history = result.get("history", []) + history: Final = result.get("history", []) for msg in reversed(history): if msg.get("role") == "agent": parts = msg.get("parts", []) @@ -368,7 +369,7 @@ class PydanticAITransformation: return full_text, message_id, parts # Fall back to message field (original format) - message = result.get("message", {}) + message: Final = result.get("message", {}) if message: parts = message.get("parts", []) message_id = message.get("messageId", str(uuid4())) @@ -382,11 +383,11 @@ class PydanticAITransformation: @staticmethod async def fake_streaming_from_response( - response_data: Dict[str, Any], + response_data: dict[str, Any], request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: """ Convert a non-streaming A2A response into fake streaming chunks. @@ -409,8 +410,8 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history - result = response_data.get("result", {}) - history = result.get("history", []) + result: Final = response_data.get("result", {}) + history: Final = result.get("history", []) input_message = {} for msg in history: if msg.get("role") == "user": @@ -418,14 +419,14 @@ class PydanticAITransformation: break # Generate IDs for streaming events - task_id = str(uuid4()) - context_id = str(uuid4()) - artifact_id = str(uuid4()) - input_message_id = input_message.get("messageId", str(uuid4())) + task_id: Final = str(uuid4()) + context_id: Final = str(uuid4()) + artifact_id: Final = str(uuid4()) + input_message_id: Final = input_message.get("messageId", str(uuid4())) # 1. Emit initial task event (kind: "task", status: "submitted") # Format matches A2ACompletionBridgeTransformation.create_task_event - task_event = { + task_event: Final = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -451,7 +452,7 @@ class PydanticAITransformation: # 2. Emit status update (kind: "status-update", status: "working") # Format matches A2ACompletionBridgeTransformation.create_status_update_event - working_event = { + working_event: Final = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -502,7 +503,7 @@ class PydanticAITransformation: await asyncio.sleep(delay_ms / 1000.0) # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event = { + completed_event: Final = { "jsonrpc": "2.0", "id": request_id, "result": { @@ -517,4 +518,4 @@ class PydanticAITransformation: } yield completed_event - verbose_logger.info(f"Pydantic AI: Fake streaming completed for request_id={request_id}") + verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id) diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py index dbd4a0558f7..ca84d3e07b4 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -2,7 +2,8 @@ A2A provider configuration for IBM watsonx Orchestrate (WXO). """ -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any, Final from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( @@ -16,12 +17,12 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): async def handle_non_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs: Any, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Handle a non-streaming A2A request via WXO runs API.""" - litellm_params = kwargs.get("litellm_params") + litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: raise ValueError( "litellm_params is required for WatsonxOrchestrateA2AConfig " @@ -36,12 +37,12 @@ class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): async def handle_streaming( self, request_id: str, - params: Dict[str, Any], - api_base: Optional[str] = None, + params: dict[str, Any], + api_base: str | None = None, **kwargs: Any, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: """Handle a streaming A2A request via WXO streaming runs API.""" - litellm_params = kwargs.get("litellm_params") + litellm_params: Final = kwargs.get("litellm_params") if not litellm_params: raise ValueError( "litellm_params is required for WatsonxOrchestrateA2AConfig " diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 07235c1118c..bb29700cd46 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -6,7 +6,8 @@ import asyncio import hashlib import json import time -from typing import Any, AsyncIterator, Dict, NamedTuple, Optional, Tuple, cast +from collections.abc import AsyncIterator +from typing import Any, Final, NamedTuple, cast import httpx @@ -20,11 +21,11 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.llms.custom_http import httpxSpecialProvider -_IBM_CLOUD_IAM_URL = "https://iam.cloud.ibm.com/identity/token" -_POLL_INTERVAL_S = 2.0 -_MAX_POLL_ATTEMPTS = 90 -_TOKEN_CACHE_TTL_BUFFER_S = 60 -_token_cache: Dict[str, Tuple[str, float]] = {} +_IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token" +_POLL_INTERVAL_S: Final = 2.0 +_MAX_POLL_ATTEMPTS: Final = 90 +_TOKEN_CACHE_TTL_BUFFER_S: Final = 60 +_token_cache: Final[dict[str, tuple[str, float]]] = {} class WXORequestParams(NamedTuple): @@ -32,9 +33,9 @@ class WXORequestParams(NamedTuple): instance_id: str wxo_agent_id: str api_key: str - username: Optional[str] + username: str | None auth_mode: str - thread_id: Optional[str] + thread_id: str | None class WatsonxOrchestrateHandler: @@ -50,16 +51,16 @@ class WatsonxOrchestrateHandler: auth_mode: str, cp4d_host: str, api_key: str, - username: Optional[str], + username: str | None, ) -> str: - material = f"{auth_mode}:{cp4d_host}:{username or ''}:{api_key}" + material: Final = f"{auth_mode}:{cp4d_host}:{username or ''}:{api_key}" return hashlib.sha256(material.encode()).hexdigest() @staticmethod - def _cp4d_token_ttl_seconds(expiration: Any, now_wall: Optional[float] = None) -> int: + def _cp4d_token_ttl_seconds(expiration: Any, now_wall: float | None = None) -> int: # CP4D returns expiration as absolute Unix epoch seconds, not a duration. - expires_at = int(expiration) - wall = now_wall if now_wall is not None else time.time() + expires_at: Final = int(expiration) + wall: Final = now_wall if now_wall is not None else time.time() return max(expires_at - int(wall), 0) @staticmethod @@ -67,12 +68,12 @@ class WatsonxOrchestrateHandler: cp4d_host: str, auth_mode: str, api_key: str, - username: Optional[str] = None, - client: Optional[AsyncHTTPHandler] = None, + username: str | None = None, + client: AsyncHTTPHandler | None = None, ) -> str: - cache_key = WatsonxOrchestrateHandler._token_cache_key(auth_mode, cp4d_host, api_key, username) - now = time.monotonic() - cached = _token_cache.get(cache_key) + cache_key: Final = WatsonxOrchestrateHandler._token_cache_key(auth_mode, cp4d_host, api_key, username) + now: Final = time.monotonic() + cached: Final = _token_cache.get(cache_key) if cached and cached[1] > now: return cached[0] @@ -95,7 +96,7 @@ class WatsonxOrchestrateHandler: else: if not username: raise ValueError("'username' is required in litellm_params when auth_mode='cp4d'") - token_url = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize" + token_url: Final = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize" response = await client.post( token_url, json={"username": username, "api_key": api_key}, @@ -104,13 +105,13 @@ class WatsonxOrchestrateHandler: response.raise_for_status() payload = response.json() token = str(payload["token"]) - expiration = payload.get("expiration") + expiration: Final = payload.get("expiration") if expiration is None: ttl_s = 3600 else: ttl_s = WatsonxOrchestrateHandler._cp4d_token_ttl_seconds(expiration) - expires_at = now + max(ttl_s - _TOKEN_CACHE_TTL_BUFFER_S, 0) + expires_at: Final = now + max(ttl_s - _TOKEN_CACHE_TTL_BUFFER_S, 0) _token_cache[cache_key] = (token, expires_at) for stale_key, (_, stale_expires_at) in list(_token_cache.items()): if stale_expires_at <= now: @@ -121,20 +122,20 @@ class WatsonxOrchestrateHandler: async def _poll_run( base_url: str, run_id: str, - auth_headers: Dict[str, str], + auth_headers: dict[str, str], client: AsyncHTTPHandler, max_attempts: int = _MAX_POLL_ATTEMPTS, interval_s: float = _POLL_INTERVAL_S, - ) -> Dict[str, Any]: - url = f"{base_url}/v1/orchestrate/runs/{run_id}" + ) -> dict[str, Any]: + url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}" for attempt in range(max_attempts): await asyncio.sleep(interval_s) response = await client.get(url, headers=auth_headers) response.raise_for_status() - result: Dict[str, Any] = response.json() + result: dict[str, Any] = response.json() status = result.get("status", "") - verbose_logger.debug(f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'") + verbose_logger.debug("WXO: Poll %s/%s run='%s' status='%s'", attempt + 1, max_attempts, run_id, status) if status in WatsonxOrchestrateTransformation.TERMINAL_STATES: return result @@ -144,14 +145,14 @@ class WatsonxOrchestrateHandler: @staticmethod async def _get_successful_run_data( - run_data: Dict[str, Any], + run_data: dict[str, Any], base_url: str, - auth_headers: Dict[str, str], + auth_headers: dict[str, str], client: AsyncHTTPHandler, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: status = run_data.get("status", "") if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES: - run_id = run_data.get("run_id") or run_data.get("id") or "" + run_id: Final = run_data.get("run_id") or run_data.get("id") or "" if not run_id: raise ValueError(f"WXO: No run_id in response: {run_data}") run_data = await WatsonxOrchestrateHandler._poll_run( @@ -186,11 +187,11 @@ class WatsonxOrchestrateHandler: return accumulated_text @staticmethod - def _extract_litellm_params(litellm_params: Dict[str, Any]) -> WXORequestParams: - cp4d_host = litellm_params.get("cp4d_host") or "" - instance_id = litellm_params.get("instance_id") or "" - wxo_agent_id = litellm_params.get("wxo_agent_id") or "" - api_key = litellm_params.get("api_key") or "" + def _extract_litellm_params(litellm_params: dict[str, Any]) -> WXORequestParams: + cp4d_host: Final = litellm_params.get("cp4d_host") or "" + instance_id: Final = litellm_params.get("instance_id") or "" + wxo_agent_id: Final = litellm_params.get("wxo_agent_id") or "" + api_key: Final = litellm_params.get("api_key") or "" if not cp4d_host: raise ValueError("'cp4d_host' is required in litellm_params for WXO agents") @@ -214,38 +215,38 @@ class WatsonxOrchestrateHandler: @staticmethod async def handle_non_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], - ) -> Dict[str, Any]: - wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) + params: dict[str, Any], + litellm_params: dict[str, Any], + ) -> dict[str, Any]: + wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client = WatsonxOrchestrateHandler._http_client(timeout=90.0) - token = await WatsonxOrchestrateHandler._get_bearer_token( + client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0) + token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, api_key=wxo.api_key, username=wxo.username, client=client, ) - base_url = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) - auth_headers = { + base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) + auth_headers: Final = { "Authorization": f"Bearer {token}", "Content-Type": "application/json", "Accept": "application/json", } - text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) - body = WatsonxOrchestrateTransformation.build_wxo_run_body( + text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) + body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body( wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id ) - run_response = await client.post( + run_response: Final = await client.post( f"{base_url}/v1/orchestrate/runs", json=body, headers=auth_headers, ) run_response.raise_for_status() - run_data: Dict[str, Any] = run_response.json() + run_data: dict[str, Any] = run_response.json() run_data = await WatsonxOrchestrateHandler._get_successful_run_data( run_data=run_data, @@ -254,40 +255,40 @@ class WatsonxOrchestrateHandler: client=client, ) - response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data) + response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(run_data) return WatsonxOrchestrateTransformation.build_a2a_message_response(request_id=request_id, text=response_text) @staticmethod async def handle_streaming( request_id: str, - params: Dict[str, Any], - litellm_params: Dict[str, Any], + params: dict[str, Any], + litellm_params: dict[str, Any], chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[Dict[str, Any]]: - wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) + ) -> AsyncIterator[dict[str, Any]]: + wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) - client = WatsonxOrchestrateHandler._http_client(timeout=120.0) - token = await WatsonxOrchestrateHandler._get_bearer_token( + client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0) + token: Final = await WatsonxOrchestrateHandler._get_bearer_token( cp4d_host=wxo.cp4d_host, auth_mode=wxo.auth_mode, api_key=wxo.api_key, username=wxo.username, client=client, ) - base_url = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) - auth_headers = { + base_url: Final = WatsonxOrchestrateTransformation.get_api_base_url(wxo.cp4d_host, wxo.instance_id) + auth_headers: Final = { "Authorization": f"Bearer {token}", "Content-Type": "application/json", "Accept": "text/event-stream, application/json", } - text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) - body = WatsonxOrchestrateTransformation.build_wxo_run_body( + text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) + body: Final = WatsonxOrchestrateTransformation.build_wxo_run_body( wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id ) try: - response = await client.post( + response: Final = await client.post( f"{base_url}/v1/orchestrate/runs/stream", json=body, headers=auth_headers, @@ -296,8 +297,8 @@ class WatsonxOrchestrateHandler: response.raise_for_status() except httpx.TransportError as exc: verbose_logger.warning( - f"WXO: Streaming request failed before a run was submitted " - f"({exc!r}), falling back to non-streaming + fake streaming", + "WXO: Streaming request failed before a run was submitted (%r), falling back to non-streaming + fake streaming", + exc, exc_info=True, ) result = await WatsonxOrchestrateHandler.handle_non_streaming( @@ -305,7 +306,7 @@ class WatsonxOrchestrateHandler: params=params, litellm_params=litellm_params, ) - response_text = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result) + response_text: Final = WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response(result) async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( text=response_text, request_id=request_id, @@ -315,9 +316,9 @@ class WatsonxOrchestrateHandler: yield chunk return - content_type = response.headers.get("content-type", "").lower() + content_type: Final = response.headers.get("content-type", "").lower() if "text/event-stream" not in content_type: - response_body = await response.aread() + response_body: Final = await response.aread() result = json.loads(response_body) result = await WatsonxOrchestrateHandler._get_successful_run_data( run_data=result, diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py index c9bda822aae..3748d8043cc 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -8,7 +8,8 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model: """ import asyncio -from typing import Any, AsyncIterator, Dict, Optional +from collections.abc import AsyncIterator +from typing import Any, Final from uuid import uuid4 from litellm._logging import verbose_logger @@ -28,15 +29,15 @@ class WatsonxOrchestrateTransformation: return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}" @staticmethod - def extract_text_from_a2a_params(params: Dict[str, Any]) -> str: + def extract_text_from_a2a_params(params: dict[str, Any]) -> str: """ Extract user message text from A2A MessageSendParams. A2A format: params.message.parts[*] where part.kind == "text" """ - message = params.get("message", {}) - parts = message.get("parts", []) - texts = [] + message: Final = params.get("message", {}) + parts: Final = message.get("parts", []) + texts: Final = [] for part in parts: if not isinstance(part, dict): continue @@ -49,10 +50,10 @@ class WatsonxOrchestrateTransformation: def build_wxo_run_body( wxo_agent_id: str, text: str, - thread_id: Optional[str] = None, - ) -> Dict[str, Any]: + thread_id: str | None = None, + ) -> dict[str, Any]: """Build the WXO POST /v1/orchestrate/runs request body.""" - body: Dict[str, Any] = { + body: Final[dict[str, Any]] = { "agent_id": wxo_agent_id, "message": { "role": "user", @@ -95,19 +96,19 @@ class WatsonxOrchestrateTransformation: pass # Tertiary: results as a raw string - results = result.get("results") + results: Final = result.get("results") if results and isinstance(results, str): return results return "" @staticmethod - def extract_text_from_a2a_message_response(a2a_response: Dict[str, Any]) -> str: - result = a2a_response.get("result") + def extract_text_from_a2a_message_response(a2a_response: dict[str, Any]) -> str: + result: Final = a2a_response.get("result") if not isinstance(result, dict): verbose_logger.warning("WXO: A2A response missing result object") return "" - parts = result.get("parts") + parts: Final = result.get("parts") if not isinstance(parts, list): verbose_logger.warning("WXO: A2A result has no parts list") return "" @@ -118,7 +119,7 @@ class WatsonxOrchestrateTransformation: return "" @staticmethod - def build_a2a_message_response(request_id: str, text: str) -> Dict[str, Any]: + def build_a2a_message_response(request_id: str, text: str) -> dict[str, Any]: """ Build a standard A2A non-streaming SendMessageResponse (kind=message). """ @@ -139,7 +140,7 @@ class WatsonxOrchestrateTransformation: request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[Dict[str, Any]]: + ) -> AsyncIterator[dict[str, Any]]: """ Emit standard A2A streaming events from a completed text response. @@ -149,9 +150,9 @@ class WatsonxOrchestrateTransformation: 3. artifact-update chunks 4. status-update (kind="status-update", state="completed", final=True) """ - task_id = str(uuid4()) - context_id = str(uuid4()) - artifact_id = str(uuid4()) + task_id: Final = str(uuid4()) + context_id: Final = str(uuid4()) + artifact_id: Final = str(uuid4()) # 1. Task submitted yield { @@ -180,7 +181,7 @@ class WatsonxOrchestrateTransformation: await asyncio.sleep(delay_ms / 1000.0) # 3. Artifact chunks (always emit at least one chunk, even for empty text) - text_to_chunk = text or "" + text_to_chunk: Final = text or "" for i in range(0, max(len(text_to_chunk), 1), chunk_size): chunk_text = text_to_chunk[i : i + chunk_size] is_last = (i + chunk_size) >= max(len(text_to_chunk), 1) @@ -213,4 +214,4 @@ class WatsonxOrchestrateTransformation: }, } - verbose_logger.debug(f"WXO: Fake streaming completed for request_id={request_id}") + verbose_logger.debug("WXO: Fake streaming completed for request_id=%s", request_id) diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 1ef174a5eee..413691f233d 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -3,8 +3,9 @@ A2A Streaming Iterator with token tracking and logging support. """ import asyncio +from collections.abc import AsyncIterator from datetime import datetime -from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_logger @@ -37,16 +38,16 @@ class A2AStreamingIterator: self.start_time = datetime.now() # Collect chunks for token counting - self.chunks: List[Any] = [] - self.collected_text_parts: List[str] = [] - self.final_chunk: Optional[Any] = None + self.chunks: list[Any] = [] + self.collected_text_parts: list[str] = [] + self.final_chunk: Any | None = None def __aiter__(self): return self async def __anext__(self) -> "SendStreamingMessageResponse": try: - chunk = await self.stream.__anext__() + chunk: Final = await self.stream.__anext__() # Store chunk self.chunks.append(chunk) @@ -70,8 +71,8 @@ class A2AStreamingIterator: def _collect_text_from_chunk(self, chunk: Any) -> None: """Extract text from a streaming chunk and add to collected parts.""" try: - chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} - text = A2ARequestUtils.extract_text_from_response(chunk_dict) + chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + text: Final = A2ARequestUtils.extract_text_from_response(chunk_dict) if text: self.collected_text_parts.append(text) except Exception: @@ -80,10 +81,10 @@ class A2AStreamingIterator: def _is_completed_chunk(self, chunk: Any) -> bool: """Check if chunk indicates stream completion.""" try: - chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} - result = chunk_dict.get("result", {}) + chunk_dict: Final = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {} + result: Final = chunk_dict.get("result", {}) if isinstance(result, dict): - status = result.get("status", {}) + status: Final = result.get("status", {}) if isinstance(status, dict): return status.get("state") == "completed" except Exception: @@ -93,21 +94,21 @@ class A2AStreamingIterator: async def _handle_stream_complete(self) -> None: """Handle logging and token counting when stream completes.""" try: - end_time = datetime.now() + end_time: Final = datetime.now() # Calculate tokens from collected text - input_message = A2ARequestUtils.get_input_message_from_request(self.request) - input_text = A2ARequestUtils.extract_text_from_message(input_message) - prompt_tokens = A2ARequestUtils.count_tokens(input_text) + input_message: Final = A2ARequestUtils.get_input_message_from_request(self.request) + input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) + prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) # Use the last (most complete) text from chunks - output_text = self.collected_text_parts[-1] if self.collected_text_parts else "" - completion_tokens = A2ARequestUtils.count_tokens(output_text) + output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else "" + completion_tokens: Final = A2ARequestUtils.count_tokens(output_text) - total_tokens = prompt_tokens + completion_tokens + total_tokens: Final = prompt_tokens + completion_tokens # Create usage object - usage = litellm.Usage( + usage: Final = litellm.Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens, @@ -119,11 +120,11 @@ class A2AStreamingIterator: self.logging_obj.model_call_details["stream"] = False # Calculate cost using A2ACostCalculator - response_cost = A2ACostCalculator.calculate_a2a_cost(self.logging_obj) + response_cost: Final = A2ACostCalculator.calculate_a2a_cost(self.logging_obj) self.logging_obj.model_call_details["response_cost"] = response_cost # Build result for logging - result = self._build_logging_result(usage) + result: Final = self._build_logging_result(usage) # Call success handlers - they will build standard_logging_object asyncio.create_task( @@ -137,17 +138,19 @@ class A2AStreamingIterator: ) verbose_logger.info( - f"A2A streaming completed: prompt_tokens={prompt_tokens}, " - f"completion_tokens={completion_tokens}, total_tokens={total_tokens}, " - f"response_cost={response_cost}" + "A2A streaming completed: prompt_tokens=%s, completion_tokens=%s, total_tokens=%s, response_cost=%s", + prompt_tokens, + completion_tokens, + total_tokens, + response_cost, ) except Exception as e: - verbose_logger.debug(f"Error in A2A streaming completion handler: {e}") + verbose_logger.debug("Error in A2A streaming completion handler: %s", e) - def _build_logging_result(self, usage: litellm.Usage) -> Dict[str, Any]: + def _build_logging_result(self, usage: litellm.Usage) -> dict[str, Any]: """Build a result dict for logging.""" - result: Dict[str, Any] = { + result: Final[dict[str, Any]] = { "id": getattr(self.request, "id", "unknown"), "jsonrpc": "2.0", "usage": (usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)), @@ -156,7 +159,7 @@ class A2AStreamingIterator: # Add final chunk result if available if self.final_chunk: try: - chunk_dict = self.final_chunk.model_dump(mode="json", exclude_none=True) + chunk_dict: Final = self.final_chunk.model_dump(mode="json", exclude_none=True) result["result"] = chunk_dict.get("result", {}) except Exception: pass diff --git a/litellm/a2a_protocol/utils.py b/litellm/a2a_protocol/utils.py index ce5a168c3ac..f2e61f66105 100644 --- a/litellm/a2a_protocol/utils.py +++ b/litellm/a2a_protocol/utils.py @@ -2,7 +2,7 @@ Utility functions for A2A protocol. """ -from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union +from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_logger @@ -34,7 +34,7 @@ class A2ARequestUtils: else: parts = getattr(message, "parts", []) or [] - text_parts: List[str] = [] + text_parts: Final[list[str]] = [] for part in parts: if isinstance(part, dict): if part.get("kind") == "text": @@ -46,7 +46,7 @@ class A2ARequestUtils: return " ".join(text_parts) @staticmethod - def extract_text_from_response(response_dict: Dict[str, Any]) -> str: + def extract_text_from_response(response_dict: dict[str, Any]) -> str: """ Extract text content from A2A response result. @@ -56,7 +56,7 @@ class A2ARequestUtils: Returns: Text from response message parts """ - result = response_dict.get("result", {}) + result: Final = response_dict.get("result", {}) if not isinstance(result, dict): return "" @@ -66,12 +66,12 @@ class A2ARequestUtils: if result.get("kind") == "message": return A2ARequestUtils.extract_text_from_message(result) - message = result.get("message", {}) + message: Final = result.get("message", {}) return A2ARequestUtils.extract_text_from_message(message) @staticmethod def get_input_message_from_request( - request: "Union[SendMessageRequest, SendStreamingMessageRequest]", + request: "SendMessageRequest | SendStreamingMessageRequest", ) -> Any: """ Extract the input message from an A2A request. @@ -82,7 +82,7 @@ class A2ARequestUtils: Returns: The message object/dict or None """ - params = getattr(request, "params", None) + params: Final = getattr(request, "params", None) if params is None: return None return getattr(params, "message", None) @@ -108,9 +108,9 @@ class A2ARequestUtils: @staticmethod def calculate_usage_from_request_response( - request: "Union[SendMessageRequest, SendStreamingMessageRequest]", - response_dict: Dict[str, Any], - ) -> Tuple[int, int, int]: + request: "SendMessageRequest | SendStreamingMessageRequest", + response_dict: dict[str, Any], + ) -> tuple[int, int, int]: """ Calculate token usage from A2A request and response. @@ -128,14 +128,14 @@ class A2ARequestUtils: input_message = A2ARequestUtils.get_input_message_from_request(request) if input_message is not None and hasattr(input_message, "model_dump"): input_message = input_message.model_dump(mode="json") - input_text = A2ARequestUtils.extract_text_from_message(input_message) - prompt_tokens = A2ARequestUtils.count_tokens(input_text) + input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) + prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) # Count output tokens - output_text = A2ARequestUtils.extract_text_from_response(response_dict) - completion_tokens = A2ARequestUtils.count_tokens(output_text) + output_text: Final = A2ARequestUtils.extract_text_from_response(response_dict) + completion_tokens: Final = A2ARequestUtils.count_tokens(output_text) - total_tokens = prompt_tokens + completion_tokens + total_tokens: Final = prompt_tokens + completion_tokens return prompt_tokens, completion_tokens, total_tokens @@ -145,5 +145,5 @@ def extract_text_from_a2a_message(message: Any) -> str: return A2ARequestUtils.extract_text_from_message(message) -def extract_text_from_a2a_response(response_dict: Dict[str, Any]) -> str: +def extract_text_from_a2a_response(response_dict: dict[str, Any]) -> str: return A2ARequestUtils.extract_text_from_response(response_dict) diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index d0082498b09..abce47c191e 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -25,14 +25,14 @@ Environment Variables: import json import os from importlib.resources import files -from typing import Dict, List, Optional, Set +from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import verbose_logger # Cache for the loaded configuration -_BETA_HEADERS_CONFIG: Optional[Dict] = None +_BETA_HEADERS_CONFIG: dict | None = None class GetAnthropicBetaHeadersConfig: @@ -44,15 +44,15 @@ class GetAnthropicBetaHeadersConfig: """ @staticmethod - def load_local_beta_headers_config() -> Dict: + def load_local_beta_headers_config() -> dict: """Load the local backup beta headers config bundled with the package.""" try: - content = json.loads( + content: Final = json.loads( files("litellm").joinpath("anthropic_beta_headers_config.json").read_text(encoding="utf-8") ) return content except Exception as e: - verbose_logger.error(f"Failed to load local beta headers config: {e}") + verbose_logger.error("Failed to load local beta headers config: %s", e) # Return empty config as fallback return { "anthropic": {}, @@ -80,14 +80,14 @@ class GetAnthropicBetaHeadersConfig: return False # Check for at least one provider key - provider_keys = [ + provider_keys: Final = [ "anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai", ] - has_provider = any(key in fetched_config for key in provider_keys) + has_provider: Final = any(key in fetched_config for key in provider_keys) if not has_provider: verbose_logger.warning( @@ -114,7 +114,7 @@ class GetAnthropicBetaHeadersConfig: Returns the parsed JSON dict. Raises on network/parse errors (caller is expected to handle). """ - response = httpx.get(url, timeout=timeout) + response: Final = httpx.get(url, timeout=timeout) response.raise_for_status() return response.json() @@ -139,7 +139,7 @@ def get_beta_headers_config(url: str) -> dict: return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config() try: - content = GetAnthropicBetaHeadersConfig.fetch_remote_beta_headers_config(url) + content: Final = GetAnthropicBetaHeadersConfig.fetch_remote_beta_headers_config(url) except Exception as e: verbose_logger.warning( "LiteLLM: Failed to fetch remote beta headers config from %s: %s. Falling back to local backup.", @@ -159,7 +159,7 @@ def get_beta_headers_config(url: str) -> dict: return content -def _load_beta_headers_config() -> Dict: +def _load_beta_headers_config() -> dict: """ Load the beta headers configuration. Uses caching to avoid repeated fetches/file reads. @@ -183,7 +183,7 @@ def _load_beta_headers_config() -> Dict: return _BETA_HEADERS_CONFIG -def reload_beta_headers_config() -> Dict: +def reload_beta_headers_config() -> dict: """ Force reload the beta headers configuration from source (remote or local). Clears the cache and fetches fresh configuration. @@ -207,15 +207,15 @@ def get_provider_name(provider: str) -> str: Returns: Canonical provider name """ - config = _load_beta_headers_config() - aliases = config.get("provider_aliases", {}) + config: Final = _load_beta_headers_config() + aliases: Final = config.get("provider_aliases", {}) return aliases.get(provider, provider) def filter_and_transform_beta_headers( - beta_headers: List[str], + beta_headers: list[str], provider: str, -) -> List[str]: +) -> list[str]: """ Filter and transform beta headers based on provider's mapping configuration. @@ -234,20 +234,22 @@ def filter_and_transform_beta_headers( if not beta_headers: return [] - config = _load_beta_headers_config() + config: Final = _load_beta_headers_config() provider = get_provider_name(provider) # Get the header mapping for this provider - provider_mapping = config.get(provider, {}) + provider_mapping: Final = config.get(provider, {}) - filtered_headers: Set[str] = set() + filtered_headers: Final[set[str]] = set() for header in beta_headers: header = header.strip() # Check if header is in the mapping if header not in provider_mapping: - verbose_logger.debug(f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)") + verbose_logger.debug( + "Dropping unknown beta header '%s' for provider '%s' (not in mapping)", header, provider + ) continue # Get the mapped header value @@ -255,7 +257,7 @@ def filter_and_transform_beta_headers( # Skip if header is unsupported (null value) if mapped_header is None: - verbose_logger.debug(f"Dropping unsupported beta header '{header}' for provider '{provider}'") + verbose_logger.debug("Dropping unsupported beta header '%s' for provider '%s'", header, provider) continue # Add the mapped header @@ -278,9 +280,9 @@ def is_beta_header_supported( Returns: True if the header is in the mapping with a non-null value, False otherwise """ - config = _load_beta_headers_config() + config: Final = _load_beta_headers_config() provider = get_provider_name(provider) - provider_mapping = config.get(provider, {}) + provider_mapping: Final = config.get(provider, {}) # Header is supported if it's in the mapping and has a non-null value return beta_header in provider_mapping and provider_mapping[beta_header] is not None @@ -289,7 +291,7 @@ def is_beta_header_supported( def get_provider_beta_header( anthropic_beta_header: str, provider: str, -) -> Optional[str]: +) -> str | None: """ Get the provider-specific beta header name for a given Anthropic beta header. @@ -302,11 +304,11 @@ def get_provider_beta_header( Returns: The provider-specific header name if supported, or None if unsupported/unknown """ - config = _load_beta_headers_config() + config: Final = _load_beta_headers_config() provider = get_provider_name(provider) # Get the header mapping for this provider - provider_mapping = config.get(provider, {}) + provider_mapping: Final = config.get(provider, {}) # Check if header is in the mapping if anthropic_beta_header not in provider_mapping: @@ -331,15 +333,15 @@ def update_headers_with_filtered_beta( Returns: Updated headers dict """ - existing_beta = headers.get("anthropic-beta") + existing_beta: Final = headers.get("anthropic-beta") if not existing_beta: return headers # Parse existing beta headers - beta_values = [b.strip() for b in existing_beta.split(",") if b.strip()] + beta_values: Final = [b.strip() for b in existing_beta.split(",") if b.strip()] # Filter and transform based on provider - filtered_beta_values = filter_and_transform_beta_headers( + filtered_beta_values: Final = filter_and_transform_beta_headers( beta_headers=beta_values, provider=provider, ) @@ -373,11 +375,11 @@ def update_request_with_filtered_beta( """ headers = update_headers_with_filtered_beta(headers=headers, provider=provider) - existing_body_betas = request_data.get("anthropic_beta") + existing_body_betas: Final = request_data.get("anthropic_beta") if not existing_body_betas: return headers, request_data - filtered_body_betas = filter_and_transform_beta_headers( + filtered_body_betas: Final = filter_and_transform_beta_headers( beta_headers=existing_body_betas, provider=provider, ) @@ -390,7 +392,7 @@ def update_request_with_filtered_beta( return headers, request_data -def get_unsupported_headers(provider: str) -> List[str]: +def get_unsupported_headers(provider: str) -> list[str]: """ Get all beta headers that are unsupported by a provider (have null values in mapping). @@ -400,9 +402,9 @@ def get_unsupported_headers(provider: str) -> List[str]: Returns: List of unsupported Anthropic beta header names """ - config = _load_beta_headers_config() + config: Final = _load_beta_headers_config() provider = get_provider_name(provider) - provider_mapping = config.get(provider, {}) + provider_mapping: Final = config.get(provider, {}) # Return headers with null values return [header for header, value in provider_mapping.items() if value is None] diff --git a/litellm/anthropic_interface/exceptions/__init__.py b/litellm/anthropic_interface/exceptions/__init__.py index 875b09e3da3..7f2de0e60dc 100644 --- a/litellm/anthropic_interface/exceptions/__init__.py +++ b/litellm/anthropic_interface/exceptions/__init__.py @@ -11,9 +11,9 @@ from .exceptions import ( ) __all__ = [ - "AnthropicErrorType", + "ANTHROPIC_ERROR_TYPE_MAP", "AnthropicErrorDetail", "AnthropicErrorResponse", - "ANTHROPIC_ERROR_TYPE_MAP", + "AnthropicErrorType", "AnthropicExceptionMapping", ] diff --git a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py index b4ec83517ee..ad8be8ec40e 100644 --- a/litellm/anthropic_interface/exceptions/exception_mapping_utils.py +++ b/litellm/anthropic_interface/exceptions/exception_mapping_utils.py @@ -4,14 +4,15 @@ Utilities for mapping exceptions to Anthropic error format. Similar to litellm/litellm_core_utils/exception_mapping_utils.py but for Anthropic response format. """ +from typing import Final + from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from typing import Dict, Optional from .exceptions import AnthropicErrorResponse, AnthropicErrorType # HTTP status code -> Anthropic error type # Source: https://docs.anthropic.com/en/api/errors -ANTHROPIC_ERROR_TYPE_MAP: Dict[int, AnthropicErrorType] = { +ANTHROPIC_ERROR_TYPE_MAP: Final[dict[int, AnthropicErrorType]] = { 400: "invalid_request_error", 401: "authentication_error", 403: "permission_error", @@ -39,7 +40,7 @@ class AnthropicExceptionMapping: def create_error_response( status_code: int, message: str, - request_id: Optional[str] = None, + request_id: str | None = None, ) -> AnthropicErrorResponse: """ Create an Anthropic-formatted error response dict. @@ -51,9 +52,9 @@ class AnthropicExceptionMapping: "request_id": "req_..." } """ - error_type = AnthropicExceptionMapping.get_error_type(status_code) + error_type: Final = AnthropicExceptionMapping.get_error_type(status_code) - response: AnthropicErrorResponse = { + response: Final[AnthropicErrorResponse] = { "type": "error", "error": { "type": error_type, @@ -77,7 +78,7 @@ class AnthropicExceptionMapping: - Generic: {"message": "..."} - Plain strings """ - parsed = safe_json_loads(raw_message) + parsed: Final = safe_json_loads(raw_message) if isinstance(parsed, dict): # Bedrock format if "detail" in parsed and isinstance(parsed["detail"], dict): @@ -124,7 +125,7 @@ class AnthropicExceptionMapping: def transform_to_anthropic_error( status_code: int, raw_message: str, - request_id: Optional[str] = None, + request_id: str | None = None, ) -> AnthropicErrorResponse: """ Transform an error message to Anthropic format. @@ -143,7 +144,7 @@ class AnthropicExceptionMapping: AnthropicErrorResponse dict """ # Try to parse as JSON once - parsed: Optional[dict] = safe_json_loads(raw_message) + parsed: dict | None = safe_json_loads(raw_message) if not isinstance(parsed, dict): parsed = None diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py index b289e493e6b..ae333d1f4ad 100644 --- a/litellm/anthropic_interface/exceptions/exceptions.py +++ b/litellm/anthropic_interface/exceptions/exceptions.py @@ -1,6 +1,8 @@ """Anthropic error format type definitions.""" -from typing_extensions import Literal, Required, TypedDict +from typing import Literal + +from typing_extensions import Required, TypedDict # Known Anthropic error types # Source: https://docs.anthropic.com/en/api/errors diff --git a/litellm/anthropic_interface/messages/__init__.py b/litellm/anthropic_interface/messages/__init__.py index 52c9ecd5aa4..2698cff5980 100644 --- a/litellm/anthropic_interface/messages/__init__.py +++ b/litellm/anthropic_interface/messages/__init__.py @@ -10,7 +10,8 @@ This is an __init__.py file to allow the following interface """ -from typing import Any, AsyncIterator, Coroutine, Dict, Iterator, List, Optional, Union +from collections.abc import AsyncIterator, Coroutine, Iterator +from typing import Any from litellm.llms.anthropic.experimental_pass_through.messages.handler import ( anthropic_messages as _async_anthropic_messages, @@ -25,21 +26,21 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( async def acreate( max_tokens: int, - messages: List[Dict], + messages: list[dict], model: str, - metadata: Optional[Dict] = None, - stop_sequences: Optional[List[str]] = None, - stream: Optional[bool] = False, - system: Optional[str] = None, - temperature: Optional[float] = None, - thinking: Optional[Dict] = None, - tool_choice: Optional[Dict] = None, - tools: Optional[List[Dict]] = None, - top_k: Optional[int] = None, - top_p: Optional[float] = None, - container: Optional[Dict] = None, + metadata: dict | None = None, + stop_sequences: list[str] | None = None, + stream: bool | None = False, + system: str | None = None, + temperature: float | None = None, + thinking: dict | None = None, + tool_choice: dict | None = None, + tools: list[dict] | None = None, + top_k: int | None = None, + top_p: float | None = None, + container: dict | None = None, **kwargs, -) -> Union[AnthropicMessagesResponse, AsyncIterator]: +) -> AnthropicMessagesResponse | AsyncIterator: """ Async wrapper for Anthropic's messages API @@ -84,26 +85,26 @@ async def acreate( def create( max_tokens: int, - messages: List[Dict], + messages: list[dict], model: str, - metadata: Optional[Dict] = None, - stop_sequences: Optional[List[str]] = None, - stream: Optional[bool] = False, - system: Optional[str] = None, - temperature: Optional[float] = None, - thinking: Optional[Dict] = None, - tool_choice: Optional[Dict] = None, - tools: Optional[List[Dict]] = None, - top_k: Optional[int] = None, - top_p: Optional[float] = None, - container: Optional[Dict] = None, + metadata: dict | None = None, + stop_sequences: list[str] | None = None, + stream: bool | None = False, + system: str | None = None, + temperature: float | None = None, + thinking: dict | None = None, + tool_choice: dict | None = None, + tools: list[dict] | None = None, + top_k: int | None = None, + top_p: float | None = None, + container: dict | None = None, **kwargs, -) -> Union[ - AnthropicMessagesResponse, - Iterator[bytes], - AsyncIterator[Any], - Coroutine[Any, Any, Union[AnthropicMessagesResponse, AsyncIterator[Any], Iterator[bytes]]], -]: +) -> ( + AnthropicMessagesResponse + | Iterator[bytes] + | AsyncIterator[Any] + | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any] | Iterator[bytes]] +): """ Async wrapper for Anthropic's messages API diff --git a/litellm/assistants/main.py b/litellm/assistants/main.py index d515cb278bc..237e35fdd5e 100644 --- a/litellm/assistants/main.py +++ b/litellm/assistants/main.py @@ -3,8 +3,9 @@ import asyncio import contextvars import os +from collections.abc import Coroutine, Iterable from functools import partial -from typing import Any, Coroutine, Dict, Iterable, List, Literal, Optional, Union +from typing import Any, Final, Literal import httpx from openai import AsyncOpenAI, OpenAI @@ -28,34 +29,34 @@ from ..types.router import * from .utils import get_optional_params_add_message ####### ENVIRONMENT VARIABLES ################### -openai_assistants_api = OpenAIAssistantsAPI() -azure_assistants_api = AzureAssistantsAPI() +openai_assistants_api: Final = OpenAIAssistantsAPI() +azure_assistants_api: Final = AzureAssistantsAPI() ### ASSISTANTS ### async def aget_assistants( custom_llm_provider: Literal["openai", "azure"], - client: Optional[AsyncOpenAI] = None, + client: AsyncOpenAI | None = None, **kwargs, ) -> AsyncCursorPage[Assistant]: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["aget_assistants"] = True try: # Use a partial function to pass your keyword arguments - func = partial(get_assistants, custom_llm_provider, client, **kwargs) + func: Final = partial(get_assistants, custom_llm_provider, client, **kwargs) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model="", custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -73,17 +74,17 @@ async def aget_assistants( def get_assistants( custom_llm_provider: Literal["openai", "azure"], - client: Optional[Any] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + client: Any | None = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, **kwargs, ) -> SyncCursorPage[Assistant]: - aget_assistants: Optional[bool] = kwargs.pop("aget_assistants", None) + aget_assistants: Final[bool | None] = kwargs.pop("aget_assistants", None) if aget_assistants is not None and not isinstance(aget_assistants, bool): raise Exception("Invalid value passed in for aget_assistants. Only bool or None allowed") - optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + optional_params: Final = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 @@ -94,14 +95,14 @@ def get_assistants( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - response: Optional[SyncCursorPage[Assistant]] = None + response: SyncCursorPage[Assistant] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -110,7 +111,7 @@ def get_assistants( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -146,8 +147,8 @@ def get_assistants( or get_secret("AZURE_API_KEY") ) # type: ignore - extra_body = optional_params.get("extra_body", {}) - azure_ad_token: Optional[str] = None + extra_body: Final = optional_params.get("extra_body", {}) + azure_ad_token: str | None = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) else: @@ -166,9 +167,7 @@ def get_assistants( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'get_assistants'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_assistants'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -180,9 +179,7 @@ def get_assistants( if response is None: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'get_assistants'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_assistants'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -197,28 +194,28 @@ def get_assistants( async def acreate_assistants( custom_llm_provider: Literal["openai", "azure"], - client: Optional[AsyncOpenAI] = None, + client: AsyncOpenAI | None = None, **kwargs, ) -> Assistant: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["async_create_assistants"] = True - model = kwargs.pop("model", None) + model: Final = kwargs.pop("model", None) try: kwargs["client"] = client # Use a partial function to pass your keyword arguments - func = partial(create_assistants, custom_llm_provider, model, **kwargs) + func: Final = partial(create_assistants, custom_llm_provider, model, **kwargs) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model=model, custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -237,26 +234,26 @@ async def acreate_assistants( def create_assistants( custom_llm_provider: Literal["openai", "azure"], model: str, - name: Optional[str] = None, - description: Optional[str] = None, - instructions: Optional[str] = None, - tools: Optional[List[Dict[str, Any]]] = None, - tool_resources: Optional[Dict[str, Any]] = None, - metadata: Optional[Dict[str, str]] = None, - temperature: Optional[float] = None, - top_p: Optional[float] = None, - response_format: Optional[Union[str, Dict[str, str]]] = None, - client: Optional[Any] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + name: str | None = None, + description: str | None = None, + instructions: str | None = None, + tools: list[dict[str, Any]] | None = None, + tool_resources: dict[str, Any] | None = None, + metadata: dict[str, str] | None = None, + temperature: float | None = None, + top_p: float | None = None, + response_format: str | dict[str, str] | None = None, + client: Any | None = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, **kwargs, -) -> Union[Assistant, Coroutine[Any, Any, Assistant]]: - async_create_assistants: Optional[bool] = kwargs.pop("async_create_assistants", None) +) -> Assistant | Coroutine[Any, Any, Assistant]: + async_create_assistants: Final[bool | None] = kwargs.pop("async_create_assistants", None) if async_create_assistants is not None and not isinstance(async_create_assistants, bool): raise ValueError("Invalid value passed in for async_create_assistants. Only bool or None allowed") - optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + optional_params: Final = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 @@ -267,7 +264,7 @@ def create_assistants( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore @@ -290,7 +287,7 @@ def create_assistants( # only send params that are not None create_assistant_data = {k: v for k, v in create_assistant_data.items() if v is not None} - response: Optional[Union[Coroutine[Any, Any, Assistant], Assistant]] = None + response: Coroutine[Any, Any, Assistant] | Assistant | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -299,7 +296,7 @@ def create_assistants( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -336,8 +333,8 @@ def create_assistants( or get_secret("AZURE_API_KEY") ) # type: ignore - extra_body = optional_params.get("extra_body", {}) - azure_ad_token: Optional[str] = None + extra_body: Final = optional_params.get("extra_body", {}) + azure_ad_token: str | None = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) else: @@ -360,9 +357,7 @@ def create_assistants( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_assistants'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_assistants'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -382,27 +377,27 @@ def create_assistants( async def adelete_assistant( custom_llm_provider: Literal["openai", "azure"], - client: Optional[AsyncOpenAI] = None, + client: AsyncOpenAI | None = None, **kwargs, ) -> AssistantDeleted: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["async_delete_assistants"] = True try: kwargs["client"] = client # Use a partial function to pass your keyword arguments - func = partial(delete_assistant, custom_llm_provider, **kwargs) + func: Final = partial(delete_assistant, custom_llm_provider, **kwargs) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model="", custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -421,17 +416,17 @@ async def adelete_assistant( def delete_assistant( custom_llm_provider: Literal["openai", "azure"], assistant_id: str, - client: Optional[Any] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + client: Any | None = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, **kwargs, -) -> Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]: - optional_params = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) +) -> AssistantDeleted | Coroutine[Any, Any, AssistantDeleted]: + optional_params: Final = GenericLiteLLMParams(api_key=api_key, api_base=api_base, api_version=api_version, **kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) - async_delete_assistants: Optional[bool] = kwargs.pop("async_delete_assistants", None) + async_delete_assistants: Final[bool | None] = kwargs.pop("async_delete_assistants", None) if async_delete_assistants is not None and not isinstance(async_delete_assistants, bool): raise ValueError("Invalid value passed in for async_delete_assistants. Only bool or None allowed") @@ -444,14 +439,14 @@ def delete_assistant( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - response: Optional[Union[AssistantDeleted, Coroutine[Any, Any, AssistantDeleted]]] = None + response: AssistantDeleted | Coroutine[Any, Any, AssistantDeleted] | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base @@ -460,7 +455,7 @@ def delete_assistant( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None ) # set API KEY @@ -489,8 +484,8 @@ def delete_assistant( or get_secret("AZURE_API_KEY") ) # type: ignore - extra_body = optional_params.get("extra_body", {}) - azure_ad_token: Optional[str] = None + extra_body: Final = optional_params.get("extra_body", {}) + azure_ad_token: str | None = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) else: @@ -513,9 +508,7 @@ def delete_assistant( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'delete_assistant'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'delete_assistant'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -537,23 +530,23 @@ def delete_assistant( async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwargs) -> Thread: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["acreate_thread"] = True try: # Use a partial function to pass your keyword arguments - func = partial(create_thread, custom_llm_provider, **kwargs) + func: Final = partial(create_thread, custom_llm_provider, **kwargs) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model="", custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -571,10 +564,10 @@ async def acreate_thread(custom_llm_provider: Literal["openai", "azure"], **kwar def create_thread( custom_llm_provider: Literal["openai", "azure"], - messages: Optional[Iterable[OpenAICreateThreadParamsMessage]] = None, - metadata: Optional[dict] = None, - tool_resources: Optional[OpenAICreateThreadParamsToolResources] = None, - client: Optional[OpenAI] = None, + messages: Iterable[OpenAICreateThreadParamsMessage] | None = None, + metadata: dict | None = None, + tool_resources: OpenAICreateThreadParamsToolResources | None = None, + client: OpenAI | None = None, **kwargs, ) -> Thread: """ @@ -599,9 +592,9 @@ def create_thread( ) ``` """ - acreate_thread = kwargs.get("acreate_thread", None) - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + acreate_thread: Final = kwargs.get("acreate_thread", None) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 @@ -612,17 +605,17 @@ def create_thread( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - api_base: Optional[str] = None - api_key: Optional[str] = None + api_base: str | None = None + api_key: str | None = None - response: Optional[Thread] = None + response: Thread | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -631,7 +624,7 @@ def create_thread( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -666,12 +659,10 @@ def create_thread( or get_secret("AZURE_API_KEY") ) # type: ignore - api_version: Optional[str] = ( - optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore - extra_body = optional_params.get("extra_body", {}) - azure_ad_token: Optional[str] = None + extra_body: Final = optional_params.get("extra_body", {}) + azure_ad_token: str | None = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) else: @@ -695,9 +686,7 @@ def create_thread( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_thread'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_thread'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -712,26 +701,26 @@ def create_thread( async def aget_thread( custom_llm_provider: Literal["openai", "azure"], thread_id: str, - client: Optional[AsyncOpenAI] = None, + client: AsyncOpenAI | None = None, **kwargs, ) -> Thread: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["aget_thread"] = True try: # Use a partial function to pass your keyword arguments - func = partial(get_thread, custom_llm_provider, thread_id, client, **kwargs) + func: Final = partial(get_thread, custom_llm_provider, thread_id, client, **kwargs) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model="", custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -754,9 +743,9 @@ def get_thread( **kwargs, ) -> Thread: """Get the thread object, given a thread_id""" - aget_thread = kwargs.pop("aget_thread", None) - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + aget_thread: Final = kwargs.pop("aget_thread", None) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -766,15 +755,15 @@ def get_thread( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - api_base: Optional[str] = None - api_key: Optional[str] = None - response: Optional[Thread] = None + api_base: str | None = None + api_key: str | None = None + response: Thread | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -783,7 +772,7 @@ def get_thread( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -810,9 +799,7 @@ def get_thread( elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version: Optional[str] = ( - optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -822,8 +809,8 @@ def get_thread( or get_secret("AZURE_API_KEY") ) # type: ignore - extra_body = optional_params.get("extra_body", {}) - azure_ad_token: Optional[str] = None + extra_body: Final = optional_params.get("extra_body", {}) + azure_ad_token: str | None = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) else: @@ -846,9 +833,7 @@ def get_thread( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'get_thread'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_thread'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -868,17 +853,17 @@ async def a_add_message( thread_id: str, role: Literal["user", "assistant"], content: str, - attachments: Optional[List[Attachment]] = None, - metadata: Optional[dict] = None, + attachments: list[Attachment] | None = None, + metadata: dict | None = None, client=None, **kwargs, ) -> OpenAIMessage: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["a_add_message"] = True try: # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( add_message, custom_llm_provider, thread_id, @@ -891,15 +876,15 @@ async def a_add_message( ) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model="", custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -921,18 +906,18 @@ def add_message( thread_id: str, role: Literal["user", "assistant"], content: str, - attachments: Optional[List[Attachment]] = None, - metadata: Optional[dict] = None, + attachments: list[Attachment] | None = None, + metadata: dict | None = None, client=None, **kwargs, ) -> OpenAIMessage: ### COMMON OBJECTS ### - a_add_message = kwargs.pop("a_add_message", None) - _message_data = MessageData(role=role, content=content, attachments=attachments, metadata=metadata) - litellm_params_dict = get_litellm_params(**kwargs) - optional_params = GenericLiteLLMParams(**kwargs) + a_add_message: Final = kwargs.pop("a_add_message", None) + _message_data: Final = MessageData(role=role, content=content, attachments=attachments, metadata=metadata) + litellm_params_dict: Final = get_litellm_params(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) - message_data = get_optional_params_add_message( + message_data: Final = get_optional_params_add_message( role=_message_data["role"], content=_message_data["content"], attachments=_message_data["attachments"], @@ -949,15 +934,15 @@ def add_message( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - api_key: Optional[str] = None - api_base: Optional[str] = None - response: Optional[OpenAIMessage] = None + api_key: str | None = None + api_base: str | None = None + response: OpenAIMessage | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -966,7 +951,7 @@ def add_message( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -993,9 +978,7 @@ def add_message( elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version: Optional[str] = ( - optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -1005,8 +988,8 @@ def add_message( or get_secret("AZURE_API_KEY") ) # type: ignore - extra_body = optional_params.get("extra_body", {}) - azure_ad_token: Optional[str] = None + extra_body: Final = optional_params.get("extra_body", {}) + azure_ad_token: str | None = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) else: @@ -1027,9 +1010,7 @@ def add_message( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_thread'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_thread'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -1045,15 +1026,15 @@ def add_message( async def aget_messages( custom_llm_provider: Literal["openai", "azure"], thread_id: str, - client: Optional[AsyncOpenAI] = None, + client: AsyncOpenAI | None = None, **kwargs, ) -> AsyncCursorPage[OpenAIMessage]: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["aget_messages"] = True try: # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( get_messages, custom_llm_provider, thread_id, @@ -1062,15 +1043,15 @@ async def aget_messages( ) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model="", custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -1090,12 +1071,12 @@ async def aget_messages( def get_messages( custom_llm_provider: Literal["openai", "azure"], thread_id: str, - client: Optional[Any] = None, + client: Any | None = None, **kwargs, ) -> SyncCursorPage[OpenAIMessage]: - aget_messages = kwargs.pop("aget_messages", None) - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + aget_messages: Final = kwargs.pop("aget_messages", None) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 @@ -1106,16 +1087,16 @@ def get_messages( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - response: Optional[SyncCursorPage[OpenAIMessage]] = None - api_key: Optional[str] = None - api_base: Optional[str] = None + response: SyncCursorPage[OpenAIMessage] | None = None + api_key: str | None = None + api_base: str | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -1124,7 +1105,7 @@ def get_messages( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -1150,9 +1131,7 @@ def get_messages( elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore - api_version: Optional[str] = ( - optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") - ) # type: ignore + api_version: str | None = optional_params.api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore api_key = ( optional_params.api_key @@ -1162,8 +1141,8 @@ def get_messages( or get_secret("AZURE_API_KEY") ) # type: ignore - extra_body = optional_params.get("extra_body", {}) - azure_ad_token: Optional[str] = None + extra_body: Final = optional_params.get("extra_body", {}) + azure_ad_token: str | None = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) else: @@ -1183,9 +1162,7 @@ def get_messages( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'get_messages'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'get_messages'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -1201,7 +1178,7 @@ def get_messages( ### RUNS ### def arun_thread_stream( *, - event_handler: Optional[AssistantEventHandler] = None, + event_handler: AssistantEventHandler | None = None, **kwargs, ) -> AsyncAssistantStreamManager[AsyncAssistantEventHandler]: kwargs["arun_thread"] = True @@ -1212,21 +1189,21 @@ async def arun_thread( custom_llm_provider: Literal["openai", "azure"], thread_id: str, assistant_id: str, - additional_instructions: Optional[str] = None, - instructions: Optional[str] = None, - metadata: Optional[dict] = None, - model: Optional[str] = None, - stream: Optional[bool] = None, - tools: Optional[Iterable[AssistantToolParam]] = None, - client: Optional[Any] = None, + additional_instructions: str | None = None, + instructions: str | None = None, + metadata: dict | None = None, + model: str | None = None, + stream: bool | None = None, + tools: Iterable[AssistantToolParam] | None = None, + client: Any | None = None, **kwargs, ) -> Run: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() ### PASS ARGS TO GET ASSISTANTS ### kwargs["arun_thread"] = True try: # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( run_thread, custom_llm_provider, thread_id, @@ -1242,15 +1219,15 @@ async def arun_thread( ) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore model="", custom_llm_provider=custom_llm_provider ) # type: ignore # Await normally - init_response = await loop.run_in_executor(None, func_with_context) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -1269,7 +1246,7 @@ async def arun_thread( def run_thread_stream( *, - event_handler: Optional[AssistantEventHandler] = None, + event_handler: AssistantEventHandler | None = None, **kwargs, ) -> AssistantStreamManager[AssistantEventHandler]: return run_thread(stream=True, event_handler=event_handler, **kwargs) # type: ignore @@ -1279,20 +1256,20 @@ def run_thread( custom_llm_provider: Literal["openai", "azure"], thread_id: str, assistant_id: str, - additional_instructions: Optional[str] = None, - instructions: Optional[str] = None, - metadata: Optional[dict] = None, - model: Optional[str] = None, - stream: Optional[bool] = None, - tools: Optional[Iterable[AssistantToolParam]] = None, - client: Optional[Any] = None, - event_handler: Optional[AssistantEventHandler] = None, # for stream=True calls + additional_instructions: str | None = None, + instructions: str | None = None, + metadata: dict | None = None, + model: str | None = None, + stream: bool | None = None, + tools: Iterable[AssistantToolParam] | None = None, + client: Any | None = None, + event_handler: AssistantEventHandler | None = None, # for stream=True calls **kwargs, ) -> Run: """Run a given thread + assistant.""" - arun_thread = kwargs.pop("arun_thread", None) - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + arun_thread: Final = kwargs.pop("arun_thread", None) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 @@ -1303,14 +1280,14 @@ def run_thread( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - response: Optional[Run] = None + response: Run | None = None if custom_llm_provider == "openai": api_base = ( optional_params.api_base # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there @@ -1319,7 +1296,7 @@ def run_thread( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -1364,7 +1341,7 @@ def run_thread( or get_secret("AZURE_API_KEY") ) # type: ignore - extra_body = optional_params.get("extra_body", {}) + extra_body: Final = optional_params.get("extra_body", {}) azure_ad_token = None if extra_body is not None: azure_ad_token = extra_body.pop("azure_ad_token", None) @@ -1392,9 +1369,7 @@ def run_thread( ) # type: ignore else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'run_thread'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'run_thread'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( diff --git a/litellm/assistants/utils.py b/litellm/assistants/utils.py index f775c1b6508..e80131a5011 100644 --- a/litellm/assistants/utils.py +++ b/litellm/assistants/utils.py @@ -1,4 +1,4 @@ -from typing import Optional, Union +from typing import Final import litellm @@ -7,21 +7,10 @@ from ..types.llms.openai import * def get_optional_params_add_message( - role: Optional[str], - content: Optional[ - Union[ - str, - List[ - Union[ - MessageContentTextObject, - MessageContentImageFileObject, - MessageContentImageURLObject, - ] - ], - ] - ], - attachments: Optional[List[Attachment]], - metadata: Optional[dict], + role: str | None, + content: str | List[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None, + attachments: List[Attachment] | None, + metadata: dict | None, custom_llm_provider: str, **kwargs, ): @@ -30,13 +19,13 @@ def get_optional_params_add_message( Reference - https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-messages?tabs=python#create-message """ - passed_params = locals() + passed_params: Final = locals() custom_llm_provider = passed_params.pop("custom_llm_provider") - special_params = passed_params.pop("kwargs") + special_params: Final = passed_params.pop("kwargs") for k, v in special_params.items(): passed_params[k] = v - default_params = { + default_params: Final = { "role": None, "content": None, "attachments": None, @@ -49,51 +38,49 @@ def get_optional_params_add_message( ## raise exception if non-default value passed for non-openai/azure embedding calls def _check_valid_arg(supported_params): if len(non_default_params.keys()) > 0: - keys = list(non_default_params.keys()) + keys: Final = list(non_default_params.keys()) for k in keys: if litellm.drop_params is True and k not in supported_params: # drop the unsupported non-default values non_default_params.pop(k, None) elif k not in supported_params: raise litellm.utils.UnsupportedParamsError( status_code=500, - message="k={}, not supported by {}. Supported params={}. To drop it from the call, set `litellm.drop_params = True`.".format( - k, custom_llm_provider, supported_params - ), + message=f"k={k}, not supported by {custom_llm_provider}. Supported params={supported_params}. To drop it from the call, set `litellm.drop_params = True`.", ) return non_default_params if custom_llm_provider == "openai": optional_params = non_default_params elif custom_llm_provider == "azure": - supported_params = litellm.AzureOpenAIAssistantsAPIConfig().get_supported_openai_create_message_params() + supported_params: Final = litellm.AzureOpenAIAssistantsAPIConfig().get_supported_openai_create_message_params() _check_valid_arg(supported_params=supported_params) optional_params = litellm.AzureOpenAIAssistantsAPIConfig().map_openai_params_create_message_params( non_default_params=non_default_params, optional_params=optional_params ) for k in passed_params.keys(): - if k not in default_params.keys(): + if k not in default_params: optional_params[k] = passed_params[k] return optional_params def get_optional_params_image_gen( - n: Optional[int] = None, - quality: Optional[str] = None, - response_format: Optional[str] = None, - size: Optional[str] = None, - style: Optional[str] = None, - user: Optional[str] = None, - custom_llm_provider: Optional[str] = None, + n: int | None = None, + quality: str | None = None, + response_format: str | None = None, + size: str | None = None, + style: str | None = None, + user: str | None = None, + custom_llm_provider: str | None = None, **kwargs, ): # retrieve all parameters passed to the function - passed_params = locals() + passed_params: Final = locals() custom_llm_provider = passed_params.pop("custom_llm_provider") - special_params = passed_params.pop("kwargs") + special_params: Final = passed_params.pop("kwargs") for k, v in special_params.items(): passed_params[k] = v - default_params = { + default_params: Final = { "n": None, "quality": None, "response_format": None, @@ -108,7 +95,7 @@ def get_optional_params_image_gen( ## raise exception if non-default value passed for non-openai/azure embedding calls def _check_valid_arg(supported_params): if len(non_default_params.keys()) > 0: - keys = list(non_default_params.keys()) + keys: Final = list(non_default_params.keys()) for k in keys: if litellm.drop_params is True and k not in supported_params: # drop the unsupported non-default values non_default_params.pop(k, None) @@ -142,6 +129,6 @@ def get_optional_params_image_gen( optional_params["sampleCount"] = int(n) for k in passed_params.keys(): - if k not in default_params.keys(): + if k not in default_params: optional_params[k] = passed_params[k] return optional_params diff --git a/litellm/batch_completion/main.py b/litellm/batch_completion/main.py index 664977dc8d6..702dd194fda 100644 --- a/litellm/batch_completion/main.py +++ b/litellm/batch_completion/main.py @@ -1,5 +1,5 @@ from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait -from typing import List, Optional +from typing import Final import litellm from litellm._logging import print_verbose @@ -11,23 +11,23 @@ from ..llms.vllm.completion import handler as vllm_handler def batch_completion( model: str, # Optional OpenAI params: see https://platform.openai.com/docs/api-reference/chat/create - messages: List = [], - functions: Optional[List] = None, - function_call: Optional[str] = None, - temperature: Optional[float] = None, - top_p: Optional[float] = None, - n: Optional[int] = None, - stream: Optional[bool] = None, + messages: list = [], + functions: list | None = None, + function_call: str | None = None, + temperature: float | None = None, + top_p: float | None = None, + n: int | None = None, + stream: bool | None = None, stop=None, - max_tokens: Optional[int] = None, - presence_penalty: Optional[float] = None, - frequency_penalty: Optional[float] = None, - logit_bias: Optional[dict] = None, - user: Optional[str] = None, + max_tokens: int | None = None, + presence_penalty: float | None = None, + frequency_penalty: float | None = None, + logit_bias: dict | None = None, + user: str | None = None, deployment_id=None, - request_timeout: Optional[int] = None, - timeout: Optional[int] = 600, - max_workers: Optional[int] = 100, + request_timeout: int | None = None, + timeout: int | None = 600, + max_workers: int | None = 100, # Optional liteLLM function params **kwargs, ): @@ -56,17 +56,17 @@ def batch_completion( Returns: list: A list of completion results. """ - args = locals() + args: Final = locals() - batch_messages = messages - completions = [] + batch_messages: Final = messages + completions: Final = [] model = model custom_llm_provider = None if model.split("/", 1)[0] in litellm.provider_list: custom_llm_provider = model.split("/", 1)[0] model = model.split("/", 1)[1] if custom_llm_provider == "vllm": - optional_params = get_optional_params( + optional_params: Final = get_optional_params( functions=functions, function_call=function_call, temperature=temperature, @@ -146,7 +146,7 @@ def batch_completion_models(*args, **kwargs): if "model" in kwargs: kwargs.pop("model") if "models" in kwargs: - models = kwargs["models"] + models: Final = kwargs["models"] kwargs.pop("models") futures = {} with ThreadPoolExecutor(max_workers=len(models)) as executor: @@ -157,14 +157,14 @@ def batch_completion_models(*args, **kwargs): if future.result() is not None: return future.result() elif "deployments" in kwargs: - deployments = kwargs["deployments"] + deployments: Final = kwargs["deployments"] kwargs.pop("deployments") kwargs.pop("model_list") - nested_kwargs = kwargs.pop("kwargs", {}) + nested_kwargs: Final = kwargs.pop("kwargs", {}) futures = {} with ThreadPoolExecutor(max_workers=len(deployments)) as executor: for deployment in deployments: - for key in kwargs.keys(): + for key in kwargs: if key not in deployment: # don't override deployment values e.g. model name, api base, etc. deployment[key] = kwargs[key] kwargs = {**deployment, **nested_kwargs} @@ -239,10 +239,10 @@ def batch_completion_models_all_responses(*args, **kwargs): if len(models) == 0: return [] - responses = [] + responses: Final = [] with concurrent.futures.ThreadPoolExecutor(max_workers=len(models)) as executor: - futures = [executor.submit(litellm.completion, *args, model=model, **kwargs) for model in models] + futures: Final = [executor.submit(litellm.completion, *args, model=model, **kwargs) for model in models] for future in futures: try: @@ -250,7 +250,7 @@ def batch_completion_models_all_responses(*args, **kwargs): if result is not None: responses.append(result) except Exception as e: - print_verbose(f"batch_completion_models_all_responses: model request failed: {str(e)}") + print_verbose(f"batch_completion_models_all_responses: model request failed: {e}") continue return responses diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index eef4cf8d87f..1ef6e06abfe 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,6 +1,7 @@ import json +from collections.abc import Iterable, Iterator from dataclasses import dataclass -from typing import Any, Iterable, Iterator, List, Literal, Optional, Tuple +from typing import Any, Final, Literal import litellm from litellm._logging import verbose_logger @@ -11,11 +12,11 @@ from litellm.utils import token_counter async def calculate_batch_cost_and_usage( - file_content_dictionary: List[dict], + file_content_dictionary: list[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], - model_name: Optional[str] = None, - model_info: Optional[ModelInfo] = None, -) -> Tuple[float, Usage, List[str]]: + model_name: str | None = None, + model_info: ModelInfo | None = None, +) -> tuple[float, Usage, list[str]]: """ Calculate the cost and usage of a batch. @@ -44,9 +45,9 @@ async def calculate_batch_cost_and_usage( async def _handle_completed_batch( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], - model_name: Optional[str] = None, - litellm_params: Optional[dict] = None, -) -> Tuple[float, Usage, List[str]]: + model_name: str | None = None, + litellm_params: dict | None = None, +) -> tuple[float, Usage, list[str]]: """Fetch a completed batch's output file and aggregate its cost, usage, and models in a single pass over the JSONL lines, so the parsed file content is never materialized in memory. @@ -84,14 +85,14 @@ class _BatchOutputLineStats: total_tokens: int cache_read_tokens: int cache_creation_tokens: int - model: Optional[str] + model: str | None def _iter_successful_output_line_stats( entries: Iterable[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], - model_name: Optional[str], - model_info: Optional[ModelInfo], + model_name: str | None, + model_info: ModelInfo | None, ) -> Iterator[_BatchOutputLineStats]: from litellm.cost_calculator import batch_cost_calculator @@ -135,14 +136,14 @@ def _iter_successful_output_line_stats( def _aggregate_batch_cost_usage_models( entries: Iterable[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], - model_name: Optional[str] = None, - model_info: Optional[ModelInfo] = None, -) -> Tuple[float, Usage, List[str]]: + model_name: str | None = None, + model_info: ModelInfo | None = None, +) -> tuple[float, Usage, list[str]]: """Aggregate cost, usage, and models from batch output entries in a single pass, holding one small stats record per line instead of the parsed file.""" - line_stats = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info)) + line_stats: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info)) - cache_token_params = { + cache_token_params: Final = { key: tokens for key, tokens in ( ("cache_read_input_tokens", sum(stats.cache_read_tokens for stats in line_stats)), @@ -150,22 +151,22 @@ def _aggregate_batch_cost_usage_models( ) if tokens > 0 } - batch_usage = Usage( + batch_usage: Final = Usage( total_tokens=sum(stats.total_tokens for stats in line_stats), prompt_tokens=sum(stats.prompt_tokens for stats in line_stats), completion_tokens=sum(stats.completion_tokens for stats in line_stats), **cache_token_params, ) - batch_models = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] - total_cost = sum((stats.cost for stats in line_stats), 0.0) + batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] + total_cost: Final = sum((stats.cost for stats in line_stats), 0.0) verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models) return total_cost, batch_usage, batch_models def calculate_vertex_ai_batch_cost_and_usage( - vertex_ai_batch_responses: List[dict], - model_name: Optional[str] = None, -) -> Tuple[float, Usage]: + vertex_ai_batch_responses: list[dict], + model_name: str | None = None, +) -> tuple[float, Usage]: """ Calculate both cost and usage from raw Vertex AI batch responses. @@ -183,7 +184,7 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 - actual_model_name = model_name or "gemini-2.0-flash-001" + actual_model_name: Final = model_name or "gemini-2.0-flash-001" for response in vertex_ai_batch_responses: response_body = response.get("response") @@ -233,7 +234,7 @@ def calculate_vertex_ai_batch_cost_and_usage( async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", - litellm_params: Optional[dict] = None, + litellm_params: dict | None = None, ) -> bytes: """ Fetch the batch output file and return its raw JSONL bytes @@ -253,31 +254,31 @@ async def _fetch_batch_output_file_content( raise ValueError("Output file id is None cannot retrieve file content") file_id = batch.output_file_id - is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id) + is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) if is_base64_unified_file_id: try: file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug(f"Extracted LLM output file ID from unified file ID: {file_id}") + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) except (IndexError, AttributeError) as e: verbose_logger.error( - f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}" + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e ) # Build kwargs for afile_content with credentials from litellm_params - file_content_kwargs = { + file_content_kwargs: Final = { "file_id": file_id, "custom_llm_provider": custom_llm_provider, } # Extract and add credentials for file access - credentials = _extract_file_access_credentials(litellm_params) + credentials: Final = _extract_file_access_credentials(litellm_params) file_content_kwargs.update(credentials) - _file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType] + _file_content: Final = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType] return _file_content.content -def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict: +def _extract_file_access_credentials(litellm_params: dict | None) -> dict: """ Extract credentials from litellm_params for file access operations. @@ -290,11 +291,11 @@ def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict: Returns: Dictionary containing only the credentials needed for file access """ - credentials = {} + credentials: Final = {} if litellm_params: # List of credential keys that should be passed to file operations - credential_keys = [ + credential_keys: Final = [ "api_key", "api_base", "api_version", @@ -316,7 +317,7 @@ def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict: return credentials -def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: +def _get_file_content_as_dictionary(file_content: bytes) -> list[dict]: """ Get the file content as a list of dictionaries from JSON Lines format """ @@ -354,7 +355,7 @@ def _iter_batch_input_entries(file_content: bytes) -> Iterator[dict]: # A batch request's input tokens scale roughly with its serialized size, so this # is a conservative per-row fallback when the token counter cannot measure a row. -_BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN = 4 +_BATCH_TOKEN_ESTIMATE_BYTES_PER_TOKEN: Final = 4 def _estimate_batch_entry_tokens(raw_line: bytes) -> int: @@ -366,21 +367,21 @@ def _estimate_batch_entry_tokens(raw_line: bytes) -> int: def _count_entry_tokens( entry: dict, - model_name: Optional[str] = None, + model_name: str | None = None, ) -> int: """Token-count a single batch input entry's body (chat / text / embedding).""" - body = entry.get("body", {}) or {} - model = body.get("model", model_name or "") + body: Final = entry.get("body", {}) or {} + model: Final = body.get("model", model_name or "") - messages = body.get("messages") + messages: Final = body.get("messages") if messages: return token_counter(model=model, messages=messages) - prompt = body.get("prompt") + prompt: Final = body.get("prompt") if prompt: return _count_prompt_or_input_tokens(model=model, value=prompt) - input_data = body.get("input") + input_data: Final = body.get("input") if input_data: return _count_prompt_or_input_tokens(model=model, value=input_data) @@ -431,8 +432,12 @@ def _get_batch_job_usage_from_response_body(response_body: dict, custom_llm_prov usage_object=response_body.get("usage", None) or {}, reasoning_content=None, ) - _usage_dict = response_body.get("usage", None) or {} - usage: Usage = Usage(**_usage_dict) + from litellm.responses.utils import ResponseAPILoggingUtils + + _usage_dict: Final = response_body.get("usage", None) or {} + if ResponseAPILoggingUtils._is_response_api_usage(_usage_dict): + return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_usage_dict) + usage: Final[Usage] = Usage(**_usage_dict) return usage @@ -454,8 +459,8 @@ def _get_response_from_batch_job_output_file(batch_job_output_file: dict, custom return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("message", None) or {} if custom_llm_provider == "bedrock": return batch_job_output_file.get("modelOutput", None) or {} - _response: dict = batch_job_output_file.get("response", None) or {} - _response_body = _response.get("body", None) or {} + _response: Final[dict] = batch_job_output_file.get("response", None) or {} + _response_body: Final = _response.get("body", None) or {} return _response_body @@ -471,5 +476,5 @@ def _batch_response_was_successful(batch_job_output_file: dict, custom_llm_provi return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded" if custom_llm_provider == "bedrock": return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None - _response: dict = batch_job_output_file.get("response", None) or {} + _response: Final[dict] = batch_job_output_file.get("response", None) or {} return _response.get("status_code", None) == 200 diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 3a2d9e13f77..61a42f515cc 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -13,8 +13,9 @@ https://platform.openai.com/docs/api-reference/batch import asyncio import contextvars import os +from collections.abc import Coroutine from functools import partial -from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast +from typing import Any, Final, Literal, cast import httpx from openai.types.batch import BatchRequestCounts @@ -53,17 +54,17 @@ from litellm.utils import ( ) ####### ENVIRONMENT VARIABLES ################### -openai_batches_instance = OpenAIBatchesAPI() -azure_batches_instance = AzureBatchesAPI() -vertex_ai_batches_instance = VertexAIBatchPrediction(gcs_bucket_name="") -anthropic_batches_instance = AnthropicBatchesHandler() +openai_batches_instance: Final = OpenAIBatchesAPI() +azure_batches_instance: Final = AzureBatchesAPI() +vertex_ai_batches_instance: Final = VertexAIBatchPrediction(gcs_bucket_name="") +anthropic_batches_instance: Final = AnthropicBatchesHandler() base_llm_http_handler = BaseLLMHTTPHandler() ################################################# def _resolve_timeout( optional_params: GenericLiteLLMParams, - kwargs: Dict[str, Any], + kwargs: dict[str, Any], custom_llm_provider: str, default_timeout: float = 600.0, ) -> float: @@ -79,13 +80,13 @@ def _resolve_timeout( Returns: Resolved timeout as float """ - timeout = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout + timeout: Final = optional_params.timeout or kwargs.get("request_timeout", default_timeout) or default_timeout # Handle httpx.Timeout objects if isinstance(timeout, httpx.Timeout): if supports_httpx_timeout(custom_llm_provider) is False: # Extract read timeout for providers that don't support httpx.Timeout - read_timeout = timeout.read or default_timeout + read_timeout: Final = timeout.read or default_timeout return float(read_timeout) else: # For providers that support httpx.Timeout, we still need to return a float @@ -103,13 +104,13 @@ def _resolve_timeout( @client async def acreate_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], input_file_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", - metadata: Optional[Dict[str, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, - output_expires_after: Optional[Dict[str, Any]] = None, + metadata: dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, + output_expires_after: dict[str, Any] | None = None, **kwargs, ) -> LiteLLMBatch: """ @@ -118,11 +119,11 @@ async def acreate_batch( LiteLLM Equivalent of POST: https://api.openai.com/v1/batches """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acreate_batch"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( create_batch, completion_window, endpoint, @@ -136,9 +137,9 @@ async def acreate_batch( ) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response @@ -153,26 +154,26 @@ async def acreate_batch( @client def create_batch( completion_window: Literal["24h"], - endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], + endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"], input_file_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", - metadata: Optional[Dict[str, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, - output_expires_after: Optional[Dict[str, Any]] = None, + metadata: dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, + output_expires_after: dict[str, Any] | None = None, **kwargs, -) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: +) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: """ Creates and executes a batch from an uploaded file of request LiteLLM Equivalent of POST: https://api.openai.com/v1/batches """ try: - optional_params = GenericLiteLLMParams(**kwargs) - litellm_call_id = kwargs.get("litellm_call_id", None) - proxy_server_request = kwargs.get("proxy_server_request", None) - model_info = kwargs.get("model_info", None) - model: Optional[str] = kwargs.get("model", None) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_call_id: Final = kwargs.get("litellm_call_id", None) + proxy_server_request: Final = kwargs.get("proxy_server_request", None) + model_info: Final = kwargs.get("model_info", None) + model: str | None = kwargs.get("model", None) try: if model is not None: model, _, _, _ = get_llm_provider( @@ -181,14 +182,14 @@ def create_batch( ) except Exception as e: verbose_logger.exception( - f"litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - {str(e)}" + "litellm.batches.main.py::create_batch() - Error inferring custom_llm_provider - %s", e ) - _is_async = kwargs.pop("acreate_batch", False) is True - litellm_params = dict(GenericLiteLLMParams(**kwargs)) - litellm_logging_obj: LiteLLMLoggingObj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None)) + _is_async: Final = kwargs.pop("acreate_batch", False) is True + litellm_params: Final = dict(GenericLiteLLMParams(**kwargs)) + litellm_logging_obj: Final[LiteLLMLoggingObj] = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None)) ### TIMEOUT LOGIC ### - timeout = _resolve_timeout(optional_params, kwargs, custom_llm_provider) + timeout: Final = _resolve_timeout(optional_params, kwargs, custom_llm_provider) litellm_logging_obj.update_from_kwargs( kwargs=kwargs, model=model, @@ -205,7 +206,7 @@ def create_batch( custom_llm_provider=custom_llm_provider, ) - _create_batch_request = CreateBatchRequest( + _create_batch_request: Final = CreateBatchRequest( completion_window=completion_window, endpoint=endpoint, input_file_id=input_file_id, @@ -237,7 +238,7 @@ def create_batch( model=model, ) return response - api_base: Optional[str] = None + api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( @@ -247,7 +248,7 @@ def create_batch( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -300,13 +301,13 @@ def create_batch( ) elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" - vertex_ai_project = ( + vertex_ai_project: Final = ( optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) - vertex_ai_location = ( + vertex_ai_location: Final = ( optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") + vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.create_batch( _is_async=_is_async, @@ -320,7 +321,7 @@ def create_batch( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support custom_llm_provider={} for 'create_batch'".format(custom_llm_provider), + message=f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'create_batch'", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -338,9 +339,9 @@ def create_batch( async def aretrieve_batch( batch_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", - metadata: Optional[Dict[str, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + metadata: dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> LiteLLMBatch: """ @@ -349,11 +350,11 @@ async def aretrieve_batch( LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id} """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["aretrieve_batch"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( retrieve_batch, batch_id, custom_llm_provider, @@ -363,9 +364,9 @@ async def aretrieve_batch( **kwargs, ) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -379,14 +380,14 @@ async def aretrieve_batch( def _handle_retrieve_batch_providers_without_provider_config( batch_id: str, optional_params: GenericLiteLLMParams, - timeout: Union[float, httpx.Timeout], + timeout: float | httpx.Timeout, litellm_params: dict, _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", - logging_obj: Optional[Any] = None, + logging_obj: Any | None = None, ): - api_base: Optional[str] = None + api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( @@ -396,7 +397,7 @@ def _handle_retrieve_batch_providers_without_provider_config( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -421,7 +422,7 @@ def _handle_retrieve_batch_providers_without_provider_config( ) elif custom_llm_provider == "azure": api_base = optional_params.api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") - api_version = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") + api_version: Final = optional_params.api_version or litellm.api_version or get_secret_str("AZURE_API_VERSION") api_key = ( optional_params.api_key @@ -431,7 +432,7 @@ def _handle_retrieve_batch_providers_without_provider_config( or get_secret_str("AZURE_API_KEY") ) - extra_body = optional_params.get("extra_body", {}) + extra_body: Final = optional_params.get("extra_body", {}) if extra_body is not None: extra_body.pop("azure_ad_token", None) else: @@ -449,13 +450,13 @@ def _handle_retrieve_batch_providers_without_provider_config( ) elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" - vertex_ai_project = ( + vertex_ai_project: Final = ( optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) - vertex_ai_location = ( + vertex_ai_location: Final = ( optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") + vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.retrieve_batch( _is_async=_is_async, @@ -488,10 +489,10 @@ def _handle_retrieve_batch_providers_without_provider_config( else: raise litellm.exceptions.BadRequestError( message=( - "LiteLLM doesn't support custom_llm_provider={} for 'retrieve_batch' without a `model` kwarg. " + f"LiteLLM doesn't support custom_llm_provider={custom_llm_provider} for 'retrieve_batch' without a `model` kwarg. " "Supported via this path: 'openai', 'azure', 'vertex_ai', 'anthropic'. " "'bedrock' is supported but requires `model` to be passed so the provider config can be loaded." - ).format(custom_llm_provider), + ), model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -507,22 +508,22 @@ def _handle_retrieve_batch_providers_without_provider_config( def retrieve_batch( batch_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai", - metadata: Optional[Dict[str, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + metadata: dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, -) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: +) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: """ Retrieves a batch. LiteLLM Equivalent of GET https://api.openai.com/v1/batches/{batch_id} """ try: - optional_params = GenericLiteLLMParams(**kwargs) - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj", None) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj", None) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 - litellm_params = get_litellm_params( + litellm_params: Final = get_litellm_params( custom_llm_provider=custom_llm_provider, **kwargs, ) @@ -541,21 +542,21 @@ def retrieve_batch( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _retrieve_batch_request = RetrieveBatchRequest( + _retrieve_batch_request: Final = RetrieveBatchRequest( batch_id=batch_id, extra_headers=extra_headers, extra_body=extra_body, ) - _is_async = kwargs.pop("aretrieve_batch", False) is True - client = kwargs.get("client", None) + _is_async: Final = kwargs.pop("aretrieve_batch", False) is True + client: Final = kwargs.get("client", None) # Bedrock has two distinct ARN families that need different APIs: # * async-invoke ARNs (Twelve Labs Marengo embeddings) -> bedrock-runtime data plane @@ -567,7 +568,7 @@ def retrieve_batch( if batch_id.startswith("arn:aws") and ":bedrock:" in batch_id: if ":async-invoke/" in batch_id: # Remove aws_region_name from kwargs to avoid duplicate parameter - async_kwargs = kwargs.copy() + async_kwargs: Final = kwargs.copy() async_kwargs.pop("aws_region_name", None) return BedrockBatchesHandler._handle_async_invoke_status( @@ -577,7 +578,7 @@ def retrieve_batch( **async_kwargs, ) if ":model-invocation-job/" in batch_id: - mij_kwargs = kwargs.copy() + mij_kwargs: Final = kwargs.copy() mij_kwargs.pop("aws_region_name", None) return BedrockBatchesHandler._handle_model_invocation_job_status( @@ -588,7 +589,7 @@ def retrieve_batch( ) # Try to use provider config first (for providers like bedrock) - model: Optional[str] = kwargs.get("model", None) + model: Final[str | None] = kwargs.get("model", None) if model is not None: provider_config = ProviderConfigManager.get_provider_batches_config( model=model, @@ -598,7 +599,7 @@ def retrieve_batch( provider_config = None if provider_config is not None: - response = base_llm_http_handler.retrieve_batch( + response: Final = base_llm_http_handler.retrieve_batch( batch_id=batch_id, provider_config=provider_config, litellm_params=litellm_params, @@ -642,12 +643,12 @@ def retrieve_batch( @client async def alist_batches( - after: Optional[str] = None, - limit: Optional[int] = None, + after: str | None = None, + limit: int | None = None, custom_llm_provider: ListBatchesSupportedProvider = "openai", - metadata: Optional[Dict[str, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + metadata: dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ): """ @@ -655,11 +656,11 @@ async def alist_batches( """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["alist_batches"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( list_batches, after, limit, @@ -670,9 +671,9 @@ async def alist_batches( ) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -685,11 +686,11 @@ async def alist_batches( @client def list_batches( - after: Optional[str] = None, - limit: Optional[int] = None, + after: str | None = None, + limit: int | None = None, custom_llm_provider: ListBatchesSupportedProvider = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ): """ @@ -699,8 +700,8 @@ def list_batches( """ try: # set API KEY - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params = get_litellm_params( + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params: Final = get_litellm_params( custom_llm_provider=custom_llm_provider, **kwargs, ) @@ -719,14 +720,14 @@ def list_batches( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _is_async = kwargs.pop("alist_batches", False) is True + _is_async: Final = kwargs.pop("alist_batches", False) is True if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( @@ -736,7 +737,7 @@ def list_batches( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) @@ -782,13 +783,13 @@ def list_batches( ) elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or "" - vertex_ai_project = ( + vertex_ai_project: Final = ( optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) - vertex_ai_location = ( + vertex_ai_location: Final = ( optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") + vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.list_batches( _is_async=_is_async, @@ -822,11 +823,11 @@ def list_batches( async def acancel_batch( batch_id: str, - model: Optional[str] = None, + model: str | None = None, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - metadata: Optional[Dict[str, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + metadata: dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> LiteLLMBatch: """ @@ -835,14 +836,14 @@ async def acancel_batch( LiteLLM Equivalent of POST https://api.openai.com/v1/batches/{batch_id}/cancel """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acancel_batch"] = True # Preserve model parameter - only pop from kwargs if it exists there # (to avoid passing it twice), otherwise keep the function parameter value model = kwargs.pop("model", None) or model # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( cancel_batch, batch_id, model, @@ -853,9 +854,9 @@ async def acancel_batch( **kwargs, ) # Add the context to the function - ctx = contextvars.copy_context() - func_with_context = partial(ctx.run, func) - init_response = await loop.run_in_executor(None, func_with_context) + ctx: Final = contextvars.copy_context() + func_with_context: Final = partial(ctx.run, func) + init_response: Final = await loop.run_in_executor(None, func_with_context) if asyncio.iscoroutine(init_response): response = await init_response else: @@ -868,13 +869,13 @@ async def acancel_batch( def cancel_batch( batch_id: str, - model: Optional[str] = None, - custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai", - metadata: Optional[Dict[str, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + model: str | None = None, + custom_llm_provider: Literal["openai", "azure", "vertex_ai"] | str = "openai", + metadata: dict[str, str] | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, -) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]: +) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: """ Cancels a batch. @@ -889,10 +890,10 @@ def cancel_batch( ) except Exception as e: verbose_logger.exception( - f"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - {str(e)}" + "litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - %s", e ) - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params = get_litellm_params( + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params: Final = get_litellm_params( custom_llm_provider=custom_llm_provider, **kwargs, ) @@ -905,21 +906,21 @@ def cancel_batch( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _cancel_batch_request = CancelBatchRequest( + _cancel_batch_request: Final = CancelBatchRequest( batch_id=batch_id, extra_headers=extra_headers, extra_body=extra_body, ) - _is_async = kwargs.pop("acancel_batch", False) is True - api_base: Optional[str] = None + _is_async: Final = kwargs.pop("acancel_batch", False) is True + api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: api_base = ( optional_params.api_base @@ -928,7 +929,7 @@ def cancel_batch( or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - organization = ( + organization: Final = ( optional_params.organization or litellm.organization or os.getenv("OPENAI_ORGANIZATION", None) or None ) api_key = optional_params.api_key or litellm.api_key or litellm.openai_key or os.getenv("OPENAI_API_KEY") @@ -972,13 +973,13 @@ def cancel_batch( ) elif custom_llm_provider == "vertex_ai": api_base = optional_params.api_base or None - vertex_ai_project = ( + vertex_ai_project: Final = ( optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT") ) - vertex_ai_location = ( + vertex_ai_location: Final = ( optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION") ) - vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") + vertex_credentials: Final = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS") response = vertex_ai_batches_instance.cancel_batch( _is_async=_is_async, @@ -992,9 +993,7 @@ def cancel_batch( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -1026,10 +1025,10 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj async def _async_get_status(): # Create embedding handler instance - embedding_handler = BedrockEmbedding() + embedding_handler: Final = BedrockEmbedding() # Get the status of the async invoke job - status_response = await embedding_handler._get_async_invoke_status( + status_response: Final = await embedding_handler._get_async_invoke_status( invocation_arn=batch_id, aws_region_name=aws_region_name, logging_obj=logging_obj, @@ -1041,16 +1040,16 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj from litellm.types.utils import LiteLLMBatch # Normalize status to lowercase (AWS returns 'Completed', 'Failed', etc.) - aws_status_raw = status_response.get("status", "") - aws_status_lower = aws_status_raw.lower() + aws_status_raw: Final = status_response.get("status", "") + aws_status_lower: Final = aws_status_raw.lower() # Map AWS status values to LiteLLM expected values - status_mapping: dict[str, BatchJobStatus] = { + status_mapping: Final[dict[str, BatchJobStatus]] = { "completed": "completed", "failed": "failed", "inprogress": "in_progress", "in_progress": "in_progress", } - normalized_status: BatchJobStatus = status_mapping.get( + normalized_status: Final[BatchJobStatus] = status_mapping.get( aws_status_lower, "failed" ) # Default to "failed" if unknown status @@ -1074,7 +1073,7 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj _, _, ) = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw) - result = LiteLLMBatch( + result: Final = LiteLLMBatch( id=status_response["invocationArn"], object="batch", status=normalized_status, @@ -1106,7 +1105,7 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj import concurrent.futures def run_in_thread(): - new_loop = asyncio.new_event_loop() + new_loop: Final = asyncio.new_event_loop() asyncio.set_event_loop(new_loop) try: return new_loop.run_until_complete(_async_get_status()) @@ -1114,5 +1113,5 @@ def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj new_loop.close() with concurrent.futures.ThreadPoolExecutor() as executor: - future = executor.submit(run_in_thread) + future: Final = executor.submit(run_in_thread) return future.result() diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index 26f888c8077..dcb5a7cc183 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -11,7 +11,7 @@ import json import os import threading import time -from typing import Literal, Optional +from typing import Final, Literal import litellm from litellm.constants import ( @@ -28,8 +28,8 @@ class BudgetManager: self, project_name: str, client_type: str = "local", - api_base: Optional[str] = None, - headers: Optional[dict] = None, + api_base: str | None = None, + headers: dict | None = None, ): self.client_type = client_type self.project_name = project_name @@ -60,8 +60,8 @@ class BudgetManager: self.print_verbose(f"user dict from local: {self.user_dict}") elif self.client_type == "hosted": # Load the user_dict from hosted db - url = self.api_base + "/get_budget" - data = {"project_name": self.project_name} + url: Final = self.api_base + "/get_budget" + data: Final = {"project_name": self.project_name} response = litellm.module_level_client.post(url, headers=self.headers, json=data) response = response.json() if response["status"] == "error": @@ -73,7 +73,7 @@ class BudgetManager: self, total_budget: float, user: str, - duration: Optional[Literal["daily", "weekly", "monthly", "yearly"]] = None, + duration: Literal["daily", "weekly", "monthly", "yearly"] | None = None, created_at: float = time.time(), ): self.user_dict[user] = {"total_budget": total_budget} @@ -100,11 +100,11 @@ class BudgetManager: return self.user_dict[user] def projected_cost(self, model: str, messages: list, user: str): - text = "".join(message["content"] for message in messages) - prompt_tokens = litellm.token_counter(model=model, text=text) + text: Final = "".join(message["content"] for message in messages) + prompt_tokens: Final = litellm.token_counter(model=model, text=text) prompt_cost, _ = litellm.cost_per_token(model=model, prompt_tokens=prompt_tokens, completion_tokens=0) - current_cost = self.user_dict[user].get("current_cost", 0) - projected_cost = prompt_cost + current_cost + current_cost: Final = self.user_dict[user].get("current_cost", 0) + projected_cost: Final = prompt_cost + current_cost return projected_cost def get_total_budget(self, user: str): @@ -113,10 +113,10 @@ class BudgetManager: def update_cost( self, user: str, - completion_obj: Optional[ModelResponse] = None, - model: Optional[str] = None, - input_text: Optional[str] = None, - output_text: Optional[str] = None, + completion_obj: ModelResponse | None = None, + model: str | None = None, + input_text: str | None = None, + output_text: str | None = None, ): if model and input_text and output_text: prompt_tokens = litellm.token_counter(model=model, messages=[{"role": "user", "content": input_text}]) @@ -178,11 +178,11 @@ class BudgetManager: def reset_on_duration(self, user: str): # Get current and creation time - last_updated_at = self.user_dict[user]["last_updated_at"] - current_time = time.time() + last_updated_at: Final = self.user_dict[user]["last_updated_at"] + current_time: Final = time.time() # Convert duration from days to seconds - duration_in_seconds = self.user_dict[user]["duration"] * HOURS_IN_A_DAY * 60 * 60 + duration_in_seconds: Final = self.user_dict[user]["duration"] * HOURS_IN_A_DAY * 60 * 60 # Check if duration has elapsed if current_time - last_updated_at >= duration_in_seconds: @@ -197,7 +197,7 @@ class BudgetManager: self.reset_on_duration(user) def _save_data_thread(self): - thread = threading.Thread(target=self.save_data) # [Non-Blocking]: saves data without blocking execution + thread: Final = threading.Thread(target=self.save_data) # [Non-Blocking]: saves data without blocking execution thread.start() def save_data(self): @@ -209,8 +209,8 @@ class BudgetManager: json.dump(self.user_dict, json_file, indent=4) # Indent for pretty formatting return {"status": "success"} elif self.client_type == "hosted": - url = self.api_base + "/set_budget" - data = {"project_name": self.project_name, "user_dict": self.user_dict} + url: Final = self.api_base + "/set_budget" + data: Final = {"project_name": self.project_name, "user_dict": self.user_dict} response = litellm.module_level_client.post(url, headers=self.headers, json=data) response = response.json() return response diff --git a/litellm/caching/__init__.py b/litellm/caching/__init__.py index bbe90b04121..87f4f7a7c63 100644 --- a/litellm/caching/__init__.py +++ b/litellm/caching/__init__.py @@ -2,10 +2,10 @@ from .azure_blob_cache import AzureBlobCache from .caching import Cache, LiteLLMCacheType from .disk_cache import DiskCache from .dual_cache import DualCache +from .gcs_cache import GCSCache from .in_memory_cache import InMemoryCache from .qdrant_semantic_cache import QdrantSemanticCache from .redis_cache import RedisCache from .redis_cluster_cache import RedisClusterCache from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache -from .gcs_cache import GCSCache diff --git a/litellm/caching/_embedding_router.py b/litellm/caching/_embedding_router.py index ec886b14020..1073b34ef25 100644 --- a/litellm/caching/_embedding_router.py +++ b/litellm/caching/_embedding_router.py @@ -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 diff --git a/litellm/caching/_internal_lru_cache.py b/litellm/caching/_internal_lru_cache.py index 54b0fe9690c..218ce2b9d79 100644 --- a/litellm/caching/_internal_lru_cache.py +++ b/litellm/caching/_internal_lru_cache.py @@ -1,11 +1,12 @@ +from collections.abc import Callable from functools import lru_cache -from typing import Callable, Optional, TypeVar +from typing import Final, TypeVar T = TypeVar("T") def lru_cache_wrapper( - maxsize: Optional[int] = None, + maxsize: int | None = None, ) -> Callable[[Callable[..., T]], Callable[..., T]]: """ Wrapper for lru_cache that caches success and exceptions @@ -20,7 +21,7 @@ def lru_cache_wrapper( return ("error", e) def wrapped(*args, **kwargs): - result = wrapper(*args, **kwargs) + result: Final = wrapper(*args, **kwargs) if result[0] == "error": raise result[1] return result[1] diff --git a/litellm/caching/azure_blob_cache.py b/litellm/caching/azure_blob_cache.py index fca7cf20313..742932731b3 100644 --- a/litellm/caching/azure_blob_cache.py +++ b/litellm/caching/azure_blob_cache.py @@ -11,6 +11,7 @@ Has 4 methods: import asyncio import json from contextlib import suppress +from typing import Final from litellm._logging import print_verbose, verbose_logger @@ -19,12 +20,12 @@ from .base_cache import BaseCache class AzureBlobCache(BaseCache): def __init__(self, account_url, container) -> None: - from azure.storage.blob import BlobServiceClient from azure.core.exceptions import ResourceExistsError from azure.identity import DefaultAzureCredential from azure.identity.aio import ( DefaultAzureCredential as AsyncDefaultAzureCredential, ) + from azure.storage.blob import BlobServiceClient from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient self.container_client = BlobServiceClient( @@ -41,7 +42,7 @@ class AzureBlobCache(BaseCache): def set_cache(self, key, value, **kwargs) -> None: print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}") - serialized_value = json.dumps(value) + serialized_value: Final = json.dumps(value) try: self.container_client.upload_blob(key, serialized_value) except Exception as e: @@ -50,7 +51,7 @@ class AzureBlobCache(BaseCache): async def async_set_cache(self, key, value, **kwargs) -> None: print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}") - serialized_value = json.dumps(value) + serialized_value: Final = json.dumps(value) try: await self.async_container_client.upload_blob(key, serialized_value, overwrite=True) except Exception as e: @@ -62,12 +63,15 @@ class AzureBlobCache(BaseCache): try: print_verbose(f"Get Azure Blob Cache: key: {key}") - as_bytes = self.container_client.download_blob(key).readall() - as_str = as_bytes.decode("utf-8") - cached_response = json.loads(as_str) + as_bytes: Final = self.container_client.download_blob(key).readall() + as_str: Final = as_bytes.decode("utf-8") + cached_response: Final = json.loads(as_str) verbose_logger.debug( - f"Got Azure Blob Cache: key: {key}, cached_response {cached_response}. Type Response {type(cached_response)}" + "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s", + key, + cached_response, + type(cached_response), ) return cached_response @@ -79,12 +83,15 @@ class AzureBlobCache(BaseCache): try: print_verbose(f"Get Azure Blob Cache: key: {key}") - blob = await self.async_container_client.download_blob(key) - as_bytes = await blob.readall() - as_str = as_bytes.decode("utf-8") - cached_response = json.loads(as_str) + blob: Final = await self.async_container_client.download_blob(key) + as_bytes: Final = await blob.readall() + as_str: Final = as_bytes.decode("utf-8") + cached_response: Final = json.loads(as_str) verbose_logger.debug( - f"Got Azure Blob Cache: key: {key}, cached_response {cached_response}. Type Response {type(cached_response)}" + "Got Azure Blob Cache: key: %s, cached_response %s. Type Response %s", + key, + cached_response, + type(cached_response), ) return cached_response except ResourceNotFoundError: @@ -99,7 +106,7 @@ class AzureBlobCache(BaseCache): await self.async_container_client.close() async def async_set_cache_pipeline(self, cache_list, **kwargs) -> None: - tasks = [] + tasks: Final = [] for val in cache_list: tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) await asyncio.gather(*tasks) diff --git a/litellm/caching/base_cache.py b/litellm/caching/base_cache.py index 81f1d61bd0d..6fe0609445f 100644 --- a/litellm/caching/base_cache.py +++ b/litellm/caching/base_cache.py @@ -9,7 +9,7 @@ Has 4 methods: """ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Final, Union if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -23,8 +23,8 @@ class BaseCache(ABC): def __init__(self, default_ttl: int = 60): self.default_ttl = default_ttl - def get_ttl(self, **kwargs) -> Optional[int]: - kwargs_ttl: Optional[int] = kwargs.get("ttl") + def get_ttl(self, **kwargs) -> int | None: + kwargs_ttl: Final[int | None] = kwargs.get("ttl") if kwargs_ttl is not None: try: return int(kwargs_ttl) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 34badaa3e8a..446b7f8be13 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -13,7 +13,7 @@ import json import time import traceback from enum import Enum -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Final from pydantic import BaseModel @@ -55,19 +55,18 @@ class CacheMode(str, Enum): class Cache: def __init__( self, - type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL, - mode: Optional[ - CacheMode - ] = CacheMode.default_on, # when default_on cache is always on, when default_off cache is opt in - host: Optional[str] = None, - port: Optional[str] = None, - password: Optional[str] = None, - namespace: Optional[str] = None, - ttl: Optional[float] = None, - default_in_memory_ttl: Optional[float] = None, - default_in_redis_ttl: Optional[float] = None, - similarity_threshold: Optional[float] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ + type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL, + mode: CacheMode + | None = CacheMode.default_on, # when default_on cache is always on, when default_off cache is opt in + host: str | None = None, + port: str | None = None, + password: str | None = None, + namespace: str | None = None, + ttl: float | None = None, + default_in_memory_ttl: float | None = None, + default_in_redis_ttl: float | None = None, + similarity_threshold: float | None = None, + supported_call_types: list[CachingSupportedCallTypes] | None = [ "completion", "acompletion", "embedding", @@ -82,38 +81,38 @@ class Cache: "aresponses", ], # s3 Bucket, boto3 configuration - azure_account_url: Optional[str] = None, - azure_blob_container: Optional[str] = None, - s3_bucket_name: Optional[str] = None, - s3_region_name: Optional[str] = None, - s3_api_version: Optional[str] = None, - s3_use_ssl: Optional[bool] = True, - s3_verify: Optional[Union[bool, str]] = None, - s3_endpoint_url: Optional[str] = None, - s3_aws_access_key_id: Optional[str] = None, - s3_aws_secret_access_key: Optional[str] = None, - s3_aws_session_token: Optional[str] = None, - s3_config: Optional[Any] = None, - s3_path: Optional[str] = None, - gcs_bucket_name: Optional[str] = None, - gcs_path_service_account: Optional[str] = None, - gcs_path: Optional[str] = None, + azure_account_url: str | None = None, + azure_blob_container: str | None = None, + s3_bucket_name: str | None = None, + s3_region_name: str | None = None, + s3_api_version: str | None = None, + s3_use_ssl: bool | None = True, + s3_verify: bool | str | None = None, + s3_endpoint_url: str | None = None, + s3_aws_access_key_id: str | None = None, + s3_aws_secret_access_key: str | None = None, + s3_aws_session_token: str | None = None, + s3_config: Any | None = None, + s3_path: str | None = None, + gcs_bucket_name: str | None = None, + gcs_path_service_account: str | None = None, + gcs_path: str | None = None, redis_semantic_cache_embedding_model: str = "text-embedding-ada-002", - redis_semantic_cache_index_name: Optional[str] = None, + redis_semantic_cache_index_name: str | None = None, valkey_semantic_cache_embedding_model: str = "text-embedding-ada-002", valkey_semantic_cache_index_name: str | None = None, - redis_flush_size: Optional[int] = None, - redis_startup_nodes: Optional[List] = None, - disk_cache_dir: Optional[str] = None, - qdrant_api_base: Optional[str] = None, - qdrant_api_key: Optional[str] = None, - qdrant_collection_name: Optional[str] = None, - qdrant_quantization_config: Optional[str] = None, + redis_flush_size: int | None = None, + redis_startup_nodes: list | None = None, + disk_cache_dir: str | None = None, + qdrant_api_base: str | None = None, + qdrant_api_key: str | None = None, + qdrant_collection_name: str | None = None, + qdrant_quantization_config: str | None = None, qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002", - qdrant_semantic_cache_vector_size: Optional[int] = None, + qdrant_semantic_cache_vector_size: int | None = None, # GCP IAM authentication parameters - gcp_service_account: Optional[str] = None, - gcp_ssl_ca_certs: Optional[str] = None, + gcp_service_account: str | None = None, + gcp_ssl_ca_certs: str | None = None, **kwargs, ): """ @@ -170,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, @@ -313,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: @@ -339,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: Optional[str] = self._get_param_value(param, kwargs) + param_value: str | None = self._get_param_value(param, kwargs) if param_value is not None: - cache_key += f"{str(param)}: {str(param_value)}" + cache_key += f"{param}: {param_value}" elif param not in litellm_param_kwargs: # check if user passed in optional param - e.g. top_k if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now if kwargs[param] is None: continue # ignore None params param_value = kwargs[param] - cache_key += f"{str(param)}: {str(param_value)}" + cache_key += f"{param}: {param_value}" if is_semantic_cache: cache_key += self._get_semantic_cache_tenant_scope(kwargs) @@ -374,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 @@ -382,7 +381,7 @@ class Cache: self, param: str, kwargs: dict, - ) -> Optional[str]: + ) -> str | None: """ Get the value for the given param from kwargs """ @@ -400,15 +399,15 @@ class Cache: 2. Else if a model_group is set, then return the model_group as the model. This is used for all requests sent through the litellm.Router() 3. Else use the `model` passed in kwargs """ - metadata: Dict = kwargs.get("metadata", {}) or {} - litellm_params: Dict = kwargs.get("litellm_params", {}) or {} - metadata_in_litellm_params: Dict = litellm_params.get("metadata", {}) or {} - model_group: Optional[str] = metadata.get("model_group") or metadata_in_litellm_params.get("model_group") - 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: Optional[str]) -> Optional[str]: - caching_groups: Optional[List] = metadata.get("caching_groups", []) + def _get_caching_group(self, metadata: dict, model_group: str | None) -> str | None: + caching_groups: Final[list | None] = metadata.get("caching_groups", []) if caching_groups: for group in caching_groups: if model_group in group: @@ -419,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) @@ -429,7 +428,7 @@ class Cache: or litellm_params.get("file_name") ) - def _get_preset_cache_key_from_kwargs(self, **kwargs) -> Optional[str]: + def _get_preset_cache_key_from_kwargs(self, **kwargs) -> str | None: """ Get the preset cache key from kwargs["litellm_params"] @@ -468,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 @@ -485,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": [ @@ -510,19 +509,19 @@ class Cache: def _get_cache_logic( self, - cached_result: Optional[Any], - max_age: Optional[float], + cached_result: Any | None, + max_age: float | None, ): """ Common get cache logic across sync + async implementations """ # 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: @@ -544,13 +543,13 @@ class Cache: return cached_result @staticmethod - def _get_safe_cache_lookup_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: - cache_lookup_kwargs: Dict[str, Any] = {} + def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + cache_lookup_kwargs: 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) @@ -558,17 +557,17 @@ class Cache: @staticmethod def _update_metadata_from_cache_lookup_kwargs( - original_kwargs: Dict[str, Any], cache_lookup_kwargs: Dict[str, Any] + original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any] ) -> None: - original_metadata = original_kwargs.get("metadata") - cache_lookup_metadata = cache_lookup_kwargs.get("metadata") + 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 if "semantic-similarity" in cache_lookup_metadata: original_metadata["semantic-similarity"] = cache_lookup_metadata["semantic-similarity"] - def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): + def get_cache(self, dynamic_cache_object: BaseCache | None = None, **kwargs): """ Retrieves the cached result for the given arguments. @@ -587,9 +586,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: @@ -603,7 +602,7 @@ class Cache: print_verbose(f"An exception occurred: {traceback.format_exc()}") return None - async def async_get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): + async def async_get_cache(self, dynamic_cache_object: BaseCache | None = None, **kwargs): """ Async get cache implementation. @@ -619,8 +618,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: @@ -647,13 +646,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") @@ -677,9 +676,9 @@ class Cache: cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) self.cache.set_cache(cache_key, cached_data, **kwargs) except Exception as e: - verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") + verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e) - async def async_add_cache(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): + async def async_add_cache(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs): """ Async implementation of add_cache """ @@ -696,14 +695,14 @@ class Cache: else: await self.cache.async_set_cache(cache_key, cached_data, **kwargs) except Exception as e: - verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") + verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e) def _convert_to_cached_embedding( self, embedding_response: Any, - model: Optional[str], - prompt_tokens: Optional[int] = None, - prompt_tokens_details: Optional[dict] = None, + model: str | None, + prompt_tokens: int | None = None, + prompt_tokens_details: dict | None = None, ) -> CachedEmbedding: """ Convert any embedding response into the standardized CachedEmbedding TypedDict format. @@ -745,7 +744,7 @@ class Cache: self, result: EmbeddingResponse, idx_in_result_data: int, - ) -> Optional[dict]: + ) -> dict | None: """ Extract per-item prompt_tokens_details from a response for caching. @@ -757,7 +756,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): @@ -768,12 +767,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) @@ -788,7 +787,7 @@ class Cache: self, result: EmbeddingResponse, idx_in_result_data: int, - ) -> Optional[int]: + ) -> int | None: """ Extract the per-item prompt_tokens from a response for caching. @@ -799,8 +798,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,24 +812,24 @@ class Cache: input: str, kwargs: dict, idx_in_result_data: int = 0, - ) -> Tuple[str, dict, dict]: - preset_cache_key = self.get_cache_key(**{**kwargs, "input": input}) + ) -> tuple[str, dict, dict]: + 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, @@ -843,7 +842,7 @@ class Cache: ) return cache_key, cached_data, kwargs - async def async_add_cache_pipeline(self, result, dynamic_cache_object: Optional[BaseCache] = None, **kwargs): + async def async_add_cache_pipeline(self, result, dynamic_cache_object: BaseCache | None = None, **kwargs): """ Async implementation of add_cache for Embedding calls @@ -857,7 +856,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"]): ( @@ -875,7 +874,7 @@ class Cache: else: await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) except Exception as e: - verbose_logger.exception(f"LiteLLM Cache: Excepton add_cache: {str(e)}") + verbose_logger.exception("LiteLLM Cache: Excepton add_cache: %s", e) def should_use_cache(self, **kwargs): """ @@ -888,7 +887,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: @@ -900,13 +899,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 @@ -926,11 +925,11 @@ class Cache: def enable_cache( - type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL, - host: Optional[str] = None, - port: Optional[str] = None, - password: Optional[str] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ + type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL, + host: str | None = None, + port: str | None = None, + password: str | None = None, + supported_call_types: list[CachingSupportedCallTypes] | None = [ "completion", "acompletion", "embedding", @@ -986,11 +985,11 @@ def enable_cache( def update_cache( - type: Optional[LiteLLMCacheType] = LiteLLMCacheType.LOCAL, - host: Optional[str] = None, - port: Optional[str] = None, - password: Optional[str] = None, - supported_call_types: Optional[List[CachingSupportedCallTypes]] = [ + type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL, + host: str | None = None, + port: str | None = None, + password: str | None = None, + supported_call_types: list[CachingSupportedCallTypes] | None = [ "completion", "acompletion", "embedding", diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index d8a2d2d76b7..4747aac54c6 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -18,18 +18,8 @@ import asyncio import datetime import inspect import time -from typing import ( - TYPE_CHECKING, - Any, - AsyncGenerator, - Callable, - Dict, - Generator, - List, - Optional, - Tuple, - Union, -) +from collections.abc import AsyncGenerator, Callable, Generator +from typing import TYPE_CHECKING, Any, Final, Optional from pydantic import BaseModel @@ -77,12 +67,12 @@ class CachingHandlerResponse(BaseModel): For embeddings there can be a cache hit for some of the inputs in the list and a cache miss for others """ - cached_result: Optional[Any] = None - final_embedding_cached_response: Optional[EmbeddingResponse] = None + cached_result: Any | None = None + final_embedding_cached_response: EmbeddingResponse | None = None embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call -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]: @@ -102,16 +92,16 @@ 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 -def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool: +def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bool: """ When stream=True, do not run success callbacks at cache-hit time. @@ -127,25 +117,24 @@ class LLMCachingHandler: def __init__( self, original_function: Callable, - request_kwargs: Dict[str, Any], + request_kwargs: dict[str, Any], start_time: datetime.datetime, ): from litellm.caching import DualCache, RedisCache - self.async_streaming_chunks: List[ModelResponse] = [] - self.sync_streaming_chunks: List[ModelResponse] = [] + self.async_streaming_chunks: list[ModelResponse] = [] + self.sync_streaming_chunks: list[ModelResponse] = [] self.request_kwargs = _drop_logging_obj_from_kwargs(request_kwargs) - self.preset_cache_key: Optional[str] = None + self.preset_cache_key: str | None = None self.original_function = original_function self.start_time = start_time if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache): - self.dual_cache: Optional[DualCache] = DualCache( + self.dual_cache: DualCache | None = DualCache( redis_cache=litellm.cache.cache, in_memory_cache=in_memory_cache_obj, ) else: self.dual_cache = None - pass async def _async_get_cache( self, @@ -154,9 +143,9 @@ class LLMCachingHandler: logging_obj: LiteLLMLoggingObj, start_time: datetime.datetime, call_type: str, - kwargs: Dict[str, Any], - args: Optional[Tuple[Any, ...]] = None, - ) -> Optional[CachingHandlerResponse]: + kwargs: dict[str, Any], + args: tuple[Any, ...] | None = None, + ) -> CachingHandlerResponse | None: """ Internal method to get from the cache. Handles different call types (embeddings, chat/completions, text_completion, transcription) @@ -184,17 +173,17 @@ class LLMCachingHandler: kwargs.get("cache", {}).get("no-cache", False) is not True ): # allow users to control returning cached responses from the completion function args = args or () - final_embedding_cached_response: Optional[EmbeddingResponse] = None + final_embedding_cached_response: EmbeddingResponse | None = None embedding_all_elements_cache_hit: bool = False - cached_result: Optional[Any] = None + cached_result: Any | None = None kwargs = kwargs.copy() ######################################################### # Init cache timing metrics ######################################################### - cache_check_start_time = time.perf_counter() - cache_check_end_time: Optional[float] = None + cache_check_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): @@ -208,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, @@ -247,7 +236,7 @@ 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) @@ -278,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, @@ -293,16 +282,16 @@ class LLMCachingHandler: logging_obj: LiteLLMLoggingObj, start_time: datetime.datetime, call_type: str, - kwargs: Dict[str, Any], - args: Optional[Tuple[Any, ...]] = None, + kwargs: dict[str, Any], + args: tuple[Any, ...] | None = None, ) -> CachingHandlerResponse: - cached_result: Optional[Any] = None + cached_result: Any | None = None # Check if caching should be performed BEFORE doing expensive kwargs copy if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): 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, @@ -333,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, @@ -361,7 +350,7 @@ 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) @@ -371,7 +360,7 @@ class LLMCachingHandler: return CachingHandlerResponse(cached_result=cached_result) return CachingHandlerResponse(cached_result=cached_result) - def handle_kwargs_input_list_or_str(self, kwargs: Dict[str, Any]) -> List[str]: + def handle_kwargs_input_list_or_str(self, kwargs: dict[str, Any]) -> list[str]: """ Handles the input of kwargs['input'] being a list or a string """ @@ -382,7 +371,7 @@ class LLMCachingHandler: else: raise ValueError("input must be a string or a list") - def _extract_model_from_cached_results(self, non_null_list: List[Tuple[int, CachedEmbedding]]) -> Optional[str]: + def _extract_model_from_cached_results(self, non_null_list: list[tuple[int, CachedEmbedding]]) -> str | None: """ Helper method to extract the model name from cached results. @@ -399,13 +388,13 @@ class LLMCachingHandler: def _process_async_embedding_cached_response( self, - final_embedding_cached_response: Optional[EmbeddingResponse], - cached_result: List[Optional[CachedEmbedding]], - kwargs: Dict[str, Any], + final_embedding_cached_response: EmbeddingResponse | None, + cached_result: list[CachedEmbedding | None], + kwargs: dict[str, Any], logging_obj: LiteLLMLoggingObj, start_time: datetime.datetime, model: str, - ) -> Tuple[Optional[EmbeddingResponse], bool]: + ) -> tuple[EmbeddingResponse | None, bool]: """ Returns the final embedding cached response and a boolean indicating if all elements in the list have a cache hit @@ -427,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]) @@ -448,7 +437,7 @@ class LLMCachingHandler: final_embedding_cached_response._hidden_params["cache_hit"] = True prompt_tokens = 0 - aggregated_details: Optional[dict] = None + aggregated_details: dict | None = None for val in non_null_list: idx, cr = val # (idx, cr) tuple if cr is not None: @@ -478,15 +467,15 @@ class LLMCachingHandler: aggregated_details[key] = value ## USAGE - prompt_tokens_details: Optional["PromptTokensDetailsWrapper"] = None - if aggregated_details: - from litellm.types.utils import PromptTokensDetailsWrapper + from litellm.types.utils import PromptTokensDetailsWrapper + prompt_tokens_details: PromptTokensDetailsWrapper | None = None + if aggregated_details: try: prompt_tokens_details = PromptTokensDetailsWrapper(**aggregated_details) except Exception: prompt_tokens_details = None - usage = Usage( + usage: Final = Usage( prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens, @@ -495,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, @@ -553,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) @@ -614,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]) @@ -676,9 +665,7 @@ class LLMCachingHandler: cache_hit=cache_hit, ) - async def _retrieve_from_cache( - self, call_type: str, kwargs: Dict[str, Any], args: Tuple[Any, ...] - ) -> Optional[Any]: + async def _retrieve_from_cache(self, call_type: str, kwargs: dict[str, Any], args: tuple[Any, ...]) -> Any | None: """ Internal method to - get cache key @@ -699,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, @@ -711,13 +698,13 @@ class LLMCachingHandler: if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) - cached_result: Optional[Any] = None + cached_result: Any | None = None if call_type == CallTypes.aembedding.value: if isinstance(new_kwargs["input"], str): new_kwargs["input"] = [new_kwargs["input"]] 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( @@ -733,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) @@ -756,21 +743,20 @@ class LLMCachingHandler: self, cached_result: Any, call_type: str, - kwargs: Dict[str, Any], + kwargs: dict[str, Any], logging_obj: LiteLLMLoggingObj, model: str, - args: Tuple[Any, ...], - custom_llm_provider: Optional[str] = None, - ) -> Optional[ - Union[ - ModelResponse, - TextCompletionResponse, - EmbeddingResponse, - RerankResponse, - TranscriptionResponse, - CustomStreamWrapper, - ] - ]: + args: tuple[Any, ...], + custom_llm_provider: str | None = None, + ) -> ( + ModelResponse + | TextCompletionResponse + | EmbeddingResponse + | RerankResponse + | TranscriptionResponse + | CustomStreamWrapper + | None + ): """ Internal method to process the cached result @@ -838,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, @@ -850,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( @@ -872,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 @@ -923,7 +909,7 @@ class LLMCachingHandler: convert_to_streaming_response_async, ) - _stream_cached_result: Union[AsyncGenerator, Generator] + _stream_cached_result: AsyncGenerator | Generator if call_type == CallTypes.acompletion.value or call_type == CallTypes.atext_completion.value: _stream_cached_result = convert_to_streaming_response_async( response_object=cached_result, @@ -943,8 +929,8 @@ class LLMCachingHandler: self, result: Any, original_function: Callable, - kwargs: Dict[str, Any], - args: Optional[Tuple[Any, ...]] = None, + kwargs: dict[str, Any], + args: tuple[Any, ...] | None = None, ): """ Internal method to check the type of the result & cache used and adds the result to the cache accordingly @@ -967,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): @@ -1009,14 +995,14 @@ class LLMCachingHandler: def sync_set_cache( self, result: Any, - kwargs: Dict[str, Any], - args: Optional[Tuple[Any, ...]] = None, + kwargs: dict[str, Any], + args: tuple[Any, ...] | None = None, ): """ Sync internal method to add the result to the cache """ - new_kwargs = kwargs.copy() + new_kwargs: Final = kwargs.copy() new_kwargs.update( convert_args_to_kwargs( self.original_function, @@ -1031,7 +1017,7 @@ class LLMCachingHandler: return - def _should_store_result_in_cache(self, original_function: Callable, kwargs: Dict[str, Any]) -> bool: + def _should_store_result_in_cache(self, original_function: Callable, kwargs: dict[str, Any]) -> bool: """ Helper function to determine if the result should be stored in the cache. @@ -1077,7 +1063,7 @@ class LLMCachingHandler: """ - complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = ( + complete_streaming_response: Final[ModelResponse | TextCompletionResponse | None] = ( _assemble_complete_response_from_streaming_chunks( result=processed_chunk, start_time=self.start_time, @@ -1099,7 +1085,7 @@ class LLMCachingHandler: """ Sync internal method to add the streaming response to the cache """ - complete_streaming_response: Optional[Union[ModelResponse, TextCompletionResponse]] = ( + complete_streaming_response: Final[ModelResponse | TextCompletionResponse | None] = ( _assemble_complete_response_from_streaming_chunks( result=processed_chunk, start_time=self.start_time, @@ -1121,12 +1107,12 @@ class LLMCachingHandler: self, logging_obj: LiteLLMLoggingObj, model: str, - kwargs: Dict[str, Any], + kwargs: dict[str, Any], cached_result: Any, is_async: bool, is_embedding: bool = False, - custom_llm_provider: Optional[str] = None, - cache_duration_ms: Optional[float] = None, + custom_llm_provider: str | None = None, + cache_duration_ms: float | None = None, ): """ Helper function to update the LiteLLMLoggingObj environment variables. @@ -1143,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", ""), @@ -1180,16 +1166,16 @@ class LLMCachingHandler: def convert_args_to_kwargs( original_function: Callable, - args: Optional[Tuple[Any, ...]] = None, -) -> Dict[str, Any]: + args: tuple[Any, ...] | None = None, +) -> dict[str, Any]: # Get the signature of the original function - signature = inspect.signature(original_function) + 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): diff --git a/litellm/caching/disk_cache.py b/litellm/caching/disk_cache.py index af8eb92849f..895f276eb20 100644 --- a/litellm/caching/disk_cache.py +++ b/litellm/caching/disk_cache.py @@ -1,5 +1,5 @@ import json -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Final, Union from .base_cache import BaseCache @@ -12,7 +12,7 @@ else: class DiskCache(BaseCache): - def __init__(self, disk_cache_dir: Optional[str] = None): + def __init__(self, disk_cache_dir: str | None = None): try: import diskcache as dc except ModuleNotFoundError as e: @@ -41,7 +41,7 @@ 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 @@ -51,7 +51,7 @@ class DiskCache(BaseCache): 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) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 0e3c93946fd..3b181ca23ff 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -13,7 +13,7 @@ import time import traceback from concurrent.futures import ThreadPoolExecutor from threading import Lock -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Final, Union if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -57,11 +57,11 @@ class DualCache(BaseCache): def __init__( self, - in_memory_cache: Optional[InMemoryCache] = None, - redis_cache: Optional[RedisCache] = None, - default_in_memory_ttl: Optional[float] = None, - default_redis_ttl: Optional[float] = None, - default_redis_batch_cache_expiry: Optional[float] = None, + in_memory_cache: InMemoryCache | None = None, + redis_cache: RedisCache | None = None, + default_in_memory_ttl: float | None = None, + default_redis_ttl: float | None = None, + default_redis_batch_cache_expiry: float | None = None, default_max_redis_batch_cache_size: int = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, ) -> None: super().__init__() @@ -77,7 +77,7 @@ class DualCache(BaseCache): self.default_in_memory_ttl = default_in_memory_ttl or litellm.default_in_memory_ttl self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl - def update_cache_ttl(self, default_in_memory_ttl: Optional[float], default_redis_ttl: Optional[float]): + def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None): if default_in_memory_ttl is not None: self.default_in_memory_ttl = default_in_memory_ttl @@ -86,9 +86,9 @@ class DualCache(BaseCache): def attach_redis_cache( self, - redis_cache: Optional[RedisCache] = None, + redis_cache: RedisCache | None = None, *, - default_redis_ttl: Optional[float] = None, + default_redis_ttl: float | None = None, ) -> None: """ Attach a Redis backend if this DualCache does not already have one. @@ -147,13 +147,13 @@ class DualCache(BaseCache): return result except Exception as e: - verbose_logger.error(f"LiteLLM Cache: Excepton async add_cache: {str(e)}") + verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e) raise e def get_cache( self, key, - parent_otel_span: Optional[Span] = None, + parent_otel_span: Span | None = None, local_only: bool = False, **kwargs, ): @@ -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 @@ -184,16 +184,16 @@ class DualCache(BaseCache): def batch_get_cache( self, keys: list, - parent_otel_span: Optional[Span] = None, + parent_otel_span: Span | None = None, local_only: bool = False, **kwargs, ): - 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: @@ -217,7 +217,7 @@ class DualCache(BaseCache): async def async_get_cache( self, key, - parent_otel_span: Optional[Span] = None, + parent_otel_span: Span | None = None, local_only: bool = False, **kwargs, ): @@ -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 @@ -250,15 +250,15 @@ class DualCache(BaseCache): def _reserve_redis_batch_keys( self, current_time: float, - keys: List[str], - result: List[Any], - ) -> Tuple[List[str], Dict[str, Optional[float]]]: + keys: list[str], + result: list[Any], + ) -> tuple[list[str], dict[str, float | None]]: """ Atomically choose keys to fetch from Redis and reserve their access time. This prevents check-then-act races under concurrent async callers. """ - sublist_keys: List[str] = [] - previous_access_times: Dict[str, Optional[float]] = {} + sublist_keys: 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): @@ -275,7 +275,7 @@ class DualCache(BaseCache): return sublist_keys, previous_access_times - def _rollback_redis_batch_key_reservations(self, previous_access_times: Dict[str, Optional[float]]) -> None: + def _rollback_redis_batch_key_reservations(self, previous_access_times: dict[str, float | None]) -> None: with self._last_redis_batch_access_time_lock: for key, previous_time in previous_access_times.items(): if previous_time is None: @@ -286,14 +286,14 @@ class DualCache(BaseCache): async def async_batch_get_cache( self, keys: list, - parent_otel_span: Optional[Span] = None, + parent_otel_span: Span | None = None, local_only: bool = False, **kwargs, ): 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: {str(e)}") + verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e) # async_batch_set_cache async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs): @@ -366,17 +366,17 @@ class DualCache(BaseCache): cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs ) except Exception as e: - verbose_logger.exception(f"LiteLLM Cache: Excepton async add_cache: {str(e)}") + verbose_logger.exception("LiteLLM Cache: Excepton async add_cache: %s", e) async def async_increment_cache( self, key, value: float, - parent_otel_span: Optional[Span] = None, + parent_otel_span: Span | None = None, local_only: bool = False, refresh_ttl: bool = False, **kwargs, - ) -> Optional[float]: + ) -> float | None: """ Key - the key in cache @@ -388,7 +388,7 @@ class DualCache(BaseCache): Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ - result: Optional[float] = None + result: float | None = None try: if self.in_memory_cache is not None: result = await self.in_memory_cache.async_increment(key, value, **kwargs) @@ -412,12 +412,12 @@ class DualCache(BaseCache): async def async_increment_cache_pipeline( self, - increment_list: List["RedisPipelineIncrementOperation"], + increment_list: list["RedisPipelineIncrementOperation"], local_only: bool = False, - parent_otel_span: Optional[Span] = None, + parent_otel_span: Span | None = None, **kwargs, - ) -> Optional[List[float]]: - result: Optional[List[float]] = None + ) -> list[float] | None: + result: list[float] | None = None try: if self.in_memory_cache is not None: result = await self.in_memory_cache.async_increment_pipeline( @@ -439,7 +439,7 @@ class DualCache(BaseCache): ) return result - async def async_set_cache_sadd(self, key, value: List, local_only: bool = False, **kwargs) -> None: + async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None: """ Add value to a set @@ -456,7 +456,7 @@ class DualCache(BaseCache): if self.redis_cache is not None and local_only is False: _ = await self.redis_cache.async_set_cache_sadd(key, value, ttl=kwargs.get("ttl", None)) - return None + return except Exception as e: raise e # don't log, if exception is raised @@ -484,7 +484,7 @@ class DualCache(BaseCache): if self.redis_cache is not None: await self.redis_cache.async_delete_cache(key) - async def async_get_ttl(self, key: str) -> Optional[int]: + async def async_get_ttl(self, key: str) -> int | None: """ Get the remaining TTL of a key in in-memory cache or redis """ diff --git a/litellm/caching/gcs_cache.py b/litellm/caching/gcs_cache.py index 3345f8fc5eb..e9922218828 100644 --- a/litellm/caching/gcs_cache.py +++ b/litellm/caching/gcs_cache.py @@ -2,27 +2,28 @@ Supports syncing responses to Google Cloud Storage Buckets using HTTP requests. """ -import json import asyncio -from typing import Optional +import json +from typing import Final from urllib.parse import quote from litellm._logging import print_verbose, verbose_logger from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.llms.custom_httpx.http_handler import ( - get_async_httpx_client, _get_httpx_client, + get_async_httpx_client, httpxSpecialProvider, ) + from .base_cache import BaseCache class GCSCache(BaseCache): def __init__( self, - bucket_name: Optional[str] = None, - path_service_account: Optional[str] = None, - gcs_path: Optional[str] = None, + bucket_name: str | None = None, + path_service_account: str | None = None, + gcs_path: str | None = None, ) -> None: super().__init__() self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME @@ -33,7 +34,7 @@ class GCSCache(BaseCache): self.sync_client = _get_httpx_client() def _construct_headers(self) -> dict: - base = GCSBucketBase(bucket_name=self.bucket_name) + base: Final = GCSBucketBase(bucket_name=self.bucket_name) base.path_service_account_json = self.path_service_account base.BUCKET_NAME = self.bucket_name return base.sync_construct_request_headers() @@ -41,55 +42,58 @@ class GCSCache(BaseCache): def set_cache(self, key, value, **kwargs): try: print_verbose(f"LiteLLM SET Cache - GCS. Key={key}. Value={value}") - headers = self._construct_headers() - object_name = self.key_prefix + key - bucket_name = self.bucket_name + headers: Final = self._construct_headers() + object_name: Final = self.key_prefix + key + bucket_name: Final = self.bucket_name url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={quote(object_name, safe='')}" - data = json.dumps(value) + data: Final = json.dumps(value) self.sync_client.post(url=url, data=data, headers=headers) except Exception as e: print_verbose(f"GCS Caching: set_cache() - Got exception from GCS: {e}") async def async_set_cache(self, key, value, **kwargs): try: - headers = self._construct_headers() - object_name = self.key_prefix + key - bucket_name = self.bucket_name + headers: Final = self._construct_headers() + object_name: Final = self.key_prefix + key + bucket_name: Final = self.bucket_name url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={quote(object_name, safe='')}" - data = json.dumps(value) + data: Final = json.dumps(value) await self.async_client.post(url=url, data=data, headers=headers) except Exception as e: print_verbose(f"GCS Caching: async_set_cache() - Got exception from GCS: {e}") def get_cache(self, key, **kwargs): try: - headers = self._construct_headers() - object_name = self.key_prefix + key - bucket_name = self.bucket_name + headers: Final = self._construct_headers() + object_name: Final = self.key_prefix + key + bucket_name: Final = self.bucket_name url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media" - response = self.sync_client.get(url=url, headers=headers) + response: Final = self.sync_client.get(url=url, headers=headers) if response.status_code == 200: - cached_response = json.loads(response.text) + cached_response: Final = json.loads(response.text) verbose_logger.debug( - f"Got GCS Cache: key: {key}, cached_response {cached_response}. Type Response {type(cached_response)}" + "Got GCS Cache: key: %s, cached_response %s. Type Response %s", + key, + cached_response, + type(cached_response), ) return cached_response return None except Exception as e: - verbose_logger.error(f"GCS Caching: get_cache() - Got exception from GCS: {e}") + verbose_logger.error("GCS Caching: get_cache() - Got exception from GCS: %s", e) async def async_get_cache(self, key, **kwargs): try: - headers = self._construct_headers() - object_name = self.key_prefix + key - bucket_name = self.bucket_name + headers: Final = self._construct_headers() + object_name: Final = self.key_prefix + key + bucket_name: Final = self.bucket_name url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media" - response = await self.async_client.get(url=url, headers=headers) + response: Final = await self.async_client.get(url=url, headers=headers) if response.status_code == 200: return json.loads(response.text) return None except Exception as e: - verbose_logger.error(f"GCS Caching: async_get_cache() - Got exception from GCS: {e}") + verbose_logger.error("GCS Caching: async_get_cache() - Got exception from GCS: %s", e) def flush_cache(self): pass @@ -98,7 +102,7 @@ class GCSCache(BaseCache): pass async def async_set_cache_pipeline(self, cache_list, **kwargs): - tasks = [] + tasks: Final = [] for val in cache_list: tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) await asyncio.gather(*tasks) diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 36b477f7a8b..38a9966f9f9 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -8,12 +8,12 @@ Has 4 methods: - async_get_cache """ +import heapq import json import sys -import time -import heapq import threading -from typing import TYPE_CHECKING, Any, List, Optional +import time +from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -28,11 +28,10 @@ from .base_cache import BaseCache class InMemoryCache(BaseCache): def __init__( self, - max_size_in_memory: Optional[int] = 200, - default_ttl: Optional[ - int - ] = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute - max_size_per_item: Optional[int] = 1024, # 1MB = 1024KB + max_size_in_memory: int | None = 200, + default_ttl: int + | None = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute + max_size_per_item: int | None = 1024, # 1MB = 1024KB ): """ max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default @@ -68,7 +67,7 @@ class InMemoryCache(BaseCache): # Handle special types without full conversion when possible if hasattr(value, "__sizeof__"): # Use __sizeof__ if available - size = value.__sizeof__() / 1024 + size: Final = value.__sizeof__() / 1024 return size <= self.max_size_per_item # Fallback for complex types @@ -112,7 +111,7 @@ class InMemoryCache(BaseCache): - 3. the size of in-memory cache is bounded """ - current_time = time.time() + current_time: Final = time.time() # Step 1: Remove expired or outdated items while self.expiration_heap: @@ -145,10 +144,8 @@ class InMemoryCache(BaseCache): """ Check if ttl is set for a key """ - ttl_time = self.ttl_dict.get(key) - if ttl_time is None: # if ttl is not set, allow override - return True - elif float(ttl_time) < time.time(): # if ttl is expired, allow override + ttl_time: Final = self.ttl_dict.get(key) + if ttl_time is None or float(ttl_time) < time.time(): # if ttl is not set, allow override return True else: return False @@ -184,12 +181,12 @@ class InMemoryCache(BaseCache): else: self.set_cache(key=cache_key, value=cache_value) - async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float]): + async def async_set_cache_sadd(self, key, value: list, ttl: float | None): """ Add value to set """ # get the value - init_value = self.get_cache(key=key) or set() + init_value: Final = self.get_cache(key=key) or set() for val in value: init_value.add(val) self.set_cache(key, init_value, ttl=ttl) @@ -210,7 +207,7 @@ class InMemoryCache(BaseCache): if key in self.cache_dict: if self.evict_element_if_expired(key): return None - original_cached_response = self.cache_dict[key] + original_cached_response: Final = self.cache_dict[key] try: cached_response = json.loads(original_cached_response) except Exception: @@ -219,7 +216,7 @@ class InMemoryCache(BaseCache): 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) @@ -228,7 +225,7 @@ class InMemoryCache(BaseCache): def increment_cache(self, key, value: float, **kwargs) -> float: with self._increment_lock: # keep read-modify-write atomic - init_value = self.get_cache(key=key) or 0 + init_value: Final = self.get_cache(key=key) or 0 value = init_value + value self.set_cache(key, value, **kwargs) return value @@ -237,7 +234,7 @@ class InMemoryCache(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) @@ -247,9 +244,9 @@ class InMemoryCache(BaseCache): return self.increment_cache(key=key, value=value, **kwargs) async def async_increment_pipeline( - self, increment_list: List["RedisPipelineIncrementOperation"], **kwargs - ) -> Optional[List[float]]: - results = [] + self, increment_list: list["RedisPipelineIncrementOperation"], **kwargs + ) -> list[float] | None: + results: Final = [] for increment in increment_list: result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs) results.append(result) @@ -266,16 +263,16 @@ class InMemoryCache(BaseCache): def delete_cache(self, key): self._remove_key(key) - async def async_get_ttl(self, key: str) -> Optional[int]: + async def async_get_ttl(self, key: str) -> int | None: """ Get the remaining TTL of a key in in-memory cache """ return self.ttl_dict.get(key, None) - async def async_get_oldest_n_keys(self, n: int) -> List[str]: + async def async_get_oldest_n_keys(self, n: int) -> list[str]: """ Get the oldest n keys in the cache """ # sorted ttl dict by ttl - sorted_ttl_dict = sorted(self.ttl_dict.items(), key=lambda x: x[1]) + sorted_ttl_dict: Final = sorted(self.ttl_dict.items(), key=lambda x: x[1]) return [key for key, _ in sorted_ttl_dict[:n]] diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index c2274713bb9..7d072a40195 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -3,6 +3,7 @@ Add the event loop to the cache key, to prevent event loop closed errors. """ import asyncio +from typing import Final from .in_memory_cache import InMemoryCache @@ -25,8 +26,8 @@ class LLMClientCache(InMemoryCache): If none, use the key as is. """ try: - event_loop = asyncio.get_running_loop() - stringified_event_loop = str(id(event_loop)) + event_loop: Final = asyncio.get_running_loop() + stringified_event_loop: Final = str(id(event_loop)) return f"{key}-{stringified_event_loop}" except RuntimeError: # handle no current running event loop return key diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 5ed1bb47eba..8f8323550f3 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -12,7 +12,7 @@ import ast import asyncio import json import os -from typing import Any, Dict, cast +from typing import Any, Final, cast import litellm from litellm._logging import print_verbose @@ -88,7 +88,7 @@ class QdrantSemanticCache(BaseCache): if quantization_config is None: print_verbose("Quantization config is not provided. Default binary quantization will be used.") - collection_exists = self.sync_client.get( + collection_exists: Final = self.sync_client.get( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/exists", headers=self.headers, ) @@ -104,7 +104,7 @@ class QdrantSemanticCache(BaseCache): print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}") self._ensure_cache_key_payload_index() else: - quantization_params: Dict[str, Any] + quantization_params: dict[str, Any] if quantization_config is None or quantization_config == "binary": quantization_params = { "binary": { @@ -124,7 +124,7 @@ class QdrantSemanticCache(BaseCache): else: raise Exception("Quantization config must be one of 'scalar', 'binary' or 'product'") - new_collection_status = self.sync_client.put( + new_collection_status: Final = self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}", json={ "vectors": {"size": self.vector_size, "distance": "Cosine"}, @@ -167,7 +167,7 @@ class QdrantSemanticCache(BaseCache): def _ensure_cache_key_payload_index(self) -> None: try: - response = self.sync_client.put( + response: Final = self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/index", headers=self.headers, json={ @@ -178,17 +178,17 @@ class QdrantSemanticCache(BaseCache): if response.status_code not in (200, 201): print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {response.text}") except Exception as exc: - print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {str(exc)}") + print_verbose(f"Qdrant semantic-cache could not create cache-key payload index: {exc}") def _payload_matches_cache_key(self, payload: dict, key: str) -> bool: # Pre-isolation points stored only prompt + response with no cache-key # payload field. Reassigning them to a caller's key would risk # cross-scope hits, so they're treated as misses and re-populated on # the next set_cache. - cached_key = payload.get(self.CACHE_KEY_FIELD_NAME) + cached_key: Final = payload.get(self.CACHE_KEY_FIELD_NAME) return cached_key is not None and str(cached_key) == str(key) - def _get_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse: + def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse: """Embed via the proxy Router when it serves the model, else direct.""" try: from litellm.proxy.proxy_server import llm_model_list, llm_router @@ -196,7 +196,7 @@ class QdrantSemanticCache(BaseCache): llm_model_list = None llm_router = None - router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) + router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) if router is not None: return router.embedding( model=self.embedding_model, @@ -210,14 +210,14 @@ class QdrantSemanticCache(BaseCache): cache={"no-store": True, "no-cache": True}, ) - async def _get_async_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> EmbeddingResponse: + async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse: try: from litellm.proxy.proxy_server import llm_model_list, llm_router except ImportError: llm_model_list = None llm_router = None - router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) + router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) if router is not None: return await router.aembedding( model=self.embedding_model, @@ -237,22 +237,22 @@ class QdrantSemanticCache(BaseCache): from litellm._uuid import uuid # get the prompt - messages = kwargs["messages"] - prompt = get_str_from_messages(messages) + messages: Final = kwargs["messages"] + prompt: Final = get_str_from_messages(messages) # create an embedding for prompt - embedding_response = cast( + embedding_response: Final = cast( EmbeddingResponse, self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding - embedding = embedding_response["data"][0]["embedding"] + embedding: Final = embedding_response["data"][0]["embedding"] value = str(value) assert isinstance(value, str) - data = { + data: Final = { "points": [ { "id": str(uuid.uuid4()), @@ -270,25 +270,24 @@ class QdrantSemanticCache(BaseCache): headers=self.headers, json=data, ) - return def get_cache(self, key, **kwargs): print_verbose(f"sync qdrant semantic-cache get_cache, kwargs: {kwargs}") # get the messages - messages = kwargs["messages"] - prompt = get_str_from_messages(messages) + messages: Final = kwargs["messages"] + prompt: Final = get_str_from_messages(messages) # convert to embedding - embedding_response = cast( + embedding_response: Final = cast( EmbeddingResponse, self._get_embedding(prompt, metadata=kwargs.get("metadata")), ) # get the embedding - embedding = embedding_response["data"][0]["embedding"] + embedding: Final = embedding_response["data"][0]["embedding"] - data = { + data: Final = { "vector": embedding, "params": { "quantization": { @@ -302,12 +301,12 @@ class QdrantSemanticCache(BaseCache): } self._add_cache_key_filter_to_search_data(data=data, key=key) - search_response = self.sync_client.post( + search_response: Final = self.sync_client.post( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points/search", headers=self.headers, json=data, ) - results = search_response.json()["result"] + results: Final = search_response.json()["result"] if results is None: kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 @@ -317,14 +316,14 @@ class QdrantSemanticCache(BaseCache): kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - similarity = results[0]["score"] - payload = results[0]["payload"] + similarity: Final = results[0]["score"] + payload: Final = results[0]["payload"] if not self._payload_matches_cache_key(payload=payload, key=key): print_verbose("Qdrant semantic-cache hit did not match cache key scope") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - cached_prompt = payload["text"] + cached_prompt: Final = payload["text"] # check similarity, if more than self.similarity_threshold, return results print_verbose( @@ -336,7 +335,7 @@ class QdrantSemanticCache(BaseCache): if similarity >= self.similarity_threshold: # cache hit ! - cached_value = payload["response"] + cached_value: Final = payload["response"] print_verbose( f"got a cache hit, similarity: {similarity}, Current prompt: {prompt}, cached_prompt: {cached_prompt}" ) @@ -344,7 +343,6 @@ class QdrantSemanticCache(BaseCache): else: # cache miss ! return None - pass async def async_set_cache(self, key, value, **kwargs): from litellm._uuid import uuid @@ -352,17 +350,17 @@ class QdrantSemanticCache(BaseCache): print_verbose(f"async qdrant semantic-cache set_cache, kwargs: {kwargs}") # get the prompt - messages = kwargs["messages"] - prompt = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + messages: Final = kwargs["messages"] + prompt: Final = get_str_from_messages(messages) + embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) # get the embedding - embedding = embedding_response["data"][0]["embedding"] + embedding: Final = embedding_response["data"][0]["embedding"] value = str(value) assert isinstance(value, str) - data = { + data: Final = { "points": [ { "id": str(uuid.uuid4()), @@ -381,21 +379,20 @@ class QdrantSemanticCache(BaseCache): headers=self.headers, json=data, ) - return async def async_get_cache(self, key, **kwargs): print_verbose(f"async qdrant semantic-cache get_cache, kwargs: {kwargs}") # get the messages - messages = kwargs["messages"] - prompt = get_str_from_messages(messages) + messages: Final = kwargs["messages"] + prompt: Final = get_str_from_messages(messages) - embedding_response = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding_response: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) # get the embedding - embedding = embedding_response["data"][0]["embedding"] + embedding: Final = embedding_response["data"][0]["embedding"] - data = { + data: Final = { "vector": embedding, "params": { "quantization": { @@ -409,13 +406,13 @@ class QdrantSemanticCache(BaseCache): } self._add_cache_key_filter_to_search_data(data=data, key=key) - search_response = await self.async_client.post( + search_response: Final = await self.async_client.post( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points/search", headers=self.headers, json=data, ) - results = search_response.json()["result"] + results: Final = search_response.json()["result"] if results is None: kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 @@ -425,14 +422,14 @@ class QdrantSemanticCache(BaseCache): kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - similarity = results[0]["score"] - payload = results[0]["payload"] + similarity: Final = results[0]["score"] + payload: Final = results[0]["payload"] if not self._payload_matches_cache_key(payload=payload, key=key): print_verbose("Qdrant semantic-cache hit did not match cache key scope") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - cached_prompt = payload["text"] + cached_prompt: Final = payload["text"] # check similarity, if more than self.similarity_threshold, return results print_verbose( @@ -444,7 +441,7 @@ class QdrantSemanticCache(BaseCache): if similarity >= self.similarity_threshold: # cache hit ! - cached_value = payload["response"] + cached_value: Final = payload["response"] print_verbose( f"got a cache hit, similarity: {similarity}, Current prompt: {prompt}, cached_prompt: {cached_prompt}" ) @@ -452,13 +449,12 @@ class QdrantSemanticCache(BaseCache): else: # cache miss ! return None - pass async def _collection_info(self): return self.collection_info async def async_set_cache_pipeline(self, cache_list, **kwargs): - tasks = [] + tasks: Final = [] for val in cache_list: tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) await asyncio.gather(*tasks) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 9e0f022262b..378260b954d 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -16,9 +16,9 @@ import inspect import json import time from collections.abc import Awaitable, Callable, Sequence -from datetime import timedelta from contextvars import ContextVar -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union, cast +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final, TypeVar, Union, cast import litellm from litellm._logging import print_verbose, verbose_logger @@ -69,18 +69,18 @@ def _get_call_stack_info(num_frames: int = 2) -> str: A string with format "current_function <- caller_function [<- grandparent_function]" """ try: - current_frame = inspect.currentframe() + current_frame: Final = inspect.currentframe() if current_frame is None: return "unknown" # Skip this function and the immediate caller (which sets call_type) - f_back = current_frame.f_back + f_back: Final = current_frame.f_back if f_back is None: return "unknown" frame = f_back.f_back if frame is None: return "unknown" - function_names = [] + function_names: Final = [] for _ in range(num_frames): if frame is None: @@ -127,7 +127,7 @@ class RedisCircuitBreaker: self.recovery_timeout = recovery_timeout self.enabled = enabled self._failure_count = 0 - self._opened_at: Optional[float] = None + self._opened_at: float | None = None self._state = self.CLOSED def is_open(self) -> bool: @@ -172,7 +172,7 @@ class RedisCircuitBreaker: _RedisCallResult = TypeVar("_RedisCallResult") -_swallowed_redis_failures: ContextVar[int] = ContextVar("litellm_swallowed_redis_failures", default=0) +_swallowed_redis_failures: Final[ContextVar[int]] = ContextVar("litellm_swallowed_redis_failures", default=0) @functools.lru_cache(maxsize=1) @@ -230,9 +230,9 @@ async def _run_under_circuit_breaker( """ if breaker.is_open(): raise Exception(f"Redis circuit breaker is open — skipping {name}") - swallowed_before = _swallowed_redis_failures.get() + swallowed_before: Final = _swallowed_redis_failures.get() try: - result = await call() + result: Final = await call() except Exception as e: if _is_redis_health_failure(e): breaker.record_failure() @@ -272,17 +272,17 @@ class RedisCache(BaseCache): host=None, port=None, password=None, - redis_flush_size: Optional[int] = 100, - namespace: Optional[str] = None, - startup_nodes: Optional[List] = None, # for redis-cluster - socket_timeout: Optional[float] = 5.0, # default 5 second timeout + redis_flush_size: int | None = 100, + namespace: str | None = None, + startup_nodes: list | None = None, # for redis-cluster + socket_timeout: float | None = 5.0, # default 5 second timeout **kwargs, ): from litellm._service_logger import ServiceLogging from .._redis import get_redis_client, get_redis_connection_pool - redis_kwargs = {} + redis_kwargs: Final = {} if host is not None: redis_kwargs["host"] = host if port is not None: @@ -304,7 +304,7 @@ class RedisCache(BaseCache): redis_kwargs.update(kwargs) self.redis_client = get_redis_client(**redis_kwargs) - self.redis_async_client: Optional[Union[async_redis_client, async_redis_cluster_client]] = None + self.redis_async_client: async_redis_client | async_redis_cluster_client | None = None self.redis_kwargs = redis_kwargs self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs) @@ -346,7 +346,8 @@ class RedisCache(BaseCache): verbose_logger.debug("Ignoring async redis ping. No running event loop.") else: verbose_logger.error( - "Error connecting to Async Redis client - {}".format(str(e)), + "Error connecting to Async Redis client - %s", + e, extra={"error": str(e)}, ) self._handle_async_ping_error(e) @@ -362,9 +363,9 @@ class RedisCache(BaseCache): def _handle_async_ping_error(self, e: Exception): """Handle async ping error with service failure hook.""" try: - loop = asyncio.get_running_loop() - start_time = time.time() - end_time = start_time + loop: Final = asyncio.get_running_loop() + start_time: Final = time.time() + end_time: Final = start_time loop.create_task( self.service_logger_obj.async_service_failure_hook( service=ServiceTypes.REDIS, @@ -379,9 +380,9 @@ class RedisCache(BaseCache): def _handle_sync_ping_error(self, e: Exception): """Handle sync ping error with service failure hook.""" try: - loop = asyncio.get_running_loop() - start_time = time.time() - end_time = start_time + loop: Final = asyncio.get_running_loop() + start_time: Final = time.time() + end_time: Final = start_time loop.create_task( self.service_logger_obj.async_service_failure_hook( service=ServiceTypes.REDIS, @@ -400,22 +401,22 @@ class RedisCache(BaseCache): """ # Create a stable representation of redis_kwargs for hashing # Sort keys to ensure consistent hash regardless of parameter order - sorted_kwargs = sorted(self.redis_kwargs.items()) - kwargs_str = json.dumps(sorted_kwargs, sort_keys=True) - kwargs_hash = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16] + sorted_kwargs: Final = sorted(self.redis_kwargs.items()) + kwargs_str: Final = json.dumps(sorted_kwargs, sort_keys=True) + kwargs_hash: Final = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16] return f"async-redis-client-{kwargs_hash}" def init_async_client( self, - ) -> Union[async_redis_client, async_redis_cluster_client]: + ) -> async_redis_client | async_redis_cluster_client: from litellm import in_memory_llm_clients_cache from .._redis import get_redis_async_client, get_redis_connection_pool - cache_key = self._get_async_client_cache_key() - cached_client = in_memory_llm_clients_cache.get_cache(key=cache_key) + cache_key: Final = self._get_async_client_cache_key() + cached_client: Final = in_memory_llm_clients_cache.get_cache(key=cache_key) if cached_client is not None: - redis_async_client = cast(Union[async_redis_client, async_redis_cluster_client], cached_client) + redis_async_client = cast(async_redis_client | async_redis_cluster_client, cached_client) else: # Create new connection pool and client for current event loop self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) @@ -453,7 +454,7 @@ class RedisCache(BaseCache): return DEFAULT_REDIS_MAJOR_VERSION try: - version_str = str(self.redis_version).strip() + version_str: Final = str(self.redis_version).strip() # Handle cases where there's no dot (e.g., "7" or 7) if "." in version_str: major_version = int(version_str.split(".")[0]) @@ -466,14 +467,14 @@ class RedisCache(BaseCache): return DEFAULT_REDIS_MAJOR_VERSION def set_cache(self, key, value, **kwargs): - ttl = self.get_ttl(**kwargs) + ttl: Final = self.get_ttl(**kwargs) print_verbose(f"Set Redis Cache: key: {key}\nValue {value}\nttl={ttl}, redis_version={self.redis_version}") key = self.check_and_fix_namespace(key=key) try: - start_time = time.time() + start_time: Final = time.time() self.redis_client.set(name=key, value=str(value), ex=ttl) - end_time = time.time() - _duration = end_time - start_time + end_time: Final = time.time() + _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, @@ -483,16 +484,16 @@ class RedisCache(BaseCache): ) except Exception as e: # NON blocking - notify users Redis is throwing an exception - print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {str(e)}") + print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") - def increment_cache(self, key, value: int, ttl: Optional[float] = None, **kwargs) -> int: - _redis_client = self.redis_client + def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: + _redis_client: Final = self.redis_client start_time = time.time() - set_ttl = self.get_ttl(ttl=ttl) + set_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) try: start_time = time.time() - result: int = _redis_client.incr(name=key, amount=value) # type: ignore + result: Final[int] = _redis_client.incr(name=key, amount=value) # type: ignore end_time = time.time() _duration = end_time - start_time self.service_logger_obj.service_success_hook( @@ -506,7 +507,7 @@ class RedisCache(BaseCache): if set_ttl is not None: # check if key already has ttl, if not -> set ttl start_time = time.time() - current_ttl = _redis_client.ttl(key) + current_ttl: Final = _redis_client.ttl(key) end_time = time.time() _duration = end_time - start_time self.service_logger_obj.service_success_hook( @@ -543,10 +544,10 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: - start_time = time.time() + start_time: Final = time.time() try: - keys = [] - _redis_client = self.init_async_client() + keys: Final = [] + _redis_client: Final = self.init_async_client() if not hasattr(_redis_client, "scan_iter"): verbose_logger.debug( "Redis client does not support scan_iter, potentially using Redis Cluster. Returning empty list." @@ -619,14 +620,14 @@ class RedisCache(BaseCache): # different key prefixes never share an executor; in_memory_llm_clients_cache # then adds the running loop, completing the per-(client, namespace, loop) # scoping. - script_cache_key = ( + script_cache_key: Final = ( f"redis-registered-script-{self._get_async_client_cache_key()}-" f"{self.namespace}-{hashlib.sha256(script.encode()).hexdigest()[:16]}" ) async def run_script(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any: async def execute() -> object: - executor: Optional[Callable[..., Awaitable[Any]]] = litellm.in_memory_llm_clients_cache.get_cache( + executor: Callable[..., Awaitable[Any]] | None = litellm.in_memory_llm_clients_cache.get_cache( key=script_cache_key ) if executor is None: @@ -645,21 +646,21 @@ class RedisCache(BaseCache): Kept separate from async_register_script so each loop caches its own executor; see that method for why the binding must be per loop. """ - _redis_client: Any = self.init_async_client() + _redis_client: Final[Any] = self.init_async_client() if hasattr(_redis_client, "register_script"): - registered_script = _redis_client.register_script(script) + registered_script: Final = _redis_client.register_script(script) async def standalone_executor(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any: - namespaced_keys = tuple(self.check_and_fix_namespace(key=key) for key in keys) + namespaced_keys: Final = tuple(self.check_and_fix_namespace(key=key) for key in keys) return await registered_script(keys=namespaced_keys, args=args, client=client) return standalone_executor if hasattr(_redis_client, "script_load"): - script_sha = _redis_client.script_load(script) + script_sha: Final = _redis_client.script_load(script) async def cluster_executor(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any: - namespaced_keys = tuple(self.check_and_fix_namespace(key=key) for key in keys) + namespaced_keys: Final = tuple(self.check_and_fix_namespace(key=key) for key in keys) return await _redis_client.evalsha(script_sha, len(namespaced_keys), *namespaced_keys, *args) return cluster_executor @@ -677,9 +678,9 @@ class RedisCache(BaseCache): ) return None - start_time = time.time() + start_time: Final = time.time() try: - _redis_client: Redis = self.init_async_client() # type: ignore + _redis_client: Final[Redis] = self.init_async_client() # type: ignore except Exception as e: end_time = time.time() _duration = end_time - start_time @@ -703,14 +704,14 @@ class RedisCache(BaseCache): raise e key = self.check_and_fix_namespace(key=key) - ttl = self.get_ttl(**kwargs) - nx = kwargs.get("nx", False) + ttl: Final = self.get_ttl(**kwargs) + nx: Final = kwargs.get("nx", False) print_verbose(f"Set ASYNC Redis Cache: key: {key}\nValue {value}\nttl={ttl}") try: if not hasattr(_redis_client, "set"): raise Exception("Redis client cannot set cache. Attribute not found.") - result = await _redis_client.set( + result: Final = await _redis_client.set( name=key, value=json.dumps(value), nx=nx, @@ -755,10 +756,10 @@ class RedisCache(BaseCache): async def _pipeline_helper( self, - pipe: Union[pipeline, cluster_pipeline], - cache_list: List[Tuple[Any, Any]], - ttl: Optional[float], - ) -> List: + pipe: pipeline | cluster_pipeline, + cache_list: list[tuple[Any, Any]], + ttl: float | None, + ) -> list: """ Helper function for executing a pipeline of set operations on Redis """ @@ -769,7 +770,7 @@ class RedisCache(BaseCache): print_verbose(f"Set ASYNC Redis Cache PIPELINE: key: {cache_key}\nValue {cache_value}\nttl={ttl}") json_cache_value = json.dumps(cache_value) # Set the value with a TTL if it's provided. - _td: Optional[timedelta] = None + _td: timedelta | None = None if ttl is not None: _td = timedelta(seconds=ttl) pipe.set( # type: ignore @@ -778,11 +779,11 @@ class RedisCache(BaseCache): ex=_td, ) # Execute the pipeline and return the results. - results = await pipe.execute() + results: Final = await pipe.execute() return results @_redis_circuit_breaker_guard - async def async_set_cache_pipeline(self, cache_list: List[Tuple[Any, Any]], ttl: Optional[float] = None, **kwargs): + async def async_set_cache_pipeline(self, cache_list: list[tuple[Any, Any]], ttl: float | None = None, **kwargs): """ Use Redis Pipelines for bulk write operations """ @@ -790,14 +791,14 @@ class RedisCache(BaseCache): if len(cache_list) == 0: return - _redis_client = self.init_async_client() - start_time = time.time() + _redis_client: Final = self.init_async_client() + start_time: Final = time.time() print_verbose(f"Set Async Redis Cache: key list: {cache_list}\nttl={ttl}, redis_version={self.redis_version}") - cache_value: Any = None + cache_value: Final[Any] = None try: async with _redis_client.pipeline(transaction=False) as pipe: - results = await self._pipeline_helper(pipe, cache_list, ttl) + results: Final = await self._pipeline_helper(pipe, cache_list, ttl) print_verbose(f"pipeline results: {results}") # Optionally, you could process 'results' to make sure that all set operations were successful. @@ -814,7 +815,7 @@ class RedisCache(BaseCache): parent_otel_span=_get_parent_otel_span_from_kwargs(kwargs), ) ) - return None + return except Exception as e: ## LOGGING ## end_time = time.time() @@ -842,26 +843,26 @@ class RedisCache(BaseCache): self, redis_client: async_redis_client, key: str, - value: List, - ttl: Optional[float], + value: list, + ttl: float | None, ) -> None: """Helper function for async_set_cache_sadd. Separated for testing.""" ttl = self.get_ttl(ttl=ttl) try: await redis_client.sadd(key, *value) # type: ignore if ttl is not None: - _td = timedelta(seconds=ttl) + _td: Final = timedelta(seconds=ttl) await redis_client.expire(key, _td) except Exception: raise @_redis_circuit_breaker_guard - async def async_set_cache_sadd(self, key, value: List, ttl: Optional[float], **kwargs): + async def async_set_cache_sadd(self, key, value: list, ttl: float | None, **kwargs): from redis.asyncio import Redis - start_time = time.time() + start_time: Final = time.time() try: - _redis_client: Redis = self.init_async_client() # type: ignore + _redis_client: Final[Redis] = self.init_async_client() # type: ignore except Exception as e: end_time = time.time() _duration = end_time - start_time @@ -938,23 +939,23 @@ class RedisCache(BaseCache): self, key, value: float, - ttl: Optional[int] = None, - parent_otel_span: Optional[Span] = None, + ttl: int | None = None, + parent_otel_span: Span | None = None, refresh_ttl: bool = False, ) -> float: from redis.asyncio import Redis - _redis_client: Redis = self.init_async_client() # type: ignore - start_time = time.time() - _used_ttl = self.get_ttl(ttl=ttl) + _redis_client: Final[Redis] = self.init_async_client() # type: ignore + start_time: Final = time.time() + _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) try: - result = await _redis_client.incrbyfloat(name=key, amount=value) + result: Final = await _redis_client.incrbyfloat(name=key, amount=value) if _used_ttl is not None: if refresh_ttl: await _redis_client.expire(key, _used_ttl) else: - current_ttl = await _redis_client.ttl(key) + current_ttl: Final = await _redis_client.ttl(key) if current_ttl == -1: await _redis_client.expire(key, _used_ttl) @@ -1011,10 +1012,10 @@ class RedisCache(BaseCache): GET/compare/SET runs in a single Lua call, so it is also atomic across racing callers and pods. Returns the resulting value. """ - _redis_client = self.init_async_client() - _used_ttl = self.get_ttl(ttl=ttl) + _redis_client: Final = self.init_async_client() + _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) - lua = ( + lua: Final = ( "local cur = redis.call('GET', KEYS[1]) " "if cur == false or tonumber(cur) < tonumber(ARGV[1]) then " "redis.call('SET', KEYS[1], ARGV[1]) " @@ -1051,14 +1052,14 @@ class RedisCache(BaseCache): cached_response = ast.literal_eval(cached_response) return cached_response - def get_cache(self, key, parent_otel_span: Optional[Span] = None, **kwargs): + def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): try: key = self.check_and_fix_namespace(key=key) print_verbose(f"Get Redis Cache: key: {key}") - start_time = time.time() - cached_response = self.redis_client.get(key) - end_time = time.time() - _duration = end_time - start_time + start_time: Final = time.time() + cached_response: Final = self.redis_client.get(key) + end_time: Final = time.time() + _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, @@ -1073,7 +1074,7 @@ class RedisCache(BaseCache): # NON blocking - notify users Redis is throwing an exception verbose_logger.error("litellm.caching.caching: get() - Got exception from REDIS: ", e) - def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]: + def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]: """ Wrapper to call `mget` on the redis client @@ -1081,19 +1082,19 @@ class RedisCache(BaseCache): """ return self.redis_client.mget(keys=keys) # type: ignore - async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]: + async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]: """ Wrapper to call `mget` on the redis client We use a wrapper so RedisCluster can override this method """ - async_redis_client = self.init_async_client() + async_redis_client: Final = self.init_async_client() return await async_redis_client.mget(keys=keys) # type: ignore def batch_get_cache( self, - key_list: Union[List[str], List[Optional[str]]], - parent_otel_span: Optional[Span] = None, + key_list: list[str] | list[str | None], + parent_otel_span: Span | None = None, ) -> dict: """ Use Redis for bulk read operations @@ -1106,17 +1107,17 @@ class RedisCache(BaseCache): dict: A dictionary mapping keys to their cached values """ key_value_dict = {} - _key_list = [key for key in key_list if key is not None] + _key_list: Final = [key for key in key_list if key is not None] try: - _keys = [] + _keys: Final = [] for cache_key in _key_list: cache_key = self.check_and_fix_namespace(key=cache_key or "") _keys.append(cache_key) - start_time = time.time() - results: List = self._run_redis_mget_operation(keys=_keys) - end_time = time.time() - _duration = end_time - start_time + start_time: Final = time.time() + results: Final[list] = self._run_redis_mget_operation(keys=_keys) + end_time: Final = time.time() + _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( service=ServiceTypes.REDIS, duration=_duration, @@ -1130,7 +1131,7 @@ class RedisCache(BaseCache): # 'results' is a list of values corresponding to the order of keys in '_key_list'. key_value_dict = dict(zip(_key_list, results)) - decoded_results = {} + decoded_results: Final = {} for k, v in key_value_dict.items(): if isinstance(k, bytes): k = k.decode("utf-8") @@ -1139,22 +1140,22 @@ class RedisCache(BaseCache): return decoded_results except Exception as e: - verbose_logger.error(f"Error occurred in batch get cache - {str(e)}") + verbose_logger.error("Error occurred in batch get cache - %s", e) return key_value_dict @_redis_circuit_breaker_guard - async def async_get_cache(self, key, parent_otel_span: Optional[Span] = None, **kwargs): + async def async_get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): from redis.asyncio import Redis - _redis_client: Redis = self.init_async_client() # type: ignore + _redis_client: Final[Redis] = self.init_async_client() # type: ignore key = self.check_and_fix_namespace(key=key) - start_time = time.time() + start_time: Final = time.time() try: print_verbose(f"Get Async Redis Cache: key: {key}") - cached_response = await _redis_client.get(key) + cached_response: Final = await _redis_client.get(key) print_verbose(f"Got Async Redis Cache: key: {key}, cached_response {cached_response}") - response = self._get_cache_logic(cached_response=cached_response) + response: Final = self._get_cache_logic(cached_response=cached_response) end_time = time.time() _duration = end_time - start_time @@ -1185,14 +1186,14 @@ class RedisCache(BaseCache): event_metadata={"key": key}, ) ) - print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {str(e)}") + print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}") _record_swallowed_redis_failure(self._circuit_breaker, e) @_redis_circuit_breaker_guard async def async_batch_get_cache( self, - key_list: Union[List[str], List[Optional[str]]], - parent_otel_span: Optional[Span] = None, + key_list: list[str] | list[str | None], + parent_otel_span: Span | None = None, ) -> dict: """ Use Redis for bulk read operations @@ -1208,14 +1209,14 @@ class RedisCache(BaseCache): """ # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `mget` key_value_dict = {} - start_time = time.time() - _key_list = [key for key in key_list if key is not None] + start_time: Final = time.time() + _key_list: Final = [key for key in key_list if key is not None] try: - _keys = [] + _keys: Final = [] for cache_key in _key_list: cache_key = self.check_and_fix_namespace(key=cache_key) _keys.append(cache_key) - results = await self._async_run_redis_mget_operation(keys=_keys) + results: Final = await self._async_run_redis_mget_operation(keys=_keys) ## LOGGING ## end_time = time.time() _duration = end_time - start_time @@ -1234,7 +1235,7 @@ class RedisCache(BaseCache): # 'results' is a list of values corresponding to the order of keys in 'key_list'. key_value_dict = dict(zip(_key_list, results)) - decoded_results = {} + decoded_results: Final = {} for k, v in key_value_dict.items(): if isinstance(k, bytes): k = k.decode("utf-8") @@ -1257,7 +1258,7 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) ) - verbose_logger.error(f"Error occurred in async batch get cache - {str(e)}") + verbose_logger.error("Error occurred in async batch get cache - %s", e) _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict @@ -1266,9 +1267,9 @@ class RedisCache(BaseCache): Tests if the sync redis client is correctly setup. """ print_verbose("Pinging Sync Redis Cache") - start_time = time.time() + start_time: Final = time.time() try: - response: bool = self.redis_client.ping() # type: ignore + response: Final[bool] = self.redis_client.ping() # type: ignore print_verbose(f"Redis Cache PING: {response}") ## LOGGING ## end_time = time.time() @@ -1292,16 +1293,16 @@ class RedisCache(BaseCache): error=e, call_type=f"sync_ping <- {_get_call_stack_info()}", ) - verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}") + verbose_logger.error("LiteLLM Redis Cache PING: - Got exception from REDIS : %s", e) raise e async def ping(self) -> bool: # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ping` - _redis_client: Any = self.init_async_client() - start_time = time.time() + _redis_client: Final[Any] = self.init_async_client() + start_time: Final = time.time() print_verbose("Pinging Async Redis Cache") try: - response = await _redis_client.ping() + response: Final = await _redis_client.ping() ## LOGGING ## end_time = time.time() _duration = end_time - start_time @@ -1326,23 +1327,23 @@ class RedisCache(BaseCache): call_type=f"async_ping <- {_get_call_stack_info()}", ) ) - verbose_logger.error(f"LiteLLM Redis Cache PING: - Got exception from REDIS : {str(e)}") + verbose_logger.error("LiteLLM Redis Cache PING: - Got exception from REDIS : %s", e) raise e @_redis_circuit_breaker_guard async def delete_cache_keys(self, keys): # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` - _redis_client: Any = self.init_async_client() + _redis_client: Final[Any] = self.init_async_client() keys = [self.check_and_fix_namespace(key=key) for key in keys] # keys is a list, unpack it so it gets passed as individual elements to delete await _redis_client.delete(*keys) - def client_list(self) -> List: - client_list: List = self.redis_client.client_list() # type: ignore + def client_list(self) -> list: + client_list: Final[list] = self.redis_client.client_list() # type: ignore return client_list def info(self): - info = self.redis_client.info() + info: Final = self.redis_client.info() return info def flush_cache(self): @@ -1372,10 +1373,10 @@ class RedisCache(BaseCache): import redis.asyncio as redis_async # Create a fresh Redis client with current settings - redis_client = redis_async.Redis(**self.redis_kwargs) + redis_client: Final = redis_async.Redis(**self.redis_kwargs) # Test the connection - ping_result = await redis_client.ping() # type: ignore[misc] + ping_result: Final = await redis_client.ping() # type: ignore[misc] # Close the connection await redis_client.aclose() # type: ignore[attr-defined] @@ -1388,17 +1389,17 @@ class RedisCache(BaseCache): else: return {"status": "failed", "message": "Redis ping returned False"} except Exception as e: - verbose_logger.error(f"Redis connection test failed: {str(e)}") + verbose_logger.error("Redis connection test failed: %s", e) return { "status": "failed", - "message": f"Redis connection failed: {str(e)}", + "message": f"Redis connection failed: {e}", "error": str(e), } @_redis_circuit_breaker_guard async def async_delete_cache(self, key: str): # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete` - _redis_client: Any = self.init_async_client() + _redis_client: Final[Any] = self.init_async_client() key = self.check_and_fix_namespace(key=key) # keys is str return await _redis_client.delete(key) @@ -1410,8 +1411,8 @@ class RedisCache(BaseCache): async def _pipeline_increment_helper( self, pipe: pipeline, - increment_list: List[RedisPipelineIncrementOperation], - ) -> Optional[List[float]]: + increment_list: list[RedisPipelineIncrementOperation], + ) -> list[float] | None: """Helper function for pipeline increment operations""" # Iterate through each increment operation and add commands to pipeline for increment_op in increment_list: @@ -1424,15 +1425,15 @@ class RedisCache(BaseCache): _td = timedelta(seconds=increment_op["ttl"]) pipe.expire(cache_key, _td) # Execute the pipeline and return results - results = await pipe.execute() + results: Final = await pipe.execute() # only return float values - verbose_logger.debug(f"Increment ASYNC Redis Cache PIPELINE: results: {results}") + verbose_logger.debug("Increment ASYNC Redis Cache PIPELINE: results: %s", results) return [r for r in results if isinstance(r, float)] @_redis_circuit_breaker_guard async def async_increment_pipeline( - self, increment_list: List[RedisPipelineIncrementOperation], **kwargs - ) -> Optional[List[float]]: + self, increment_list: list[RedisPipelineIncrementOperation], **kwargs + ) -> list[float] | None: """ Use Redis Pipelines for bulk increment operations Args: @@ -1447,14 +1448,14 @@ class RedisCache(BaseCache): from redis.asyncio import Redis - _redis_client: Redis = self.init_async_client() # type: ignore - start_time = time.time() + _redis_client: Final[Redis] = self.init_async_client() # type: ignore + start_time: Final = time.time() print_verbose(f"Increment Async Redis Cache Pipeline: increment list: {increment_list}") try: async with _redis_client.pipeline(transaction=False) as pipe: - results = await self._pipeline_increment_helper(pipe, increment_list) + results: Final = await self._pipeline_increment_helper(pipe, increment_list) ## LOGGING ## end_time = time.time() @@ -1492,7 +1493,7 @@ class RedisCache(BaseCache): raise e @_redis_circuit_breaker_guard - async def async_get_ttl(self, key: str) -> Optional[int]: + async def async_get_ttl(self, key: str) -> int | None: """ Get the remaining TTL of a key in Redis @@ -1506,14 +1507,14 @@ class RedisCache(BaseCache): """ try: # typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl` - _redis_client: Any = self.init_async_client() + _redis_client: Final[Any] = self.init_async_client() key = self.check_and_fix_namespace(key=key) - ttl = await _redis_client.ttl(key) + ttl: Final = await _redis_client.ttl(key) if ttl <= -1: # -1 means the key does not exist, -2 key does not exist return None return ttl except Exception as e: - verbose_logger.debug(f"Redis TTL Error: {e}") + verbose_logger.debug("Redis TTL Error: %s", e) _record_swallowed_redis_failure(self._circuit_breaker, e) return None @@ -1521,8 +1522,8 @@ class RedisCache(BaseCache): async def async_rpush( self, key: str, - values: List[Any], - parent_otel_span: Optional[Span] = None, + values: list[Any], + parent_otel_span: Span | None = None, **kwargs, ) -> int: """ @@ -1536,11 +1537,11 @@ class RedisCache(BaseCache): Returns: int: The length of the list after the push operation """ - _redis_client: Any = self.init_async_client() + _redis_client: Final[Any] = self.init_async_client() key = self.check_and_fix_namespace(key=key) - start_time = time.time() + start_time: Final = time.time() try: - response = await _redis_client.rpush(key, *values) + response: Final = await _redis_client.rpush(key, *values) ## LOGGING ## end_time = time.time() _duration = end_time - start_time @@ -1565,19 +1566,19 @@ class RedisCache(BaseCache): call_type=f"async_rpush <- {_get_call_stack_info()}", ) ) - verbose_logger.error(f"LiteLLM Redis Cache RPUSH: - Got exception from REDIS : {str(e)}") + verbose_logger.error("LiteLLM Redis Cache RPUSH: - Got exception from REDIS : %s", e) raise e async def _pipeline_rpush_helper( self, pipe: pipeline, - rpush_list: List[RedisPipelineRpushOperation], - ) -> List[int]: + rpush_list: list[RedisPipelineRpushOperation], + ) -> list[int]: """Helper function for pipeline rpush operations""" for rpush_op in rpush_list: key = self.check_and_fix_namespace(key=rpush_op["key"]) pipe.rpush(key, *rpush_op["values"]) - results = await pipe.execute() + results: Final = await pipe.execute() # Preserve positional correspondence — raise on per-command errors for r in results: if isinstance(r, Exception): @@ -1587,8 +1588,8 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_rpush_pipeline( self, - rpush_list: List[RedisPipelineRpushOperation], - ) -> List[int]: + rpush_list: list[RedisPipelineRpushOperation], + ) -> list[int]: """ Use Redis Pipelines for bulk RPUSH operations @@ -1603,12 +1604,12 @@ class RedisCache(BaseCache): if len(rpush_list) == 0: return [] - _redis_client: Any = self.init_async_client() - start_time = time.time() + _redis_client: Final[Any] = self.init_async_client() + start_time: Final = time.time() try: async with _redis_client.pipeline(transaction=False) as pipe: - results = await self._pipeline_rpush_helper(pipe, rpush_list) + results: Final = await self._pipeline_rpush_helper(pipe, rpush_list) ## LOGGING ## end_time = time.time() @@ -1639,8 +1640,8 @@ class RedisCache(BaseCache): ) raise e - async def handle_lpop_count_for_older_redis_versions(self, pipe: pipeline, key: str, count: int) -> List[bytes]: - result: List[bytes] = [] + async def handle_lpop_count_for_older_redis_versions(self, pipe: pipeline, key: str, count: int) -> list[bytes]: + result: Final[list[bytes]] = [] for _ in range(count): pipe.lpop(key) results = await pipe.execute() @@ -1656,16 +1657,16 @@ class RedisCache(BaseCache): async def async_lpop( self, key: str, - count: Optional[int] = None, - parent_otel_span: Optional[Span] = None, + count: int | None = None, + parent_otel_span: Span | None = None, **kwargs, - ) -> Union[Any, List[Any]]: - _redis_client: Any = self.init_async_client() + ) -> Any | list[Any]: + _redis_client: Final[Any] = self.init_async_client() key = self.check_and_fix_namespace(key=key) - start_time = time.time() + start_time: Final = time.time() print_verbose(f"LPOP from Redis list: key: {key}, count: {count}") try: - major_version = self._parse_redis_major_version() + major_version: Final = self._parse_redis_major_version() if count is not None and major_version < 7: # For Redis < 7.0, use pipeline to execute multiple LPOP commands @@ -1711,20 +1712,20 @@ class RedisCache(BaseCache): call_type=f"async_lpop <- {_get_call_stack_info()}", ) ) - verbose_logger.error(f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}") + verbose_logger.error("LiteLLM Redis Cache LPOP: - Got exception from REDIS : %s", e) raise e async def _pipeline_lpop_helper( self, pipe: pipeline, - lpop_list: List[RedisPipelineLpopOperation], - ) -> List[Optional[List[str]]]: + lpop_list: list[RedisPipelineLpopOperation], + ) -> list[list[str] | None]: """Helper function for pipeline lpop operations. For Redis >= 7, queues one LPOP(key, count) per operation. For Redis < 7, queues `count` individual LPOP(key) commands per operation. """ - major_version = self._parse_redis_major_version() + major_version: Final = self._parse_redis_major_version() if major_version >= 7: for lpop_op in lpop_list: @@ -1734,14 +1735,14 @@ class RedisCache(BaseCache): else: # For Redis < 7, LPOP doesn't support count param. # Issue `count` individual LPOP commands per key, all in one pipeline. - counts: List[int] = [] + counts: Final[list[int]] = [] for lpop_op in lpop_list: key = self.check_and_fix_namespace(key=lpop_op["key"]) count = lpop_op["count"] or 1 counts.append(count) for _ in range(count): pipe.lpop(key) - flat_results = await pipe.execute() + flat_results: Final = await pipe.execute() # Re-group the flat results back into per-key lists raw_results = [] @@ -1757,7 +1758,7 @@ class RedisCache(BaseCache): raise r # Decode bytes -> str for each result set - decoded_results: List[Optional[List[str]]] = [] + decoded_results: Final[list[list[str] | None]] = [] for r in raw_results: if r is None: decoded_results.append(None) @@ -1776,8 +1777,8 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_lpop_pipeline( self, - lpop_list: List[RedisPipelineLpopOperation], - ) -> List[Optional[List[str]]]: + lpop_list: list[RedisPipelineLpopOperation], + ) -> list[list[str] | None]: """ Use Redis Pipelines for bulk LPOP operations @@ -1792,12 +1793,12 @@ class RedisCache(BaseCache): if len(lpop_list) == 0: return [] - _redis_client: Any = self.init_async_client() - start_time = time.time() + _redis_client: Final[Any] = self.init_async_client() + start_time: Final = time.time() try: async with _redis_client.pipeline(transaction=False) as pipe: - results = await self._pipeline_lpop_helper(pipe, lpop_list) + results: Final = await self._pipeline_lpop_helper(pipe, lpop_list) ## LOGGING ## end_time = time.time() diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 0698ebdcf2a..c275e3c1bf7 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -5,7 +5,7 @@ Key differences: - RedisClient NEEDs to be re-used across requests, adds 3000ms latency if it's re-created """ -from typing import TYPE_CHECKING, Any, List, Optional, Union +from typing import TYPE_CHECKING, Any, Final, Union from litellm.caching.redis_cache import RedisCache @@ -26,8 +26,8 @@ else: class RedisClusterCache(RedisCache): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.redis_async_redis_cluster_client: Optional[RedisCluster] = None - self.redis_sync_redis_cluster_client: Optional[RedisCluster] = None + self.redis_async_redis_cluster_client: RedisCluster | None = None + self.redis_sync_redis_cluster_client: RedisCluster | None = None def init_async_client(self): from redis.asyncio import RedisCluster @@ -37,23 +37,23 @@ class RedisClusterCache(RedisCache): if self.redis_async_redis_cluster_client: return self.redis_async_redis_cluster_client - _redis_client = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs) + _redis_client: Final = get_redis_async_client(connection_pool=self.async_redis_conn_pool, **self.redis_kwargs) if isinstance(_redis_client, RedisCluster): self.redis_async_redis_cluster_client = _redis_client return _redis_client - def _run_redis_mget_operation(self, keys: List[str]) -> List[Any]: + def _run_redis_mget_operation(self, keys: list[str]) -> list[Any]: """ Overrides `_run_redis_mget_operation` in redis_cache.py """ return self.redis_client.mget_nonatomic(keys=keys) # type: ignore - async def _async_run_redis_mget_operation(self, keys: List[str]) -> List[Any]: + async def _async_run_redis_mget_operation(self, keys: list[str]) -> list[Any]: """ Overrides `_async_run_redis_mget_operation` in redis_cache.py """ - async_redis_cluster_client = self.init_async_client() + async_redis_cluster_client: Final = self.init_async_client() return await async_redis_cluster_client.mget_nonatomic(keys=keys) # type: ignore async def test_connection(self) -> dict: @@ -68,21 +68,21 @@ class RedisClusterCache(RedisCache): from redis.cluster import ClusterNode # Create ClusterNode objects from startup_nodes - cluster_kwargs = self.redis_kwargs.copy() - startup_nodes = cluster_kwargs.pop("startup_nodes", []) + cluster_kwargs: Final = self.redis_kwargs.copy() + startup_nodes: Final = cluster_kwargs.pop("startup_nodes", []) - new_startup_nodes: List[ClusterNode] = [] + new_startup_nodes: Final[list[ClusterNode]] = [] for item in startup_nodes: new_startup_nodes.append(ClusterNode(**item)) # Create a fresh Redis Cluster client with current settings - redis_client = redis_async.RedisCluster( + redis_client: Final = redis_async.RedisCluster( startup_nodes=new_startup_nodes, **cluster_kwargs, # type: ignore ) # Test the connection - ping_result = await redis_client.ping() # type: ignore[attr-defined, misc] + ping_result: Final = await redis_client.ping() # type: ignore[attr-defined, misc] # Close the connection await redis_client.aclose() # type: ignore[attr-defined] @@ -100,9 +100,9 @@ class RedisClusterCache(RedisCache): except Exception as e: from litellm._logging import verbose_logger - verbose_logger.error(f"Redis Cluster connection test failed: {str(e)}") + verbose_logger.error("Redis Cluster connection test failed: %s", e) return { "status": "failed", - "message": f"Redis Cluster connection failed: {str(e)}", + "message": f"Redis Cluster connection failed: {e}", "error": str(e), } diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index d4288cc777c..b0c8fa963ee 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -13,7 +13,7 @@ import ast import asyncio import json import os -from typing import Any, Dict, List, Optional, Tuple, cast +from typing import Any, Final, cast import litellm from litellm._logging import print_verbose, verbose_logger @@ -40,13 +40,13 @@ class RedisSemanticCache(BaseCache): def __init__( self, - host: Optional[str] = None, - port: Optional[str] = None, - password: Optional[str] = None, - redis_url: Optional[str] = None, - similarity_threshold: Optional[float] = None, + host: str | None = None, + port: str | None = None, + password: str | None = None, + redis_url: str | None = None, + similarity_threshold: float | None = None, embedding_model: str = "text-embedding-ada-002", - index_name: Optional[str] = None, + index_name: str | None = None, **kwargs, ): """ @@ -95,7 +95,7 @@ class RedisSemanticCache(BaseCache): password = password or os.environ["REDIS_PASSWORD"] except KeyError as e: # Raise a more informative exception if any of the required keys are missing - missing_var = e.args[0] + missing_var: Final = e.args[0] raise ValueError( f"Missing required Redis configuration: {missing_var}. Provide {missing_var} or redis_url." ) from e @@ -130,7 +130,7 @@ class RedisSemanticCache(BaseCache): from redisvl.utils.vectorize import CustomTextVectorizer # type: ignore[import-not-found, import-untyped] try: - cache_vectorizer = CustomTextVectorizer(self._get_embedding) + cache_vectorizer: Final = CustomTextVectorizer(self._get_embedding) return self._init_semantic_cache( semantic_cache_cls=SemanticCache, index_name=self._index_name, @@ -138,11 +138,11 @@ class RedisSemanticCache(BaseCache): cache_vectorizer=cache_vectorizer, ) except Exception as e: - verbose_logger.error(f"Redis semantic-cache index build failed: {e}") + verbose_logger.error("Redis semantic-cache index build failed: %s", e) raise @classmethod - def _cache_key_filterable_field(cls) -> Dict[str, str]: + def _cache_key_filterable_field(cls) -> dict[str, str]: return { "name": cls.CACHE_KEY_FIELD_NAME, "type": "tag", @@ -156,7 +156,7 @@ class RedisSemanticCache(BaseCache): cache_vectorizer: Any, ) -> Any: def _is_schema_mismatch(exc: ValueError) -> bool: - error_message = str(exc).lower() + error_message: Final = str(exc).lower() return any(phrase in error_message for phrase in ("schema does not match", "index schema")) try: @@ -172,7 +172,7 @@ class RedisSemanticCache(BaseCache): if not _is_schema_mismatch(exc): raise - isolated_index_name = f"{index_name}_isolated" + isolated_index_name: Final = f"{index_name}_isolated" print_verbose( "Redis semantic-cache existing index schema is not isolated; " f"using isolated index - {isolated_index_name}" @@ -203,7 +203,7 @@ class RedisSemanticCache(BaseCache): overwrite=True, ) - def _get_cache_filters(self, key: str) -> Dict[str, str]: + def _get_cache_filters(self, key: str) -> dict[str, str]: return {self.CACHE_KEY_FIELD_NAME: str(key)} def _get_cache_key_filter_expression(self, key: str) -> Any: @@ -211,7 +211,7 @@ class RedisSemanticCache(BaseCache): return Tag(self.CACHE_KEY_FIELD_NAME) == str(key) - def _cache_hit_matches_key(self, cache_hit: Dict[str, Any], key: str) -> bool: + def _cache_hit_matches_key(self, cache_hit: dict[str, Any], key: str) -> bool: # Pre-isolation entries with no ``litellm_cache_key`` field cannot be # safely reassigned to a caller's scope and are treated as misses. cached_key = cache_hit.get(self.CACHE_KEY_FIELD_NAME) @@ -219,7 +219,7 @@ class RedisSemanticCache(BaseCache): cached_key = cached_key.decode("utf-8") return cached_key is not None and str(cached_key) == str(key) - def _get_ttl(self, **kwargs) -> Optional[int]: + def _get_ttl(self, **kwargs) -> int | None: """ Get the TTL (time-to-live) value for cache entries. @@ -235,30 +235,30 @@ class RedisSemanticCache(BaseCache): return ttl @classmethod - def _get_prompt_from_kwargs(cls, **kwargs) -> Optional[str]: + def _get_prompt_from_kwargs(cls, **kwargs) -> str | None: """ Extract a semantic-cache prompt from chat or Responses API request kwargs. """ - messages = kwargs.get("messages") + messages: Final = kwargs.get("messages") if messages: return get_str_from_messages(messages) if "input" not in kwargs: return None - prompt_parts: List[str] = [] + prompt_parts: Final[list[str]] = [] cls._collect_responses_input_text(kwargs.get("input"), prompt_parts) - prompt = "\n".join(prompt_parts).strip() + prompt: Final = "\n".join(prompt_parts).strip() return prompt or None @classmethod - def _collect_responses_input_text(cls, value: Any, prompt_parts: List[str]) -> None: + def _collect_responses_input_text(cls, value: Any, prompt_parts: list[str]) -> None: value = cls._coerce_response_input_value(value) if value is None: return if isinstance(value, str): - stripped_value = value.strip() + stripped_value: Final = value.strip() if stripped_value: prompt_parts.append(stripped_value) return @@ -298,15 +298,15 @@ class RedisSemanticCache(BaseCache): @staticmethod def _coerce_response_input_value(value: Any) -> Any: - model_dump = getattr(value, "model_dump", None) + model_dump: Final = getattr(value, "model_dump", None) if callable(model_dump): return model_dump() - dict_method = getattr(value, "dict", None) + dict_method: Final = getattr(value, "dict", None) if callable(dict_method): return dict_method() return value - def _get_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> List[float]: + def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]: """ Routes through the proxy Router when the embedding model is a Router deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies, @@ -318,7 +318,7 @@ class RedisSemanticCache(BaseCache): llm_model_list = None llm_router = None - router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) + router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) if router is not None: embedding_response = cast( EmbeddingResponse, @@ -364,7 +364,7 @@ class RedisSemanticCache(BaseCache): try: cached_response = ast.literal_eval(cached_response) except (ValueError, SyntaxError) as e: - print_verbose(f"Error parsing cached response: {str(e)}") + print_verbose(f"Error parsing cached response: {e}") return None return cached_response @@ -381,29 +381,29 @@ class RedisSemanticCache(BaseCache): """ print_verbose(f"Redis semantic-cache set_cache, kwargs: {kwargs}") - value_str: Optional[str] = None + value_str: str | None = None try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: print_verbose("No prompt provided for semantic caching") return value_str = str(value) - prompt_embedding = self._get_embedding(prompt, metadata=kwargs.get("metadata")) + prompt_embedding: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata")) - store_kwargs: dict[str, Any] = { + store_kwargs: Final[dict[str, Any]] = { "vector": prompt_embedding, "filters": self._get_cache_filters(key), } # Get TTL and store in Redis semantic cache - ttl = self._get_ttl(**kwargs) + ttl: Final = self._get_ttl(**kwargs) if ttl is not None: store_kwargs["ttl"] = int(ttl) self.llmcache.store(prompt, value_str, **store_kwargs) except Exception as e: - print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {str(e)}") + print_verbose(f"Error setting {value_str or value} in the Redis semantic cache: {e}") def get_cache(self, key: str, **kwargs) -> Any: """ @@ -419,7 +419,7 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Redis semantic-cache get_cache, kwargs: {kwargs}") try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: print_verbose("No prompt provided for semantic cache lookup") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 @@ -427,13 +427,13 @@ class RedisSemanticCache(BaseCache): # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - prompt_embedding = self._get_embedding(prompt, metadata=kwargs.get("metadata")) - check_kwargs: dict[str, Any] = { + prompt_embedding: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata")) + check_kwargs: Final[dict[str, Any]] = { "prompt": prompt, "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), } - results = self.llmcache.check(**check_kwargs) + results: Final = self.llmcache.check(**check_kwargs) # Return None if no similar prompts found if not results: @@ -441,20 +441,20 @@ class RedisSemanticCache(BaseCache): return None # Process the best matching result - cache_hit = results[0] + cache_hit: Final = results[0] if not self._cache_hit_matches_key(cache_hit=cache_hit, key=key): print_verbose("Redis semantic-cache hit did not match cache key scope") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - vector_distance = float(cache_hit["vector_distance"]) + vector_distance: Final = float(cache_hit["vector_distance"]) # Convert vector distance back to similarity score # For cosine distance: 0 = most similar, 2 = least similar # While similarity: 1 = most similar, 0 = least similar - similarity = 1 - vector_distance + similarity: Final = 1 - vector_distance - cached_prompt = cache_hit["prompt"] - cached_response = cache_hit["response"] + cached_prompt: Final = cache_hit["prompt"] + cached_response: Final = cache_hit["response"] # update kwargs["metadata"] with similarity, don't rewrite the original metadata kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity @@ -468,10 +468,10 @@ class RedisSemanticCache(BaseCache): return self._get_cache_logic(cached_response=cached_response) except Exception as e: - print_verbose(f"Error retrieving from Redis semantic cache: {str(e)}") + print_verbose(f"Error retrieving from Redis semantic cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _get_async_embedding(self, prompt: str, metadata: Dict[str, Any] | None = None) -> List[float]: + async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]: """ Asynchronously generate an embedding for the given prompt. @@ -488,7 +488,7 @@ class RedisSemanticCache(BaseCache): llm_model_list = None llm_router = None - router = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) + router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) try: if router is not None: embedding_response = await router.aembedding( @@ -505,8 +505,8 @@ class RedisSemanticCache(BaseCache): ) return embedding_response["data"][0]["embedding"] except Exception as e: - print_verbose(f"Error generating async embedding: {str(e)}") - raise ValueError(f"Failed to generate embedding: {str(e)}") from e + print_verbose(f"Error generating async embedding: {e}") + raise ValueError(f"Failed to generate embedding: {e}") from e async def async_set_cache(self, key: str, value: Any, **kwargs) -> None: """ @@ -521,23 +521,23 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Async Redis semantic-cache set_cache, kwargs: {kwargs}") try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: print_verbose("No prompt provided for semantic caching") return - value_str = str(value) + value_str: Final = str(value) # Generate embedding for the value (response) to cache - prompt_embedding = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + prompt_embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) - store_kwargs: dict[str, Any] = { + store_kwargs: Final[dict[str, Any]] = { "vector": prompt_embedding, "filters": self._get_cache_filters(key), } # Get TTL and store in Redis semantic cache - ttl = self._get_ttl(**kwargs) + ttl: Final = self._get_ttl(**kwargs) if ttl is not None: store_kwargs["ttl"] = ttl await self.llmcache.astore( @@ -546,7 +546,7 @@ class RedisSemanticCache(BaseCache): **store_kwargs, ) except Exception as e: - print_verbose(f"Error in async_set_cache: {str(e)}") + print_verbose(f"Error in async_set_cache: {e}") async def async_get_cache(self, key: str, **kwargs) -> Any: """ @@ -562,43 +562,43 @@ class RedisSemanticCache(BaseCache): print_verbose(f"Async Redis semantic-cache get_cache, kwargs: {kwargs}") try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: print_verbose("No prompt provided for semantic cache lookup") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None # Generate embedding for the prompt - prompt_embedding = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + prompt_embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) # Check the cache for semantically similar prompts in this exact # LiteLLM cache-key scope. - check_kwargs: dict[str, Any] = { + check_kwargs: Final[dict[str, Any]] = { "prompt": prompt, "vector": prompt_embedding, "filter_expression": self._get_cache_key_filter_expression(key), } - results = await self.llmcache.acheck(**check_kwargs) + results: Final = await self.llmcache.acheck(**check_kwargs) # handle results / cache hit if not results: kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - cache_hit = results[0] + cache_hit: Final = results[0] if not self._cache_hit_matches_key(cache_hit=cache_hit, key=key): print_verbose("Redis semantic-cache hit did not match cache key scope") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - vector_distance = float(cache_hit["vector_distance"]) + vector_distance: Final = float(cache_hit["vector_distance"]) # Convert vector distance back to similarity # For cosine distance: 0 = most similar, 2 = least similar # While similarity: 1 = most similar, 0 = least similar - similarity = 1 - vector_distance + similarity: Final = 1 - vector_distance - cached_prompt = cache_hit["prompt"] - cached_response = cache_hit["response"] + cached_prompt: Final = cache_hit["prompt"] + cached_response: Final = cache_hit["response"] # update kwargs["metadata"] with similarity, don't rewrite the original metadata kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity @@ -612,20 +612,20 @@ class RedisSemanticCache(BaseCache): return self._get_cache_logic(cached_response=cached_response) except Exception as e: - print_verbose(f"Error in async_get_cache: {str(e)}") + print_verbose(f"Error in async_get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 - async def _index_info(self) -> Dict[str, Any]: + async def _index_info(self) -> dict[str, Any]: """ Get information about the Redis index. Returns: Dict[str, Any]: Information about the Redis index """ - aindex = await self.llmcache._get_async_index() + aindex: Final = await self.llmcache._get_async_index() return await aindex.info() - async def async_set_cache_pipeline(self, cache_list: List[Tuple[str, Any]], **kwargs) -> None: + async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs) -> None: """ Asynchronously store multiple values in the semantic cache. @@ -634,9 +634,9 @@ class RedisSemanticCache(BaseCache): **kwargs: Additional arguments """ try: - tasks = [] + tasks: Final = [] for val in cache_list: tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) await asyncio.gather(*tasks) except Exception as e: - print_verbose(f"Error in async_set_cache_pipeline: {str(e)}") + print_verbose(f"Error in async_set_cache_pipeline: {e}") diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index 1ada940a9c9..e953c9d67b0 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -11,9 +11,9 @@ Has 4 methods: import ast import asyncio import json +from datetime import datetime, timedelta, timezone from functools import partial -from typing import Optional -from datetime import datetime, timezone, timedelta +from typing import Final from litellm._logging import print_verbose, verbose_logger @@ -26,7 +26,7 @@ class S3Cache(BaseCache): s3_bucket_name, s3_region_name=None, s3_api_version=None, - s3_use_ssl: Optional[bool] = True, + s3_use_ssl: bool | None = True, s3_verify=None, s3_endpoint_url=None, s3_aws_access_key_id=None, @@ -63,16 +63,16 @@ class S3Cache(BaseCache): def set_cache(self, key, value, **kwargs): try: print_verbose(f"LiteLLM SET Cache - S3. Key={key}. Value={value}") - ttl = kwargs.get("ttl", None) + ttl: Final = kwargs.get("ttl", None) # Convert value to JSON before storing in S3 - serialized_value = json.dumps(value) + serialized_value: Final = json.dumps(value) key = self._to_s3_key(key) if ttl is not None: cache_control = f"immutable, max-age={ttl}, s-maxage={ttl}" # Calculate expiration time - expiration_time = datetime.now(timezone.utc) + timedelta(seconds=ttl) + expiration_time: Final = datetime.now(timezone.utc) + timedelta(seconds=ttl) # Upload the data to S3 with the calculated expiration time self.s3_client.put_object( Bucket=self.bucket_name, @@ -105,12 +105,12 @@ class S3Cache(BaseCache): Compatible with Python 3.8+. """ try: - verbose_logger.debug(f"Set ASYNC S3 Cache: Key={key}. Value={value}") - loop = asyncio.get_event_loop() - func = partial(self.set_cache, key, value, **kwargs) + verbose_logger.debug("Set ASYNC S3 Cache: Key=%s. Value=%s", key, value) + loop: Final = asyncio.get_event_loop() + func: Final = partial(self.set_cache, key, value, **kwargs) await loop.run_in_executor(None, func) except Exception as e: - verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}") + verbose_logger.error("S3 Caching: async_set_cache() - Got exception from S3: %s", e) def get_cache(self, key, **kwargs): import botocore @@ -124,8 +124,8 @@ class S3Cache(BaseCache): if cached_response is not None: if "Expires" in cached_response: - expires_time = cached_response["Expires"] - current_time = datetime.now(expires_time.tzinfo) + expires_time: Final = cached_response["Expires"] + current_time: Final = datetime.now(expires_time.tzinfo) if current_time > expires_time: return None @@ -139,17 +139,20 @@ class S3Cache(BaseCache): if not isinstance(cached_response, dict): cached_response = dict(cached_response) verbose_logger.debug( - f"Got S3 Cache: key: {key}, cached_response {cached_response}. Type Response {type(cached_response)}" + "Got S3 Cache: key: %s, cached_response %s. Type Response %s", + key, + cached_response, + type(cached_response), ) return cached_response except botocore.exceptions.ClientError as e: # type: ignore if e.response["Error"]["Code"] == "NoSuchKey": - verbose_logger.debug(f"S3 Cache: The specified key '{key}' does not exist in the S3 bucket.") + verbose_logger.debug("S3 Cache: The specified key '%s' does not exist in the S3 bucket.", key) return None except Exception as e: - verbose_logger.error(f"S3 Caching: get_cache() - Got exception from S3: {e}") + verbose_logger.error("S3 Caching: get_cache() - Got exception from S3: %s", e) async def async_get_cache(self, key, **kwargs): """ @@ -157,13 +160,13 @@ class S3Cache(BaseCache): Compatible with Python 3.8+. """ try: - verbose_logger.debug(f"Get ASYNC S3 Cache: key: {key}") - loop = asyncio.get_event_loop() - func = partial(self.get_cache, key, **kwargs) - result = await loop.run_in_executor(None, func) + verbose_logger.debug("Get ASYNC S3 Cache: key: %s", key) + loop: Final = asyncio.get_event_loop() + func: Final = partial(self.get_cache, key, **kwargs) + result: Final = await loop.run_in_executor(None, func) return result except Exception as e: - verbose_logger.error(f"S3 Caching: async_get_cache() - Got exception from S3: {e}") + verbose_logger.error("S3 Caching: async_get_cache() - Got exception from S3: %s", e) return None def flush_cache(self): @@ -173,7 +176,7 @@ class S3Cache(BaseCache): pass async def async_set_cache_pipeline(self, cache_list, **kwargs): - tasks = [] + tasks: Final = [] for val in cache_list: tasks.append(self.async_set_cache(val[0], val[1], **kwargs)) await asyncio.gather(*tasks) diff --git a/litellm/caching/valkey_semantic_cache.py b/litellm/caching/valkey_semantic_cache.py index 76b7f7d5b87..0fe8581df86 100644 --- a/litellm/caching/valkey_semantic_cache.py +++ b/litellm/caching/valkey_semantic_cache.py @@ -19,7 +19,7 @@ import hashlib import os import struct from dataclasses import dataclass -from typing import Any +from typing import Any, Final from redis import Redis from redis.asyncio import Redis as AsyncRedis @@ -106,8 +106,8 @@ class ValkeySemanticCache(RedisSemanticCache): "(or VALKEY_HOST/VALKEY_PORT), or pass redis_url." ) - credentials = f":{password}@" if password else "" - scheme = "rediss" if ssl else "redis" + credentials: Final = f":{password}@" if password else "" + scheme: Final = "rediss" if ssl else "redis" return f"{scheme}://{credentials}{host}:{port}" @classmethod @@ -154,7 +154,7 @@ class ValkeySemanticCache(RedisSemanticCache): return None def _assert_dim_matches(self, info: dict, dim: int) -> None: - existing_dim = self._extract_index_dim(info) + existing_dim: Final = self._extract_index_dim(info) if existing_dim is not None and existing_dim != dim: raise ValueError( f"Valkey semantic-cache index '{self.index_name}' already exists with " @@ -186,7 +186,7 @@ class ValkeySemanticCache(RedisSemanticCache): except Exception as exc: if not self._is_index_exists_error(exc): raise - info = await self.async_client.ft(self.index_name).info() + info: Final = await self.async_client.ft(self.index_name).info() self._assert_dim_matches(info, dim) self._index_dim = dim @@ -202,8 +202,8 @@ class ValkeySemanticCache(RedisSemanticCache): } def _knn_query(self, key: str) -> Query: - scope = self._scope_tag(key) - query_string = ( + scope: Final = self._scope_tag(key) + query_string: Final = ( f"(@{self.CACHE_KEY_FIELD_NAME}:{{{scope}}})" f"=>[KNN 1 @{self.EMBEDDING_FIELD_NAME} $vec AS {self.DISTANCE_FIELD_NAME}]" ) @@ -211,10 +211,10 @@ class ValkeySemanticCache(RedisSemanticCache): @classmethod def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None: - docs = getattr(search_result, "docs", []) + docs: Final = getattr(search_result, "docs", []) if not docs: return None - doc = docs[0] + doc: Final = docs[0] return _ValkeyCacheHit( response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)), distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)), @@ -225,7 +225,7 @@ class ValkeySemanticCache(RedisSemanticCache): kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - similarity = 1 - hit.distance + similarity: Final = 1 - hit.distance kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity if similarity < self.similarity_threshold: @@ -235,86 +235,86 @@ class ValkeySemanticCache(RedisSemanticCache): def set_cache(self, key: str, value: Any, **kwargs: Any) -> None: print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}") try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: print_verbose("No prompt provided for semantic caching") return - embedding = self._get_embedding(prompt) + embedding: Final = self._get_embedding(prompt) self._ensure_index_sync(len(embedding)) - doc_key = self._doc_key(key) + doc_key: Final = self._doc_key(key) self.sync_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding)) - ttl = self._get_ttl(**kwargs) + ttl: Final = self._get_ttl(**kwargs) if ttl is not None: self.sync_client.expire(doc_key, ttl) except Exception as e: - print_verbose(f"Error in Valkey semantic-cache set_cache: {str(e)}") + print_verbose(f"Error in Valkey semantic-cache set_cache: {e}") def get_cache(self, key: str, **kwargs: Any) -> Any: print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}") try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - embedding = self._get_embedding(prompt) + embedding: Final = self._get_embedding(prompt) self._ensure_index_sync(len(embedding)) - search_result = self.sync_client.ft(self.index_name).search( + search_result: Final = self.sync_client.ft(self.index_name).search( self._knn_query(key), query_params={"vec": self._embedding_to_bytes(embedding)}, ) return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: - print_verbose(f"Error in Valkey semantic-cache get_cache: {str(e)}") + print_verbose(f"Error in Valkey semantic-cache get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None: print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}") try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: print_verbose("No prompt provided for semantic caching") return - embedding = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) await self._ensure_index_async(len(embedding)) - doc_key = self._doc_key(key) + doc_key: Final = self._doc_key(key) await self.async_client.hset(doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding)) - ttl = self._get_ttl(**kwargs) + ttl: Final = self._get_ttl(**kwargs) if ttl is not None: await self.async_client.expire(doc_key, ttl) except Exception as e: - print_verbose(f"Error in async Valkey semantic-cache set_cache: {str(e)}") + print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}") async def async_get_cache(self, key: str, **kwargs: Any) -> Any: print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}") try: - prompt = self._get_prompt_from_kwargs(**kwargs) + prompt: Final = self._get_prompt_from_kwargs(**kwargs) if prompt is None: kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 return None - embedding = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) + embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata")) await self._ensure_index_async(len(embedding)) - search_result = await self.async_client.ft(self.index_name).search( + search_result: Final = await self.async_client.ft(self.index_name).search( self._knn_query(key), query_params={"vec": self._embedding_to_bytes(embedding)}, ) return self._resolve_hit(self._first_hit(search_result), key, **kwargs) except Exception as e: - print_verbose(f"Error in async Valkey semantic-cache get_cache: {str(e)}") + print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}") kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0 async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None: try: await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list]) except Exception as e: - print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {str(e)}") + print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}") async def _index_info(self) -> dict: return await self.async_client.ft(self.index_name).info() diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 15f5b28e30e..1e5cccaf23f 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -2,7 +2,8 @@ Handler for transforming /chat/completions api requests to litellm.responses requests """ -from typing import TYPE_CHECKING, Any, Coroutine, Optional, Union +from collections.abc import Coroutine +from typing import TYPE_CHECKING, Any, Final, Union from typing_extensions import TypedDict @@ -46,7 +47,7 @@ class ResponsesToCompletionBridgeHandler: @staticmethod def _coerce_response_object( response_obj: Any, - hidden_params: Optional[dict], + hidden_params: dict | None, ) -> "ResponsesAPIResponse": if isinstance(response_obj, ResponsesAPIResponse): response = response_obj @@ -59,7 +60,7 @@ class ResponsesToCompletionBridgeHandler: raise ValueError("Unexpected responses stream payload") if hidden_params: - existing = getattr(response, "_hidden_params", None) + existing: Final = getattr(response, "_hidden_params", None) if not isinstance(existing, dict) or not existing: setattr(response, "_hidden_params", dict(hidden_params)) else: @@ -71,13 +72,13 @@ class ResponsesToCompletionBridgeHandler: for _ in stream_iter: pass - completed = getattr(stream_iter, "completed_response", None) - response_obj = getattr(completed, "response", None) if completed else None + completed: Final = getattr(stream_iter, "completed_response", None) + response_obj: Final = getattr(completed, "response", None) if completed else None if response_obj is None: raise ValueError("Stream ended without a completed response") - hidden_params = getattr(stream_iter, "_hidden_params", None) - response = self._coerce_response_object(response_obj, hidden_params) + hidden_params: Final = getattr(stream_iter, "_hidden_params", None) + response: Final = self._coerce_response_object(response_obj, hidden_params) if not isinstance(response, ResponsesAPIResponse): raise ValueError("Stream completed response is invalid") return response @@ -86,13 +87,13 @@ class ResponsesToCompletionBridgeHandler: async for _ in stream_iter: pass - completed = getattr(stream_iter, "completed_response", None) - response_obj = getattr(completed, "response", None) if completed else None + completed: Final = getattr(stream_iter, "completed_response", None) + response_obj: Final = getattr(completed, "response", None) if completed else None if response_obj is None: raise ValueError("Stream ended without a completed response") - hidden_params = getattr(stream_iter, "_hidden_params", None) - response = self._coerce_response_object(response_obj, hidden_params) + hidden_params: Final = getattr(stream_iter, "_hidden_params", None) + response: Final = self._coerce_response_object(response_obj, hidden_params) if not isinstance(response, ResponsesAPIResponse): raise ValueError("Stream completed response is invalid") return response @@ -101,35 +102,35 @@ class ResponsesToCompletionBridgeHandler: from litellm import LiteLLMLoggingObj from litellm.types.utils import ModelResponse - model = kwargs.get("model") + model: Final = kwargs.get("model") if model is None or not isinstance(model, str): raise ValueError("model is required") - custom_llm_provider = kwargs.get("custom_llm_provider") + custom_llm_provider: Final = kwargs.get("custom_llm_provider") if custom_llm_provider is None or not isinstance(custom_llm_provider, str): raise ValueError("custom_llm_provider is required") - messages = kwargs.get("messages") + messages: Final = kwargs.get("messages") if messages is None or not isinstance(messages, list): raise ValueError("messages is required") - optional_params = kwargs.get("optional_params") + optional_params: Final = kwargs.get("optional_params") if optional_params is None or not isinstance(optional_params, dict): raise ValueError("optional_params is required") - litellm_params = kwargs.get("litellm_params") + litellm_params: Final = kwargs.get("litellm_params") if litellm_params is None or not isinstance(litellm_params, dict): raise ValueError("litellm_params is required") - headers = kwargs.get("headers") + headers: Final = kwargs.get("headers") if headers is None or not isinstance(headers, dict): raise ValueError("headers is required") - model_response = kwargs.get("model_response") + model_response: Final = kwargs.get("model_response") if model_response is None or not isinstance(model_response, ModelResponse): raise ValueError("model_response is required") - logging_obj = kwargs.get("logging_obj") + logging_obj: Final = kwargs.get("logging_obj") if logging_obj is None or not isinstance(logging_obj, LiteLLMLoggingObj): raise ValueError("logging_obj is required") @@ -157,19 +158,19 @@ class ResponsesToCompletionBridgeHandler: from litellm import responses from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - validated_kwargs = self.validate_input_kwargs(kwargs) - model = validated_kwargs["model"] - messages = validated_kwargs["messages"] + validated_kwargs: Final = self.validate_input_kwargs(kwargs) + model: Final = validated_kwargs["model"] + messages: Final = validated_kwargs["messages"] optional_params = validated_kwargs["optional_params"] - litellm_params = validated_kwargs["litellm_params"] - headers = validated_kwargs["headers"] - model_response = validated_kwargs["model_response"] - logging_obj = validated_kwargs["logging_obj"] - custom_llm_provider = validated_kwargs["custom_llm_provider"] + litellm_params: Final = validated_kwargs["litellm_params"] + headers: Final = validated_kwargs["headers"] + model_response: Final = validated_kwargs["model_response"] + logging_obj: Final = validated_kwargs["logging_obj"] + custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} - request_data = self.transformation_handler.transform_request( + request_data: Final = self.transformation_handler.transform_request( model=model, messages=messages, optional_params=optional_params, @@ -187,13 +188,13 @@ class ResponsesToCompletionBridgeHandler: # than adding an explicit kwarg) avoids the duplicate-keyword # TypeError that would otherwise fire on the real bridge path. request_data["custom_llm_provider"] = custom_llm_provider - result = responses( + result: Final = responses( **request_data, ) from litellm.types.utils import ModelResponse - stream = self._resolve_stream_flag(optional_params, litellm_params) + stream: Final = self._resolve_stream_flag(optional_params, litellm_params) if isinstance(result, ResponsesAPIResponse): return self.transformation_handler.transform_response( model=model, @@ -219,7 +220,7 @@ class ResponsesToCompletionBridgeHandler: json_mode=kwargs.get("json_mode"), ) elif not stream: - responses_api_response = self._collect_response_from_stream(result) + responses_api_response: Final = self._collect_response_from_stream(result) return self.transformation_handler.transform_response( model=model, raw_response=responses_api_response, @@ -236,12 +237,12 @@ class ResponsesToCompletionBridgeHandler: else: if self._is_preformatted_cached_chat_stream(result): return self._apply_post_stream_processing(result, model, custom_llm_provider) - completion_stream = self.transformation_handler.get_model_response_iterator( + completion_stream: Final = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore sync_stream=True, json_mode=kwargs.get("json_mode"), ) - streamwrapper = CustomStreamWrapper( + streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, model=model, custom_llm_provider=custom_llm_provider, @@ -253,20 +254,20 @@ class ResponsesToCompletionBridgeHandler: from litellm import aresponses from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - validated_kwargs = self.validate_input_kwargs(kwargs) - model = validated_kwargs["model"] - messages = validated_kwargs["messages"] + validated_kwargs: Final = self.validate_input_kwargs(kwargs) + model: Final = validated_kwargs["model"] + messages: Final = validated_kwargs["messages"] optional_params = validated_kwargs["optional_params"] - litellm_params = validated_kwargs["litellm_params"] - headers = validated_kwargs["headers"] - model_response = validated_kwargs["model_response"] - logging_obj = validated_kwargs["logging_obj"] - custom_llm_provider = validated_kwargs["custom_llm_provider"] + litellm_params: Final = validated_kwargs["litellm_params"] + headers: Final = validated_kwargs["headers"] + model_response: Final = validated_kwargs["model_response"] + logging_obj: Final = validated_kwargs["logging_obj"] + custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} try: - request_data = self.transformation_handler.transform_request( + request_data: Final = self.transformation_handler.transform_request( model=model, messages=messages, optional_params=optional_params, @@ -284,14 +285,14 @@ class ResponsesToCompletionBridgeHandler: # keyword TypeError when `sanitized_litellm_params` already # carries `custom_llm_provider`. request_data["custom_llm_provider"] = custom_llm_provider - result = await aresponses( + result: Final = await aresponses( **request_data, aresponses=True, ) from litellm.types.utils import ModelResponse - stream = self._resolve_stream_flag(optional_params, litellm_params) + stream: Final = self._resolve_stream_flag(optional_params, litellm_params) if isinstance(result, ResponsesAPIResponse): return self.transformation_handler.transform_response( model=model, @@ -317,7 +318,7 @@ class ResponsesToCompletionBridgeHandler: json_mode=kwargs.get("json_mode"), ) elif not stream: - responses_api_response = await self._collect_response_from_stream_async(result) + responses_api_response: Final = await self._collect_response_from_stream_async(result) return self.transformation_handler.transform_response( model=model, raw_response=responses_api_response, @@ -334,12 +335,12 @@ class ResponsesToCompletionBridgeHandler: else: if self._is_preformatted_cached_chat_stream(result): return self._apply_post_stream_processing(result, model, custom_llm_provider) - completion_stream = self.transformation_handler.get_model_response_iterator( + completion_stream: Final = self.transformation_handler.get_model_response_iterator( streaming_response=result, # type: ignore sync_stream=False, json_mode=kwargs.get("json_mode"), ) - streamwrapper = CustomStreamWrapper( + streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, model=model, custom_llm_provider=custom_llm_provider, @@ -358,7 +359,7 @@ class ResponsesToCompletionBridgeHandler: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.base_model_iterator import MockResponseIterator - streamwrapper = CustomStreamWrapper( + streamwrapper: Final = CustomStreamWrapper( completion_stream=MockResponseIterator(model_response=response, json_mode=json_mode), model=model, custom_llm_provider=custom_llm_provider, @@ -377,7 +378,7 @@ class ResponsesToCompletionBridgeHandler: from litellm.utils import ProviderConfigManager try: - provider_config = ProviderConfigManager.get_provider_chat_config( + provider_config: Final = ProviderConfigManager.get_provider_chat_config( model=model, provider=LlmProviders(custom_llm_provider) ) except (ValueError, KeyError): @@ -388,4 +389,4 @@ class ResponsesToCompletionBridgeHandler: return stream -responses_api_bridge = ResponsesToCompletionBridgeHandler() +responses_api_bridge: Final = ResponsesToCompletionBridgeHandler() diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 89a44fcdeef..4c6112952cc 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -4,22 +4,17 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os -from typing import ( - TYPE_CHECKING, - Any, - AsyncIterator, - Callable, - Dict, - Iterable, - Iterator, - List, - Literal, - Optional, - Tuple, - Union, - cast, -) +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping +from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast +from openai.types.responses.custom_tool_param import CustomToolParam +from openai.types.responses.response_input_param import ( + FunctionCallOutput, + ResponseCustomToolCallOutputParam, + ResponseCustomToolCallParam, +) +from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam +from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam from openai.types.responses.tool_param import FunctionToolParam from pydantic import BaseModel @@ -39,6 +34,8 @@ from litellm.responses.utils import normalize_responses_api_stream_options from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionReasoningItem, + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, ChatCompletionToolParamFunctionChunk, Reasoning, ResponsesAPIOptionalRequestParams, @@ -64,9 +61,9 @@ if TYPE_CHECKING: def _get_reasoning_items( msg: "AllMessageValues", -) -> List[ChatCompletionReasoningItem]: +) -> list[ChatCompletionReasoningItem]: """Extract reasoning_items from a message dict with proper typing.""" - items = msg.get("reasoning_items") # type: ignore[union-attr] + items: Final = msg.get("reasoning_items") # type: ignore[union-attr] if items: return items # type: ignore[return-value] return [] @@ -74,14 +71,14 @@ def _get_reasoning_items( def _build_reasoning_item( item_id: str, - encrypted_content: Optional[str], + encrypted_content: str | None, summary_raw: Any, -) -> Dict[str, Any]: +) -> dict[str, Any]: """Build a ChatCompletionReasoningItem-shaped dict from raw response data. Handles both pydantic objects (attribute access) and plain dicts. """ - summary: List[Dict[str, Any]] = [] + summary: Final[list[dict[str, Any]]] = [] for s in summary_raw or []: if isinstance(s, dict): summary.append({"type": s.get("type", "summary_text"), "text": s.get("text", "")}) @@ -100,11 +97,55 @@ def _build_reasoning_item( } +class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False): + provider_specific_fields: Mapping[str, Any] + + +def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: + """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat + completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw + string payload in ``input`` rather than ``arguments``; both map to + ``function.arguments`` so chat clients (e.g. Cursor agent mode) receive them like + any other tool call. The single conversion rule shared by the non-streaming + accumulator and the streaming ``output_item.added`` branch.""" + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + is_custom: Final = item.get("type") == "custom_tool_call" + arguments: Final = (item.get("input") if is_custom else item.get("arguments")) or "" + name: Final = item.get("name") or ("custom_tool" if is_custom else "") + function_chunk: Final = ChatCompletionToolCallFunctionChunk(name=name, arguments=arguments) + tool_call_dict: Final = _ChatToolCallDict( + id=LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item(item.get("id"), item.get("call_id")), + type="function", + function=function_chunk, + index=index, + ) + raw_provider_fields: Final = item.get("provider_specific_fields") + if isinstance(raw_provider_fields, dict): + provider_specific_fields = raw_provider_fields + elif raw_provider_fields and hasattr(raw_provider_fields, "__dict__"): + provider_specific_fields = vars(raw_provider_fields) + else: + provider_specific_fields = None + if provider_specific_fields: + tool_call_dict["provider_specific_fields"] = provider_specific_fields + function_chunk["provider_specific_fields"] = provider_specific_fields + return tool_call_dict + + +def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFunctionParam | ToolChoiceCustomParam: + if choice_type == "custom": + return ToolChoiceCustomParam(type="custom", name=name) + return ToolChoiceFunctionParam(type="function", name=name) + + def _reasoning_item_to_response_input( - r_item: Union[ChatCompletionReasoningItem, Dict[str, Any]], -) -> Dict[str, Any]: + r_item: ChatCompletionReasoningItem | dict[str, Any], +) -> dict[str, Any]: """Convert a stored ChatCompletionReasoningItem back to a Responses API input item.""" - r_input: Dict[str, Any] = { + r_input: Final[dict[str, Any]] = { "type": "reasoning", "id": r_item.get("id") or f"rs_{id(r_item)}", # summary is always required by the Responses API, even when empty @@ -124,20 +165,23 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): pass def _normalize_tool_choice_for_responses_api(self, tool_choice: Any) -> Any: - """Chat tool_choice uses function.name; Responses API expects top-level name.""" - if not isinstance(tool_choice, dict) or tool_choice.get("type") != "function": + """Chat tool_choice nests the name under function/custom; Responses API expects top-level name.""" + if not isinstance(tool_choice, dict): + return tool_choice + choice_type: Final = tool_choice.get("type") + if choice_type not in ("function", "custom"): return tool_choice if isinstance(tool_choice.get("name"), str) and tool_choice.get("name"): - # Return only Responses shape so stray chat ``function`` key is not sent upstream. - return {"type": "function", "name": tool_choice["name"]} - fn = tool_choice.get("function") - if isinstance(fn, dict): - fn_name = fn.get("name") - if isinstance(fn_name, str) and fn_name: - return {"type": "function", "name": fn_name} + # Return only Responses shape so stray chat ``function``/``custom`` keys are not sent upstream. + return _flat_responses_tool_choice(choice_type, tool_choice["name"]) + nested: Final = tool_choice.get(choice_type) + if isinstance(nested, dict): + nested_name: Final = nested.get("name") + if isinstance(nested_name, str) and nested_name: + return _flat_responses_tool_choice(choice_type, nested_name) return tool_choice - def _handle_raw_dict_response_item(self, item: Dict[str, Any], index: int) -> Tuple[Optional[Any], int]: + def _handle_raw_dict_response_item(self, item: dict[str, Any], index: int) -> tuple[Any | None, int]: """ Handle raw dict response items from Responses API (e.g., GPT-5 Codex format). @@ -150,7 +194,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): """ from litellm.types.utils import Choices, Message - item_type = item.get("type") + item_type: Final = item.get("type") # Ignore reasoning items for now if item_type == "reasoning": @@ -158,7 +202,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Handle message items with output_text content if item_type == "message": - content_list = item.get("content", []) + content_list: Final = item.get("content", []) for content_item in content_list: if isinstance(content_item, dict): content_type = content_item.get("type") @@ -176,45 +220,26 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): choice = Choices(message=msg, finish_reason="stop", index=index) return choice, index + 1 - # Handle function_call items (e.g., from GPT-5 Codex format) - if item_type == "function_call": - # Extract provider_specific_fields if present and pass through as-is - provider_specific_fields = item.get("provider_specific_fields") - if provider_specific_fields and not isinstance(provider_specific_fields, dict): - provider_specific_fields = ( - dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} - ) - - tool_call_dict = { - "id": item.get("call_id") or item.get("id", ""), - "function": { - "name": item.get("name", ""), - "arguments": item.get("arguments", ""), - }, - "type": "function", - } - - # Pass through provider_specific_fields as-is if present - if provider_specific_fields: - tool_call_dict["provider_specific_fields"] = provider_specific_fields - # Also add to function's provider_specific_fields for consistency - tool_call_dict["function"]["provider_specific_fields"] = provider_specific_fields - - msg = Message( - content=None, - tool_calls=[tool_call_dict], - ) - choice = Choices(message=msg, finish_reason="tool_calls", index=index) - return choice, index + 1 + # function_call / custom_tool_call dicts are intercepted and accumulated by + # _convert_response_output_to_choices before this callback is reached # Unknown or unsupported type return None, index def convert_chat_completion_messages_to_responses_api( - self, messages: List["AllMessageValues"] - ) -> Tuple[List[Any], Optional[str]]: - input_items: List[Any] = [] - instructions: Optional[str] = None + self, messages: list["AllMessageValues"] + ) -> tuple[list[Any], str | None]: + input_items: Final[list[Any]] = [] + instructions: str | None = None + custom_tool_call_ids: Final = frozenset( + tool_call["id"] + for msg in messages + if msg.get("role") == "assistant" and isinstance(msg.get("tool_calls"), list) + for tool_call in msg.get("tool_calls") or () + if isinstance(tool_call, dict) + and not tool_call.get("function") + and isinstance(tool_call.get("custom"), dict) + ) for msg in messages: role = msg.get("role") @@ -245,7 +270,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Convert tool message to function call output format # The Responses API expects 'output' to be a list with input_text/input_image types # Using list format for consistency across text and multimodal content - tool_output: List[Dict[str, Any]] + tool_output: list[dict[str, Any]] if content is None: tool_output = [] elif isinstance(content, str): @@ -260,20 +285,30 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: # Fallback: convert unexpected types to input_text tool_output = [{"type": "input_text", "text": str(content)}] - input_items.append( - { - "type": "function_call_output", - "call_id": tool_call_id, - "output": tool_output, - } - ) + if tool_call_id in custom_tool_call_ids: + input_items.append( + ResponseCustomToolCallOutputParam( + type="custom_tool_call_output", + call_id=tool_call_id, + output=content if isinstance(content, str) else tool_output, + ) + ) + else: + input_items.append( + FunctionCallOutput( + type="function_call_output", + call_id=tool_call_id, + output=tool_output, + ) + ) elif role == "assistant" and tool_calls and isinstance(tool_calls, list): for r_item in _get_reasoning_items(msg): input_items.append(_reasoning_item_to_response_input(r_item)) for tool_call in tool_calls: function = tool_call.get("function") + custom = tool_call.get("custom") if function: - input_tool_call: Dict[str, Any] = { + input_tool_call: dict[str, Any] = { "type": "function_call", "call_id": tool_call["id"], } @@ -282,6 +317,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if "arguments" in function: input_tool_call["arguments"] = function["arguments"] input_items.append(input_tool_call) + elif isinstance(custom, dict): + input_items.append( + ResponseCustomToolCallParam( + type="custom_tool_call", + call_id=tool_call["id"], + name=custom.get("name", ""), + input=custom.get("input", ""), + ) + ) else: raise ValueError(f"tool call not supported: {tool_call}") elif content is not None: @@ -311,7 +355,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): responses_api_request["max_output_tokens"] = value elif key == "tools" and value is not None: responses_api_request["tools"] = self._convert_tools_to_responses_format( - cast(List[Dict[str, Any]], value) + cast(list[dict[str, Any]], value) ) elif key == "response_format": text_format = self._transform_response_format_to_text_format(value) @@ -334,15 +378,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif key == "web_search_options": self._add_web_search_tool(responses_api_request, value) - def _build_sanitized_litellm_params(self, litellm_params: dict) -> Dict[str, Any]: + def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, Any]: """Build sanitized litellm_params with merged metadata.""" - responses_optional_param_keys = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) - sanitized: Dict[str, Any] = { + responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) + sanitized: Final[dict[str, Any]] = { key: value for key, value in litellm_params.items() if key not in responses_optional_param_keys } - legacy_metadata = litellm_params.get("metadata") - existing_litellm_metadata = litellm_params.get("litellm_metadata") - merged_litellm_metadata: Dict[str, Any] = {} + legacy_metadata: Final = litellm_params.get("metadata") + existing_litellm_metadata: Final = litellm_params.get("litellm_metadata") + merged_litellm_metadata: Final[dict[str, Any]] = {} if isinstance(legacy_metadata, dict): merged_litellm_metadata.update(legacy_metadata) if isinstance(existing_litellm_metadata, dict): @@ -355,9 +399,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _merge_responses_api_request_into_request_data( self, - request_data: Dict[str, Any], + request_data: dict[str, Any], responses_api_request: "ResponsesAPIOptionalRequestParams", - instructions: Optional[str], + instructions: str | None, ) -> None: """Add non-None values from responses_api_request into request_data.""" for key, value in responses_api_request.items(): @@ -377,12 +421,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def transform_request( self, model: str, - messages: List["AllMessageValues"], + messages: list["AllMessageValues"], optional_params: dict, litellm_params: dict, headers: dict, litellm_logging_obj: "LiteLLMLoggingObj", - client: Optional[Any] = None, + client: Any | None = None, ) -> dict: ( input_items, @@ -406,7 +450,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): optional_params = self._extract_extra_body_params(optional_params) # Build responses API request using the reverse transformation logic - responses_api_request = ResponsesAPIOptionalRequestParams() + responses_api_request: Final = ResponsesAPIOptionalRequestParams() # Set instructions if we found a system message if instructions: @@ -414,30 +458,30 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): self._map_optional_params_to_responses_api_request(optional_params, responses_api_request) - stream = optional_params.get("stream") or litellm_params.get("stream", False) - verbose_logger.debug(f"Chat provider: Stream parameter: {stream}") + stream: Final = optional_params.get("stream") or litellm_params.get("stream", False) + verbose_logger.debug("Chat provider: Stream parameter: %s", stream) # Ensure stream is properly set in the request if stream: responses_api_request["stream"] = True # Handle session management if previous_response_id is provided - previous_response_id = optional_params.get("previous_response_id") + previous_response_id: Final = optional_params.get("previous_response_id") if previous_response_id: # Use the existing session handler for responses API - verbose_logger.debug(f"Chat provider: Warning ignoring previous response ID: {previous_response_id}") + verbose_logger.debug("Chat provider: Warning ignoring previous response ID: %s", previous_response_id) # Convert back to responses API format for the actual request - api_model = model + api_model: Final = model from litellm.types.utils import CallTypes setattr(litellm_logging_obj, "call_type", CallTypes.responses.value) - sanitized_litellm_params = self._build_sanitized_litellm_params(litellm_params) + sanitized_litellm_params: Final = self._build_sanitized_litellm_params(litellm_params) - request_data = { + request_data: Final = { "model": api_model, "input": input_items, "litellm_logging_obj": litellm_logging_obj, @@ -445,7 +489,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): "client": client, } - verbose_logger.debug(f"Chat provider: Final request model={api_model}, input_items={len(input_items)}") + verbose_logger.debug("Chat provider: Final request model=%s, input_items=%s", api_model, len(input_items)) self._merge_responses_api_request_into_request_data(request_data, responses_api_request, instructions) @@ -456,9 +500,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): @staticmethod def _convert_response_output_to_choices( - output_items: List[Any], - handle_raw_dict_callback: Optional[Callable] = None, - ) -> List[Any]: + output_items: list[Any], + handle_raw_dict_callback: Callable | None = None, + ) -> list[Any]: """ Convert Responses API output items to chat completion choices. @@ -484,14 +528,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): from litellm.types.utils import Choices, Message - choices: List[Choices] = [] + choices: Final[list[Choices]] = [] index = 0 - reasoning_content: Optional[str] = None - pending_reasoning_item: Optional[Dict[str, Any]] = None + reasoning_content: str | None = None + pending_reasoning_item: dict[str, Any] | None = None # Collect all tool calls to put them in a single choice # (Chat Completions API expects all tool calls in one message) - accumulated_tool_calls: List[Dict[str, Any]] = [] + accumulated_tool_calls: Final[list[dict[str, Any]]] = [] tool_call_index = 0 for item in output_items: @@ -517,7 +561,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): reasoning_content=reasoning_content, annotations=annotations, reasoning_items=cast( - Optional[List[ChatCompletionReasoningItem]], + list[ChatCompletionReasoningItem] | None, ([pending_reasoning_item] if pending_reasoning_item is not None else None), ), ) @@ -562,11 +606,21 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): accumulated_tool_calls.append(tool_call_dict) tool_call_index += 1 - elif isinstance(item, dict) and handle_raw_dict_callback is not None: - # Handle raw dict responses (e.g., from GPT-5 Codex) - choice, index = handle_raw_dict_callback(item=item, index=index) - if choice is not None: - choices.append(choice) + elif isinstance(item, (dict, BaseModel)): + # Raw dict items (e.g., from GPT-5 Codex) and pydantic items matching no + # openai SDK class above: typed ResponseCustomToolCall and litellm's own + # GenericResponseOutputItem from the completion bridge both land here + raw_item = item if isinstance(item, dict) else item.model_dump() + if raw_item.get("type") in ("function_call", "custom_tool_call"): + # Tool calls accumulate into the single trailing tool_calls choice + # like the typed branches above; a choice per call would hide every + # call after choices[0] from chat clients + accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index)) + tool_call_index += 1 + elif handle_raw_dict_callback is not None: + choice, index = handle_raw_dict_callback(item=raw_item, index=index) + if choice is not None: + choices.append(choice) else: pass # don't fail request if item in list is not supported @@ -577,7 +631,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): tool_calls=accumulated_tool_calls, reasoning_content=reasoning_content, reasoning_items=cast( - Optional[List[ChatCompletionReasoningItem]], + list[ChatCompletionReasoningItem] | None, ([pending_reasoning_item] if pending_reasoning_item is not None else None), ), ) @@ -588,22 +642,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return choices @classmethod - def _extract_output_from_completed_event(cls, parsed_chunk: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]: - response_payload = parsed_chunk.get("response") + def _extract_output_from_completed_event(cls, parsed_chunk: dict[str, Any]) -> list[dict[str, Any]] | None: + response_payload: Final = parsed_chunk.get("response") if not isinstance(response_payload, dict): return None - response_output = response_payload.get("output") + response_output: Final = response_payload.get("output") if not isinstance(response_output, list) or len(response_output) == 0: return None - return cast(List[Dict[str, Any]], response_output) + return cast(list[dict[str, Any]], response_output) @classmethod - def _recover_output_items_from_raw_sse(cls, raw_sse: Optional[str]) -> List[Dict[str, Any]]: + def _recover_output_items_from_raw_sse(cls, raw_sse: str | None) -> list[dict[str, Any]]: if not raw_sse or not isinstance(raw_sse, str): return [] - recovered_output_items: Dict[int, Dict[str, Any]] = {} - recovered_text_only_items: Dict[int, Dict[str, Any]] = {} + recovered_output_items: Final[dict[int, dict[str, Any]]] = {} + recovered_text_only_items: Final[dict[int, dict[str, Any]]] = {} for chunk in raw_sse.splitlines(): parsed_chunk = parse_sse_json_chunk(chunk) @@ -638,7 +692,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # but text-only items at indices without a matching OUTPUT_ITEM_DONE # must still be preserved (e.g. multi-output responses where some # indices only emitted OUTPUT_TEXT_DONE). - merged_items: Dict[int, Dict[str, Any]] = {**recovered_text_only_items} + merged_items: Final[dict[int, dict[str, Any]]] = {**recovered_text_only_items} merged_items.update(recovered_output_items) if merged_items: @@ -647,9 +701,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return [] @classmethod - def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> List[Dict[str, Any]]: - model_call_details = getattr(logging_obj, "model_call_details", {}) or {} - original_response = model_call_details.get("original_response") + def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> list[dict[str, Any]]: + model_call_details: Final = getattr(logging_obj, "model_call_details", {}) or {} + original_response: Final = model_call_details.get("original_response") return cls._recover_output_items_from_raw_sse(original_response) def transform_response( @@ -659,12 +713,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): model_response: "ModelResponse", logging_obj: "LiteLLMLoggingObj", request_data: dict, - messages: List["AllMessageValues"], + messages: list["AllMessageValues"], optional_params: dict, litellm_params: dict, encoding: Any, - api_key: Optional[str] = None, - json_mode: Optional[bool] = None, + api_key: str | None = None, + json_mode: bool | None = None, ) -> "ModelResponse": """Transform Responses API response to chat completion response""" from litellm.responses.utils import ResponseAPILoggingUtils @@ -678,7 +732,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): output_items = raw_response.output if len(output_items) == 0: - recovered_output_items = self._recover_output_items_from_logging(logging_obj) + recovered_output_items: Final = self._recover_output_items_from_logging(logging_obj) if recovered_output_items: output_items = cast(Any, recovered_output_items) raw_response.output = cast(Any, recovered_output_items) @@ -688,7 +742,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) # Convert response output to choices using the static helper - choices = self._convert_response_output_to_choices( + choices: Final = self._convert_response_output_to_choices( output_items=output_items, handle_raw_dict_callback=self._handle_raw_dict_response_item, ) @@ -711,7 +765,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id - raw_response_hidden_params = getattr(raw_response, "_hidden_params", {}) + raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {}) if raw_response_hidden_params: if not hasattr(model_response, "_hidden_params") or model_response._hidden_params is None: model_response._hidden_params = {} @@ -732,11 +786,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): self, streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"], sync_stream: bool, - json_mode: Optional[bool] = False, + json_mode: bool | None = False, ) -> BaseModelResponseIterator: return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode) - def _convert_content_str_to_input_text(self, content: str, role: str) -> Dict[str, Any]: + def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, Any]: if role == "user" or role == "system" or role == "tool": return {"type": "input_text", "text": content} else: @@ -747,23 +801,23 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) -> "ResponseInputImageParam": from openai.types.responses import ResponseInputImageParam - content_image_url = content.get("image_url") - actual_image_url: Optional[str] = None - detail: Optional[Literal["low", "high", "auto"]] = None + content_image_url: Final = content.get("image_url") + actual_image_url: str | None = None + detail: Literal["low", "high", "auto"] | None = None if isinstance(content_image_url, str): actual_image_url = content_image_url elif isinstance(content_image_url, dict): actual_image_url = content_image_url.get("url") detail = cast( - Optional[Literal["low", "high", "auto"]], + Literal["low", "high", "auto"] | None, content_image_url.get("detail"), ) if actual_image_url is None: raise ValueError(f"Invalid image URL: {content_image_url}") - image_param = ResponseInputImageParam(image_url=actual_image_url, detail="auto", type="input_image") + image_param: Final = ResponseInputImageParam(image_url=actual_image_url, detail="auto", type="input_image") if detail: image_param["detail"] = detail @@ -772,47 +826,40 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _convert_content_to_responses_format( self, - content: Optional[ - Union[ - str, - List[Any], - Iterable[ - Union[ - "OpenAIMessageContentListBlock", - "ChatCompletionThinkingBlock", - "ChatCompletionRedactedThinkingBlock", - ] - ], - ] - ], + content: str + | list[Any] + | Iterable[ + Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"] + ] + | None, role: str, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Convert chat completion content to responses API format""" from litellm.types.llms.openai import ChatCompletionImageObject - verbose_logger.debug(f"Chat provider: Converting content to responses format - input type: {type(content)}") + verbose_logger.debug("Chat provider: Converting content to responses format - input type: %s", type(content)) if content is None: return [self._convert_content_str_to_input_text("", role)] elif isinstance(content, str): result = [self._convert_content_str_to_input_text(content, role)] - verbose_logger.debug(f"Chat provider: String content -> {result}") + verbose_logger.debug("Chat provider: String content -> %s", result) return result elif isinstance(content, list): result = [] for i, item in enumerate(content): - verbose_logger.debug(f"Chat provider: Processing content item {i}: {type(item)} = {item}") + verbose_logger.debug("Chat provider: Processing content item %s: %s = %s", i, type(item), item) if isinstance(item, str): converted = self._convert_content_str_to_input_text(item, role) result.append(converted) - verbose_logger.debug(f"Chat provider: -> {converted}") + verbose_logger.debug("Chat provider: -> %s", converted) elif isinstance(item, dict): # Handle multimodal content original_type = item.get("type") if original_type == "text": converted = self._convert_content_str_to_input_text(item.get("text", ""), role) result.append(converted) - verbose_logger.debug(f"Chat provider: text -> {converted}") + verbose_logger.debug("Chat provider: text -> %s", converted) elif original_type == "image_url": # Map to responses API image format converted = cast( @@ -822,14 +869,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ), ) result.append(converted) - verbose_logger.debug(f"Chat provider: image_url -> {converted}") + verbose_logger.debug("Chat provider: image_url -> %s", converted) else: # Try to map other types to responses API format item_type = original_type or "input_text" if item_type == "image": converted = {"type": "input_image", **item} result.append(converted) - verbose_logger.debug(f"Chat provider: image -> {converted}") + verbose_logger.debug("Chat provider: image -> %s", converted) elif item_type == "file": # Map Chat Completion file to Responses API input_file # {"type": "file", "file": {"file_data": "...", "filename": "..."}} @@ -841,7 +888,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if key in file_data: converted[key] = file_data[key] result.append(converted) - verbose_logger.debug(f"Chat provider: file -> {converted}") + verbose_logger.debug("Chat provider: file -> %s", converted) elif item_type in [ "input_text", "input_image", @@ -853,22 +900,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ]: # Already in responses API format result.append(item) - verbose_logger.debug(f"Chat provider: passthrough -> {item}") + verbose_logger.debug("Chat provider: passthrough -> %s", item) else: # Default to input_text for unknown types converted = self._convert_content_str_to_input_text(str(item.get("text", item)), role) result.append(converted) - verbose_logger.debug(f"Chat provider: unknown({original_type}) -> {converted}") - verbose_logger.debug(f"Chat provider: Final converted content: {result}") + verbose_logger.debug("Chat provider: unknown(%s) -> %s", original_type, converted) + verbose_logger.debug("Chat provider: Final converted content: %s", result) return result else: result = [self._convert_content_str_to_input_text(str(content), role)] - verbose_logger.debug(f"Chat provider: Other content type -> {result}") + verbose_logger.debug("Chat provider: Other content type -> %s", result) return result - def _convert_tools_to_responses_format(self, tools: List[Dict[str, Any]]) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]: + def _convert_tools_to_responses_format(self, tools: list[dict[str, Any]]) -> list["ALL_RESPONSES_API_TOOL_PARAMS"]: """Convert chat completion tools to responses API tools format""" - responses_tools: List["ALL_RESPONSES_API_TOOL_PARAMS"] = [] + responses_tools: Final[list[ALL_RESPONSES_API_TOOL_PARAMS]] = [] for tool in tools: # convert function tool from chat completion to responses API format if tool.get("type") == "function": @@ -882,10 +929,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): description=function_tool.get("description"), ) ) + elif tool.get("type") == "custom" and isinstance(tool.get("custom"), dict): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_custom_tool_format_to_responses_shape, + ) + + custom_payload = tool["custom"] + flat_custom = CustomToolParam(type="custom", name=custom_payload.get("name", "")) + if custom_payload.get("description") is not None: + flat_custom["description"] = custom_payload["description"] + if isinstance(custom_payload.get("format"), dict): + flat_custom["format"] = convert_custom_tool_format_to_responses_shape(custom_payload["format"]) + responses_tools.append(flat_custom) else: responses_tools.append(tool) # type: ignore - return cast(List["ALL_RESPONSES_API_TOOL_PARAMS"], responses_tools) + return cast(list["ALL_RESPONSES_API_TOOL_PARAMS"], responses_tools) def _extract_extra_body_params(self, optional_params: dict): """ @@ -894,11 +953,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): unsupported params remain in extra_body. """ # Extract extra_body and separate supported params from unsupported ones - extra_body = optional_params.pop("extra_body", None) or {} + extra_body: Final = optional_params.pop("extra_body", None) or {} if not extra_body: return optional_params - supported_responses_api_params = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) + supported_responses_api_params: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) # Also include params we handle specially supported_responses_api_params.update( { @@ -908,7 +967,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) # Extract supported params from extra_body and merge into optional_params - extra_body_copy = extra_body.copy() + extra_body_copy: Final = extra_body.copy() for key, value in extra_body_copy.items(): if key in supported_responses_api_params: # Prefer extra_body value if it exists (may have more complete info like summary in reasoning_effort) @@ -916,14 +975,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return optional_params - def _map_reasoning_effort(self, reasoning_effort: Union[str, Dict[str, Any]]) -> Optional[Reasoning]: + def _map_reasoning_effort(self, reasoning_effort: str | dict[str, Any]) -> Reasoning | None: # If dict is passed, convert it directly to Reasoning object if isinstance(reasoning_effort, dict): return Reasoning(**reasoning_effort) # type: ignore[typeddict-item] # Check if auto-summary is enabled via flag or environment variable # Priority: litellm.reasoning_auto_summary flag > LITELLM_REASONING_AUTO_SUMMARY env var - auto_summary_enabled = ( + auto_summary_enabled: Final = ( litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" ) @@ -967,16 +1026,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): tools = [] responses_api_request["tools"] = tools - web_search_tool: Dict[str, Any] = {"type": "web_search"} + web_search_tool: Final[dict[str, Any]] = {"type": "web_search"} if isinstance(web_search_options, dict): web_search_tool.update(web_search_options) # Cast to Any to match the expected union type for tools list items tools.append(cast(Any, web_search_tool)) - def _transform_response_format_to_text_format( - self, response_format: Union[Dict[str, Any], Any] - ) -> Optional[Dict[str, Any]]: + def _transform_response_format_to_text_format(self, response_format: dict[str, Any] | Any) -> dict[str, Any] | None: """ Transform Chat Completion response_format parameter to Responses API text.format parameter. @@ -1004,10 +1061,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return None if isinstance(response_format, dict): - format_type = response_format.get("type") + format_type: Final = response_format.get("type") if format_type == "json_schema": - json_schema = response_format.get("json_schema", {}) + json_schema: Final = response_format.get("json_schema", {}) return { "format": { "type": "json_schema", @@ -1025,8 +1082,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): @staticmethod def _convert_annotations_to_chat_format( - annotations: Optional[List[Any]], - ) -> Optional[List[ChatCompletionAnnotation]]: + annotations: list[Any] | None, + ) -> list[ChatCompletionAnnotation] | None: """ Convert annotations from Responses API to Chat Completions format. @@ -1036,7 +1093,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if not annotations: return None - result: List[ChatCompletionAnnotation] = [] + result: Final[list[ChatCompletionAnnotation]] = [] for annotation in annotations: try: # Convert Pydantic models to dicts (handles both v1 and v2) @@ -1048,23 +1105,23 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): annotation_dict = annotation else: # Skip unsupported annotation types - verbose_logger.debug(f"Skipping unsupported annotation type: {type(annotation)}") + verbose_logger.debug("Skipping unsupported annotation type: %s", type(annotation)) continue result.append(annotation_dict) # type: ignore except Exception as e: # Skip malformed annotations - verbose_logger.debug(f"Skipping malformed annotation: {annotation}, error: {e}") + verbose_logger.debug("Skipping malformed annotation: %s, error: %s", annotation, e) continue return result if result else None - def _map_responses_status_to_finish_reason(self, status: Optional[str]) -> str: + def _map_responses_status_to_finish_reason(self, status: str | None) -> str: """Map responses API status to chat completion finish_reason""" if not status: return "stop" - status_mapping = { + status_mapping: Final = { "completed": "stop", "incomplete": "length", "failed": "stop", @@ -1075,9 +1132,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): - def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False): + def __init__(self, streaming_response, sync_stream: bool, json_mode: bool | None = False): super().__init__(streaming_response, sync_stream, json_mode) self._chat_completion_id: str | None = None + self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state def _handle_string_chunk( self, str_line: Union[str, "BaseModel"] @@ -1090,21 +1148,41 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): if not str_line or str_line.startswith("event:"): # ignore. return GenericStreamingChunk(text="", tool_use=None, is_finished=False, finish_reason="", usage=None) - index = str_line.find("data:") + index: Final = str_line.find("data:") if index != -1: str_line = str_line[index + 5 :] return self.chunk_parser(json.loads(str_line)) + @staticmethod + def _sequential_tool_call_index( + tool_call_index_map: dict[int, int] | None, # mutable-ok: per-stream state, remapped in place + output_index: int, + ) -> int: + """Chat-completions tool_call indices must be 0-based and sequential, but + Responses API ``output_index`` counts every output item (reasoning, + message, ...), so the first tool call of a reasoning model arrives at + output_index >= 1 and strict SSE accumulators (e.g. Cursor agent mode) + misplace it. When a per-stream map is provided, remap each distinct + output_index to the next sequential slot; without a map (stateless + callers), fall back to the raw output_index.""" + if tool_call_index_map is None: + return output_index + if output_index not in tool_call_index_map: + tool_call_index_map[output_index] = len(tool_call_index_map) # mutable-ok: per-stream accumulator state + return tool_call_index_map[output_index] + @staticmethod def translate_responses_chunk_to_openai_stream( - parsed_chunk: Union[dict, BaseModel], + parsed_chunk: dict | BaseModel, + tool_call_index_map: dict[int, int] | None = None, # mutable-ok: per-stream state, remapped in place ) -> "ModelResponseStream": """ Translate a Responses API streaming chunk to OpenAI chat completion streaming format. Args: parsed_chunk: Dict containing the Responses API event chunk + tool_call_index_map: Per-stream output_index -> sequential tool_call index map Returns: ModelResponseStream: OpenAI-formatted streaming chunk @@ -1138,11 +1216,11 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ): return ModelResponseStream(**parsed_chunk) - verbose_logger.debug(f"Chat provider: Processing event type: {event_type}") + verbose_logger.debug("Chat provider: Processing event type: %s", event_type) if event_type == "response.created": # Initial response creation event - verbose_logger.debug(f"Chat provider: response.created -> {parsed_chunk}") + verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk) return ModelResponseStream( choices=[ StreamingChoices( @@ -1155,37 +1233,26 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): elif event_type == "response.output_item.added": # New output item added output_item = parsed_chunk.get("item", {}) - if output_item.get("type") == "function_call": - # Extract provider_specific_fields if present - provider_specific_fields = output_item.get("provider_specific_fields") - if provider_specific_fields and not isinstance(provider_specific_fields, dict): - provider_specific_fields = ( - dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} - ) + if output_item.get("type") in ("function_call", "custom_tool_call"): + converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) + provider_specific_fields: Final = converted.get("provider_specific_fields") - function_chunk = ChatCompletionToolCallFunctionChunk( - name=output_item.get("name", None), - arguments=parsed_chunk.get("arguments", ""), + function_chunk: Final = ChatCompletionToolCallFunctionChunk( + name=converted["function"]["name"] or None, + arguments=converted["function"]["arguments"] or parsed_chunk.get("arguments") or "", ) - if provider_specific_fields: function_chunk["provider_specific_fields"] = provider_specific_fields - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, + tool_call_index = OpenAiResponsesToChatCompletionStreamIterator._sequential_tool_call_index( + tool_call_index_map, parsed_chunk.get("output_index", 0) ) - - tool_call_index = parsed_chunk.get("output_index", 0) - tool_call_chunk = ChatCompletionToolCallChunk( - id=LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item( - output_item.get("id"), output_item.get("call_id") - ), + tool_call_chunk: Final = ChatCompletionToolCallChunk( + id=converted["id"], index=tool_call_index, type="function", function=function_chunk, ) - - # Add provider_specific_fields if present if provider_specific_fields: tool_call_chunk.provider_specific_fields = provider_specific_fields # type: ignore @@ -1198,10 +1265,15 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) ] ) - elif event_type == "response.function_call_arguments.delta": - content_part: Optional[str] = parsed_chunk.get("delta", None) + elif event_type in ( + ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA, + ResponsesAPIStreamEvents.CUSTOM_TOOL_CALL_INPUT_DELTA, + ): + content_part: str | None = parsed_chunk.get("delta", None) if content_part: - tool_call_index = parsed_chunk.get("output_index", 0) + tool_call_index = OpenAiResponsesToChatCompletionStreamIterator._sequential_tool_call_index( + tool_call_index_map, parsed_chunk.get("output_index", 0) + ) return ModelResponseStream( choices=[ StreamingChoices( @@ -1225,39 +1297,32 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: # New output item added output_item = parsed_chunk.get("item", {}) - if output_item.get("type") == "function_call": - # Extract provider_specific_fields if present - provider_specific_fields = output_item.get("provider_specific_fields") - if provider_specific_fields and not isinstance(provider_specific_fields, dict): - provider_specific_fields = ( - dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {} + if output_item.get("type") in ("function_call", "custom_tool_call"): + if tool_call_index_map is None: + # Stateless callers (the responses guardrail handler extracting + # tool calls from a buffered output_item.done) get the complete + # tool call; per-stream callers already received it via + # output_item.added and the argument delta events + return ModelResponseStream( + choices=[ # mutable-ok: ModelResponseStream coerces only list choices + StreamingChoices( + index=0, + delta=Delta( + tool_calls=( + _tool_call_dict_from_output_item( + output_item, parsed_chunk.get("output_index", 0) + ), + ) + ), + finish_reason=None, + ) + ] ) - - function_chunk = ChatCompletionToolCallFunctionChunk( - name=output_item.get("name", None), - arguments="", # responses API sends everything again, we don't - ) - - # Add provider_specific_fields to function if present - if provider_specific_fields: - function_chunk["provider_specific_fields"] = provider_specific_fields - - tool_call_index = parsed_chunk.get("output_index", 0) - tool_call_chunk = ChatCompletionToolCallChunk( - id=output_item.get("call_id"), - index=tool_call_index, - type="function", - function=function_chunk, - ) - - # Add provider_specific_fields if present - if provider_specific_fields: - tool_call_chunk.provider_specific_fields = provider_specific_fields # type: ignore - # Do NOT emit finish_reason here — response.completed handles the terminal # finish_reason. Emitting "tool_calls" here would prematurely terminate # the stream before subsequent tool calls arrive (same fix as #17246 for - # the message-type branch). + # the message-type branch). The item's fields were already streamed via + # output_item.added and the argument delta events. return ModelResponseStream( choices=[ StreamingChoices( @@ -1312,17 +1377,19 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): # Check if response contains function_call items in output # to determine correct finish_reason - response_data = parsed_chunk.get("response", {}) - output_items = response_data.get("output", []) if response_data else [] + response_data: Final = parsed_chunk.get("response", {}) + output_items: Final = response_data.get("output", []) if response_data else [] - has_function_calls = any( - item.get("type") == "function_call" for item in output_items if isinstance(item, dict) + has_function_calls: Final = any( + item.get("type") in ("function_call", "custom_tool_call") + for item in output_items + if isinstance(item, dict) ) - finish_reason = "tool_calls" if has_function_calls else "stop" + finish_reason: Final = "tool_calls" if has_function_calls else "stop" # Extract reasoning items with encrypted_content for round-tripping - completed_reasoning_items: Optional[List[Dict[str, Any]]] = None + completed_reasoning_items: list[dict[str, Any]] | None = None for item in output_items: if not isinstance(item, dict) or item.get("type") != "reasoning": continue @@ -1335,8 +1402,8 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): summary_raw=item.get("summary"), ) ) - completed_reasoning_items_typed = cast( - Optional[List[ChatCompletionReasoningItem]], + completed_reasoning_items_typed: Final = cast( + list[ChatCompletionReasoningItem] | None, completed_reasoning_items, ) @@ -1361,7 +1428,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): else: pass # For any unhandled event types, create a minimal valid chunk or skip - verbose_logger.debug(f"Chat provider: Unhandled event type '{event_type}', creating empty chunk") + verbose_logger.debug("Chat provider: Unhandled event type '%s', creating empty chunk", event_type) # Return a minimal valid chunk for unknown events return ModelResponseStream( @@ -1384,9 +1451,11 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): Returns: ModelResponseStream: OpenAI-formatted streaming chunk """ - verbose_logger.debug(f"Chat provider: transform_streaming_response called with chunk: {chunk}") + verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk) return self._with_stream_scoped_id( - OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(chunk) + OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( + chunk, tool_call_index_map=self._tool_call_index_map + ) ) def _with_stream_scoped_id(self, chunk: "ModelResponseStream") -> "ModelResponseStream": diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 9c5f57bc98f..f844b3a3d7f 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -4,7 +4,7 @@ scoring, message stubbing, and retrieval tool injection. """ from collections.abc import Mapping, Sequence -from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast +from typing import Any, Final, cast from litellm.caching.dual_cache import DualCache from litellm.compression.message_stubbing import ( @@ -20,11 +20,11 @@ from litellm.types.utils import CallTypes # CallTypes that produce Anthropic-shaped messages (structured content blocks). # Everything else is treated as OpenAI chat-completions shape. -_ANTHROPIC_CALL_TYPES = frozenset({CallTypes.anthropic_messages.value}) +_ANTHROPIC_CALL_TYPES: Final = frozenset({CallTypes.anthropic_messages.value}) # CallTypes that are valid targets for compression. Compression operates on # message-shaped inputs, so we only accept call types whose payload is a list # of role/content messages. -_SUPPORTED_CALL_TYPES = frozenset( +_SUPPORTED_CALL_TYPES: Final = frozenset( { CallTypes.completion.value, CallTypes.acompletion.value, @@ -33,7 +33,7 @@ _SUPPORTED_CALL_TYPES = frozenset( ) -def _normalize_call_type(call_type: Union[CallTypes, str]) -> str: +def _normalize_call_type(call_type: CallTypes | str) -> str: """Return the string value for a ``CallTypes`` enum or a raw string.""" if isinstance(call_type, CallTypes): return call_type.value @@ -44,7 +44,7 @@ def _is_anthropic_call_type(call_type: str) -> bool: return call_type in _ANTHROPIC_CALL_TYPES -def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]: +def _build_retrieval_tools(keys: list[str], call_type: str) -> list[dict]: """ Build retrieval tool definitions in the target request schema. @@ -54,7 +54,7 @@ def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]: if not keys: return [] - openai_tools = [build_retrieval_tool(keys)] + openai_tools: Final = [build_retrieval_tool(keys)] if not _is_anthropic_call_type(call_type): return openai_tools @@ -63,7 +63,7 @@ def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]: from litellm.llms.anthropic.chat.transformation import AnthropicConfig anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools) - return cast(List[dict], anthropic_tools) + return cast(list[dict], anthropic_tools) def _content_to_text(content: Any) -> str: @@ -77,8 +77,8 @@ def _content_to_text(content: Any) -> str: Implemented iteratively (stack-based) to avoid unbounded recursion. """ - parts: List[str] = [] - stack: List[Any] = [content] + parts: Final[list[str]] = [] + stack: Final[list[Any]] = [content] while stack: item = stack.pop() if isinstance(item, str): @@ -97,9 +97,9 @@ def _content_to_text(content: Any) -> str: def _normalize_messages_for_compression( - messages: List[dict], + messages: list[dict], call_type: str, -) -> Tuple[List[dict], List[dict]]: +) -> tuple[list[dict], list[dict]]: """ Normalize each original message to a text-surrogate content for scoring. @@ -111,9 +111,9 @@ def _normalize_messages_for_compression( f"Unsupported call_type={call_type!r} for compression. Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}." ) - original_messages: List[Dict[str, Any]] = [dict(m) for m in messages] + original_messages: Final[list[dict[str, Any]]] = [dict(m) for m in messages] - normalized_messages: List[dict] = [] + normalized_messages: Final[list[dict]] = [] for msg in original_messages: normalized_messages.append( { @@ -124,7 +124,7 @@ def _normalize_messages_for_compression( return normalized_messages, original_messages -def _extract_last_user_message(messages: List[dict]) -> str: +def _extract_last_user_message(messages: list[dict]) -> str: """Return the text content of the last user message.""" for msg in reversed(messages): if msg.get("role") == "user": @@ -132,10 +132,10 @@ def _extract_last_user_message(messages: List[dict]) -> str: return "" -def _extract_tool_use_ids(content: Any) -> List[str]: +def _extract_tool_use_ids(content: Any) -> list[str]: if not isinstance(content, list): return [] - tool_use_ids: List[str] = [] + tool_use_ids: Final[list[str]] = [] for part in content: if not isinstance(part, dict): continue @@ -147,10 +147,10 @@ def _extract_tool_use_ids(content: Any) -> List[str]: return tool_use_ids -def _extract_tool_result_ids(content: Any) -> Set[str]: +def _extract_tool_result_ids(content: Any) -> set[str]: if not isinstance(content, list): return set() - tool_result_ids: Set[str] = set() + tool_result_ids: Final[set[str]] = set() for part in content: if not isinstance(part, dict): continue @@ -163,15 +163,15 @@ def _extract_tool_result_ids(content: Any) -> Set[str]: def _extract_anthropic_tool_exchange_spans( - messages: List[dict], -) -> Tuple[List[Set[int]], Optional[str]]: + messages: list[dict], +) -> tuple[list[set[int]], str | None]: """ Return atomic 2-message spans for Anthropic tool exchanges. Each assistant message containing `tool_use` must be immediately followed by a user message containing matching `tool_result` blocks for all tool_use ids. """ - spans: List[Set[int]] = [] + spans: Final[list[set[int]]] = [] i = 0 while i < len(messages): current = messages[i] @@ -216,44 +216,44 @@ def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int so compressing it replaces the live instruction with a marker. Compression guardrails share this policy; see the Headroom guardrail. """ - system_indices = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system") - last_user = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:] + system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system") + last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:] last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:] return system_indices + last_user + last_assistant def _combine_scores( - bm25_scores: List[float], - emb_scores: List[float], + bm25_scores: list[float], + emb_scores: list[float], bm25_weight: float = 0.4, -) -> List[float]: +) -> list[float]: """Weighted average of BM25 and embedding scores, with min-max normalization.""" - def _normalize(scores: List[float]) -> List[float]: - min_s = min(scores) if scores else 0.0 - max_s = max(scores) if scores else 0.0 - rng = max_s - min_s + def _normalize(scores: list[float]) -> list[float]: + min_s: Final = min(scores) if scores else 0.0 + max_s: Final = max(scores) if scores else 0.0 + rng: Final = max_s - min_s if rng == 0: return [0.0] * len(scores) return [(s - min_s) / rng for s in scores] - norm_bm25 = _normalize(bm25_scores) - norm_emb = _normalize(emb_scores) - emb_weight = 1.0 - bm25_weight + norm_bm25: Final = _normalize(bm25_scores) + norm_emb: Final = _normalize(emb_scores) + emb_weight: Final = 1.0 - bm25_weight return [bm25_weight * b + emb_weight * e for b, e in zip(norm_bm25, norm_emb)] def _select_kept_indices_for_budget( - normalized_messages: List[dict], - original_messages: List[dict], - combined_scores: List[float], + normalized_messages: list[dict], + original_messages: list[dict], + combined_scores: list[float], compression_target: int, model: str, - initial_kept_indices: Set[int], - tool_exchange_spans: List[Set[int]], -) -> Tuple[Set[int], Dict[int, dict]]: - kept_indices = set(initial_kept_indices) + initial_kept_indices: set[int], + tool_exchange_spans: list[set[int]], +) -> tuple[set[int], dict[int, dict]]: + kept_indices: Final = set(initial_kept_indices) current_tokens = 0 for i in kept_indices: current_tokens += token_counter( @@ -265,14 +265,14 @@ def _select_kept_indices_for_budget( # A unit is either: # 1) a single message index, or # 2) an Anthropic tool-exchange span that must be kept/dropped atomically. - truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict - span_id_by_index: Dict[int, int] = {} + truncated_overrides: Final[dict[int, dict]] = {} # idx -> truncated message dict + span_id_by_index: Final[dict[int, int]] = {} for span_id, span in enumerate(tool_exchange_spans): for idx in span: span_id_by_index[idx] = span_id # Build single-message candidate units (non-span messages). - candidate_units: List[Tuple[float, Tuple[int, ...], bool]] = [] + candidate_units: Final[list[tuple[float, tuple[int, ...], bool]]] = [] for idx in range(len(normalized_messages)): if idx in span_id_by_index or idx in kept_indices: continue @@ -322,8 +322,8 @@ def _select_kept_indices_for_budget( return kept_indices, truncated_overrides -def _get_dropped_tool_span_indices(kept_indices: Set[int], tool_exchange_spans: List[Set[int]]) -> Set[int]: - dropped_tool_span_indices: Set[int] = set() +def _get_dropped_tool_span_indices(kept_indices: set[int], tool_exchange_spans: list[set[int]]) -> set[int]: + dropped_tool_span_indices: Final[set[int]] = set() for span in tool_exchange_spans: if not any(idx in kept_indices for idx in span): dropped_tool_span_indices.update(span) @@ -331,14 +331,14 @@ def _get_dropped_tool_span_indices(kept_indices: Set[int], tool_exchange_spans: def compress( - messages: List[dict], + messages: list[dict], model: str, - call_type: Union[CallTypes, str] = CallTypes.completion, + call_type: CallTypes | str = CallTypes.completion, compression_trigger: int = 200_000, - compression_target: Optional[int] = None, - embedding_model: Optional[str] = None, - embedding_model_params: Optional[Dict[str, Any]] = None, - compression_cache: Optional[DualCache] = None, + compression_target: int | None = None, + embedding_model: str | None = None, + embedding_model_params: dict[str, Any] | None = None, + compression_cache: DualCache | None = None, ) -> CompressedResult: """ Compress a list of messages by replacing low-relevance content with stubs. @@ -372,7 +372,7 @@ def compress( A ``CompressedResult`` dict containing compressed messages, token counts, a cache of original content, and the retrieval tool definition. """ - call_type_str = _normalize_call_type(call_type) + call_type_str: Final = _normalize_call_type(call_type) normalized_messages, original_messages = _normalize_messages_for_compression( messages=messages, call_type=call_type_str, @@ -381,9 +381,9 @@ def compress( if compression_target is None: compression_target = compression_trigger * 7 // 10 - original_tokens = token_counter( + original_tokens: Final = token_counter( model=model, - messages=cast(List[Any], original_messages), + messages=cast(list[Any], original_messages), ) # Pass through if below trigger @@ -399,17 +399,17 @@ def compress( ) # Extract query for relevance scoring - query = _extract_last_user_message(normalized_messages) + query: Final = _extract_last_user_message(normalized_messages) # Score each message - bm25_scores = bm25_score_messages(query, normalized_messages) + bm25_scores: Final = bm25_score_messages(query, normalized_messages) if embedding_model: from litellm.compression.scoring.embedding_scorer import ( embedding_score_messages, ) - emb_scores = embedding_score_messages( + emb_scores: Final = embedding_score_messages( query, normalized_messages, model=embedding_model, @@ -421,10 +421,10 @@ def compress( combined_scores = bm25_scores # Protected messages are never compressed - protected_indices = get_protected_indices(normalized_messages) - kept_indices: Set[int] = set(protected_indices) + protected_indices: Final = get_protected_indices(normalized_messages) + kept_indices: set[int] = set(protected_indices) - tool_exchange_spans: List[Set[int]] = [] + tool_exchange_spans: list[set[int]] = [] if _is_anthropic_call_type(call_type_str): tool_exchange_spans, tool_sequence_error = _extract_anthropic_tool_exchange_spans(original_messages) if tool_sequence_error is not None: @@ -454,10 +454,10 @@ def compress( ) # Build compressed messages and cache - compressed_messages: List[dict] = [] - cache: Dict[str, str] = {} - used_keys: Set[str] = set() - dropped_tool_span_indices = _get_dropped_tool_span_indices( + compressed_messages: Final[list[dict]] = [] + cache: Final[dict[str, str]] = {} + used_keys: Final[set[str]] = set() + dropped_tool_span_indices: Final = _get_dropped_tool_span_indices( kept_indices=kept_indices, tool_exchange_spans=tool_exchange_spans ) @@ -474,11 +474,11 @@ def compress( compressed_messages.append(stub_message(msg, key)) # Build retrieval tool in the target request schema - tools = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str) + tools: Final = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str) - compressed_tokens = token_counter( + compressed_tokens: Final = token_counter( model=model, - messages=cast(List[Any], compressed_messages), + messages=cast(list[Any], compressed_messages), ) return CompressedResult( diff --git a/litellm/compression/content_detection.py b/litellm/compression/content_detection.py index 4a072b63f2c..9cd475d662b 100644 --- a/litellm/compression/content_detection.py +++ b/litellm/compression/content_detection.py @@ -4,8 +4,9 @@ Auto-detect content type per message: code, JSON, or text. import json import re +from typing import Final -_CODE_KEYWORDS = re.compile( +_CODE_KEYWORDS: Final = re.compile( r"\b(?:def |function |class |import |from |require\(|#include|fn |func |const |let |var |public |private |static )\b" ) @@ -16,7 +17,7 @@ def detect_content_type(content: str) -> str: Returns one of: "code", "json", "text" """ - stripped = content.strip() + stripped: Final = content.strip() if not stripped: return "text" @@ -30,10 +31,10 @@ def detect_content_type(content: str) -> str: # Check code indicators # Sample first 5000 chars for performance - sample = stripped[:5000] - keyword_matches = len(_CODE_KEYWORDS.findall(sample)) - lines = sample.split("\n") - indented_lines = sum(1 for line in lines if line.startswith((" ", "\t")) and line.strip()) + sample: Final = stripped[:5000] + keyword_matches: Final = len(_CODE_KEYWORDS.findall(sample)) + lines: Final = sample.split("\n") + indented_lines: Final = sum(1 for line in lines if line.startswith((" ", "\t")) and line.strip()) # If we see multiple code keywords or significant indentation, it's likely code if keyword_matches >= 3 or (indented_lines > len(lines) * 0.3 and len(lines) > 5): diff --git a/litellm/compression/message_stubbing.py b/litellm/compression/message_stubbing.py index 8d4e65752c1..a9a74646614 100644 --- a/litellm/compression/message_stubbing.py +++ b/litellm/compression/message_stubbing.py @@ -3,12 +3,12 @@ Replace messages with compact stubs and extract human-readable keys. """ import re -from typing import Set +from typing import Final from litellm.compression.content_detection import detect_content_type # Patterns for extracting file paths from content -_FILE_PATH_PATTERNS = [ +_FILE_PATH_PATTERNS: Final = [ re.compile(r"^#\s*(\S+\.\w+)", re.MULTILINE), # # filename.py re.compile(r"^//\s*(\S+\.\w+)", re.MULTILINE), # // filename.js re.compile(r"^File:\s*(\S+)", re.MULTILINE), # File: path/to/file @@ -17,7 +17,7 @@ _FILE_PATH_PATTERNS = [ ] -def extract_key(message: dict, fallback_index: int, used_keys: Set[str]) -> str: +def extract_key(message: dict, fallback_index: int, used_keys: set[str]) -> str: """ Extract a human-readable key for the message. @@ -41,7 +41,7 @@ def extract_key(message: dict, fallback_index: int, used_keys: Set[str]) -> str: key = f"message_{fallback_index}" # Handle duplicates - base_key = key + base_key: Final = key counter = 2 while key in used_keys: key = f"{base_key}_{counter}" @@ -62,10 +62,10 @@ def stub_message(message: dict, key: str) -> dict: if isinstance(content, list): content = " ".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in content) - line_count = content.count("\n") + 1 - content_type = detect_content_type(content) + line_count: Final = content.count("\n") + 1 + content_type: Final = detect_content_type(content) - stub_content = ( + stub_content: Final = ( f"[Compressed: {key} — {line_count} lines, {content_type}. " f"Use litellm_content_retrieve tool to get full content.]" ) @@ -90,23 +90,23 @@ def truncate_message(message: dict, max_tokens: int) -> dict: content = " ".join(p.get("text", "") if isinstance(p, dict) else str(p) for p in content) # Rough conversion: 1 token ≈ 3 characters - target_chars = max(100, max_tokens * 3) + target_chars: Final = max(100, max_tokens * 3) if len(content) <= target_chars: return {**message, "content": content} - lines = content.split("\n") + lines: Final = content.split("\n") # Estimate target line count from character budget - avg_line_len = max(1, len(content) // max(1, len(lines))) - target_lines = max(2, target_chars // avg_line_len) + avg_line_len: Final = max(1, len(content) // max(1, len(lines))) + target_lines: Final = max(2, target_chars // avg_line_len) if len(lines) <= target_lines: return {**message, "content": content} - first_count = (target_lines * 7) // 10 - last_count = target_lines - first_count - truncated = ( + first_count: Final = (target_lines * 7) // 10 + last_count: Final = target_lines - first_count + truncated: Final = ( "\n".join(lines[:first_count]) + "\n...[truncated for context window]...\n" + "\n".join(lines[-last_count:]) ) return {**message, "content": truncated} diff --git a/litellm/compression/retrieval_tool.py b/litellm/compression/retrieval_tool.py index 99431a2a15d..ed8c486ad10 100644 --- a/litellm/compression/retrieval_tool.py +++ b/litellm/compression/retrieval_tool.py @@ -2,10 +2,8 @@ Build the litellm_content_retrieve tool definition for the LLM. """ -from typing import List - -def build_retrieval_tool(available_keys: List[str]) -> dict: +def build_retrieval_tool(available_keys: list[str]) -> dict: """ Return an OpenAI-format tool definition that lets the model retrieve the full content of a compressed message. diff --git a/litellm/compression/scoring/bm25.py b/litellm/compression/scoring/bm25.py index 7f919ef16fb..a42ab7919f9 100644 --- a/litellm/compression/scoring/bm25.py +++ b/litellm/compression/scoring/bm25.py @@ -7,21 +7,21 @@ No external dependencies — uses only stdlib. import math import re from collections import Counter -from typing import Dict, List +from typing import Final -def _tokenize(text: str) -> List[str]: +def _tokenize(text: str) -> list[str]: """Split text into lowercase tokens on word boundaries.""" return re.findall(r"[a-z0-9_]+", text.lower()) def _extract_content(message: dict) -> str: """Extract text content from a message dict.""" - content = message.get("content", "") + content: Final = message.get("content", "") if isinstance(content, str): return content if isinstance(content, list): - parts = [] + parts: Final = [] for part in content: if isinstance(part, dict) and part.get("type") == "text": parts.append(part.get("text", "")) @@ -33,10 +33,10 @@ def _extract_content(message: dict) -> str: def bm25_score_messages( query: str, - messages: List[dict], + messages: list[dict], k1: float = 1.5, b: float = 0.75, -) -> List[float]: +) -> list[float]: """ Score each message's relevance to the query using BM25 (Okapi BM25). @@ -49,32 +49,32 @@ def bm25_score_messages( Returns: List of float scores, one per message. Higher = more relevant. """ - query_terms = _tokenize(query) + query_terms: Final = _tokenize(query) if not query_terms: return [0.0] * len(messages) # Tokenize all documents - doc_tokens: List[List[str]] = [] + doc_tokens: Final[list[list[str]]] = [] for msg in messages: doc_tokens.append(_tokenize(_extract_content(msg))) - n = len(doc_tokens) + n: Final = len(doc_tokens) if n == 0: return [] # Average document length - doc_lengths = [len(dt) for dt in doc_tokens] - avgdl = sum(doc_lengths) / n if n > 0 else 1.0 + doc_lengths: Final = [len(dt) for dt in doc_tokens] + avgdl: Final = sum(doc_lengths) / n if n > 0 else 1.0 # Document frequency for each term - df: Dict[str, int] = {} + df: Final[dict[str, int]] = {} for dt in doc_tokens: seen = set(dt) for term in seen: df[term] = df.get(term, 0) + 1 # IDF for query terms - idf: Dict[str, float] = {} + idf: Final[dict[str, float]] = {} for term in set(query_terms): term_df = df.get(term, 0) # Standard BM25 IDF: log((N - df + 0.5) / (df + 0.5) + 1) @@ -86,7 +86,7 @@ def bm25_score_messages( # stemmer dependency. def _expand_tf(query_term: str, tf_counts: Counter) -> int: # type: ignore[type-arg] """Sum TF across all doc tokens that are prefixed by query_term.""" - exact = tf_counts.get(query_term, 0) + exact: Final = tf_counts.get(query_term, 0) if exact: return exact if len(query_term) < 4: @@ -94,7 +94,7 @@ def bm25_score_messages( return sum(count for token, count in tf_counts.items() if token != query_term and token.startswith(query_term)) # Score each document - scores: List[float] = [] + scores: Final[list[float]] = [] for i, dt in enumerate(doc_tokens): if not dt: scores.append(0.0) diff --git a/litellm/compression/scoring/embedding_scorer.py b/litellm/compression/scoring/embedding_scorer.py index f3558ae8f5c..aab1371e097 100644 --- a/litellm/compression/scoring/embedding_scorer.py +++ b/litellm/compression/scoring/embedding_scorer.py @@ -5,18 +5,18 @@ Computes cosine similarity between the query embedding and each message embeddin """ import math -from typing import Any, Dict, List, Optional +from typing import Any, Final from litellm.caching.dual_cache import DualCache def _extract_content(message: dict) -> str: """Extract text content from a message dict.""" - content = message.get("content", "") + content: Final = message.get("content", "") if isinstance(content, str): return content if isinstance(content, list): - parts = [] + parts: Final = [] for part in content: if isinstance(part, dict) and part.get("type") == "text": parts.append(part.get("text", "")) @@ -30,15 +30,15 @@ def _truncate_text(text: str, max_chars: int = 30000) -> str: """Truncate long text, keeping first and last portions.""" if len(text) <= max_chars: return text - half = max_chars // 2 + half: Final = max_chars // 2 return text[:half] + "\n...\n" + text[-half:] -def _cosine_similarity(a: List[float], b: List[float]) -> float: +def _cosine_similarity(a: list[float], b: list[float]) -> float: """Compute cosine similarity between two vectors.""" - dot = sum(x * y for x, y in zip(a, b)) - norm_a = math.sqrt(sum(x * x for x in a)) - norm_b = math.sqrt(sum(x * x for x in b)) + dot: Final = sum(x * y for x, y in zip(a, b)) + norm_a: Final = math.sqrt(sum(x * x for x in a)) + norm_b: Final = math.sqrt(sum(x * x for x in b)) if norm_a == 0 or norm_b == 0: return 0.0 return dot / (norm_a * norm_b) @@ -46,11 +46,11 @@ def _cosine_similarity(a: List[float], b: List[float]) -> float: def embedding_score_messages( query: str, - messages: List[dict], + messages: list[dict], model: str, - cache: Optional[DualCache] = None, - embedding_model_params: Optional[Dict[str, Any]] = None, -) -> List[float]: + cache: DualCache | None = None, + embedding_model_params: dict[str, Any] | None = None, +) -> list[float]: """ Score each message's semantic similarity to the query using embeddings. @@ -67,14 +67,14 @@ def embedding_score_messages( """ import litellm - texts = [_truncate_text(query)] + texts: Final = [_truncate_text(query)] for msg in messages: texts.append(_truncate_text(_extract_content(msg))) # Filter out empty texts — replace with a placeholder to maintain indexing - processed_texts = [t if t.strip() else "empty" for t in texts] + processed_texts: Final = [t if t.strip() else "empty" for t in texts] - kwargs: Dict[str, Any] = { + kwargs: dict[str, Any] = { "model": model, "input": processed_texts, "caching": cache is not None, @@ -82,13 +82,13 @@ def embedding_score_messages( if embedding_model_params: kwargs = {**kwargs, **embedding_model_params} - response = litellm.embedding(**kwargs) + response: Final = litellm.embedding(**kwargs) # Extract embedding vectors - embeddings = [item["embedding"] for item in response.data] + embeddings: Final = [item["embedding"] for item in response.data] - query_embedding = embeddings[0] - scores: List[float] = [] + query_embedding: Final = embeddings[0] + scores: Final[list[float]] = [] for i in range(1, len(embeddings)): scores.append(_cosine_similarity(query_embedding, embeddings[i])) diff --git a/litellm/constants.py b/litellm/constants.py index 78bfc6501e8..264f595027f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1,222 +1,222 @@ import os import sys -from typing import List, Literal, Optional +from typing import Final, Literal from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none -DEFAULT_HEALTH_CHECK_PROMPT = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) -AZURE_DEFAULT_RESPONSES_API_VERSION = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) -ROUTER_MAX_FALLBACKS = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) -DEFAULT_BATCH_SIZE = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) -DEFAULT_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) -DEFAULT_S3_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) -DEFAULT_S3_BATCH_SIZE = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) -DEFAULT_SQS_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) -DEFAULT_NUM_WORKERS_LITELLM_PROXY = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) +DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) +AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) +ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) +DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) +DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) +DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) +DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) +DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) +DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) -DEFAULT_SQS_BATCH_SIZE = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) -SQS_SEND_MESSAGE_ACTION = "SendMessage" -SQS_API_VERSION = "2012-11-05" -DEFAULT_MAX_RETRIES = int(os.getenv("DEFAULT_MAX_RETRIES", 2)) +DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) +SQS_SEND_MESSAGE_ACTION: Final = "SendMessage" +SQS_API_VERSION: Final = "2012-11-05" +DEFAULT_MAX_RETRIES: Final = int(os.getenv("DEFAULT_MAX_RETRIES", 2)) # Max records accepted in one POST /v1/callbacks/logs batch. Bounds the blast # radius: each record fans out to spend logs + every callback integration. -MAX_CALLBACK_LOG_RECORDS = 1000 -DEFAULT_MAX_RECURSE_DEPTH = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100)) +MAX_CALLBACK_LOG_RECORDS: Final = 1000 +DEFAULT_MAX_RECURSE_DEPTH: Final = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100)) DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER", 10)) -DEFAULT_FAILURE_THRESHOLD_PERCENT = float( +DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = float( os.getenv("DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5) ) # default cooldown a deployment if 50% of requests fail in a given minute -DEFAULT_MAX_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) -DEFAULT_ALLOWED_FAILS = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) -DEFAULT_REDIS_SYNC_INTERVAL = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) -DEFAULT_COOLDOWN_TIME_SECONDS = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) -DEFAULT_REPLICATE_POLLING_RETRIES = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) -DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) -DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) +DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) +DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) +DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) +DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) +DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) +DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) +DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) # Maximum wall-clock seconds a streaming response is allowed to run. # Streams exceeding this duration are terminated with a Timeout error. # None (default) = no limit. Set env var to a number of seconds to enable globally. -_max_stream_duration_env = os.getenv("LITELLM_MAX_STREAMING_DURATION_SECONDS", None) -LITELLM_MAX_STREAMING_DURATION_SECONDS = ( +_max_stream_duration_env: Final = os.getenv("LITELLM_MAX_STREAMING_DURATION_SECONDS", None) +LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = ( float(_max_stream_duration_env) if _max_stream_duration_env is not None else None ) # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. # Set to 0 to disable truncation. -MAX_BASE64_LENGTH_FOR_LOGGING = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)) +MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)) # When true, adds detailed per-phase timing breakdown headers to responses. # Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms -LITELLM_DETAILED_TIMING = os.getenv("LITELLM_DETAILED_TIMING", "false").lower() == "true" +LITELLM_DETAILED_TIMING: Final = os.getenv("LITELLM_DETAILED_TIMING", "false").lower() == "true" # Model cost map validation constants -MODEL_COST_MAP_MIN_MODEL_COUNT = int( +MODEL_COST_MAP_MIN_MODEL_COUNT: Final = int( os.getenv("MODEL_COST_MAP_MIN_MODEL_COUNT", 50) ) # Minimum number of models a fetched cost map must contain to be considered valid -MODEL_COST_MAP_MAX_SHRINK_RATIO = float( +MODEL_COST_MAP_MAX_SHRINK_RATIO: Final = float( os.getenv("MODEL_COST_MAP_MAX_SHRINK_RATIO", 0.5) ) # Maximum allowed shrinkage ratio vs local backup (0.5 = reject if fetched map is <50% of backup) -DEFAULT_IMAGE_WIDTH = int(os.getenv("DEFAULT_IMAGE_WIDTH", 300)) -DEFAULT_IMAGE_HEIGHT = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300)) +DEFAULT_IMAGE_WIDTH: Final = int(os.getenv("DEFAULT_IMAGE_WIDTH", 300)) +DEFAULT_IMAGE_HEIGHT: Final = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300)) # Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit) # This prevents memory issues from downloading very large images # Maps to OpenAI's 50 MB payload limit - requests with images exceeding this size will be rejected # Set MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0 to disable image URL handling entirely -MAX_IMAGE_URL_DOWNLOAD_SIZE_MB = float(os.getenv("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 50)) -MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int( +MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: Final = float(os.getenv("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 50)) +MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB: Final = int( os.getenv("MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB", 1024) ) # 1MB = 1024KB # Surrogate-repair fallback in _read_request_body runs two full-body re.sub passes # that block the event loop on multi-MB malformed bodies. Skip the repair above this # size and raise the existing 400 immediately. Set to 0 to disable the cap. -MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB = get_env_int("MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 1) -SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD = int( +MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB: Final = get_env_int("MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 1) +SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD: Final = int( os.getenv("SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD", 1000) ) # Minimum number of requests to consider "reasonable traffic". Used for single-deployment cooldown logic. -DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS = int( +DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS: Final = int( os.getenv("DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS", 5) ) # Minimum number of requests before applying error rate cooldown. Prevents cooldown from triggering on first failure. DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_DISABLE_THINKING_BUDGET", 0)) # MCP Semantic Tool Filter Defaults -DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL = str( +DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL: Final = str( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_EMBEDDING_MODEL", "text-embedding-3-small") ) -DEFAULT_MCP_SEMANTIC_FILTER_TOP_K = int(os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_TOP_K", 10)) -DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD = float( +DEFAULT_MCP_SEMANTIC_FILTER_TOP_K: Final = int(os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_TOP_K", 10)) +DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float( os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3) ) -MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)) +MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150)) # Semantic Guard Defaults -DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL = str( +DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL: Final = str( os.getenv("DEFAULT_SEMANTIC_GUARD_EMBEDDING_MODEL", "text-embedding-3-small") ) DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD", 0.75)) # MCP OAuth2 Client Credentials Defaults -MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) -MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200")) -MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600")) +MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) +MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200")) +MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600")) # Default npm cache directory for STDIO MCP servers. # npm/npx needs a writable cache dir; in containers the default (~/.npm) # may not exist or be read-only. /tmp is always writable. -MCP_NPM_CACHE_DIR = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache") -MCP_OAUTH2_TOKEN_CACHE_MIN_TTL = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10")) +MCP_NPM_CACHE_DIR: Final = os.getenv("MCP_NPM_CACHE_DIR", "/tmp/.npm_mcp_cache") +MCP_OAUTH2_TOKEN_CACHE_MIN_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MIN_TTL", "10")) # Per-user OAuth token Redis cache (for server-side token storage) -MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX = "mcp:per_user_token" -MCP_PER_USER_TOKEN_DEFAULT_TTL = int( +MCP_PER_USER_TOKEN_REDIS_KEY_PREFIX: Final = "mcp:per_user_token" +MCP_PER_USER_TOKEN_DEFAULT_TTL: Final = int( os.getenv("MCP_PER_USER_TOKEN_DEFAULT_TTL", "43200") # 12 hours ) -MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS = int(os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) +MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", "60")) # MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers. -MCP_CLIENT_TIMEOUT = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0")) -MCP_TOOL_LISTING_TIMEOUT = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) -MCP_METADATA_TIMEOUT = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) -MCP_HEALTH_CHECK_TIMEOUT = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) +MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0")) +MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) +MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) +MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. # Note: allowlisted runtimes can still execute code via args (e.g. python -c "..."). # This is an accepted residual risk since these endpoints require PROXY_ADMIN. # Extend via LITELLM_MCP_STDIO_EXTRA_COMMANDS env var (comma-separated). -_MCP_STDIO_EXTRA_COMMANDS = os.getenv("LITELLM_MCP_STDIO_EXTRA_COMMANDS", "") -MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset( +_MCP_STDIO_EXTRA_COMMANDS: Final = os.getenv("LITELLM_MCP_STDIO_EXTRA_COMMANDS", "") +MCP_STDIO_ALLOWED_COMMANDS: Final[frozenset] = frozenset( {"npx", "uvx", "python", "python3", "node", "docker", "deno"} | (set(_MCP_STDIO_EXTRA_COMMANDS.split(",")) - {""}) ) # MCP OAuth2 Token Exchange (OBO) Defaults -MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE = int(os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500")) +MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE", "500")) -LITELLM_UI_ALLOW_HEADERS = [ +LITELLM_UI_ALLOW_HEADERS: Final = [ "x-litellm-semantic-filter", "x-litellm-semantic-filter-tools", "x-litellm-adaptive-router-model", ] # Gemini model-specific minimal thinking budget constants -DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH = int( +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH", 1) ) -DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO = int( +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO", 128) ) -DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int( +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE", 512) ) # Maximum number of callbacks that can be registered # This prevents callbacks from exponentially growing and consuming CPU resources # Override with LITELLM_MAX_CALLBACKS env var for large deployments (e.g., many teams with guardrails) -MAX_CALLBACKS = get_env_int("LITELLM_MAX_CALLBACKS", 100) +MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100) # Metadata key recording which pre_call guardrails the proxy loop already ran, # so the deployment-level hook does not re-run them for the same request -PRE_CALL_EXECUTED_GUARDRAILS_KEY = "_pre_call_executed_guardrails" +PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails" # Generic fallback for unknown models -DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int( +DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128) ) # Provider-specific API base URLs -XAI_API_BASE = "https://api.x.ai/v1" -OPEN_SANDBOX_API_BASE_ENV_VAR = "OPEN_SANDBOX_API_BASE" -OPEN_SANDBOX_API_KEY_ENV_VAR = "OPEN_SANDBOX_API_KEY" -OPEN_SANDBOX_DEFAULT_TEMPLATE = "opensandbox/code-interpreter:v1.1.0" -_OPEN_SANDBOX_FALLBACK_ENTRYPOINT = "/opt/code-interpreter/code-interpreter.sh" -OPEN_SANDBOX_DEFAULT_ENTRYPOINT = (_OPEN_SANDBOX_FALLBACK_ENTRYPOINT,) -OPEN_SANDBOX_DEFAULT_LANGUAGE = "python" -OPEN_SANDBOX_DEFAULT_CPU_LIMIT = "1" -OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT = "2Gi" -OPEN_SANDBOX_EXECD_PORT = 44772 -OPEN_SANDBOX_DEFAULT_TIMEOUT = 300 -OPEN_SANDBOX_READY_TIMEOUT = 30.0 -OPEN_SANDBOX_POLL_INTERVAL = 0.2 +XAI_API_BASE: Final = "https://api.x.ai/v1" +OPEN_SANDBOX_API_BASE_ENV_VAR: Final = "OPEN_SANDBOX_API_BASE" +OPEN_SANDBOX_API_KEY_ENV_VAR: Final = "OPEN_SANDBOX_API_KEY" +OPEN_SANDBOX_DEFAULT_TEMPLATE: Final = "opensandbox/code-interpreter:v1.1.0" +_OPEN_SANDBOX_FALLBACK_ENTRYPOINT: Final = "/opt/code-interpreter/code-interpreter.sh" +OPEN_SANDBOX_DEFAULT_ENTRYPOINT: Final = (_OPEN_SANDBOX_FALLBACK_ENTRYPOINT,) +OPEN_SANDBOX_DEFAULT_LANGUAGE: Final = "python" +OPEN_SANDBOX_DEFAULT_CPU_LIMIT: Final = "1" +OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT: Final = "2Gi" +OPEN_SANDBOX_EXECD_PORT: Final = 44772 +OPEN_SANDBOX_DEFAULT_TIMEOUT: Final = 300 +OPEN_SANDBOX_READY_TIMEOUT: Final = 30.0 +OPEN_SANDBOX_POLL_INTERVAL: Final = 0.2 DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024)) -DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET = int( +DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET: Final = int( os.getenv("DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET", 2048) ) DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET", 4096)) DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET", 8192)) DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET = int(os.getenv("DEFAULT_REASONING_EFFORT_MAX_THINKING_BUDGET", 16384)) -MAX_TOKEN_TRIMMING_ATTEMPTS = int( +MAX_TOKEN_TRIMMING_ATTEMPTS: Final = int( os.getenv("MAX_TOKEN_TRIMMING_ATTEMPTS", 10) ) # Maximum number of attempts to trim the message -RUNWAYML_DEFAULT_API_VERSION = str(os.getenv("RUNWAYML_DEFAULT_API_VERSION", "2024-11-06")) +RUNWAYML_DEFAULT_API_VERSION: Final = str(os.getenv("RUNWAYML_DEFAULT_API_VERSION", "2024-11-06")) RUNWAYML_POLLING_TIMEOUT = int(os.getenv("RUNWAYML_POLLING_TIMEOUT", 600)) # 10 minutes default for image generation ########## Networking constants ############################################################## -_DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour +_DEFAULT_TTL_FOR_HTTPX_CLIENTS: Final = 3600 # 1 hour, re-use the same httpx client for 1 hour # Aiohttp connection pooling - prevents memory leaks from unbounded connection growth # Set to 0 for unlimited (not recommended for production) -AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000)) -AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500)) -AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) -AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) +AIOHTTP_CONNECTOR_LIMIT: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000)) +AIOHTTP_CONNECTOR_LIMIT_PER_HOST: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500)) +AIOHTTP_KEEPALIVE_TIMEOUT: Final = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) +AIOHTTP_TTL_DNS_CACHE: Final = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) # TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs # whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT # Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing # during a long provider call and the NAT reaps the flow before the response # arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset # the NAT idle timer. -AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true" -AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60)) -AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30)) -AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5)) +AIOHTTP_SO_KEEPALIVE: Final = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true" +AIOHTTP_TCP_KEEPIDLE: Final = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60)) +AIOHTTP_TCP_KEEPINTVL: Final = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30)) +AIOHTTP_TCP_KEEPCNT: Final = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5)) # enable_cleanup_closed is only needed for Python versions with the SSL leak bug # Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960) # Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78 -AIOHTTP_NEEDS_CLEANUP_CLOSED = (3, 13, 0) <= sys.version_info < ( +AIOHTTP_NEEDS_CLEANUP_CLOSED: Final = (3, 13, 0) <= sys.version_info < ( 3, 13, 1, @@ -225,13 +225,13 @@ AIOHTTP_NEEDS_CLEANUP_CLOSED = (3, 13, 0) <= sys.version_info < ( # WebSocket constants # Default to None (unlimited) to match OpenAI's official agents SDK behavior # https://github.com/openai/openai-agents-python/blob/cf1b933660e44fd37b4350c41febab8221801409/src/agents/realtime/openai_realtime.py#L235 -_max_size_env = os.getenv("REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES") -REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES = int(_max_size_env) if _max_size_env is not None else None +_max_size_env: Final = os.getenv("REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES") +REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = int(_max_size_env) if _max_size_env is not None else None # SSL/TLS cipher configuration for faster handshakes # Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones # This balances performance with broad compatibility -DEFAULT_SSL_CIPHERS = os.getenv( +DEFAULT_SSL_CIPHERS: Final = os.getenv( "LITELLM_SSL_CIPHERS", # Priority 1: TLS 1.3 ciphers (fastest, ~50ms handshake) "TLS_AES_256_GCM_SHA384:" # Fastest observed in testing @@ -253,131 +253,141 @@ DEFAULT_SSL_CIPHERS = os.getenv( ) ########### v2 Architecture constants for managing writing updates to the database ########### -REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" -REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer" -REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer" -REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer" -REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer" -REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_buffer" -REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer" -MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) +REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer" +REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer" +REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer" +REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer" +REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer" +REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" +REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" +MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth -LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) -TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) -GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS = int( +LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) +TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) +GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int( os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) ) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. -MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) -MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = int(os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000)) +MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) +MAX_IN_MEMORY_QUEUE_FLUSH_COUNT: Final = int(os.getenv("MAX_IN_MEMORY_QUEUE_FLUSH_COUNT", 1000)) ############################################################################################### # Providers will not cache a prefix below a minimum size. That minimum is per-model, not global: # Anthropic's ranges from 512 to 4096 depending on the model, and can differ per platform for the # same model. The real minimum is resolved from `prompt_cache_min_tokens` in the model cost map; # this value is only the fallback for models the cost map has no entry for, and doubles as a global # escape hatch when `MINIMUM_PROMPT_CACHE_TOKEN_COUNT` is explicitly set. -MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE: int | None = get_env_int_or_none("MINIMUM_PROMPT_CACHE_TOKEN_COUNT") -DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT = 1024 -MINIMUM_PROMPT_CACHE_TOKEN_COUNT = ( +MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE: Final[int | None] = get_env_int_or_none("MINIMUM_PROMPT_CACHE_TOKEN_COUNT") +DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = 1024 +MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = ( MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT ) -DEFAULT_TRIM_RATIO = float( +DEFAULT_TRIM_RATIO: Final = float( os.getenv("DEFAULT_TRIM_RATIO", 0.75) ) # default ratio of tokens to trim from the end of a prompt -HOURS_IN_A_DAY = int(os.getenv("HOURS_IN_A_DAY", 24)) -DAYS_IN_A_WEEK = int(os.getenv("DAYS_IN_A_WEEK", 7)) -DAYS_IN_A_MONTH = int(os.getenv("DAYS_IN_A_MONTH", 28)) -DAYS_IN_A_YEAR = int(os.getenv("DAYS_IN_A_YEAR", 365)) -REPLICATE_MODEL_NAME_WITH_ID_LENGTH = int(os.getenv("REPLICATE_MODEL_NAME_WITH_ID_LENGTH", 64)) +HOURS_IN_A_DAY: Final = int(os.getenv("HOURS_IN_A_DAY", 24)) +DAYS_IN_A_WEEK: Final = int(os.getenv("DAYS_IN_A_WEEK", 7)) +DAYS_IN_A_MONTH: Final = int(os.getenv("DAYS_IN_A_MONTH", 28)) +DAYS_IN_A_YEAR: Final = int(os.getenv("DAYS_IN_A_YEAR", 365)) +REPLICATE_MODEL_NAME_WITH_ID_LENGTH: Final = int(os.getenv("REPLICATE_MODEL_NAME_WITH_ID_LENGTH", 64)) #### TOKEN COUNTING #### -FUNCTION_DEFINITION_TOKEN_COUNT = int(os.getenv("FUNCTION_DEFINITION_TOKEN_COUNT", 9)) -SYSTEM_MESSAGE_TOKEN_COUNT = int(os.getenv("SYSTEM_MESSAGE_TOKEN_COUNT", 4)) -TOOL_CHOICE_OBJECT_TOKEN_COUNT = int(os.getenv("TOOL_CHOICE_OBJECT_TOKEN_COUNT", 4)) -DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT = int(os.getenv("DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT", 10)) -DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT = int(os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20)) -MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES = int(os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768)) -MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES = int(os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000)) -MAX_TILE_WIDTH = int(os.getenv("MAX_TILE_WIDTH", 512)) -MAX_TILE_HEIGHT = int(os.getenv("MAX_TILE_HEIGHT", 512)) -OPENAI_FILE_SEARCH_COST_PER_1K_CALLS = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) +FUNCTION_DEFINITION_TOKEN_COUNT: Final = int(os.getenv("FUNCTION_DEFINITION_TOKEN_COUNT", 9)) +SYSTEM_MESSAGE_TOKEN_COUNT: Final = int(os.getenv("SYSTEM_MESSAGE_TOKEN_COUNT", 4)) +TOOL_CHOICE_OBJECT_TOKEN_COUNT: Final = int(os.getenv("TOOL_CHOICE_OBJECT_TOKEN_COUNT", 4)) +DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT", 10)) +DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20)) +MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768)) +MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000)) +MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512)) +MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512)) +OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000)) +GROQ_BROWSER_VISIT_WEBSITE_COST_PER_CALL: Final = 1.0 / 1000 # Azure OpenAI Assistants feature costs # Source: https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/ -AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY = float( +AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY: Final = float( os.getenv("AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY", 0.1) # $0.1 USD per 1 GB/Day ) -AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS = float( +AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS: Final = float( os.getenv("AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS", 3.0) # $0.003 USD per 1K Tokens ) -AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS = float( +AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS: Final = float( os.getenv("AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS", 12.0) # $0.012 USD per 1K Tokens ) -AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY = float( +AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY: Final = float( os.getenv("AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY", 0.1) # $0.1 USD per 1 GB/Day (same as file search) ) -MIN_NON_ZERO_TEMPERATURE = float(os.getenv("MIN_NON_ZERO_TEMPERATURE", 0.0001)) +MIN_NON_ZERO_TEMPERATURE: Final = float(os.getenv("MIN_NON_ZERO_TEMPERATURE", 0.0001)) #### RELIABILITY #### -REPEATED_STREAMING_CHUNK_LIMIT = int( +REPEATED_STREAMING_CHUNK_LIMIT: Final = int( os.getenv("REPEATED_STREAMING_CHUNK_LIMIT", 100) ) # catch if model starts looping the same chunk while streaming. Uses high default to prevent false positives. # Shared maxsize for functools.lru_cache usage across hot paths. # Defaulted to 64 to avoid cache thrash in multi-model production workloads. -DEFAULT_MAX_LRU_CACHE_SIZE = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 64)) +DEFAULT_MAX_LRU_CACHE_SIZE: Final = int(os.getenv("DEFAULT_MAX_LRU_CACHE_SIZE", 64)) _REALTIME_BODY_CACHE_SIZE = 1000 # Keep realtime helper caches bounded; workloads rarely exceed 1k models/intents -INITIAL_RETRY_DELAY = float(os.getenv("INITIAL_RETRY_DELAY", 0.5)) -MAX_RETRY_DELAY = float(os.getenv("MAX_RETRY_DELAY", 8.0)) -JITTER = float(os.getenv("JITTER", 0.75)) +INITIAL_RETRY_DELAY: Final = float(os.getenv("INITIAL_RETRY_DELAY", 0.5)) +MAX_RETRY_DELAY: Final = float(os.getenv("MAX_RETRY_DELAY", 8.0)) +JITTER: Final = float(os.getenv("JITTER", 0.75)) DEFAULT_IN_MEMORY_TTL = int(os.getenv("DEFAULT_IN_MEMORY_TTL", 5)) # default time to live for the in-memory cache -DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE = int( +DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE: Final = int( os.getenv("DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 1000) ) # default max size for redis batch cache -DEFAULT_POLLING_INTERVAL = float( +DEFAULT_POLLING_INTERVAL: Final = float( os.getenv("DEFAULT_POLLING_INTERVAL", 0.03) ) # default polling interval for the scheduler -AZURE_OPERATION_POLLING_TIMEOUT = int(os.getenv("AZURE_OPERATION_POLLING_TIMEOUT", 120)) -AZURE_DOCUMENT_INTELLIGENCE_API_VERSION = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30")) -AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96)) -REDIS_SOCKET_TIMEOUT = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1)) -REDIS_CONNECTION_POOL_TIMEOUT = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5)) -REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5)) -REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)) -REDIS_CIRCUIT_BREAKER_ENABLED = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true" +AZURE_OPERATION_POLLING_TIMEOUT: Final = int(os.getenv("AZURE_OPERATION_POLLING_TIMEOUT", 120)) +AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30")) +AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96)) +REDIS_SOCKET_TIMEOUT: Final = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1)) +REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5)) +REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5)) +REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)) +REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true" # Seconds of idle before a Redis cluster connection is validated with a PING and # reconnected if dead, so a connection silently dropped by a cluster restart # (e.g. ElastiCache Serverless maintenance) is not reused while broken -REDIS_CLUSTER_HEALTH_CHECK_INTERVAL = 25 +REDIS_CLUSTER_HEALTH_CHECK_INTERVAL: Final = 25 # Default Redis major version to assume when version cannot be determined # Using 7 as it's the modern version that supports LPOP with count parameter -DEFAULT_REDIS_MAJOR_VERSION = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7)) -NON_LLM_CONNECTION_TIMEOUT = int( +DEFAULT_REDIS_MAJOR_VERSION: Final = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7)) +NON_LLM_CONNECTION_TIMEOUT: Final = int( os.getenv("NON_LLM_CONNECTION_TIMEOUT", 15) ) # timeout for adjacent services (e.g. jwt auth) -MAX_EXCEPTION_MESSAGE_LENGTH = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000)) -MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048)) -BEDROCK_MAX_POLICY_SIZE = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75)) -BEDROCK_MIN_THINKING_BUDGET_TOKENS = int(os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024)) +MAX_EXCEPTION_MESSAGE_LENGTH: Final = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000)) +MAX_STRING_LENGTH_PROMPT_IN_DB: Final = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048)) +BEDROCK_MAX_POLICY_SIZE: Final = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75)) +# One entry per distinct AWS credential-argument set. Per-user cost attribution passes the attributed +# identity as aws_session_name, so this bounds how many attributed identities keep a cached STS session. +BEDROCK_IAM_CACHE_MAX_ENTRIES: Final = 1000 +# Single-flight lock stripes over that cache. Only keys landing on the same stripe wait for each +# other, so a burst of distinct identities still resolves its credentials in parallel. +BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES: Final = 64 +# Retire a cached STS credential this many seconds before AWS expires it, so a request that reads it +# still has a usable credential for the whole call. +STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS: Final = 60 +BEDROCK_MIN_THINKING_BUDGET_TOKENS: Final = int(os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024)) # Anthropic's Messages API rejects thinking.budget_tokens < 1024. -ANTHROPIC_MIN_THINKING_BUDGET_TOKENS = 1024 -REPLICATE_POLLING_DELAY_SECONDS = float(os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5)) -DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS = int(os.getenv("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096)) -DEFAULT_OCI_CHAT_MAX_TOKENS = 4096 -TOGETHER_AI_4_B = int(os.getenv("TOGETHER_AI_4_B", 4)) -TOGETHER_AI_8_B = int(os.getenv("TOGETHER_AI_8_B", 8)) -TOGETHER_AI_21_B = int(os.getenv("TOGETHER_AI_21_B", 21)) -TOGETHER_AI_41_B = int(os.getenv("TOGETHER_AI_41_B", 41)) -TOGETHER_AI_80_B = int(os.getenv("TOGETHER_AI_80_B", 80)) -TOGETHER_AI_110_B = int(os.getenv("TOGETHER_AI_110_B", 110)) -TOGETHER_AI_EMBEDDING_150_M = int(os.getenv("TOGETHER_AI_EMBEDDING_150_M", 150)) -TOGETHER_AI_EMBEDDING_350_M = int(os.getenv("TOGETHER_AI_EMBEDDING_350_M", 350)) -QDRANT_SCALAR_QUANTILE = float(os.getenv("QDRANT_SCALAR_QUANTILE", 0.99)) -QDRANT_VECTOR_SIZE = int(os.getenv("QDRANT_VECTOR_SIZE", 1536)) -CACHED_STREAMING_CHUNK_DELAY = float(os.getenv("CACHED_STREAMING_CHUNK_DELAY", 0.02)) -AUDIO_SPEECH_CHUNK_SIZE = int( +ANTHROPIC_MIN_THINKING_BUDGET_TOKENS: Final = 1024 +REPLICATE_POLLING_DELAY_SECONDS: Final = float(os.getenv("REPLICATE_POLLING_DELAY_SECONDS", 0.5)) +DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096)) +DEFAULT_OCI_CHAT_MAX_TOKENS: Final = 4096 +TOGETHER_AI_4_B: Final = int(os.getenv("TOGETHER_AI_4_B", 4)) +TOGETHER_AI_8_B: Final = int(os.getenv("TOGETHER_AI_8_B", 8)) +TOGETHER_AI_21_B: Final = int(os.getenv("TOGETHER_AI_21_B", 21)) +TOGETHER_AI_41_B: Final = int(os.getenv("TOGETHER_AI_41_B", 41)) +TOGETHER_AI_80_B: Final = int(os.getenv("TOGETHER_AI_80_B", 80)) +TOGETHER_AI_110_B: Final = int(os.getenv("TOGETHER_AI_110_B", 110)) +TOGETHER_AI_EMBEDDING_150_M: Final = int(os.getenv("TOGETHER_AI_EMBEDDING_150_M", 150)) +TOGETHER_AI_EMBEDDING_350_M: Final = int(os.getenv("TOGETHER_AI_EMBEDDING_350_M", 350)) +QDRANT_SCALAR_QUANTILE: Final = float(os.getenv("QDRANT_SCALAR_QUANTILE", 0.99)) +QDRANT_VECTOR_SIZE: Final = int(os.getenv("QDRANT_VECTOR_SIZE", 1536)) +CACHED_STREAMING_CHUNK_DELAY: Final = float(os.getenv("CACHED_STREAMING_CHUNK_DELAY", 0.02)) +AUDIO_SPEECH_CHUNK_SIZE: Final = int( os.getenv("AUDIO_SPEECH_CHUNK_SIZE", 8192) ) # chunk_size for audio speech streaming. Balance between latency and memory usage -DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000)) +DEFAULT_MAX_TOKENS_FOR_TRITON: Final = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2000)) #### Networking settings #### # Sentinel used when `REQUEST_TIMEOUT` is unset: `litellm.request_timeout` keeps this # value so longer-running surfaces (Router `timeout or litellm.request_timeout`, @@ -385,70 +395,70 @@ DEFAULT_MAX_TOKENS_FOR_TRITON = int(os.getenv("DEFAULT_MAX_TOKENS_FOR_TRITON", 2 # `completion()` maps this sentinel down to 600s when the caller did not set a # per-request/model timeout—see ``CompletionTimeout.resolve`` in completion_timeout.py. MCP uses # dedicated timeouts (e.g. `MCP_CLIENT_TIMEOUT`), not `request_timeout`. -DEFAULT_REQUEST_TIMEOUT_SECONDS: float = 6000.0 +DEFAULT_REQUEST_TIMEOUT_SECONDS: Final[float] = 6000.0 # Pair used for default httpx clients when no custom timeout is passed: read/write # deadline and connect handshake (see ``http_handler`` cached handler paths). -COMPLETION_HTTP_FALLBACK_SECONDS: float = 600.0 -HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: float = 5.0 +COMPLETION_HTTP_FALLBACK_SECONDS: Final[float] = 600.0 +HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: Final[float] = 5.0 request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS)))) request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ -DEFAULT_A2A_AGENT_TIMEOUT: float = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes +DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes # Patterns that indicate a localhost/internal URL in A2A agent cards that should be # replaced with the original base_url. This is a common misconfiguration where # developers deploy agents with development URLs in their agent cards. -LOCALHOST_URL_PATTERNS: List[str] = [ +LOCALHOST_URL_PATTERNS: Final[list[str]] = [ "localhost", "127.0.0.1", "0.0.0.0", "[::1]", # IPv6 localhost ] # Patterns in error messages that indicate a connection failure -CONNECTION_ERROR_PATTERNS: List[str] = [ +CONNECTION_ERROR_PATTERNS: Final[list[str]] = [ "connect", "connection", "network", "refused", ] -STREAM_SSE_DONE_STRING: str = "[DONE]" -STREAM_SSE_DATA_PREFIX: str = "data: " +STREAM_SSE_DONE_STRING: Final[str] = "[DONE]" +STREAM_SSE_DATA_PREFIX: Final[str] = "data: " ### SPEND TRACKING ### -DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND = float( +DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = float( os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400) ) # price per second for a100 80GB -FIREWORKS_AI_56_B_MOE = int(os.getenv("FIREWORKS_AI_56_B_MOE", 56)) -FIREWORKS_AI_176_B_MOE = int(os.getenv("FIREWORKS_AI_176_B_MOE", 176)) -FIREWORKS_AI_4_B = int(os.getenv("FIREWORKS_AI_4_B", 4)) -FIREWORKS_AI_16_B = int(os.getenv("FIREWORKS_AI_16_B", 16)) -FIREWORKS_AI_80_B = int(os.getenv("FIREWORKS_AI_80_B", 80)) +FIREWORKS_AI_56_B_MOE: Final = int(os.getenv("FIREWORKS_AI_56_B_MOE", 56)) +FIREWORKS_AI_176_B_MOE: Final = int(os.getenv("FIREWORKS_AI_176_B_MOE", 176)) +FIREWORKS_AI_4_B: Final = int(os.getenv("FIREWORKS_AI_4_B", 4)) +FIREWORKS_AI_16_B: Final = int(os.getenv("FIREWORKS_AI_16_B", 16)) +FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80)) #### Logging callback constants #### -REDACTED_BY_LITELM_STRING = "REDACTED_BY_LITELM" -MAX_LANGFUSE_INITIALIZED_CLIENTS = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)) -LOGGING_WORKER_CONCURRENCY = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 -LOGGING_WORKER_MAX_QUEUE_SIZE = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) -LOGGING_WORKER_MAX_TIME_PER_COROUTINE = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) -LOGGING_WORKER_CLEAR_PERCENTAGE = int( +REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM" +MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)) +LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 +LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) +LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) +LOGGING_WORKER_CLEAR_PERCENTAGE: Final = int( os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50) ) # Percentage of queue to clear (default: 50%) -MAX_ITERATIONS_TO_CLEAR_QUEUE = int(os.getenv("MAX_ITERATIONS_TO_CLEAR_QUEUE", 200)) -MAX_TIME_TO_CLEAR_QUEUE = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0)) -LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS = float( +MAX_ITERATIONS_TO_CLEAR_QUEUE: Final = int(os.getenv("MAX_ITERATIONS_TO_CLEAR_QUEUE", 200)) +MAX_TIME_TO_CLEAR_QUEUE: Final = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0)) +LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float( os.getenv("LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS", 0.5) ) # Cooldown time in seconds before allowing another aggressive clear (default: 0.5s) -DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv( +DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) -LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED = 499 +LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED: Final = 499 -EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds -EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float( +EMAIL_BUDGET_ALERT_TTL: Final = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds +EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float( os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8) ) # 80% of max budget ############### LLM Provider Constants ############### ### ANTHROPIC CONSTANTS ### ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01") -ANTHROPIC_SKILLS_API_BETA_VERSION = "skills-2025-10-02" -ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES = { +ANTHROPIC_SKILLS_API_BETA_VERSION: Final = "skills-2025-10-02" +ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = { "low": 1, "medium": 5, "high": 10, @@ -456,19 +466,19 @@ ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES = { # LiteLLM standard web search tool name # Used for web search interception across providers -LITELLM_WEB_SEARCH_TOOL_NAME = "litellm_web_search" +LITELLM_WEB_SEARCH_TOOL_NAME: Final = "litellm_web_search" -DEFAULT_IMAGE_ENDPOINT_MODEL = "dall-e-2" -DEFAULT_VIDEO_ENDPOINT_MODEL = "sora-2" +DEFAULT_IMAGE_ENDPOINT_MODEL: Final = "dall-e-2" +DEFAULT_VIDEO_ENDPOINT_MODEL: Final = "sora-2" -DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS = int(os.getenv("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8)) +DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS: Final = int(os.getenv("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8)) ### DATAFORSEO CONSTANTS ### -DEFAULT_DATAFORSEO_LOCATION_CODE = int( +DEFAULT_DATAFORSEO_LOCATION_CODE: Final = int( os.getenv("DEFAULT_DATAFORSEO_LOCATION_CODE", 2250) ) # Default to France (2250) - lower number, commonly used location -LITELLM_CHAT_PROVIDERS = [ +LITELLM_CHAT_PROVIDERS: Final = [ "openai", "openai_like", "bytez", @@ -565,7 +575,7 @@ LITELLM_CHAT_PROVIDERS = [ "amazon_nova", ] -LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS = [ +LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [ "openai", "azure", "hosted_vllm", @@ -573,7 +583,7 @@ LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS = [ ] -OPENAI_CHAT_COMPLETION_PARAMS = [ +OPENAI_CHAT_COMPLETION_PARAMS: Final = [ "functions", "function_call", "temperature", @@ -622,22 +632,22 @@ OPENAI_CHAT_COMPLETION_PARAMS = [ "store", ] -OPENAI_TRANSCRIPTION_PARAMS = [ +OPENAI_TRANSCRIPTION_PARAMS: Final = [ "language", "response_format", "timestamp_granularities", ] -OPENAI_EMBEDDING_PARAMS = ["dimensions", "encoding_format", "user"] +OPENAI_EMBEDDING_PARAMS: Final = ["dimensions", "encoding_format", "user"] -DEFAULT_EMBEDDING_PARAM_VALUES = { +DEFAULT_EMBEDDING_PARAM_VALUES: Final = { **{k: None for k in OPENAI_EMBEDDING_PARAMS}, "model": None, "custom_llm_provider": "", "input": None, } -DEFAULT_CHAT_COMPLETION_PARAM_VALUES = { +DEFAULT_CHAT_COMPLETION_PARAM_VALUES: Final = { "functions": None, "function_call": None, "temperature": None, @@ -685,7 +695,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES = { "context_management": None, } -openai_compatible_endpoints: List = [ +openai_compatible_endpoints: Final[list] = [ "api.perplexity.ai", "api.endpoints.anyscale.com/v1", "api.deepinfra.com/v1/openai", @@ -731,7 +741,7 @@ openai_compatible_endpoints: List = [ ] -openai_compatible_providers: List = [ +openai_compatible_providers: Final[list] = [ "anyscale", "groq", "nvidia_nim", @@ -796,7 +806,7 @@ openai_compatible_providers: List = [ "darkbloom", "meta", # Meta Model API (Muse Spark) - JSON-configured provider ] -openai_text_completion_compatible_providers: List = [ # providers that support `/v1/completions` +openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions` "together_ai", "fireworks_ai", "hosted_vllm", @@ -819,14 +829,14 @@ openai_text_completion_compatible_providers: List = [ # providers that support "hyperbolic", "wandb", ] -_openai_like_providers: List = [ +_openai_like_providers: Final[list] = [ "predibase", "databricks", "lemonade", "watsonx", ] # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk # well supported replicate llms -replicate_models: set = set( +replicate_models: Final[set] = set( [ # llama replicate supported LLMs "replicate/llama-2-70b-chat:2796ee9483c3fd7aa2e171d38f4ca12251a30609463dcfd4cd76703f22e96cdf", @@ -843,7 +853,7 @@ replicate_models: set = set( ] ) -clarifai_models: set = set( +clarifai_models: Final[set] = set( [ "clarifai/openai.chat-completion.gpt-oss-20b", "clarifai/qwen.qwenLM.Qwen3-30B-A3B-Instruct-2507", @@ -879,7 +889,7 @@ clarifai_models: set = set( ) -huggingface_models: set = set( +huggingface_models: Final[set] = set( [ "meta-llama/Llama-2-7b-hf", "meta-llama/Llama-2-7b-chat-hf", @@ -895,14 +905,14 @@ huggingface_models: set = set( "meta-llama/Llama-2-70b-chat", ] ) # these have been tested on extensively. But by default all text2text-generation and text-generation models are supported by liteLLM. - https://docs.litellm.ai/docs/providers -empower_models = set( +empower_models: Final = set( [ "empower/empower-functions", "empower/empower-functions-small", ] ) -together_ai_models: set = set( +together_ai_models: Final[set] = set( [ # llama llms - chat "togethercomputer/llama-2-70b-chat", @@ -936,7 +946,7 @@ together_ai_models: set = set( # supports all together ai models, just pass in the model id e.g. completion(model="together_computer/replit_code_3b",...) -baseten_models: set = set( +baseten_models: Final[set] = set( [ "qvv0xeq", "q841o8w", @@ -944,7 +954,7 @@ baseten_models: set = set( ] ) # FALCON 7B # WizardLM # Mosaic ML -featherless_ai_models: set = set( +featherless_ai_models: Final[set] = set( [ "featherless-ai/Qwerky-72B", "featherless-ai/Qwerky-QwQ-32B", @@ -958,7 +968,7 @@ featherless_ai_models: set = set( ] ) -nebius_models: set = set( +nebius_models: Final[set] = set( [ # deepseek models "deepseek-ai/DeepSeek-R1-0528", @@ -1012,7 +1022,7 @@ nebius_models: set = set( ] ) -dashscope_models: set = set( +dashscope_models: Final[set] = set( [ "qwen-turbo", "qwen-plus", @@ -1027,7 +1037,7 @@ dashscope_models: set = set( ] ) -nebius_embedding_models: set = set( +nebius_embedding_models: Final[set] = set( [ "BAAI/bge-en-icl", "BAAI/bge-multilingual-gemma2", @@ -1035,7 +1045,7 @@ nebius_embedding_models: set = set( ] ) -WANDB_MODELS: set = set( +WANDB_MODELS: Final[set] = set( [ # openai models "openai/gpt-oss-120b", @@ -1064,7 +1074,7 @@ WANDB_MODELS: set = set( ] ) -modelscope_models: set = set( +modelscope_models: Final[set] = set( [ # Qwen series models "Qwen/Qwen3-0.6B", @@ -1131,7 +1141,7 @@ BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[ "nova", ] -BEDROCK_CONVERSE_MODELS = [ +BEDROCK_CONVERSE_MODELS: Final = [ "qwen.qwen3-coder-480b-a35b-v1:0", "qwen.qwen3-coder-next", "qwen.qwen3-235b-a22b-2507-v1:0", @@ -1192,8 +1202,8 @@ BEDROCK_CONVERSE_MODELS = [ ] -open_ai_embedding_models: set = set(["text-embedding-ada-002"]) -cohere_embedding_models: set = set( +open_ai_embedding_models: Final[set] = set(["text-embedding-ada-002"]) +cohere_embedding_models: Final[set] = set( [ "embed-v4.0", "embed-english-v3.0", @@ -1204,7 +1214,7 @@ cohere_embedding_models: set = set( "embed-multilingual-v2.0", ] ) -bedrock_embedding_models: set = set( +bedrock_embedding_models: Final[set] = set( [ "amazon.titan-embed-text-v1", "amazon.nova-2-multimodal-embeddings-v1:0", @@ -1215,7 +1225,7 @@ bedrock_embedding_models: set = set( ] ) -known_tokenizer_config = { +known_tokenizer_config: Final = { "mistralai/Mistral-7B-Instruct-v0.1": { "tokenizer": { "chat_template": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}", @@ -1273,33 +1283,34 @@ known_tokenizer_config = { } -OPENAI_FINISH_REASONS = [ +OPENAI_FINISH_REASONS: Final = [ "stop", "length", "function_call", "tool_calls", "content_filter", ] -HUMANLOOP_PROMPT_CACHE_TTL_SECONDS = int(os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60)) # 1 minute +HUMANLOOP_PROMPT_CACHE_TTL_SECONDS: Final = int(os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60)) # 1 minute RESPONSE_FORMAT_TOOL_NAME = "json_tool_call" # default tool name used when converting response format to tool call ########################### Logging Callback Constants ########################### -AZURE_STORAGE_MSFT_VERSION = "2019-07-07" -PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES = int( +AZURE_STORAGE_MSFT_VERSION: Final = "2019-07-07" +AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX: Final = "core.windows.net" +PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = int( os.getenv("PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5) ) -CLOUDZERO_EXPORT_INTERVAL_MINUTES = int(os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60)) -MCP_TOOL_NAME_PREFIX = "mcp_tool" -MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)) +CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = int(os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60)) +MCP_TOOL_NAME_PREFIX: Final = "mcp_tool" +MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)) # Headers to control callbacks -X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks" -LITELLM_METADATA_FIELD = "litellm_metadata" -OLD_LITELLM_METADATA_FIELD = "metadata" -RETURN_RAW_MODEL_NAME_METADATA_KEY = "_complexity_router_return_raw_model_name" -INTERNAL_CALL_ORIGIN_METADATA_KEY = "internal_call_origin" -LITELLM_TRUNCATED_PAYLOAD_FIELD = "litellm_truncated" -LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE = ( +X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks" +LITELLM_METADATA_FIELD: Final = "litellm_metadata" +OLD_LITELLM_METADATA_FIELD: Final = "metadata" +RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name" +INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" +LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" +LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " "Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.). " "To increase the truncation limit, set `MAX_STRING_LENGTH_PROMPT_IN_DB` in your env." @@ -1310,25 +1321,25 @@ LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE = ( # Standard headers that are always checked for customer/end-user ID (no configuration required) # These headers work out-of-the-box for tools like Claude Code that support custom headers -STANDARD_CUSTOMER_ID_HEADERS = [ +STANDARD_CUSTOMER_ID_HEADERS: Final = [ "x-litellm-customer-id", "x-litellm-end-user-id", ] -MAX_SPENDLOG_ROWS_TO_QUERY = int( +MAX_SPENDLOG_ROWS_TO_QUERY: Final = int( os.getenv("MAX_SPENDLOG_ROWS_TO_QUERY", 1_000_000) ) # if spendLogs has more than 1M rows, do not query the DB -DEFAULT_SOFT_BUDGET = float( +DEFAULT_SOFT_BUDGET: Final = float( os.getenv("DEFAULT_SOFT_BUDGET", 50.0) ) # by default all litellm proxy keys have a soft budget of 50.0 # makes it clear this is a rate limit error for a litellm virtual key -RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY = "LiteLLM Virtual Key user_api_key_hash" +RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY: Final = "LiteLLM Virtual Key user_api_key_hash" # Python garbage collection threshold configuration # Format: "gen0,gen1,gen2" e.g., "1000,50,50" -PYTHON_GC_THRESHOLD = os.getenv("PYTHON_GC_THRESHOLD") +PYTHON_GC_THRESHOLD: Final = os.getenv("PYTHON_GC_THRESHOLD") # pass through route constansts -BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES = [ +BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES: Final = [ "agents/", "knowledgebases/", "flows/", @@ -1341,7 +1352,7 @@ BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES = [ # Headers that are safe to forward from incoming requests to Vertex AI # Using an allowlist approach for security - only forward headers we explicitly trust -ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS = { +ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: Final = { "anthropic-beta", # Required for Anthropic features like extended context windows "content-type", # Required for request body parsing } @@ -1349,34 +1360,34 @@ ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS = { # Prefix for headers that should be forwarded to the provider with the prefix stripped # e.g., 'x-pass-anthropic-beta: value' becomes 'anthropic-beta: value' # Works for all LLM pass-through endpoints (Vertex AI, Anthropic, Bedrock, etc.) -PASS_THROUGH_HEADER_PREFIX = "x-pass-" +PASS_THROUGH_HEADER_PREFIX: Final = "x-pass-" -BASE_MCP_ROUTE = "/mcp" +BASE_MCP_ROUTE: Final = "/mcp" -BATCH_STATUS_POLL_INTERVAL_SECONDS = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour -BATCH_STATUS_POLL_MAX_ATTEMPTS = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours +BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour +BATCH_STATUS_POLL_MAX_ATTEMPTS: Final = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours -HEALTH_CHECK_TIMEOUT_SECONDS = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds -_background_health_check_max_tokens_env = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS") +HEALTH_CHECK_TIMEOUT_SECONDS: Final = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds +_background_health_check_max_tokens_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS") try: - _raw_background_health_check_max_tokens = ( + _raw_background_health_check_max_tokens: Final = ( _background_health_check_max_tokens_env.strip() if _background_health_check_max_tokens_env is not None else "" ) - BACKGROUND_HEALTH_CHECK_MAX_TOKENS: Optional[int] = ( + BACKGROUND_HEALTH_CHECK_MAX_TOKENS: int | None = ( int(_raw_background_health_check_max_tokens) if _raw_background_health_check_max_tokens else None ) except (ValueError, TypeError): BACKGROUND_HEALTH_CHECK_MAX_TOKENS = None -_background_health_check_max_tokens_reasoning_env = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING") +_background_health_check_max_tokens_reasoning_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING") try: - _raw_background_health_check_max_tokens_reasoning = ( + _raw_background_health_check_max_tokens_reasoning: Final = ( _background_health_check_max_tokens_reasoning_env.strip() if _background_health_check_max_tokens_reasoning_env is not None else "" ) - BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: Optional[int] = ( + BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING: int | None = ( int(_raw_background_health_check_max_tokens_reasoning) if _raw_background_health_check_max_tokens_reasoning else None @@ -1384,140 +1395,141 @@ try: except (ValueError, TypeError): BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING = None -LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check" -LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli" -LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs" +LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME: Final = "litellm-internal-health-check" +LITTELM_CLI_SERVICE_ACCOUNT_NAME: Final = "litellm-cli" +LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME: Final = "litellm_internal_jobs" # Stable identifier substituted in place of the master key on UserAPIKeyAuth # objects so the master key (or its hash) never propagates to spend logs, # Prometheus metrics, audit trails, or any other downstream consumer. -LITELLM_PROXY_MASTER_KEY_ALIAS = "litellm_proxy_master_key" +LITELLM_PROXY_MASTER_KEY_ALIAS: Final = "litellm_proxy_master_key" # Marker placed in ``model_call_details`` on a synthetic ``Logging`` object that # records a proxy-gate error (auth/rate-limit rejection) for a request that never # reached an upstream provider. Tracing callbacks key off it to avoid fabricating # an LLM-call span for a call that did not happen. See # ``ProxyLogging._handle_logging_proxy_only_error``. -LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL = "litellm_no_upstream_llm_call" +LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: Final = "litellm_no_upstream_llm_call" # Key Rotation Constants -LITELLM_KEY_ROTATION_ENABLED = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false") -LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int( +LITELLM_KEY_ROTATION_ENABLED: Final = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false") +LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS: Final = int( os.getenv("LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS", 86400) ) # 24 hours default -LITELLM_KEY_ROTATION_GRACE_PERIOD: str = os.getenv( +LITELLM_KEY_ROTATION_GRACE_PERIOD: Final[str] = os.getenv( "LITELLM_KEY_ROTATION_GRACE_PERIOD", "" ) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default) -LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int( +LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS: Final = int( os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600) ) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation -UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" +UI_SESSION_TOKEN_TEAM_ID: Final = "litellm-dashboard" LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED = os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED", "false") -LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS = int( +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS: Final = int( os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_INTERVAL_SECONDS", 86400) ) # 24 hours default -LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE = int( +LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = int( os.getenv("LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE", 1000) ) -LITELLM_PROXY_ADMIN_NAME = "default_user_id" -LITELLM_PROXY_BUDGET_NAME = "litellm-proxy-budget" -GLOBAL_PROXY_SPEND_CACHE_KEY = f"{LITELLM_PROXY_ADMIN_NAME}:spend" +LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id" +LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget" +GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend" ########################### CLI SSO AUTHENTICATION CONSTANTS ########################### -LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli" -LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token" -CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session" -CLI_SSO_SESSION_TTL_SECONDS = 600 -CLI_SESSION_KEY_PREFIX = "cli-session" +LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli" +LITELLM_CLI_SESSION_TOKEN_PREFIX: Final = "litellm-session-token" +CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session" +CLI_SSO_SESSION_TTL_SECONDS: Final = 600 +CLI_SESSION_KEY_PREFIX: Final = "cli-session" # Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility -CLI_JWT_EXPIRATION_HOURS = int( +CLI_JWT_EXPIRATION_HOURS: Final = int( os.getenv("CLI_JWT_EXPIRATION_HOURS") or os.getenv("LITELLM_CLI_JWT_EXPIRATION_HOURS") or 24 ) # Comma-separated allowlisted OIDC claim map for CLI SSO polling, e.g. # "employment_type->acme_employment_type,org_info.department->department" -CLI_SSO_CLAIM_MAP = os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or "" -CLI_SSO_CLAIM_MAX_SCALAR_LENGTH = 1024 +CLI_SSO_CLAIM_MAP: Final = os.getenv("CLI_SSO_CLAIM_MAP") or os.getenv("LITELLM_CLI_SSO_CLAIM_MAP") or "" +CLI_SSO_CLAIM_MAX_SCALAR_LENGTH: Final = 1024 ########################### UI SESSION DURATION ########################### # Duration for UI login session (username/password, SSO, invitation links). Format: "30s", "30m", "24h", "7d" # Does NOT apply to EXPERIMENTAL_UI_LOGIN flow, which intentionally uses a fixed 10-minute expiry for security. -LITELLM_UI_SESSION_DURATION = os.getenv("LITELLM_UI_SESSION_DURATION", "24h") +LITELLM_UI_SESSION_DURATION: Final = os.getenv("LITELLM_UI_SESSION_DURATION", "24h") ########################### DB CRON JOB NAMES ########################### -DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job" -DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job" -PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics" -CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data" -MAVVRIK_FOCUS_EXPORT_JOB_NAME = "mavvrik_focus_export_usage_data" -CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)) -SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" -KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" -EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME = "litellm_expired_ui_session_key_cleanup_job" -SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) -SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) +DB_SPEND_UPDATE_JOB_NAME: Final = "db_spend_update_job" +DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME: Final = "db_daily_tag_spend_update_job" +PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME: Final = "prometheus_emit_budget_metrics" +CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data" +MAVVRIK_FOCUS_EXPORT_JOB_NAME: Final = "mavvrik_focus_export_usage_data" +CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)) +SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup" +KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job" +EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job" +SPEND_LOG_RUN_LOOPS: Final = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) +SPEND_LOG_CLEANUP_BATCH_SIZE: Final = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)) -SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float( +SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float( os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5) ) -TOOL_SPEND_TOP_TOOLS = 100 -SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") -SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) -SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) -SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) -SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) -DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute -PROXY_BUDGET_RESCHEDULER_MIN_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) -PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) -MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) -MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) -STALE_OBJECT_CLEANUP_BATCH_SIZE = max(1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))) +TOOL_SPEND_TOP_TOOLS: Final = 100 +SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") +SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) +SPEND_LOG_WRITE_BATCH_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_BYTES", 2_000_000))) +SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) +SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0)) +SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) +DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute +PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) +PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) +MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) +MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) +STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000))) # Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and # CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on # installations with large numbers of stale managed objects). -_batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() -PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true" -PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) -PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 -PROXY_CONFIG_RELOAD_INTERVAL_SECONDS = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) +_batch_polling_env: Final = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() +PROXY_BATCH_POLLING_ENABLED: Final = _batch_polling_env == "true" +PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) +PROXY_BATCH_WRITE_AT: Final = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 +PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) # APScheduler Configuration - MEMORY LEAK FIX # These settings prevent memory leaks in APScheduler's normalize() and _apply_jitter() functions -APSCHEDULER_COALESCE = os.getenv("APSCHEDULER_COALESCE", "True").lower() in [ +APSCHEDULER_COALESCE: Final = os.getenv("APSCHEDULER_COALESCE", "True").lower() in [ "true", "1", ] # collapse many missed runs into one -APSCHEDULER_MISFIRE_GRACE_TIME = int( +APSCHEDULER_MISFIRE_GRACE_TIME: Final = int( os.getenv("APSCHEDULER_MISFIRE_GRACE_TIME", 3600) ) # ignore runs older than 1 hour (was 120) -APSCHEDULER_MAX_INSTANCES = int(os.getenv("APSCHEDULER_MAX_INSTANCES", 1)) # prevent concurrent job instances -APSCHEDULER_REPLACE_EXISTING = os.getenv("APSCHEDULER_REPLACE_EXISTING", "True").lower() in [ +APSCHEDULER_MAX_INSTANCES: Final = int(os.getenv("APSCHEDULER_MAX_INSTANCES", 1)) # prevent concurrent job instances +APSCHEDULER_REPLACE_EXISTING: Final = os.getenv("APSCHEDULER_REPLACE_EXISTING", "True").lower() in [ "true", "1", ] # always replace existing jobs # The number of tag entries are higher than number of user, team entries. This leads to a higher QPS. # This will run tag spcific tasks at a later time to smooth QPS -DAILY_TAG_SPEND_BATCH_MULTIPLIER = 2.3 +DAILY_TAG_SPEND_BATCH_MULTIPLIER: Final = 2.3 -DEFAULT_HEALTH_CHECK_INTERVAL = int(os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300)) # 5 minutes -DEFAULT_SHARED_HEALTH_CHECK_TTL = int( +DEFAULT_HEALTH_CHECK_INTERVAL: Final = int(os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300)) # 5 minutes +DEFAULT_SHARED_HEALTH_CHECK_TTL: Final = int( os.getenv("DEFAULT_SHARED_HEALTH_CHECK_TTL", 300) ) # 5 minutes - TTL for cached health check results -DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL = int( +DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL: Final = int( os.getenv("DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL", 60) ) # 1 minute - TTL for health check lock -DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER = 2 # health state is stale after interval * this -PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS = int(os.getenv("PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS", 9)) -DEFAULT_MODEL_CREATED_AT_TIME = int( +DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER: Final = 2 # health state is stale after interval * this +PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS: Final = int(os.getenv("PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS", 9)) +DEFAULT_MODEL_CREATED_AT_TIME: Final = int( os.getenv("DEFAULT_MODEL_CREATED_AT_TIME", 1677610602) ) # returns on `/models` endpoint -DEFAULT_SLACK_ALERTING_THRESHOLD = int(os.getenv("DEFAULT_SLACK_ALERTING_THRESHOLD", 300)) -MAX_TEAM_LIST_LIMIT = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20)) -MAX_POLICY_ESTIMATE_IMPACT_ROWS = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000)) +DEFAULT_SLACK_ALERTING_THRESHOLD: Final = int(os.getenv("DEFAULT_SLACK_ALERTING_THRESHOLD", 300)) +MAX_TEAM_LIST_LIMIT: Final = int(os.getenv("MAX_TEAM_LIST_LIMIT", 20)) +MAX_POLICY_ESTIMATE_IMPACT_ROWS: Final = int(os.getenv("MAX_POLICY_ESTIMATE_IMPACT_ROWS", 1000)) DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD", 0.7)) -LENGTH_OF_LITELLM_GENERATED_KEY = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16)) -MINIMUM_CUSTOM_KEY_LENGTH = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16)) -SECRET_MANAGER_REFRESH_INTERVAL = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400)) -LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ +LENGTH_OF_LITELLM_GENERATED_KEY: Final = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY", 16)) +MINIMUM_CUSTOM_KEY_LENGTH: Final = int(os.getenv("MINIMUM_CUSTOM_KEY_LENGTH", 16)) +SECRET_MANAGER_REFRESH_INTERVAL: Final = int(os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400)) +LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [ "default_internal_user_params", "default_team_params", "public_mcp_servers", @@ -1535,21 +1547,21 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES = [ "anthropic_prompt_caching_ttl", "max_ui_session_budget", ] -SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] +SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"] DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)) -DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) +DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)) # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding # staleness so a transient DB error (which surfaces as an empty list) cannot # hide a real group for long. -DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL = 10 +DEFAULT_MCP_ACCESS_GROUP_NEGATIVE_CACHE_TTL: Final = 10 # Maximum number of comma-separated MCP server / access-group tokens accepted # in a single ``/{name1,name2,...}/mcp`` URL. Bounds the per-request DB / cache # fan-out an authenticated caller can trigger by stuffing the path with tokens. -DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS = 16 +DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16 # Sentry Scrubbing Configuration -SENTRY_DENYLIST = [ +SENTRY_DENYLIST: Final = [ # API Keys and Tokens "api_key", "token", @@ -1605,7 +1617,7 @@ SENTRY_DENYLIST = [ "proxy_key", "environment_variables", ] -SENTRY_PII_DENYLIST = [ +SENTRY_PII_DENYLIST: Final = [ "user_id", "email", "phone", @@ -1616,42 +1628,79 @@ SENTRY_PII_DENYLIST = [ ] # CoroutineChecker cache configuration -COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int(os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000)) +COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY: Final = int(os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000)) ########################### RAG Text Splitter Constants ########################### -DEFAULT_CHUNK_SIZE = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000)) -DEFAULT_CHUNK_OVERLAP = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200)) +DEFAULT_CHUNK_SIZE: Final = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000)) +DEFAULT_CHUNK_OVERLAP: Final = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200)) ########################### S3 Vectors RAG Constants ########################### -S3_VECTORS_DEFAULT_DIMENSION = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 1024)) -S3_VECTORS_DEFAULT_DISTANCE_METRIC = str(os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine")) -S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS = ["source_text"] +S3_VECTORS_DEFAULT_DIMENSION: Final = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 1024)) +S3_VECTORS_DEFAULT_DISTANCE_METRIC: Final = str(os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine")) +S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS: Final = ["source_text"] ########################### Microsoft SSO Constants ########################### -MICROSOFT_USER_EMAIL_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName")) -MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName")) -MICROSOFT_USER_ID_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id")) -MICROSOFT_USER_FIRST_NAME_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "givenName")) -MICROSOFT_USER_LAST_NAME_ATTRIBUTE = str(os.getenv("MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "surname")) +MICROSOFT_USER_EMAIL_ATTRIBUTE: Final = str(os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName")) +MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE: Final = str(os.getenv("MICROSOFT_USER_DISPLAY_NAME_ATTRIBUTE", "displayName")) +MICROSOFT_USER_ID_ATTRIBUTE: Final = str(os.getenv("MICROSOFT_USER_ID_ATTRIBUTE", "id")) +MICROSOFT_USER_FIRST_NAME_ATTRIBUTE: Final = str(os.getenv("MICROSOFT_USER_FIRST_NAME_ATTRIBUTE", "givenName")) +MICROSOFT_USER_LAST_NAME_ATTRIBUTE: Final = str(os.getenv("MICROSOFT_USER_LAST_NAME_ATTRIBUTE", "surname")) # Maximum payload size (in bytes) to fully serialize for DEBUG logging. # Payloads larger than this are truncated to avoid multi-second json.dumps blocking the response. -MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG = int(os.getenv("MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG", 102400)) # 100 KB +MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG: Final = int(os.getenv("MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG", 102400)) # 100 KB # Policy template enrichment -MAX_COMPETITOR_NAMES = int(os.getenv("MAX_COMPETITOR_NAMES", 100)) -COMPETITOR_LLM_TEMPERATURE = float(os.getenv("COMPETITOR_LLM_TEMPERATURE", 0.3)) -DEFAULT_COMPETITOR_DISCOVERY_MODEL = "gpt-4o-mini" +MAX_COMPETITOR_NAMES: Final = int(os.getenv("MAX_COMPETITOR_NAMES", 100)) +COMPETITOR_LLM_TEMPERATURE: Final = float(os.getenv("COMPETITOR_LLM_TEMPERATURE", 0.3)) +DEFAULT_COMPETITOR_DISCOVERY_MODEL: Final = "gpt-4o-mini" # Advisor tool orchestration # Providers that support advisor_20260301 natively (no LiteLLM orchestration needed). # Add vertex_ai here once verified. -ADVISOR_NATIVE_PROVIDERS: frozenset = frozenset({"anthropic"}) +ADVISOR_NATIVE_PROVIDERS: Final[frozenset] = frozenset({"anthropic"}) # Hard cap on advisor iterations per request to prevent runaway loops. -ADVISOR_MAX_USES: int = 5 +ADVISOR_MAX_USES: Final[int] = 5 # Description injected into the synthetic advisor tool definition sent to non-native providers. -ADVISOR_TOOL_DESCRIPTION: str = ( +ADVISOR_TOOL_DESCRIPTION: Final[str] = ( "Consult a highly intelligent advisor model when you need expert guidance, " "want to verify your reasoning, or face a complex decision. " "Describe your question or challenge clearly in the 'question' field." ) + +# Headers that must be stripped from a provider exception before it's forwarded as +# the proxy's own HTTP response, or they conflict with the framing the proxy sets. +HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset( + { + "content-length", + "transfer-encoding", + "content-encoding", + "content-type", + "set-cookie", + "cookie", + "proxy-authenticate", + "proxy-authorization", + } +) + +# Browser-facing security headers that a malicious or misconfigured upstream +# provider must not be able to set on the proxy's own response. +BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( + { + "access-control-allow-origin", + "access-control-allow-credentials", + "access-control-allow-methods", + "access-control-allow-headers", + "access-control-expose-headers", + "content-security-policy", + "content-security-policy-report-only", + "clear-site-data", + "strict-transport-security", + "x-frame-options", + "cross-origin-opener-policy", + "cross-origin-embedder-policy", + "cross-origin-resource-policy", + } +) + +UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS diff --git a/litellm/containers/endpoint_factory.py b/litellm/containers/endpoint_factory.py index bebdfa2f9e6..09bc7eda41f 100644 --- a/litellm/containers/endpoint_factory.py +++ b/litellm/containers/endpoint_factory.py @@ -8,9 +8,10 @@ that use the generic container handler. import asyncio import contextvars import json +from collections.abc import Callable from functools import partial from pathlib import Path -from typing import Any, Callable, Dict, List, Literal, Optional, Type +from typing import Any, Final, Literal import litellm from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT @@ -27,45 +28,45 @@ from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client # Response type mapping -RESPONSE_TYPES: Dict[str, Type] = { +RESPONSE_TYPES: Final[dict[str, type]] = { "ContainerFileListResponse": ContainerFileListResponse, "ContainerFileObject": ContainerFileObject, "DeleteContainerFileResponse": DeleteContainerFileResponse, } -def _load_endpoints_config() -> Dict: +def _load_endpoints_config() -> dict: """Load the endpoints configuration from JSON file.""" - config_path = Path(__file__).parent / "endpoints.json" + config_path: Final = Path(__file__).parent / "endpoints.json" with open(config_path) as f: return json.load(f) -def create_sync_endpoint_function(endpoint_config: Dict) -> Callable: +def create_sync_endpoint_function(endpoint_config: dict) -> Callable: """ Create a sync SDK function from endpoint config. Uses the generic container handler instead of individual handler methods. """ - endpoint_name = endpoint_config["name"] - response_type = RESPONSE_TYPES.get(endpoint_config["response_type"]) - path_params = endpoint_config.get("path_params", []) + endpoint_name: Final = endpoint_config["name"] + response_type: Final = RESPONSE_TYPES.get(endpoint_config["response_type"]) + path_params: Final = endpoint_config.get("path_params", []) @client def endpoint_func( timeout: int = 600, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ): - local_vars = locals() + local_vars: Final = locals() try: resolved_custom_llm_provider: str = custom_llm_provider - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id") - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id") + _is_async: Final = kwargs.pop("async_call", False) is True # Check for mock response mock_response = kwargs.get("mock_response") @@ -90,17 +91,15 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable: custom_llm_provider=resolved_custom_llm_provider, litellm_params=litellm_params, ) - container_provider_config: Optional[BaseContainerConfig] = ( - ProviderConfigManager.get_provider_container_config( - provider=litellm.LlmProviders(resolved_custom_llm_provider), - ) + container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config( + provider=litellm.LlmProviders(resolved_custom_llm_provider), ) if container_provider_config is None: raise ValueError(f"Container provider config not found for: {resolved_custom_llm_provider}") # Build optional params for logging - optional_params = {k: kwargs.get(k) for k in path_params if k in kwargs} + optional_params: Final = {k: kwargs.get(k) for k in path_params if k in kwargs} # Pre-call logging litellm_logging_obj.update_from_kwargs( @@ -138,7 +137,7 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable: def create_async_endpoint_function( sync_func: Callable, - endpoint_config: Dict, + endpoint_config: dict, ) -> Callable: """Create an async SDK function that wraps the sync function.""" @@ -146,17 +145,17 @@ def create_async_endpoint_function( async def async_endpoint_func( timeout: int = 600, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ): - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( + func: Final = partial( sync_func, timeout=timeout, custom_llm_provider=custom_llm_provider, @@ -166,9 +165,9 @@ def create_async_endpoint_function( **kwargs, ) - 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 @@ -188,14 +187,14 @@ def create_async_endpoint_function( return async_endpoint_func -def generate_container_endpoints() -> Dict[str, Callable]: +def generate_container_endpoints() -> dict[str, Callable]: """ Generate all container endpoint functions from the JSON config. Returns a dict mapping function names to their implementations. """ - config = _load_endpoints_config() - endpoints = {} + config: Final = _load_endpoints_config() + endpoints: Final = {} for endpoint_config in config["endpoints"]: # Create sync function @@ -209,33 +208,33 @@ def generate_container_endpoints() -> Dict[str, Callable]: return endpoints -def get_all_endpoint_names() -> List[str]: +def get_all_endpoint_names() -> list[str]: """Get all endpoint names (sync and async) from config.""" - config = _load_endpoints_config() - names = [] + config: Final = _load_endpoints_config() + names: Final = [] for endpoint in config["endpoints"]: names.append(endpoint["name"]) names.append(endpoint["async_name"]) return names -def get_async_endpoint_names() -> List[str]: +def get_async_endpoint_names() -> list[str]: """Get all async endpoint names for router registration.""" - config = _load_endpoints_config() + config: Final = _load_endpoints_config() return [endpoint["async_name"] for endpoint in config["endpoints"]] # Generate endpoints on module load -_generated_endpoints = generate_container_endpoints() +_generated_endpoints: Final = generate_container_endpoints() # Export generated functions dynamically -list_container_files = _generated_endpoints.get("list_container_files") -alist_container_files = _generated_endpoints.get("alist_container_files") -upload_container_file = _generated_endpoints.get("upload_container_file") -aupload_container_file = _generated_endpoints.get("aupload_container_file") -retrieve_container_file = _generated_endpoints.get("retrieve_container_file") -aretrieve_container_file = _generated_endpoints.get("aretrieve_container_file") -delete_container_file = _generated_endpoints.get("delete_container_file") -adelete_container_file = _generated_endpoints.get("adelete_container_file") -retrieve_container_file_content = _generated_endpoints.get("retrieve_container_file_content") -aretrieve_container_file_content = _generated_endpoints.get("aretrieve_container_file_content") +list_container_files: Final = _generated_endpoints.get("list_container_files") +alist_container_files: Final = _generated_endpoints.get("alist_container_files") +upload_container_file: Final = _generated_endpoints.get("upload_container_file") +aupload_container_file: Final = _generated_endpoints.get("aupload_container_file") +retrieve_container_file: Final = _generated_endpoints.get("retrieve_container_file") +aretrieve_container_file: Final = _generated_endpoints.get("aretrieve_container_file") +delete_container_file: Final = _generated_endpoints.get("delete_container_file") +adelete_container_file: Final = _generated_endpoints.get("adelete_container_file") +retrieve_container_file_content: Final = _generated_endpoints.get("retrieve_container_file_content") +aretrieve_container_file_content: Final = _generated_endpoints.get("aretrieve_container_file_content") diff --git a/litellm/containers/main.py b/litellm/containers/main.py index caf6c684844..c13f8bc75a6 100644 --- a/litellm/containers/main.py +++ b/litellm/containers/main.py @@ -1,8 +1,9 @@ import asyncio import contextvars import json +from collections.abc import Coroutine from functools import partial -from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, overload +from typing import Any, Final, Literal, overload import litellm from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT @@ -47,16 +48,16 @@ __all__ = [ @client async def acreate_container( name: str, - expires_after: Optional[Dict[str, Any]] = None, - file_ids: Optional[List[str]] = None, + expires_after: dict[str, Any] | None = None, + file_ids: list[str] | None = None, timeout=600, # default to 10 minutes # LiteLLM specific params, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ) -> ContainerObject: """Asynchronously calls the `create_container` function with the given arguments and keyword arguments. @@ -75,12 +76,12 @@ async def acreate_container( Returns: - `response` (ContainerObject): The created container object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( + func: Final = partial( create_container, name=name, expires_after=expires_after, @@ -93,9 +94,9 @@ async def acreate_container( **kwargs, ) - 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 @@ -119,12 +120,12 @@ async def acreate_container( @overload def create_container( name: str, - expires_after: Optional[Dict[str, Any]] = None, - file_ids: Optional[List[str]] = None, + expires_after: dict[str, Any] | None = None, + file_ids: list[str] | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, acreate_container: Literal[True], @@ -136,12 +137,12 @@ def create_container( @overload def create_container( name: str, - expires_after: Optional[Dict[str, Any]] = None, - file_ids: Optional[List[str]] = None, + expires_after: dict[str, Any] | None = None, + file_ids: list[str] | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, acreate_container: Literal[False] = False, @@ -155,23 +156,20 @@ def create_container( @client def create_container( name: str, - expires_after: Optional[Dict[str, Any]] = None, - file_ids: Optional[List[str]] = None, + expires_after: dict[str, Any] | None = None, + file_ids: list[str] | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, -) -> Union[ - ContainerObject, - Coroutine[Any, Any, ContainerObject], -]: +) -> ContainerObject | Coroutine[Any, Any, ContainerObject]: """Create a container using the OpenAI Container API. Currently supports OpenAI @@ -187,11 +185,11 @@ def create_container( print(response) ``` """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id") - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id") + _is_async: Final = kwargs.pop("async_call", False) is True # Check for mock response first mock_response = kwargs.get("mock_response") @@ -199,19 +197,19 @@ def create_container( if isinstance(mock_response, str): mock_response = json.loads(mock_response) - response = ContainerObject(**mock_response) + response: Final = ContainerObject(**mock_response) return response # get llm provider logic # Pass credential params explicitly since they're named args, not in kwargs - litellm_params = GenericLiteLLMParams( + litellm_params: Final = GenericLiteLLMParams( api_key=api_key, api_base=api_base, api_version=api_version, **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config( provider=litellm.LlmProviders(custom_llm_provider), ) @@ -220,12 +218,12 @@ def create_container( local_vars.update(kwargs) # Get ContainerCreateOptionalRequestParams with only valid parameters - container_create_optional_params: ContainerCreateOptionalRequestParams = ( + container_create_optional_params: Final[ContainerCreateOptionalRequestParams] = ( ContainerRequestUtils.get_requested_container_create_optional_param(local_vars) ) # Get optional parameters for the container API - container_create_request_params: Dict = ContainerRequestUtils.get_optional_params_container_create( + container_create_request_params: Final[dict] = ContainerRequestUtils.get_optional_params_container_create( container_provider_config=container_provider_config, container_create_optional_params=container_create_optional_params, ) @@ -280,16 +278,16 @@ def create_container( ##### Container List ####################### @client async def alist_containers( - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, # default to 10 minutes custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ) -> ContainerListResponse: """Asynchronously list containers. @@ -308,12 +306,12 @@ async def alist_containers( Returns: - `response` (ContainerListResponse): The list of containers """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( + func: Final = partial( list_containers, after=after, limit=limit, @@ -326,9 +324,9 @@ async def alist_containers( **kwargs, ) - 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 @@ -350,13 +348,13 @@ async def alist_containers( @overload def list_containers( - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, alist_containers: Literal[True], @@ -367,13 +365,13 @@ def list_containers( @overload def list_containers( - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, alist_containers: Literal[False] = False, @@ -386,33 +384,30 @@ def list_containers( @client def list_containers( - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, -) -> Union[ - ContainerListResponse, - Coroutine[Any, Any, ContainerListResponse], -]: +) -> ContainerListResponse | Coroutine[Any, Any, ContainerListResponse]: """List containers using the OpenAI Container API. Currently supports OpenAI """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id") - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id") + _is_async: Final = kwargs.pop("async_call", False) is True # Check for mock response first mock_response = kwargs.get("mock_response") @@ -420,19 +415,19 @@ def list_containers( if isinstance(mock_response, str): mock_response = json.loads(mock_response) - response = ContainerListResponse(**mock_response) + response: Final = ContainerListResponse(**mock_response) return response # get llm provider logic # Pass credential params explicitly since they're named args, not in kwargs - litellm_params = GenericLiteLLMParams( + litellm_params: Final = GenericLiteLLMParams( api_key=api_key, api_base=api_base, api_version=api_version, **kwargs, ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config( provider=litellm.LlmProviders(custom_llm_provider), ) @@ -440,7 +435,7 @@ def list_containers( raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}") # Get container list request parameters - container_list_optional_params: ContainerListOptionalRequestParams = ( + container_list_optional_params: Final[ContainerListOptionalRequestParams] = ( ContainerRequestUtils.get_requested_container_list_optional_param(local_vars) ) @@ -490,9 +485,9 @@ async def aretrieve_container( custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ) -> ContainerObject: """Asynchronously retrieve a container. @@ -509,12 +504,12 @@ async def aretrieve_container( Returns: - `response` (ContainerObject): The container object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( + func: Final = partial( retrieve_container, container_id=container_id, timeout=timeout, @@ -525,9 +520,9 @@ async def aretrieve_container( **kwargs, ) - 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 @@ -551,9 +546,9 @@ async def aretrieve_container( def retrieve_container( container_id: str, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, aretrieve_container: Literal[True], @@ -566,9 +561,9 @@ def retrieve_container( def retrieve_container( container_id: str, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, aretrieve_container: Literal[False] = False, @@ -583,30 +578,27 @@ def retrieve_container( def retrieve_container( container_id: str, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, -) -> Union[ - ContainerObject, - Coroutine[Any, Any, ContainerObject], -]: +) -> ContainerObject | Coroutine[Any, Any, ContainerObject]: """Retrieve a container using the OpenAI Container API. Currently supports OpenAI """ - local_vars = locals() + local_vars: Final = locals() try: resolved_custom_llm_provider: str = custom_llm_provider - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id") - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id") + _is_async: Final = kwargs.pop("async_call", False) is True # Check for mock response first mock_response = kwargs.get("mock_response") @@ -614,7 +606,7 @@ def retrieve_container( if isinstance(mock_response, str): mock_response = json.loads(mock_response) - response = ContainerObject(**mock_response) + response: Final = ContainerObject(**mock_response) return response # get llm provider logic @@ -633,10 +625,10 @@ def retrieve_container( litellm_params=litellm_params, ) # True when input was a LiteLLM-managed ID (any length); needed to re-encode output for routing affinity - was_encoded = original_container_id != container_id + was_encoded: Final = original_container_id != container_id # get provider config - container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config( provider=litellm.LlmProviders(resolved_custom_llm_provider), ) @@ -708,9 +700,9 @@ async def adelete_container( custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ) -> DeleteContainerResult: """Asynchronously delete a container. @@ -727,12 +719,12 @@ async def adelete_container( Returns: - `response` (DeleteContainerResult): The deletion result """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( + func: Final = partial( delete_container, container_id=container_id, timeout=timeout, @@ -743,9 +735,9 @@ async def adelete_container( **kwargs, ) - 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 @@ -769,9 +761,9 @@ async def adelete_container( def delete_container( container_id: str, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, adelete_container: Literal[True], @@ -784,9 +776,9 @@ def delete_container( def delete_container( container_id: str, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, adelete_container: Literal[False] = False, @@ -801,30 +793,27 @@ def delete_container( def delete_container( container_id: str, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, -) -> Union[ - DeleteContainerResult, - Coroutine[Any, Any, DeleteContainerResult], -]: +) -> DeleteContainerResult | Coroutine[Any, Any, DeleteContainerResult]: """Delete a container using the OpenAI Container API. Currently supports OpenAI """ - local_vars = locals() + local_vars: Final = locals() try: resolved_custom_llm_provider: str = custom_llm_provider - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id") - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id") + _is_async: Final = kwargs.pop("async_call", False) is True # Check for mock response first mock_response = kwargs.get("mock_response") @@ -832,7 +821,7 @@ def delete_container( if isinstance(mock_response, str): mock_response = json.loads(mock_response) - response = DeleteContainerResult(**mock_response) + response: Final = DeleteContainerResult(**mock_response) return response # get llm provider logic @@ -851,10 +840,10 @@ def delete_container( litellm_params=litellm_params, ) # True when input was a LiteLLM-managed ID (any length); needed to re-encode output for routing affinity - was_encoded = original_container_id != container_id + was_encoded: Final = original_container_id != container_id # get provider config - container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config( provider=litellm.LlmProviders(resolved_custom_llm_provider), ) @@ -922,14 +911,14 @@ def delete_container( @client async def alist_container_files( container_id: str, - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, # default to 10 minutes custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ) -> ContainerFileListResponse: """Asynchronously list files in a container. @@ -949,12 +938,12 @@ async def alist_container_files( Returns: - `response` (ContainerFileListResponse): The list of container files """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( + func: Final = partial( list_container_files, container_id=container_id, after=after, @@ -968,9 +957,9 @@ async def alist_container_files( **kwargs, ) - 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 @@ -993,13 +982,13 @@ async def alist_container_files( @overload def list_container_files( container_id: str, - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, alist_container_files: Literal[True], @@ -1011,13 +1000,13 @@ def list_container_files( @overload def list_container_files( container_id: str, - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, alist_container_files: Literal[False] = False, @@ -1031,32 +1020,29 @@ def list_container_files( @client def list_container_files( container_id: str, - after: Optional[str] = None, - limit: Optional[int] = None, - order: Optional[str] = None, + after: str | None = None, + limit: int | None = None, + order: str | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, -) -> Union[ - ContainerFileListResponse, - Coroutine[Any, Any, ContainerFileListResponse], -]: +) -> ContainerFileListResponse | Coroutine[Any, Any, ContainerFileListResponse]: """List files in a container using the OpenAI Container API. Currently supports OpenAI """ - local_vars = locals() + local_vars: Final = locals() try: resolved_custom_llm_provider: str = custom_llm_provider - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id") - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id") + _is_async: Final = kwargs.pop("async_call", False) is True # Check for mock response first mock_response = kwargs.get("mock_response") @@ -1064,7 +1050,7 @@ def list_container_files( if isinstance(mock_response, str): mock_response = json.loads(mock_response) - response = ContainerFileListResponse(**mock_response) + response: Final = ContainerFileListResponse(**mock_response) return response # get llm provider logic @@ -1084,7 +1070,7 @@ def list_container_files( ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config( provider=litellm.LlmProviders(resolved_custom_llm_provider), ) @@ -1141,9 +1127,9 @@ async def aupload_container_file( file: FileTypes, timeout=600, # default to 10 minutes custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, ) -> ContainerFileObject: """Asynchronously upload a file to a container. @@ -1182,12 +1168,12 @@ async def aupload_container_file( print(response) ``` """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True - func = partial( + func: Final = partial( upload_container_file, container_id=container_id, file=file, @@ -1199,9 +1185,9 @@ async def aupload_container_file( **kwargs, ) - 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 @@ -1226,9 +1212,9 @@ def upload_container_file( container_id: str, file: FileTypes, timeout=600, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, aupload_container_file: Literal[True], @@ -1242,9 +1228,9 @@ def upload_container_file( container_id: str, file: FileTypes, timeout=600, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", *, aupload_container_file: Literal[False] = False, @@ -1260,18 +1246,15 @@ def upload_container_file( container_id: str, file: FileTypes, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai", - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, **kwargs, -) -> Union[ - ContainerFileObject, - Coroutine[Any, Any, ContainerFileObject], -]: +) -> ContainerFileObject | Coroutine[Any, Any, ContainerFileObject]: """Upload a file to a container using the OpenAI Container API. This endpoint allows uploading files directly to a container session, @@ -1305,12 +1288,12 @@ def upload_container_file( """ from litellm.llms.custom_httpx.container_handler import generic_container_handler - local_vars = locals() + local_vars: Final = locals() try: resolved_custom_llm_provider: str = custom_llm_provider - litellm_logging_obj: LiteLLMLoggingObj = kwargs.pop("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id") - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.pop("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id") + _is_async: Final = kwargs.pop("async_call", False) is True # Check for mock response first mock_response = kwargs.get("mock_response") @@ -1318,7 +1301,7 @@ def upload_container_file( if isinstance(mock_response, str): mock_response = json.loads(mock_response) - response = ContainerFileObject(**mock_response) + response: Final = ContainerFileObject(**mock_response) return response # get llm provider logic @@ -1338,7 +1321,7 @@ def upload_container_file( ) # get provider config - container_provider_config: Optional[BaseContainerConfig] = ProviderConfigManager.get_provider_container_config( + container_provider_config: BaseContainerConfig | None = ProviderConfigManager.get_provider_container_config( provider=litellm.LlmProviders(resolved_custom_llm_provider), ) diff --git a/litellm/containers/utils.py b/litellm/containers/utils.py index 2b115c6b3c4..f07820602bf 100644 --- a/litellm/containers/utils.py +++ b/litellm/containers/utils.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional, TypeVar +from typing import Any, Final, TypeVar from litellm.llms.base_llm.containers.transformation import BaseContainerConfig from litellm.responses.utils import ResponsesAPIRequestUtils @@ -19,14 +19,14 @@ def decode_managed_container_id_for_request( Returns: (original_container_id, resolved_provider, updated_litellm_params) """ - decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) - original_container_id = decoded.get("response_id", container_id) + decoded: Final = ResponsesAPIRequestUtils._decode_container_id(container_id) + original_container_id: Final = decoded.get("response_id", container_id) - decoded_provider = decoded.get("custom_llm_provider") + decoded_provider: Final = decoded.get("custom_llm_provider") if decoded_provider and custom_llm_provider == "openai": custom_llm_provider = decoded_provider - decoded_model_id = decoded.get("model_id") + decoded_model_id: Final = decoded.get("model_id") if decoded_model_id and not litellm_params.get("model_id"): litellm_params["model_id"] = decoded_model_id @@ -42,9 +42,9 @@ class ContainerRequestUtils: passed_params: dict, ) -> ContainerCreateOptionalRequestParams: """Extract only valid container creation parameters from the passed parameters.""" - container_create_optional_params = ContainerCreateOptionalRequestParams() + container_create_optional_params: Final = ContainerCreateOptionalRequestParams() - valid_params = [ + valid_params: Final = [ "expires_after", "file_ids", "extra_headers", @@ -61,12 +61,12 @@ class ContainerRequestUtils: def get_optional_params_container_create( container_provider_config: BaseContainerConfig, container_create_optional_params: ContainerCreateOptionalRequestParams, - ) -> Dict: + ) -> dict: """Get the optional parameters for container creation.""" - supported_params = container_provider_config.get_supported_openai_params() + supported_params: Final = container_provider_config.get_supported_openai_params() # Filter out unsupported parameters - filtered_params = {k: v for k, v in container_create_optional_params.items() if k in supported_params} + filtered_params: Final = {k: v for k, v in container_create_optional_params.items() if k in supported_params} return container_provider_config.map_openai_params( container_create_optional_params=filtered_params, # type: ignore @@ -78,9 +78,9 @@ class ContainerRequestUtils: passed_params: dict, ) -> ContainerListOptionalRequestParams: """Extract only valid container list parameters from the passed parameters.""" - container_list_optional_params = ContainerListOptionalRequestParams() + container_list_optional_params: Final = ContainerListOptionalRequestParams() - valid_params = [ + valid_params: Final = [ "after", "limit", "order", @@ -97,9 +97,9 @@ class ContainerRequestUtils: @staticmethod def encode_container_id_in_response( response_obj: T, - custom_llm_provider: Optional[str], - litellm_metadata: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, + custom_llm_provider: str | None, + litellm_metadata: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, ) -> T: """ Encode container_id in response object with provider/model metadata for routing. @@ -124,7 +124,7 @@ class ContainerRequestUtils: """ # Extract model_id from litellm_metadata litellm_metadata = litellm_metadata or {} - model_info: Dict[str, Any] = litellm_metadata.get("model_info", {}) or {} + model_info: Final[dict[str, Any]] = litellm_metadata.get("model_info", {}) or {} model_id = model_info.get("id") # Check if we should encode based on routing metadata @@ -139,7 +139,7 @@ class ContainerRequestUtils: should_encode = True # Extract model_id from target_model_names if not already set if model_id is None: - target_models = extra_body["target_model_names"] + target_models: Final = extra_body["target_model_names"] # Use first model as model_id for encoding if isinstance(target_models, str): model_id = target_models.split(",")[0].strip() @@ -148,7 +148,7 @@ class ContainerRequestUtils: # Only encode if we have routing metadata if should_encode and response_obj and hasattr(response_obj, "id"): - encoded_id = ResponsesAPIRequestUtils._build_container_id( + encoded_id: Final = ResponsesAPIRequestUtils._build_container_id( custom_llm_provider=custom_llm_provider, model_id=model_id, container_id=response_obj.id, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 96aed20529f..025d400509b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -3,7 +3,7 @@ import logging import time from functools import lru_cache -from typing import TYPE_CHECKING, Any, List, Literal, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response from pydantic import BaseModel @@ -29,8 +29,8 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _parse_prompt_tokens_details, calculate_cost_component, generic_cost_per_token, - get_token_type_cost_breakdown, get_billable_input_tokens, + get_token_type_cost_breakdown, select_cost_metric_for_model, ) from litellm.llms.anthropic.cost_calculation import ( @@ -52,9 +52,6 @@ from litellm.llms.databricks.cost_calculator import ( from litellm.llms.deepseek.cost_calculator import ( cost_per_token as deepseek_cost_per_token, ) -from litellm.llms.tencent.cost_calculator import ( - cost_per_token as tencent_cost_per_token, -) from litellm.llms.fireworks_ai.cost_calculator import ( cost_per_token as fireworks_ai_cost_per_token, ) @@ -64,12 +61,19 @@ from litellm.llms.lemonade.cost_calculator import ( ) from litellm.llms.openai.cost_calculation import ( _video_output_cost_per_second, +) +from litellm.llms.openai.cost_calculation import ( cost_per_second as openai_cost_per_second, +) +from litellm.llms.openai.cost_calculation import ( cost_per_token as openai_cost_per_token, ) from litellm.llms.perplexity.cost_calculator import ( cost_per_token as perplexity_cost_per_token, ) +from litellm.llms.tencent.cost_calculator import ( + cost_per_token as tencent_cost_per_token, +) from litellm.llms.together_ai.cost_calculator import get_model_params_and_category from litellm.llms.vertex_ai.cost_calculator import ( cost_per_character as google_cost_per_character, @@ -127,14 +131,14 @@ else: LitellmLoggingObject = Any # Pre-resolved CallTypes enum values for fast membership checks -_A2A_CALL_TYPES = frozenset( +_A2A_CALL_TYPES: Final = frozenset( { CallTypes.asend_message.value, CallTypes.send_message.value, } ) -_VIDEO_CALL_TYPES = frozenset( +_VIDEO_CALL_TYPES: Final = frozenset( { CallTypes.create_video.value, CallTypes.acreate_video.value, @@ -145,48 +149,48 @@ _VIDEO_CALL_TYPES = frozenset( } ) -_SPEECH_CALL_TYPES = frozenset( +_SPEECH_CALL_TYPES: Final = frozenset( { CallTypes.speech.value, CallTypes.aspeech.value, } ) -_TRANSCRIPTION_CALL_TYPES = frozenset( +_TRANSCRIPTION_CALL_TYPES: Final = frozenset( { CallTypes.atranscription.value, CallTypes.transcription.value, } ) -_RERANK_CALL_TYPES = frozenset( +_RERANK_CALL_TYPES: Final = frozenset( { CallTypes.rerank.value, CallTypes.arerank.value, } ) -_SEARCH_CALL_TYPES = frozenset( +_SEARCH_CALL_TYPES: Final = frozenset( { CallTypes.search.value, CallTypes.asearch.value, } ) -_AREALTIME_CALL_TYPE = CallTypes.arealtime.value -_MCP_CALL_TYPE = CallTypes.call_mcp_tool.value +_AREALTIME_CALL_TYPE: Final = CallTypes.arealtime.value +_MCP_CALL_TYPE: Final = CallTypes.call_mcp_tool.value def _cost_per_token_custom_pricing_helper( prompt_tokens: float = 0, completion_tokens: float = 0, - response_time_ms: Optional[float] = 0.0, + response_time_ms: float | None = 0.0, cached_tokens: float = 0, cache_creation_tokens: float = 0, ### CUSTOM PRICING ### - custom_cost_per_token: Optional[CostPerToken] = None, - custom_cost_per_second: Optional[float] = None, -) -> Optional[Tuple[float, float]]: + custom_cost_per_token: CostPerToken | None = None, + custom_cost_per_second: float | None = None, +) -> tuple[float, float] | None: """Internal helper function for calculating cost, if custom pricing given. prompt_tokens is assumed to include both cached_tokens and cache_creation_tokens @@ -197,24 +201,24 @@ def _cost_per_token_custom_pricing_helper( return None if custom_cost_per_token is not None: - input_cost_per_token = custom_cost_per_token["input_cost_per_token"] - output_cost_per_token = custom_cost_per_token["output_cost_per_token"] + input_cost_per_token: Final = custom_cost_per_token["input_cost_per_token"] + output_cost_per_token: Final = custom_cost_per_token["output_cost_per_token"] - cache_read_input_token_cost = custom_cost_per_token.get( + cache_read_input_token_cost: Final = custom_cost_per_token.get( "cache_read_input_token_cost", input_cost_per_token, ) - cache_creation_input_token_cost = custom_cost_per_token.get( + cache_creation_input_token_cost: Final = custom_cost_per_token.get( "cache_creation_input_token_cost", input_cost_per_token, ) - regular_prompt_tokens = max( + regular_prompt_tokens: Final = max( prompt_tokens - cached_tokens - cache_creation_tokens, 0, ) - input_cost = ( + input_cost: Final = ( regular_prompt_tokens * input_cost_per_token + cached_tokens * cache_read_input_token_cost + cache_creation_tokens * cache_creation_input_token_cost @@ -230,10 +234,10 @@ def _cost_per_token_custom_pricing_helper( def _get_additional_costs( model: str, - custom_llm_provider: Optional[str], + custom_llm_provider: str | None, prompt_tokens: int, completion_tokens: int, -) -> Optional[dict]: +) -> dict | None: """ Calculate additional costs beyond standard token costs. @@ -269,24 +273,24 @@ def _get_additional_costs( completion_tokens=completion_tokens, ) except Exception as e: - verbose_logger.debug(f"Error calculating additional costs: {e}") + verbose_logger.debug("Error calculating additional costs: %s", e) return None def _transcription_usage_has_token_details( - usage_block: Optional[Usage], + usage_block: Usage | None, ) -> bool: if usage_block is None: return False - prompt_tokens_val = getattr(usage_block, "prompt_tokens", 0) or 0 - completion_tokens_val = getattr(usage_block, "completion_tokens", 0) or 0 - prompt_details = getattr(usage_block, "prompt_tokens_details", None) + prompt_tokens_val: Final = getattr(usage_block, "prompt_tokens", 0) or 0 + completion_tokens_val: Final = getattr(usage_block, "completion_tokens", 0) or 0 + prompt_details: Final = getattr(usage_block, "prompt_tokens_details", None) if prompt_details is not None: - audio_token_count = getattr(prompt_details, "audio_tokens", 0) or 0 - text_token_count = getattr(prompt_details, "text_tokens", 0) or 0 + audio_token_count: Final = getattr(prompt_details, "audio_tokens", 0) or 0 + text_token_count: Final = getattr(prompt_details, "text_tokens", 0) or 0 if audio_token_count > 0 or text_token_count > 0: return True @@ -297,35 +301,35 @@ def cost_per_token( model: str = "", prompt_tokens: int = 0, completion_tokens: int = 0, - response_time_ms: Optional[float] = 0.0, - custom_llm_provider: Optional[str] = None, + response_time_ms: float | None = 0.0, + custom_llm_provider: str | None = None, region_name=None, ### CHARACTER PRICING ### - prompt_characters: Optional[int] = None, - completion_characters: Optional[int] = None, + prompt_characters: int | None = None, + completion_characters: int | None = None, ### PROMPT CACHING PRICING ### - used for anthropic - cache_creation_input_tokens: Optional[int] = 0, - cache_read_input_tokens: Optional[int] = 0, + cache_creation_input_tokens: int | None = 0, + cache_read_input_tokens: int | None = 0, ### CUSTOM PRICING ### - custom_cost_per_token: Optional[CostPerToken] = None, - custom_cost_per_second: Optional[float] = None, + custom_cost_per_token: CostPerToken | None = None, + custom_cost_per_second: float | None = None, ### NUMBER OF QUERIES ### - number_of_queries: Optional[int] = None, + number_of_queries: int | None = None, ### USAGE OBJECT ### - usage_object: Optional[Usage] = None, # just read the usage object if provided + usage_object: Usage | None = None, # just read the usage object if provided ### BILLED UNITS ### - rerank_billed_units: Optional[RerankBilledUnits] = None, + rerank_billed_units: RerankBilledUnits | None = None, ### CALL TYPE ### call_type: CallTypesLiteral = "completion", audio_transcription_file_duration: float = 0.0, # for audio transcription calls - the file time in seconds ### SERVICE TIER ### - service_tier: Optional[str] = None, # for OpenAI service tier pricing + service_tier: str | None = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### - data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") - response: Optional[Any] = None, + data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + response: Any | None = None, ### REQUEST MODEL ### - request_model: Optional[str] = None, # original request model for router detection -) -> Tuple[float, float]: # type: ignore + request_model: str | None = None, # original request model for router detection +) -> tuple[float, float]: # type: ignore """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -371,7 +375,7 @@ def cost_per_token( _is_anthropic_style = False if usage_object is not None: - _pt_details = getattr(usage_object, "prompt_tokens_details", None) + _pt_details: Final = getattr(usage_object, "prompt_tokens_details", None) if _pt_details is not None: _cache_read_tokens = float(getattr(_pt_details, "cached_tokens", 0) or 0) # OpenAI-compatible providers report cache-write tokens under @@ -381,8 +385,8 @@ def cost_per_token( getattr(_pt_details, "cache_write_tokens", 0) or getattr(_pt_details, "cache_creation_tokens", 0) or 0 ) - _anthropic_read = getattr(usage_object, "cache_read_input_tokens", None) - _anthropic_create = getattr(usage_object, "cache_creation_input_tokens", None) + _anthropic_read: Final = getattr(usage_object, "cache_read_input_tokens", None) + _anthropic_create: Final = getattr(usage_object, "cache_creation_input_tokens", None) if _anthropic_read is not None or _anthropic_create is not None: _is_anthropic_style = True if _anthropic_read is not None: @@ -403,7 +407,7 @@ def cost_per_token( if _is_anthropic_style: _normalized_prompt_tokens += _cache_read_tokens + _cache_creation_tokens - response_cost = _cost_per_token_custom_pricing_helper( + response_cost: Final = _cost_per_token_custom_pricing_helper( prompt_tokens=_normalized_prompt_tokens, completion_tokens=completion_tokens, response_time_ms=response_time_ms, @@ -419,23 +423,23 @@ def cost_per_token( # given prompt_tokens_cost_usd_dollar: float = 0 completion_tokens_cost_usd_dollar: float = 0 - model_cost_ref = litellm.model_cost + model_cost_ref: Final = litellm.model_cost # Only callers that explicitly pass `custom_llm_provider` get the # dedup/prefix-join treatment. When provider is omitted, preserve legacy # behavior: `model_with_provider` stays equal to the raw `model` string # (provider is detected below for downstream use only). - caller_supplied_provider = custom_llm_provider is not None + caller_supplied_provider: Final = custom_llm_provider is not None # `model` is normally a string, but callers that mock the transport can pass # non-string objects. Only run the string-based dedup/prefix-join when it is # actually a string — e.g. a MagicMock's `.startswith()` is always truthy and # its slices return new mocks, which would spin the dedup loop forever. - model_is_str = isinstance(model, str) + model_is_str: Final = isinstance(model, str) # Router/proxy deployments may repeat the provider segment (e.g. model_name # "openai/openai/gpt-5.5"). Strip duplicated `{provider}/` chains before joining. if caller_supplied_provider and model_is_str: - _dup_prefix = f"{custom_llm_provider}/" + _dup_prefix: Final = f"{custom_llm_provider}/" while model.startswith(_dup_prefix): _remainder = model[len(_dup_prefix) :] if _remainder.startswith(_dup_prefix): @@ -445,13 +449,13 @@ def cost_per_token( model_with_provider = model if caller_supplied_provider: - _prov_prefix = f"{custom_llm_provider}/" + _prov_prefix: Final = f"{custom_llm_provider}/" if model_is_str and model.startswith(_prov_prefix): model_with_provider = model else: model_with_provider = f"{custom_llm_provider}/{model}" if region_name is not None: - model_with_provider_and_region = f"{custom_llm_provider}/{region_name}/{model}" + model_with_provider_and_region: Final = f"{custom_llm_provider}/{region_name}/{model}" if model_with_provider_and_region in model_cost_ref: # use region based pricing, if it's available model_with_provider = model_with_provider_and_region else: @@ -460,7 +464,7 @@ def cost_per_token( assert custom_llm_provider is not None # caller-supplied or get_llm_provider model_without_prefix = model - model_parts = model.split("/", 1) + model_parts: Final = model.split("/", 1) if len(model_parts) > 1: model_without_prefix = model_parts[1] else: @@ -483,18 +487,13 @@ def cost_per_token( # see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models if call_type == "speech" or call_type == "aspeech": speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) - cost_metric = select_cost_metric_for_model(speech_model_info) + cost_metric: Final = select_cost_metric_for_model(speech_model_info) prompt_cost: float = 0.0 completion_cost: float = 0.0 if cost_metric == "cost_per_character": if prompt_characters is None: raise ValueError( - "prompt_characters must be provided for tts calls. prompt_characters={}, model={}, custom_llm_provider={}, call_type={}".format( - prompt_characters, - model, - custom_llm_provider, - call_type, - ) + f"prompt_characters must be provided for tts calls. prompt_characters={prompt_characters}, model={model}, custom_llm_provider={custom_llm_provider}, call_type={call_type}" ) _prompt_cost, _completion_cost = _generic_cost_per_character( model=model_without_prefix, @@ -506,14 +505,7 @@ def cost_per_token( ) if _prompt_cost is None or _completion_cost is None: raise ValueError( - "cost for tts call is None. prompt_cost={}, completion_cost={}, model={}, custom_llm_provider={}, prompt_characters={}, completion_characters={}".format( - _prompt_cost, - _completion_cost, - model_without_prefix, - custom_llm_provider, - prompt_characters, - completion_characters, - ) + f"cost for tts call is None. prompt_cost={_prompt_cost}, completion_cost={_completion_cost}, model={model_without_prefix}, custom_llm_provider={custom_llm_provider}, prompt_characters={prompt_characters}, completion_characters={completion_characters}" ) prompt_cost = _prompt_cost completion_cost = _completion_cost @@ -582,7 +574,7 @@ def cost_per_token( optional_params=(response._hidden_params if response and hasattr(response, "_hidden_params") else None), ) elif custom_llm_provider == "vertex_ai": - cost_router = google_cost_router( + cost_router: Final = google_cost_router( model=model_without_prefix, custom_llm_provider=custom_llm_provider, call_type=call_type, @@ -651,7 +643,7 @@ def cost_per_token( service_tier=service_tier, ) else: - model_info = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) + model_info: Final = _cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider) if (model_info.get("input_cost_per_token") or 0.0) > 0 or (model_info.get("output_cost_per_token") or 0.0) > 0: return generic_cost_per_token( @@ -662,7 +654,7 @@ def cost_per_token( data_residency=data_residency, ) - input_cost_per_second = model_info.get("input_cost_per_second") + input_cost_per_second: Final = model_info.get("input_cost_per_second") if input_cost_per_second is not None and response_time_ms is not None: verbose_logger.debug( "For model=%s - input_cost_per_second: %s; response time: %s", @@ -673,7 +665,7 @@ def cost_per_token( ## COST PER SECOND ## prompt_tokens_cost_usd_dollar = input_cost_per_second * response_time_ms / 1000 - output_cost_per_second = model_info.get("output_cost_per_second") + output_cost_per_second: Final = model_info.get("output_cost_per_second") if output_cost_per_second is not None and response_time_ms is not None: verbose_logger.debug( "For model=%s - output_cost_per_second: %s; response time: %s", @@ -696,12 +688,12 @@ def cost_per_token( def get_replicate_completion_pricing(completion_response: dict, total_time=0.0): # see https://replicate.com/pricing # for all litellm currently supported LLMs, almost all requests go to a100_80gb - a100_80gb_price_per_second_public = ( + a100_80gb_price_per_second_public: Final = ( DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND # assume all calls sent to A100 80GB for now ) if total_time == 0.0: # total time is in ms - start_time = completion_response.get("created", time.time()) - end_time = getattr(completion_response, "ended", time.time()) + start_time: Final = completion_response.get("created", time.time()) + end_time: Final = getattr(completion_response, "ended", time.time()) total_time = end_time - start_time return a100_80gb_price_per_second_public * total_time / 1000 @@ -712,9 +704,9 @@ def has_hidden_params(obj: Any) -> bool: def _get_provider_for_cost_calc( - model: Optional[str], - custom_llm_provider: Optional[str] = None, -) -> Optional[str]: + model: str | None, + custom_llm_provider: str | None = None, +) -> str | None: if custom_llm_provider is not None: return custom_llm_provider if model is None: @@ -723,7 +715,7 @@ def _get_provider_for_cost_calc( _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) except Exception as e: verbose_logger.debug( - f"litellm.cost_calculator.py::_get_provider_for_cost_calc() - Error inferring custom_llm_provider - {str(e)}" + "litellm.cost_calculator.py::_get_provider_for_cost_calc() - Error inferring custom_llm_provider - %s", e ) return None @@ -731,13 +723,13 @@ def _get_provider_for_cost_calc( def _select_model_name_for_cost_calc( - model: Optional[str], - completion_response: Optional[Any], - base_model: Optional[str] = None, - custom_pricing: Optional[bool] = None, - custom_llm_provider: Optional[str] = None, - router_model_id: Optional[str] = None, -) -> Optional[str]: + model: str | None, + completion_response: Any | None, + base_model: str | None = None, + custom_pricing: bool | None = None, + custom_llm_provider: str | None = None, + router_model_id: str | None = None, +) -> str | None: """ 1. If custom pricing is true, return received model name 2. If base_model is set (e.g. for azure models), return that @@ -745,21 +737,21 @@ def _select_model_name_for_cost_calc( 4. Check if model is passed in return that """ - return_model: Optional[str] = None - region_name: Optional[str] = None + return_model: str | None = None + region_name: str | None = None custom_llm_provider = _get_provider_for_cost_calc(model=model, custom_llm_provider=custom_llm_provider) - completion_response_model: Optional[str] = None + completion_response_model: str | None = None if completion_response is not None: if isinstance(completion_response, BaseModel): completion_response_model = getattr(completion_response, "model", None) elif isinstance(completion_response, dict): completion_response_model = completion_response.get("model", None) - hidden_params: Optional[dict] = getattr(completion_response, "_hidden_params", None) + hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None) if custom_pricing is True: if router_model_id is not None and router_model_id in litellm.model_cost: - entry = litellm.model_cost[router_model_id] + entry: Final = litellm.model_cost[router_model_id] if ( entry.get("input_cost_per_token") is not None or entry.get("input_cost_per_second") is not None @@ -804,11 +796,11 @@ def _model_contains_known_llm_provider(model: str) -> bool: """ Check if the model contains a known llm provider """ - _provider_prefix = model.split("/")[0] + _provider_prefix: Final = model.split("/")[0] return _provider_prefix in LlmProvidersSet -def _get_response_model(completion_response: Any) -> Optional[str]: +def _get_response_model(completion_response: Any) -> str | None: """ Extract the model name from a completion response object. @@ -826,7 +818,7 @@ def _get_response_model(completion_response: Any) -> Optional[str]: return None -_GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: dict = { +_GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = { # ON_DEMAND_PRIORITY maps to "priority" — selects input_cost_per_token_priority, etc. "ON_DEMAND_PRIORITY": "priority", # FLEX / BATCH maps to "flex" — selects input_cost_per_token_flex, etc. @@ -837,7 +829,7 @@ _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: dict = { } -def _map_traffic_type_to_service_tier(traffic_type: Optional[str]) -> Optional[str]: +def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None: """ Map a Gemini usageMetadata.trafficType value to a LiteLLM service_tier string. @@ -852,7 +844,7 @@ def _map_traffic_type_to_service_tier(traffic_type: Optional[str]) -> Optional[s """ if traffic_type is None: return None - service_tier = _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER.get(str(traffic_type).upper()) + service_tier: Final = _GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER.get(str(traffic_type).upper()) return service_tier @@ -872,9 +864,9 @@ def _normalize_service_tier(service_tier: object) -> str | None: def _get_usage_object( completion_response: Any, -) -> Optional[Usage]: - usage_obj = cast( - Union[Usage, ResponseAPIUsage, dict, BaseModel], +) -> Usage | None: + usage_obj: Final = cast( + Usage | ResponseAPIUsage | dict | BaseModel, ( completion_response.get("usage") if isinstance(completion_response, dict) @@ -886,6 +878,8 @@ def _get_usage_object( return None if isinstance(usage_obj, Usage): return usage_obj + elif isinstance(usage_obj, dict) and litellm.AnthropicConfig.is_anthropic_usage_object(usage_obj): + return litellm.AnthropicConfig().calculate_usage(usage_object=usage_obj, reasoning_content=None) elif ( usage_obj is not None and (isinstance(usage_obj, dict) or isinstance(usage_obj, ResponseAPIUsage)) @@ -895,7 +889,7 @@ def _get_usage_object( elif TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj): return TranscriptionUsageObjectTransformation.transform_transcription_usage_object( cast( - Union[TranscriptionUsageDurationObject, TranscriptionUsageTokensObject], + TranscriptionUsageDurationObject | TranscriptionUsageTokensObject, usage_obj, ) ) @@ -904,7 +898,7 @@ def _get_usage_object( elif isinstance(usage_obj, BaseModel): return Usage(**usage_obj.model_dump()) else: - verbose_logger.debug(f"Unknown usage object type: {type(usage_obj)}, usage_obj: {usage_obj}") + verbose_logger.debug("Unknown usage object type: %s, usage_obj: %s", type(usage_obj), usage_obj) return None @@ -917,7 +911,7 @@ def _is_known_usage_objects(usage_obj): ) -def _infer_call_type(call_type: Optional[CallTypesLiteral], completion_response: Any) -> Optional[CallTypesLiteral]: +def _infer_call_type(call_type: CallTypesLiteral | None, completion_response: Any) -> CallTypesLiteral | None: if call_type is not None: return call_type @@ -946,8 +940,8 @@ def _infer_call_type(call_type: Optional[CallTypesLiteral], completion_response: def _apply_cost_discount( base_cost: float, - custom_llm_provider: Optional[str], -) -> Tuple[float, float, float]: + custom_llm_provider: str | None, +) -> tuple[float, float, float]: """ Apply provider-specific cost discount from module-level config. @@ -958,14 +952,14 @@ def _apply_cost_discount( Returns: Tuple of (final_cost, discount_percent, discount_amount) """ - original_cost = base_cost + original_cost: Final = base_cost discount_percent = 0.0 discount_amount = 0.0 if custom_llm_provider and custom_llm_provider in litellm.cost_discount_config: discount_percent = litellm.cost_discount_config[custom_llm_provider] discount_amount = original_cost * discount_percent - final_cost = original_cost - discount_amount + final_cost: Final = original_cost - discount_amount if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug( @@ -980,8 +974,8 @@ def _apply_cost_discount( def _apply_cost_margin( base_cost: float, - custom_llm_provider: Optional[str], -) -> Tuple[float, float, float, float]: + custom_llm_provider: str | None, +) -> tuple[float, float, float, float]: """ Apply provider-specific or global cost margin from module-level config. @@ -992,7 +986,7 @@ def _apply_cost_margin( Returns: Tuple of (final_cost, margin_percent, margin_fixed_amount, margin_total_amount) """ - original_cost = base_cost + original_cost: Final = base_cost margin_percent = 0.0 margin_fixed_amount = 0.0 margin_total_amount = 0.0 @@ -1002,16 +996,17 @@ def _apply_cost_margin( if custom_llm_provider and custom_llm_provider in litellm.cost_margin_config: margin_config = litellm.cost_margin_config[custom_llm_provider] if verbose_logger.isEnabledFor(logging.DEBUG): - verbose_logger.debug(f"Found provider-specific margin config for {custom_llm_provider}: {margin_config}") + verbose_logger.debug("Found provider-specific margin config for %s: %s", custom_llm_provider, margin_config) elif "global" in litellm.cost_margin_config: margin_config = litellm.cost_margin_config["global"] if verbose_logger.isEnabledFor(logging.DEBUG): - verbose_logger.debug(f"Using global margin config: {margin_config}") + verbose_logger.debug("Using global margin config: %s", margin_config) else: if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug( - f"No margin config found. Provider: {custom_llm_provider}, " - f"Available configs: {list(litellm.cost_margin_config.keys())}" + "No margin config found. Provider: %s, Available configs: %s", + custom_llm_provider, + list(litellm.cost_margin_config.keys()), ) if margin_config is not None: @@ -1029,7 +1024,7 @@ def _apply_cost_margin( margin_fixed_amount = float(margin_config["fixed_amount"]) margin_total_amount += margin_fixed_amount - final_cost = original_cost + margin_total_amount + final_cost: Final = original_cost + margin_total_amount if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug( @@ -1044,21 +1039,23 @@ def _apply_cost_margin( def _store_cost_breakdown_in_logging_obj( - litellm_logging_obj: Optional[LitellmLoggingObject], + litellm_logging_obj: LitellmLoggingObject | None, prompt_tokens_cost_usd_dollar: float, completion_tokens_cost_usd_dollar: float, cost_for_built_in_tools_cost_usd_dollar: float, total_cost_usd_dollar: float, - additional_costs: Optional[dict] = None, - original_cost: Optional[float] = None, - discount_percent: Optional[float] = None, - discount_amount: Optional[float] = None, - margin_percent: Optional[float] = None, - margin_fixed_amount: Optional[float] = None, - margin_total_amount: Optional[float] = None, - cache_read_cost: Optional[float] = None, - cache_creation_cost: Optional[float] = None, - reasoning_cost: Optional[float] = None, + additional_costs: dict | None = None, + original_cost: float | None = None, + discount_percent: float | None = None, + discount_amount: float | None = None, + margin_percent: float | None = None, + margin_fixed_amount: float | None = None, + margin_total_amount: float | None = None, + cache_read_cost: float | None = None, + cache_creation_cost: float | None = None, + reasoning_cost: float | None = None, + service_tier: str | None = None, + data_residency: str | None = None, ) -> None: """ Helper function to store cost breakdown in the logging object. @@ -1076,6 +1073,8 @@ def _store_cost_breakdown_in_logging_obj( margin_percent: Margin percentage applied (0.10 = 10%) margin_fixed_amount: Fixed margin amount in USD margin_total_amount: Total margin added in USD + service_tier: Tier the costs above were priced on, already resolved + data_residency: Region uplift the costs above were priced on, already resolved """ if litellm_logging_obj is None: return @@ -1097,43 +1096,44 @@ def _store_cost_breakdown_in_logging_obj( cache_read_cost=cache_read_cost, cache_creation_cost=cache_creation_cost, reasoning_cost=reasoning_cost, + service_tier=service_tier, + data_residency=data_residency, ) except Exception as breakdown_error: - verbose_logger.debug(f"Error storing cost breakdown: {str(breakdown_error)}") + verbose_logger.debug("Error storing cost breakdown: %s", breakdown_error) # Don't fail the main cost calculation if breakdown storage fails - pass def completion_cost( completion_response=None, - model: Optional[str] = None, + model: str | None = None, prompt="", - messages: List = [], + messages: list = [], completion="", - total_time: Optional[float] = 0.0, # used for replicate, sagemaker - call_type: Optional[CallTypesLiteral] = None, + total_time: float | None = 0.0, # used for replicate, sagemaker + call_type: CallTypesLiteral | None = None, ### REGION ### custom_llm_provider=None, region_name=None, # used for bedrock pricing ### IMAGE GEN ### - size: Optional[str] = None, - quality: Optional[str] = None, - n: Optional[int] = None, # number of images + size: str | None = None, + quality: str | None = None, + n: int | None = None, # number of images ### CUSTOM PRICING ### - custom_cost_per_token: Optional[CostPerToken] = None, - custom_cost_per_second: Optional[float] = None, - optional_params: Optional[dict] = None, - custom_pricing: Optional[bool] = None, - base_model: Optional[str] = None, - standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, - litellm_model_name: Optional[str] = None, - router_model_id: Optional[str] = None, - litellm_logging_obj: Optional[LitellmLoggingObject] = None, + custom_cost_per_token: CostPerToken | None = None, + custom_cost_per_second: float | None = None, + optional_params: dict | None = None, + custom_pricing: bool | None = None, + base_model: str | None = None, + standard_built_in_tools_params: StandardBuiltInToolsParams | None = None, + litellm_model_name: str | None = None, + router_model_id: str | None = None, + litellm_logging_obj: LitellmLoggingObject | None = None, ### SERVICE TIER ### - service_tier: Optional[str] = None, # for OpenAI service tier pricing + service_tier: str | None = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### - data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ) -> float: """ Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. @@ -1176,14 +1176,14 @@ def completion_cost( model = "dall-e-2" # for dall-e-2, azure expects an empty model name # Handle Inputs to completion_cost prompt_tokens = 0 - prompt_characters: Optional[int] = None + prompt_characters: int | None = None completion_tokens = 0 - completion_characters: Optional[int] = None - cache_creation_input_tokens: Optional[int] = None - cache_read_input_tokens: Optional[int] = None + completion_characters: int | None = None + cache_creation_input_tokens: int | None = None + cache_read_input_tokens: int | None = None audio_transcription_file_duration: float = 0.0 - cost_per_token_usage_object: Optional[Usage] = _get_usage_object(completion_response=completion_response) - rerank_billed_units: Optional[RerankBilledUnits] = None + cost_per_token_usage_object: Final[Usage | None] = _get_usage_object(completion_response=completion_response) + rerank_billed_units: RerankBilledUnits | None = None # Extract service_tier from optional_params if not provided directly if service_tier is None and optional_params is not None: @@ -1209,7 +1209,7 @@ def completion_cost( service_tier = _normalize_service_tier(service_tier) - selected_model = _select_model_name_for_cost_calc( + selected_model: Final = _select_model_name_for_cost_calc( model=model, completion_response=completion_response, custom_llm_provider=custom_llm_provider, @@ -1218,7 +1218,7 @@ def completion_cost( router_model_id=router_model_id, ) - potential_model_names = [ + potential_model_names: Final = [ selected_model, _get_response_model(completion_response), ] @@ -1228,13 +1228,13 @@ def completion_cost( for idx, model in enumerate(potential_model_names): try: if verbose_logger.isEnabledFor(logging.DEBUG): - verbose_logger.debug(f"selected model name for cost calculation: {model}") + verbose_logger.debug("selected model name for cost calculation: %s", model) if completion_response is not None and ( isinstance(completion_response, BaseModel) or isinstance(completion_response, dict) ): # tts returns a custom class if isinstance(completion_response, dict): - usage_obj: Optional[Union[dict, Usage]] = completion_response.get("usage", {}) + usage_obj: dict | Usage | None = completion_response.get("usage", {}) else: usage_obj = getattr(completion_response, "usage", {}) if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects(usage_obj=usage_obj): @@ -1251,17 +1251,20 @@ def completion_cost( else: _usage = usage_obj - if ResponseAPILoggingUtils._is_response_api_usage(_usage): + if litellm.AnthropicConfig.is_anthropic_usage_object(_usage): + _usage = ( + litellm.AnthropicConfig() + .calculate_usage(usage_object=_usage, reasoning_content=None) + .model_dump() + ) + elif ResponseAPILoggingUtils._is_response_api_usage(_usage): _usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( _usage ).model_dump() elif TranscriptionUsageObjectTransformation.is_transcription_usage_object(_usage): tr_usage = TranscriptionUsageObjectTransformation.transform_transcription_usage_object( cast( - Union[ - TranscriptionUsageDurationObject, - TranscriptionUsageTokensObject, - ], + TranscriptionUsageDurationObject | TranscriptionUsageTokensObject, _usage, ) ) @@ -1327,9 +1330,8 @@ def completion_cost( ) # strip the llm provider from the model name -> for image gen cost calculation except Exception as e: verbose_logger.debug( - "litellm.cost_calculator.py::completion_cost() - Error inferring custom_llm_provider - {}".format( - str(e) - ) + "litellm.cost_calculator.py::completion_cost() - Error inferring custom_llm_provider - %s", + e, ) if CostCalculatorUtils._call_type_has_image_response(call_type) and isinstance( completion_response, ImageResponse @@ -1348,7 +1350,7 @@ def completion_cost( elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### # Extract custom model_info for deployment-specific pricing - _video_model_info: Optional[ModelInfo] = None + _video_model_info: ModelInfo | None = None if custom_pricing and litellm_logging_obj is not None: _litellm_params = getattr(litellm_logging_obj, "litellm_params", None) if _litellm_params is not None: @@ -1356,8 +1358,8 @@ def completion_cost( _video_model_info = _metadata.get("model_info", None) usage_obj = getattr(completion_response, "usage", None) - duration_seconds: Optional[float] = None - video_resolution: Optional[str] = None + duration_seconds: float | None = None + video_resolution: str | None = None if completion_response is not None and usage_obj: # Handle both dict and Pydantic Usage object if isinstance(usage_obj, dict): @@ -1483,6 +1485,8 @@ def completion_cost( margin_percent=margin_percent, margin_fixed_amount=margin_fixed_amount, margin_total_amount=margin_total_amount, + service_tier=service_tier, + data_residency=data_residency, ) return _final_cost @@ -1491,10 +1495,7 @@ def completion_cost( ): if cost_per_token_usage_object is None or custom_llm_provider is None: raise ValueError( - "usage object and custom_llm_provider must be provided for realtime stream cost calculation. Got cost_per_token_usage_object={}, custom_llm_provider={}".format( - cost_per_token_usage_object, - custom_llm_provider, - ) + f"usage object and custom_llm_provider must be provided for realtime stream cost calculation. Got cost_per_token_usage_object={cost_per_token_usage_object}, custom_llm_provider={custom_llm_provider}" ) return handle_realtime_stream_cost_calculation( results=completion_response.results, @@ -1577,13 +1578,15 @@ def completion_cost( if completion_response is not None: hidden_params = getattr(completion_response, "_hidden_params", None) or {} hidden_model = hidden_params.get("model") or hidden_params.get("litellm_model_name") - if hidden_model and ( - "model_router" in (hidden_model or "").lower() - or "model-router" in (hidden_model or "").lower() + if ( + hidden_model + and ( + "model_router" in (hidden_model or "").lower() + or "model-router" in (hidden_model or "").lower() + ) + or model_for_additional_costs is None ): model_for_additional_costs = hidden_model - elif model_for_additional_costs is None: - model_for_additional_costs = hidden_model if model_for_additional_costs is None: model_for_additional_costs = model additional_costs = _get_additional_costs( @@ -1639,11 +1642,11 @@ def completion_cost( # Store cost breakdown in logging object if available if litellm_logging_obj is not None: - _reasoning_cost: Optional[float] = None - _cache_read_cost: Optional[float] = None - _cache_creation_cost: Optional[float] = None + _reasoning_cost: float | None = None + _cache_read_cost: float | None = None + _cache_creation_cost: float | None = None if cost_per_token_usage_object is not None and model: - _breakdown_provider: Optional[str] = ( + _breakdown_provider: str | None = ( custom_llm_provider if isinstance(custom_llm_provider, str) else None ) _token_type_breakdown = get_token_type_cost_breakdown( @@ -1672,33 +1675,33 @@ def completion_cost( cache_read_cost=_cache_read_cost, cache_creation_cost=_cache_creation_cost, reasoning_cost=_reasoning_cost, + service_tier=service_tier, + data_residency=data_residency, ) return _final_cost except Exception as e: verbose_logger.debug( - "litellm.cost_calculator.py::completion_cost() - Error calculating cost for model={} - {}".format( - model, str(e) - ) + "litellm.cost_calculator.py::completion_cost() - Error calculating cost for model=%s - %s", model, e ) if idx == len(potential_model_names) - 1: raise e - raise Exception("Unable to calculat cost for received potential model names - {}".format(potential_model_names)) + raise Exception(f"Unable to calculat cost for received potential model names - {potential_model_names}") except Exception as e: raise e def get_response_cost_from_hidden_params( - hidden_params: Union[dict, BaseModel], -) -> Optional[float]: + hidden_params: dict | BaseModel, +) -> float | None: if isinstance(hidden_params, BaseModel): _hidden_params_dict = cast(BaseModel, hidden_params).model_dump() else: _hidden_params_dict = hidden_params - additional_headers = _hidden_params_dict.get("additional_headers", {}) + additional_headers: Final = _hidden_params_dict.get("additional_headers", {}) if additional_headers and "llm_provider-x-litellm-response-cost" in additional_headers: - response_cost = additional_headers["llm_provider-x-litellm-response-cost"] + response_cost: Final = additional_headers["llm_provider-x-litellm-response-cost"] if response_cost is None: return None return float(additional_headers["llm_provider-x-litellm-response-cost"]) @@ -1706,22 +1709,20 @@ def get_response_cost_from_hidden_params( def response_cost_calculator( - response_object: Union[ - ModelResponse, - EmbeddingResponse, - ImageResponse, - TranscriptionResponse, - TextCompletionResponse, - HttpxBinaryResponseContent, - RerankResponse, - ResponsesAPIResponse, - LiteLLMRealtimeStreamLoggingObject, - OpenAIModerationResponse, - Response, - SearchResponse, - ], + response_object: ModelResponse + | EmbeddingResponse + | ImageResponse + | TranscriptionResponse + | TextCompletionResponse + | HttpxBinaryResponseContent + | RerankResponse + | ResponsesAPIResponse + | LiteLLMRealtimeStreamLoggingObject + | OpenAIModerationResponse + | Response + | SearchResponse, model: str, - custom_llm_provider: Optional[str], + custom_llm_provider: str | None, call_type: Literal[ "embedding", "aembedding", @@ -1743,18 +1744,18 @@ def response_cost_calculator( "asearch", ], optional_params: dict, - cache_hit: Optional[bool] = None, - base_model: Optional[str] = None, - custom_pricing: Optional[bool] = None, + cache_hit: bool | None = None, + base_model: str | None = None, + custom_pricing: bool | None = None, prompt: str = "", - standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, - litellm_model_name: Optional[str] = None, - router_model_id: Optional[str] = None, - litellm_logging_obj: Optional[LitellmLoggingObject] = None, + standard_built_in_tools_params: StandardBuiltInToolsParams | None = None, + litellm_model_name: str | None = None, + router_model_id: str | None = None, + litellm_logging_obj: LitellmLoggingObject | None = None, ### SERVICE TIER ### - service_tier: Optional[str] = None, # for OpenAI service tier pricing + service_tier: str | None = None, # for OpenAI service tier pricing ### DATA RESIDENCY ### - data_residency: Optional[str] = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") + data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ) -> float: """ Returns @@ -1768,7 +1769,7 @@ def response_cost_calculator( if isinstance(response_object, BaseModel): if hasattr(response_object, "_hidden_params"): response_object._hidden_params["optional_params"] = optional_params - provider_response_cost = get_response_cost_from_hidden_params(response_object._hidden_params) + provider_response_cost: Final = get_response_cost_from_hidden_params(response_object._hidden_params) if provider_response_cost is not None: return provider_response_cost @@ -1795,9 +1796,9 @@ def response_cost_calculator( def ocr_cost( model: str, - custom_llm_provider: Optional[str], - response: Optional[Any] = None, -) -> Tuple[float, float]: + custom_llm_provider: str | None, + response: Any | None = None, +) -> tuple[float, float]: """ Args: model: str - model name @@ -1821,22 +1822,22 @@ def ocr_cost( raise ValueError("OCR response usage_info is None") try: - model_info: Optional[ModelInfo] = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: model_info = None - credits = getattr(response.usage_info, "credits", None) + credits: Final = getattr(response.usage_info, "credits", None) cost_per_credit = None if model_info is not None: cost_per_credit = model_info.get("ocr_cost_per_credit") if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: Optional[float] = None + ocr_cost_per_page: float | None = None if model_info is not None: ocr_cost_per_page = model_info.get("ocr_cost_per_page") - pages_processed = response.usage_info.pages_processed + pages_processed: Final = response.usage_info.pages_processed if pages_processed is None: if cost_per_credit is not None or ocr_cost_per_page is None: # Surface missing usage data instead of silently under-reporting @@ -1869,20 +1870,20 @@ def ocr_cost( ) return 0.0, 0.0 - total_ocr_processing_cost: float = ocr_cost_per_page * pages_processed + total_ocr_processing_cost: Final[float] = ocr_cost_per_page * pages_processed return total_ocr_processing_cost, 0.0 def vector_store_search_cost( - model: Optional[str], + model: str | None, custom_llm_provider: str, response: VectorStoreSearchResponse, -) -> Tuple[float, float]: +) -> tuple[float, float]: """ Returns - float or None: cost of vector store search """ - api_type: Optional[str] = None + api_type: str | None = None if custom_llm_provider is None: custom_llm_provider = "openai" @@ -1891,13 +1892,13 @@ def vector_store_search_cost( model=model, ) - config = ProviderConfigManager.get_provider_vector_stores_config( + config: Final = ProviderConfigManager.get_provider_vector_stores_config( provider=LlmProviders(custom_llm_provider), api_type=api_type, ) if config is None: - verbose_logger.debug(f"Vector store search is not supported for {custom_llm_provider}") + verbose_logger.debug("Vector store search is not supported for %s", custom_llm_provider) return 0.0, 0.0 return config.calculate_vector_store_cost( @@ -1907,9 +1908,9 @@ def vector_store_search_cost( def rerank_cost( model: str, - custom_llm_provider: Optional[str], - billed_units: Optional[RerankBilledUnits] = None, -) -> Tuple[float, float]: + custom_llm_provider: str | None, + billed_units: RerankBilledUnits | None = None, +) -> tuple[float, float]: """ Returns - float or None: cost of response OR none if error. @@ -1917,7 +1918,7 @@ def rerank_cost( _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) try: - config = ProviderConfigManager.get_provider_rerank_config( + config: Final = ProviderConfigManager.get_provider_rerank_config( model=model, api_base=None, present_version_params=[], @@ -1925,9 +1926,7 @@ def rerank_cost( ) try: - model_info: Optional[ModelInfo] = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) + model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) except Exception: model_info = None @@ -1941,17 +1940,17 @@ def rerank_cost( raise e -def transcription_cost(model: str, custom_llm_provider: Optional[str], duration: float) -> Tuple[float, float]: +def transcription_cost(model: str, custom_llm_provider: str | None, duration: float) -> tuple[float, float]: return openai_cost_per_second(model=model, custom_llm_provider=custom_llm_provider, duration=duration) def default_image_cost_calculator( model: str, - custom_llm_provider: Optional[str] = None, - quality: Optional[str] = None, - n: Optional[int] = 1, # Default to 1 image - size: Optional[str] = "1024-x-1024", # OpenAI default - optional_params: Optional[dict] = None, + custom_llm_provider: str | None = None, + quality: str | None = None, + n: int | None = 1, # Default to 1 image + size: str | None = "1024-x-1024", # OpenAI default + optional_params: dict | None = None, ) -> float: """ Default image cost calculator for image generation @@ -1978,23 +1977,23 @@ def default_image_cost_calculator( # Build model names for cost lookup base_model_name = f"{size_str}/{model}" - model_name_without_custom_llm_provider: Optional[str] = None + model_name_without_custom_llm_provider: str | None = None if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "") base_model_name = f"{custom_llm_provider}/{size_str}/{model_name_without_custom_llm_provider}" - model_name_with_quality = f"{quality}/{base_model_name}" if quality else base_model_name + model_name_with_quality: Final = f"{quality}/{base_model_name}" if quality else base_model_name # gpt-image-1 models use low, medium, high quality. If user did not specify quality, use medium fot gpt-image-1 model family - model_name_with_v2_quality = f"{ImageGenerationRequestQuality.HIGH.value}/{base_model_name}" + model_name_with_v2_quality: Final = f"{ImageGenerationRequestQuality.HIGH.value}/{base_model_name}" - verbose_logger.debug(f"Looking up cost for models: {model_name_with_quality}, {base_model_name}") + verbose_logger.debug("Looking up cost for models: %s, %s", model_name_with_quality, base_model_name) - model_without_provider = f"{size_str}/{model.split('/')[-1]}" + model_without_provider: Final = f"{size_str}/{model.split('/')[-1]}" model_with_quality_without_provider = f"{quality}/{model_without_provider}" if quality else model_without_provider # Try model with quality first, fall back to base model name - cost_info: Optional[dict] = None - models_to_check: List[Optional[str]] = [ + cost_info: dict | None = None + models_to_check: Final[list[str | None]] = [ model_name_with_quality, base_model_name, model_name_with_v2_quality, @@ -2023,9 +2022,9 @@ def default_image_cost_calculator( def default_video_cost_calculator( model: str, duration_seconds: float, - custom_llm_provider: Optional[str] = None, - model_info: Optional[ModelInfo] = None, - video_resolution: Optional[str] = None, + custom_llm_provider: str | None = None, + model_info: ModelInfo | None = None, + video_resolution: str | None = None, ) -> float: """ Default video cost calculator for video generation @@ -2046,23 +2045,23 @@ def default_video_cost_calculator( Exception: If model pricing not found in cost map """ # Use custom model_info pricing if provided (deployment-specific pricing) - cost_info: Optional[dict] = None + cost_info: dict | None = None if model_info is not None: cost_info = dict(model_info) else: # Build model names for cost lookup base_model_name = model - model_name_without_custom_llm_provider: Optional[str] = None + model_name_without_custom_llm_provider: str | None = None if custom_llm_provider and model.startswith(f"{custom_llm_provider}/"): model_name_without_custom_llm_provider = model.replace(f"{custom_llm_provider}/", "") base_model_name = f"{custom_llm_provider}/{model_name_without_custom_llm_provider}" - verbose_logger.debug(f"Looking up cost for video model: {base_model_name}") + verbose_logger.debug("Looking up cost for video model: %s", base_model_name) - model_without_provider = model.split("/")[-1] + model_without_provider: Final = model.split("/")[-1] # Try model with provider first, fall back to base model name - models_to_check: List[Optional[str]] = [ + models_to_check: Final[list[str | None]] = [ base_model_name, model, model_without_provider, @@ -2075,7 +2074,7 @@ def default_video_cost_calculator( # If still not found, try with custom_llm_provider prefix if cost_info is None and custom_llm_provider: - prefixed_model = f"{custom_llm_provider}/{model}" + prefixed_model: Final = f"{custom_llm_provider}/{model}" if prefixed_model in litellm.model_cost: cost_info = litellm.model_cost[prefixed_model] @@ -2083,17 +2082,18 @@ def default_video_cost_calculator( raise Exception(f"Model not found in cost map for model={model}") # Check for video-specific cost per second first - video_cost_per_second = cost_info.get("output_cost_per_video_per_second") + video_cost_per_second: Final = cost_info.get("output_cost_per_video_per_second") if video_cost_per_second is not None: return video_cost_per_second * duration_seconds - output_cost_per_second = _video_output_cost_per_second(cost_info, video_resolution) + output_cost_per_second: Final = _video_output_cost_per_second(cost_info, video_resolution) if output_cost_per_second is not None: return output_cost_per_second * duration_seconds # If no cost information found, return 0 verbose_logger.info( - f"No cost information found for video model {model}. Please add pricing to model_prices_and_context_window.json" + "No cost information found for video model %s. Please add pricing to model_prices_and_context_window.json", + model, ) return 0.0 @@ -2101,10 +2101,10 @@ def default_video_cost_calculator( def batch_cost_calculator( usage: Usage, model: str, - custom_llm_provider: Optional[str] = None, - model_info: Optional[ModelInfo] = None, - data_residency: Optional[str] = None, -) -> Tuple[float, float]: + custom_llm_provider: str | None = None, + model_info: ModelInfo | None = None, + data_residency: str | None = None, +) -> tuple[float, float]: """ Calculate the cost of a batch job. @@ -2141,7 +2141,7 @@ def batch_cost_calculator( # but carries no pricing fields. Fall back to the global pricing table so # that standard model pricing is used instead of silently returning $0. try: - global_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + global_info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) if global_info: model_info = global_info except Exception: @@ -2150,31 +2150,31 @@ def batch_cost_calculator( if not model_info: return 0.0, 0.0 - input_cost_per_token_batches = model_info.get("input_cost_per_token_batches") - input_cost_per_token = model_info.get("input_cost_per_token") - output_cost_per_token_batches = model_info.get("output_cost_per_token_batches") - output_cost_per_token = model_info.get("output_cost_per_token") + input_cost_per_token_batches: Final = model_info.get("input_cost_per_token_batches") + input_cost_per_token: Final = model_info.get("input_cost_per_token") + output_cost_per_token_batches: Final = model_info.get("output_cost_per_token_batches") + output_cost_per_token: Final = model_info.get("output_cost_per_token") total_prompt_cost = 0.0 total_completion_cost = 0.0 if input_cost_per_token_batches: total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches elif input_cost_per_token: - details = _parse_prompt_tokens_details(usage) - cache_read_tokens = details["cache_hit_tokens"] - cache_creation_tokens = details["cache_creation_tokens"] + details: Final = _parse_prompt_tokens_details(usage) + cache_read_tokens: Final = details["cache_hit_tokens"] + cache_creation_tokens: Final = details["cache_creation_tokens"] # Subtract cached tokens from prompt_tokens before calculating cost # Fixes issue where cached tokens are being charged again - base_input_tokens = get_billable_input_tokens(usage) - cache_creation_tokens + base_input_tokens: Final = get_billable_input_tokens(usage) - cache_creation_tokens total_prompt_cost = ( base_input_tokens * (input_cost_per_token) / 2 ) # batch cost is usually half of the regular token cost # Add cache read cost if applicable - cache_read_cost_key = _get_service_tier_cost_key("cache_read_input_token_cost", None) + cache_read_cost_key: Final = _get_service_tier_cost_key("cache_read_input_token_cost", None) total_prompt_cost += calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) / 2 - cache_creation_cost = model_info.get("cache_creation_input_token_cost") or input_cost_per_token + cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") or input_cost_per_token total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2 if output_cost_per_token_batches: total_completion_cost = usage.completion_tokens * output_cost_per_token_batches @@ -2183,7 +2183,7 @@ def batch_cost_calculator( usage.completion_tokens * (output_cost_per_token) / 2 ) # batch cost is usually half of the regular token cost - uplift = _get_regional_uplift_multiplier(model_info, data_residency) + uplift: Final = _get_regional_uplift_multiplier(model_info, data_residency) if uplift != 1.0: total_prompt_cost *= uplift total_completion_cost *= uplift @@ -2191,8 +2191,8 @@ def batch_cost_calculator( return total_prompt_cost, total_completion_cost -def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> List[str]: - field_names = list(type(prompt_tokens_details).model_fields) +def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]: + field_names: Final = list(type(prompt_tokens_details).model_fields) if getattr(prompt_tokens_details, "cache_write_tokens", None) is None: return field_names return [attr for attr in field_names if attr != "cache_creation_tokens"] @@ -2200,7 +2200,7 @@ def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> List[str] class BaseTokenUsageProcessor: @staticmethod - def combine_usage_objects(usage_objects: List[Usage]) -> Usage: + def combine_usage_objects(usage_objects: list[Usage]) -> Usage: """ Combine multiple Usage objects into a single Usage object, checking model keys for nested values. """ @@ -2210,7 +2210,7 @@ class BaseTokenUsageProcessor: Usage, ) - combined = Usage() + combined: Final = Usage() # Sum basic token counts for usage in usage_objects: @@ -2272,15 +2272,15 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): @staticmethod def collect_usage_from_realtime_stream_results( results: OpenAIRealtimeStreamList, - ) -> List[Usage]: + ) -> list[Usage]: """ Collect usage from realtime stream results """ - response_done_events: List[OpenAIRealtimeStreamResponseBaseObject] = cast( - List[OpenAIRealtimeStreamResponseBaseObject], + response_done_events: Final[list[OpenAIRealtimeStreamResponseBaseObject]] = cast( + list[OpenAIRealtimeStreamResponseBaseObject], [result for result in results if result["type"] == "response.done"], ) - usage_objects: List[Usage] = [] + usage_objects: Final[list[Usage]] = [] for result in response_done_events: usage_object = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( result["response"].get("usage", {}) @@ -2296,7 +2296,7 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): Collect and combine usage from realtime stream results """ collected_usage_objects = RealtimeAPITokenUsageProcessor.collect_usage_from_realtime_stream_results(results) - combined_usage_object = RealtimeAPITokenUsageProcessor.combine_usage_objects(collected_usage_objects) + combined_usage_object: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects(collected_usage_objects) return combined_usage_object @staticmethod @@ -2309,7 +2309,7 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): ) -_TRANSCRIPTION_COMPLETED_EVENT_TYPE = "conversation.item.input_audio_transcription.completed" +_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" def handle_realtime_stream_cost_calculation( @@ -2317,8 +2317,8 @@ def handle_realtime_stream_cost_calculation( combined_usage_object: Usage, custom_llm_provider: str, litellm_model_name: str, - data_residency: Optional[str] = None, - litellm_logging_obj: Optional[LitellmLoggingObject] = None, + data_residency: str | None = None, + litellm_logging_obj: LitellmLoggingObject | None = None, ) -> float: """ Handles the cost calculation for realtime stream responses. @@ -2329,7 +2329,7 @@ def handle_realtime_stream_cost_calculation( results: A list of OpenAIRealtimeStreamBaseObject objects """ received_model = None - potential_model_names = [] + potential_model_names: Final = [] for result in results: if result["type"] == "session.created": received_model = cast(OpenAIRealtimeStreamSessionEvents, result)["session"].get("model", None) @@ -2354,7 +2354,7 @@ def handle_realtime_stream_cost_calculation( input_cost_per_token += _input_cost_per_token output_cost_per_token += _output_cost_per_token break # exit if we find a valid model - transcription_cost = ( + transcription_cost: Final = ( handle_realtime_transcription_cost_calculation( results=results, custom_llm_provider=custom_llm_provider, @@ -2363,7 +2363,7 @@ def handle_realtime_stream_cost_calculation( if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results) else 0.0 ) - total_cost = input_cost_per_token + output_cost_per_token + transcription_cost + total_cost: Final = input_cost_per_token + output_cost_per_token + transcription_cost _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, @@ -2372,6 +2372,7 @@ def handle_realtime_stream_cost_calculation( cost_for_built_in_tools_cost_usd_dollar=0.0, total_cost_usd_dollar=total_cost, additional_costs={"transcription_cost": transcription_cost} if transcription_cost > 0 else None, + data_residency=data_residency, ) return total_cost @@ -2391,13 +2392,13 @@ def handle_realtime_transcription_cost_calculation( - {"type": "duration", "seconds": } → priced via input_cost_per_second - {"type": "tokens", "input_tokens": ...} → priced via input/audio token cost """ - completed_events = [ + completed_events: Final = [ cast(dict, result) for result in results if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE ] if not completed_events: return 0.0 - model_name = _get_transcription_model_name_from_results(results) or litellm_model_name + model_name: Final = _get_transcription_model_name_from_results(results) or litellm_model_name try: model_info = litellm.get_model_info(model=model_name, custom_llm_provider=custom_llm_provider) except Exception: @@ -2412,7 +2413,7 @@ def handle_realtime_transcription_cost_calculation( def _get_transcription_model_name_from_results( results: OpenAIRealtimeStreamList, -) -> Optional[str]: +) -> str | None: """Resolve the ASR model from a transcription_session.* / session.* event.""" for result in results: if result.get("type") in ( @@ -2431,23 +2432,23 @@ def _get_transcription_model_name_from_results( return None -def _transcription_usage_cost(usage: dict, model_info: Optional[ModelInfo]) -> float: +def _transcription_usage_cost(usage: dict, model_info: ModelInfo | None) -> float: if model_info is None: return 0.0 - usage_type = usage.get("type") + usage_type: Final = usage.get("type") if usage_type == "duration": - seconds = usage.get("seconds") or 0.0 - per_second = model_info.get("input_cost_per_second") or 0.0 + seconds: Final = usage.get("seconds") or 0.0 + per_second: Final = model_info.get("input_cost_per_second") or 0.0 return float(seconds) * float(per_second) if usage_type == "tokens": - input_token_details = usage.get("input_token_details") or {} - audio_tokens = input_token_details.get("audio_tokens") or 0 - text_tokens = input_token_details.get("text_tokens") or 0 - output_tokens = usage.get("output_tokens") or 0 - audio_cost = float(audio_tokens) * float( + input_token_details: Final = usage.get("input_token_details") or {} + audio_tokens: Final = input_token_details.get("audio_tokens") or 0 + text_tokens: Final = input_token_details.get("text_tokens") or 0 + output_tokens: Final = usage.get("output_tokens") or 0 + audio_cost: Final = float(audio_tokens) * float( model_info.get("input_cost_per_audio_token") or model_info.get("input_cost_per_token") or 0.0 ) - text_cost = float(text_tokens) * float(model_info.get("input_cost_per_token") or 0.0) - output_cost = float(output_tokens) * float(model_info.get("output_cost_per_token") or 0.0) + text_cost: Final = float(text_tokens) * float(model_info.get("input_cost_per_token") or 0.0) + output_cost: Final = float(output_tokens) * float(model_info.get("output_cost_per_token") or 0.0) return audio_cost + text_cost + output_cost return 0.0 diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index f2b443eb7bf..9e949db625a 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -2,7 +2,7 @@ Handler for transforming /chat/completions api requests to litellm.responses requests """ -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Final from typing_extensions import TypedDict @@ -14,7 +14,7 @@ if TYPE_CHECKING: class SpeechToCompletionBridgeHandlerInputKwargs(TypedDict): model: str input: str - voice: Optional[Union[str, dict]] + voice: str | dict | None optional_params: dict litellm_params: dict logging_obj: "LiteLLMLoggingObj" @@ -32,23 +32,23 @@ class SpeechToCompletionBridgeHandler: def validate_input_kwargs(self, kwargs: dict) -> SpeechToCompletionBridgeHandlerInputKwargs: from litellm import LiteLLMLoggingObj - model = kwargs.get("model") + model: Final = kwargs.get("model") if model is None or not isinstance(model, str): raise ValueError("model is required") - custom_llm_provider = kwargs.get("custom_llm_provider") + custom_llm_provider: Final = kwargs.get("custom_llm_provider") if custom_llm_provider is None or not isinstance(custom_llm_provider, str): raise ValueError("custom_llm_provider is required") - input = kwargs.get("input") + input: Final = kwargs.get("input") if input is None or not isinstance(input, str): raise ValueError("input is required") - optional_params = kwargs.get("optional_params") + optional_params: Final = kwargs.get("optional_params") if optional_params is None or not isinstance(optional_params, dict): raise ValueError("optional_params is required") - litellm_params = kwargs.get("litellm_params") + litellm_params: Final = kwargs.get("litellm_params") if litellm_params is None or not isinstance(litellm_params, dict): raise ValueError("litellm_params is required") @@ -60,7 +60,7 @@ class SpeechToCompletionBridgeHandler: if headers is None or not isinstance(headers, dict): raise ValueError("headers is required") - logging_obj = kwargs.get("logging_obj") + logging_obj: Final = kwargs.get("logging_obj") if logging_obj is None or not isinstance(logging_obj, LiteLLMLoggingObj): raise ValueError("logging_obj is required") @@ -79,18 +79,18 @@ class SpeechToCompletionBridgeHandler: self, model: str, input: str, - voice: Optional[Union[str, dict]], + voice: str | dict | None, optional_params: dict, litellm_params: dict, headers: dict, logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, ) -> "HttpxBinaryResponseContent": - received_args = locals() + received_args: Final = locals() from litellm import completion from litellm.types.utils import ModelResponse - validated_kwargs = self.validate_input_kwargs(received_args) + validated_kwargs: Final = self.validate_input_kwargs(received_args) model = validated_kwargs["model"] input = validated_kwargs["input"] optional_params = validated_kwargs["optional_params"] @@ -100,7 +100,7 @@ class SpeechToCompletionBridgeHandler: custom_llm_provider = validated_kwargs["custom_llm_provider"] voice = validated_kwargs["voice"] - request_data = self.transformation_handler.transform_request( + request_data: Final = self.transformation_handler.transform_request( model=model, input=input, optional_params=optional_params, @@ -111,7 +111,7 @@ class SpeechToCompletionBridgeHandler: voice=voice, ) - result = completion( + result: Final = completion( **request_data, ) @@ -120,7 +120,7 @@ class SpeechToCompletionBridgeHandler: model_response=result, ) else: - raise Exception("Unmapped response type. Got type: {}".format(type(result))) + raise Exception(f"Unmapped response type. Got type: {type(result)}") speech_to_completion_bridge_handler = SpeechToCompletionBridgeHandler() diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 94de4878b65..a9429b673e4 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Optional, Union, cast +from typing import TYPE_CHECKING, Final, cast from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS @@ -13,14 +13,14 @@ class SpeechToCompletionBridgeTransformationHandler: self, model: str, input: str, - voice: Optional[Union[str, dict]], + voice: str | dict | None, optional_params: dict, litellm_params: dict, headers: dict, litellm_logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, ) -> dict: - passed_optional_params = {} + passed_optional_params: Final = {} for op in optional_params: if op in OPENAI_CHAT_COMPLETION_PARAMS: passed_optional_params[op] = optional_params[op] @@ -66,13 +66,13 @@ class SpeechToCompletionBridgeTransformationHandler: import struct # WAV header parameters - byte_rate = sample_rate * channels * 2 # 2 bytes per sample (16-bit) - block_align = channels * 2 - data_size = len(pcm_data) - file_size = 36 + data_size + byte_rate: Final = sample_rate * channels * 2 # 2 bytes per sample (16-bit) + block_align: Final = channels * 2 + data_size: Final = len(pcm_data) + file_size: Final = 36 + data_size # Create WAV header - wav_header = struct.pack( + wav_header: Final = struct.pack( "<4sI4s4sIHHIIHH4sI", b"RIFF", # Chunk ID file_size, # Chunk Size @@ -103,17 +103,17 @@ class SpeechToCompletionBridgeTransformationHandler: from litellm.types.llms.openai import HttpxBinaryResponseContent from litellm.types.utils import Choices - audio_part = cast(Choices, model_response.choices[0]).message.audio + audio_part: Final = cast(Choices, model_response.choices[0]).message.audio if audio_part is None: raise ValueError("No audio part found in the response") - audio_content = audio_part.data + audio_content: Final = audio_part.data # Decode base64 to get binary content binary_data = base64.b64decode(audio_content) # Check if this is a Gemini TTS model that returns raw PCM16 data - model = getattr(model_response, "model", "") - headers = {} + model: Final = getattr(model_response, "model", "") + headers: Final = {} if self._is_gemini_tts_model(model): # Convert PCM16 to WAV format for proper audio file playback binary_data = self._convert_pcm16_to_wav(binary_data) @@ -122,5 +122,5 @@ class SpeechToCompletionBridgeTransformationHandler: headers["Content-Type"] = "audio/mpeg" # Create an httpx.Response object - response = httpx.Response(status_code=200, content=binary_data, headers=headers) + response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers) return HttpxBinaryResponseContent(response) diff --git a/litellm/evals/__init__.py b/litellm/evals/__init__.py index 89dfb62b2b7..14311ded659 100644 --- a/litellm/evals/__init__.py +++ b/litellm/evals/__init__.py @@ -18,16 +18,16 @@ from .main import ( ) __all__ = [ - "acreate_eval", - "alist_evals", - "aget_eval", - "aupdate_eval", - "adelete_eval", "acancel_eval", - "create_eval", - "list_evals", - "get_eval", - "update_eval", - "delete_eval", + "acreate_eval", + "adelete_eval", + "aget_eval", + "alist_evals", + "aupdate_eval", "cancel_eval", + "create_eval", + "delete_eval", + "get_eval", + "list_evals", + "update_eval", ] diff --git a/litellm/evals/main.py b/litellm/evals/main.py index d4e9d638583..bf6337bd234 100644 --- a/litellm/evals/main.py +++ b/litellm/evals/main.py @@ -5,8 +5,9 @@ Provides create, list, get, update, delete, and cancel operations for evals import asyncio import contextvars +from collections.abc import Coroutine from functools import partial -from typing import Any, Coroutine, Dict, List, Optional, Union +from typing import Any, Final import httpx @@ -35,20 +36,20 @@ from litellm.utils import ProviderConfigManager, client # Initialize HTTP handler base_llm_http_handler = BaseLLMHTTPHandler() -DEFAULT_OPENAI_API_BASE = "https://api.openai.com" +DEFAULT_OPENAI_API_BASE: Final = "https://api.openai.com" @client async def acreate_eval( - data_source_config: Dict[str, Any], - testing_criteria: List[Dict[str, Any]], - name: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + data_source_config: dict[str, Any], + testing_criteria: list[dict[str, Any]], + name: str | None = None, + metadata: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Eval: """ @@ -69,12 +70,12 @@ async def acreate_eval( Returns: Eval object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acreate_eval"] = True - func = partial( + func: Final = partial( create_eval, data_source_config=data_source_config, testing_criteria=testing_criteria, @@ -88,9 +89,9 @@ async def acreate_eval( **kwargs, ) - 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 @@ -109,17 +110,17 @@ async def acreate_eval( @client def create_eval( - data_source_config: Dict[str, Any], - testing_criteria: List[Dict[str, Any]], - name: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + data_source_config: dict[str, Any], + testing_criteria: list[dict[str, Any]], + name: str | None = None, + metadata: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[Eval, Coroutine[Any, Any, Eval]]: +) -> Eval | Coroutine[Any, Any, Eval]: """ Create a new evaluation @@ -138,21 +139,21 @@ def create_eval( Returns: Eval object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("acreate_eval", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("acreate_eval", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -160,7 +161,7 @@ def create_eval( raise ValueError(f"CREATE eval is not supported for {custom_llm_provider}") # Build create request - create_request: CreateEvalRequest = { + create_request: Final[CreateEvalRequest] = { "data_source_config": data_source_config, # type: ignore "testing_criteria": testing_criteria, # type: ignore } @@ -176,15 +177,15 @@ def create_eval( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - request_body = evals_api_provider_config.transform_create_eval_request( + request_body: Final = evals_api_provider_config.transform_create_eval_request( create_request=create_request, litellm_params=litellm_params, headers=headers, ) # Get API base and URL - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE - url = evals_api_provider_config.get_complete_url(api_base=api_base, endpoint="evals") + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + url: Final = evals_api_provider_config.get_complete_url(api_base=api_base, endpoint="evals") # Pre-call logging litellm_logging_obj.update_from_kwargs( @@ -198,7 +199,7 @@ def create_eval( ) # Make HTTP request - response = base_llm_http_handler.create_eval_handler( # type: ignore + response: Final = base_llm_http_handler.create_eval_handler( # type: ignore url=url, request_body=request_body, evals_api_provider_config=evals_api_provider_config, @@ -225,15 +226,15 @@ def create_eval( @client async def alist_evals( - limit: Optional[int] = None, - after: Optional[str] = None, - before: Optional[str] = None, - order: Optional[str] = None, - order_by: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + limit: int | None = None, + after: str | None = None, + before: str | None = None, + order: str | None = None, + order_by: str | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> ListEvalsResponse: """ @@ -254,12 +255,12 @@ async def alist_evals( Returns: ListEvalsResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["alist_evals"] = True - func = partial( + func: Final = partial( list_evals, limit=limit, after=after, @@ -273,9 +274,9 @@ async def alist_evals( **kwargs, ) - 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 @@ -294,17 +295,17 @@ async def alist_evals( @client def list_evals( - limit: Optional[int] = None, - after: Optional[str] = None, - before: Optional[str] = None, - order: Optional[str] = None, - order_by: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + limit: int | None = None, + after: str | None = None, + before: str | None = None, + order: str | None = None, + order_by: str | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[ListEvalsResponse, Coroutine[Any, Any, ListEvalsResponse]]: +) -> ListEvalsResponse | Coroutine[Any, Any, ListEvalsResponse]: """ List all evaluations @@ -323,21 +324,21 @@ def list_evals( Returns: ListEvalsResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("alist_evals", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("alist_evals", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -345,7 +346,7 @@ def list_evals( raise ValueError(f"LIST evals is not supported for {custom_llm_provider}") # Build list parameters - list_params: ListEvalsParams = {} + list_params: Final[ListEvalsParams] = {} if limit is not None: list_params["limit"] = limit if after is not None: @@ -384,7 +385,7 @@ def list_evals( ) # Make HTTP request - response = base_llm_http_handler.list_evals_handler( # type: ignore + response: Final = base_llm_http_handler.list_evals_handler( # type: ignore url=url, query_params=query_params, evals_api_provider_config=evals_api_provider_config, @@ -412,10 +413,10 @@ def list_evals( @client async def aget_eval( eval_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Eval: """ @@ -432,12 +433,12 @@ async def aget_eval( Returns: Eval object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["aget_eval"] = True - func = partial( + func: Final = partial( get_eval, eval_id=eval_id, extra_headers=extra_headers, @@ -447,9 +448,9 @@ async def aget_eval( **kwargs, ) - 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 @@ -469,12 +470,12 @@ async def aget_eval( @client def get_eval( eval_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[Eval, Coroutine[Any, Any, Eval]]: +) -> Eval | Coroutine[Any, Any, Eval]: """ Get an evaluation by ID @@ -489,21 +490,21 @@ def get_eval( Returns: Eval object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("aget_eval", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("aget_eval", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -515,7 +516,7 @@ def get_eval( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE url, headers = evals_api_provider_config.transform_get_eval_request( eval_id=eval_id, api_base=api_base, @@ -535,7 +536,7 @@ def get_eval( ) # Make HTTP request - response = base_llm_http_handler.get_eval_handler( # type: ignore + response: Final = base_llm_http_handler.get_eval_handler( # type: ignore url=url, evals_api_provider_config=evals_api_provider_config, custom_llm_provider=custom_llm_provider, @@ -562,13 +563,13 @@ def get_eval( @client async def aupdate_eval( eval_id: str, - name: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + name: str | None = None, + metadata: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Eval: """ @@ -588,12 +589,12 @@ async def aupdate_eval( Returns: Eval object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["aupdate_eval"] = True - func = partial( + func: Final = partial( update_eval, eval_id=eval_id, name=name, @@ -606,9 +607,9 @@ async def aupdate_eval( **kwargs, ) - 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 @@ -628,15 +629,15 @@ async def aupdate_eval( @client def update_eval( eval_id: str, - name: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + name: str | None = None, + metadata: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[Eval, Coroutine[Any, Any, Eval]]: +) -> Eval | Coroutine[Any, Any, Eval]: """ Update an evaluation @@ -654,21 +655,21 @@ def update_eval( Returns: Eval object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("aupdate_eval", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("aupdate_eval", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -676,14 +677,14 @@ def update_eval( raise ValueError(f"UPDATE eval is not supported for {custom_llm_provider}") # Build update request - update_request: UpdateEvalRequest = {} + update_request: Final[UpdateEvalRequest] = {} if name is not None: update_request["name"] = name # Filter metadata to exclude internal LiteLLM fields if metadata is not None: # List of internal LiteLLM metadata keys that should NOT be sent to OpenAI - internal_keys = { + internal_keys: Final = { "headers", "requester_metadata", "user_api_key_hash", @@ -716,7 +717,7 @@ def update_eval( "user_agent", } # Only include user-provided metadata keys - filtered_metadata = {k: v for k, v in metadata.items() if k not in internal_keys} + filtered_metadata: Final = {k: v for k, v in metadata.items() if k not in internal_keys} if filtered_metadata: # Only add if there's user metadata update_request["metadata"] = filtered_metadata @@ -729,7 +730,7 @@ def update_eval( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE ( url, headers, @@ -754,7 +755,7 @@ def update_eval( ) # Make HTTP request - response = base_llm_http_handler.update_eval_handler( # type: ignore + response: Final = base_llm_http_handler.update_eval_handler( # type: ignore url=url, request_body=request_body, evals_api_provider_config=evals_api_provider_config, @@ -782,10 +783,10 @@ def update_eval( @client async def adelete_eval( eval_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> DeleteEvalResponse: """ @@ -802,12 +803,12 @@ async def adelete_eval( Returns: DeleteEvalResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["adelete_eval"] = True - func = partial( + func: Final = partial( delete_eval, eval_id=eval_id, extra_headers=extra_headers, @@ -817,9 +818,9 @@ async def adelete_eval( **kwargs, ) - 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 @@ -839,12 +840,12 @@ async def adelete_eval( @client def delete_eval( eval_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[DeleteEvalResponse, Coroutine[Any, Any, DeleteEvalResponse]]: +) -> DeleteEvalResponse | Coroutine[Any, Any, DeleteEvalResponse]: """ Delete an evaluation @@ -859,21 +860,21 @@ def delete_eval( Returns: DeleteEvalResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("adelete_eval", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("adelete_eval", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -885,7 +886,7 @@ def delete_eval( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE url, headers = evals_api_provider_config.transform_delete_eval_request( eval_id=eval_id, api_base=api_base, @@ -905,7 +906,7 @@ def delete_eval( ) # Make HTTP request - response = base_llm_http_handler.delete_eval_handler( # type: ignore + response: Final = base_llm_http_handler.delete_eval_handler( # type: ignore url=url, evals_api_provider_config=evals_api_provider_config, custom_llm_provider=custom_llm_provider, @@ -932,10 +933,10 @@ def delete_eval( @client async def acancel_eval( eval_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> CancelEvalResponse: """ @@ -952,12 +953,12 @@ async def acancel_eval( Returns: CancelEvalResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acancel_eval"] = True - func = partial( + func: Final = partial( cancel_eval, eval_id=eval_id, extra_headers=extra_headers, @@ -967,9 +968,9 @@ async def acancel_eval( **kwargs, ) - 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 @@ -989,12 +990,12 @@ async def acancel_eval( @client def cancel_eval( eval_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[CancelEvalResponse, Coroutine[Any, Any, CancelEvalResponse]]: +) -> CancelEvalResponse | Coroutine[Any, Any, CancelEvalResponse]: """ Cancel a running evaluation @@ -1009,21 +1010,21 @@ def cancel_eval( Returns: CancelEvalResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("acancel_eval", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("acancel_eval", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -1035,7 +1036,7 @@ def cancel_eval( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE ( url, headers, @@ -1059,7 +1060,7 @@ def cancel_eval( ) # Make HTTP request - response = base_llm_http_handler.cancel_eval_handler( # type: ignore + response: Final = base_llm_http_handler.cancel_eval_handler( # type: ignore url=url, evals_api_provider_config=evals_api_provider_config, custom_llm_provider=custom_llm_provider, @@ -1091,14 +1092,14 @@ def cancel_eval( @client async def acreate_run( eval_id: str, - data_source: Dict[str, Any], - name: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + data_source: dict[str, Any], + name: str | None = None, + metadata: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Run: """ @@ -1119,12 +1120,12 @@ async def acreate_run( Returns: Run object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acreate_run"] = True - func = partial( + func: Final = partial( create_run, eval_id=eval_id, data_source=data_source, @@ -1138,9 +1139,9 @@ async def acreate_run( **kwargs, ) - 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 @@ -1160,16 +1161,16 @@ async def acreate_run( @client def create_run( eval_id: str, - data_source: Dict[str, Any], - name: Optional[str] = None, - metadata: Optional[Dict[str, Any]] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + data_source: dict[str, Any], + name: str | None = None, + metadata: dict[str, Any] | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[Run, Coroutine[Any, Any, Run]]: +) -> Run | Coroutine[Any, Any, Run]: """ Create a new run for an evaluation @@ -1188,21 +1189,21 @@ def create_run( Returns: Run object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("acreate_run", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("acreate_run", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -1210,7 +1211,7 @@ def create_run( raise ValueError(f"CREATE run is not supported for {custom_llm_provider}") # Build create request - create_request: CreateRunRequest = { + create_request: Final[CreateRunRequest] = { "data_source": data_source, # type: ignore } if name is not None: @@ -1227,7 +1228,7 @@ def create_run( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE url, request_body = evals_api_provider_config.transform_create_run_request( eval_id=eval_id, create_request=create_request, @@ -1247,7 +1248,7 @@ def create_run( ) # Make HTTP request (default 600s timeout for long-running operations) - response = base_llm_http_handler.create_run_handler( # type: ignore + response: Final = base_llm_http_handler.create_run_handler( # type: ignore url=url, request_body=request_body, evals_api_provider_config=evals_api_provider_config, @@ -1275,14 +1276,14 @@ def create_run( @client async def alist_runs( eval_id: str, - limit: Optional[int] = None, - after: Optional[str] = None, - before: Optional[str] = None, - order: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + limit: int | None = None, + after: str | None = None, + before: str | None = None, + order: str | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> ListRunsResponse: """ @@ -1303,12 +1304,12 @@ async def alist_runs( Returns: ListRunsResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["alist_runs"] = True - func = partial( + func: Final = partial( list_runs, eval_id=eval_id, limit=limit, @@ -1322,9 +1323,9 @@ async def alist_runs( **kwargs, ) - 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 @@ -1344,16 +1345,16 @@ async def alist_runs( @client def list_runs( eval_id: str, - limit: Optional[int] = None, - after: Optional[str] = None, - before: Optional[str] = None, - order: Optional[str] = None, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + limit: int | None = None, + after: str | None = None, + before: str | None = None, + order: str | None = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[ListRunsResponse, Coroutine[Any, Any, ListRunsResponse]]: +) -> ListRunsResponse | Coroutine[Any, Any, ListRunsResponse]: """ List all runs for an evaluation @@ -1372,21 +1373,21 @@ def list_runs( Returns: ListRunsResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("alist_runs", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("alist_runs", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -1394,7 +1395,7 @@ def list_runs( raise ValueError(f"LIST runs is not supported for {custom_llm_provider}") # Build list parameters - list_params: ListRunsParams = {} + list_params: Final[ListRunsParams] = {} if limit is not None: list_params["limit"] = limit if after is not None: @@ -1432,7 +1433,7 @@ def list_runs( ) # Make HTTP request - response = base_llm_http_handler.list_runs_handler( # type: ignore + response: Final = base_llm_http_handler.list_runs_handler( # type: ignore url=url, query_params=query_params, evals_api_provider_config=evals_api_provider_config, @@ -1461,10 +1462,10 @@ def list_runs( async def aget_run( eval_id: str, run_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Run: """ @@ -1482,12 +1483,12 @@ async def aget_run( Returns: Run object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["aget_run"] = True - func = partial( + func: Final = partial( get_run, eval_id=eval_id, run_id=run_id, @@ -1498,9 +1499,9 @@ async def aget_run( **kwargs, ) - 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 @@ -1521,12 +1522,12 @@ async def aget_run( def get_run( eval_id: str, run_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[Run, Coroutine[Any, Any, Run]]: +) -> Run | Coroutine[Any, Any, Run]: """ Get a specific run @@ -1542,21 +1543,21 @@ def get_run( Returns: Run object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("aget_run", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("aget_run", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -1568,7 +1569,7 @@ def get_run( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE url, headers = evals_api_provider_config.transform_get_run_request( eval_id=eval_id, run_id=run_id, @@ -1589,7 +1590,7 @@ def get_run( ) # Make HTTP request - response = base_llm_http_handler.get_run_handler( # type: ignore + response: Final = base_llm_http_handler.get_run_handler( # type: ignore url=url, evals_api_provider_config=evals_api_provider_config, custom_llm_provider=custom_llm_provider, @@ -1617,10 +1618,10 @@ def get_run( async def acancel_run( eval_id: str, run_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> CancelRunResponse: """ @@ -1638,12 +1639,12 @@ async def acancel_run( Returns: CancelRunResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acancel_run"] = True - func = partial( + func: Final = partial( cancel_run, eval_id=eval_id, run_id=run_id, @@ -1654,9 +1655,9 @@ async def acancel_run( **kwargs, ) - 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 @@ -1677,12 +1678,12 @@ async def acancel_run( def cancel_run( eval_id: str, run_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[CancelRunResponse, Coroutine[Any, Any, CancelRunResponse]]: +) -> CancelRunResponse | Coroutine[Any, Any, CancelRunResponse]: """ Cancel a running run @@ -1698,21 +1699,21 @@ def cancel_run( Returns: CancelRunResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("acancel_run", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("acancel_run", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -1724,7 +1725,7 @@ def cancel_run( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE ( url, headers, @@ -1749,7 +1750,7 @@ def cancel_run( ) # Make HTTP request - response = base_llm_http_handler.cancel_run_handler( # type: ignore + response: Final = base_llm_http_handler.cancel_run_handler( # type: ignore url=url, evals_api_provider_config=evals_api_provider_config, custom_llm_provider=custom_llm_provider, @@ -1782,10 +1783,10 @@ def cancel_run( async def adelete_run( eval_id: str, run_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, ) -> RunDeleteResponse: """ @@ -1803,12 +1804,12 @@ async def adelete_run( Returns: RunDeleteResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["adelete_run"] = True - func = partial( + func: Final = partial( delete_run, eval_id=eval_id, run_id=run_id, @@ -1819,9 +1820,9 @@ async def adelete_run( **kwargs, ) - 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 @@ -1842,12 +1843,12 @@ async def adelete_run( def delete_run( eval_id: str, run_id: str, - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - custom_llm_provider: Optional[str] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[RunDeleteResponse, Coroutine[Any, Any, RunDeleteResponse]]: +) -> RunDeleteResponse | Coroutine[Any, Any, RunDeleteResponse]: """ Delete a run @@ -1863,21 +1864,21 @@ def delete_run( Returns: RunDeleteResponse object """ - local_vars = locals() + local_vars: Final = locals() try: - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - _is_async = kwargs.pop("adelete_run", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + _is_async: Final = kwargs.pop("adelete_run", False) is True # Get LiteLLM parameters - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) # Determine provider if custom_llm_provider is None: custom_llm_provider = "openai" # Get provider config - evals_api_provider_config: Optional[BaseEvalsAPIConfig] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore + evals_api_provider_config: BaseEvalsAPIConfig | None = ProviderConfigManager.get_provider_evals_api_config( # type: ignore provider=litellm.LlmProviders(custom_llm_provider), ) @@ -1889,7 +1890,7 @@ def delete_run( headers = evals_api_provider_config.validate_environment(headers=headers, litellm_params=litellm_params) # Transform request - api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE + api_base: Final = litellm_params.api_base or DEFAULT_OPENAI_API_BASE ( url, headers, @@ -1914,7 +1915,7 @@ def delete_run( ) # Make HTTP request - response = base_llm_http_handler.delete_run_handler( # type: ignore + response: Final = base_llm_http_handler.delete_run_handler( # type: ignore url=url, evals_api_provider_config=evals_api_provider_config, custom_llm_provider=custom_llm_provider, diff --git a/litellm/exceptions.py b/litellm/exceptions.py index fd0a2afb3e8..dfb0fc32f5f 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,7 +10,7 @@ ## LiteLLM versions of the OpenAI Exception Types import enum -from typing import Any, Dict, Optional, Union +from typing import Any, Final import httpx import openai @@ -81,11 +81,11 @@ class RateLimitType(str, enum.Enum): """Per-session max-iterations cap reached (agent-style flows).""" -_RATE_LIMIT_CATEGORY_VALUES = frozenset(c.value for c in RateLimitErrorCategory) -_RATE_LIMIT_TYPE_VALUES = frozenset(t.value for t in RateLimitType) +_RATE_LIMIT_CATEGORY_VALUES: Final = frozenset(c.value for c in RateLimitErrorCategory) +_RATE_LIMIT_TYPE_VALUES: Final = frozenset(t.value for t in RateLimitType) -def validate_rate_limit_category(value: Any) -> Optional[str]: +def validate_rate_limit_category(value: Any) -> str | None: """Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`. Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus @@ -100,7 +100,7 @@ def validate_rate_limit_category(value: Any) -> Optional[str]: return None -def validate_rate_limit_type(value: Any) -> Optional[str]: +def validate_rate_limit_type(value: Any) -> str | None: """Return ``value`` only if it matches a known :class:`RateLimitType`. See :func:`validate_rate_limit_category` for the rationale. @@ -112,7 +112,7 @@ def validate_rate_limit_type(value: Any) -> Optional[str]: return None -_MINIMAL_ERROR_RESPONSE: Optional[httpx.Response] = None +_MINIMAL_ERROR_RESPONSE: httpx.Response | None = None def _get_minimal_error_response() -> httpx.Response: @@ -132,13 +132,13 @@ class AuthenticationError(openai.AuthenticationError): # type: ignore message, llm_provider, model, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 401 - self.message = "litellm.AuthenticationError: {}".format(message) + self.message = f"litellm.AuthenticationError: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info @@ -176,13 +176,13 @@ class NotFoundError(openai.NotFoundError): # type: ignore message, model, llm_provider, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 404 - self.message = "litellm.NotFoundError: {}".format(message) + self.message = f"litellm.NotFoundError: {message}" self.model = model self.llm_provider = llm_provider self.litellm_debug_info = litellm_debug_info @@ -219,14 +219,14 @@ class BadRequestError(openai.BadRequestError): # type: ignore message, model, llm_provider, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, - body: Optional[dict] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, + body: dict | None = None, ): self.status_code = 400 - self.message = "litellm.BadRequestError: {}".format(message) + self.message = f"litellm.BadRequestError: {message}" self.model = model self.llm_provider = llm_provider self.litellm_debug_info = litellm_debug_info @@ -270,11 +270,11 @@ class ImageFetchError(BadRequestError): message, model=None, llm_provider=None, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, - body: Optional[dict] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, + body: dict | None = None, ): super().__init__( message=message, @@ -295,12 +295,12 @@ class UnprocessableEntityError(openai.UnprocessableEntityError): # type: ignore model, llm_provider, response: httpx.Response, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 422 - self.message = "litellm.UnprocessableEntityError: {}".format(message) + self.message = f"litellm.UnprocessableEntityError: {message}" self.model = model self.llm_provider = llm_provider self.litellm_debug_info = litellm_debug_info @@ -333,19 +333,19 @@ class Timeout(openai.APITimeoutError): # type: ignore message, model, llm_provider, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, - headers: Optional[dict] = None, - exception_status_code: Optional[int] = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, + headers: dict | None = None, + exception_status_code: int | None = None, ): - request = httpx.Request( + request: Final = httpx.Request( method="POST", url="https://api.openai.com/v1", ) super().__init__(request=request) # Call the base class constructor with the parameters it needs self.status_code = exception_status_code or 408 - self.message = "litellm.Timeout: {}".format(message) + self.message = f"litellm.Timeout: {message}" self.model = model self.llm_provider = llm_provider self.litellm_debug_info = litellm_debug_info @@ -378,12 +378,12 @@ class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore llm_provider, model, response: httpx.Response, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 403 - self.message = "litellm.PermissionDeniedError: {}".format(message) + self.message = f"litellm.PermissionDeniedError: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info @@ -428,17 +428,17 @@ class RateLimitError(openai.RateLimitError): # type: ignore message, llm_provider, model, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, - category: Union[str, RateLimitErrorCategory] = (RateLimitErrorCategory.VENDOR_RATE_LIMIT), - rate_limit_type: Optional[Union[str, RateLimitType]] = None, - headers: Optional[Dict[str, str]] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, + category: str | RateLimitErrorCategory = (RateLimitErrorCategory.VENDOR_RATE_LIMIT), + rate_limit_type: str | RateLimitType | None = None, + headers: dict[str, str] | None = None, detail: Any = None, ): self.status_code = 429 - self.message = "litellm.RateLimitError: {}".format(message) + self.message = f"litellm.RateLimitError: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info @@ -448,7 +448,7 @@ class RateLimitError(openai.RateLimitError): # type: ignore # Which dimension was exceeded — request count, token count, parallel # requests, budget, max iterations. None when the source didn't # classify the failure (e.g. legacy vendor 429 with no header hints). - self.rate_limit_type: Optional[str] = ( + self.rate_limit_type: str | None = ( rate_limit_type.value if isinstance(rate_limit_type, RateLimitType) else rate_limit_type ) # Headers explicitly attached to the error (e.g. retry-after, @@ -464,8 +464,8 @@ class RateLimitError(openai.RateLimitError): # type: ignore # headers stay reachable on `e.response.headers` for callers that # explicitly want them; only the proxy-supplied `headers=` kwarg # makes it onto `self.headers`. - _response_headers = getattr(response, "headers", None) if response is not None else None - self.headers: Optional[Dict[str, str]] = {k: str(v) for k, v in headers.items()} if headers else None + _response_headers: Final = getattr(response, "headers", None) if response is not None else None + self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None # Mirrors FastAPI HTTPException.detail so the same instance can be # serialized through both the ProxyException and HTTPException paths. self.detail = detail if detail is not None else self.message @@ -507,8 +507,8 @@ class ContextWindowExceededError(BadRequestError): # type: ignore message, model, llm_provider, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, ): self.status_code = 400 self.model = model @@ -523,7 +523,7 @@ class ContextWindowExceededError(BadRequestError): # type: ignore ) # Call the base class constructor with the parameters it needs # set after, to make it clear the raised error is a context window exceeded error - self.message = "litellm.ContextWindowExceededError: {}".format(self.message) + self.message = f"litellm.ContextWindowExceededError: {self.message}" def __str__(self): _message = self.message @@ -550,16 +550,16 @@ class RejectedRequestError(BadRequestError): # type: ignore model, llm_provider, request_data: dict, - litellm_debug_info: Optional[str] = None, + litellm_debug_info: str | None = None, ): self.status_code = 400 - self.message = "litellm.RejectedRequestError: {}".format(message) + self.message = f"litellm.RejectedRequestError: {message}" self.model = model self.llm_provider = llm_provider self.litellm_debug_info = litellm_debug_info self.request_data = request_data - request = httpx.Request(method="POST", url="https://api.openai.com/v1") - response = httpx.Response(status_code=400, request=request) + request: Final = httpx.Request(method="POST", url="https://api.openai.com/v1") + response: Final = httpx.Response(status_code=400, request=request) super().__init__( message=self.message, model=self.model, # type: ignore @@ -592,13 +592,13 @@ class ContentPolicyViolationError(BadRequestError): # type: ignore message, model, llm_provider, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - provider_specific_fields: Optional[dict] = None, - body: Optional[dict] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + provider_specific_fields: dict | None = None, + body: dict | None = None, ): self.status_code = 400 - self.message = "litellm.ContentPolicyViolationError: {}".format(message) + self.message = f"litellm.ContentPolicyViolationError: {message}" self.model = model self.llm_provider = llm_provider self.litellm_debug_info = litellm_debug_info @@ -636,19 +636,19 @@ class ServiceUnavailableError(openai.APIStatusError): # type: ignore message, llm_provider, model, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 503 - self.message = "litellm.ServiceUnavailableError: {}".format(message) + self.message = f"litellm.ServiceUnavailableError: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers = getattr(response, "headers", None) if response is not None else None + _response_headers: Final = getattr(response, "headers", None) if response is not None else None self.response = httpx.Response( status_code=self.status_code, headers=_response_headers, @@ -684,19 +684,19 @@ class BadGatewayError(openai.APIStatusError): # type: ignore message, llm_provider, model, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 502 - self.message = "litellm.BadGatewayError: {}".format(message) + self.message = f"litellm.BadGatewayError: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers = getattr(response, "headers", None) if response is not None else None + _response_headers: Final = getattr(response, "headers", None) if response is not None else None self.response = httpx.Response( status_code=self.status_code, headers=_response_headers, @@ -732,19 +732,19 @@ class InternalServerError(openai.InternalServerError): # type: ignore message, llm_provider, model, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 500 - self.message = "litellm.InternalServerError: {}".format(message) + self.message = f"litellm.InternalServerError: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - _response_headers = getattr(response, "headers", None) if response is not None else None + _response_headers: Final = getattr(response, "headers", None) if response is not None else None self.response = httpx.Response( status_code=self.status_code, headers=_response_headers, @@ -782,13 +782,13 @@ class APIError(openai.APIError): # type: ignore message, llm_provider, model, - request: Optional[httpx.Request] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + request: httpx.Request | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = status_code - self.message = "litellm.APIError: {}".format(message) + self.message = f"litellm.APIError: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info @@ -822,12 +822,12 @@ class APIConnectionError(openai.APIConnectionError): # type: ignore message, llm_provider, model, - request: Optional[httpx.Request] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + request: httpx.Request | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): - self.message = "litellm.APIConnectionError: {}".format(message) + self.message = f"litellm.APIConnectionError: {message}" self.llm_provider = llm_provider self.model = model self.status_code = 500 @@ -861,15 +861,15 @@ class APIResponseValidationError(openai.APIResponseValidationError): # type: ig message, llm_provider, model, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): - self.message = "litellm.APIResponseValidationError: {}".format(message) + self.message = f"litellm.APIResponseValidationError: {message}" self.llm_provider = llm_provider self.model = model - request = httpx.Request(method="POST", url="https://api.openai.com/v1") - response = httpx.Response(status_code=500, request=request) + request: Final = httpx.Request(method="POST", url="https://api.openai.com/v1") + response: Final = httpx.Response(status_code=500, request=request) self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries @@ -897,9 +897,7 @@ class JSONSchemaValidationError(APIResponseValidationError): self.raw_response = raw_response self.schema = schema self.model = model - message = "litellm.JSONSchemaValidationError: model={}, returned an invalid response={}, for schema={}.\nAccess raw response with `e.raw_response`".format( - model, raw_response, schema - ) + message = f"litellm.JSONSchemaValidationError: model={model}, returned an invalid response={raw_response}, for schema={schema}.\nAccess raw response with `e.raw_response`" self.message = message super().__init__(model=model, message=message, llm_provider=llm_provider) @@ -914,16 +912,16 @@ class UnsupportedParamsError(BadRequestError): def __init__( self, message, - llm_provider: Optional[str] = None, - model: Optional[str] = None, + llm_provider: str | None = None, + model: str | None = None, status_code: int = 400, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = 400 - self.message = "litellm.UnsupportedParamsError: {}".format(message) + self.message = f"litellm.UnsupportedParamsError: {message}" self.model = model self.llm_provider = llm_provider self.litellm_debug_info = litellm_debug_info @@ -935,7 +933,7 @@ class UnsupportedParamsError(BadRequestError): self.num_retries = num_retries -LITELLM_EXCEPTION_TYPES = [ +LITELLM_EXCEPTION_TYPES: Final = [ AuthenticationError, NotFoundError, BadRequestError, @@ -964,10 +962,10 @@ class BudgetExceededError(Exception): self, current_cost: float, max_budget: float, - message: Optional[str] = None, - llm_provider: Optional[str] = None, - entity_type: Optional[str] = None, - entity_id: Optional[str] = None, + message: str | None = None, + llm_provider: str | None = None, + entity_type: str | None = None, + entity_id: str | None = None, ): self.current_cost = current_cost self.max_budget = max_budget @@ -1012,13 +1010,13 @@ class MockException(openai.APIError): message, llm_provider, model, - request: Optional[httpx.Request] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + request: httpx.Request | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, ): self.status_code = status_code - self.message = "litellm.MockException: {}".format(message) + self.message = f"litellm.MockException: {message}" self.llm_provider = llm_provider self.model = model self.litellm_debug_info = litellm_debug_info @@ -1030,7 +1028,7 @@ class MockException(openai.APIError): class LiteLLMUnknownProvider(BadRequestError): - def __init__(self, model: str, custom_llm_provider: Optional[str] = None): + def __init__(self, model: str, custom_llm_provider: str | None = None): self.message = LiteLLMCommonStrings.llm_provider_not_provided.value.format( model=model, custom_llm_provider=custom_llm_provider ) @@ -1043,12 +1041,12 @@ class LiteLLMUnknownProvider(BadRequestError): class GuardrailRaisedException(Exception): def __init__( self, - guardrail_name: Optional[str] = None, + guardrail_name: str | None = None, message: str = "", should_wrap_with_default_message: bool = True, status_code: int = 400, ): - default_message = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}" + default_message: Final = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}" self.guardrail_name = guardrail_name self.status_code = status_code self.message = default_message if should_wrap_with_default_message else message @@ -1059,7 +1057,7 @@ class BlockedPiiEntityError(Exception): def __init__( self, entity_type: str, - guardrail_name: Optional[str] = None, + guardrail_name: str | None = None, status_code: int = 400, ): """ @@ -1078,15 +1076,15 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore message: str, model: str, llm_provider: str, - original_exception: Optional[Exception] = None, - response: Optional[httpx.Response] = None, - litellm_debug_info: Optional[str] = None, - max_retries: Optional[int] = None, - num_retries: Optional[int] = None, + original_exception: Exception | None = None, + response: httpx.Response | None = None, + litellm_debug_info: str | None = None, + max_retries: int | None = None, + num_retries: int | None = None, generated_content: str = "", is_pre_first_chunk: bool = False, ): - original_status = getattr(original_exception, "status_code", None) + original_status: Final = getattr(original_exception, "status_code", None) self.status_code = int(original_status) if original_status is not None else 503 self.message = f"litellm.MidStreamFallbackError: {message}" self.model = model @@ -1111,11 +1109,11 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore self.response = response # Save the original attributes before they are overridden by ServiceUnavailableError - _saved_response = self.response - _saved_request = getattr(self.response, "request", None) or httpx.Request( + _saved_response: Final = self.response + _saved_request: Final = getattr(self.response, "request", None) or httpx.Request( method="POST", url=f"https://{llm_provider}.com/v1/" ) - _saved_message = self.message + _saved_message: Final = self.message # Call the parent constructor (which hardcodes status_code=503 and modifies the response object) super().__init__( @@ -1142,7 +1140,7 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore if self.max_retries: _message += f", LiteLLM Max Retries: {self.max_retries}" if self.original_exception: - _message += f" Original exception: {type(self.original_exception).__name__}: {str(self.original_exception)}" + _message += f" Original exception: {type(self.original_exception).__name__}: {self.original_exception}" return _message def __repr__(self): @@ -1166,10 +1164,10 @@ class ModifyResponseException(Exception): self, message: str, model: str, - request_data: Dict[str, Any], - guardrail_name: Optional[str] = None, - detection_info: Optional[Dict[str, Any]] = None, - original_response: Optional[Any] = None, + request_data: dict[str, Any], + guardrail_name: str | None = None, + detection_info: dict[str, Any] | None = None, + original_response: Any | None = None, ): self.message = message self.model = model @@ -1201,9 +1199,9 @@ class SensitiveDataRouteException(Exception): self, route_to_model: str, session_id: str, - guardrail_name: Optional[str] = None, - detection_info: Optional[Dict[str, Any]] = None, - message: Optional[str] = None, + guardrail_name: str | None = None, + detection_info: dict[str, Any] | None = None, + message: str | None = None, sticky_session_routing: bool = True, ): self.route_to_model = route_to_model diff --git a/litellm/experimental_mcp_client/__init__.py b/litellm/experimental_mcp_client/__init__.py index 7110d5375e4..5399968ff74 100644 --- a/litellm/experimental_mcp_client/__init__.py +++ b/litellm/experimental_mcp_client/__init__.py @@ -1,3 +1,3 @@ from .tools import call_openai_tool, load_mcp_tools -__all__ = ["load_mcp_tools", "call_openai_tool"] +__all__ = ["call_openai_tool", "load_mcp_tools"] diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index da711463a44..64f4a773901 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -5,24 +5,15 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 import os -from typing import ( - Any, - Awaitable, - Callable, - Dict, - Generator, - List, - Optional, - Tuple, - TypeVar, - Union, -) +from collections.abc import Awaitable, Callable, Generator +from typing import Any, Final, TypeVar + import httpx from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters from mcp.client.sse import sse_client from mcp.client.stdio import stdio_client -streamable_http_client: Optional[Any] = None +streamable_http_client: Any | None = None try: import mcp.client.streamable_http as streamable_http_module # type: ignore @@ -40,6 +31,7 @@ from mcp.types import ( ) from mcp.types import Tool as MCPTool from pydantic import AnyUrl + from litellm._logging import verbose_logger from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR from litellm.llms.custom_httpx.http_handler import get_ssl_configuration @@ -58,15 +50,15 @@ def to_basic_auth(auth_value: str) -> str: return base64.b64encode(auth_value.encode("utf-8")).decode() -def _strip_header_whitespace(headers: Dict[str, str]) -> Dict[str, str]: +def _strip_header_whitespace(headers: dict[str, str]) -> dict[str, str]: return { (key.strip() if isinstance(key, str) else key): (value.strip() if isinstance(value, str) else value) for key, value in headers.items() } -def _first_non_cancelled_cause(exc: BaseException) -> Optional[BaseException]: - queue: List[BaseException] = [exc] +def _first_non_cancelled_cause(exc: BaseException) -> BaseException | None: + queue: Final[list[BaseException]] = [exc] while queue: current = queue.pop(0) nested = getattr(current, "exceptions", None) @@ -92,13 +84,13 @@ class MCPSigV4Auth(httpx.Auth): def __init__( self, - aws_access_key_id: Optional[str] = None, - aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None, - aws_region_name: Optional[str] = None, - aws_service_name: Optional[str] = None, - aws_role_name: Optional[str] = None, - aws_session_name: Optional[str] = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + aws_session_token: str | None = None, + aws_region_name: str | None = None, + aws_service_name: str | None = None, + aws_role_name: str | None = None, + aws_session_name: str | None = None, ): try: from botocore.credentials import Credentials @@ -128,7 +120,7 @@ class MCPSigV4Auth(httpx.Auth): # Fall back to default boto3 credential chain import botocore.session - session = botocore.session.get_session() + session: Final = botocore.session.get_session() self.credentials = session.get_credentials() if self.credentials is None: raise ValueError( @@ -140,29 +132,29 @@ class MCPSigV4Auth(httpx.Auth): @staticmethod def _assume_role( aws_role_name: str, - aws_session_name: Optional[str], - aws_access_key_id: Optional[str], - aws_secret_access_key: Optional[str], - aws_session_token: Optional[str], + aws_session_name: str | None, + aws_access_key_id: str | None, + aws_secret_access_key: str | None, + aws_session_token: str | None, aws_region_name: str, ): """Call STS AssumeRole and return temporary credentials.""" import boto3 from botocore.credentials import Credentials - session_name = aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" - sts_kwargs: dict = {"region_name": aws_region_name} + session_name: Final = aws_session_name or f"litellm-mcp-{int(__import__('time').time())}" + sts_kwargs: Final[dict] = {"region_name": aws_region_name} if aws_access_key_id and aws_secret_access_key: sts_kwargs["aws_access_key_id"] = aws_access_key_id sts_kwargs["aws_secret_access_key"] = aws_secret_access_key if aws_session_token: sts_kwargs["aws_session_token"] = aws_session_token - sts_client = boto3.client("sts", **sts_kwargs) - sts_response = sts_client.assume_role( + sts_client: Final = boto3.client("sts", **sts_kwargs) + sts_response: Final = sts_client.assume_role( RoleArn=aws_role_name, RoleSessionName=session_name, ) - sts_creds = sts_response["Credentials"] + sts_creds: Final = sts_response["Credentials"] return Credentials( access_key=sts_creds["AccessKeyId"], secret_key=sts_creds["SecretAccessKey"], @@ -175,7 +167,7 @@ class MCPSigV4Auth(httpx.Auth): # Build AWSRequest from the httpx Request. # Pass all request headers so the canonical SigV4 signature covers them. - aws_request = AWSRequest( + aws_request: Final = AWSRequest( method=request.method, url=str(request.url), data=request.content, @@ -184,7 +176,7 @@ class MCPSigV4Auth(httpx.Auth): # Sign the request — SigV4Auth.add_auth() adds Authorization, # X-Amz-Date, and X-Amz-Security-Token (if session token present). # Host header is derived automatically from the URL. - sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name) + sigv4: Final = SigV4Auth(self.credentials, self.service_name, self.region_name) sigv4.add_auth(aws_request) # Copy SigV4 headers back to the httpx request for header_name, header_value in aws_request.headers.items(): @@ -207,51 +199,51 @@ class MCPClient: server_url: str = "", transport_type: MCPTransportType = MCPTransport.http, auth_type: MCPAuthType = None, - auth_value: Optional[Union[str, Dict[str, str]]] = None, - timeout: Optional[float] = None, - stdio_config: Optional[MCPStdioConfig] = None, - extra_headers: Optional[Dict[str, str]] = None, - ssl_verify: Optional[VerifyTypes] = None, - aws_auth: Optional[httpx.Auth] = None, - resolved_auth: Optional[httpx.Auth] = None, - sampling_callback: Optional[Callable] = None, - elicitation_callback: Optional[Callable] = None, - logging_callback: Optional[Callable] = None, + auth_value: str | dict[str, str] | None = None, + timeout: float | None = None, + stdio_config: MCPStdioConfig | None = None, + extra_headers: dict[str, str] | None = None, + ssl_verify: VerifyTypes | None = None, + aws_auth: httpx.Auth | None = None, + resolved_auth: httpx.Auth | None = None, + sampling_callback: Callable | None = None, + elicitation_callback: Callable | None = None, + logging_callback: Callable | None = None, ): self.server_url: str = server_url self.transport_type: MCPTransport = transport_type self.auth_type: MCPAuthType = auth_type self.timeout: float = timeout if timeout is not None else MCP_CLIENT_TIMEOUT - self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None - self.stdio_config: Optional[MCPStdioConfig] = stdio_config - self.extra_headers: Optional[Dict[str, str]] = extra_headers - self.ssl_verify: Optional[VerifyTypes] = ssl_verify - self._aws_auth: Optional[httpx.Auth] = aws_auth + self._mcp_auth_value: str | dict[str, str] | None = None + self.stdio_config: MCPStdioConfig | None = stdio_config + self.extra_headers: dict[str, str] | None = extra_headers + self.ssl_verify: VerifyTypes | None = ssl_verify + self._aws_auth: httpx.Auth | None = aws_auth # A pre-resolved httpx.Auth (e.g. from the v2 credential resolver) attached to the # upstream client's auth= slot, taking precedence over the SigV4 aws_auth. - self._resolved_auth: Optional[httpx.Auth] = resolved_auth - self._last_initialize_instructions: Optional[str] = None - self._sampling_callback: Optional[Callable] = sampling_callback - self._elicitation_callback: Optional[Callable] = elicitation_callback - self._logging_callback: Optional[Callable] = logging_callback + self._resolved_auth: httpx.Auth | None = resolved_auth + self._last_initialize_instructions: str | None = None + self._sampling_callback: Callable | None = sampling_callback + self._elicitation_callback: Callable | None = elicitation_callback + self._logging_callback: Callable | None = logging_callback # handle the basic auth value if provided if auth_value: self.update_auth_value(auth_value) def _create_transport_context( self, - ) -> Tuple[Any, Optional[httpx.AsyncClient]]: + ) -> tuple[Any, httpx.AsyncClient | None]: """ Create the appropriate transport context based on transport type. Returns: Tuple of (transport_context, http_client). http_client is only set for HTTP transport and needs cleanup. """ - http_client: Optional[httpx.AsyncClient] = None + http_client: httpx.AsyncClient | None = None if self.transport_type == MCPTransport.stdio: if not self.stdio_config: raise ValueError("stdio_config is required for stdio transport") - server_params = StdioServerParameters( + server_params: Final = StdioServerParameters( command=self.stdio_config.get("command", ""), args=self.stdio_config.get("args", []), env=self._get_safe_stdio_env(self.stdio_config.get("env")), @@ -279,13 +271,13 @@ class MCPClient: headers=headers, timeout=httpx.Timeout(self.timeout), ) - transport_ctx = streamable_http_client( + transport_ctx: Final = streamable_http_client( url=self.server_url, http_client=http_client, ) return transport_ctx, http_client - def _get_safe_stdio_env(self, provided_env: Optional[Dict[str, str]]) -> Optional[Dict[str, str]]: + def _get_safe_stdio_env(self, provided_env: dict[str, str] | None) -> dict[str, str] | None: """ Return a safe environment for the stdio subprocess. @@ -297,7 +289,7 @@ class MCPClient: return provided_env # Minimal allowlist of safe/standard environment variables - safe_keys = { + safe_keys: Final = { "PATH", "HOME", "USER", @@ -321,7 +313,7 @@ class MCPClient: "WINDIR", } - safe_env = {} + safe_env: Final = {} for key in safe_keys: if key in os.environ: safe_env[key] = os.environ[key] @@ -343,25 +335,25 @@ class MCPClient: so that upstream MCP servers can request LLM inference (sampling), user input (elicitation), or send log messages. """ - transport = await transport_ctx.__aenter__() - in_flight_error: Optional[BaseException] = None + transport: Final = await transport_ctx.__aenter__() + in_flight_error: BaseException | None = None try: read_stream, write_stream = transport[0], transport[1] # Build session kwargs with optional callbacks - session_kwargs: Dict[str, Any] = {} + session_kwargs: Final[dict[str, Any]] = {} if self._sampling_callback is not None: session_kwargs["sampling_callback"] = self._sampling_callback if self._elicitation_callback is not None: session_kwargs["elicitation_callback"] = self._elicitation_callback if self._logging_callback is not None: session_kwargs["logging_callback"] = self._logging_callback - session_ctx = ClientSession(read_stream, write_stream, **session_kwargs) - session = await session_ctx.__aenter__() + session_ctx: Final = ClientSession(read_stream, write_stream, **session_kwargs) + session: Final = await session_ctx.__aenter__() try: - init_result = await session.initialize() + init_result: Final = await session.initialize() self._last_initialize_instructions = None if init_result is not None: - ins = getattr(init_result, "instructions", None) + ins: Final = getattr(init_result, "instructions", None) if isinstance(ins, str) and ins.strip(): self._last_initialize_instructions = ins.strip() return await operation(session) @@ -369,7 +361,7 @@ class MCPClient: try: await session_ctx.__aexit__(None, None, None) except BaseException as e: - verbose_logger.debug(f"Error during session context exit: {e}") + verbose_logger.debug("Error during session context exit: %s", e) except BaseException as e: in_flight_error = e raise @@ -377,8 +369,8 @@ class MCPClient: try: await transport_ctx.__aexit__(None, None, None) except BaseException as exit_error: - verbose_logger.debug(f"Error during transport context exit: {exit_error}") - root_cause = _first_non_cancelled_cause(exit_error) + verbose_logger.debug("Error during transport context exit: %s", exit_error) + root_cause: Final = _first_non_cancelled_cause(exit_error) if root_cause is not None and isinstance(in_flight_error, asyncio.CancelledError): raise root_cause from in_flight_error @@ -393,13 +385,13 @@ class MCPClient: quiet_on_error demotes the failure line to debug for callers that own the exception (call_tool / list_tools under raise_on_error), so an expected pass-through re-auth does not emit a warning per call; every other caller keeps the operator-visible warning.""" - http_client: Optional[httpx.AsyncClient] = None + http_client: httpx.AsyncClient | None = None try: self._last_initialize_instructions = None transport_ctx, http_client = self._create_transport_context() return await self._execute_session_operation(transport_ctx, operation) except Exception: - _log = verbose_logger.debug if quiet_on_error else verbose_logger.warning + _log: Final = verbose_logger.debug if quiet_on_error else verbose_logger.warning _log("MCP client run_with_session failed for %s", self.server_url or "stdio") raise finally: @@ -407,9 +399,9 @@ class MCPClient: try: await http_client.aclose() except BaseException as e: - verbose_logger.debug(f"Error during http_client cleanup: {e}") + verbose_logger.debug("Error during http_client cleanup: %s", e) - def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]): + def update_auth_value(self, mcp_auth_value: str | dict[str, str]): """ Set the authentication header for the MCP client. """ @@ -423,7 +415,7 @@ class MCPClient: def _get_auth_headers(self) -> dict: """Generate authentication headers based on auth type.""" - headers = {} + headers: Final = {} if self._mcp_auth_value: if isinstance(self._mcp_auth_value, str): if self.auth_type == MCPAuth.bearer_token: @@ -462,19 +454,19 @@ class MCPClient: def factory( *, - headers: Optional[Dict[str, str]] = None, - timeout: Optional[httpx.Timeout] = None, - auth: Optional[httpx.Auth] = None, + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, ) -> httpx.AsyncClient: """Create an httpx.AsyncClient with LiteLLM's SSL configuration.""" # Get unified SSL configuration using the same logic as http_handler.py - ssl_config = get_ssl_configuration(self.ssl_verify) - verbose_logger.debug(f"MCP client using SSL configuration: {type(ssl_config).__name__}") + ssl_config: Final = get_ssl_configuration(self.ssl_verify) + verbose_logger.debug("MCP client using SSL configuration: %s", type(ssl_config).__name__) # The MCP SDK's sse_client and streamable_http_client call this factory without # passing auth=, so the fallback is used: a v2-resolved auth if present, else the # SigV4 aws_auth. Both are None for the common case — no behavior change. - fallback_auth = self._resolved_auth if self._resolved_auth is not None else self._aws_auth - effective_auth = auth if auth is not None else fallback_auth + fallback_auth: Final = self._resolved_auth if self._resolved_auth is not None else self._aws_auth + effective_auth: Final = auth if auth is not None else fallback_auth return httpx.AsyncClient( headers=headers, timeout=timeout, @@ -485,7 +477,7 @@ class MCPClient: return factory - async def list_tools(self, raise_on_error: bool = False) -> List[MCPTool]: + async def list_tools(self, raise_on_error: bool = False) -> list[MCPTool]: """List available tools from the server. Args: @@ -495,38 +487,40 @@ class MCPClient: MCP client (triggering the upstream OAuth flow) rather than masking them as "connected, no tools". """ - verbose_logger.debug(f"MCP client listing tools from {self.server_url or 'stdio'}") + verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") async def _list_tools_operation(session: ClientSession): return await session.list_tools() try: - result = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) - tool_count = len(result.tools) - tool_names = [tool.name for tool in result.tools] - verbose_logger.info(f"MCP client listed {tool_count} tools from {self.server_url or 'stdio'}: {tool_names}") + result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) + tool_count: Final = len(result.tools) + tool_names: Final = [tool.name for tool in result.tools] + verbose_logger.info( + "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names + ) return result.tools except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") raise except Exception as e: - error_type = type(e).__name__ + error_type: Final = type(e).__name__ # Mirror call_tool: when the caller opted into raise_on_error it owns the exception and # logs it at the fitting level (an expected pass-through re-auth 401 is info, not an # error), so log at debug here to avoid an error-level line + traceback that would trip # error-rate alerts on that expected signal. The swallow path still logs the full # exception because nothing downstream will surface the failure. - _log = verbose_logger.debug if raise_on_error else verbose_logger.exception + _log: Final = verbose_logger.debug if raise_on_error else verbose_logger.exception _log( f"MCP client list_tools failed - " f"Error Type: {error_type}, " - f"Error: {str(e)}, " + f"Error: {e}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" ) # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: - _log_broken = verbose_logger.debug if raise_on_error else verbose_logger.error + _log_broken: Final = verbose_logger.debug if raise_on_error else verbose_logger.error _log_broken( "MCP client detected broken connection/stream during list_tools - " "the MCP server may have crashed, disconnected, or timed out" @@ -541,14 +535,14 @@ class MCPClient: def error_tool_result(exc: Exception) -> MCPCallToolResult: """The error result ``call_tool`` returns when it swallows a failure (no re-execution).""" return MCPCallToolResult( - content=[TextContent(type="text", text=f"{type(exc).__name__}: {str(exc)}")], + content=[TextContent(type="text", text=f"{type(exc).__name__}: {exc}")], isError=True, ) async def call_tool( self, call_tool_request_params: MCPCallToolRequestParams, - host_progress_callback: Optional[Callable] = None, + host_progress_callback: Callable | None = None, raise_on_error: bool = False, ) -> MCPCallToolResult: """ @@ -560,10 +554,10 @@ class MCPClient: an upstream 401 so it can re-mint the exchanged token and retry once; every other caller keeps the default and gets graceful ``isError`` degradation. """ - verbose_logger.info(f"MCP client calling tool '{call_tool_request_params.name}'") + verbose_logger.info("MCP client calling tool '%s'", call_tool_request_params.name) async def on_progress(progress: float, total: float | None, message: str | None): - percentage = (progress / total * 100) if total else 0 + percentage: Final = (progress / total * 100) if total else 0 verbose_logger.info( f"MCP Tool '{call_tool_request_params.name}' progress: " f"{progress}/{total} ({percentage:.0f}%) - {message or ''}" @@ -573,7 +567,7 @@ class MCPClient: try: await host_progress_callback(progress, total) except Exception as e: - verbose_logger.warning(f"Failed to forward to Host: {e}") + verbose_logger.warning("Failed to forward to Host: %s", e) async def _call_tool_operation(session: ClientSession): verbose_logger.debug("MCP client sending tool call to session") @@ -584,29 +578,29 @@ class MCPClient: ) try: - tool_result = await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error) - verbose_logger.info(f"MCP client tool call '{call_tool_request_params.name}' completed successfully") + tool_result: Final = await self.run_with_session(_call_tool_operation, quiet_on_error=raise_on_error) + verbose_logger.info("MCP client tool call '%s' completed successfully", call_tool_request_params.name) return tool_result except asyncio.CancelledError: - verbose_logger.warning(f"MCP client tool call timed out after {self.timeout}s for {self.server_url}") + verbose_logger.warning("MCP client tool call timed out after %ss for %s", self.timeout, self.server_url) raise except Exception as e: import traceback - error_trace = traceback.format_exc() - verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}") + error_trace: Final = traceback.format_exc() + verbose_logger.debug("MCP client tool call traceback:\n%s", error_trace) # Log detailed error information - error_type = type(e).__name__ + error_type: Final = type(e).__name__ # When the caller opted into raise_on_error it owns the exception and logs it at the # level that fits (an expected pass-through re-auth 401 is info, not an operator-actionable # error), so log at debug here to avoid an error-level line that would trip error-rate # alerts on that expected signal. The swallow path (raise_on_error=False) still logs at # error because nothing downstream will surface the failure. - _log = verbose_logger.debug if raise_on_error else verbose_logger.error + _log: Final = verbose_logger.debug if raise_on_error else verbose_logger.error _log( f"MCP client call_tool failed - " f"Error Type: {error_type}, " - f"Error: {str(e)}, " + f"Error: {e}, " f"Tool: {call_tool_request_params.name}, " f"Server: {self.server_url or 'stdio'}, " f"Transport: {self.transport_type}" @@ -622,32 +616,32 @@ class MCPClient: # Return a default error result instead of raising return self.error_tool_result(e) - async def list_prompts(self) -> List[Prompt]: + async def list_prompts(self) -> list[Prompt]: """List available prompts from the server.""" - verbose_logger.debug(f"MCP client listing tools from {self.server_url or 'stdio'}") + verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") async def _list_prompts_operation(session: ClientSession): return await session.list_prompts() try: - result = await self.run_with_session(_list_prompts_operation) - prompt_count = len(result.prompts) - prompt_names = [prompt.name for prompt in result.prompts] + result: Final = await self.run_with_session(_list_prompts_operation) + prompt_count: Final = len(result.prompts) + prompt_names: Final = [prompt.name for prompt in result.prompts] verbose_logger.info( - f"MCP client listed {prompt_count} tools from {self.server_url or 'stdio'}: {prompt_names}" + "MCP client listed %s tools from %s: %s", prompt_count, self.server_url or "stdio", prompt_names ) return result.prompts except asyncio.CancelledError: verbose_logger.warning("MCP client list_prompts was cancelled") raise except Exception as e: - error_type = type(e).__name__ + error_type: Final = type(e).__name__ verbose_logger.error( - f"MCP client list_prompts failed - " - f"Error Type: {error_type}, " - f"Error: {str(e)}, " - f"Server: {self.server_url or 'stdio'}, " - f"Transport: {self.transport_type}" + "MCP client list_prompts failed - Error Type: %s, Error: %s, Server: %s, Transport: %s", + error_type, + e, + self.server_url or "stdio", + self.transport_type, ) # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: @@ -660,7 +654,7 @@ class MCPClient: async def get_prompt(self, get_prompt_request_params: GetPromptRequestParams) -> GetPromptResult: """Fetch a prompt definition from the MCP server.""" - verbose_logger.info(f"MCP client fetching prompt '{get_prompt_request_params.name}'") + verbose_logger.info("MCP client fetching prompt '%s'", get_prompt_request_params.name) async def _get_prompt_operation(session: ClientSession): verbose_logger.debug("MCP client sending get_prompt request to session") @@ -670,8 +664,8 @@ class MCPClient: ) try: - get_prompt_result = await self.run_with_session(_get_prompt_operation) - verbose_logger.info(f"MCP client get_prompt '{get_prompt_request_params.name}' completed successfully") + get_prompt_result: Final = await self.run_with_session(_get_prompt_operation) + verbose_logger.info("MCP client get_prompt '%s' completed successfully", get_prompt_request_params.name) return get_prompt_result except asyncio.CancelledError: verbose_logger.warning("MCP client get_prompt was cancelled") @@ -679,17 +673,17 @@ class MCPClient: except Exception as e: import traceback - error_trace = traceback.format_exc() - verbose_logger.debug(f"MCP client get_prompt traceback:\n{error_trace}") + error_trace: Final = traceback.format_exc() + verbose_logger.debug("MCP client get_prompt traceback:\n%s", error_trace) # Log detailed error information - error_type = type(e).__name__ + error_type: Final = type(e).__name__ verbose_logger.error( - f"MCP client get_prompt failed - " - f"Error Type: {error_type}, " - f"Error: {str(e)}, " - f"Prompt: {get_prompt_request_params.name}, " - f"Server: {self.server_url or 'stdio'}, " - f"Transport: {self.transport_type}" + "MCP client get_prompt failed - Error Type: %s, Error: %s, Prompt: %s, Server: %s, Transport: %s", + error_type, + e, + get_prompt_request_params.name, + self.server_url or "stdio", + self.transport_type, ) # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: @@ -701,30 +695,30 @@ class MCPClient: async def list_resources(self) -> list[Resource]: """List available resources from the server.""" - verbose_logger.debug(f"MCP client listing resources from {self.server_url or 'stdio'}") + verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio") async def _list_resources_operation(session: ClientSession): return await session.list_resources() try: - result = await self.run_with_session(_list_resources_operation) - resource_count = len(result.resources) - resource_names = [resource.name for resource in result.resources] + result: Final = await self.run_with_session(_list_resources_operation) + resource_count: Final = len(result.resources) + resource_names: Final = [resource.name for resource in result.resources] verbose_logger.info( - f"MCP client listed {resource_count} resources from {self.server_url or 'stdio'}: {resource_names}" + "MCP client listed %s resources from %s: %s", resource_count, self.server_url or "stdio", resource_names ) return result.resources except asyncio.CancelledError: verbose_logger.warning("MCP client list_resources was cancelled") raise except Exception as e: - error_type = type(e).__name__ + error_type: Final = type(e).__name__ verbose_logger.error( - f"MCP client list_resources failed - " - f"Error Type: {error_type}, " - f"Error: {str(e)}, " - f"Server: {self.server_url or 'stdio'}, " - f"Transport: {self.transport_type}" + "MCP client list_resources failed - Error Type: %s, Error: %s, Server: %s, Transport: %s", + error_type, + e, + self.server_url or "stdio", + self.transport_type, ) # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: @@ -737,30 +731,33 @@ class MCPClient: async def list_resource_templates(self) -> list[ResourceTemplate]: """List available resource templates from the server.""" - verbose_logger.debug(f"MCP client listing resource templates from {self.server_url or 'stdio'}") + verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio") async def _list_resource_templates_operation(session: ClientSession): return await session.list_resource_templates() try: - result = await self.run_with_session(_list_resource_templates_operation) - resource_template_count = len(result.resourceTemplates) - resource_template_names = [resourceTemplate.name for resourceTemplate in result.resourceTemplates] + result: Final = await self.run_with_session(_list_resource_templates_operation) + resource_template_count: Final = len(result.resourceTemplates) + resource_template_names: Final = [resourceTemplate.name for resourceTemplate in result.resourceTemplates] verbose_logger.info( - f"MCP client listed {resource_template_count} resource templates from {self.server_url or 'stdio'}: {resource_template_names}" + "MCP client listed %s resource templates from %s: %s", + resource_template_count, + self.server_url or "stdio", + resource_template_names, ) return result.resourceTemplates except asyncio.CancelledError: verbose_logger.warning("MCP client list_resource_templates was cancelled") raise except Exception as e: - error_type = type(e).__name__ + error_type: Final = type(e).__name__ verbose_logger.error( - f"MCP client list_resource_templates failed - " - f"Error Type: {error_type}, " - f"Error: {str(e)}, " - f"Server: {self.server_url or 'stdio'}, " - f"Transport: {self.transport_type}" + "MCP client list_resource_templates failed - Error Type: %s, Error: %s, Server: %s, Transport: %s", + error_type, + e, + self.server_url or "stdio", + self.transport_type, ) # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: @@ -773,15 +770,15 @@ class MCPClient: async def read_resource(self, url: AnyUrl) -> ReadResourceResult: """Fetch resource contents from the MCP server.""" - verbose_logger.info(f"MCP client fetching resource '{url}'") + verbose_logger.info("MCP client fetching resource '%s'", url) async def _read_resource_operation(session: ClientSession): verbose_logger.debug("MCP client sending read_resource request to session") return await session.read_resource(url) try: - read_resource_result = await self.run_with_session(_read_resource_operation) - verbose_logger.info(f"MCP client read_resource '{url}' completed successfully") + read_resource_result: Final = await self.run_with_session(_read_resource_operation) + verbose_logger.info("MCP client read_resource '%s' completed successfully", url) return read_resource_result except asyncio.CancelledError: verbose_logger.warning("MCP client read_resource was cancelled") @@ -789,17 +786,17 @@ class MCPClient: except Exception as e: import traceback - error_trace = traceback.format_exc() - verbose_logger.debug(f"MCP client read_resource traceback:\n{error_trace}") + error_trace: Final = traceback.format_exc() + verbose_logger.debug("MCP client read_resource traceback:\n%s", error_trace) # Log detailed error information - error_type = type(e).__name__ + error_type: Final = type(e).__name__ verbose_logger.error( - f"MCP client read_resource failed - " - f"Error Type: {error_type}, " - f"Error: {str(e)}, " - f"Url: {url}, " - f"Server: {self.server_url or 'stdio'}, " - f"Transport: {self.transport_type}" + "MCP client read_resource failed - Error Type: %s, Error: %s, Url: %s, Server: %s, Transport: %s", + error_type, + e, + url, + self.server_url or "stdio", + self.transport_type, ) # Check if it's a stream/connection error if "BrokenResourceError" in error_type or "Broken" in error_type: diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 500d226752b..30d50e2a74b 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -1,5 +1,5 @@ import json -from typing import Dict, List, Literal, Union +from typing import Final, Literal from mcp import ClientSession from mcp.types import CallToolRequestParams as MCPCallToolRequestParams @@ -18,7 +18,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall ######################################################## def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam: """Convert an MCP tool to an OpenAI tool.""" - normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema) + normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.inputSchema) return ChatCompletionToolParam( type="function", @@ -44,7 +44,7 @@ def _normalize_mcp_input_schema(input_schema: dict) -> dict: return {"type": "object", "properties": {}, "additionalProperties": False} # Make a copy to avoid modifying the original - normalized_schema = dict(input_schema) + normalized_schema: Final = dict(input_schema) # Ensure type is 'object' if "type" not in normalized_schema: @@ -65,7 +65,7 @@ def transform_mcp_tool_to_openai_responses_api_tool( mcp_tool: MCPTool, ) -> FunctionToolParam: """Convert an MCP tool to an OpenAI Responses API tool.""" - normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema) + normalized_parameters: Final = _normalize_mcp_input_schema(mcp_tool.inputSchema) return FunctionToolParam( name=mcp_tool.name, @@ -92,7 +92,7 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages async def load_mcp_tools( session: ClientSession, format: Literal["mcp", "openai"] = "mcp" -) -> Union[List[MCPTool], List[ChatCompletionToolParam]]: +) -> list[MCPTool] | list[ChatCompletionToolParam]: """ Load all available MCP tools @@ -103,7 +103,7 @@ async def load_mcp_tools( If format is set to "openai", the tools are converted to OpenAI API compatible tools. """ - tools = await session.list_tools() + tools: Final = await session.list_tools() if format == "openai": return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools] return tools.tools @@ -119,7 +119,7 @@ async def call_mcp_tool( call_tool_request_params: MCPCallToolRequestParams, ) -> MCPCallToolResult: """Call an MCP tool.""" - tool_result = await session.call_tool( + tool_result: Final = await session.call_tool( name=call_tool_request_params.name, arguments=call_tool_request_params.arguments, ) @@ -138,10 +138,10 @@ def _get_function_arguments(function: FunctionDefinition) -> dict: def transform_openai_tool_call_request_to_mcp_tool_call_request( - openai_tool: Union[ChatCompletionMessageToolCall, Dict], + openai_tool: ChatCompletionMessageToolCall | dict, ) -> MCPCallToolRequestParams: """Convert an OpenAI ChatCompletionMessageToolCall to an MCP CallToolRequestParams.""" - function = openai_tool["function"] + function: Final = openai_tool["function"] return MCPCallToolRequestParams( name=function["name"], arguments=_get_function_arguments(function), @@ -161,7 +161,7 @@ async def call_openai_tool( Returns: The result of the MCP tool call. """ - mcp_tool_call_request_params = transform_openai_tool_call_request_to_mcp_tool_call_request( + mcp_tool_call_request_params: Final = transform_openai_tool_call_request_to_mcp_tool_call_request( openai_tool=openai_tool, ) return await call_mcp_tool( diff --git a/litellm/files/main.py b/litellm/files/main.py index 3b359b55fe3..e137c7587c0 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -9,9 +9,10 @@ import asyncio import contextvars import time import uuid as uuid_module +from collections.abc import Coroutine from functools import partial from types import MappingProxyType -from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast +from typing import Any, Final, Literal, cast import httpx @@ -69,7 +70,7 @@ base_llm_http_handler = BaseLLMHTTPHandler() def _should_sdk_support_streaming( - custom_llm_provider: Optional[Union[FileContentProvider, str]], + custom_llm_provider: FileContentProvider | str | None, ) -> bool: """ Return whether file content streaming is supported for the provider. @@ -77,17 +78,17 @@ def _should_sdk_support_streaming( return custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS -openai_files_instance = OpenAIFilesAPI() -azure_files_instance = AzureOpenAIFilesAPI() -vertex_ai_files_instance = VertexAIFilesHandler() -bedrock_files_instance = BedrockFilesHandler() +openai_files_instance: Final = OpenAIFilesAPI() +azure_files_instance: Final = AzureOpenAIFilesAPI() +vertex_ai_files_instance: Final = VertexAIFilesHandler() +bedrock_files_instance: Final = BedrockFilesHandler() ################################################# def _add_trusted_model_credentials_to_litellm_params( - litellm_params_dict: Dict[str, Any], kwargs: Dict[str, Any] + litellm_params_dict: dict[str, Any], kwargs: dict[str, Any] ) -> None: - trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials") + trusted_model_credentials: Final = kwargs.get("_litellm_internal_model_credentials") if isinstance(trusted_model_credentials, type(MappingProxyType({}))): litellm_params_dict["_litellm_internal_model_credentials"] = trusted_model_credentials @@ -96,10 +97,10 @@ def _add_trusted_model_credentials_to_litellm_params( async def acreate_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune", "messages"], - expires_after: Optional[FileExpiresAfter] = None, + expires_after: FileExpiresAfter | None = None, custom_llm_provider: FileCreateProvider = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> OpenAIFileObject: """ @@ -108,10 +109,10 @@ async def acreate_file( LiteLLM Equivalent of POST: POST https://api.openai.com/v1/files """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acreate_file"] = True - call_args = { + call_args: Final = { "file": file, "purpose": purpose, "expires_after": expires_after, @@ -122,11 +123,11 @@ async def acreate_file( } # Use a partial function to pass your keyword arguments - func = partial(create_file, **call_args) + func: Final = partial(create_file, **call_args) # 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: @@ -141,12 +142,12 @@ async def acreate_file( def create_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune", "messages"], - expires_after: Optional[FileExpiresAfter] = None, - custom_llm_provider: Optional[FileCreateProvider] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + expires_after: FileExpiresAfter | None = None, + custom_llm_provider: FileCreateProvider | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, -) -> Union[OpenAIFileObject, Coroutine[Any, Any, OpenAIFileObject]]: +) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]: """ Files are used to upload documents that can be used with features like Assistants, Fine-tuning, and Batch API. @@ -155,13 +156,13 @@ def create_file( Specify either provider_list or custom_llm_provider. """ try: - _is_async = kwargs.pop("acreate_file", False) is True - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = dict(**kwargs) - logging_obj = cast(Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")) + _is_async: Final = kwargs.pop("acreate_file", False) is True + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params_dict: Final = dict(**kwargs) + logging_obj: Final = cast(LiteLLMLoggingObj | None, kwargs.get("litellm_logging_obj")) if logging_obj is None: raise ValueError("logging_obj is required") - client = kwargs.get("client") + client: Final = kwargs.get("client") ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 @@ -172,7 +173,7 @@ def create_file( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(cast(str, 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 @@ -195,7 +196,7 @@ def create_file( extra_body=extra_body, ) - provider_config = ProviderConfigManager.get_provider_files_config( + provider_config: Final = ProviderConfigManager.get_provider_files_config( model="", provider=LlmProviders(custom_llm_provider), ) @@ -213,7 +214,7 @@ def create_file( timeout=timeout, ) elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = get_openai_credentials( + openai_creds: Final = get_openai_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, organization=optional_params.organization, @@ -228,7 +229,7 @@ def create_file( create_file_data=_create_file_request, ) elif custom_llm_provider == "azure": - azure_creds = get_azure_credentials( + azure_creds: Final = get_azure_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, api_version=optional_params.api_version, @@ -245,9 +246,7 @@ def create_file( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus', 'anthropic'] are supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus', 'anthropic'] are supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -265,8 +264,8 @@ def create_file( async def afile_retrieve( file_id: str, custom_llm_provider: FileRetrieveProvider = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> OpenAIFileObject: """ @@ -275,11 +274,11 @@ async def afile_retrieve( LiteLLM Equivalent of GET https://api.openai.com/v1/files """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["is_async"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( file_retrieve, file_id, custom_llm_provider, @@ -289,9 +288,9 @@ async def afile_retrieve( ) # 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: @@ -306,8 +305,8 @@ async def afile_retrieve( def file_retrieve( file_id: str, custom_llm_provider: FileRetrieveProvider = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> FileObject: """ @@ -316,7 +315,7 @@ def file_retrieve( LiteLLM Equivalent of POST: POST https://api.openai.com/v1/files """ try: - optional_params = GenericLiteLLMParams(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -326,17 +325,17 @@ def file_retrieve( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _is_async = kwargs.pop("is_async", False) is True + _is_async: Final = kwargs.pop("is_async", False) is True if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = get_openai_credentials( + openai_creds: Final = get_openai_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, organization=optional_params.organization, @@ -351,7 +350,7 @@ def file_retrieve( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = get_azure_credentials( + azure_creds: Final = get_azure_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, api_version=optional_params.api_version, @@ -367,12 +366,12 @@ def file_retrieve( ) else: # Try using provider config pattern (for Manus, Bedrock, etc.) - provider_config = ProviderConfigManager.get_provider_files_config( + provider_config: Final = ProviderConfigManager.get_provider_files_config( model="", provider=LlmProviders(custom_llm_provider), ) if provider_config is not None: - litellm_params_dict = get_litellm_params(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) _add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, @@ -396,7 +395,7 @@ def file_retrieve( function_id=str(kwargs.get("id") or ""), ) - client = kwargs.get("client") + client: Final = kwargs.get("client") response = base_llm_http_handler.retrieve_file( file_id=file_id, provider_config=provider_config, @@ -411,9 +410,7 @@ def file_retrieve( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_retrieve'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -436,8 +433,8 @@ def file_retrieve( async def afile_delete( file_id: str, custom_llm_provider: FileDeleteProvider = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> Coroutine[Any, Any, FileObject]: """ @@ -446,12 +443,12 @@ async def afile_delete( LiteLLM Equivalent of DELETE https://api.openai.com/v1/files """ try: - loop = asyncio.get_event_loop() - model = kwargs.pop("model", None) + loop: Final = asyncio.get_event_loop() + model: Final = kwargs.pop("model", None) kwargs["is_async"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( file_delete, file_id, model, @@ -462,9 +459,9 @@ async def afile_delete( ) # 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: @@ -478,10 +475,10 @@ async def afile_delete( @client def file_delete( file_id: str, - model: Optional[str] = None, - custom_llm_provider: Union[FileDeleteProvider, str] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + model: str | None = None, + custom_llm_provider: FileDeleteProvider | str = "openai", + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> FileDeleted: """ @@ -495,8 +492,8 @@ def file_delete( _, custom_llm_provider, _, _ = get_llm_provider(model, custom_llm_provider) except Exception: pass - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) _add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, @@ -504,22 +501,22 @@ def file_delete( ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default - client = kwargs.get("client") + client: Final = kwargs.get("client") if ( timeout is not None and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _is_async = kwargs.pop("is_async", False) is True + _is_async: Final = kwargs.pop("is_async", False) is True if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = get_openai_credentials( + openai_creds: Final = get_openai_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, organization=optional_params.organization, @@ -534,7 +531,7 @@ def file_delete( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = get_azure_credentials( + azure_creds: Final = get_azure_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, api_version=optional_params.api_version, @@ -552,7 +549,7 @@ def file_delete( ) else: # Try using provider config pattern (for Manus, Bedrock, etc.) - provider_config = ProviderConfigManager.get_provider_files_config( + provider_config: Final = ProviderConfigManager.get_provider_files_config( model="", provider=LlmProviders(custom_llm_provider), ) @@ -590,9 +587,7 @@ def file_delete( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', 'manus', and 'anthropic' are supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_delete'. Only 'openai', 'azure', 'gemini', 'manus', and 'anthropic' are supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -613,9 +608,9 @@ def file_delete( @client async def afile_list( custom_llm_provider: FileListProvider = "openai", - purpose: Optional[str] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + purpose: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ): """ @@ -624,11 +619,11 @@ async def afile_list( LiteLLM Equivalent of GET https://api.openai.com/v1/files """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["is_async"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( file_list, custom_llm_provider, purpose, @@ -638,9 +633,9 @@ async def afile_list( ) # 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: @@ -654,9 +649,9 @@ async def afile_list( @client def file_list( custom_llm_provider: FileListProvider = "openai", - purpose: Optional[str] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + purpose: str | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ): """ @@ -665,7 +660,7 @@ def file_list( LiteLLM Equivalent of GET https://api.openai.com/v1/files """ try: - optional_params = GenericLiteLLMParams(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -675,22 +670,22 @@ def file_list( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _is_async = kwargs.pop("is_async", False) is True + _is_async: Final = kwargs.pop("is_async", False) is True # Check if provider has a custom files config (e.g., Manus, Bedrock, Vertex AI) - provider_config = ProviderConfigManager.get_provider_files_config( + provider_config: Final = ProviderConfigManager.get_provider_files_config( model="", provider=LlmProviders(custom_llm_provider), ) if provider_config is not None: - litellm_params_dict = get_litellm_params(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) litellm_params_dict["api_key"] = optional_params.api_key litellm_params_dict["api_base"] = optional_params.api_base @@ -710,7 +705,7 @@ def file_list( function_id=str(kwargs.get("id", "")), ) - client = kwargs.get("client") + client: Final = kwargs.get("client") response = base_llm_http_handler.list_files( purpose=purpose, provider_config=provider_config, @@ -723,7 +718,7 @@ def file_list( ) return response elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = get_openai_credentials( + openai_creds: Final = get_openai_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, organization=optional_params.organization, @@ -738,7 +733,7 @@ def file_list( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = get_azure_credentials( + azure_creds: Final = get_azure_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, api_version=optional_params.api_version, @@ -754,9 +749,7 @@ def file_list( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_list'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -774,24 +767,24 @@ def file_list( async def afile_content( file_id: str, custom_llm_provider: FileContentProvider = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, chunk_size: int = 1024 * 1024, stream: bool = False, **kwargs, -) -> Union[HttpxBinaryResponseContent, FileContentStreamingResult]: +) -> HttpxBinaryResponseContent | FileContentStreamingResult: """ Async: Get file contents LiteLLM Equivalent of GET https://api.openai.com/v1/files """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["afile_content"] = True - model = kwargs.pop("model", None) + model: Final = kwargs.pop("model", None) # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( file_content, file_id=file_id, model=model, @@ -804,9 +797,9 @@ async def afile_content( ) # 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: @@ -820,34 +813,34 @@ async def afile_content( @client def file_content( file_id: str, - model: Optional[str] = None, - custom_llm_provider: Optional[Union[FileContentProvider, str]] = None, - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + model: str | None = None, + custom_llm_provider: FileContentProvider | str | None = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, chunk_size: int = 1024 * 1024, stream: bool = False, **kwargs, -) -> Union[ - HttpxBinaryResponseContent, - FileContentStreamingResult, - Coroutine[Any, Any, HttpxBinaryResponseContent], - Coroutine[Any, Any, FileContentStreamingResult], -]: +) -> ( + HttpxBinaryResponseContent + | FileContentStreamingResult + | Coroutine[Any, Any, HttpxBinaryResponseContent] + | Coroutine[Any, Any, FileContentStreamingResult] +): """ Returns the contents of the specified file. LiteLLM Equivalent of POST: POST https://api.openai.com/v1/files """ try: - optional_params = GenericLiteLLMParams(**kwargs) - litellm_params_dict = get_litellm_params(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) _add_trusted_model_credentials_to_litellm_params( litellm_params_dict=litellm_params_dict, kwargs=kwargs, ) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 - client = kwargs.get("client") + client: Final = kwargs.get("client") # set timeout for 10 minutes by default try: @@ -861,20 +854,20 @@ def file_content( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(cast(str, custom_llm_provider)) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _file_content_request = FileContentRequest( + _file_content_request: Final = FileContentRequest( file_id=file_id, extra_headers=extra_headers, extra_body=extra_body, ) - _is_async = kwargs.pop("afile_content", False) is True + _is_async: Final = kwargs.pop("afile_content", False) is True if stream and _should_sdk_support_streaming(custom_llm_provider): return file_content_streaming( @@ -886,13 +879,13 @@ def file_content( chunk_size=chunk_size, optional_params=optional_params, timeout=timeout, - logging_obj=cast(Optional[LiteLLMLoggingObj], kwargs.get("litellm_logging_obj")), + logging_obj=cast(LiteLLMLoggingObj | None, kwargs.get("litellm_logging_obj")), _is_async=_is_async, client=client, ) # Check if provider has a custom files config (e.g., Anthropic, Manus) - provider_config = ProviderConfigManager.get_provider_files_config( + provider_config: Final = ProviderConfigManager.get_provider_files_config( model="", provider=LlmProviders(custom_llm_provider), ) @@ -925,7 +918,7 @@ def file_content( return response if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = get_openai_credentials( + openai_creds: Final = get_openai_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, organization=optional_params.organization, @@ -940,7 +933,7 @@ def file_content( organization=openai_creds.organization, ) elif custom_llm_provider == "azure": - azure_creds = get_azure_credentials( + azure_creds: Final = get_azure_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, api_version=optional_params.api_version, @@ -957,14 +950,14 @@ def file_content( litellm_params=litellm_params_dict, ) elif custom_llm_provider == "vertex_ai": - api_base = optional_params.api_base or "" - vertex_ai_project = ( + api_base: Final = optional_params.api_base or "" + 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_files_instance.file_content( _is_async=_is_async, @@ -988,9 +981,7 @@ def file_content( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -1007,23 +998,23 @@ def file_content( def file_content_streaming( *, file_id: str, - model: Optional[str], - custom_llm_provider: Optional[Union[FileContentProvider, str]], - extra_headers: Optional[Dict[str, str]], - extra_body: Optional[Dict[str, str]], + model: str | None, + custom_llm_provider: FileContentProvider | str | None, + extra_headers: dict[str, str] | None, + extra_body: dict[str, str] | None, chunk_size: int, optional_params: GenericLiteLLMParams, - timeout: Union[float, httpx.Timeout], - logging_obj: Optional[LiteLLMLoggingObj], + timeout: float | httpx.Timeout, + logging_obj: LiteLLMLoggingObj | None, _is_async: bool, - client: Optional[Any], -) -> Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]]: + client: Any | None, +) -> FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult]: if logging_obj is not None: logging_obj.model = model or "" logging_obj.model_call_details["model"] = model or "" logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider - litellm_params = logging_obj.model_call_details.get("litellm_params", {}) or {} + litellm_params: Final = logging_obj.model_call_details.get("litellm_params", {}) or {} if optional_params.api_base is not None: litellm_params["api_base"] = optional_params.api_base logging_obj.model_call_details["litellm_params"] = litellm_params @@ -1042,11 +1033,11 @@ def file_content_streaming( headers=response.headers, ) - response: Union[FileContentStreamingResult, Coroutine[Any, Any, FileContentStreamingResult]] = ( - FileContentStreamingResult(stream_iterator=iter(()), headers={}) + response: FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult] = FileContentStreamingResult( + stream_iterator=iter(()), headers={} ) if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: - openai_creds = get_openai_credentials( + openai_creds: Final = get_openai_credentials( api_base=optional_params.api_base, api_key=optional_params.api_key, organization=optional_params.organization, @@ -1068,10 +1059,7 @@ def file_content_streaming( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for streaming 'file_content'. Supported providers are {}.".format( - custom_llm_provider, - sorted(OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS), - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for streaming 'file_content'. Supported providers are {sorted(OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS)}.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( diff --git a/litellm/files/streaming.py b/litellm/files/streaming.py index 6d84f73dcfe..d9df05e7135 100644 --- a/litellm/files/streaming.py +++ b/litellm/files/streaming.py @@ -1,17 +1,10 @@ import datetime import traceback -from typing import ( - TYPE_CHECKING, - Any, - AsyncIterator, - Dict, - Iterator, - Optional, - Union, - cast, -) +from collections.abc import AsyncIterator, Iterator +from typing import TYPE_CHECKING, Any, Final, Optional, cast import anyio + from litellm.files.types import FileContentProvider if TYPE_CHECKING: @@ -29,10 +22,10 @@ class FileContentStreamingResponse: def __init__( self, - stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]], + stream_iterator: Iterator[bytes] | AsyncIterator[bytes], file_id: str, - model: Optional[str], - custom_llm_provider: Optional[Union[FileContentProvider, str]], + model: str | None, + custom_llm_provider: FileContentProvider | str | None, logging_obj: Optional["LiteLLMLoggingObj"], ) -> None: self.stream_iterator = stream_iterator @@ -40,8 +33,8 @@ class FileContentStreamingResponse: self.model = model self.custom_llm_provider = custom_llm_provider self.logging_obj = logging_obj - self.standard_logging_object: Optional["StandardLoggingPayload"] = None - self._hidden_params: Dict[str, Any] = {} + self.standard_logging_object: StandardLoggingPayload | None = None + self._hidden_params: dict[str, Any] = {} self._logging_completed = False self._close_completed = False self._start_time = ( @@ -93,8 +86,8 @@ class FileContentStreamingResponse: self._close_completed = True self._logging_completed = True - stream_to_close = self.stream_iterator - self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(())) + stream_to_close: Final = self.stream_iterator + self.stream_iterator = cast(Iterator[bytes] | AsyncIterator[bytes], iter(())) # Shield cleanup from request cancellation so upstream HTTP connections # are released promptly on client disconnects. @@ -102,7 +95,7 @@ class FileContentStreamingResponse: if hasattr(stream_to_close, "aclose"): await cast(AsyncIterator[bytes], stream_to_close).aclose() # type: ignore[attr-defined] elif hasattr(stream_to_close, "close"): - result = cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined] + result: Final = cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined] if result is not None: await result @@ -112,14 +105,14 @@ class FileContentStreamingResponse: self._close_completed = True self._logging_completed = True - stream_to_close = self.stream_iterator - self.stream_iterator = cast(Union[Iterator[bytes], AsyncIterator[bytes]], iter(())) + stream_to_close: Final = self.stream_iterator + self.stream_iterator = cast(Iterator[bytes] | AsyncIterator[bytes], iter(())) if hasattr(stream_to_close, "close"): cast(Iterator[bytes], stream_to_close).close() # type: ignore[attr-defined] - def _build_logging_response(self) -> Dict[str, str]: - response = { + def _build_logging_response(self) -> dict[str, str]: + response: Final = { "id": self.file_id, "object": "file.content", } @@ -156,7 +149,7 @@ class FileContentStreamingResponse: ) self._sync_hidden_params() - payload = get_standard_logging_object_payload( + payload: Final = get_standard_logging_object_payload( kwargs=self.logging_obj.model_call_details, init_response_obj=self._build_logging_response(), start_time=self._start_time, @@ -167,10 +160,10 @@ class FileContentStreamingResponse: if payload is None: return None - merged_hidden_params = cast( + merged_hidden_params: Final = cast( "StandardLoggingHiddenParams", { - **cast(Dict[str, Any], payload.get("hidden_params") or {}), + **cast(dict[str, Any], payload.get("hidden_params") or {}), **self._hidden_params, }, ) @@ -191,8 +184,8 @@ class FileContentStreamingResponse: return self._logging_completed = True - end_time = datetime.datetime.now() - standard_logging_object = self._build_standard_logging_object(end_time=end_time) + end_time: Final = datetime.datetime.now() + standard_logging_object: Final = self._build_standard_logging_object(end_time=end_time) await self.logging_obj.async_success_handler( result=self._build_logging_response(), start_time=self._start_time, @@ -210,8 +203,8 @@ class FileContentStreamingResponse: return self._logging_completed = True - end_time = datetime.datetime.now() - standard_logging_object = self._build_standard_logging_object(end_time=end_time) + end_time: Final = datetime.datetime.now() + standard_logging_object: Final = self._build_standard_logging_object(end_time=end_time) self.logging_obj.success_handler( result=self._build_logging_response(), start_time=self._start_time, @@ -224,8 +217,8 @@ class FileContentStreamingResponse: return self._logging_completed = True - end_time = datetime.datetime.now() - traceback_str = traceback.format_exc() + end_time: Final = datetime.datetime.now() + traceback_str: Final = traceback.format_exc() self.logging_obj.failure_handler(error, traceback_str, self._start_time, end_time) await self.logging_obj.async_failure_handler(error, traceback_str, self._start_time, end_time) @@ -234,5 +227,5 @@ class FileContentStreamingResponse: return self._logging_completed = True - end_time = datetime.datetime.now() + end_time: Final = datetime.datetime.now() self.logging_obj.failure_handler(error, traceback.format_exc(), self._start_time, end_time) diff --git a/litellm/files/types.py b/litellm/files/types.py index 6bf7b1a1cc2..8cadd69f024 100644 --- a/litellm/files/types.py +++ b/litellm/files/types.py @@ -1,8 +1,9 @@ -from typing import AsyncIterator, Dict, Iterator, Literal, NamedTuple, Union +from collections.abc import AsyncIterator, Iterator +from typing import Literal, NamedTuple FileContentProvider = Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"] class FileContentStreamingResult(NamedTuple): - stream_iterator: Union[Iterator[bytes], AsyncIterator[bytes]] - headers: Dict[str, str] + stream_iterator: Iterator[bytes] | AsyncIterator[bytes] + headers: dict[str, str] diff --git a/litellm/files/utils.py b/litellm/files/utils.py index 3ee4953bfef..f470931115f 100644 --- a/litellm/files/utils.py +++ b/litellm/files/utils.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Final from litellm.types.llms.openai import CreateFileRequest from litellm.types.utils import ExtractedFileData @@ -8,7 +8,7 @@ from litellm.types.utils import ExtractedFileData # batch file must not silently bypass the streaming path just because of its # declared type. ``purpose == "batch"`` is the authoritative signal; non-JSONL # content still fails loudly when the rows are parsed. -_BATCH_JSONL_CONTENT_TYPES = frozenset( +_BATCH_JSONL_CONTENT_TYPES: Final = frozenset( { "application/jsonl", "application/json", @@ -37,7 +37,7 @@ class FilesAPIUtils: ) @staticmethod - def is_batch_jsonl_request(create_file_data: CreateFileRequest, content_type: Optional[str]) -> bool: + def is_batch_jsonl_request(create_file_data: CreateFileRequest, content_type: str | None) -> bool: """ Batch-jsonl check from metadata only, so the body can stay a streamable Path/handle instead of being read into memory. @@ -49,7 +49,7 @@ class FilesAPIUtils: ) @staticmethod - def valid_content_type(content_type: Optional[str]) -> bool: + def valid_content_type(content_type: str | None) -> bool: """ Whether the upload's MIME type is one a batch JSONL file is plausibly sent as (see ``_BATCH_JSONL_CONTENT_TYPES``). diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index ce5074cdaf5..e89defedabe 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -11,8 +11,9 @@ https://platform.openai.com/docs/api-reference/fine-tuning import asyncio import contextvars import os +from collections.abc import Coroutine from functools import partial -from typing import Any, Coroutine, Dict, Literal, Optional, Union +from typing import Any, Final, Literal import httpx @@ -28,17 +29,17 @@ from litellm.types.utils import LiteLLMFineTuningJob from litellm.utils import client, supports_httpx_timeout ####### ENVIRONMENT VARIABLES ################### -openai_fine_tuning_apis_instance = OpenAIFineTuningAPI() -azure_fine_tuning_apis_instance = AzureOpenAIFineTuningAPI() -vertex_fine_tuning_apis_instance = VertexFineTuningAPI() +openai_fine_tuning_apis_instance: Final = OpenAIFineTuningAPI() +azure_fine_tuning_apis_instance: Final = AzureOpenAIFineTuningAPI() +vertex_fine_tuning_apis_instance: Final = VertexFineTuningAPI() ################################################# def _prepare_azure_extra_body( - extra_body: Optional[Dict[str, Any]], - kwargs: Dict[str, Any], - azure_specific_hyperparams: Dict[str, Any], -) -> Dict[str, Any]: + extra_body: dict[str, Any] | None, + kwargs: dict[str, Any], + azure_specific_hyperparams: dict[str, Any], +) -> dict[str, Any]: """ Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters. @@ -60,7 +61,7 @@ def _prepare_azure_extra_body( extra_body = {} # Azure-specific root-level parameters - azure_specific_params = ["trainingType"] + azure_specific_params: Final = ["trainingType"] for param in azure_specific_params: if param in kwargs: extra_body[param] = kwargs[param] @@ -76,14 +77,14 @@ def _prepare_azure_extra_body( async def acreate_fine_tuning_job( model: str, training_file: str, - hyperparameters: Optional[dict] = {}, - suffix: Optional[str] = None, - validation_file: Optional[str] = None, - integrations: Optional[List[str]] = None, - seed: Optional[int] = None, + hyperparameters: dict | None = {}, + suffix: str | None = None, + validation_file: str | None = None, + integrations: List[str] | None = None, + seed: int | None = None, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> LiteLLMFineTuningJob: """ @@ -92,11 +93,11 @@ async def acreate_fine_tuning_job( """ verbose_logger.debug("inside acreate_fine_tuning_job model=%s and kwargs=%s", model, kwargs) try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acreate_fine_tuning_job"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( create_fine_tuning_job, model, training_file, @@ -112,9 +113,9 @@ async def acreate_fine_tuning_job( ) # 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: @@ -139,7 +140,7 @@ def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, v def _resolve_fine_tuning_timeout( timeout: Any, custom_llm_provider: str, -) -> Union[float, httpx.Timeout]: +) -> float | httpx.Timeout: """Normalise a raw timeout value to a float (seconds) or httpx.Timeout for fine-tuning calls.""" timeout = timeout or 600.0 if isinstance(timeout, httpx.Timeout): @@ -153,16 +154,16 @@ def _resolve_fine_tuning_timeout( def create_fine_tuning_job( model: str, training_file: str, - hyperparameters: Optional[dict] = {}, - suffix: Optional[str] = None, - validation_file: Optional[str] = None, - integrations: Optional[List[str]] = None, - seed: Optional[int] = None, + hyperparameters: dict | None = {}, + suffix: str | None = None, + validation_file: str | None = None, + integrations: List[str] | None = None, + seed: int | None = None, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, -) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: +) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: """ Creates a fine-tuning job which begins the process of creating a new model from a given dataset. @@ -170,24 +171,24 @@ def create_fine_tuning_job( """ try: - _is_async = kwargs.pop("acreate_fine_tuning_job", False) is True - optional_params = GenericLiteLLMParams(**kwargs) + _is_async: Final = kwargs.pop("acreate_fine_tuning_job", False) is True + optional_params: Final = GenericLiteLLMParams(**kwargs) # handle hyperparameters hyperparameters = hyperparameters or {} # original hyperparameters # For Azure, extract Azure-specific hyperparameters before creating OpenAI-spec hyperparameters - azure_specific_hyperparams = {} + azure_specific_hyperparams: Final = {} if custom_llm_provider == "azure": - azure_hyperparameter_keys = ["prompt_loss_weight"] + azure_hyperparameter_keys: Final = ["prompt_loss_weight"] for key in azure_hyperparameter_keys: if key in hyperparameters: azure_specific_hyperparams[key] = hyperparameters.pop(key) - _oai_hyperparameters: Hyperparameters = Hyperparameters( + _oai_hyperparameters: Final[Hyperparameters] = Hyperparameters( **hyperparameters ) # Typed Hyperparameters for OpenAI Spec - timeout = _resolve_fine_tuning_timeout( + timeout: Final = _resolve_fine_tuning_timeout( optional_params.timeout or kwargs.get("request_timeout", 600), custom_llm_provider, ) @@ -202,7 +203,7 @@ def create_fine_tuning_job( 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) @@ -286,13 +287,13 @@ def create_fine_tuning_job( ) 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_fine_tuning_apis_instance.create_fine_tuning_job( _is_async=_is_async, create_fine_tuning_job_data=_build_fine_tuning_job_data( @@ -314,9 +315,7 @@ def create_fine_tuning_job( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_batch'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -335,19 +334,19 @@ def create_fine_tuning_job( async def acancel_fine_tuning_job( fine_tuning_job_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> LiteLLMFineTuningJob: """ Async: Immediately cancel a fine-tune job. """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["acancel_fine_tuning_job"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( cancel_fine_tuning_job, fine_tuning_job_id, custom_llm_provider, @@ -357,9 +356,9 @@ async def acancel_fine_tuning_job( ) # 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: @@ -373,10 +372,10 @@ async def acancel_fine_tuning_job( def cancel_fine_tuning_job( fine_tuning_job_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, -) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: +) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: """ Immediately cancel a fine-tune job. @@ -384,7 +383,7 @@ def cancel_fine_tuning_job( """ try: - optional_params = GenericLiteLLMParams(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -394,14 +393,14 @@ def cancel_fine_tuning_job( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _is_async = kwargs.pop("acancel_fine_tuning_job", False) is True + _is_async: Final = kwargs.pop("acancel_fine_tuning_job", False) is True # OpenAI if custom_llm_provider == "openai": @@ -413,7 +412,7 @@ def cancel_fine_tuning_job( 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) @@ -468,9 +467,7 @@ def cancel_fine_tuning_job( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_batch'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -485,22 +482,22 @@ def cancel_fine_tuning_job( async def alist_fine_tuning_jobs( - after: Optional[str] = None, - limit: Optional[int] = None, + after: str | None = None, + limit: int | None = None, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ): """ Async: List your organization's fine-tuning jobs """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["alist_fine_tuning_jobs"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( list_fine_tuning_jobs, after, limit, @@ -511,9 +508,9 @@ async def alist_fine_tuning_jobs( ) # 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: @@ -524,11 +521,11 @@ async def alist_fine_tuning_jobs( def list_fine_tuning_jobs( - after: Optional[str] = None, - limit: Optional[int] = None, + after: str | None = None, + limit: int | None = None, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ): """ @@ -540,7 +537,7 @@ def list_fine_tuning_jobs( - limit: Optional[int] = None, Number of fine-tuning jobs to retrieve. Defaults to 20 """ try: - optional_params = GenericLiteLLMParams(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -550,14 +547,14 @@ def list_fine_tuning_jobs( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _is_async = kwargs.pop("alist_fine_tuning_jobs", False) is True + _is_async: Final = kwargs.pop("alist_fine_tuning_jobs", False) is True # OpenAI if custom_llm_provider == "openai": @@ -569,7 +566,7 @@ def list_fine_tuning_jobs( 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) @@ -626,9 +623,7 @@ def list_fine_tuning_jobs( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'create_batch'. Only 'openai' is supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( @@ -646,19 +641,19 @@ def list_fine_tuning_jobs( async def aretrieve_fine_tuning_job( fine_tuning_job_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, ) -> LiteLLMFineTuningJob: """ Async: Get info about a fine-tuning job. """ try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["aretrieve_fine_tuning_job"] = True # Use a partial function to pass your keyword arguments - func = partial( + func: Final = partial( retrieve_fine_tuning_job, fine_tuning_job_id, custom_llm_provider, @@ -668,9 +663,9 @@ async def aretrieve_fine_tuning_job( ) # 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: @@ -684,15 +679,15 @@ async def aretrieve_fine_tuning_job( def retrieve_fine_tuning_job( fine_tuning_job_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - extra_headers: Optional[Dict[str, str]] = None, - extra_body: Optional[Dict[str, str]] = None, + extra_headers: dict[str, str] | None = None, + extra_body: dict[str, str] | None = None, **kwargs, -) -> Union[LiteLLMFineTuningJob, Coroutine[Any, Any, LiteLLMFineTuningJob]]: +) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]: """ Get info about a fine-tuning job. """ try: - optional_params = GenericLiteLLMParams(**kwargs) + optional_params: Final = GenericLiteLLMParams(**kwargs) ### TIMEOUT LOGIC ### timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600 # set timeout for 10 minutes by default @@ -702,14 +697,14 @@ def retrieve_fine_tuning_job( and isinstance(timeout, httpx.Timeout) and supports_httpx_timeout(custom_llm_provider) is False ): - read_timeout = timeout.read or 600 + read_timeout: Final = timeout.read or 600 timeout = read_timeout # default 10 min timeout elif timeout is not None and not isinstance(timeout, httpx.Timeout): timeout = float(timeout) # type: ignore elif timeout is None: timeout = 600.0 - _is_async = kwargs.pop("aretrieve_fine_tuning_job", False) is True + _is_async: Final = kwargs.pop("aretrieve_fine_tuning_job", False) is True # OpenAI if custom_llm_provider == "openai": @@ -720,7 +715,7 @@ def retrieve_fine_tuning_job( 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") @@ -766,9 +761,7 @@ def retrieve_fine_tuning_job( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'retrieve_fine_tuning_job'. Only 'openai' and 'azure' are supported.".format( - custom_llm_provider - ), + message=f"LiteLLM doesn't support {custom_llm_provider} for 'retrieve_fine_tuning_job'. Only 'openai' and 'azure' are supported.", model="n/a", llm_provider=custom_llm_provider, response=httpx.Response( diff --git a/litellm/google_genai/__init__.py b/litellm/google_genai/__init__.py index ca7b547c440..eeff6a5fd65 100644 --- a/litellm/google_genai/__init__.py +++ b/litellm/google_genai/__init__.py @@ -12,8 +12,8 @@ from .main import ( ) __all__ = [ - "generate_content", "agenerate_content", - "generate_content_stream", "agenerate_content_stream", + "generate_content", + "generate_content_stream", ] diff --git a/litellm/google_genai/adapters/__init__.py b/litellm/google_genai/adapters/__init__.py index 6fbe7d95a55..796ddce8831 100644 --- a/litellm/google_genai/adapters/__init__.py +++ b/litellm/google_genai/adapters/__init__.py @@ -13,7 +13,7 @@ from .handler import GenerateContentToCompletionHandler from .transformation import GoogleGenAIAdapter, GoogleGenAIStreamWrapper __all__ = [ + "GenerateContentToCompletionHandler", "GoogleGenAIAdapter", "GoogleGenAIStreamWrapper", - "GenerateContentToCompletionHandler", ] diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 82777fb1378..5dafe2befee 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -1,4 +1,5 @@ -from typing import Any, AsyncIterator, Coroutine, Dict, List, Optional, Union, cast +from collections.abc import AsyncIterator, Coroutine +from typing import Any, Final, cast import litellm from litellm.types.router import GenericLiteLLMParams @@ -7,7 +8,7 @@ from litellm.types.utils import ModelResponse from .transformation import GoogleGenAIAdapter # Initialize adapter -GOOGLE_GENAI_ADAPTER = GoogleGenAIAdapter() +GOOGLE_GENAI_ADAPTER: Final = GoogleGenAIAdapter() class GenerateContentToCompletionHandler: @@ -16,16 +17,16 @@ class GenerateContentToCompletionHandler: @staticmethod def _prepare_completion_kwargs( model: str, - contents: Union[List[Dict[str, Any]], Dict[str, Any]], - config: Optional[Dict[str, Any]] = None, + contents: list[dict[str, Any]] | dict[str, Any], + config: dict[str, Any] | None = None, stream: bool = False, - litellm_params: Optional[GenericLiteLLMParams] = None, - extra_kwargs: Optional[Dict[str, Any]] = None, - ) -> Dict[str, Any]: + litellm_params: GenericLiteLLMParams | None = None, + extra_kwargs: dict[str, Any] | None = None, + ) -> dict[str, Any]: """Prepare kwargs for litellm.completion/acompletion""" # Transform generate_content request to completion format - completion_request = GOOGLE_GENAI_ADAPTER.translate_generate_content_to_completion( + completion_request: Final = GOOGLE_GENAI_ADAPTER.translate_generate_content_to_completion( model=model, contents=contents, config=config, @@ -33,7 +34,7 @@ class GenerateContentToCompletionHandler: **(extra_kwargs or {}), ) - completion_kwargs: Dict[str, Any] = dict(completion_request) + completion_kwargs: Final[dict[str, Any]] = dict(completion_request) # Forward extra_kwargs that should be passed to completion call if extra_kwargs is not None: @@ -52,15 +53,15 @@ class GenerateContentToCompletionHandler: @staticmethod async def async_generate_content_handler( model: str, - contents: Union[List[Dict[str, Any]], Dict[str, Any]], + contents: list[dict[str, Any]] | dict[str, Any], litellm_params: GenericLiteLLMParams, - config: Optional[Dict[str, Any]] = None, + config: dict[str, Any] | None = None, stream: bool = False, **kwargs, - ) -> Union[Dict[str, Any], AsyncIterator[bytes]]: + ) -> dict[str, Any] | AsyncIterator[bytes]: """Handle generate_content call asynchronously using completion adapter""" - completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs( + completion_kwargs: Final = GenerateContentToCompletionHandler._prepare_completion_kwargs( model=model, contents=contents, config=config, @@ -70,7 +71,7 @@ class GenerateContentToCompletionHandler: ) try: - completion_response = await litellm.acompletion(**completion_kwargs) + completion_response: Final = await litellm.acompletion(**completion_kwargs) if stream: # Check if completion_response is actually a stream or a ModelResponse @@ -83,7 +84,7 @@ class GenerateContentToCompletionHandler: return generate_content_response else: # Transform streaming completion response to generate_content format - transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( + transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( completion_response ) if transformed_stream is not None: @@ -97,22 +98,18 @@ class GenerateContentToCompletionHandler: return generate_content_response except Exception as e: - raise ValueError(f"Error calling litellm.acompletion for generate_content: {str(e)}") + raise ValueError(f"Error calling litellm.acompletion for generate_content: {e}") @staticmethod def generate_content_handler( model: str, - contents: Union[List[Dict[str, Any]], Dict[str, Any]], + contents: list[dict[str, Any]] | dict[str, Any], litellm_params: GenericLiteLLMParams, - config: Optional[Dict[str, Any]] = None, + config: dict[str, Any] | None = None, stream: bool = False, _is_async: bool = False, **kwargs, - ) -> Union[ - Dict[str, Any], - AsyncIterator[bytes], - Coroutine[Any, Any, Union[Dict[str, Any], AsyncIterator[bytes]]], - ]: + ) -> dict[str, Any] | AsyncIterator[bytes] | Coroutine[Any, Any, dict[str, Any] | AsyncIterator[bytes]]: """Handle generate_content call using completion adapter""" if _is_async: @@ -125,7 +122,7 @@ class GenerateContentToCompletionHandler: **kwargs, ) - completion_kwargs = GenerateContentToCompletionHandler._prepare_completion_kwargs( + completion_kwargs: Final = GenerateContentToCompletionHandler._prepare_completion_kwargs( model=model, contents=contents, config=config, @@ -135,7 +132,7 @@ class GenerateContentToCompletionHandler: ) try: - completion_response = litellm.completion(**completion_kwargs) + completion_response: Final = litellm.completion(**completion_kwargs) if stream: # Check if completion_response is actually a stream or a ModelResponse @@ -148,7 +145,7 @@ class GenerateContentToCompletionHandler: return generate_content_response else: # Transform streaming completion response to generate_content format - transformed_stream = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( + transformed_stream: Final = GOOGLE_GENAI_ADAPTER.translate_completion_output_params_streaming( completion_response ) if transformed_stream is not None: @@ -162,4 +159,4 @@ class GenerateContentToCompletionHandler: return generate_content_response except Exception as e: - raise ValueError(f"Error calling litellm.completion for generate_content: {str(e)}") + raise ValueError(f"Error calling litellm.completion for generate_content: {e}") diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 02dde12a30d..4f127f476c3 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,5 +1,6 @@ import json -from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union, cast +from collections.abc import AsyncIterator, Iterator +from typing import Any, Final, cast from litellm import verbose_logger from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema @@ -35,7 +36,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): sent_first_chunk: bool = False # State tracking for accumulating partial tool calls - accumulated_tool_calls: Dict[str, Dict[str, Any]] + accumulated_tool_calls: dict[str, dict[str, Any]] def __init__(self, completion_stream: Any): self.sent_first_chunk = False @@ -84,7 +85,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): # After the stream is exhausted, check for any remaining accumulated tool calls if self.accumulated_tool_calls: try: - parts = [] + parts: Final = [] for ( tool_call_index, tool_call_data, @@ -103,13 +104,13 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): except json.JSONDecodeError: # This can happen if the stream is abruptly cut off mid-argument string. verbose_logger.warning( - f"Could not parse tool call arguments at end of stream for index {tool_call_index}. " - f"Name: {tool_call_data['name']}. " - f"Partial args: {tool_call_data['arguments']}" + "Could not parse tool call arguments at end of stream for index %s. Name: %s. Partial args: %s", + tool_call_index, + tool_call_data["name"], + tool_call_data["arguments"], ) - pass if parts: - final_chunk = { + final_chunk: Final = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -177,11 +178,11 @@ class GoogleGenAIAdapter: def translate_generate_content_to_completion( self, model: str, - contents: Union[List[Dict[str, Any]], Dict[str, Any]], - config: Optional[Dict[str, Any]] = None, - litellm_params: Optional[GenericLiteLLMParams] = None, + contents: list[dict[str, Any]] | dict[str, Any], + config: dict[str, Any] | None = None, + litellm_params: GenericLiteLLMParams | None = None, **kwargs, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Transform generate_content request to litellm completion format @@ -196,9 +197,9 @@ class GoogleGenAIAdapter: """ # Extract top-level fields from kwargs - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") - tools = kwargs.get("tools") - tool_config = kwargs.get("toolConfig") or kwargs.get("tool_config") + system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + tools: Final = kwargs.get("tools") + tool_config: Final = kwargs.get("toolConfig") or kwargs.get("tool_config") # Normalize contents to list format if isinstance(contents, dict): @@ -207,10 +208,10 @@ class GoogleGenAIAdapter: contents_list = contents # Transform contents to OpenAI messages format - messages = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction) + messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction) # Create base request as dict (which is compatible with ChatCompletionRequest) - completion_request: ChatCompletionRequest = { + completion_request: Final[ChatCompletionRequest] = { "model": model, "messages": messages, } @@ -247,14 +248,14 @@ class GoogleGenAIAdapter: # Check if tools are already in OpenAI format or Google GenAI format if isinstance(tools, list) and len(tools) > 0: # Tools are in Google GenAI format, transform them - openai_tools = self._transform_google_genai_tools_to_openai(tools) + openai_tools: Final = self._transform_google_genai_tools_to_openai(tools) if openai_tools: completion_request["tools"] = openai_tools # Handle tool_config (tool choice) if tool_config: - tool_choice = self._transform_google_genai_tool_config_to_openai(tool_config) + tool_choice: Final = self._transform_google_genai_tool_config_to_openai(tool_config) if tool_choice: completion_request["tool_choice"] = tool_choice @@ -272,8 +273,8 @@ class GoogleGenAIAdapter: def _add_generic_litellm_params_to_request( self, - completion_request_dict: Dict[str, Any], - litellm_params: Optional[GenericLiteLLMParams] = None, + completion_request_dict: dict[str, Any], + litellm_params: GenericLiteLLMParams | None = None, ) -> dict: """Add generic litellm params to request. e.g add api_base, api_key, api_version, etc. @@ -284,9 +285,9 @@ class GoogleGenAIAdapter: Returns: Dict[str, Any] """ - allowed_fields = GenericLiteLLMParams.model_fields.keys() + allowed_fields: Final = GenericLiteLLMParams.model_fields.keys() if litellm_params: - litellm_dict = litellm_params.model_dump(exclude_none=True) + litellm_dict: Final = litellm_params.model_dump(exclude_none=True) for key, value in litellm_dict.items(): if key in allowed_fields: completion_request_dict[key] = value @@ -295,23 +296,23 @@ class GoogleGenAIAdapter: def translate_completion_output_params_streaming( self, completion_stream: Any, - ) -> Union[AsyncIterator[bytes], None]: + ) -> AsyncIterator[bytes] | None: """Transform streaming completion output to Google GenAI format""" - google_genai_wrapper = GoogleGenAIStreamWrapper(completion_stream=completion_stream) + google_genai_wrapper: Final = GoogleGenAIStreamWrapper(completion_stream=completion_stream) # Return the SSE-wrapped version for proper event formatting return google_genai_wrapper.async_google_genai_sse_wrapper() def _transform_google_genai_tools_to_openai( self, - tools: List[Dict[str, Any]], - ) -> List[ChatCompletionToolParam]: + tools: list[dict[str, Any]], + ) -> list[ChatCompletionToolParam]: """Transform Google GenAI tools to OpenAI tools format""" - openai_tools: List[Dict[str, Any]] = [] + openai_tools: Final[list[dict[str, Any]]] = [] for tool in tools: if "functionDeclarations" in tool: for func_decl in tool["functionDeclarations"]: - function_chunk: Dict[str, Any] = { + function_chunk: dict[str, Any] = { "name": func_decl.get("name", ""), } @@ -324,34 +325,34 @@ class GoogleGenAIAdapter: openai_tools.append(openai_tool) # normalize the tool schemas - normalized_tools = [normalize_tool_schema(tool) for tool in openai_tools] + normalized_tools: Final = [normalize_tool_schema(tool) for tool in openai_tools] - return cast(List[ChatCompletionToolParam], normalized_tools) + return cast(list[ChatCompletionToolParam], normalized_tools) def _transform_google_genai_tool_config_to_openai( self, - tool_config: Dict[str, Any], - ) -> Optional[ChatCompletionToolChoiceValues]: + tool_config: dict[str, Any], + ) -> ChatCompletionToolChoiceValues | None: """Transform Google GenAI tool_config to OpenAI tool_choice""" - function_calling_config = tool_config.get("functionCallingConfig", {}) - mode = function_calling_config.get("mode", "AUTO") + function_calling_config: Final = tool_config.get("functionCallingConfig", {}) + mode: Final = function_calling_config.get("mode", "AUTO") - mode_mapping = {"AUTO": "auto", "ANY": "required", "NONE": "none"} + mode_mapping: Final = {"AUTO": "auto", "ANY": "required", "NONE": "none"} - tool_choice = mode_mapping.get(mode, "auto") + tool_choice: Final = mode_mapping.get(mode, "auto") return cast(ChatCompletionToolChoiceValues, tool_choice) def _transform_contents_to_messages( self, - contents: List[Dict[str, Any]], - system_instruction: Optional[Dict[str, Any]] = None, - ) -> List[AllMessageValues]: + contents: list[dict[str, Any]], + system_instruction: dict[str, Any] | None = None, + ) -> list[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" - messages: List[AllMessageValues] = [] + messages: Final[list[AllMessageValues]] = [] # Handle system instruction if system_instruction: - system_parts = system_instruction.get("parts", []) + system_parts: Final = system_instruction.get("parts", []) if system_parts and "text" in system_parts[0]: messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"])) @@ -361,8 +362,8 @@ class GoogleGenAIAdapter: if role == "user": # Handle user messages with potential function responses - content_parts: List[Union[ChatCompletionTextObject, ChatCompletionImageObject]] = [] - tool_messages: List[ChatCompletionToolMessage] = [] + content_parts: list[ChatCompletionTextObject | ChatCompletionImageObject] = [] + tool_messages: list[ChatCompletionToolMessage] = [] for part in parts: if isinstance(part, dict): @@ -419,7 +420,7 @@ class GoogleGenAIAdapter: elif role == "model": # Handle assistant messages with potential function calls combined_text = "" - tool_calls: List[ChatCompletionAssistantToolCall] = [] + tool_calls: list[ChatCompletionAssistantToolCall] = [] for part in parts: if isinstance(part, dict): @@ -460,7 +461,7 @@ class GoogleGenAIAdapter: def translate_completion_to_generate_content( self, response: ModelResponse, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """ Transform litellm completion response to Google GenAI generate_content format @@ -472,7 +473,7 @@ class GoogleGenAIAdapter: """ # Extract the main response content - choice = response.choices[0] if response.choices else None + choice: Final = response.choices[0] if response.choices else None if not choice: raise ValueError("Invalid completion response: no choices found") @@ -489,7 +490,7 @@ class GoogleGenAIAdapter: parts = [{"text": message_content}] if message_content else [] # Create Google GenAI format response - generate_content_response: Dict[str, Any] = { + generate_content_response: Final[dict[str, Any]] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -521,9 +522,9 @@ class GoogleGenAIAdapter: def translate_streaming_completion_to_generate_content( self, - response: Union[ModelResponse, ModelResponseStream], + response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper, - ) -> Optional[Dict[str, Any]]: + ) -> dict[str, Any] | None: """ Transform streaming litellm completion chunk to Google GenAI generate_content format @@ -536,7 +537,7 @@ class GoogleGenAIAdapter: """ # Extract the main response content from streaming chunk - choice = response.choices[0] if response.choices else None + choice: Final = response.choices[0] if response.choices else None if not choice: # Return empty chunk if no choices return None @@ -550,7 +551,7 @@ class GoogleGenAIAdapter: finish_reason = getattr(choice, "finish_reason", None) else: # Fallback for generic choice objects - message_content = getattr(choice, "delta", {}).get("content", "") + message_content: Final = getattr(choice, "delta", {}).get("content", "") parts = [{"text": message_content}] if message_content else [] finish_reason = getattr(choice, "finish_reason", None) @@ -559,7 +560,7 @@ class GoogleGenAIAdapter: return None # Create Google GenAI streaming format response - streaming_chunk: Dict[str, Any] = { + streaming_chunk: Final[dict[str, Any]] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -572,7 +573,7 @@ class GoogleGenAIAdapter: # Add usage metadata only in the final chunk (when finish_reason is present) if finish_reason: - usage_metadata = ( + usage_metadata: Final = ( self._map_usage(getattr(response, "usage", None)) if hasattr(response, "usage") and getattr(response, "usage", None) else { @@ -596,9 +597,9 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, message: Any, - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Transform OpenAI message to Google GenAI parts format""" - parts: List[Dict[str, Any]] = [] + parts: Final[list[dict[str, Any]]] = [] # Add text content if present if hasattr(message, "content") and message.content: @@ -625,20 +626,20 @@ class GoogleGenAIAdapter: def _transform_openai_delta_to_google_genai_parts_with_accumulation( self, delta: Any, wrapper: GoogleGenAIStreamWrapper - ) -> List[Dict[str, Any]]: + ) -> list[dict[str, Any]]: """Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls.""" # 1. Initialize wrapper state if it doesn't exist if not hasattr(wrapper, "accumulated_tool_calls"): wrapper.accumulated_tool_calls = {} - parts: List[Dict[str, Any]] = [] + parts: Final[list[dict[str, Any]]] = [] if hasattr(delta, "content") and delta.content: parts.append({"text": delta.content}) # 2. Ensure tool_calls is iterable - tool_calls = delta.tool_calls or [] + tool_calls: Final = delta.tool_calls or [] for tool_call in tool_calls: if not hasattr(tool_call, "function"): @@ -662,7 +663,7 @@ class GoogleGenAIAdapter: # Optimization: Skip chunks that have no new data if not function_name and not args_chunk: - verbose_logger.debug(f"Skipping empty tool call chunk for index: {tool_call_index}") + verbose_logger.debug("Skipping empty tool call chunk for index: %s", tool_call_index) continue if function_name: @@ -698,12 +699,12 @@ class GoogleGenAIAdapter: return parts - def _map_finish_reason(self, finish_reason: Optional[str]) -> str: + def _map_finish_reason(self, finish_reason: str | None) -> str: """Map OpenAI finish reasons to Google GenAI finish reasons""" if not finish_reason: return "STOP" - mapping = { + mapping: Final = { "stop": "STOP", "length": "MAX_TOKENS", "content_filter": "SAFETY", @@ -713,7 +714,7 @@ class GoogleGenAIAdapter: return mapping.get(finish_reason, "STOP") - def _map_usage(self, usage: Any) -> Dict[str, int]: + def _map_usage(self, usage: Any) -> dict[str, int]: """Map OpenAI usage to Google GenAI usage format""" return { "promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0, diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 3b1e712342f..634739d86f8 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -1,7 +1,8 @@ import asyncio import contextvars +from collections.abc import Iterator from functools import partial -from typing import TYPE_CHECKING, Any, ClassVar, Dict, Iterator, Optional, Union +from typing import TYPE_CHECKING, Any, ClassVar, Final import httpx from pydantic import BaseModel, ConfigDict @@ -51,14 +52,14 @@ class GenerateContentSetupResult(BaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) model: str - request_body: Dict[str, Any] + request_body: dict[str, Any] custom_llm_provider: str - generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] - generate_content_config_dict: Dict[str, Any] + generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None + generate_content_config_dict: dict[str, Any] native_request_fields: dict[str, object] litellm_params: GenericLiteLLMParams litellm_logging_obj: LiteLLMLoggingObj - litellm_call_id: Optional[str] + litellm_call_id: str | None class GenerateContentHelper: @@ -67,7 +68,7 @@ class GenerateContentHelper: @staticmethod def mock_generate_content_response( mock_response: str = "This is a mock response from Google GenAI generate_content.", - ) -> Dict[str, Any]: + ) -> dict[str, Any]: """Mock response for generate_content for testing purposes""" return { "text": mock_response, @@ -90,9 +91,9 @@ class GenerateContentHelper: def setup_generate_content_call( model: str, contents: GenerateContentContentListUnionDict, - config: Optional[GenerateContentConfigDict] = None, - custom_llm_provider: Optional[str] = None, - tools: Optional[ToolConfigDict] = None, + config: GenerateContentConfigDict | None = None, + custom_llm_provider: str | None = None, + tools: ToolConfigDict | None = None, **kwargs, ) -> GenerateContentSetupResult: """ @@ -109,11 +110,11 @@ class GenerateContentHelper: Returns: GenerateContentSetupResult containing all setup information """ - litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) + litellm_logging_obj: Final[LiteLLMLoggingObj | None] = kwargs.get("litellm_logging_obj") + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) # get llm provider logic - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) ## MOCK RESPONSE LOGIC (only for non-streaming) if ( @@ -139,7 +140,7 @@ class GenerateContentHelper: litellm_params.custom_llm_provider = custom_llm_provider # get provider config - generate_content_provider_config: Optional[BaseGoogleGenAIGenerateContentConfig] = ( + generate_content_provider_config: Final[BaseGoogleGenAIGenerateContentConfig | None] = ( ProviderConfigManager.get_provider_google_genai_generate_content_config( model=model, provider=litellm.LlmProviders(custom_llm_provider), @@ -167,19 +168,19 @@ class GenerateContentHelper: # Construct request body ######################################################################################### # Create Google Optional Params Config - generate_content_config_dict = generate_content_provider_config.map_generate_content_optional_params( + generate_content_config_dict: Final = generate_content_provider_config.map_generate_content_optional_params( generate_content_config_dict=config or {}, model=model, ) # Extract systemInstruction from kwargs to pass to transform - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Native top-level REST fields arrive as loose kwargs and are otherwise dropped. - native_request_fields: dict[str, object] = { + native_request_fields: Final[dict[str, object]] = { field: kwargs[field] for field in generate_content_provider_config.get_generate_content_request_top_level_fields() if field in kwargs } - request_body = generate_content_provider_config.transform_generate_content_request( + request_body: Final = generate_content_provider_config.transform_generate_content_request( model=model, contents=contents, tools=tools, @@ -234,24 +235,24 @@ def _merge_native_request_fields( async def agenerate_content( model: str, contents: GenerateContentContentListUnionDict, - config: Optional[GenerateContentConfigDict] = None, - tools: Optional[ToolConfigDict] = None, + config: GenerateContentConfigDict | None = None, + tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Any: """ Async: Generate content using Google GenAI """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["agenerate_content"] = True # Handle generationConfig parameter from kwargs for backward compatibility @@ -264,7 +265,7 @@ async def agenerate_content( custom_llm_provider=custom_llm_provider, ) - func = partial( + func: Final = partial( generate_content, model=model, contents=contents, @@ -278,9 +279,9 @@ async def agenerate_content( **kwargs, ) - 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 @@ -302,24 +303,24 @@ async def agenerate_content( def generate_content( model: str, contents: GenerateContentContentListUnionDict, - config: Optional[GenerateContentConfigDict] = None, - tools: Optional[ToolConfigDict] = None, + config: GenerateContentConfigDict | None = None, + tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Any: """ Generate content using Google GenAI """ - local_vars = locals() + local_vars: Final = locals() try: - _is_async = kwargs.pop("agenerate_content", False) + _is_async: Final = kwargs.pop("agenerate_content", False) _mark_async_entrypoint(kwargs.get("litellm_logging_obj"), CallTypes.agenerate_content.value, _is_async) @@ -327,12 +328,12 @@ def generate_content( if "generationConfig" in kwargs and config is None: config = kwargs.pop("generationConfig") # Check for mock response first - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) if litellm_params.mock_response and isinstance(litellm_params.mock_response, str): return GenerateContentHelper.mock_generate_content_response(mock_response=litellm_params.mock_response) # Setup the call - setup_result = GenerateContentHelper.setup_generate_content_call( + setup_result: Final = GenerateContentHelper.setup_generate_content_call( model=model, contents=contents, config=config, @@ -342,7 +343,7 @@ def generate_content( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -359,7 +360,7 @@ def generate_content( ) # Call the standard handler - response = base_llm_http_handler.generate_content_handler( + response: Final = base_llm_http_handler.generate_content_handler( model=setup_result.model, contents=contents, tools=tools, @@ -392,22 +393,22 @@ def generate_content( async def agenerate_content_stream( model: str, contents: GenerateContentContentListUnionDict, - config: Optional[GenerateContentConfigDict] = None, - tools: Optional[ToolConfigDict] = None, + config: GenerateContentConfigDict | None = None, + tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Any: """ Async: Generate content using Google GenAI with streaming response """ - local_vars = locals() + local_vars: Final = locals() try: kwargs["agenerate_content_stream"] = True @@ -423,7 +424,7 @@ async def agenerate_content_stream( ) # Setup the call - setup_result = GenerateContentHelper.setup_generate_content_call( + setup_result: Final = GenerateContentHelper.setup_generate_content_call( model=model, contents=contents, config=config, @@ -433,7 +434,7 @@ async def agenerate_content_stream( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -487,25 +488,25 @@ async def agenerate_content_stream( def generate_content_stream( model: str, contents: GenerateContentContentListUnionDict, - config: Optional[GenerateContentConfigDict] = None, - tools: Optional[ToolConfigDict] = None, + config: GenerateContentConfigDict | None = None, + tools: ToolConfigDict | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, **kwargs, ) -> Iterator[Any]: """ Generate content using Google GenAI with streaming response """ - local_vars = locals() + local_vars: Final = locals() try: # Remove any async-related flags since this is the sync function - _is_async = kwargs.pop("agenerate_content_stream", False) + _is_async: Final = kwargs.pop("agenerate_content_stream", False) _mark_async_entrypoint(kwargs.get("litellm_logging_obj"), CallTypes.agenerate_content_stream.value, _is_async) @@ -513,7 +514,7 @@ def generate_content_stream( if "generationConfig" in kwargs and config is None: config = kwargs.pop("generationConfig") # Setup the call - setup_result = GenerateContentHelper.setup_generate_content_call( + setup_result: Final = GenerateContentHelper.setup_generate_content_call( model=model, contents=contents, config=config, @@ -523,7 +524,7 @@ def generate_content_stream( ) # Extract systemInstruction from kwargs to pass to handler - system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index 900a171640b..e03f7ee745f 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Final from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.success_handler import ( @@ -15,15 +15,15 @@ if TYPE_CHECKING: else: BaseGoogleGenAIGenerateContentConfig = Any -GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging() +GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() -def _encode_google_genai_sse_event(event_lines: List[str]) -> bytes: +def _encode_google_genai_sse_event(event_lines: list[str]) -> bytes: return ("\n".join(event_lines) + "\n\n").encode("utf-8") def _next_google_genai_sse_chunk(line_iter) -> bytes: - event_lines: List[str] = [] + event_lines: Final[list[str]] = [] while True: try: line = next(line_iter) @@ -39,7 +39,7 @@ def _next_google_genai_sse_chunk(line_iter) -> bytes: async def _anext_google_genai_sse_chunk(line_iter) -> bytes: - event_lines: List[str] = [] + event_lines: Final[list[str]] = [] while True: try: line = await line_iter.__anext__() @@ -65,14 +65,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: litellm_logging_obj: LiteLLMLoggingObj, request_body: dict, model: str, - hidden_params: Optional[Dict[str, Any]] = None, + hidden_params: dict[str, Any] | None = None, ): self.litellm_logging_obj = litellm_logging_obj self.request_body = request_body self.start_time = datetime.now() - self.collected_chunks: List[bytes] = [] + self.collected_chunks: list[bytes] = [] self.model = model - self._hidden_params: Dict[str, Any] = hidden_params or {} + self._hidden_params: dict[str, Any] = hidden_params or {} async def _handle_async_streaming_logging( self, @@ -82,7 +82,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: PassThroughStreamingHandler, ) - end_time = datetime.now() + end_time: Final = datetime.now() asyncio.create_task( PassThroughStreamingHandler._route_streaming_logging_to_handler( litellm_logging_obj=self.litellm_logging_obj, @@ -111,8 +111,8 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, litellm_metadata: dict, custom_llm_provider: str, - request_body: Optional[dict] = None, - hidden_params: Optional[Dict[str, Any]] = None, + request_body: dict | None = None, + hidden_params: dict[str, Any] | None = None, ): super().__init__( litellm_logging_obj=logging_obj, @@ -134,7 +134,7 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent def __next__(self): try: - chunk = _next_google_genai_sse_chunk(self.stream_iterator) + chunk: Final = _next_google_genai_sse_chunk(self.stream_iterator) self.collected_chunks.append(chunk) return chunk except StopIteration: @@ -162,8 +162,8 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig, litellm_metadata: dict, custom_llm_provider: str, - request_body: Optional[dict] = None, - hidden_params: Optional[Dict[str, Any]] = None, + request_body: dict | None = None, + hidden_params: dict[str, Any] | None = None, ): super().__init__( litellm_logging_obj=logging_obj, @@ -185,7 +185,7 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo async def __anext__(self): try: - chunk = await _anext_google_genai_sse_chunk(self.stream_iterator) + chunk: Final = await _anext_google_genai_sse_chunk(self.stream_iterator) self.collected_chunks.append(chunk) return chunk except StopAsyncIteration: diff --git a/litellm/images/main.py b/litellm/images/main.py index 17ea9aa177b..4430bb5beb4 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1,19 +1,9 @@ import asyncio import contextvars import importlib +from collections.abc import Coroutine from functools import partial -from typing import ( - TYPE_CHECKING, - Any, - Coroutine, - Dict, - List, - Literal, - Optional, - Union, - cast, - overload, -) +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload if TYPE_CHECKING: from litellm.images.utils import ImageEditRequestUtils @@ -79,7 +69,7 @@ def _get_ImageEditRequestUtils() -> "ImageEditRequestUtils": global _ImageEditRequestUtils_cache if _ImageEditRequestUtils_cache is None: # Access via module to trigger __getattr__ if not cached - module = importlib.import_module(__name__) + module: Final = importlib.import_module(__name__) _ImageEditRequestUtils_cache = module.ImageEditRequestUtils assert _ImageEditRequestUtils_cache is not None # Type narrowing for type checker return _ImageEditRequestUtils_cache @@ -98,25 +88,25 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: Returns: - `response` (Any): The response returned by the `image_generation` function. """ - loop = asyncio.get_event_loop() - model = args[0] if len(args) > 0 else kwargs["model"] + loop: Final = asyncio.get_event_loop() + model: Final = args[0] if len(args) > 0 else kwargs["model"] ### PASS ARGS TO Image Generation ### kwargs["aimg_generation"] = True custom_llm_provider = None try: # Use a partial function to pass your keyword arguments - func = partial(image_generation, *args, **kwargs) + func: Final = partial(image_generation, *args, **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(model=model, api_base=kwargs.get("api_base", None)) # 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) - response: Optional[ImageResponse] = None + response: ImageResponse | None = None if isinstance(init_response, dict): response = ImageResponse(**init_response) elif isinstance(init_response, ImageResponse): ## CACHING SCENARIO @@ -145,17 +135,17 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse: @overload def image_generation( prompt: str, - model: Optional[str] = None, - n: Optional[int] = None, - quality: Optional[Union[str, ImageGenerationRequestQuality]] = None, - response_format: Optional[str] = None, - size: Optional[str] = None, - style: Optional[str] = None, - user: Optional[str] = None, + model: str | None = None, + n: int | None = None, + quality: str | ImageGenerationRequestQuality | None = None, + response_format: str | None = None, + size: str | None = None, + style: str | None = None, + user: str | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider=None, *, aimg_generation: Literal[True], @@ -169,17 +159,17 @@ def image_generation( @overload def image_generation( prompt: str, - model: Optional[str] = None, - n: Optional[int] = None, - quality: Optional[Union[str, ImageGenerationRequestQuality]] = None, - response_format: Optional[str] = None, - size: Optional[str] = None, - style: Optional[str] = None, - user: Optional[str] = None, + model: str | None = None, + n: int | None = None, + quality: str | ImageGenerationRequestQuality | None = None, + response_format: str | None = None, + size: str | None = None, + style: str | None = None, + user: str | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider=None, *, aimg_generation: Literal[False] = False, @@ -193,47 +183,44 @@ def image_generation( @client def image_generation( prompt: str, - model: Optional[str] = None, - n: Optional[int] = None, - quality: Optional[Union[str, ImageGenerationRequestQuality]] = None, - response_format: Optional[str] = None, - size: Optional[str] = None, - style: Optional[str] = None, - user: Optional[str] = None, + model: str | None = None, + n: int | None = None, + quality: str | ImageGenerationRequestQuality | None = None, + response_format: str | None = None, + size: str | None = None, + style: str | None = None, + user: str | None = None, timeout=600, # default to 10 minutes - api_key: Optional[str] = None, - api_base: Optional[str] = None, - api_version: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + api_version: str | None = None, custom_llm_provider=None, **kwargs, -) -> Union[ - ImageResponse, - Coroutine[Any, Any, ImageResponse], -]: +) -> ImageResponse | Coroutine[Any, Any, ImageResponse]: """ Maps the https://api.openai.com/v1/images/generations endpoint. Currently supports just Azure + OpenAI. """ try: - args = locals() - aimg_generation = kwargs.get("aimg_generation", False) - litellm_call_id = kwargs.get("litellm_call_id", None) - logger_fn = kwargs.get("logger_fn", None) - mock_response: Optional[str] = kwargs.get("mock_response", None) # type: ignore - proxy_server_request = kwargs.get("proxy_server_request", None) + args: Final = locals() + aimg_generation: Final = kwargs.get("aimg_generation", False) + litellm_call_id: Final = kwargs.get("litellm_call_id", None) + logger_fn: Final = kwargs.get("logger_fn", None) + mock_response: Final[str | None] = kwargs.get("mock_response", None) # type: ignore + proxy_server_request: Final = kwargs.get("proxy_server_request", None) azure_ad_token_provider = kwargs.get("azure_ad_token_provider", None) - model_info = kwargs.get("model_info", None) - metadata = kwargs.get("metadata", {}) - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - client = kwargs.get("client", None) - extra_headers = kwargs.get("extra_headers", None) - headers: dict = kwargs.get("headers", None) or {} - base_model = kwargs.get("base_model", None) + model_info: Final = kwargs.get("model_info", None) + metadata: Final = kwargs.get("metadata", {}) + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + client: Final = kwargs.get("client", None) + extra_headers: Final = kwargs.get("extra_headers", None) + headers: Final[dict] = kwargs.get("headers", None) or {} + base_model: Final = kwargs.get("base_model", None) if extra_headers is not None: headers.update(extra_headers) model_response: ImageResponse = litellm.utils.ImageResponse() - dynamic_api_key: Optional[str] = None + dynamic_api_key: str | None = None if model is not None or custom_llm_provider is not None: model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, # type: ignore @@ -244,7 +231,7 @@ def image_generation( model = "dall-e-2" custom_llm_provider = "openai" # default to dall-e-2 on openai model_response._hidden_params["model"] = model - openai_params = [ + openai_params: Final = [ "user", "request_timeout", "api_base", @@ -261,20 +248,20 @@ def image_generation( "size", "style", ] - litellm_params = all_litellm_params - default_params = openai_params + litellm_params - non_default_params = { + litellm_params: Final = all_litellm_params + default_params: Final = openai_params + litellm_params + non_default_params: Final = { k: v for k, v in kwargs.items() if k not in default_params } # model-specific params - pass them straight to the model/provider - image_generation_config: Optional[BaseImageGenerationConfig] = None + image_generation_config: BaseImageGenerationConfig | None = None if custom_llm_provider is not None and custom_llm_provider in LlmProviders._member_map_.values(): image_generation_config = ProviderConfigManager.get_provider_image_generation_config( model=base_model or model, provider=LlmProviders(custom_llm_provider), ) - optional_params = get_optional_params_image_gen( + optional_params: Final = get_optional_params_image_gen( model=base_model or model, n=n, quality=quality, @@ -287,9 +274,9 @@ def image_generation( **non_default_params, ) - litellm_params_dict = get_litellm_params(**kwargs) + litellm_params_dict: Final = get_litellm_params(**kwargs) - logging: Logging = litellm_logging_obj + logging: Final[Logging] = litellm_logging_obj logging.update_from_kwargs( kwargs=kwargs, model=model, @@ -314,7 +301,7 @@ def image_generation( if custom_llm_provider == "azure": # azure configs - api_type = get_secret_str("AZURE_API_TYPE") or "azure" + api_type: Final = get_secret_str("AZURE_API_TYPE") or "azure" api_base = api_base or litellm.api_base or get_secret_str("AZURE_API_BASE") @@ -328,7 +315,7 @@ def image_generation( or get_secret_str("AZURE_API_KEY") ) - azure_ad_token = optional_params.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") + azure_ad_token: Final = optional_params.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") # Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided if azure_ad_token_provider is None: @@ -337,9 +324,9 @@ def image_generation( ) # Extract Azure AD credentials from litellm_params - tenant_id = litellm_params_dict.get("tenant_id") - client_id = litellm_params_dict.get("client_id") - client_secret = litellm_params_dict.get("client_secret") + tenant_id: Final = litellm_params_dict.get("tenant_id") + client_id: Final = litellm_params_dict.get("client_id") + client_secret: Final = litellm_params_dict.get("client_secret") azure_scope = litellm_params_dict.get("azure_scope") or "https://cognitiveservices.azure.com/.default" # Create token provider if credentials are available @@ -398,7 +385,7 @@ def image_generation( raise ValueError(f"image generation config is not supported for {custom_llm_provider}") # Resolve api_base from litellm.api_base if not explicitly provided - _api_base = api_base or litellm.api_base + _api_base: Final = api_base or litellm.api_base litellm_params_dict["api_base"] = _api_base return llm_http_handler.image_generation_handler( @@ -474,7 +461,7 @@ def image_generation( if extra_headers is not None: optional_params["extra_headers"] = extra_headers # Forward OpenAI organization if present (set by proxy pre-call utils) - organization: Optional[str] = kwargs.get("organization", None) + organization: Final[str | None] = kwargs.get("organization", None) model_response = openai_chat_completions.image_generation( model=model, prompt=prompt, @@ -506,7 +493,7 @@ def image_generation( ) elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider # Get the Custom Handler - custom_handler: Optional[CustomLLM] = None + custom_handler: CustomLLM | None = None for item in litellm.custom_provider_map: if item["provider"] == custom_llm_provider: custom_handler = item["custom_handler"] @@ -516,7 +503,7 @@ def image_generation( ## ROUTE LLM CALL ## if aimg_generation is True: - async_custom_client: Optional[AsyncHTTPHandler] = None + async_custom_client: AsyncHTTPHandler | None = None if client is not None and isinstance(client, AsyncHTTPHandler): async_custom_client = client @@ -533,7 +520,7 @@ def image_generation( client=async_custom_client, ) else: - custom_client: Optional[HTTPHandler] = None + custom_client: HTTPHandler | None = None if client is not None and isinstance(client, HTTPHandler): custom_client = client @@ -574,18 +561,18 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse: Returns: - `response` (Any): The response returned by the `image_variation` function. """ - loop = asyncio.get_event_loop() - model = kwargs.get("model", None) + loop: Final = asyncio.get_event_loop() + model: Final = kwargs.get("model", None) custom_llm_provider = kwargs.get("custom_llm_provider", None) ### PASS ARGS TO Image Generation ### kwargs["async_call"] = True try: # Use a partial function to pass your keyword arguments - func = partial(image_variation, *args, **kwargs) + func: Final = partial(image_variation, *args, **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) if custom_llm_provider is None and model is not None: _, custom_llm_provider, _, _ = get_llm_provider(model=model, api_base=kwargs.get("api_base", None)) @@ -619,17 +606,17 @@ def image_variation( model: str = "dall-e-2", # set to dall-e-2 by default - like OpenAI. n: int = 1, response_format: Literal["url", "b64_json"] = "url", - size: Optional[str] = None, - user: Optional[str] = None, + size: str | None = None, + user: str | None = None, **kwargs, ) -> ImageResponse: # get non-default params - client = kwargs.get("client", None) + client: Final = kwargs.get("client", None) # get logging object - litellm_logging_obj = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj")) + litellm_logging_obj: Final = cast(LiteLLMLoggingObj, kwargs.get("litellm_logging_obj")) # get the litellm params - litellm_params = get_litellm_params(**kwargs) + litellm_params: Final = get_litellm_params(**kwargs) # get the custom llm provider model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( model=model, @@ -640,17 +627,17 @@ def image_variation( # route to the correct provider w/ the params try: - llm_provider = LlmProviders(custom_llm_provider) - image_variation_provider = LITELLM_IMAGE_VARIATION_PROVIDERS(llm_provider) + llm_provider: Final = LlmProviders(custom_llm_provider) + image_variation_provider: Final = LITELLM_IMAGE_VARIATION_PROVIDERS(llm_provider) except ValueError: raise ValueError( f"Invalid image variation provider: {custom_llm_provider}. Supported providers are: {LITELLM_IMAGE_VARIATION_PROVIDERS}" ) - model_response = ImageResponse() + model_response: Final = ImageResponse() - response: Optional[ImageResponse] = None + response: ImageResponse | None = None - provider_config = ProviderConfigManager.get_provider_model_info( + provider_config: Final = ProviderConfigManager.get_provider_model_info( model=model or "", # openai defaults to dall-e-2 provider=llm_provider, ) @@ -660,7 +647,7 @@ def image_variation( f"image variation provider has no known model info config - required for getting api keys, etc.: {custom_llm_provider}. Supported providers are: {LITELLM_IMAGE_VARIATION_PROVIDERS}" ) - api_key = provider_config.get_api_key(litellm_params.get("api_key", None)) + api_key: Final = provider_config.get_api_key(litellm_params.get("api_key", None)) api_base = provider_config.get_api_base(litellm_params.get("api_base", None)) if image_variation_provider == LITELLM_IMAGE_VARIATION_PROVIDERS.OPENAI: @@ -711,31 +698,31 @@ def image_variation( @client def image_edit( - image: Optional[Union[FileTypes, List[FileTypes]]] = None, - prompt: Optional[str] = None, - model: Optional[str] = None, - mask: Optional[str] = None, - n: Optional[int] = None, - quality: Optional[Union[str, ImageGenerationRequestQuality]] = None, - response_format: Optional[str] = None, - size: Optional[str] = None, - user: Optional[str] = None, + image: FileTypes | list[FileTypes] | None = None, + prompt: str | None = None, + model: str | None = None, + mask: str | None = None, + n: int | None = None, + quality: str | ImageGenerationRequestQuality | None = None, + response_format: str | None = None, + size: str | None = None, + user: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, **kwargs, -) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse]]: +) -> ImageResponse | Coroutine[Any, Any, ImageResponse]: """ Maps the image edit functionality, similar to OpenAI's images/edits endpoint. """ - local_vars = locals() + local_vars: Final = locals() try: - openai_params = [ + openai_params: Final = [ "user", "request_timeout", "api_base", @@ -753,22 +740,22 @@ def image_edit( "style", "async_call", ] - litellm_params_list = all_litellm_params - default_params = openai_params + litellm_params_list - non_default_params = { + litellm_params_list: Final = all_litellm_params + default_params: Final = openai_params + litellm_params_list + non_default_params: Final = { k: v for k, v in kwargs.items() if k not in default_params } # model-specific params - pass them straight to the model/provider - litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore - litellm_call_id: Optional[str] = kwargs.get("litellm_call_id", None) - model_info = kwargs.get("model_info", None) - metadata = kwargs.get("metadata", {}) - _is_async = kwargs.pop("async_call", False) is True + litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") # type: ignore + litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) + model_info: Final = kwargs.get("model_info", None) + metadata: Final = kwargs.get("metadata", {}) + _is_async: Final = kwargs.pop("async_call", False) is True # add images / or return a single image - images = image if isinstance(image, list) else ([image] if image is not None else []) + images: Final = image if isinstance(image, list) else ([image] if image is not None else []) - headers_from_kwargs = kwargs.get("headers") - merged_extra_headers: Dict[str, Any] = {} + headers_from_kwargs: Final = kwargs.get("headers") + merged_extra_headers: Final[dict[str, Any]] = {} if isinstance(headers_from_kwargs, dict): merged_extra_headers.update(headers_from_kwargs) if isinstance(extra_headers, dict): @@ -778,7 +765,7 @@ def image_edit( extra_headers = dict(merged_extra_headers) # get llm provider logic - litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params: Final = GenericLiteLLMParams(**kwargs) model, custom_llm_provider, _, _ = get_llm_provider( model=model or DEFAULT_IMAGE_ENDPOINT_MODEL, custom_llm_provider=custom_llm_provider, @@ -786,7 +773,7 @@ def image_edit( # Check for custom provider if custom_llm_provider in litellm._custom_providers: - custom_handler: Optional[CustomLLM] = None + custom_handler: CustomLLM | None = None for item in litellm.custom_provider_map: if item["provider"] == custom_llm_provider: custom_handler = item["custom_handler"] @@ -794,10 +781,10 @@ def image_edit( if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) - model_response = ImageResponse() + model_response: Final = ImageResponse() if _is_async: - async_custom_client: Optional[AsyncHTTPHandler] = None + async_custom_client: AsyncHTTPHandler | None = None if kwargs.get("client") is not None and isinstance(kwargs.get("client"), AsyncHTTPHandler): async_custom_client = kwargs.get("client") @@ -814,7 +801,7 @@ def image_edit( client=async_custom_client, ) else: - custom_client: Optional[HTTPHandler] = None + custom_client: HTTPHandler | None = None if kwargs.get("client") is not None and isinstance(kwargs.get("client"), HTTPHandler): custom_client = kwargs.get("client") @@ -832,11 +819,9 @@ def image_edit( ) # get provider config - image_edit_provider_config: Optional[BaseImageEditConfig] = ( - ProviderConfigManager.get_provider_image_edit_config( - model=model, - provider=litellm.LlmProviders(custom_llm_provider), - ) + image_edit_provider_config: BaseImageEditConfig | None = ProviderConfigManager.get_provider_image_edit_config( + model=model, + provider=litellm.LlmProviders(custom_llm_provider), ) if image_edit_provider_config is None: @@ -844,11 +829,11 @@ def image_edit( local_vars.update(kwargs) # Get ImageEditOptionalRequestParams with only valid parameters - image_edit_optional_params: ImageEditOptionalRequestParams = ( + image_edit_optional_params: Final[ImageEditOptionalRequestParams] = ( _get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars) ) # Get optional parameters for the responses API - image_edit_request_params: Dict = _get_ImageEditRequestUtils().get_optional_params_image_edit( + image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit( model=model, image_edit_provider_config=image_edit_provider_config, image_edit_optional_params=image_edit_optional_params, @@ -952,23 +937,23 @@ def image_edit( @client async def aimage_edit( - image: Union[FileTypes, List[FileTypes]], + image: FileTypes | list[FileTypes], model: str, prompt: str, - mask: Optional[str] = None, - n: Optional[int] = None, - quality: Optional[Union[str, ImageGenerationRequestQuality]] = None, - response_format: Optional[str] = None, - size: Optional[str] = None, - user: Optional[str] = None, + mask: str | None = None, + n: int | None = None, + quality: str | ImageGenerationRequestQuality | None = None, + response_format: str | None = None, + size: str | None = None, + user: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: Optional[Dict[str, Any]] = None, - extra_query: Optional[Dict[str, Any]] = None, - extra_body: Optional[Dict[str, Any]] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + extra_headers: dict[str, Any] | None = None, + extra_query: dict[str, Any] | None = None, + extra_body: dict[str, Any] | None = None, + timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, - custom_llm_provider: Optional[str] = None, + custom_llm_provider: str | None = None, **kwargs, ) -> ImageResponse: """ @@ -981,9 +966,9 @@ async def aimage_edit( Returns: - `response` (Any): The response returned by the `image_edit` function. """ - local_vars = locals() + local_vars: Final = locals() try: - loop = asyncio.get_event_loop() + loop: Final = asyncio.get_event_loop() kwargs["async_call"] = True # get custom llm provider so we can use this for mapping exceptions @@ -992,9 +977,9 @@ async def aimage_edit( model=model, api_base=local_vars.get("base_url", None) ) - images = image if isinstance(image, list) else [image] + images: Final = image if isinstance(image, list) else [image] - func = partial( + func: Final = partial( image_edit, image=images, prompt=prompt, @@ -1010,9 +995,9 @@ async def aimage_edit( **kwargs, ) - 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 @@ -1037,7 +1022,7 @@ def __getattr__(name: str) -> Any: from .utils import ImageEditRequestUtils as _ImageEditRequestUtils # Cache it in the module's __dict__ for subsequent accesses - module = importlib.import_module(__name__) + module: Final = importlib.import_module(__name__) module.__dict__["ImageEditRequestUtils"] = _ImageEditRequestUtils return _ImageEditRequestUtils raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/images/utils.py b/litellm/images/utils.py index f0d4c985c01..2f080d88de4 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -1,5 +1,5 @@ from io import BufferedReader, BytesIO -from typing import Any, Dict, List, Optional, cast, get_type_hints +from typing import Any, Final, cast, get_type_hints import litellm from litellm.litellm_core_utils.token_counter import get_image_type @@ -14,9 +14,9 @@ class ImageEditRequestUtils: model: str, image_edit_provider_config: BaseImageEditConfig, image_edit_optional_params: ImageEditOptionalRequestParams, - drop_params: Optional[bool] = None, - additional_drop_params: Optional[List[str]] = None, - ) -> Dict: + drop_params: bool | None = None, + additional_drop_params: list[str] | None = None, + ) -> dict: """ Get optional parameters for the image edit API. @@ -30,16 +30,16 @@ class ImageEditRequestUtils: Returns: A dictionary of supported parameters for the image edit API """ - supported_params = image_edit_provider_config.get_supported_openai_params(model) + supported_params: Final = image_edit_provider_config.get_supported_openai_params(model) - should_drop = litellm.drop_params is True or drop_params is True + should_drop: Final = litellm.drop_params is True or drop_params is True - filtered_optional_params = dict(image_edit_optional_params) + filtered_optional_params: Final = dict(image_edit_optional_params) if additional_drop_params: for param in additional_drop_params: filtered_optional_params.pop(param, None) - unsupported_params = [param for param in filtered_optional_params if param not in supported_params] + unsupported_params: Final = [param for param in filtered_optional_params if param not in supported_params] if unsupported_params: if should_drop: @@ -51,7 +51,7 @@ class ImageEditRequestUtils: message=f"The following parameters are not supported for model {model}: {', '.join(unsupported_params)}", ) - mapped_params = image_edit_provider_config.map_openai_params( + mapped_params: Final = image_edit_provider_config.map_openai_params( image_edit_optional_params=cast(ImageEditOptionalRequestParams, filtered_optional_params), model=model, drop_params=should_drop, @@ -61,7 +61,7 @@ class ImageEditRequestUtils: @staticmethod def get_requested_image_edit_optional_param( - params: Dict[str, Any], + params: dict[str, Any], ) -> ImageEditOptionalRequestParams: """ Filter parameters to only include those defined in ImageEditOptionalRequestParams. @@ -72,8 +72,8 @@ class ImageEditRequestUtils: Returns: ImageEditOptionalRequestParams instance with only the valid parameters """ - valid_keys = get_type_hints(ImageEditOptionalRequestParams).keys() - filtered_params = {k: v for k, v in params.items() if k in valid_keys and v is not None} + valid_keys: Final = get_type_hints(ImageEditOptionalRequestParams).keys() + filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None} return cast(ImageEditOptionalRequestParams, filtered_params) @staticmethod @@ -118,13 +118,13 @@ class ImageEditRequestUtils: return FILE_MIME_TYPES[FileType.PNG] # Default fallback # Use the existing get_image_type function to detect image type - image_type_str = get_image_type(bytes_data) + image_type_str: Final = get_image_type(bytes_data) if image_type_str is None: return FILE_MIME_TYPES[FileType.PNG] # Default if detection fails # Map detected type string to FileType enum and get MIME type - type_mapping = { + type_mapping: Final = { "png": FileType.PNG, "jpeg": FileType.JPEG, "gif": FileType.GIF, @@ -132,7 +132,7 @@ class ImageEditRequestUtils: "heic": FileType.HEIC, } - file_type = type_mapping.get(image_type_str) + file_type: Final = type_mapping.get(image_type_str) if file_type is None: return FILE_MIME_TYPES[FileType.PNG] # Default to PNG if unknown diff --git a/litellm/integrations/SlackAlerting/batching_handler.py b/litellm/integrations/SlackAlerting/batching_handler.py index 42f4f562422..a7febdadacd 100644 --- a/litellm/integrations/SlackAlerting/batching_handler.py +++ b/litellm/integrations/SlackAlerting/batching_handler.py @@ -6,7 +6,7 @@ Slack alerts are sent every 10s or when events are greater than X events see custom_batch_logger.py for more details / defaults """ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger @@ -19,7 +19,7 @@ else: def squash_payloads(queue): - squashed = {} + squashed: Final = {} if len(queue) == 0: return squashed if len(queue) == 1: @@ -57,19 +57,19 @@ async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count) """ import json - payload = item.get("payload", {}) + payload: Final = item.get("payload", {}) try: if count > 1: payload["text"] = f"[Num Alerts: {count}]\n\n{payload['text']}" - response = await slackAlertingInstance.async_http_handler.post( + response: Final = await slackAlertingInstance.async_http_handler.post( url=item["url"], headers=item["headers"], data=json.dumps(payload), ) if response.status_code != 200: - verbose_proxy_logger.debug(f"Error sending slack alert to url={item['url']}. Error={response.text}") + verbose_proxy_logger.debug("Error sending slack alert to url=%s. Error=%s", item["url"], response.text) except Exception as e: - verbose_proxy_logger.debug(f"Error sending slack alert: {str(e)}") + verbose_proxy_logger.debug("Error sending slack alert: %s", e) finally: _print_alerting_payload_warning(payload, slackAlertingInstance=slackAlertingInstance) diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index 2a19ec0b7fa..f35ff7b5f82 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Literal +from typing import Final, Literal from litellm.proxy._types import CallInfo, Litellm_EntityType @@ -10,12 +10,10 @@ class BaseBudgetAlertType(ABC): @abstractmethod def get_event_message(self) -> str: """Return the event message for this alert type""" - pass @abstractmethod def get_id(self, user_info: CallInfo) -> str: """Return the ID to use for caching/tracking this alert""" - pass class ProxyBudgetAlert(BaseBudgetAlertType): @@ -99,7 +97,7 @@ def get_budget_alert_type( ) -> BaseBudgetAlertType: """Factory function to get the appropriate budget alert type class""" - alert_types = { + alert_types: Final = { "proxy_budget": ProxyBudgetAlert(), "soft_budget": SoftBudgetAlert(), "user_budget": UserBudgetAlert(), diff --git a/litellm/integrations/SlackAlerting/hanging_request_check.py b/litellm/integrations/SlackAlerting/hanging_request_check.py index 136b6583f38..4d7cbfe8fd1 100644 --- a/litellm/integrations/SlackAlerting/hanging_request_check.py +++ b/litellm/integrations/SlackAlerting/hanging_request_check.py @@ -9,7 +9,7 @@ Notes: import asyncio import time -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_proxy_logger @@ -49,7 +49,7 @@ class AlertingHangingRequestCheck: async def add_request_to_hanging_request_check( self, - request_data: Optional[dict] = None, + request_data: dict | None = None, ): """ Add a request to the hanging request cache. This is the list of request_ids that gets periodicall checked for hanging requests @@ -57,9 +57,9 @@ class AlertingHangingRequestCheck: if request_data is None: return - request_metadata = get_litellm_metadata_from_kwargs(kwargs=request_data) - model = request_data.get("model", "") - api_base: Optional[str] = None + request_metadata: Final = get_litellm_metadata_from_kwargs(kwargs=request_data) + model: Final = request_data.get("model", "") + api_base: str | None = None if request_data.get("deployment", None) is not None and isinstance(request_data["deployment"], dict): api_base = litellm.get_api_base( @@ -67,7 +67,7 @@ class AlertingHangingRequestCheck: optional_params=request_data["deployment"].get("litellm_params", {}), ) - hanging_request_data = HangingRequestData( + hanging_request_data: Final = HangingRequestData( request_id=request_data.get("litellm_call_id", ""), model=model, api_base=api_base, @@ -96,12 +96,12 @@ class AlertingHangingRequestCheck: if proxy_logging_obj.internal_usage_cache is None: return - hanging_requests = await self.hanging_request_cache.async_get_oldest_n_keys( + hanging_requests: Final = await self.hanging_request_cache.async_get_oldest_n_keys( n=MAX_OLDEST_HANGING_REQUESTS_TO_CHECK, ) for request_id in hanging_requests: - hanging_request_data: Optional[HangingRequestData] = await self.hanging_request_cache.async_get_cache( + hanging_request_data: HangingRequestData | None = await self.hanging_request_cache.async_get_cache( key=request_id, ) @@ -112,7 +112,7 @@ class AlertingHangingRequestCheck: continue request_status = await proxy_logging_obj.internal_usage_cache.async_get_cache( - key="request_status:{}".format(hanging_request_data.request_id), + key=f"request_status:{hanging_request_data.request_id}", litellm_parent_otel_span=None, local_only=True, ) @@ -166,7 +166,7 @@ class AlertingHangingRequestCheck: ################ # Send the Alert on Slack ################ - request_info = f"""Request Model: `{hanging_request_data.model}` + request_info: Final = f"""Request Model: `{hanging_request_data.model}` API Base: `{hanging_request_data.api_base}` Key Alias: `{hanging_request_data.key_alias}` Team Alias: `{hanging_request_data.team_alias}`""" diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index e93c650ed97..3e81e7fa92b 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -6,7 +6,7 @@ import os import random import time from datetime import timedelta -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Final, Literal from openai import APIError @@ -61,16 +61,15 @@ class SlackAlerting(CustomBatchLogger): # Class variables or attributes def __init__( self, - internal_usage_cache: Optional[DualCache] = None, - alerting_threshold: Optional[float] = None, # threshold for slow / hanging llm responses (in seconds) - alerting: Optional[List] = [], - alert_types: List[AlertType] = DEFAULT_ALERT_TYPES, - alert_to_webhook_url: Optional[ - Dict[AlertType, Union[List[str], str]] - ] = None, # if user wants to separate alerts to diff channels + internal_usage_cache: DualCache | None = None, + alerting_threshold: float | None = None, # threshold for slow / hanging llm responses (in seconds) + alerting: list | None = [], + alert_types: list[AlertType] = DEFAULT_ALERT_TYPES, + alert_to_webhook_url: dict[AlertType, list[str] | str] + | None = None, # if user wants to separate alerts to diff channels alerting_args={}, - default_webhook_url: Optional[str] = None, - alert_type_config: Optional[Dict[str, dict]] = None, + default_webhook_url: str | None = None, + alert_type_config: dict[str, dict] | None = None, **kwargs, ): if alerting_threshold is None: @@ -89,23 +88,23 @@ class SlackAlerting(CustomBatchLogger): self.hanging_request_check = AlertingHangingRequestCheck( slack_alerting_object=self, ) - self.alert_type_config: Dict[str, AlertTypeConfig] = {} + self.alert_type_config: dict[str, AlertTypeConfig] = {} if alert_type_config: for key, val in alert_type_config.items(): self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val - self.digest_buckets: Dict[str, DigestEntry] = {} + self.digest_buckets: dict[str, DigestEntry] = {} self.digest_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) def update_values( self, - alerting: Optional[List] = None, - alerting_threshold: Optional[float] = None, - alert_types: Optional[List[AlertType]] = None, - alert_to_webhook_url: Optional[Dict[AlertType, Union[List[str], str]]] = None, - alerting_args: Optional[Dict] = None, - llm_router: Optional[Router] = None, - alert_type_config: Optional[Dict[str, dict]] = None, + alerting: list | None = None, + alerting_threshold: float | None = None, + alert_types: list[AlertType] | None = None, + alert_to_webhook_url: dict[AlertType, list[str] | str] | None = None, + alerting_args: dict | None = None, + llm_router: Router | None = None, + alert_type_config: dict[str, dict] | None = None, ): if alerting is not None: self.alerting = alerting @@ -129,26 +128,24 @@ class SlackAlerting(CustomBatchLogger): if self.alert_to_webhook_url is None: self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) else: - _new_values = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) or {} + _new_values: Final = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) or {} self.alert_to_webhook_url.update(_new_values) if llm_router is not None: self.llm_router = llm_router - def _prepare_outage_value_for_cache( - self, outage_value: Union[dict, ProviderRegionOutageModel, OutageModel] - ) -> dict: + def _prepare_outage_value_for_cache(self, outage_value: dict | ProviderRegionOutageModel | OutageModel) -> dict: """ Helper method to prepare outage value for Redis caching. Converts set objects to lists for JSON serialization. """ # Convert to dict for processing - cache_value = dict(outage_value) + cache_value: Final = dict(outage_value) if "deployment_ids" in cache_value and isinstance(cache_value["deployment_ids"], set): cache_value["deployment_ids"] = list(cache_value["deployment_ids"]) return cache_value - def _restore_outage_value_from_cache(self, outage_value: Optional[dict]) -> Optional[dict]: + def _restore_outage_value_from_cache(self, outage_value: dict | None) -> dict | None: """ Helper method to restore outage value after retrieving from cache. Converts list objects back to sets for proper handling. @@ -176,19 +173,19 @@ class SlackAlerting(CustomBatchLogger): end_time, # start/end time ): try: - time_difference = end_time - start_time + time_difference: Final = end_time - start_time # Convert the timedelta to float (in seconds) - time_difference_float = time_difference.total_seconds() - litellm_params = kwargs.get("litellm_params", {}) - model = kwargs.get("model", "") - api_base = litellm.get_api_base(model=model, optional_params=litellm_params) + time_difference_float: Final = time_difference.total_seconds() + litellm_params: Final = kwargs.get("litellm_params", {}) + model: Final = kwargs.get("model", "") + api_base: Final = litellm.get_api_base(model=model, optional_params=litellm_params) messages = kwargs.get("messages", None) # if messages does not exist fallback to "input" if messages is None: messages = kwargs.get("input", None) # only use first 100 chars for alerting - _messages = str(messages)[:100] + _messages: Final = str(messages)[:100] return time_difference_float, model, api_base, _messages except Exception as e: @@ -210,7 +207,7 @@ class SlackAlerting(CustomBatchLogger): _deployment_latencies = metadata["_latency_per_deployment"] if len(_deployment_latencies) == 0: return None - _deployment_latency_map: Optional[dict] = None + _deployment_latency_map: dict | None = None try: # try sorting deployments by latency _deployment_latencies = sorted(_deployment_latencies.items(), key=lambda x: x[1]) @@ -254,10 +251,10 @@ class SlackAlerting(CustomBatchLogger): if time_difference_float > self.alerting_threshold: # add deployment latencies to alert if kwargs is not None and "litellm_params" in kwargs and "metadata" in kwargs["litellm_params"]: - _metadata: dict = kwargs["litellm_params"]["metadata"] + _metadata: Final[dict] = kwargs["litellm_params"]["metadata"] request_info = _add_key_name_and_team_to_alert(request_info=request_info, metadata=_metadata) - _deployment_latency_map = self._get_deployment_latencies_to_alert(metadata=_metadata) + _deployment_latency_map: Final = self._get_deployment_latencies_to_alert(metadata=_metadata) if _deployment_latency_map is not None: request_info += f"\nAvailable Deployment Latencies\n{_deployment_latency_map}" @@ -290,10 +287,7 @@ class SlackAlerting(CustomBatchLogger): ## FAILED REQUESTS ## if deployment_metrics.failed_request: await self.internal_usage_cache.async_increment_cache( - key="{}:{}".format( - deployment_metrics.id, - SlackAlertingCacheKeys.failed_requests_key.value, - ), + key=f"{deployment_metrics.id}:{SlackAlertingCacheKeys.failed_requests_key.value}", value=1, parent_otel_span=None, # no attached request, this is a background operation ) @@ -303,7 +297,7 @@ class SlackAlerting(CustomBatchLogger): ## LATENCY ## if deployment_metrics.latency_per_output_token is not None: await self.internal_usage_cache.async_increment_cache( - key="{}:{}".format(deployment_metrics.id, SlackAlertingCacheKeys.latency_key.value), + key=f"{deployment_metrics.id}:{SlackAlertingCacheKeys.latency_key.value}", value=deployment_metrics.latency_per_output_token, parent_otel_span=None, # no attached request, this is a background operation ) @@ -330,15 +324,15 @@ class SlackAlerting(CustomBatchLogger): False -> if not sent """ - ids = router.get_model_ids() + ids: Final = router.get_model_ids() # get keys - failed_request_keys = ["{}:{}".format(id, SlackAlertingCacheKeys.failed_requests_key.value) for id in ids] - latency_keys = ["{}:{}".format(id, SlackAlertingCacheKeys.latency_key.value) for id in ids] + failed_request_keys: Final = [f"{id}:{SlackAlertingCacheKeys.failed_requests_key.value}" for id in ids] + latency_keys: Final = [f"{id}:{SlackAlertingCacheKeys.latency_key.value}" for id in ids] - combined_metrics_keys = failed_request_keys + latency_keys # reduce cache calls + combined_metrics_keys: Final = failed_request_keys + latency_keys # reduce cache calls - combined_metrics_values = await self.internal_usage_cache.async_batch_get_cache( + combined_metrics_values: Final = await self.internal_usage_cache.async_batch_get_cache( keys=combined_metrics_keys ) # [1, 2, None, ..] @@ -354,8 +348,8 @@ class SlackAlerting(CustomBatchLogger): if all_none: return False - failed_request_values = combined_metrics_values[: len(failed_request_keys)] # # [1, 2, None, ..] - latency_values = combined_metrics_values[len(failed_request_keys) :] + failed_request_values: Final = combined_metrics_values[: len(failed_request_keys)] # # [1, 2, None, ..] + latency_values: Final = combined_metrics_values[len(failed_request_keys) :] # find top 5 failed ## Replace None values with a placeholder value (-1 in this case) @@ -373,7 +367,7 @@ class SlackAlerting(CustomBatchLogger): # find top 5 slowest # Replace None values with a placeholder value (-1 in this case) placeholder_value = 0 - replaced_slowest_values = [value if value is not None else placeholder_value for value in latency_values] + replaced_slowest_values: Final = [value if value is not None else placeholder_value for value in latency_values] # Get the indices of top 5 values with the highest numerical values (ignoring None and 0 values) top_5_slowest = sorted( @@ -426,9 +420,9 @@ class SlackAlerting(CustomBatchLogger): message += f"\t{i + 1}. Deployment: `{deployment_name}`, Latency per output token: `{value}s/token`, API Base: `{api_base}`\n\n" # cache cleanup -> reset values to 0 - latency_cache_keys = [(key, 0) for key in latency_keys] - failed_request_cache_keys = [(key, 0) for key in failed_request_keys] - combined_metrics_cache_keys = latency_cache_keys + failed_request_cache_keys + latency_cache_keys: Final = [(key, 0) for key in latency_keys] + failed_request_cache_keys: Final = [(key, 0) for key in failed_request_keys] + combined_metrics_cache_keys: Final = latency_cache_keys + failed_request_cache_keys await self.internal_usage_cache.async_set_cache_pipeline(cache_list=combined_metrics_cache_keys) message += f"\n\nNext Run is at: `{time.time() + self.alerting_args.daily_report_frequency}`s" @@ -445,7 +439,7 @@ class SlackAlerting(CustomBatchLogger): async def response_taking_too_long( self, - request_data: Optional[dict] = None, + request_data: dict | None = None, ): if self.alerting is None or self.alert_types is None: return @@ -469,10 +463,10 @@ class SlackAlerting(CustomBatchLogger): if "failed_tracking_spend" not in self.alert_types: return - _cache: DualCache = self.internal_usage_cache - message = "Failed Tracking Cost for " + error_message - _cache_key = "budget_alerts:failed_tracking:{}".format(failing_model) - result = await _cache.async_get_cache(key=_cache_key) + _cache: Final[DualCache] = self.internal_usage_cache + message: Final = "Failed Tracking Cost for " + error_message + _cache_key: Final = f"budget_alerts:failed_tracking:{failing_model}" + result: Final = await _cache.async_get_cache(key=_cache_key) if result is None: await self.send_alert( message=message, @@ -512,7 +506,7 @@ class SlackAlerting(CustomBatchLogger): # - Alert once within 24hr period # - Cache this information # - Don't re-alert, if alert already sent - _cache: DualCache = self.internal_usage_cache + _cache: Final[DualCache] = self.internal_usage_cache if self.alerting is None or self.alert_types is None: # do nothing if alerting is not switched on @@ -521,23 +515,18 @@ class SlackAlerting(CustomBatchLogger): return # Get the appropriate budget alert type handler - budget_alert_class = get_budget_alert_type(type) - _id = budget_alert_class.get_id(user_info) - user_info_json = user_info.model_dump(exclude_none=True) - user_info_str = self._get_user_info_str(user_info) + budget_alert_class: Final = get_budget_alert_type(type) + _id: Final = budget_alert_class.get_id(user_info) + user_info_json: Final = user_info.model_dump(exclude_none=True) + user_info_str: Final = self._get_user_info_str(user_info) event_message = budget_alert_class.get_event_message() # Set default event unless we're in projected_limit_exceeded - event: Optional[ - Literal[ - "budget_crossed", - "threshold_crossed", - "projected_limit_exceeded", - "soft_budget_crossed", - ] - ] = "projected_limit_exceeded" if type == "projected_limit_exceeded" else None + event: ( + Literal["budget_crossed", "threshold_crossed", "projected_limit_exceeded", "soft_budget_crossed"] | None + ) = "projected_limit_exceeded" if type == "projected_limit_exceeded" else None - webhook_event: Optional[WebhookEvent] = None + webhook_event: WebhookEvent | None = None # percent of max_budget left to spend if user_info.max_budget is None and user_info.soft_budget is None: @@ -552,8 +541,8 @@ class SlackAlerting(CustomBatchLogger): # send alert if event is not None and user_info.event_group is not None: - _cache_key = "budget_alerts:{}:{}".format(event, _id) - result = await _cache.async_get_cache(key=_cache_key) + _cache_key: Final = f"budget_alerts:{event}:{_id}" + result: Final = await _cache.async_get_cache(key=_cache_key) if result is None: webhook_event = WebhookEvent( event=event, @@ -579,24 +568,10 @@ class SlackAlerting(CustomBatchLogger): def _get_event_and_event_message( self, user_info: CallInfo, - event: Optional[ - Literal[ - "budget_crossed", - "threshold_crossed", - "soft_budget_crossed", - "projected_limit_exceeded", - ] - ], + event: Literal["budget_crossed", "threshold_crossed", "soft_budget_crossed", "projected_limit_exceeded"] | None, event_message: str, - ) -> Tuple[ - Optional[ - Literal[ - "budget_crossed", - "threshold_crossed", - "soft_budget_crossed", - "projected_limit_exceeded", - ] - ], + ) -> tuple[ + Literal["budget_crossed", "threshold_crossed", "soft_budget_crossed", "projected_limit_exceeded"] | None, str, ]: """ @@ -606,7 +581,7 @@ class SlackAlerting(CustomBatchLogger): Handles Max Budget and Soft Budget Alerts """ - percent_left: float = self._get_percent_of_max_budget_left(user_info=user_info) + percent_left: Final[float] = self._get_percent_of_max_budget_left(user_info=user_info) ##################################################################### # SOFT BUDGET CHECK @@ -641,8 +616,8 @@ class SlackAlerting(CustomBatchLogger): Get the percent of the max budget that is left """ percent_left: float = 0.0 - current_spend: float = user_info.spend - max_budget: Optional[float] = user_info.max_budget + current_spend: Final[float] = user_info.spend + max_budget: Final[float | None] = user_info.max_budget if max_budget is None: return percent_left if max_budget <= 0: @@ -654,7 +629,7 @@ class SlackAlerting(CustomBatchLogger): """ Create a standard message for a budget alert """ - _all_fields_as_dict = user_info.model_dump(exclude_none=True) + _all_fields_as_dict: Final = user_info.model_dump(exclude_none=True) _all_fields_as_dict.pop("token") msg = "" for k, v in _all_fields_as_dict.items(): @@ -666,11 +641,11 @@ class SlackAlerting(CustomBatchLogger): async def customer_spend_alert( self, - token: Optional[str], - key_alias: Optional[str], - end_user_id: Optional[str], - response_cost: Optional[float], - max_budget: Optional[float], + token: str | None, + key_alias: str | None, + end_user_id: str | None, + response_cost: float | None, + max_budget: float | None, ): if ( self.alerting is not None @@ -680,7 +655,7 @@ class SlackAlerting(CustomBatchLogger): and response_cost is not None ): # log customer spend - event = WebhookEvent( + event: Final = WebhookEvent( spend=response_cost, max_budget=max_budget, token=token, @@ -693,12 +668,12 @@ class SlackAlerting(CustomBatchLogger): projected_spend=None, event="spend_tracked", event_group=Litellm_EntityType.END_USER, - event_message="Customer spend tracked. Customer={}, spend={}".format(end_user_id, response_cost), + event_message=f"Customer spend tracked. Customer={end_user_id}, spend={response_cost}", ) await self.send_webhook_alert(webhook_event=event) - def _count_outage_alerts(self, alerts: List[int]) -> str: + def _count_outage_alerts(self, alerts: list[int]) -> str: """ Parameters: - alerts: List[int] -> list of error codes (either 408 or 500+) @@ -706,7 +681,7 @@ class SlackAlerting(CustomBatchLogger): Returns: - str -> formatted string. This is an alert message, giving a human-friendly description of the errors. """ - error_breakdown = {"Timeout Errors": 0, "API Errors": 0, "Unknown Errors": 0} + error_breakdown: Final = {"Timeout Errors": 0, "API Errors": 0, "Unknown Errors": 0} for alert in alerts: if alert == 408: error_breakdown["Timeout Errors"] += 1 @@ -718,7 +693,7 @@ class SlackAlerting(CustomBatchLogger): error_msg = "" for key, value in error_breakdown.items(): if value > 0: - error_msg += "\n{}: {}\n".format(key, value) + error_msg += f"\n{key}: {value}\n" return error_msg @@ -728,11 +703,11 @@ class SlackAlerting(CustomBatchLogger): key: Literal["Model", "Region"], key_val: str, provider: str, - api_base: Optional[str], + api_base: str | None, outage_value: BaseOutageModel, ) -> str: """Format an alert message for slack""" - headers = {f"{key} Name": key_val, "Provider": provider} + headers: Final = {f"{key} Name": key_val, "Provider": provider} if api_base is not None: headers["API Base"] = api_base # type: ignore @@ -764,7 +739,7 @@ class SlackAlerting(CustomBatchLogger): if self.llm_router is None: return - deployment = self.llm_router.get_deployment(model_id=deployment_id) + deployment: Final = self.llm_router.get_deployment(model_id=deployment_id) if deployment is None: return @@ -786,11 +761,9 @@ class SlackAlerting(CustomBatchLogger): return ### UNIQUE CACHE KEY ### - cache_key = provider + region_name + cache_key: Final = provider + region_name - outage_value: Optional[ProviderRegionOutageModel] = await self.internal_usage_cache.async_get_cache( - key=cache_key - ) + outage_value: ProviderRegionOutageModel | None = await self.internal_usage_cache.async_get_cache(key=cache_key) # Convert deployment_ids back to set if it was stored as a list if outage_value is not None: @@ -911,7 +884,7 @@ class SlackAlerting(CustomBatchLogger): max_alerts_size = 10 """ try: - outage_value: Optional[OutageModel] = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore + outage_value: OutageModel | None = await self.internal_usage_cache.async_get_cache(key=deployment_id) # type: ignore if ( getattr(exception, "status_code", None) is None or ( @@ -923,7 +896,7 @@ class SlackAlerting(CustomBatchLogger): return ### EXTRACT MODEL DETAILS ### - deployment = self.llm_router.get_deployment(model_id=deployment_id) + deployment: Final = self.llm_router.get_deployment(model_id=deployment_id) if deployment is None: return @@ -934,7 +907,7 @@ class SlackAlerting(CustomBatchLogger): model, provider, _, _ = litellm.get_llm_provider(model=model) except Exception: provider = "" - api_base = litellm.get_api_base(model=model, optional_params=deployment.litellm_params) + api_base: Final = litellm.get_api_base(model=model, optional_params=deployment.litellm_params) if outage_value is None: outage_value = OutageModel( @@ -1006,13 +979,13 @@ class SlackAlerting(CustomBatchLogger): ## update cache ## # Convert set to list for JSON serialization - cache_value = self._prepare_outage_value_for_cache(outage_value) + cache_value: Final = self._prepare_outage_value_for_cache(outage_value) await self.internal_usage_cache.async_set_cache(key=deployment_id, value=cache_value) except Exception: pass async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: Any): - base_model_from_user = getattr(passed_model_info, "base_model", None) + base_model_from_user: Final = getattr(passed_model_info, "base_model", None) model_info = {} base_model = "" if base_model_from_user is not None: @@ -1024,11 +997,11 @@ class SlackAlerting(CustomBatchLogger): for k, v in model_info.items(): if k == "input_cost_per_token" or k == "output_cost_per_token": # when converting to string it should not be 1.63e-06 - v = "{:.8f}".format(v) + v = f"{v:.8f}" model_info_str += f"{k}: {v}\n" - message = f""" + message: Final = f""" *🚅 New Model Added* Model Name: `{model_name}` {base_model} @@ -1058,7 +1031,7 @@ Model Info: ``` """ - alert_val = self.send_alert( + alert_val: Final = self.send_alert( message=message, level="Low", alert_type=AlertType.new_model_added, @@ -1083,14 +1056,14 @@ Model Info: - if WEBHOOK_URL is not set """ - webhook_url = os.getenv("WEBHOOK_URL", None) + webhook_url: Final = os.getenv("WEBHOOK_URL", None) if webhook_url is None: raise Exception("Missing webhook_url from environment") - payload = webhook_event.model_dump_json() - headers = {"Content-type": "application/json"} + payload: Final = webhook_event.model_dump_json() + headers: Final = {"Content-type": "application/json"} - response = await self.async_http_handler.post( + response: Final = await self.async_http_handler.post( url=webhook_url, headers=headers, data=payload, @@ -1105,15 +1078,14 @@ Model Info: async def _check_if_using_premium_email_feature( self, premium_user: bool, - email_logo_url: Optional[str] = None, - email_support_contact: Optional[str] = None, + email_logo_url: str | None = None, + email_support_contact: str | None = None, ): from litellm.proxy.proxy_server import CommonProxyErrors, premium_user if premium_user is not True: if email_logo_url is not None or email_support_contact is not None: raise ValueError(f"Trying to Customize Email Alerting\n {CommonProxyErrors.not_premium_user.value}") - return async def send_key_created_or_user_invited_email(self, webhook_event: WebhookEvent) -> bool: try: @@ -1136,18 +1108,18 @@ Model Info: if email_support_contact is None: email_support_contact = LITELLM_SUPPORT_CONTACT - event_name = webhook_event.event_message + event_name: Final = webhook_event.event_message recipient_email = webhook_event.user_email - recipient_user_id = webhook_event.user_id + recipient_user_id: Final = webhook_event.user_id if recipient_email is None and recipient_user_id is not None and prisma_client is not None: user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": recipient_user_id}) if user_row is not None: recipient_email = user_row.user_email - key_token = webhook_event.token - key_budget = webhook_event.max_budget - base_url = os.getenv("PROXY_BASE_URL", "http://0.0.0.0:4000") + key_token: Final = webhook_event.token + key_budget: Final = webhook_event.max_budget + base_url: Final = os.getenv("PROXY_BASE_URL", "http://0.0.0.0:4000") email_html_content = "Alert from LiteLLM Server" if recipient_email is None: @@ -1167,10 +1139,10 @@ Model Info: ) elif webhook_event.event == "internal_user_created": # GET TEAM NAME - team_id = webhook_event.team_id + team_id: Final = webhook_event.team_id team_name = "Default Team" if team_id is not None and prisma_client is not None: - team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + team_row: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) if team_row is not None: team_name = team_row.team_alias or "-" email_html_content = USER_INVITED_EMAIL_TEMPLATE.format( @@ -1187,7 +1159,7 @@ Model Info: ) webhook_event.model_dump_json() - email_event = { + email_event: Final = { "to": recipient_email, "subject": f"LiteLLM: {event_name}", "html": email_html_content, @@ -1225,10 +1197,10 @@ Model Info: if email_support_contact is None: email_support_contact = LITELLM_SUPPORT_CONTACT - event_name = webhook_event.event_message - recipient_email = webhook_event.user_email - user_name = webhook_event.user_id - max_budget = webhook_event.max_budget + event_name: Final = webhook_event.event_message + recipient_email: Final = webhook_event.user_email + user_name: Final = webhook_event.user_id + max_budget: Final = webhook_event.max_budget email_html_content = "Alert from LiteLLM Server" if recipient_email is None: verbose_proxy_logger.error("Trying to send email alert to no recipient", extra=webhook_event.dict()) @@ -1250,7 +1222,7 @@ Model Info: """ webhook_event.model_dump_json() - email_event = { + email_event: Final = { "to": recipient_email, "subject": f"LiteLLM: {event_name}", "html": email_html_content, @@ -1274,9 +1246,9 @@ Model Info: level: Literal["Low", "Medium", "High"], alert_type: AlertType, alerting_metadata: dict, - user_info: Optional[WebhookEvent] = None, - request_model: Optional[str] = None, - api_base: Optional[str] = None, + user_info: WebhookEvent | None = None, + request_model: str | None = None, + api_base: str | None = None, **kwargs, ): """ @@ -1318,12 +1290,12 @@ Model Info: from datetime import datetime # Check if digest mode is enabled for this alert type - alert_type_name_str = getattr(alert_type, "value", str(alert_type)) - _atc = self.alert_type_config.get(alert_type_name_str) + alert_type_name_str: Final = getattr(alert_type, "value", str(alert_type)) + _atc: Final = self.alert_type_config.get(alert_type_name_str) if _atc is not None and _atc.digest: # Resolve webhook URL for this alert type (needed for digest entry) if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url: - _digest_webhook: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type] + _digest_webhook: str | list[str] | None = self.alert_to_webhook_url[alert_type] elif self.default_webhook_url is not None: _digest_webhook = self.default_webhook_url else: @@ -1331,10 +1303,10 @@ Model Info: if _digest_webhook is None: raise ValueError("Missing SLACK_WEBHOOK_URL from environment") - digest_key = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}" + digest_key: Final = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}" async with self.digest_lock: - now = datetime.now() + now: Final = datetime.now() if digest_key in self.digest_buckets: self.digest_buckets[digest_key]["count"] += 1 self.digest_buckets[digest_key]["last_time"] = now @@ -1353,11 +1325,11 @@ Model Info: return # Suppress immediate alert; will be emitted by _flush_digest_buckets # Get the current timestamp - current_time = datetime.now().strftime("%H:%M:%S") - _proxy_base_url = os.getenv("PROXY_BASE_URL", None) + current_time: Final = datetime.now().strftime("%H:%M:%S") + _proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None) # Use .name if it's an enum, otherwise use as is - alert_type_name = getattr(alert_type, "name", alert_type) - alert_type_formatted = f"Alert type: `{alert_type_name}`" + alert_type_name: Final = getattr(alert_type, "name", alert_type) + alert_type_formatted: Final = f"Alert type: `{alert_type_name}`" if alert_type == "daily_reports" or alert_type == "new_model_added": formatted_message = alert_type_formatted + message else: @@ -1376,7 +1348,7 @@ Model Info: # check if we find the slack webhook url in self.alert_to_webhook_url if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url: - slack_webhook_url: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type] + slack_webhook_url: str | list[str] | None = self.alert_to_webhook_url[alert_type] elif self.default_webhook_url is not None: slack_webhook_url = self.default_webhook_url else: @@ -1384,8 +1356,8 @@ Model Info: if slack_webhook_url is None: raise ValueError("Missing SLACK_WEBHOOK_URL from environment") - payload = {"text": formatted_message} - headers = {"Content-type": "application/json"} + payload: Final = {"text": formatted_message} + headers: Final = {"Content-type": "application/json"} if isinstance(slack_webhook_url, list): for url in slack_webhook_url: @@ -1414,8 +1386,8 @@ Model Info: if not self.log_queue: return - squashed_queue = squash_payloads(self.log_queue) - tasks = [ + squashed_queue: Final = squash_payloads(self.log_queue) + tasks: Final = [ send_to_webhook(slackAlertingInstance=self, item=item["item"], count=item["count"]) for item in squashed_queue.values() ] @@ -1430,8 +1402,8 @@ Model Info: """ from datetime import datetime - now = datetime.now() - flushed_keys: List[str] = [] + now: Final = datetime.now() + flushed_keys: Final[list[str]] = [] async with self.digest_lock: for key, entry in self.digest_buckets.items(): @@ -1495,17 +1467,17 @@ Model Info: try: await self._flush_digest_buckets() except Exception as e: - verbose_proxy_logger.debug(f"Error flushing digest buckets: {str(e)}") + verbose_proxy_logger.debug("Error flushing digest buckets: %s", e) await self.flush_queue() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Log deployment latency""" try: if "daily_reports" in self.alert_types: - litellm_params = kwargs.get("litellm_params", {}) or {} - model_info = litellm_params.get("model_info", {}) or {} - model_id = model_info.get("id", "") or "" - response_s: timedelta = end_time - start_time + litellm_params: Final = kwargs.get("litellm_params", {}) or {} + model_info: Final = litellm_params.get("model_info", {}) or {} + model_id: Final = model_info.get("id", "") or "" + response_s: Final[timedelta] = end_time - start_time final_value = response_s @@ -1514,7 +1486,7 @@ Model Info: and response_obj.usage is not None # type: ignore and hasattr(response_obj.usage, "completion_tokens") # type: ignore ): - completion_tokens = response_obj.usage.completion_tokens # type: ignore + completion_tokens: Final = response_obj.usage.completion_tokens # type: ignore if completion_tokens is not None and completion_tokens > 0: final_value = float(response_s.total_seconds() / completion_tokens) if isinstance(final_value, timedelta): @@ -1530,15 +1502,14 @@ Model Info: ) except Exception as e: verbose_proxy_logger.error( - f"[Non-Blocking Error] Slack Alerting: Got error in logging LLM deployment latency: {str(e)}" + "[Non-Blocking Error] Slack Alerting: Got error in logging LLM deployment latency: %s", e ) - pass async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """Log failure + deployment latency""" - _litellm_params = kwargs.get("litellm_params", {}) - _model_info = _litellm_params.get("model_info", {}) or {} - model_id = _model_info.get("id", "") + _litellm_params: Final = kwargs.get("litellm_params", {}) + _model_info: Final = _litellm_params.get("model_info", {}) or {} + model_id: Final = _model_info.get("id", "") try: if "daily_reports" in self.alert_types: try: @@ -1551,7 +1522,7 @@ Model Info: ) ) except Exception as e: - verbose_logger.debug(f"Exception raises -{str(e)}") + verbose_logger.debug("Exception raises -%s", e) if isinstance(kwargs.get("exception", ""), APIError): if "outage_alerts" in self.alert_types: @@ -1573,12 +1544,12 @@ Model Info: """ report_sent_bool = False - report_sent = await self.internal_usage_cache.async_get_cache( + report_sent: Final = await self.internal_usage_cache.async_get_cache( key=SlackAlertingCacheKeys.report_sent_key.value, parent_otel_span=None, ) # None | float - current_time = time.time() + current_time: Final = time.time() if report_sent is None: await self.internal_usage_cache.async_set_cache( @@ -1587,7 +1558,7 @@ Model Info: ) elif isinstance(report_sent, float): # Check if current time - interval >= time last sent - interval_seconds = self.alerting_args.daily_report_frequency + interval_seconds: Final = self.alerting_args.daily_report_frequency if current_time - report_sent >= interval_seconds: # Sneak in the reporting logic here @@ -1601,7 +1572,7 @@ Model Info: return report_sent_bool - async def _run_scheduled_daily_report(self, llm_router: Optional[Any] = None): + async def _run_scheduled_daily_report(self, llm_router: Any | None = None): """ If 'daily_reports' enabled @@ -1641,20 +1612,20 @@ Model Info: ) # Parse the time range - days = int(time_range[:-1]) + days: Final = int(time_range[:-1]) if time_range[-1].lower() != "d": raise ValueError("Time range must be specified in days, e.g., '7d'") - todays_date = datetime.datetime.now().date() - start_date = todays_date - datetime.timedelta(days=days) + todays_date: Final = datetime.datetime.now().date() + start_date: Final = todays_date - datetime.timedelta(days=days) - _event_cache_key = ( + _event_cache_key: Final = ( f"weekly_spend_report_sent_{start_date.strftime('%Y-%m-%d')}_{todays_date.strftime('%Y-%m-%d')}" ) if await self.internal_usage_cache.async_get_cache(key=_event_cache_key): return - _resp = await _get_spend_report_for_time_range( + _resp: Final = await _get_spend_report_for_time_range( start_date=start_date.strftime("%Y-%m-%d"), end_date=todays_date.strftime("%Y-%m-%d"), ) @@ -1691,9 +1662,9 @@ Model Info: ) except ValueError as ve: - verbose_proxy_logger.error(f"Invalid time range format: {ve}") + verbose_proxy_logger.error("Invalid time range format: %s", ve) except Exception as e: - verbose_proxy_logger.error(f"Error sending spend report: {e}") + verbose_proxy_logger.error("Error sending spend report: %s", e) async def send_monthly_spend_report(self): """ """ @@ -1704,8 +1675,8 @@ Model Info: _get_spend_report_for_time_range, ) - todays_date = datetime.datetime.now().date() - first_day_of_month = todays_date.replace(day=1) + todays_date: Final = datetime.datetime.now().date() + first_day_of_month: Final = todays_date.replace(day=1) _, last_day_of_month = monthrange(todays_date.year, todays_date.month) last_day_of_month = first_day_of_month + datetime.timedelta(days=last_day_of_month - 1) @@ -1713,7 +1684,7 @@ Model Info: if await self.internal_usage_cache.async_get_cache(key=_event_cache_key): return - _resp = await _get_spend_report_for_time_range( + _resp: Final = await _get_spend_report_for_time_range( start_date=first_day_of_month.strftime("%Y-%m-%d"), end_date=last_day_of_month.strftime("%Y-%m-%d"), ) @@ -1771,9 +1742,9 @@ Model Info: ) # call prometheuslogger. - falllback_success_info_prometheus = await get_fallback_metric_from_prometheus() + falllback_success_info_prometheus: Final = await get_fallback_metric_from_prometheus() - fallback_message = f"*Fallback Statistics:*\n{falllback_success_info_prometheus}" + fallback_message: Final = f"*Fallback Statistics:*\n{falllback_success_info_prometheus}" await self.send_alert( message=fallback_message, @@ -1785,8 +1756,6 @@ Model Info: except Exception as e: verbose_proxy_logger.error("Error sending weekly spend report %s", e) - pass - async def send_virtual_key_event_slack( self, key_event: VirtualKeyEvent, @@ -1804,7 +1773,7 @@ Model Info: try: message = f"`{event_name}`\n" - key_event_dict = key_event.model_dump() + key_event_dict: Final = key_event.model_dump() # Add Created by information first message += "*Action Done by:*\n" @@ -1814,7 +1783,7 @@ Model Info: # Add args sent to function in the alert message += "\n*Arguments passed:*\n" - request_kwargs = key_event.request_kwargs + request_kwargs: Final = key_event.request_kwargs for key, value in request_kwargs.items(): if key == "user_api_key_dict": continue @@ -1830,9 +1799,7 @@ Model Info: except Exception as e: verbose_proxy_logger.error("Error sending send_virtual_key_event_slack %s", e) - return - - async def _request_is_completed(self, request_data: Optional[dict]) -> bool: + async def _request_is_completed(self, request_data: dict | None) -> bool: """ Returns True if the request is completed - either as a success or failure """ @@ -1841,9 +1808,9 @@ Model Info: if request_data.get("litellm_status", "") != "success" and request_data.get("litellm_status", "") != "fail": ## CHECK IF CACHE IS UPDATED - litellm_call_id = request_data.get("litellm_call_id", "") - status: Optional[str] = await self.internal_usage_cache.async_get_cache( - key="request_status:{}".format(litellm_call_id), local_only=True + litellm_call_id: Final = request_data.get("litellm_call_id", "") + status: Final[str | None] = await self.internal_usage_cache.async_get_cache( + key=f"request_status:{litellm_call_id}", local_only=True ) if status is not None and (status == "success" or status == "fail"): return True diff --git a/litellm/integrations/SlackAlerting/utils.py b/litellm/integrations/SlackAlerting/utils.py index 4424bedba81..297d069a868 100644 --- a/litellm/integrations/SlackAlerting/utils.py +++ b/litellm/integrations/SlackAlerting/utils.py @@ -3,7 +3,7 @@ Utils used for slack alerting """ import asyncio -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Any, Final import litellm from litellm.proxy._types import AlertType @@ -18,8 +18,8 @@ else: def process_slack_alerting_variables( - alert_to_webhook_url: Optional[Dict[AlertType, Union[List[str], str]]], -) -> Optional[Dict[AlertType, Union[List[str], str]]]: + alert_to_webhook_url: dict[AlertType, list[str] | str] | None, +) -> dict[AlertType, list[str] | str] | None: """ process alert_to_webhook_url - check if any urls are set as os.environ/SLACK_WEBHOOK_URL_1 read env var and set the correct value @@ -29,7 +29,7 @@ def process_slack_alerting_variables( for alert_type, webhook_urls in alert_to_webhook_url.items(): if isinstance(webhook_urls, list): - _webhook_values: List[str] = [] + _webhook_values: list[str] = [] for webhook_url in webhook_urls: if "os.environ/" in webhook_url: _env_value = get_secret(secret_name=webhook_url) @@ -56,8 +56,8 @@ def process_slack_alerting_variables( async def _add_langfuse_trace_id_to_alert( - request_data: Optional[dict] = None, -) -> Optional[str]: + request_data: dict | None = None, +) -> str | None: """ Returns langfuse trace url @@ -73,8 +73,8 @@ async def _add_langfuse_trace_id_to_alert( ######################################################### if request_data is not None and request_data.get("litellm_logging_obj", None) is not None: - trace_id: Optional[str] = None - litellm_logging_obj: Logging = request_data["litellm_logging_obj"] + trace_id: str | None = None + litellm_logging_obj: Final[Logging] = request_data["litellm_logging_obj"] for _ in range(3): trace_id = litellm_logging_obj._get_trace_id(service_name="langfuse") @@ -82,9 +82,9 @@ async def _add_langfuse_trace_id_to_alert( break await asyncio.sleep(3) # wait 3s before retrying for trace id ######################################################### - langfuse_object = litellm_logging_obj._get_callback_object(service_name="langfuse") + langfuse_object: Final = litellm_logging_obj._get_callback_object(service_name="langfuse") if langfuse_object is not None: - base_url = langfuse_object.Langfuse.base_url + base_url: Final = langfuse_object.Langfuse.base_url return f"{base_url}/trace/{trace_id}" return None diff --git a/litellm/integrations/additional_logging_utils.py b/litellm/integrations/additional_logging_utils.py index 59319140a18..3f79a8ac007 100644 --- a/litellm/integrations/additional_logging_utils.py +++ b/litellm/integrations/additional_logging_utils.py @@ -7,7 +7,6 @@ Base class for Additional Logging Utils for CustomLoggers from abc import ABC, abstractmethod from datetime import datetime -from typing import Optional from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus @@ -21,15 +20,14 @@ class AdditionalLoggingUtils(ABC): """ Check if the service is healthy """ - pass @abstractmethod async def get_request_response_payload( self, request_id: str, - start_time_utc: Optional[datetime], - end_time_utc: Optional[datetime], - ) -> Optional[dict]: + start_time_utc: datetime | None, + end_time_utc: datetime | None, + ) -> dict | None: """ Get the request and response payload for a given `request_id` """ diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py index c60e5cb0e2a..399ee49238c 100644 --- a/litellm/integrations/agentops/agentops.py +++ b/litellm/integrations/agentops/agentops.py @@ -4,7 +4,8 @@ AgentOps integration for LiteLLM - Provides OpenTelemetry tracing for LLM calls import os from dataclasses import dataclass -from typing import Optional, Dict, Any +from typing import Any, Final + from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -12,9 +13,9 @@ from litellm.llms.custom_httpx.http_handler import _get_httpx_client @dataclass class AgentOpsConfig: endpoint: str = "https://otlp.agentops.cloud/v1/traces" - api_key: Optional[str] = None - service_name: Optional[str] = None - deployment_environment: Optional[str] = None + api_key: str | None = None + service_name: str | None = None + deployment_environment: str | None = None auth_endpoint: str = "https://api.agentops.ai/v3/auth/token" @classmethod @@ -47,7 +48,7 @@ class AgentOps(OpenTelemetry): def __init__( self, - config: Optional[AgentOpsConfig] = None, + config: AgentOpsConfig | None = None, ): if config is None: config = AgentOpsConfig.from_env() @@ -57,21 +58,21 @@ class AgentOps(OpenTelemetry): project_id = None if config.api_key: try: - response = self._fetch_auth_token(config.api_key, config.auth_endpoint) + response: Final = self._fetch_auth_token(config.api_key, config.auth_endpoint) jwt_token = response.get("token") project_id = response.get("project_id") except Exception: pass - headers = f"Authorization=Bearer {jwt_token}" if jwt_token else None + headers: Final = f"Authorization=Bearer {jwt_token}" if jwt_token else None - otel_config = OpenTelemetryConfig(exporter="otlp_http", endpoint=config.endpoint, headers=headers) + otel_config: Final = OpenTelemetryConfig(exporter="otlp_http", endpoint=config.endpoint, headers=headers) # Initialize OpenTelemetry with our config super().__init__(config=otel_config, callback_name="agentops") # Set AgentOps-specific resource attributes - resource_attrs = { + resource_attrs: Final = { "service.name": config.service_name or "litellm", "deployment.environment": config.deployment_environment or "production", "telemetry.sdk.name": "agentops", @@ -82,7 +83,7 @@ class AgentOps(OpenTelemetry): self.resource_attributes = resource_attrs - def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> Dict[str, Any]: + def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> dict[str, Any]: """ Fetch JWT authentication token from AgentOps API @@ -93,14 +94,14 @@ class AgentOps(OpenTelemetry): Returns: Dict containing JWT token and project ID """ - headers = { + headers: Final = { "Content-Type": "application/json", "Connection": "keep-alive", } - client = _get_httpx_client() + client: Final = _get_httpx_client() try: - response = client.post( + response: Final = client.post( url=auth_endpoint, headers=headers, json={"api_key": api_key}, diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index faedf8ae1a3..34b3c4dacde 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -10,7 +10,7 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and """ import copy -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger @@ -32,24 +32,24 @@ else: # Anthropic (and Bedrock Claude) reject requests with more than 4 cache_control # breakpoints: "A maximum of 4 blocks with cache_control may be provided." -MAX_CACHE_CONTROL_BLOCKS = 4 +MAX_CACHE_CONTROL_BLOCKS: Final = 4 class AnthropicCacheControlHook(CustomPromptManagement): def get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Apply cache control directives based on specified injection points. @@ -59,7 +59,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): - non_default_params: dict - params with any global cache controls """ # Extract cache control injection points - injection_points: List[CacheControlInjectionPoint] = non_default_params.pop( + injection_points: Final[list[CacheControlInjectionPoint]] = non_default_params.pop( "cache_control_injection_points", [] ) if not injection_points: @@ -69,8 +69,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): processed_messages = copy.deepcopy(messages) # Separate message-level and non-message-level injection points - message_points: List[CacheControlMessageInjectionPoint] = [] - remaining_points: List[CacheControlInjectionPoint] = [] + message_points: Final[list[CacheControlMessageInjectionPoint]] = [] + remaining_points: Final[list[CacheControlInjectionPoint]] = [] for point in injection_points: if point.get("location") == "message": message_points.append(cast(CacheControlMessageInjectionPoint, point)) @@ -81,7 +81,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): # provider transform, where each tool_config point appends at most one # cachePoint to the tools. That block also counts toward Anthropic's # limit, so reserve a slot for it here to leave room. - reserved_blocks = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0 + reserved_blocks: Final = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0 processed_messages = self._apply_message_injections( points=message_points, @@ -99,10 +99,10 @@ class AnthropicCacheControlHook(CustomPromptManagement): @staticmethod def _apply_message_injections( - points: List[CacheControlMessageInjectionPoint], - messages: List[AllMessageValues], + points: list[CacheControlMessageInjectionPoint], + messages: list[AllMessageValues], max_blocks: int, - ) -> List[AllMessageValues]: + ) -> list[AllMessageValues]: """Apply message-level cache control injection points in order. Anthropic allows at most ``MAX_CACHE_CONTROL_BLOCKS`` cache_control @@ -143,19 +143,19 @@ class AnthropicCacheControlHook(CustomPromptManagement): if limit_reached: verbose_logger.warning( - f"AnthropicCacheControlHook: Reached the Anthropic limit of " - f"{MAX_CACHE_CONTROL_BLOCKS} cache_control blocks. Skipping further injection." + "AnthropicCacheControlHook: Reached the Anthropic limit of %s cache_control blocks. Skipping further injection.", + MAX_CACHE_CONTROL_BLOCKS, ) return messages @staticmethod def _resolve_target_indices( - point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues] - ) -> List[int]: + point: CacheControlMessageInjectionPoint, messages: list[AllMessageValues] + ) -> list[int]: """Resolve which message indices an injection point targets.""" - _targetted_index: Optional[Union[int, str]] = point.get("index", None) - targetted_index: Optional[int] = None + _targetted_index: Final[int | str | None] = point.get("index", None) + targetted_index: int | None = None if isinstance(_targetted_index, str): try: targetted_index = int(_targetted_index) @@ -166,7 +166,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): # Case 1: Target by specific index if targetted_index is not None: - original_index = targetted_index + original_index: Final = targetted_index if targetted_index < 0: targetted_index += len(messages) @@ -174,13 +174,15 @@ class AnthropicCacheControlHook(CustomPromptManagement): return [targetted_index] verbose_logger.warning( - f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. " - f"Targeted index was {targetted_index}. Skipping cache control injection for this point." + "AnthropicCacheControlHook: Provided index %s is out of bounds for message list of length %s. Targeted index was %s. Skipping cache control injection for this point.", + original_index, + len(messages), + targetted_index, ) return [] # Case 2: Target by role - targetted_role = point.get("role", None) + targetted_role: Final = point.get("role", None) if targetted_role is not None: return [idx for idx, msg in enumerate(messages) if msg.get("role") == targetted_role] @@ -192,7 +194,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): count = 0 if message.get("cache_control") is not None: count += 1 - content = message.get("content") + content: Final = message.get("content") if isinstance(content, list): for block in content: if isinstance(block, dict) and block.get("cache_control") is not None: @@ -219,7 +221,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): Per Anthropic's API specification, when using multiple content blocks, only the last content block can have cache_control. """ - message_content = message.get("content", None) + message_content: Final = message.get("content", None) # 1. if string, insert cache control in the message if isinstance(message_content, str): @@ -232,10 +234,10 @@ class AnthropicCacheControlHook(CustomPromptManagement): @staticmethod def apply_to_anthropic_messages_request( - messages: List[Dict], + messages: list[dict], system: str | list | None, - injection_points: List[CacheControlInjectionPoint], - ) -> Tuple[List[Dict], str | list | None, List[CacheControlInjectionPoint]]: + injection_points: list[CacheControlInjectionPoint], + ) -> tuple[list[dict], str | list | None, list[CacheControlInjectionPoint]]: """Apply cache control injection for the Anthropic-native v1/messages endpoint. Returns (messages, system, remaining_non_message_points). @@ -243,12 +245,12 @@ class AnthropicCacheControlHook(CustomPromptManagement): if not injection_points: return messages, system, [] - processed_messages: List[Dict] = copy.deepcopy(messages) + processed_messages: list[dict] = copy.deepcopy(messages) processed_system = copy.deepcopy(system) if system is not None else None - message_points: List[CacheControlMessageInjectionPoint] = [] - system_points: List[CacheControlMessageInjectionPoint] = [] - remaining_points: List[CacheControlInjectionPoint] = [] + message_points: Final[list[CacheControlMessageInjectionPoint]] = [] + system_points: Final[list[CacheControlMessageInjectionPoint]] = [] + remaining_points: Final[list[CacheControlInjectionPoint]] = [] for point in injection_points: if point.get("location") == "message": @@ -260,8 +262,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): else: remaining_points.append(point) - reserved_blocks = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0 - max_blocks = MAX_CACHE_CONTROL_BLOCKS - reserved_blocks + reserved_blocks: Final = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0 + max_blocks: Final = MAX_CACHE_CONTROL_BLOCKS - reserved_blocks used_blocks = sum( AnthropicCacheControlHook._count_cache_control_blocks(cast(AllMessageValues, msg)) @@ -273,11 +275,11 @@ class AnthropicCacheControlHook(CustomPromptManagement): ) if system_points and processed_system is not None and used_blocks < max_blocks: - system_already_has_cc = isinstance(processed_system, list) and any( + system_already_has_cc: Final = isinstance(processed_system, list) and any( isinstance(b, dict) and b.get("cache_control") is not None for b in processed_system ) if not system_already_has_cc: - control = system_points[0].get("control") or ChatCompletionCachedContent(type="ephemeral") + control: Final = system_points[0].get("control") or ChatCompletionCachedContent(type="ephemeral") if isinstance(processed_system, str): processed_system = [{"type": "text", "text": processed_system, "cache_control": control}] used_blocks += 1 @@ -292,7 +294,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): processed_messages = AnthropicCacheControlHook._apply_message_injections( points=message_points, - messages=cast(List[AllMessageValues], processed_messages), + messages=cast(list[AllMessageValues], processed_messages), max_blocks=max_blocks - used_blocks, ) @@ -307,7 +309,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): """ import litellm - ttl = litellm.anthropic_prompt_caching_ttl + ttl: Final = litellm.anthropic_prompt_caching_ttl if ttl == "5m" or ttl == "1h": return ChatCompletionCachedContent(type="ephemeral", ttl=ttl) return ChatCompletionCachedContent(type="ephemeral") @@ -417,8 +419,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools): return [] - control = AnthropicCacheControlHook._default_control() - points: list[CacheControlInjectionPoint] = [ + control: Final = AnthropicCacheControlHook._default_control() + points: Final[list[CacheControlInjectionPoint]] = [ CacheControlMessageInjectionPoint(location="message", role="system", index=None, control=control), CacheControlMessageInjectionPoint(location="message", role=None, index=-1, control=control), ] @@ -450,7 +452,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): ): non_default_params.pop("cache_control_injection_points") return - points = AnthropicCacheControlHook.get_default_injection_points( + points: Final = AnthropicCacheControlHook.get_default_injection_points( messages=messages, system=None, model=model, @@ -462,13 +464,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): @staticmethod def maybe_inject_cache_control( - messages: List[Dict], + messages: list[dict], system: str | list | None, - kwargs: Dict[str, Any], + kwargs: dict[str, Any], model: str | None = None, custom_llm_provider: str | None = None, tools: list[dict] | None = None, - ) -> Tuple[List[Dict], str | list | None]: + ) -> tuple[list[dict], str | list | None]: """Extract cache_control_injection_points from kwargs and apply if present. Configured points stand down entirely when the client already marked @@ -482,7 +484,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): downstream transforms can handle them. """ typed_messages = cast(list[AllMessageValues], messages) # cast-ok: Anthropic-shaped dicts from v1/messages - configured = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list + configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) ) if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools): @@ -515,8 +517,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): def should_run_prompt_management( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], + prompt_id: str | None, + prompt_spec: PromptSpec | None, dynamic_callback_params: StandardCallbackDynamicParams, ) -> bool: """Always return False since this is not a true prompt management system.""" @@ -524,12 +526,12 @@ class AnthropicCacheControlHook(CustomPromptManagement): def _compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_spec: PromptSpec | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """Not used - this hook only modifies messages, doesn't fetch prompts.""" return PromptManagementClient( @@ -542,12 +544,12 @@ class AnthropicCacheControlHook(CustomPromptManagement): async def async_compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """Not used - this hook only modifies messages, doesn't fetch prompts.""" return self._compile_prompt_helper( @@ -562,19 +564,19 @@ class AnthropicCacheControlHook(CustomPromptManagement): async def async_get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, litellm_logging_obj: LiteLLMLoggingObj, - prompt_spec: Optional[PromptSpec] = None, - tools: Optional[List[Dict]] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + tools: list[dict] | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """Async version - delegates to sync since no async operations needed.""" return self.get_chat_completion_prompt( model=model, @@ -591,15 +593,15 @@ class AnthropicCacheControlHook(CustomPromptManagement): ) @staticmethod - def should_use_anthropic_cache_control_hook(non_default_params: Dict) -> bool: + def should_use_anthropic_cache_control_hook(non_default_params: dict) -> bool: if non_default_params.get("cache_control_injection_points", None): return True return False @staticmethod def get_custom_logger_for_anthropic_cache_control_hook( - non_default_params: Dict, - ) -> Optional[CustomLogger]: + non_default_params: dict, + ) -> CustomLogger | None: from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, ) diff --git a/litellm/integrations/argilla.py b/litellm/integrations/argilla.py index a86b6f9e388..76a63f75897 100644 --- a/litellm/integrations/argilla.py +++ b/litellm/integrations/argilla.py @@ -7,7 +7,7 @@ import json import os import random import types -from typing import Any, Dict, List, Optional +from typing import Any, Final import httpx from pydantic import BaseModel # type: ignore @@ -29,7 +29,7 @@ from litellm.types.utils import StandardLoggingPayload def is_serializable(value): - non_serializable_types = ( + non_serializable_types: Final = ( types.CoroutineType, types.FunctionType, types.GeneratorType, @@ -41,9 +41,9 @@ def is_serializable(value): class ArgillaLogger(CustomBatchLogger): def __init__( self, - argilla_api_key: Optional[str] = None, - argilla_dataset_name: Optional[str] = None, - argilla_base_url: Optional[str] = None, + argilla_api_key: str | None = None, + argilla_dataset_name: str | None = None, + argilla_base_url: str | None = None, **kwargs, ): if litellm.argilla_transformation_object is None: @@ -62,14 +62,14 @@ class ArgillaLogger(CustomBatchLogger): ) self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - _batch_size = os.getenv("ARGILLA_BATCH_SIZE", None) or litellm.argilla_batch_size + _batch_size: Final = os.getenv("ARGILLA_BATCH_SIZE", None) or litellm.argilla_batch_size if _batch_size: self.batch_size = int(_batch_size) asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) - def validate_argilla_transformation_object(self, argilla_transformation_object: Dict[str, Any]): + def validate_argilla_transformation_object(self, argilla_transformation_object: dict[str, Any]): if not isinstance(argilla_transformation_object, dict): raise Exception("'argilla_transformation_object' must be a dictionary, to log your payload to Argilla.") @@ -81,15 +81,15 @@ class ArgillaLogger(CustomBatchLogger): def get_credentials_from_env( self, - argilla_api_key: Optional[str], - argilla_dataset_name: Optional[str], - argilla_base_url: Optional[str], + argilla_api_key: str | None, + argilla_dataset_name: str | None, + argilla_base_url: str | None, ) -> ArgillaCredentialsObject: - _credentials_api_key = argilla_api_key or os.getenv("ARGILLA_API_KEY") + _credentials_api_key: Final = argilla_api_key or os.getenv("ARGILLA_API_KEY") if _credentials_api_key is None: raise Exception("Invalid Argilla API Key given. _credentials_api_key=None.") - _credentials_base_url = argilla_base_url or os.getenv("ARGILLA_BASE_URL") or "http://localhost:6900/" + _credentials_base_url: Final = argilla_base_url or os.getenv("ARGILLA_BASE_URL") or "http://localhost:6900/" if _credentials_base_url is None: raise Exception("Invalid Argilla Base URL given. _credentials_base_url=None.") @@ -97,11 +97,11 @@ class ArgillaLogger(CustomBatchLogger): if _credentials_dataset_name is None: raise Exception("Invalid Argilla Dataset give. Value=None.") else: - dataset_response = litellm.module_level_client.get( + dataset_response: Final = litellm.module_level_client.get( url=f"{_credentials_base_url}/api/v1/me/datasets?name={_credentials_dataset_name}", headers={"X-Argilla-Api-Key": _credentials_api_key}, ) - json_response = dataset_response.json() + json_response: Final = dataset_response.json() if ( "items" in json_response and isinstance(json_response["items"], list) @@ -115,8 +115,8 @@ class ArgillaLogger(CustomBatchLogger): ARGILLA_DATASET_NAME=_credentials_dataset_name, ) - def get_chat_messages(self, payload: StandardLoggingPayload) -> List[Dict[str, Any]]: - payload_messages = payload.get("messages", None) + def get_chat_messages(self, payload: StandardLoggingPayload) -> list[dict[str, Any]]: + payload_messages: Final = payload.get("messages", None) if payload_messages is None: raise Exception("No chat messages found in payload.") @@ -129,7 +129,7 @@ class ArgillaLogger(CustomBatchLogger): raise Exception(f"Invalid chat messages format: {payload_messages}") def get_str_response(self, payload: StandardLoggingPayload) -> str: - response = payload["response"] + response: Final = payload["response"] if response is None: raise Exception("No response found in payload.") @@ -141,17 +141,17 @@ class ArgillaLogger(CustomBatchLogger): else: raise Exception(f"Invalid response format: {response}") - def _prepare_log_data(self, kwargs, response_obj, start_time, end_time) -> Optional[ArgillaItem]: + def _prepare_log_data(self, kwargs, response_obj, start_time, end_time) -> ArgillaItem | None: try: # Ensure everything in the payload is converted to str - payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) + payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if payload is None: raise Exception("Error logging request payload. Payload=none.") - argilla_message = self.get_chat_messages(payload) - argilla_response = self.get_str_response(payload) - argilla_item: ArgillaItem = {"fields": {}} + argilla_message: Final = self.get_chat_messages(payload) + argilla_response: Final = self.get_str_response(payload) + argilla_item: Final[ArgillaItem] = {"fields": {}} for k, v in self.argilla_transformation_object.items(): if v == "messages": argilla_item["fields"][k] = argilla_message @@ -168,26 +168,26 @@ class ArgillaLogger(CustomBatchLogger): if not self.log_queue: return - argilla_api_base = self.default_credentials["ARGILLA_BASE_URL"] - argilla_dataset_name = self.default_credentials["ARGILLA_DATASET_NAME"] + argilla_api_base: Final = self.default_credentials["ARGILLA_BASE_URL"] + argilla_dataset_name: Final = self.default_credentials["ARGILLA_DATASET_NAME"] - url = f"{argilla_api_base}/api/v1/datasets/{argilla_dataset_name}/records/bulk" + url: Final = f"{argilla_api_base}/api/v1/datasets/{argilla_dataset_name}/records/bulk" - argilla_api_key = self.default_credentials["ARGILLA_API_KEY"] + argilla_api_key: Final = self.default_credentials["ARGILLA_API_KEY"] - headers = {"X-Argilla-Api-Key": argilla_api_key} + headers: Final = {"X-Argilla-Api-Key": argilla_api_key} try: - response = litellm.module_level_client.post( + response: Final = litellm.module_level_client.post( url=url, json=self.log_queue, headers=headers, ) if response.status_code >= 300: - verbose_logger.error(f"Argilla Error: {response.status_code} - {response.text}") + verbose_logger.error("Argilla Error: %s - %s", response.status_code, response.text) else: - verbose_logger.debug(f"Batch of {len(self.log_queue)} runs successfully created") + verbose_logger.debug("Batch of %s runs successfully created", len(self.log_queue)) self.log_queue.clear() except Exception: @@ -195,18 +195,16 @@ class ArgillaLogger(CustomBatchLogger): def log_success_event(self, kwargs, response_obj, start_time, end_time): try: - sampling_rate = ( + sampling_rate: Final = ( float(os.getenv("LANGSMITH_SAMPLING_RATE")) # type: ignore if os.getenv("LANGSMITH_SAMPLING_RATE") is not None and os.getenv("LANGSMITH_SAMPLING_RATE").strip().isdigit() # type: ignore else 1.0 ) - random_sample = random.random() + random_sample: Final = random.random() if random_sample > sampling_rate: verbose_logger.info( - "Skipping Langsmith logging. Sampling rate={}, random_sample={}".format( - sampling_rate, random_sample - ) + "Skipping Langsmith logging. Sampling rate=%s, random_sample=%s", sampling_rate, random_sample ) return # Skip logging verbose_logger.debug( @@ -214,12 +212,12 @@ class ArgillaLogger(CustomBatchLogger): kwargs, response_obj, ) - data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) + data: Final = self._prepare_log_data(kwargs, response_obj, start_time, end_time) if data is None: return self.log_queue.append(data) - verbose_logger.debug(f"Langsmith, event added to queue. Will flush in {self.flush_interval} seconds...") + verbose_logger.debug("Langsmith, event added to queue. Will flush in %s seconds...", self.flush_interval) if len(self.log_queue) >= self.batch_size: self._send_batch() @@ -229,13 +227,11 @@ class ArgillaLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - sampling_rate = self.sampling_rate - random_sample = random.random() + sampling_rate: Final = self.sampling_rate + random_sample: Final = random.random() if random_sample > sampling_rate: verbose_logger.info( - "Skipping Langsmith logging. Sampling rate={}, random_sample={}".format( - sampling_rate, random_sample - ) + "Skipping Langsmith logging. Sampling rate=%s, random_sample=%s", sampling_rate, random_sample ) return # Skip logging verbose_logger.debug( @@ -243,7 +239,7 @@ class ArgillaLogger(CustomBatchLogger): kwargs, response_obj, ) - payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) + payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) @@ -272,16 +268,16 @@ class ArgillaLogger(CustomBatchLogger): verbose_logger.exception("Argilla Layer Error - error logging async success event.") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - sampling_rate = self.sampling_rate - random_sample = random.random() + sampling_rate: Final = self.sampling_rate + random_sample: Final = random.random() if random_sample > sampling_rate: verbose_logger.info( - "Skipping Langsmith logging. Sampling rate={}, random_sample={}".format(sampling_rate, random_sample) + "Skipping Langsmith logging. Sampling rate=%s, random_sample=%s", sampling_rate, random_sample ) return # Skip logging verbose_logger.info("Langsmith Failure Event Logging!") try: - data = self._prepare_log_data(kwargs, response_obj, start_time, end_time) + data: Final = self._prepare_log_data(kwargs, response_obj, start_time, end_time) self.log_queue.append(data) verbose_logger.debug( "Langsmith logging: queue length %s, batch size %s", @@ -306,17 +302,17 @@ class ArgillaLogger(CustomBatchLogger): if not self.log_queue: return - argilla_api_base = self.default_credentials["ARGILLA_BASE_URL"] - argilla_dataset_name = self.default_credentials["ARGILLA_DATASET_NAME"] + argilla_api_base: Final = self.default_credentials["ARGILLA_BASE_URL"] + argilla_dataset_name: Final = self.default_credentials["ARGILLA_DATASET_NAME"] - url = f"{argilla_api_base}/api/v1/datasets/{argilla_dataset_name}/records/bulk" + url: Final = f"{argilla_api_base}/api/v1/datasets/{argilla_dataset_name}/records/bulk" - argilla_api_key = self.default_credentials["ARGILLA_API_KEY"] + argilla_api_key: Final = self.default_credentials["ARGILLA_API_KEY"] - headers = {"X-Argilla-Api-Key": argilla_api_key} + headers: Final = {"X-Argilla-Api-Key": argilla_api_key} try: - response = await self.async_httpx_client.put( + response: Final = await self.async_httpx_client.put( url=url, data=json.dumps( { @@ -329,7 +325,7 @@ class ArgillaLogger(CustomBatchLogger): response.raise_for_status() if response.status_code >= 300: - verbose_logger.error(f"Argilla Error: {response.status_code} - {response.text}") + verbose_logger.error("Argilla Error: %s - %s", response.status_code, response.text) else: verbose_logger.debug("Batch of %s runs successfully created", len(self.log_queue)) except httpx.HTTPStatusError: diff --git a/litellm/integrations/arize/__init__.py b/litellm/integrations/arize/__init__.py index ab2627801e6..21da835ee16 100644 --- a/litellm/integrations/arize/__init__.py +++ b/litellm/integrations/arize/__init__.py @@ -1,31 +1,31 @@ import os -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final if TYPE_CHECKING: - from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec from litellm.integrations.custom_prompt_management import CustomPromptManagement + from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec from litellm.types.prompts.init_prompts import SupportedPromptIntegrations from .arize_phoenix_prompt_manager import ArizePhoenixPromptManager # Global instances -global_arize_config: Optional[dict] = None +global_arize_config: Final[dict | None] = None def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec") -> "CustomPromptManagement": """ Initialize a prompt from Arize Phoenix. """ - api_key = getattr(litellm_params, "api_key", None) or os.environ.get("PHOENIX_API_KEY") - api_base = getattr(litellm_params, "api_base", None) - prompt_id = getattr(litellm_params, "prompt_id", None) + api_key: Final = getattr(litellm_params, "api_key", None) or os.environ.get("PHOENIX_API_KEY") + api_base: Final = getattr(litellm_params, "api_base", None) + prompt_id: Final = getattr(litellm_params, "prompt_id", None) if not api_key or not api_base: raise ValueError("api_key and api_base are required for Arize Phoenix prompt integration") try: - arize_prompt_manager = ArizePhoenixPromptManager( + arize_prompt_manager: Final = ArizePhoenixPromptManager( **{ "api_key": api_key, "api_base": api_base, @@ -39,6 +39,6 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom raise e -prompt_initializer_registry = { +prompt_initializer_registry: Final = { SupportedPromptIntegrations.ARIZE_PHOENIX.value: prompt_initializer, } diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 44fd7a0d01a..8c494794858 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -1,5 +1,5 @@ import json -from typing import TYPE_CHECKING, Any, Dict, Optional, Type +from typing import TYPE_CHECKING, Any, Final from typing_extensions import override @@ -31,13 +31,13 @@ from litellm.integrations._types.open_inference import ( class ArizeOTELAttributes(BaseLLMObsOTELAttributes): @staticmethod @override - def set_messages(span: "Span", kwargs: Dict[str, Any]): - messages = kwargs.get("messages") + def set_messages(span: "Span", kwargs: dict[str, Any]): + messages: Final = kwargs.get("messages") # for /chat/completions # https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions if messages: - last_message = messages[-1] + last_message: Final = messages[-1] safe_set_attribute( span, SpanAttributes.INPUT_VALUE, @@ -129,7 +129,7 @@ def _set_choice_outputs(span: "Span", response_obj, msg_attrs, span_attrs): def _set_image_outputs(span: "Span", response_obj, image_attrs, span_attrs): - images = response_obj.get("data", []) + images: Final = response_obj.get("data", []) for i, image in enumerate(images): img_url = image.get("url") if img_url is None and image.get("b64_json"): @@ -145,7 +145,7 @@ def _set_image_outputs(span: "Span", response_obj, image_attrs, span_attrs): def _set_audio_outputs(span: "Span", response_obj, audio_attrs, span_attrs): - audio = response_obj.get("audio", []) + audio: Final = response_obj.get("audio", []) for i, audio_item in enumerate(audio): audio_url = audio_item.get("url") if audio_url is None and audio_item.get("b64_json"): @@ -166,7 +166,7 @@ def _set_audio_outputs(span: "Span", response_obj, audio_attrs, span_attrs): def _set_embedding_outputs(span: "Span", response_obj, embedding_attrs, span_attrs): - embeddings = response_obj.get("data", []) + embeddings: Final = response_obj.get("data", []) for i, embedding_item in enumerate(embeddings): embedding_vector = embedding_item.get("embedding") if embedding_vector: @@ -193,7 +193,7 @@ def _set_embedding_outputs(span: "Span", response_obj, embedding_attrs, span_att def _set_structured_outputs(span: "Span", response_obj, msg_attrs, span_attrs): - output_items = response_obj.get("output", []) + output_items: Final = response_obj.get("output", []) for i, item in enumerate(output_items): prefix = f"{span_attrs.LLM_OUTPUT_MESSAGES}.{i}" if not hasattr(item, "type"): @@ -232,7 +232,7 @@ def _safe_get(obj, key, default=None): """ if obj is None: return default - getter = getattr(obj, "get", None) + getter: Final = getattr(obj, "get", None) if callable(getter): try: return getter(key, default) @@ -243,15 +243,15 @@ def _safe_get(obj, key, default=None): def _set_usage_outputs(span: "Span", response_obj, span_attrs): - usage = response_obj and response_obj.get("usage") + usage: Final = response_obj and response_obj.get("usage") if not usage: return safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_TOTAL, _safe_get(usage, "total_tokens")) - completion_tokens = _safe_get(usage, "completion_tokens") or _safe_get(usage, "output_tokens") + completion_tokens: Final = _safe_get(usage, "completion_tokens") or _safe_get(usage, "output_tokens") if completion_tokens: safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens) - prompt_tokens = _safe_get(usage, "prompt_tokens") or _safe_get(usage, "input_tokens") + prompt_tokens: Final = _safe_get(usage, "prompt_tokens") or _safe_get(usage, "input_tokens") if prompt_tokens: safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens) @@ -259,8 +259,8 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): # API (Usage) and in `output_tokens_details` for Responses API # (ResponseAPIUsage). Both nested objects may be plain Pydantic models # without `.get`. - token_details = _safe_get(usage, "completion_tokens_details") or _safe_get(usage, "output_tokens_details") - reasoning_tokens = _safe_get(token_details, "reasoning_tokens") + token_details: Final = _safe_get(usage, "completion_tokens_details") or _safe_get(usage, "output_tokens_details") + reasoning_tokens: Final = _safe_get(token_details, "reasoning_tokens") if reasoning_tokens: safe_set_attribute( span, @@ -275,8 +275,8 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): # `cache_creation_input_tokens` # All emits are conditional, so when none of these fields exist (the # situation in the existing test fixtures) no extra attributes are set. - prompt_token_details = _safe_get(usage, "prompt_tokens_details") or _safe_get(usage, "input_tokens_details") - cache_read = _safe_get(prompt_token_details, "cached_tokens") or _safe_get(usage, "cache_read_input_tokens") + prompt_token_details: Final = _safe_get(usage, "prompt_tokens_details") or _safe_get(usage, "input_tokens_details") + cache_read: Final = _safe_get(prompt_token_details, "cached_tokens") or _safe_get(usage, "cache_read_input_tokens") if cache_read: safe_set_attribute( span, @@ -285,7 +285,7 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): ) # Anthropic / Bedrock-Anthropic only — OpenAI's `prompt_tokens_details` # does not expose a cache-write count, so we read straight off `usage`. - cache_write = _safe_get(usage, "cache_creation_input_tokens") + cache_write: Final = _safe_get(usage, "cache_creation_input_tokens") if cache_write: safe_set_attribute( span, @@ -293,7 +293,7 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): cache_write, ) - audio_prompt_tokens = _safe_get(prompt_token_details, "audio_tokens") + audio_prompt_tokens: Final = _safe_get(prompt_token_details, "audio_tokens") if audio_prompt_tokens: safe_set_attribute( span, @@ -302,7 +302,7 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs): ) -def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: +def _infer_open_inference_span_kind(call_type: str | None) -> str: """ Map LiteLLM call types to OpenInference span kinds. """ @@ -310,7 +310,7 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: if not call_type: return OpenInferenceSpanKindValues.UNKNOWN.value - lowered = str(call_type).lower() + lowered: Final = str(call_type).lower() if "embed" in lowered: return OpenInferenceSpanKindValues.EMBEDDING.value @@ -360,7 +360,7 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: return OpenInferenceSpanKindValues.UNKNOWN.value -def _set_tool_attributes(span: "Span", optional_tools: Optional[list], metadata_tools: Optional[list]): +def _set_tool_attributes(span: "Span", optional_tools: list | None, metadata_tools: list | None): """set tool attributes on span from optional_params or tool call metadata""" if optional_tools: for idx, tool in enumerate(optional_tools): @@ -408,7 +408,7 @@ def _set_tool_attributes(span: "Span", optional_tools: Optional[list], metadata_ ) -def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMObsOTELAttributes]): +def set_attributes(span: "Span", kwargs, response_obj, attributes: type[BaseLLMObsOTELAttributes]): """ Populates span with OpenInference-compliant LLM attributes for Arize and Phoenix tracing. """ @@ -416,7 +416,7 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMO # routes) into a dict so downstream `.get()` calls don't crash. Existing # dict / `.get()`-bearing objects (incl. Pydantic OpenAI Responses API # models) are returned unchanged, preserving the existing test behavior. - response_obj_for_attrs = _coerce_response_obj_for_attrs(response_obj) + response_obj_for_attrs: Final = _coerce_response_obj_for_attrs(response_obj) # Set span.kind defensively before anything else. If a downstream step # throws, the span still has a kind so Arize can render it correctly @@ -425,17 +425,17 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMO _safe_emit("early span kind", _set_early_span_kind, span, kwargs) try: - optional_params = _sanitize_optional_params(kwargs.get("optional_params")) - litellm_params = kwargs.get("litellm_params", {}) or {} - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") + optional_params: Final = _sanitize_optional_params(kwargs.get("optional_params")) + litellm_params: Final = kwargs.get("litellm_params", {}) or {} + standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - metadata = standard_logging_payload.get("metadata") if standard_logging_payload else None + metadata: Final = standard_logging_payload.get("metadata") if standard_logging_payload else None _set_metadata_attributes(span, metadata, SpanAttributes) - metadata_tools = _extract_metadata_tools(metadata) - optional_tools = _extract_optional_tools(optional_params) + metadata_tools: Final = _extract_metadata_tools(metadata) + optional_tools: Final = _extract_optional_tools(optional_params) _set_request_attributes( span=span, @@ -455,20 +455,20 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMO _set_tool_attributes(span, optional_tools, metadata_tools) attributes.set_messages(span, kwargs) - model_params = standard_logging_payload.get("model_parameters") if standard_logging_payload else None + model_params: Final = standard_logging_payload.get("model_parameters") if standard_logging_payload else None _set_model_params(span, model_params, SpanAttributes) _set_response_attributes(span=span, response_obj=response_obj_for_attrs) except Exception as e: - verbose_logger.error(f"[Arize/Phoenix] Failed to set OpenInference span attributes: {e}") + verbose_logger.error("[Arize/Phoenix] Failed to set OpenInference span attributes: %s", e) if hasattr(span, "record_exception"): span.record_exception(e) # Additive emitters. Each is independently guarded so a failure can never # blank the attributes set by the main try-block above. New attributes are # written under new keys; existing attributes are not overwritten. - slp = kwargs.get("standard_logging_object") + slp: Final = kwargs.get("standard_logging_object") _safe_emit("session/user attrs", _set_session_and_user_attrs, span, kwargs, slp) _safe_emit("response cost", _set_response_cost_attr, span, slp) _safe_emit( @@ -482,28 +482,28 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: Type[BaseLLMO ) -def _sanitize_optional_params(optional_params: Optional[dict]) -> dict: +def _sanitize_optional_params(optional_params: dict | None) -> dict: if not isinstance(optional_params, dict): return {} optional_params.pop("secret_fields", None) return optional_params -def _set_metadata_attributes(span: "Span", metadata: Optional[Any], span_attrs) -> None: +def _set_metadata_attributes(span: "Span", metadata: Any | None, span_attrs) -> None: if metadata is not None: safe_set_attribute(span, span_attrs.METADATA, safe_dumps(metadata)) -def _extract_metadata_tools(metadata: Optional[Any]) -> Optional[list]: +def _extract_metadata_tools(metadata: Any | None) -> list | None: if not isinstance(metadata, dict): return None - llm_obj = metadata.get("llm") + llm_obj: Final = metadata.get("llm") if isinstance(llm_obj, dict): return llm_obj.get("tools") return None -def _extract_optional_tools(optional_params: dict) -> Optional[list]: +def _extract_optional_tools(optional_params: dict) -> list | None: return optional_params.get("tools") if isinstance(optional_params, dict) else None @@ -544,13 +544,13 @@ def _set_request_attributes( safe_set_attribute(span, "llm.response.model", response_obj.get("model")) -def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) -> None: +def _set_model_params(span: "Span", model_params: dict | None, span_attrs) -> None: if not model_params: return safe_set_attribute(span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params)) if model_params.get("user"): - user_id = model_params.get("user") + user_id: Final = model_params.get("user") if user_id is not None: safe_set_attribute(span, span_attrs.USER_ID, user_id) @@ -573,8 +573,8 @@ def _safe_emit(label: str, fn, *args, **kwargs) -> None: def _set_early_span_kind(span: "Span", kwargs: dict) -> None: """Defensively set OPENINFERENCE_SPAN_KIND before any other logic runs.""" - slp = kwargs.get("standard_logging_object") - call_type = slp.get("call_type") if isinstance(slp, dict) else None + slp: Final = kwargs.get("standard_logging_object") + call_type: Final = slp.get("call_type") if isinstance(slp, dict) else None safe_set_attribute( span, SpanAttributes.OPENINFERENCE_SPAN_KIND, @@ -595,10 +595,10 @@ def _coerce_response_obj_for_attrs(response_obj): """ if response_obj is None or hasattr(response_obj, "get"): return response_obj - text = getattr(response_obj, "text", None) + text: Final = getattr(response_obj, "text", None) if isinstance(text, str) and text: try: - parsed = json.loads(text) + parsed: Final = json.loads(text) if isinstance(parsed, dict): return parsed except Exception: @@ -606,7 +606,7 @@ def _coerce_response_obj_for_attrs(response_obj): return response_obj -def _coerce_text(value) -> Optional[str]: +def _coerce_text(value) -> str | None: """Best-effort text extraction from a message-content value. Returns None when no textual portion can be derived. Handles: @@ -620,7 +620,7 @@ def _coerce_text(value) -> Optional[str]: if isinstance(value, str): return value if isinstance(value, list): - parts = [] + parts: Final = [] for part in value: if isinstance(part, str): parts.append(part) @@ -641,7 +641,7 @@ def _to_plain_dict(value): """ if value is None or isinstance(value, dict): return value - model_dump = getattr(value, "model_dump", None) + model_dump: Final = getattr(value, "model_dump", None) if callable(model_dump): try: return model_dump() @@ -650,16 +650,16 @@ def _to_plain_dict(value): return value -def _get_tool_calls(message) -> Optional[list]: +def _get_tool_calls(message) -> list | None: """Return ``message.tool_calls`` only when it's a non-empty list. Works for dicts and Pydantic message objects via ``_safe_get``. """ - tool_calls = _safe_get(message, "tool_calls") + tool_calls: Final = _safe_get(message, "tool_calls") return tool_calls if isinstance(tool_calls, list) and tool_calls else None -def _normalize_tool_call(raw_tc) -> Optional[Dict[str, Any]]: +def _normalize_tool_call(raw_tc) -> dict[str, Any] | None: """Normalize a single tool_call (dict or Pydantic) into a stable shape: {"id": str|None, "type": str, "function": {"name": str|None, "arguments": str|None}} @@ -667,11 +667,11 @@ def _normalize_tool_call(raw_tc) -> Optional[Dict[str, Any]]: Arguments are coerced to a JSON string per OpenInference convention. Returns ``None`` when ``raw_tc`` cannot be coerced to a dict. """ - tc = _to_plain_dict(raw_tc) + tc: Final = _to_plain_dict(raw_tc) if not isinstance(tc, dict): return None - function = _to_plain_dict(tc.get("function")) - name = function.get("name") if isinstance(function, dict) else None + function: Final = _to_plain_dict(tc.get("function")) + name: Final = function.get("name") if isinstance(function, dict) else None args = function.get("arguments") if isinstance(function, dict) else None if args is not None and not isinstance(args, str): try: @@ -692,7 +692,7 @@ def _summarize_tool_calls_for_output(tool_calls) -> str: so OUTPUT_VALUE is never blanked on a malformed payload. """ try: - normalized = [n for n in (_normalize_tool_call(tc) for tc in tool_calls) if n] + normalized: Final = [n for n in (_normalize_tool_call(tc) for tc in tool_calls) if n] return json.dumps({"tool_calls": normalized}) except Exception: return str(tool_calls) @@ -705,7 +705,7 @@ def _emit_message_tool_calls(span: "Span", prefix: str, message) -> None: Accepts dicts or Pydantic message objects (e.g. ``litellm.Message``); the same applies to each tool_call entry. """ - tool_calls = _get_tool_calls(message) + tool_calls: Final = _get_tool_calls(message) if not tool_calls: return for tc_idx, raw_tc in enumerate(tool_calls): @@ -744,11 +744,11 @@ def _emit_input_message_extras(span: "Span", prefix: str, message: dict) -> None if not isinstance(message, dict): return - name = message.get("name") + name: Final = message.get("name") if name: safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_NAME}", name) - tool_call_id = message.get("tool_call_id") + tool_call_id: Final = message.get("tool_call_id") if tool_call_id: safe_set_attribute( span, @@ -758,9 +758,9 @@ def _emit_input_message_extras(span: "Span", prefix: str, message: dict) -> None _emit_message_tool_calls(span, prefix, message) - content = message.get("content") + content: Final = message.get("content") if isinstance(content, list): - contents_prefix = f"{prefix}.{MessageAttributes.MESSAGE_CONTENTS}" + contents_prefix: Final = f"{prefix}.{MessageAttributes.MESSAGE_CONTENTS}" for part_idx, part in enumerate(content): if not isinstance(part, dict): continue @@ -823,36 +823,36 @@ def _set_session_and_user_attrs(span: "Span", kwargs: dict, standard_logging_pay """ if not isinstance(standard_logging_payload, dict): return - metadata = standard_logging_payload.get("metadata") or {} + metadata: Final = standard_logging_payload.get("metadata") or {} if not isinstance(metadata, dict): return - session_id = metadata.get("user_api_key_end_user_id") + session_id: Final = metadata.get("user_api_key_end_user_id") if session_id: safe_set_attribute(span, SpanAttributes.SESSION_ID, str(session_id)) - trace_id = standard_logging_payload.get("trace_id") + trace_id: Final = standard_logging_payload.get("trace_id") if trace_id: safe_set_attribute(span, "litellm.trace_id", str(trace_id)) - optional_params = kwargs.get("optional_params") or {} - model_params = standard_logging_payload.get("model_parameters") or {} - has_user_already = bool( + optional_params: Final = kwargs.get("optional_params") or {} + model_params: Final = standard_logging_payload.get("model_parameters") or {} + has_user_already: Final = bool( (isinstance(optional_params, dict) and optional_params.get("user")) or (isinstance(model_params, dict) and model_params.get("user")) ) if not has_user_already: - user_id = metadata.get("user_api_key_user_id") + user_id: Final = metadata.get("user_api_key_user_id") if user_id: safe_set_attribute(span, SpanAttributes.USER_ID, str(user_id)) - team_id = metadata.get("user_api_key_team_id") + team_id: Final = metadata.get("user_api_key_team_id") if team_id: safe_set_attribute(span, "litellm.team_id", str(team_id)) - team_alias = metadata.get("user_api_key_team_alias") + team_alias: Final = metadata.get("user_api_key_team_alias") if team_alias: safe_set_attribute(span, "litellm.team_alias", str(team_alias)) - key_alias = metadata.get("user_api_key_alias") + key_alias: Final = metadata.get("user_api_key_alias") if key_alias: safe_set_attribute(span, "litellm.key_alias", str(key_alias)) @@ -868,21 +868,21 @@ def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None: """ if not isinstance(standard_logging_payload, dict): return - cost = standard_logging_payload.get("response_cost") + cost: Final = standard_logging_payload.get("response_cost") if cost is None: return try: - cost_value = float(cost) + cost_value: Final = float(cost) except (TypeError, ValueError): return safe_set_attribute(span, "llm.cost.total", cost_value) safe_set_attribute(span, "llm.response.cost", cost_value) -def _is_passthrough_call_type(call_type: Optional[str]) -> bool: +def _is_passthrough_call_type(call_type: str | None) -> bool: if not call_type: return False - lowered = str(call_type).lower() + lowered: Final = str(call_type).lower() return "passthrough" in lowered or "pass_through" in lowered @@ -913,7 +913,7 @@ def _maybe_normalize_passthrough( passthrough I/O (with central redaction) for free and this helper's `complete_input_dict` fallback can be deleted. See follow-up issue. """ - call_type = standard_logging_payload.get("call_type") if isinstance(standard_logging_payload, dict) else None + call_type: Final = standard_logging_payload.get("call_type") if isinstance(standard_logging_payload, dict) else None if not _is_passthrough_call_type(call_type): return @@ -927,13 +927,13 @@ def _maybe_normalize_passthrough( return # --- INPUT -------------------------------------------------------------- - additional_args = kwargs.get("additional_args") or {} + additional_args: Final = kwargs.get("additional_args") or {} complete_input_dict = additional_args.get("complete_input_dict") if isinstance(additional_args, dict) else None if isinstance(complete_input_dict, dict): _set_passthrough_input_attributes(span, complete_input_dict.get("messages")) # --- OUTPUT ------------------------------------------------------------- - parsed_response = _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs) + parsed_response: Final = _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs) if not isinstance(parsed_response, dict): return @@ -977,13 +977,13 @@ def _set_passthrough_input_attributes(span: "Span", messages) -> None: def _set_passthrough_output_attributes(span: "Span", parsed_response: dict) -> None: """Render passthrough response into OUTPUT_VALUE + LLM_OUTPUT_MESSAGES.""" # Anthropic / Bedrock-Anthropic: `content` is a list of typed parts. - content_list = parsed_response.get("content") + content_list: Final = parsed_response.get("content") if isinstance(content_list, list) and content_list: - texts = [] + texts: Final = [] for part in content_list: if isinstance(part, dict) and isinstance(part.get("text"), str): texts.append(part["text"]) - joined = "\n\n".join(t for t in texts if t) + joined: Final = "\n\n".join(t for t in texts if t) if joined: safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, joined) prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0" @@ -999,13 +999,13 @@ def _set_passthrough_output_attributes(span: "Span", parsed_response: dict) -> N ) # OpenAI-style passthrough: `choices[0].message.content` - choices = parsed_response.get("choices") + choices: Final = parsed_response.get("choices") if isinstance(choices, list) and choices: - first = choices[0] + first: Final = choices[0] if isinstance(first, dict): - msg = first.get("message") + msg: Final = first.get("message") if isinstance(msg, dict): - text = _coerce_text(msg.get("content")) + text: Final = _coerce_text(msg.get("content")) if text: safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, text) prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0" @@ -1024,7 +1024,7 @@ def _set_passthrough_output_attributes(span: "Span", parsed_response: dict) -> N def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs): """Return a dict view of the provider response for passthrough routes.""" # Prefer the coerced view (already JSON-parsed for httpx.Response). - candidates = [] + candidates: Final = [] if isinstance(coerced_response_obj, dict): candidates.append(coerced_response_obj) if isinstance(raw_response_obj, dict) and raw_response_obj is not coerced_response_obj: @@ -1047,7 +1047,7 @@ def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs): return candidate # Fallback: kwargs["original_response"] from the OTel base path. - original = kwargs.get("original_response") if isinstance(kwargs, dict) else None + original: Final = kwargs.get("original_response") if isinstance(kwargs, dict) else None if isinstance(original, dict): return original if isinstance(original, str): diff --git a/litellm/integrations/arize/arize.py b/litellm/integrations/arize/arize.py index e5fdb231933..bcab610835c 100644 --- a/litellm/integrations/arize/arize.py +++ b/litellm/integrations/arize/arize.py @@ -6,7 +6,7 @@ this file has Arize ai specific helper functions import os from datetime import datetime -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Final, Union from litellm.integrations.arize import _utils from litellm.integrations.arize._utils import ArizeOTELAttributes @@ -50,7 +50,7 @@ class ArizeLogger(OpenTelemetry): self.span_kind = SpanKind return - provider = TracerProvider(resource=self._get_litellm_resource(self.config)) + provider: Final = TracerProvider(resource=self._get_litellm_resource(self.config)) provider.add_span_processor(self._get_span_processor()) self.tracer = provider.get_tracer("litellm") self.span_kind = SpanKind @@ -61,16 +61,13 @@ class ArizeLogger(OpenTelemetry): ``open_telemetry_logger``. That attribute is reserved for the primary ``otel`` callback which handles proxy-level parent spans. """ - pass - def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): + def set_attributes(self, span: Span, kwargs, response_obj: Any | None): ArizeLogger.set_arize_attributes(span, kwargs, response_obj) - return @staticmethod def set_arize_attributes(span: Span, kwargs, response_obj): _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) - return @staticmethod def get_arize_config() -> ArizeConfig: @@ -83,13 +80,13 @@ class ArizeLogger(OpenTelemetry): Raises: ValueError: If required environment variables are not set. """ - space_id = os.environ.get("ARIZE_SPACE_ID") - space_key = os.environ.get("ARIZE_SPACE_KEY") - api_key = os.environ.get("ARIZE_API_KEY") - project_name = os.environ.get("ARIZE_PROJECT_NAME") + space_id: Final = os.environ.get("ARIZE_SPACE_ID") + space_key: Final = os.environ.get("ARIZE_SPACE_KEY") + api_key: Final = os.environ.get("ARIZE_API_KEY") + project_name: Final = os.environ.get("ARIZE_PROJECT_NAME") - grpc_endpoint = os.environ.get("ARIZE_ENDPOINT") - http_endpoint = os.environ.get("ARIZE_HTTP_ENDPOINT") + grpc_endpoint: Final = os.environ.get("ARIZE_ENDPOINT") + http_endpoint: Final = os.environ.get("ARIZE_HTTP_ENDPOINT") endpoint = None protocol: Protocol = "otlp_grpc" @@ -116,25 +113,23 @@ class ArizeLogger(OpenTelemetry): async def async_service_success_hook( self, payload: ServiceLoggerPayload, - parent_otel_span: Optional[Span] = None, - start_time: Optional[Union[datetime, float]] = None, - end_time: Optional[Union[datetime, float]] = None, - event_metadata: Optional[dict] = None, + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: datetime | float | None = None, + event_metadata: dict | None = None, ): """Arize is used mainly for LLM I/O tracing, sending router+caching metrics adds bloat to arize logs""" - pass async def async_service_failure_hook( self, payload: ServiceLoggerPayload, - error: Optional[str] = "", - parent_otel_span: Optional[Span] = None, - start_time: Optional[Union[datetime, float]] = None, - end_time: Optional[Union[float, datetime]] = None, - event_metadata: Optional[dict] = None, + error: str | None = "", + parent_otel_span: Span | None = None, + start_time: datetime | float | None = None, + end_time: float | datetime | None = None, + event_metadata: dict | None = None, ): """Arize is used mainly for LLM I/O tracing, sending router+caching metrics adds bloat to arize logs""" - pass # def create_litellm_proxy_request_started_span( # self, @@ -152,7 +147,7 @@ class ArizeLogger(OpenTelemetry): dict: Health check result with status and message """ try: - config = self.get_arize_config() + config: Final = self.get_arize_config() if not config.space_id and not config.space_key: return { @@ -174,12 +169,12 @@ class ArizeLogger(OpenTelemetry): except Exception as e: return { "status": "unhealthy", - "error_message": f"Arize health check failed: {str(e)}", + "error_message": f"Arize health check failed: {e}", } def construct_dynamic_otel_headers( self, standard_callback_dynamic_params: StandardCallbackDynamicParams - ) -> Optional[dict]: + ) -> dict | None: """ Construct dynamic Arize headers from standard callback dynamic params @@ -188,7 +183,7 @@ class ArizeLogger(OpenTelemetry): Returns: dict: A dictionary of dynamic Arize headers """ - dynamic_headers = {} + dynamic_headers: Final = {} ######################################################### # `arize-space-id` handling diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index db7aed1a71c..41011a6ee98 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -1,7 +1,7 @@ import os import threading from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Final, Union from litellm._logging import verbose_logger from litellm.integrations.arize import _utils @@ -12,8 +12,7 @@ if TYPE_CHECKING: from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SpanProcessor from opentelemetry.trace import Span as _Span - from opentelemetry.trace import SpanKind - from opentelemetry.trace import Tracer + from opentelemetry.trace import SpanKind, Tracer from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry from litellm.integrations.opentelemetry import ( @@ -44,8 +43,8 @@ else: OpenTelemetry = None # type: ignore -ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces" -_MAX_PROJECT_PROVIDERS = 64 +ARIZE_HOSTED_PHOENIX_ENDPOINT: Final = "https://otlp.arize.com/v1/traces" +_MAX_PROJECT_PROVIDERS: Final = 64 class ArizePhoenixLogger(OpenTelemetry): # type: ignore @@ -81,7 +80,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore self._shared_span_processor = self._get_span_processor() self.span_kind = SpanKind - default_project = self._resolve_project_name({}) + default_project: Final = self._resolve_project_name({}) self.tracer = self._get_tracer_for(default_project) verbose_logger.debug( "ArizePhoenixLogger: Initialized per-project TracerProvider cache " @@ -101,7 +100,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore if getattr(self, "_use_injected_tracer_provider", False): return - shared_processor = getattr(self, "_shared_span_processor", None) + shared_processor: Final = getattr(self, "_shared_span_processor", None) if shared_processor is not None: try: shared_processor.force_flush() @@ -112,7 +111,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore ) with getattr(self, "_project_providers_lock", threading.Lock()): - providers = list(getattr(self, "_project_providers", {}).values()) + providers: Final = list(getattr(self, "_project_providers", {}).values()) for provider in providers: try: @@ -130,24 +129,24 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore """ from opentelemetry.sdk.resources import OTELResourceDetector, Resource - project_attributes: dict[str, str] = { + project_attributes: Final[dict[str, str]] = { "openinference.project.name": project_name, "model_id": project_name, "service.name": project_name, } - deployment_environment = getattr(self.config, "deployment_environment", None) + deployment_environment: Final = getattr(self.config, "deployment_environment", None) if deployment_environment is not None: project_attributes["deployment.environment"] = deployment_environment - env_resource = OTELResourceDetector().detect() - project_resource = Resource.create(project_attributes) # type: ignore[arg-type] + env_resource: Final = OTELResourceDetector().detect() + project_resource: Final = Resource.create(project_attributes) # type: ignore[arg-type] return env_resource.merge(project_resource) def _build_tracer_provider_for_project(self, project_name: str) -> TracerProvider: """Create a TracerProvider for *project_name* (caller holds no cache lock).""" from opentelemetry.sdk.trace import TracerProvider - provider = TracerProvider(resource=self._get_litellm_resource_for_project(project_name)) + provider: Final = TracerProvider(resource=self._get_litellm_resource_for_project(project_name)) provider.add_span_processor(self._shared_span_processor) return provider @@ -163,7 +162,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore # OTELResourceDetector().detect() is synchronous; build outside the lock so # concurrent requests for other projects are not blocked on cache misses. - new_provider = self._build_tracer_provider_for_project(project_name) + new_provider: Final = self._build_tracer_provider_for_project(project_name) with self._project_providers_lock: if project_name in self._project_providers: @@ -176,9 +175,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore self._project_providers[project_name] = new_provider return new_provider.get_tracer(LITELLM_TRACER_NAME) - def _resolve_tracer_for_kwargs(self, kwargs: dict) -> Tuple[str, Tracer]: + def _resolve_tracer_for_kwargs(self, kwargs: dict) -> tuple[str, Tracer]: """Resolve project name once and return the matching tracer.""" - project_name = self._resolve_project_name(kwargs) + project_name: Final = self._resolve_project_name(kwargs) return project_name, self._get_tracer_for(project_name) def get_tracer_to_use_for_request(self, kwargs: dict) -> Tracer: @@ -193,22 +192,19 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore ``open_telemetry_logger``. That attribute is reserved for the primary ``otel`` callback which handles proxy-level parent spans. """ - pass - def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): + def set_attributes(self, span: Span, kwargs, response_obj: Any | None): ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj) - return @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) - return @staticmethod - def _normalize_project_name(name: Optional[str]) -> Optional[str]: + def _normalize_project_name(name: str | None) -> str | None: if name is None: return None - normalized = str(name).strip() + normalized: Final = str(name).strip() return normalized if normalized else None @staticmethod @@ -232,11 +228,11 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore user-supplied and would let an authenticated caller fake proxy-mode detection to route their telemetry into arbitrary Arize/Phoenix projects. """ - litellm_params = kwargs.get("litellm_params") + litellm_params: Final = kwargs.get("litellm_params") return isinstance(litellm_params, dict) and bool(litellm_params.get("proxy_server_request")) @staticmethod - def _project_from_metadata_dict(metadata: dict, metadata_key: str, *, proxy_mode: bool) -> Optional[str]: + def _project_from_metadata_dict(metadata: dict, metadata_key: str, *, proxy_mode: bool) -> str | None: """ Read a Phoenix project field from proxy/SDK metadata. @@ -244,9 +240,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore select the project. SDK callers may still set project fields directly on ``metadata``. """ - auth_metadata = metadata.get("user_api_key_auth_metadata") + auth_metadata: Final = metadata.get("user_api_key_auth_metadata") if isinstance(auth_metadata, dict): - project = ArizePhoenixLogger._normalize_project_name(auth_metadata.get(metadata_key)) + project: Final = ArizePhoenixLogger._normalize_project_name(auth_metadata.get(metadata_key)) if project: return project @@ -255,8 +251,8 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore return None @staticmethod - def _metadata_project_from_kwargs(kwargs: dict, metadata_key: str) -> Optional[str]: - proxy_mode = ArizePhoenixLogger._is_proxy_request(kwargs) + def _metadata_project_from_kwargs(kwargs: dict, metadata_key: str) -> str | None: + proxy_mode: Final = ArizePhoenixLogger._is_proxy_request(kwargs) for metadata in ArizePhoenixLogger._iter_metadata_dicts_from_kwargs(kwargs): project = ArizePhoenixLogger._project_from_metadata_dict(metadata, metadata_key, proxy_mode=proxy_mode) if project: @@ -272,15 +268,15 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore ``user_api_key_auth_metadata.phoenix_project_name``, env, then ``default``. SDK priority: request metadata fields, then env, then ``default``. """ - override = ArizePhoenixLogger._metadata_project_from_kwargs(kwargs, "phoenix_project_name_override") + override: Final = ArizePhoenixLogger._metadata_project_from_kwargs(kwargs, "phoenix_project_name_override") if override: return override - phoenix_name = ArizePhoenixLogger._metadata_project_from_kwargs(kwargs, "phoenix_project_name") + phoenix_name: Final = ArizePhoenixLogger._metadata_project_from_kwargs(kwargs, "phoenix_project_name") if phoenix_name: return phoenix_name - env_name = ArizePhoenixLogger._normalize_project_name( + env_name: Final = ArizePhoenixLogger._normalize_project_name( os.environ.get("PHOENIX_PROJECT_NAME") or os.environ.get("ARIZE_PROJECT_NAME") ) if env_name: @@ -288,7 +284,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore return "default" - def _get_phoenix_context(self, kwargs, tracer: Optional[Tracer] = None): + def _get_phoenix_context(self, kwargs, tracer: Tracer | None = None): """ Build a trace context for Phoenix's dedicated TracerProvider. @@ -308,23 +304,23 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore if tracer is None: tracer = self._resolve_tracer_for_kwargs(kwargs)[1] - litellm_params = kwargs.get("litellm_params", {}) or {} - proxy_server_request = litellm_params.get("proxy_server_request", {}) or {} - headers = proxy_server_request.get("headers", {}) or {} + litellm_params: Final = kwargs.get("litellm_params", {}) or {} + proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {} + headers: Final = proxy_server_request.get("headers", {}) or {} traceparent_ctx = self.get_traceparent_from_header(headers=headers) if headers.get("traceparent") else None - is_proxy_mode = bool(proxy_server_request) + is_proxy_mode: Final = bool(proxy_server_request) if is_proxy_mode: - start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time")) - parent_span = tracer.start_span( + start_time_val: Final = kwargs.get("start_time", kwargs.get("api_call_start_time")) + parent_span: Final = tracer.start_span( name="litellm_proxy_request", start_time=(self._to_ns(start_time_val) if start_time_val is not None else None), context=traceparent_ctx, kind=self.span_kind.SERVER, ) - ctx = trace.set_span_in_context(parent_span) + ctx: Final = trace.set_span_in_context(parent_span) return ctx, parent_span return traceparent_ctx, None @@ -356,9 +352,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore _project_name, tracer = self._resolve_tracer_for_kwargs(kwargs) ctx, parent_span = self._get_phoenix_context(kwargs, tracer=tracer) - status = Status(StatusCode.OK if success else StatusCode.ERROR) + status: Final = Status(StatusCode.OK if success else StatusCode.ERROR) - span = tracer.start_span( + span: Final = tracer.start_span( name=self._get_span_name(kwargs), start_time=self._to_ns(start_time), context=ctx, @@ -393,13 +389,13 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore Retrieves the Arize Phoenix configuration based on environment variables. Returns: """ - api_key = os.environ.get("PHOENIX_API_KEY", None) + api_key: Final = os.environ.get("PHOENIX_API_KEY", None) collector_endpoint = os.environ.get("PHOENIX_COLLECTOR_HTTP_ENDPOINT", None) if not collector_endpoint: - grpc_endpoint = os.environ.get("PHOENIX_COLLECTOR_ENDPOINT", None) - http_endpoint = os.environ.get("PHOENIX_COLLECTOR_HTTP_ENDPOINT", None) + grpc_endpoint: Final = os.environ.get("PHOENIX_COLLECTOR_ENDPOINT", None) + http_endpoint: Final = os.environ.get("PHOENIX_COLLECTOR_HTTP_ENDPOINT", None) collector_endpoint = http_endpoint or grpc_endpoint endpoint = None @@ -429,7 +425,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore endpoint = "http://localhost:6006/v1/traces" protocol = "otlp_http" verbose_logger.debug( - f"No PHOENIX_COLLECTOR_ENDPOINT found, using default local Phoenix endpoint: {endpoint}" + "No PHOENIX_COLLECTOR_ENDPOINT found, using default local Phoenix endpoint: %s", endpoint ) otlp_auth_headers = None @@ -438,7 +434,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore elif "app.phoenix.arize.com" in endpoint: raise ValueError("PHOENIX_API_KEY must be set when using Phoenix Cloud (app.phoenix.arize.com).") - project_name = os.environ.get("PHOENIX_PROJECT_NAME") or "default" + project_name: Final = os.environ.get("PHOENIX_PROJECT_NAME") or "default" return ArizePhoenixConfig( otlp_auth_headers=otlp_auth_headers, @@ -448,7 +444,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore ) async def async_health_check(self): - config = self.get_arize_phoenix_config() + config: Final = self.get_arize_phoenix_config() if not config.otlp_auth_headers: return { diff --git a/litellm/integrations/arize/arize_phoenix_client.py b/litellm/integrations/arize/arize_phoenix_client.py index 7c0715d2e1e..18d35fee34b 100644 --- a/litellm/integrations/arize/arize_phoenix_client.py +++ b/litellm/integrations/arize/arize_phoenix_client.py @@ -3,7 +3,7 @@ Arize Phoenix API client for fetching prompt versions from Arize Phoenix. """ import urllib.parse -from typing import Any, Dict, Optional +from typing import Any, Final from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -27,7 +27,7 @@ class ArizePhoenixClient: - Direct API base URL configuration """ - def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None): + def __init__(self, api_key: str | None = None, api_base: str | None = None): """ Initialize the Arize Phoenix client. @@ -53,7 +53,7 @@ class ArizePhoenixClient: # Initialize HTTPHandler self.http_handler = HTTPHandler(disable_default_headers=True) - def get_prompt_version(self, prompt_version_id: str) -> Optional[Dict[str, Any]]: + def get_prompt_version(self, prompt_version_id: str) -> dict[str, Any] | None: """ Fetch a prompt version from Arize Phoenix. @@ -63,15 +63,15 @@ class ArizePhoenixClient: Returns: Dictionary containing prompt version data, or None if not found """ - safe_id = _sanitize_id(prompt_version_id) - url = f"{self.api_base}/v1/prompt_versions/{safe_id}" + safe_id: Final = _sanitize_id(prompt_version_id) + url: Final = f"{self.api_base}/v1/prompt_versions/{safe_id}" try: # Use the underlying httpx client directly to avoid query param extraction response = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data = response.json() + data: Final = response.json() return data.get("data") except Exception as e: @@ -100,8 +100,8 @@ class ArizePhoenixClient: """ try: # Try to access the prompt_versions endpoint to test connection - url = f"{self.api_base}/prompt_versions" - response = self.http_handler.client.get(url, headers=self.headers) + url: Final = f"{self.api_base}/prompt_versions" + response: Final = self.http_handler.client.get(url, headers=self.headers) response.raise_for_status() return True except Exception: diff --git a/litellm/integrations/arize/arize_phoenix_prompt_manager.py b/litellm/integrations/arize/arize_phoenix_prompt_manager.py index 4053b725a0f..a541817ca8e 100644 --- a/litellm/integrations/arize/arize_phoenix_prompt_manager.py +++ b/litellm/integrations/arize/arize_phoenix_prompt_manager.py @@ -3,7 +3,7 @@ Arize Phoenix prompt manager that integrates with LiteLLM's prompt management sy Fetches prompt versions from Arize Phoenix and provides workspace-based access control. """ -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Final from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -28,9 +28,9 @@ class ArizePhoenixPromptTemplate: def __init__( self, template_id: str, - messages: List[Dict[str, Any]], - metadata: Dict[str, Any], - model: Optional[str] = None, + messages: list[dict[str, Any]], + metadata: dict[str, Any], + model: str | None = None, ): self.template_id = template_id self.messages = messages @@ -61,14 +61,14 @@ class ArizePhoenixTemplateManager: def __init__( self, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - prompt_id: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + prompt_id: str | None = None, ): self.api_key = api_key self.api_base = api_base self.prompt_id = prompt_id - self.prompts: Dict[str, ArizePhoenixPromptTemplate] = {} + self.prompts: dict[str, ArizePhoenixPromptTemplate] = {} self.arize_client = ArizePhoenixClient(api_key=self.api_key, api_base=self.api_base) # Templates fetched from Arize Phoenix come from external workspace @@ -97,23 +97,23 @@ class ArizePhoenixTemplateManager: """Load a specific prompt version from Arize Phoenix.""" try: # Fetch the prompt version from Arize Phoenix - prompt_data = self.arize_client.get_prompt_version(prompt_version_id) + prompt_data: Final = self.arize_client.get_prompt_version(prompt_version_id) if prompt_data: - template = self._parse_prompt_data(prompt_data, prompt_version_id) + template: Final = self._parse_prompt_data(prompt_data, prompt_version_id) self.prompts[prompt_version_id] = template else: raise ValueError(f"Prompt version '{prompt_version_id}' not found") except Exception as e: raise Exception(f"Failed to load prompt version '{prompt_version_id}' from Arize Phoenix: {e}") - def _parse_prompt_data(self, data: Dict[str, Any], prompt_version_id: str) -> ArizePhoenixPromptTemplate: + def _parse_prompt_data(self, data: dict[str, Any], prompt_version_id: str) -> ArizePhoenixPromptTemplate: """Parse Arize Phoenix prompt data and extract messages and metadata.""" - template_data = data.get("template", {}) - messages = template_data.get("messages", []) + template_data: Final = data.get("template", {}) + messages: Final = template_data.get("messages", []) # Extract invocation parameters - invocation_params = data.get("invocation_parameters", {}) + invocation_params: Final = data.get("invocation_parameters", {}) provider_params = {} # Extract provider-specific parameters @@ -129,7 +129,7 @@ class ArizePhoenixTemplateManager: break # Build metadata dictionary - metadata = { + metadata: Final = { "model_name": data.get("model_name"), "model_provider": data.get("model_provider"), "description": data.get("description", ""), @@ -146,13 +146,13 @@ class ArizePhoenixTemplateManager: metadata=metadata, ) - def render_template(self, template_id: str, variables: Optional[Dict[str, Any]] = None) -> List[AllMessageValues]: + def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> list[AllMessageValues]: """Render a template with the given variables and return formatted messages.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") - template = self.prompts[template_id] - rendered_messages: List[AllMessageValues] = [] + template: Final = self.prompts[template_id] + rendered_messages: Final[list[AllMessageValues]] = [] for message in template.messages: role = message.get("role", "user") @@ -180,11 +180,11 @@ class ArizePhoenixTemplateManager: return rendered_messages - def get_template(self, template_id: str) -> Optional[ArizePhoenixPromptTemplate]: + def get_template(self, template_id: str) -> ArizePhoenixPromptTemplate | None: """Get a template by ID.""" return self.prompts.get(template_id) - def list_templates(self) -> List[str]: + def list_templates(self) -> list[str]: """List all available template IDs.""" return list(self.prompts.keys()) @@ -215,16 +215,16 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def __init__( self, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - prompt_id: Optional[str] = None, + api_key: str | None = None, + api_base: str | None = None, + prompt_id: str | None = None, **kwargs, ): super().__init__(**kwargs) self.api_key = api_key self.api_base = api_base self.prompt_id = prompt_id - self._prompt_manager: Optional[ArizePhoenixTemplateManager] = None + self._prompt_manager: ArizePhoenixTemplateManager | None = None @property def integration_name(self) -> str: @@ -245,8 +245,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: Optional[Dict[str, Any]] = None, - ) -> Tuple[List[AllMessageValues], Dict[str, Any]]: + prompt_variables: dict[str, Any] | None = None, + ) -> tuple[list[AllMessageValues], dict[str, Any]]: """ Get a prompt template and render it with variables. @@ -257,22 +257,22 @@ class ArizePhoenixPromptManager(CustomPromptManagement): Returns: Tuple of (rendered_messages, metadata) """ - template = self.prompt_manager.get_template(prompt_id) + template: Final = self.prompt_manager.get_template(prompt_id) if not template: raise ValueError(f"Prompt template '{prompt_id}' not found") # Render the template - rendered_messages = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) + rendered_messages: Final = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) # Extract metadata - metadata = { + metadata: Final = { "model": template.model, "temperature": template.temperature, "max_tokens": template.max_tokens, } # Add additional invocation parameters - invocation_params = template.invocation_parameters + invocation_params: Final = template.invocation_parameters provider_params = {} if "openai" in invocation_params: @@ -289,14 +289,14 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def pre_call_hook( self, - user_id: Optional[str], - messages: List[AllMessageValues], - function_call: Optional[Union[Dict[str, Any], str]] = None, - litellm_params: Optional[Dict[str, Any]] = None, - prompt_id: Optional[str] = None, - prompt_variables: Optional[Dict[str, Any]] = None, + user_id: str | None, + messages: list[AllMessageValues], + function_call: dict[str, Any] | str | None = None, + litellm_params: dict[str, Any] | None = None, + prompt_id: str | None = None, + prompt_variables: dict[str, Any] | None = None, **kwargs, - ) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]: + ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: """ Pre-call hook that processes the prompt template before making the LLM call. """ @@ -339,10 +339,10 @@ class ArizePhoenixPromptManager(CustomPromptManagement): # Log error but don't fail the call import litellm - litellm._logging.verbose_proxy_logger.error(f"Error in Arize Phoenix prompt pre_call_hook: {e}") + litellm._logging.verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e) return messages, litellm_params - def get_available_prompts(self) -> List[str]: + def get_available_prompts(self) -> list[str]: """Get list of available prompt IDs.""" return self.prompt_manager.list_templates() @@ -354,8 +354,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def should_run_prompt_management( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], + prompt_id: str | None, + prompt_spec: PromptSpec | None, dynamic_callback_params: StandardCallbackDynamicParams, ) -> bool: """ @@ -368,12 +368,12 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def _compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_spec: PromptSpec | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """ Compile an Arize Phoenix prompt template into a PromptManagementClient structure. @@ -395,10 +395,10 @@ class ArizePhoenixPromptManager(CustomPromptManagement): rendered_messages, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Extract model from metadata (if specified) - template_model = prompt_metadata.get("model") + template_model: Final = prompt_metadata.get("model") # Extract optional parameters from metadata - optional_params = {} + optional_params: Final = {} for param in [ "temperature", "max_tokens", @@ -422,12 +422,12 @@ class ArizePhoenixPromptManager(CustomPromptManagement): async def async_compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """ Async version of compile prompt helper. Since Arize Phoenix operations are synchronous, @@ -447,17 +447,17 @@ class ArizePhoenixPromptManager(CustomPromptManagement): def get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Get chat completion prompt from Arize Phoenix and return processed model, messages, and parameters. """ diff --git a/litellm/integrations/athina.py b/litellm/integrations/athina.py index d1bf8e68624..acc7d1003a2 100644 --- a/litellm/integrations/athina.py +++ b/litellm/integrations/athina.py @@ -1,4 +1,5 @@ import datetime +from typing import Final import litellm @@ -34,11 +35,11 @@ class AthinaLogger: import traceback try: - is_stream = kwargs.get("stream", False) + is_stream: Final = kwargs.get("stream", False) if is_stream: if "complete_streaming_response" in kwargs: # Log the completion response in streaming mode - completion_response = kwargs["complete_streaming_response"] + completion_response: Final = kwargs["complete_streaming_response"] response_json = completion_response.model_dump() if completion_response else {} else: # Skip logging if the completion response is not available @@ -46,7 +47,7 @@ class AthinaLogger: else: # Log the completion response in non streaming mode response_json = response_obj.model_dump() if response_obj else {} - data = { + data: Final = { "language_model_id": kwargs.get("model"), "request": kwargs, "response": response_json, @@ -62,16 +63,16 @@ class AthinaLogger: data["prompt"] = kwargs.get("messages", None) # Directly add tools or functions if present - optional_params = kwargs.get("optional_params", {}) + optional_params: Final = kwargs.get("optional_params", {}) data.update((k, v) for k, v in optional_params.items() if k in ["tools", "functions"]) # Add additional metadata keys - metadata = kwargs.get("litellm_params", {}).get("metadata", {}) + metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {}) if metadata: for key in self.additional_keys: if key in metadata: data[key] = metadata[key] - response = litellm.module_level_client.post( + response: Final = litellm.module_level_client.post( self.athina_logging_url, headers=self.headers, data=json.dumps(data, default=str), @@ -82,4 +83,3 @@ class AthinaLogger: print_verbose(f"Athina Logger Succeeded - {response.text}") except Exception as e: print_verbose(f"Athina Logger Error - {e}, Stack trace: {traceback.format_exc()}") - pass diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 5f8afe58cb0..f23317ae9df 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -16,7 +16,7 @@ import asyncio import os import time import traceback -from typing import List, Optional, Union +from typing import Final from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -35,13 +35,13 @@ class AzureSentinelLogger(CustomBatchLogger): def __init__( self, - dcr_immutable_id: Optional[str] = None, - stream_name: Optional[str] = None, - endpoint: Optional[str] = None, - tenant_id: Optional[str] = None, - client_id: Optional[str] = None, - client_secret: Optional[str] = None, - audit_stream_name: Optional[str] = None, + dcr_immutable_id: str | None = None, + stream_name: str | None = None, + endpoint: str | None = None, + tenant_id: str | None = None, + client_id: str | None = None, + client_secret: str | None = None, + audit_stream_name: str | None = None, **kwargs, ): """ @@ -65,15 +65,15 @@ class AzureSentinelLogger(CustomBatchLogger): """ self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - resolved_dcr_immutable_id = dcr_immutable_id or os.getenv("AZURE_SENTINEL_DCR_IMMUTABLE_ID") - resolved_stream_name = stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM" - resolved_audit_stream_name = ( + resolved_dcr_immutable_id: Final = dcr_immutable_id or os.getenv("AZURE_SENTINEL_DCR_IMMUTABLE_ID") + resolved_stream_name: Final = stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM" + resolved_audit_stream_name: Final = ( audit_stream_name or os.getenv("AZURE_SENTINEL_AUDIT_STREAM_NAME") or resolved_stream_name ) - resolved_endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT") - resolved_tenant_id = tenant_id or os.getenv("AZURE_SENTINEL_TENANT_ID") or os.getenv("AZURE_TENANT_ID") - resolved_client_id = client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv("AZURE_CLIENT_ID") - resolved_client_secret = ( + resolved_endpoint: Final = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT") + resolved_tenant_id: Final = tenant_id or os.getenv("AZURE_SENTINEL_TENANT_ID") or os.getenv("AZURE_TENANT_ID") + resolved_client_id: Final = client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv("AZURE_CLIENT_ID") + resolved_client_secret: Final = ( client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET") ) @@ -120,14 +120,14 @@ class AzureSentinelLogger(CustomBatchLogger): # OAuth2 scope for Azure Monitor self.oauth_scope = "https://monitor.azure.com/.default" - self.oauth_token: Optional[str] = None - self.oauth_token_expires_at: Optional[float] = None + self.oauth_token: str | None = None + self.oauth_token_expires_at: float | None = None self.flush_lock = asyncio.Lock() super().__init__(**kwargs, flush_lock=self.flush_lock) asyncio.create_task(self.periodic_flush()) - self.log_queue: List[StandardLoggingPayload] = [] - self.audit_log_queue: List[StandardAuditLogPayload] = [] + self.log_queue: list[StandardLoggingPayload] = [] + self.audit_log_queue: list[StandardAuditLogPayload] = [] @staticmethod def _build_api_endpoint(endpoint: str, dcr_immutable_id: str, stream_name: str) -> str: @@ -150,16 +150,16 @@ class AzureSentinelLogger(CustomBatchLogger): assert self.client_id is not None, "client_id is required" assert self.client_secret is not None, "client_secret is required" - token_url = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" + token_url: Final = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" - token_data = { + token_data: Final = { "client_id": self.client_id, "client_secret": self.client_secret, "scope": self.oauth_scope, "grant_type": "client_credentials", } - response = await self.async_httpx_client.post( + response: Final = await self.async_httpx_client.post( url=token_url, data=token_data, headers={"Content-Type": "application/x-www-form-urlencoded"}, @@ -168,9 +168,9 @@ class AzureSentinelLogger(CustomBatchLogger): if response.status_code != 200: raise Exception(f"Failed to get OAuth2 token: {response.status_code} - {response.text}") - token_response = response.json() + token_response: Final = response.json() self.oauth_token = token_response.get("access_token") - expires_in = token_response.get("expires_in", 3600) + expires_in: Final = token_response.get("expires_in", 3600) if not self.oauth_token: raise Exception("OAuth2 token response did not contain access_token") @@ -192,7 +192,7 @@ class AzureSentinelLogger(CustomBatchLogger): """ try: verbose_logger.debug("Azure Sentinel: Logging - Enters logging function for model %s", kwargs) - standard_logging_payload = kwargs.get("standard_logging_object", None) + standard_logging_payload: Final = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: verbose_logger.warning("Azure Sentinel: standard_logging_object not found in kwargs") @@ -204,8 +204,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"Azure Sentinel Layer Error - {str(e)}\n{traceback.format_exc()}") - pass + verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc()) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ @@ -223,7 +222,7 @@ class AzureSentinelLogger(CustomBatchLogger): "Azure Sentinel: Logging - Enters failure logging function for model %s", kwargs, ) - standard_logging_payload = kwargs.get("standard_logging_object", None) + standard_logging_payload: Final = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: verbose_logger.warning("Azure Sentinel: standard_logging_object not found in kwargs") @@ -235,8 +234,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"Azure Sentinel Layer Error - {str(e)}\n{traceback.format_exc()}") - pass + verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc()) async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: """ @@ -259,8 +257,7 @@ class AzureSentinelLogger(CustomBatchLogger): await self.async_send_audit_batch() except Exception as e: - verbose_logger.exception(f"Azure Sentinel Audit Log Layer Error - {str(e)}\n{traceback.format_exc()}") - pass + verbose_logger.exception("Azure Sentinel Audit Log Layer Error - %s\n%s", e, traceback.format_exc()) async def async_send_batch(self): """ @@ -287,7 +284,7 @@ class AzureSentinelLogger(CustomBatchLogger): async def _async_send_batch_to_api( self, - log_queue: List[Union[StandardLoggingPayload, StandardAuditLogPayload]], + log_queue: list[StandardLoggingPayload | StandardAuditLogPayload], api_endpoint: str, log_type: str, ) -> None: @@ -298,14 +295,14 @@ class AzureSentinelLogger(CustomBatchLogger): verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type) # Get OAuth2 token - bearer_token = await self._get_oauth_token() + bearer_token: Final = await self._get_oauth_token() # Convert log queue to JSON array format expected by Logs Ingestion API # Each log entry should be a JSON object in the array - body = safe_dumps(log_queue) + body: Final = safe_dumps(log_queue) # Set headers for Logs Ingestion API - headers = { + headers: Final = { "Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json", } @@ -327,7 +324,7 @@ class AzureSentinelLogger(CustomBatchLogger): ) except Exception as e: - verbose_logger.exception(f"Azure Sentinel Error sending batch API - {str(e)}\n{traceback.format_exc()}") + verbose_logger.exception("Azure Sentinel Error sending batch API - %s\n%s", e, traceback.format_exc()) finally: log_queue.clear() diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index 5ccd1a86bff..cb7175691df 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -1,20 +1,24 @@ import asyncio import os import time -from litellm._uuid import uuid from datetime import datetime, timedelta -from typing import List, Optional +from typing import Final from litellm._logging import verbose_logger -from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, AZURE_STORAGE_MSFT_VERSION +from litellm._uuid import uuid +from litellm.constants import ( + _DEFAULT_TTL_FOR_HTTPX_CLIENTS, + AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX, + AZURE_STORAGE_MSFT_VERSION, +) from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) -from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.types.utils import StandardLoggingPayload @@ -30,35 +34,46 @@ class AzureBlobStorageLogger(CustomBatchLogger): self.tenant_id = os.getenv("AZURE_STORAGE_TENANT_ID") self.client_id = os.getenv("AZURE_STORAGE_CLIENT_ID") self.client_secret = os.getenv("AZURE_STORAGE_CLIENT_SECRET") - self.azure_storage_account_key: Optional[str] = os.getenv("AZURE_STORAGE_ACCOUNT_KEY") + self.azure_storage_account_key: str | None = os.getenv("AZURE_STORAGE_ACCOUNT_KEY") # Required Env Variables for Azure Storage - _azure_storage_account_name = os.getenv("AZURE_STORAGE_ACCOUNT_NAME") + _azure_storage_account_name: Final = os.getenv("AZURE_STORAGE_ACCOUNT_NAME") if not _azure_storage_account_name: raise ValueError("Missing required environment variable: AZURE_STORAGE_ACCOUNT_NAME") self.azure_storage_account_name: str = _azure_storage_account_name - _azure_storage_file_system = os.getenv("AZURE_STORAGE_FILE_SYSTEM") + _azure_storage_file_system: Final = os.getenv("AZURE_STORAGE_FILE_SYSTEM") if not _azure_storage_file_system: raise ValueError("Missing required environment variable: AZURE_STORAGE_FILE_SYSTEM") self.azure_storage_file_system: str = _azure_storage_file_system + self.azure_storage_endpoint_suffix: str = ( + os.getenv("AZURE_STORAGE_ENDPOINT_SUFFIX") or AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX + ) self._service_client = None # Time that the azure service client expires, in order to reset the connection pool and keep it fresh - self._service_client_timeout: Optional[float] = None + self._service_client_timeout: float | None = None # Internal variables used for Token based authentication - self.azure_auth_token: Optional[str] = None # the Azure AD token to use for Azure Storage API requests - self.token_expiry: Optional[datetime] = None # the expiry time of the currentAzure AD token + self.azure_auth_token: str | None = None # the Azure AD token to use for Azure Storage API requests + self.token_expiry: datetime | None = None # the expiry time of the currentAzure AD token asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - self.log_queue: List[StandardLoggingPayload] = [] + self.log_queue: list[StandardLoggingPayload] = [] super().__init__(**kwargs, flush_lock=self.flush_lock) except Exception as e: verbose_logger.exception( - f"AzureBlobStorageLogger: Got exception on init AzureBlobStorageLogger client {str(e)}" + "AzureBlobStorageLogger: Got exception on init AzureBlobStorageLogger client %s", e ) raise e + @property + def azure_storage_dfs_endpoint(self) -> str: + return f"https://{self.azure_storage_account_name}.dfs.{self.azure_storage_endpoint_suffix}" + + @property + def azure_storage_blob_endpoint(self) -> str: + return f"https://{self.azure_storage_account_name}.blob.{self.azure_storage_endpoint_suffix}" + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Azure Blob Storage @@ -72,7 +87,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): "AzureBlobStorageLogger: Logging - Enters logging function for model %s", kwargs, ) - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") + standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") @@ -80,8 +95,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self.log_queue.append(standard_logging_payload) except Exception as e: - verbose_logger.exception(f"AzureBlobStorageLogger Layer Error - {str(e)}") - pass + verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ @@ -96,15 +110,14 @@ class AzureBlobStorageLogger(CustomBatchLogger): "AzureBlobStorageLogger: Logging - Enters logging function for model %s", kwargs, ) - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") + standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise ValueError("standard_logging_payload is not set") self.log_queue.append(standard_logging_payload) except Exception as e: - verbose_logger.exception(f"AzureBlobStorageLogger Layer Error - {str(e)}") - pass + verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e) async def async_send_batch(self): """ @@ -127,7 +140,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): await self.async_upload_payload_to_azure_blob_storage(payload=payload) except Exception as e: - verbose_logger.exception(f"AzureBlobStorageLogger Error sending batch API - {str(e)}") + verbose_logger.exception("AzureBlobStorageLogger Error sending batch API - %s", e) async def async_upload_payload_to_azure_blob_storage(self, payload: StandardLoggingPayload): """ @@ -142,49 +155,49 @@ class AzureBlobStorageLogger(CustomBatchLogger): else: # Get a valid token instead of always requesting a new one await self.set_valid_azure_ad_token() - async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) - json_payload = safe_dumps(payload) + "\n" # Add newline for each log entry - payload_bytes = json_payload.encode("utf-8") - filename = f"{payload.get('id') or str(uuid.uuid4())}.json" - base_url = f"https://{self.azure_storage_account_name}.dfs.core.windows.net/{self.azure_storage_file_system}/{filename}" + async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + json_payload: Final = safe_dumps(payload) + "\n" # Add newline for each log entry + payload_bytes: Final = json_payload.encode("utf-8") + filename: Final = f"{payload.get('id') or str(uuid.uuid4())}.json" + base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{filename}" # Execute the 3-step upload process await self._create_file(async_client, base_url) await self._append_data(async_client, base_url, json_payload) await self._flush_data(async_client, base_url, len(payload_bytes)) - verbose_logger.debug(f"Successfully uploaded log to Azure Blob Storage: {filename}") + verbose_logger.debug("Successfully uploaded log to Azure Blob Storage: %s", filename) except Exception as e: - verbose_logger.exception(f"Error uploading to Azure Blob Storage: {str(e)}") + verbose_logger.exception("Error uploading to Azure Blob Storage: %s", e) raise e async def _create_file(self, client: AsyncHTTPHandler, base_url: str): """Helper method to create the file resource""" try: - verbose_logger.debug(f"Creating file resource at: {base_url}") - headers = { + verbose_logger.debug("Creating file resource at: %s", base_url) + headers: Final = { "x-ms-version": AZURE_STORAGE_MSFT_VERSION, "Content-Length": "0", "Authorization": f"Bearer {self.azure_auth_token}", } - response = await client.put(f"{base_url}?resource=file", headers=headers) + response: Final = await client.put(f"{base_url}?resource=file", headers=headers) response.raise_for_status() verbose_logger.debug("Successfully created file resource") except Exception as e: - verbose_logger.exception(f"Error creating file resource: {str(e)}") + verbose_logger.exception("Error creating file resource: %s", e) raise async def _append_data(self, client: AsyncHTTPHandler, base_url: str, json_payload: str): """Helper method to append data to the file""" try: - verbose_logger.debug(f"Appending data to file: {base_url}") - headers = { + verbose_logger.debug("Appending data to file: %s", base_url) + headers: Final = { "x-ms-version": AZURE_STORAGE_MSFT_VERSION, "Content-Type": "application/json", "Authorization": f"Bearer {self.azure_auth_token}", } - response = await client.patch( + response: Final = await client.patch( f"{base_url}?action=append&position=0", headers=headers, data=json_payload, @@ -192,23 +205,23 @@ class AzureBlobStorageLogger(CustomBatchLogger): response.raise_for_status() verbose_logger.debug("Successfully appended data") except Exception as e: - verbose_logger.exception(f"Error appending data: {str(e)}") + verbose_logger.exception("Error appending data: %s", e) raise async def _flush_data(self, client: AsyncHTTPHandler, base_url: str, position: int): """Helper method to flush the data""" try: - verbose_logger.debug(f"Flushing data at position {position}") - headers = { + verbose_logger.debug("Flushing data at position %s", position) + headers: Final = { "x-ms-version": AZURE_STORAGE_MSFT_VERSION, "Content-Length": "0", "Authorization": f"Bearer {self.azure_auth_token}", } - response = await client.patch(f"{base_url}?action=flush&position={position}", headers=headers) + response: Final = await client.patch(f"{base_url}?action=flush&position={position}", headers=headers) response.raise_for_status() verbose_logger.debug("Successfully flushed data") except Exception as e: - verbose_logger.exception(f"Error flushing data: {str(e)}") + verbose_logger.exception("Error flushing data: %s", e) raise ####### Helper methods to managing Authentication to Azure Storage ####### @@ -232,13 +245,13 @@ class AzureBlobStorageLogger(CustomBatchLogger): ) # Token typically expires in 1 hour self.token_expiry = datetime.now() + timedelta(hours=1) - verbose_logger.debug(f"New token will expire at {self.token_expiry}") + verbose_logger.debug("New token will expire at %s", self.token_expiry) def get_azure_ad_token_from_azure_storage( self, - tenant_id: Optional[str], - client_id: Optional[str], - client_secret: Optional[str], + tenant_id: str | None, + client_id: str | None, + client_secret: str | None, ) -> str: """ Gets Azure AD token to use for Azure Storage API requests @@ -256,13 +269,13 @@ class AzureBlobStorageLogger(CustomBatchLogger): if client_secret is None: raise ValueError("Missing required environment variable: AZURE_STORAGE_CLIENT_SECRET") - token_provider = get_azure_ad_token_from_entra_id( + token_provider: Final = get_azure_ad_token_from_entra_id( tenant_id=tenant_id, client_id=client_id, client_secret=client_secret, scope="https://storage.azure.com/.default", ) - token = token_provider() + token: Final = token_provider() verbose_logger.debug("azure auth token %s", token) @@ -298,7 +311,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client = None if not self._service_client: self._service_client = DataLakeServiceClient( - account_url=f"https://{self.azure_storage_account_name}.dfs.core.windows.net", + account_url=self.azure_storage_dfs_endpoint, credential=self.azure_storage_account_key, ) self._service_client_timeout = time.time() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS @@ -313,31 +326,31 @@ class AzureBlobStorageLogger(CustomBatchLogger): # Create an async service client - service_client = await self.get_service_client() + service_client: Final = await self.get_service_client() # Get file system client - file_system_client = service_client.get_file_system_client(file_system=self.azure_storage_file_system) + file_system_client: Final = service_client.get_file_system_client(file_system=self.azure_storage_file_system) try: # Create directory with today's date from datetime import datetime - today = datetime.now().strftime("%Y-%m-%d") - directory_client = file_system_client.get_directory_client(today) + today: Final = datetime.now().strftime("%Y-%m-%d") + directory_client: Final = file_system_client.get_directory_client(today) # check if the directory exists if not await directory_client.exists(): await directory_client.create_directory() - verbose_logger.debug(f"Created directory: {today}") + verbose_logger.debug("Created directory: %s", today) # Create a file client - file_name = f"{payload.get('id') or str(uuid.uuid4())}.json" - file_client = directory_client.get_file_client(file_name) + file_name: Final = f"{payload.get('id') or str(uuid.uuid4())}.json" + file_client: Final = directory_client.get_file_client(file_name) # Create the file await file_client.create_file() # Content to append - content = safe_dumps(payload).encode("utf-8") + content: Final = safe_dumps(payload).encode("utf-8") # Append content to the file await file_client.append_data(data=content, offset=0, length=len(content)) @@ -345,7 +358,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): # Flush the content to finalize the file await file_client.flush_data(position=len(content), offset=0) - verbose_logger.debug(f"Successfully uploaded and wrote to {today}/{file_name}") + verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name) except Exception as e: - verbose_logger.exception(f"Error occurred: {str(e)}") + verbose_logger.exception("Error occurred: %s", e) diff --git a/litellm/integrations/bitbucket/__init__.py b/litellm/integrations/bitbucket/__init__.py index 2b9bd568e32..17ef5f65eb5 100644 --- a/litellm/integrations/bitbucket/__init__.py +++ b/litellm/integrations/bitbucket/__init__.py @@ -1,16 +1,17 @@ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final if TYPE_CHECKING: - from .bitbucket_prompt_manager import BitBucketPromptManager - from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec from litellm.integrations.custom_prompt_management import CustomPromptManagement + from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec + + from .bitbucket_prompt_manager import BitBucketPromptManager from litellm.types.prompts.init_prompts import SupportedPromptIntegrations from .bitbucket_prompt_manager import BitBucketPromptManager # Global instances -global_bitbucket_config: Optional[dict] = None +global_bitbucket_config: Final[dict | None] = None def set_global_bitbucket_config(config: dict) -> None: @@ -33,14 +34,14 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom """ Initialize a prompt from a BitBucket repository. """ - bitbucket_config = getattr(litellm_params, "bitbucket_config", None) - prompt_id = getattr(litellm_params, "prompt_id", None) + bitbucket_config: Final = getattr(litellm_params, "bitbucket_config", None) + prompt_id: Final = getattr(litellm_params, "prompt_id", None) if not bitbucket_config: raise ValueError("bitbucket_config is required for BitBucket prompt integration") try: - bitbucket_prompt_manager = BitBucketPromptManager( + bitbucket_prompt_manager: Final = BitBucketPromptManager( bitbucket_config=bitbucket_config, prompt_id=prompt_id, ) @@ -50,13 +51,13 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom raise e -prompt_initializer_registry = { +prompt_initializer_registry: Final = { SupportedPromptIntegrations.BITBUCKET.value: prompt_initializer, } # Export public API __all__ = [ "BitBucketPromptManager", - "set_global_bitbucket_config", "global_bitbucket_config", + "set_global_bitbucket_config", ] diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index c02d56811a7..e06e5ab358f 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -4,7 +4,7 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. import base64 import urllib.parse -from typing import Any, Dict, List, Optional +from typing import Any, Final from litellm.llms.custom_httpx.http_handler import HTTPHandler @@ -13,7 +13,7 @@ def _sanitize_file_path(file_path: str) -> str: """Reject path traversal and URL-encode each path segment.""" if "#" in file_path or "?" in file_path: raise ValueError(f"Invalid file path {file_path!r}: contains URL special characters") - parts = file_path.split("/") + parts: Final = file_path.split("/") for part in parts: if part == "..": raise ValueError(f"Invalid file path {file_path!r}: path traversal detected") @@ -31,7 +31,7 @@ class BitBucketClient: - Branch-specific file fetching """ - def __init__(self, config: Dict[str, Any]): + def __init__(self, config: dict[str, Any]): """ Initialize the BitBucket client. @@ -64,8 +64,8 @@ class BitBucketClient: if self.auth_method == "basic" and self.username: # Use basic auth with username and app password - credentials = f"{self.username}:{self.access_token}" - encoded_credentials = base64.b64encode(credentials.encode()).decode() + credentials: Final = f"{self.username}:{self.access_token}" + encoded_credentials: Final = base64.b64encode(credentials.encode()).decode() self.headers["Authorization"] = f"Basic {encoded_credentials}" else: # Use token-based authentication (default) @@ -74,7 +74,7 @@ class BitBucketClient: # Initialize HTTPHandler self.http_handler = HTTPHandler() - def get_file_content(self, file_path: str) -> Optional[str]: + def get_file_content(self, file_path: str) -> str | None: """ Fetch the content of a file from the BitBucket repository. @@ -84,11 +84,11 @@ class BitBucketClient: Returns: File content as string, or None if file not found """ - safe_path = _sanitize_file_path(file_path) - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}" + safe_path: Final = _sanitize_file_path(file_path) + url: Final = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}" try: - response = self.http_handler.get(url, headers=self.headers) + response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() # BitBucket returns file content as base64 encoded @@ -117,7 +117,7 @@ class BitBucketClient: else: raise Exception(f"Error fetching file '{file_path}': {e}") - def list_files(self, directory_path: str = "", file_extension: str = ".prompt") -> List[str]: + def list_files(self, directory_path: str = "", file_extension: str = ".prompt") -> list[str]: """ List files in a directory with a specific extension. @@ -128,15 +128,15 @@ class BitBucketClient: Returns: List of file paths """ - safe_dir = _sanitize_file_path(directory_path) if directory_path else "" - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_dir}" + safe_dir: Final = _sanitize_file_path(directory_path) if directory_path else "" + url: Final = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_dir}" try: - response = self.http_handler.get(url, headers=self.headers) + response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data = response.json() - files = [] + data: Final = response.json() + files: Final = [] for item in data.get("values", []): if item.get("type") == "commit_file": @@ -162,17 +162,17 @@ class BitBucketClient: else: raise Exception(f"Error listing files in '{directory_path}': {e}") - def get_repository_info(self) -> Dict[str, Any]: + def get_repository_info(self) -> dict[str, Any]: """ Get information about the repository. Returns: Dictionary containing repository information """ - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}" + url: Final = f"{self.base_url}/repositories/{self.workspace}/{self.repository}" try: - response = self.http_handler.get(url, headers=self.headers) + response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() return response.json() except Exception as e: @@ -191,25 +191,25 @@ class BitBucketClient: except Exception: return False - def get_branches(self) -> List[Dict[str, Any]]: + def get_branches(self) -> list[dict[str, Any]]: """ Get list of branches in the repository. Returns: List of branch information dictionaries """ - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/refs/branches" + url: Final = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/refs/branches" try: - response = self.http_handler.get(url, headers=self.headers) + response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data = response.json() + data: Final = response.json() return data.get("values", []) except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str) -> Optional[Dict[str, Any]]: + def get_file_metadata(self, file_path: str) -> dict[str, Any] | None: """ Get metadata about a file (size, last modified, etc.). @@ -219,15 +219,15 @@ class BitBucketClient: Returns: Dictionary containing file metadata, or None if file not found """ - safe_path = _sanitize_file_path(file_path) - url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}" + safe_path: Final = _sanitize_file_path(file_path) + url: Final = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}" try: # Use GET with Range header to get just the headers (HEAD equivalent) - headers = self.headers.copy() + headers: Final = self.headers.copy() headers["Range"] = "bytes=0-0" # Request only first byte to get headers - response = self.http_handler.get(url, headers=headers) + response: Final = self.http_handler.get(url, headers=headers) response.raise_for_status() return { diff --git a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py index 6dca4d76c04..88fd7dc55dc 100644 --- a/litellm/integrations/bitbucket/bitbucket_prompt_manager.py +++ b/litellm/integrations/bitbucket/bitbucket_prompt_manager.py @@ -3,7 +3,7 @@ BitBucket prompt manager that integrates with LiteLLM's prompt management system Fetches .prompt files from BitBucket repositories and provides team-based access control. """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Final from jinja2 import DictLoader, select_autoescape from jinja2.sandbox import ImmutableSandboxedEnvironment @@ -34,8 +34,8 @@ class BitBucketPromptTemplate: self, template_id: str, content: str, - metadata: Dict[str, Any], - model: Optional[str] = None, + metadata: dict[str, Any], + model: str | None = None, ): self.template_id = template_id self.content = content @@ -65,12 +65,12 @@ class BitBucketTemplateManager: def __init__( self, - bitbucket_config: Dict[str, Any], - prompt_id: Optional[str] = None, + bitbucket_config: dict[str, Any], + prompt_id: str | None = None, ): self.bitbucket_config = bitbucket_config self.prompt_id = prompt_id - self.prompts: Dict[str, BitBucketPromptTemplate] = {} + self.prompts: dict[str, BitBucketPromptTemplate] = {} self.bitbucket_client = BitBucketClient(bitbucket_config) # Templates fetched from a BitBucket repo are not trustworthy: @@ -99,10 +99,10 @@ class BitBucketTemplateManager: """Load a specific .prompt file from BitBucket.""" try: # Fetch the .prompt file from BitBucket - prompt_content = self.bitbucket_client.get_file_content(f"{prompt_id}.prompt") + prompt_content: Final = self.bitbucket_client.get_file_content(f"{prompt_id}.prompt") if prompt_content: - template = self._parse_prompt_file(prompt_content, prompt_id) + template: Final = self._parse_prompt_file(prompt_content, prompt_id) self.prompts[prompt_id] = template except Exception as e: raise Exception(f"Failed to load prompt '{prompt_id}' from BitBucket: {e}") @@ -111,7 +111,7 @@ class BitBucketTemplateManager: """Parse a .prompt file content and extract metadata and template.""" # Split frontmatter and content if content.startswith("---"): - parts = content.split("---", 2) + parts: Final = content.split("---", 2) if len(parts) >= 3: frontmatter_str = parts[1].strip() template_content = parts[2].strip() @@ -123,7 +123,7 @@ class BitBucketTemplateManager: template_content = content # Parse YAML frontmatter - metadata: Dict[str, Any] = {} + metadata: dict[str, Any] = {} if frontmatter_str: try: import yaml @@ -141,9 +141,9 @@ class BitBucketTemplateManager: metadata=metadata, ) - def _parse_yaml_basic(self, yaml_str: str) -> Dict[str, Any]: + def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]: """Basic YAML parser for simple cases when PyYAML is not available.""" - result: Dict[str, Any] = {} + result: Final[dict[str, Any]] = {} for line in yaml_str.split("\n"): line = line.strip() if ":" in line and not line.startswith("#"): @@ -162,21 +162,21 @@ class BitBucketTemplateManager: result[key] = value.strip("\"'") return result - def render_template(self, template_id: str, variables: Optional[Dict[str, Any]] = None) -> str: + def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str: """Render a template with the given variables.""" if template_id not in self.prompts: raise ValueError(f"Template '{template_id}' not found") - template = self.prompts[template_id] - jinja_template = self.jinja_env.from_string(template.content) + template: Final = self.prompts[template_id] + jinja_template: Final = self.jinja_env.from_string(template.content) return jinja_template.render(**(variables or {})) - def get_template(self, template_id: str) -> Optional[BitBucketPromptTemplate]: + def get_template(self, template_id: str) -> BitBucketPromptTemplate | None: """Get a template by ID.""" return self.prompts.get(template_id) - def list_templates(self) -> List[str]: + def list_templates(self) -> list[str]: """List all available template IDs.""" return list(self.prompts.keys()) @@ -209,12 +209,12 @@ class BitBucketPromptManager(CustomPromptManagement): def __init__( self, - bitbucket_config: Dict[str, Any], - prompt_id: Optional[str] = None, + bitbucket_config: dict[str, Any], + prompt_id: str | None = None, ): self.bitbucket_config = bitbucket_config self.prompt_id = prompt_id - self._prompt_manager: Optional[BitBucketTemplateManager] = None + self._prompt_manager: BitBucketTemplateManager | None = None @property def integration_name(self) -> str: @@ -234,8 +234,8 @@ class BitBucketPromptManager(CustomPromptManagement): def get_prompt_template( self, prompt_id: str, - prompt_variables: Optional[Dict[str, Any]] = None, - ) -> Tuple[str, Dict[str, Any]]: + prompt_variables: dict[str, Any] | None = None, + ) -> tuple[str, dict[str, Any]]: """ Get a prompt template and render it with variables. @@ -246,15 +246,15 @@ class BitBucketPromptManager(CustomPromptManagement): Returns: Tuple of (rendered_prompt, metadata) """ - template = self.prompt_manager.get_template(prompt_id) + template: Final = self.prompt_manager.get_template(prompt_id) if not template: raise ValueError(f"Prompt template '{prompt_id}' not found") # Render the template - rendered_prompt = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) + rendered_prompt: Final = self.prompt_manager.render_template(prompt_id, prompt_variables or {}) # Extract metadata - metadata = { + metadata: Final = { "model": template.model, "temperature": template.temperature, "max_tokens": template.max_tokens, @@ -265,14 +265,14 @@ class BitBucketPromptManager(CustomPromptManagement): def pre_call_hook( self, - user_id: Optional[str], - messages: List[AllMessageValues], - function_call: Optional[Union[Dict[str, Any], str]] = None, - litellm_params: Optional[Dict[str, Any]] = None, - prompt_id: Optional[str] = None, - prompt_variables: Optional[Dict[str, Any]] = None, + user_id: str | None, + messages: list[AllMessageValues], + function_call: dict[str, Any] | str | None = None, + litellm_params: dict[str, Any] | None = None, + prompt_id: str | None = None, + prompt_variables: dict[str, Any] | None = None, **kwargs, - ) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]: + ) -> tuple[list[AllMessageValues], dict[str, Any] | None]: """ Pre-call hook that processes the prompt template before making the LLM call. """ @@ -284,12 +284,12 @@ class BitBucketPromptManager(CustomPromptManagement): rendered_prompt, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Parse the rendered prompt into messages - parsed_messages = self._parse_prompt_to_messages(rendered_prompt) + parsed_messages: Final = self._parse_prompt_to_messages(rendered_prompt) # Merge with existing messages if parsed_messages: # If we have parsed messages, use them instead of the original messages - final_messages: List[AllMessageValues] = parsed_messages + final_messages: list[AllMessageValues] = parsed_messages else: # If no messages were parsed, prepend the prompt to existing messages final_messages = [ @@ -320,16 +320,16 @@ class BitBucketPromptManager(CustomPromptManagement): # Log error but don't fail the call import litellm - litellm._logging.verbose_proxy_logger.error(f"Error in BitBucket prompt pre_call_hook: {e}") + litellm._logging.verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e) return messages, litellm_params - def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]: + def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]: """ Parse prompt content into a list of messages. Handles both simple prompts and multi-role conversations. """ messages = [] - lines = prompt_content.strip().split("\n") + lines: Final = prompt_content.strip().split("\n") current_role = None current_content = [] @@ -385,13 +385,13 @@ class BitBucketPromptManager(CustomPromptManagement): def post_call_hook( self, - user_id: Optional[str], + user_id: str | None, response: Any, - input_messages: List[AllMessageValues], - function_call: Optional[Union[Dict[str, Any], str]] = None, - litellm_params: Optional[Dict[str, Any]] = None, - prompt_id: Optional[str] = None, - prompt_variables: Optional[Dict[str, Any]] = None, + input_messages: list[AllMessageValues], + function_call: dict[str, Any] | str | None = None, + litellm_params: dict[str, Any] | None = None, + prompt_id: str | None = None, + prompt_variables: dict[str, Any] | None = None, **kwargs, ) -> Any: """ @@ -399,7 +399,7 @@ class BitBucketPromptManager(CustomPromptManagement): """ return response - def get_available_prompts(self) -> List[str]: + def get_available_prompts(self) -> list[str]: """Get list of available prompt IDs.""" return self.prompt_manager.list_templates() @@ -411,8 +411,8 @@ class BitBucketPromptManager(CustomPromptManagement): def should_run_prompt_management( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], + prompt_id: str | None, + prompt_spec: PromptSpec | None, dynamic_callback_params: StandardCallbackDynamicParams, ) -> bool: """ @@ -425,12 +425,12 @@ class BitBucketPromptManager(CustomPromptManagement): def _compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_spec: PromptSpec | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """ Compile a BitBucket prompt template into a PromptManagementClient structure. @@ -453,13 +453,13 @@ class BitBucketPromptManager(CustomPromptManagement): rendered_prompt, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables) # Convert rendered content to chat messages - messages = self._parse_prompt_to_messages(rendered_prompt) + messages: Final = self._parse_prompt_to_messages(rendered_prompt) # Extract model from metadata (if specified) - template_model = prompt_metadata.get("model") + template_model: Final = prompt_metadata.get("model") # Extract optional parameters from metadata - optional_params = {} + optional_params: Final = {} for param in [ "temperature", "max_tokens", @@ -483,12 +483,12 @@ class BitBucketPromptManager(CustomPromptManagement): async def async_compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """ Async version of compile prompt helper. Since BitBucket operations use sync client, @@ -509,17 +509,17 @@ class BitBucketPromptManager(CustomPromptManagement): def get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Get chat completion prompt from BitBucket and return processed model, messages, and parameters. """ @@ -539,19 +539,19 @@ class BitBucketPromptManager(CustomPromptManagement): async def async_get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, litellm_logging_obj: LiteLLMLoggingObj, - prompt_spec: Optional[PromptSpec] = None, - tools: Optional[List[Dict]] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + tools: list[dict] | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Async version - delegates to PromptManagementBase async implementation. """ diff --git a/litellm/integrations/braintrust_logging.py b/litellm/integrations/braintrust_logging.py index 686c37d3e17..cc87b217dd0 100644 --- a/litellm/integrations/braintrust_logging.py +++ b/litellm/integrations/braintrust_logging.py @@ -3,15 +3,15 @@ import os from datetime import datetime -from typing import Dict, Optional +from typing import Final import httpx import litellm from litellm import verbose_logger from litellm.integrations.braintrust_mock_client import ( - should_use_braintrust_mock, create_mock_braintrust_client, + should_use_braintrust_mock, ) from litellm.integrations.custom_logger import CustomLogger from litellm.llms.custom_httpx.http_handler import ( @@ -21,7 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.utils import print_verbose -API_BASE = "https://api.braintrustdata.com/v1" +API_BASE: Final = "https://api.braintrustdata.com/v1" def get_utc_datetime(): @@ -34,7 +34,7 @@ def get_utc_datetime(): class BraintrustLogger(CustomLogger): - def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None) -> None: + def __init__(self, api_key: str | None = None, api_base: str | None = None) -> None: super().__init__() self.is_mock_mode = should_use_braintrust_mock() if self.is_mock_mode: @@ -48,23 +48,23 @@ class BraintrustLogger(CustomLogger): "Authorization": "Bearer " + self.api_key, "Content-Type": "application/json", } - self._project_id_cache: Dict[str, str] = {} # Cache mapping project names to IDs + self._project_id_cache: dict[str, str] = {} # Cache mapping project names to IDs self.global_braintrust_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.global_braintrust_sync_http_handler = HTTPHandler() - def validate_environment(self, api_key: Optional[str]): + def validate_environment(self, api_key: str | None): """ Expects BRAINTRUST_API_KEY in the environment """ - missing_keys = [] + missing_keys: Final = [] if api_key is None and os.getenv("BRAINTRUST_API_KEY", None) is None: missing_keys.append("BRAINTRUST_API_KEY") if len(missing_keys) > 0: - raise Exception("Missing keys={} in environment.".format(missing_keys)) + raise Exception(f"Missing keys={missing_keys} in environment.") def get_project_id_sync(self, project_name: str) -> str: """ @@ -75,13 +75,13 @@ class BraintrustLogger(CustomLogger): return self._project_id_cache[project_name] try: - response = self.global_braintrust_sync_http_handler.post( + response: Final = self.global_braintrust_sync_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": project_name}, ) - project_dict = response.json() - project_id = project_dict["id"] + project_dict: Final = response.json() + project_id: Final = project_dict["id"] self._project_id_cache[project_name] = project_id return project_id except httpx.HTTPStatusError as e: @@ -95,42 +95,42 @@ class BraintrustLogger(CustomLogger): return self._project_id_cache[project_name] try: - response = await self.global_braintrust_http_handler.post( + response: Final = await self.global_braintrust_http_handler.post( f"{self.api_base}/project/register", headers=self.headers, json={"name": project_name}, ) - project_dict = response.json() - project_id = project_dict["id"] + project_dict: Final = response.json() + project_id: Final = project_dict["id"] self._project_id_cache[project_name] = project_id return project_id except httpx.HTTPStatusError as e: raise Exception(f"Failed to register project: {e.response.text}") async def create_default_project_and_experiment(self): - project = await self.global_braintrust_http_handler.post( + project: Final = await self.global_braintrust_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) - project_dict = project.json() + project_dict: Final = project.json() self.default_project_id = project_dict["id"] def create_sync_default_project_and_experiment(self): - project = self.global_braintrust_sync_http_handler.post( + project: Final = self.global_braintrust_sync_http_handler.post( f"{self.api_base}/project", headers=self.headers, json={"name": "litellm"} ) - project_dict = project.json() + project_dict: Final = project.json() self.default_project_id = project_dict["id"] def log_success_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") try: - litellm_call_id = kwargs.get("litellm_call_id") - standard_logging_object = kwargs.get("standard_logging_object", {}) - prompt = {"messages": kwargs.get("messages")} + litellm_call_id: Final = kwargs.get("litellm_call_id") + standard_logging_object: Final = kwargs.get("standard_logging_object", {}) + prompt: Final = {"messages": kwargs.get("messages")} output = None choices = [] @@ -147,13 +147,13 @@ class BraintrustLogger(CustomLogger): elif response_obj is not None and isinstance(response_obj, litellm.ImageResponse): output = response_obj["data"] - litellm_params = kwargs.get("litellm_params", {}) or {} - dynamic_metadata = litellm_params.get("metadata", {}) or {} + litellm_params: Final = kwargs.get("litellm_params", {}) or {} + dynamic_metadata: Final = litellm_params.get("metadata", {}) or {} # Get project_id from metadata or create default if needed project_id = dynamic_metadata.get("project_id") if project_id is None: - project_name = dynamic_metadata.get("project_name") + project_name: Final = dynamic_metadata.get("project_name") project_id = self.get_project_id_sync(project_name) if project_name else None if project_id is None: @@ -161,7 +161,7 @@ class BraintrustLogger(CustomLogger): self.create_sync_default_project_and_experiment() project_id = self.default_project_id - tags = [] + tags: Final = [] if isinstance(dynamic_metadata, dict): for key, value in dynamic_metadata.items(): @@ -178,10 +178,10 @@ class BraintrustLogger(CustomLogger): ): # support logging dynamic metadata to braintrust standard_logging_object[key] = value - cost = kwargs.get("response_cost", None) + cost: Final = kwargs.get("response_cost", None) - metrics: Optional[dict] = None - usage_obj = getattr(response_obj, "usage", None) + metrics: dict | None = None + usage_obj: Final = getattr(response_obj, "usage", None) if usage_obj and isinstance(usage_obj, litellm.Usage): litellm.utils.get_logging_id(start_time, response_obj) metrics = { @@ -195,7 +195,7 @@ class BraintrustLogger(CustomLogger): } # Allow metadata override for span name - span_name = dynamic_metadata.get("span_name", "Chat Completion") + span_name: Final = dynamic_metadata.get("span_name", "Chat Completion") # Span parents is a special case span_parents = dynamic_metadata.get("span_parents") @@ -205,13 +205,13 @@ class BraintrustLogger(CustomLogger): span_parents = [s.strip() for s in span_parents.split(",") if s.strip()] # Add optional span attributes only if present - span_attributes = { + span_attributes: Final = { "span_id": dynamic_metadata.get("span_id"), "root_span_id": dynamic_metadata.get("root_span_id"), "span_parents": span_parents, } - request_data = { + request_data: Final = { "id": litellm_call_id, "input": prompt["messages"], "metadata": standard_logging_object, @@ -254,9 +254,9 @@ class BraintrustLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): verbose_logger.debug("REACHES BRAINTRUST SUCCESS") try: - litellm_call_id = kwargs.get("litellm_call_id") - standard_logging_object = kwargs.get("standard_logging_object", {}) - prompt = {"messages": kwargs.get("messages")} + litellm_call_id: Final = kwargs.get("litellm_call_id") + standard_logging_object: Final = kwargs.get("standard_logging_object", {}) + prompt: Final = {"messages": kwargs.get("messages")} output = None choices = [] if response_obj is not None and ( @@ -272,13 +272,13 @@ class BraintrustLogger(CustomLogger): elif response_obj is not None and isinstance(response_obj, litellm.ImageResponse): output = response_obj["data"] - litellm_params = kwargs.get("litellm_params", {}) - dynamic_metadata = litellm_params.get("metadata", {}) or {} + litellm_params: Final = kwargs.get("litellm_params", {}) + dynamic_metadata: Final = litellm_params.get("metadata", {}) or {} # Get project_id from metadata or create default if needed project_id = dynamic_metadata.get("project_id") if project_id is None: - project_name = dynamic_metadata.get("project_name") + project_name: Final = dynamic_metadata.get("project_name") project_id = await self.get_project_id_async(project_name) if project_name else None if project_id is None: @@ -286,7 +286,7 @@ class BraintrustLogger(CustomLogger): await self.create_default_project_and_experiment() project_id = self.default_project_id - tags = [] + tags: Final = [] if isinstance(dynamic_metadata, dict): for key, value in dynamic_metadata.items(): @@ -303,10 +303,10 @@ class BraintrustLogger(CustomLogger): ): # support logging dynamic metadata to braintrust standard_logging_object[key] = value - cost = kwargs.get("response_cost", None) + cost: Final = kwargs.get("response_cost", None) - metrics: Optional[dict] = None - usage_obj = getattr(response_obj, "usage", None) + metrics: dict | None = None + usage_obj: Final = getattr(response_obj, "usage", None) if usage_obj and isinstance(usage_obj, litellm.Usage): litellm.utils.get_logging_id(start_time, response_obj) metrics = { @@ -318,14 +318,14 @@ class BraintrustLogger(CustomLogger): "end": end_time.timestamp(), } - api_call_start_time = kwargs.get("api_call_start_time") - completion_start_time = kwargs.get("completion_start_time") + api_call_start_time: Final = kwargs.get("api_call_start_time") + completion_start_time: Final = kwargs.get("completion_start_time") if api_call_start_time is not None and completion_start_time is not None: metrics["time_to_first_token"] = completion_start_time.timestamp() - api_call_start_time.timestamp() # Allow metadata override for span name - span_name = dynamic_metadata.get("span_name", "Chat Completion") + span_name: Final = dynamic_metadata.get("span_name", "Chat Completion") # Span parents is a special case span_parents = dynamic_metadata.get("span_parents") @@ -335,13 +335,13 @@ class BraintrustLogger(CustomLogger): span_parents = [s.strip() for s in span_parents.split(",") if s.strip()] # Add optional span attributes only if present - span_attributes = { + span_attributes: Final = { "span_id": dynamic_metadata.get("span_id"), "root_span_id": dynamic_metadata.get("root_span_id"), "span_parents": span_parents, } - request_data = { + request_data: Final = { "id": litellm_call_id, "input": prompt["messages"], "output": output, diff --git a/litellm/integrations/braintrust_mock_client.py b/litellm/integrations/braintrust_mock_client.py index e2b732d6e9c..07c01c58305 100644 --- a/litellm/integrations/braintrust_mock_client.py +++ b/litellm/integrations/braintrust_mock_client.py @@ -10,6 +10,7 @@ Usage: import os import time +from typing import Final from urllib.parse import urlparse from litellm._logging import verbose_logger @@ -22,7 +23,7 @@ from litellm.integrations.mock_client_factory import ( # Use factory for should_use_mock and MockResponse # Braintrust uses both HTTPHandler (sync) and AsyncHTTPHandler (async) # Braintrust needs endpoint-specific responses, so we use custom HTTPHandler.post patching -_config = MockClientConfig( +_config: Final = MockClientConfig( "BRAINTRUST", "BRAINTRUST_MOCK", default_latency_ms=100, @@ -51,7 +52,7 @@ _original_http_handler_post = None _mocks_initialized = False # Default mock latency in seconds -_MOCK_LATENCY_SECONDS = float(os.getenv("BRAINTRUST_MOCK_LATENCY_MS", "100")) / 1000.0 +_MOCK_LATENCY_SECONDS: Final = float(os.getenv("BRAINTRUST_MOCK_LATENCY_MS", "100")) / 1000.0 def _is_braintrust_url(url: str) -> bool: @@ -59,8 +60,8 @@ def _is_braintrust_url(url: str) -> bool: if not isinstance(url, str): return False - parsed = urlparse(url) - host = (parsed.hostname or "").lower() + parsed: Final = urlparse(url) + host: Final = (parsed.hostname or "").lower() if not host: return False @@ -89,12 +90,12 @@ def _mock_http_handler_post( """Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses.""" # Only mock Braintrust API calls if isinstance(url, str) and _is_braintrust_url(url): - verbose_logger.info(f"[BRAINTRUST MOCK] POST to {url}") + verbose_logger.info("[BRAINTRUST MOCK] POST to %s", url) time.sleep(_MOCK_LATENCY_SECONDS) # Return appropriate mock response based on endpoint if "/project" in url: # Project creation/retrieval/register endpoint - project_name = json.get("name", "litellm") if json else "litellm" + project_name: Final = json.get("name", "litellm") if json else "litellm" mock_data = {"id": f"mock-project-id-{project_name}", "name": project_name} elif "/project_logs" in url: # Log insertion endpoint diff --git a/litellm/integrations/cloudzero/cloudzero.py b/litellm/integrations/cloudzero/cloudzero.py index 121b1dc6967..f41bd885442 100644 --- a/litellm/integrations/cloudzero/cloudzero.py +++ b/litellm/integrations/cloudzero/cloudzero.py @@ -1,6 +1,6 @@ import os from datetime import datetime -from typing import TYPE_CHECKING, Any, List, Optional, cast +from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_logger @@ -25,9 +25,9 @@ class CloudZeroLogger(CustomLogger): def __init__( self, - api_key: Optional[str] = None, - connection_id: Optional[str] = None, - timezone: Optional[str] = None, + api_key: str | None = None, + connection_id: str | None = None, + timezone: str | None = None, **kwargs, ): """Initialize CloudZero logger with configuration from parameters or environment variables.""" @@ -38,7 +38,7 @@ class CloudZeroLogger(CustomLogger): self.connection_id = connection_id or os.getenv("CLOUDZERO_CONNECTION_ID") self.timezone = timezone or os.getenv("CLOUDZERO_TIMEZONE", "UTC") verbose_logger.debug( - f"CloudZero Logger initialized with connection ID: {self.connection_id}, timezone: {self.timezone}" + "CloudZero Logger initialized with connection ID: %s, timezone: %s", self.connection_id, self.timezone ) async def initialize_cloudzero_export_job(self): @@ -56,7 +56,7 @@ class CloudZeroLogger(CustomLogger): ) from litellm.proxy.proxy_server import proxy_logging_obj - pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager + pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager # if using redis, ensure only one pod exports the data at a time if pod_lock_manager and pod_lock_manager.redis_cache: @@ -80,9 +80,9 @@ class CloudZeroLogger(CustomLogger): from litellm.constants import CLOUDZERO_MAX_FETCHED_DATA_RECORDS - current_time_utc = datetime.now(timezone.utc) + current_time_utc: Final = datetime.now(timezone.utc) # Mitigates the possibility of missing spend if an hour is skipped due to a restart in an ephemeral environment - one_hour_ago_utc = current_time_utc - timedelta(minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES * 2) + one_hour_ago_utc: Final = current_time_utc - timedelta(minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES * 2) await self.export_usage_data( limit=CLOUDZERO_MAX_FETCHED_DATA_RECORDS, operation="replace_hourly", @@ -92,10 +92,10 @@ class CloudZeroLogger(CustomLogger): async def export_usage_data( self, - limit: Optional[int] = None, + limit: int | None = None, operation: str = "replace_hourly", - start_time_utc: Optional[datetime] = None, - end_time_utc: Optional[datetime] = None, + start_time_utc: datetime | None = None, + end_time_utc: datetime | None = None, ): """ Exports the usage data to CloudZero. @@ -122,7 +122,7 @@ class CloudZeroLogger(CustomLogger): ) # Initialize database connection and load data - database = LiteLLMDatabase() + database: Final = LiteLLMDatabase() verbose_logger.debug("CloudZero Logger: Loading usage data from database") data = await database.get_usage_data(limit=limit, start_time_utc=start_time_utc, end_time_utc=end_time_utc) @@ -130,33 +130,33 @@ class CloudZeroLogger(CustomLogger): verbose_logger.debug("CloudZero Logger: No usage data found to export") return - verbose_logger.debug(f"CloudZero Logger: Processing {len(data)} records") + verbose_logger.debug("CloudZero Logger: Processing %s records", len(data)) # Transform data to CloudZero CBF format - transformer = CBFTransformer() - cbf_data = transformer.transform(data) + transformer: Final = CBFTransformer() + cbf_data: Final = transformer.transform(data) if cbf_data.is_empty(): verbose_logger.warning("CloudZero Logger: No valid data after transformation") return # Send data to CloudZero - streamer = CloudZeroStreamer( + streamer: Final = CloudZeroStreamer( api_key=self.api_key, connection_id=self.connection_id, user_timezone=self.timezone, ) - verbose_logger.debug(f"CloudZero Logger: Transmitting {len(cbf_data)} records to CloudZero") + verbose_logger.debug("CloudZero Logger: Transmitting %s records to CloudZero", len(cbf_data)) streamer.send_batched(cbf_data, operation=operation) - verbose_logger.debug(f"CloudZero Logger: Successfully exported {len(cbf_data)} records to CloudZero") + verbose_logger.debug("CloudZero Logger: Successfully exported %s records to CloudZero", len(cbf_data)) except Exception as e: - verbose_logger.error(f"CloudZero Logger: Error exporting usage data: {str(e)}") + verbose_logger.error("CloudZero Logger: Error exporting usage data: %s", e) raise - async def dry_run_export_usage_data(self, limit: Optional[int] = 10000): + async def dry_run_export_usage_data(self, limit: int | None = 10000): """ Returns the data that would be exported to CloudZero without actually sending it. @@ -173,9 +173,9 @@ class CloudZeroLogger(CustomLogger): verbose_logger.debug("CloudZero Logger: Starting dry run export") # Initialize database connection and load data - database = LiteLLMDatabase() + database: Final = LiteLLMDatabase() verbose_logger.debug("CloudZero Logger: Loading usage data for dry run") - data = await database.get_usage_data(limit=limit) + data: Final = await database.get_usage_data(limit=limit) if data.is_empty(): verbose_logger.warning("CloudZero Dry Run: No usage data found") @@ -191,14 +191,14 @@ class CloudZeroLogger(CustomLogger): }, } - verbose_logger.debug(f"CloudZero Dry Run: Processing {len(data)} records...") + verbose_logger.debug("CloudZero Dry Run: Processing %s records...", len(data)) # Convert usage data to dict format for response - usage_data_sample = data.head(50).to_dicts() # Return first 50 rows + usage_data_sample: Final = data.head(50).to_dicts() # Return first 50 rows # Transform data to CloudZero CBF format - transformer = CBFTransformer() - cbf_data = transformer.transform(data) + transformer: Final = CBFTransformer() + cbf_data: Final = transformer.transform(data) if cbf_data.is_empty(): verbose_logger.warning("CloudZero Dry Run: No valid data after transformation") @@ -217,19 +217,19 @@ class CloudZeroLogger(CustomLogger): } # Convert CBF data to dict format for response - cbf_data_dict = cbf_data.to_dicts() + cbf_data_dict: Final = cbf_data.to_dicts() # Calculate summary statistics - total_cost = sum(record.get("cost/cost", 0) for record in cbf_data_dict) - unique_accounts = len( + total_cost: Final = sum(record.get("cost/cost", 0) for record in cbf_data_dict) + unique_accounts: Final = len( set(record.get("resource/account", "") for record in cbf_data_dict if record.get("resource/account")) ) - unique_services = len( + unique_services: Final = len( set(record.get("resource/service", "") for record in cbf_data_dict if record.get("resource/service")) ) - total_tokens = sum(record.get("usage/amount", 0) for record in cbf_data_dict) + total_tokens: Final = sum(record.get("usage/amount", 0) for record in cbf_data_dict) - verbose_logger.debug(f"CloudZero Logger: Dry run completed for {len(cbf_data)} records") + verbose_logger.debug("CloudZero Logger: Dry run completed for %s records", len(cbf_data)) return { "usage_data": usage_data_sample, @@ -244,8 +244,8 @@ class CloudZeroLogger(CustomLogger): } except Exception as e: - verbose_logger.error(f"CloudZero Logger: Error in dry run export: {str(e)}") - verbose_logger.error(f"CloudZero Dry Run Error: {str(e)}") + verbose_logger.error("CloudZero Logger: Error in dry run export: %s", e) + verbose_logger.error("CloudZero Dry Run Error: %s", e) raise def _display_cbf_data_on_screen(self, cbf_data): @@ -254,7 +254,7 @@ class CloudZeroLogger(CustomLogger): from rich.console import Console from rich.table import Table - console = Console() + console: Final = Console() if cbf_data.is_empty(): console.print("[yellow]No CBF data to display[/yellow]") @@ -263,10 +263,10 @@ class CloudZeroLogger(CustomLogger): console.print(f"\n[bold green]💰 CloudZero CBF Transformed Data ({len(cbf_data)} records)[/bold green]") # Convert to dicts for easier processing - records = cbf_data.to_dicts() + records: Final = cbf_data.to_dicts() # Create main CBF table - cbf_table = Table(show_header=True, header_style="bold cyan", box=SIMPLE, padding=(0, 1)) + cbf_table: Final = Table(show_header=True, header_style="bold cyan", box=SIMPLE, padding=(0, 1)) cbf_table.add_column("time/usage_start", style="blue", no_wrap=False) cbf_table.add_column("cost/cost", style="green", justify="right", no_wrap=False) cbf_table.add_column("entity_type", style="magenta", justify="right", no_wrap=False) @@ -316,16 +316,16 @@ class CloudZeroLogger(CustomLogger): console.print(cbf_table) # Show summary statistics - total_cost = sum(record.get("cost/cost", 0) for record in records) - unique_accounts = len( + total_cost: Final = sum(record.get("cost/cost", 0) for record in records) + unique_accounts: Final = len( set(record.get("resource/account", "") for record in records if record.get("resource/account")) ) - unique_services = len( + unique_services: Final = len( set(record.get("resource/service", "") for record in records if record.get("resource/service")) ) # Count total tokens from usage metrics - total_tokens = sum(record.get("usage/amount", 0) for record in records) + total_tokens: Final = sum(record.get("usage/amount", 0) for record in records) console.print("\n[bold blue]📊 CBF Summary[/bold blue]") console.print(f" Records: {len(records):,}") @@ -346,13 +346,13 @@ class CloudZeroLogger(CustomLogger): from litellm.constants import CLOUDZERO_EXPORT_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger - prometheus_loggers: List[CustomLogger] = litellm.logging_callback_manager.get_custom_loggers_for_type( + prometheus_loggers: Final[list[CustomLogger]] = litellm.logging_callback_manager.get_custom_loggers_for_type( callback_type=CloudZeroLogger ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s cloudzero loggers", len(prometheus_loggers)) if len(prometheus_loggers) > 0: - cloudzero_logger = cast(CloudZeroLogger, prometheus_loggers[0]) + cloudzero_logger: Final = cast(CloudZeroLogger, prometheus_loggers[0]) verbose_logger.debug( "Initializing remaining budget metrics as a cron job executing every %s minutes" % CLOUDZERO_EXPORT_INTERVAL_MINUTES diff --git a/litellm/integrations/cloudzero/cz_resource_names.py b/litellm/integrations/cloudzero/cz_resource_names.py index 15cb66002f7..9b685996a41 100644 --- a/litellm/integrations/cloudzero/cz_resource_names.py +++ b/litellm/integrations/cloudzero/cz_resource_names.py @@ -18,7 +18,7 @@ import re from enum import Enum -from typing import Any, cast +from typing import Any, Final, cast import litellm @@ -34,7 +34,6 @@ class CZRNGenerator: def __init__(self): """Initialize CZRN generator.""" - pass def create_from_litellm_data(self, row: dict[str, Any]) -> str: """Create a CZRN from LiteLLM daily spend data. @@ -49,20 +48,20 @@ class CZRNGenerator: - resource-type: 'llm-usage' (represents LLM usage/inference) - cloud-local-id: model """ - service_type = "litellm" - provider = self._normalize_provider(row.get("custom_llm_provider", "unknown")) - region = "cross-region" + service_type: Final = "litellm" + provider: Final = self._normalize_provider(row.get("custom_llm_provider", "unknown")) + region: Final = "cross-region" # Use the actual entity_id (team_id or user_id) as the owner account - team_id = row.get("team_id", "unknown") - owner_account_id = self._normalize_component(team_id) + team_id: Final = row.get("team_id", "unknown") + owner_account_id: Final = self._normalize_component(team_id) - resource_type = "llm-usage" + resource_type: Final = "llm-usage" # Create a unique identifier with just the model (entity info already in owner_account_id) - model = row.get("model", "unknown") + model: Final = row.get("model", "unknown") - cloud_local_id = model + cloud_local_id: Final = model return self.create_from_components( service_type=service_type, @@ -91,7 +90,7 @@ class CZRNGenerator: resource_type = self._normalize_component(resource_type) # cloud_local_id can contain pipes and other characters, so don't normalize it - czrn = f"czrn:{service_type}:{provider}:{region}:{owner_account_id}:{resource_type}:{cloud_local_id}" + czrn: Final = f"czrn:{service_type}:{provider}:{region}:{owner_account_id}:{resource_type}:{cloud_local_id}" if not self.is_valid(czrn): raise ValueError(f"Generated CZRN is invalid: {czrn}") @@ -107,7 +106,7 @@ class CZRNGenerator: Returns: (service_type, provider, region, owner_account_id, resource_type, cloud_local_id) """ - match = self.CZRN_REGEX.match(czrn) + match: Final = self.CZRN_REGEX.match(czrn) if not match: raise ValueError(f"Invalid CZRN format: {czrn}") @@ -116,7 +115,7 @@ class CZRNGenerator: def _normalize_provider(self, provider: str) -> str: """Normalize provider names to standard CZRN format.""" # Map common provider names to CZRN standards - provider_map = { + provider_map: Final = { litellm.LlmProviders.AZURE.value: "azure", litellm.LlmProviders.AZURE_AI.value: "azure", litellm.LlmProviders.ANTHROPIC.value: "anthropic", @@ -129,7 +128,7 @@ class CZRNGenerator: litellm.LlmProviders.TOGETHER_AI.value: "together-ai", } - normalized = provider.lower().replace("_", "-") + normalized: Final = provider.lower().replace("_", "-") # use litellm custom llm provider if not in provider_map if normalized not in provider_map: diff --git a/litellm/integrations/cloudzero/cz_stream_api.py b/litellm/integrations/cloudzero/cz_stream_api.py index 47d6f7474a2..1e2fa318786 100644 --- a/litellm/integrations/cloudzero/cz_stream_api.py +++ b/litellm/integrations/cloudzero/cz_stream_api.py @@ -20,7 +20,7 @@ import zoneinfo from datetime import datetime, timezone -from typing import Any, Optional, Union +from typing import Any, Final import httpx import polars as pl @@ -30,7 +30,7 @@ from rich.console import Console class CloudZeroStreamer: """Stream CBF data to CloudZero AnyCost API with proper batching and timezone handling.""" - def __init__(self, api_key: str, connection_id: str, user_timezone: Optional[str] = None): + def __init__(self, api_key: str, connection_id: str, user_timezone: str | None = None): """Initialize CloudZero streamer with credentials.""" self.api_key = api_key self.connection_id = connection_id @@ -38,7 +38,7 @@ class CloudZeroStreamer: self.console = Console() # Set timezone - default to UTC - self.user_timezone: Union[zoneinfo.ZoneInfo, timezone] + self.user_timezone: zoneinfo.ZoneInfo | timezone if user_timezone: try: self.user_timezone = zoneinfo.ZoneInfo(user_timezone) @@ -55,7 +55,7 @@ class CloudZeroStreamer: return # Group data by date and send each day as a batch - daily_batches = self._group_by_date(data) + daily_batches: Final = self._group_by_date(data) if not daily_batches: self.console.print("[yellow]No valid daily batches to send[/yellow]") @@ -68,14 +68,14 @@ class CloudZeroStreamer: def _group_by_date(self, data: pl.DataFrame) -> dict[str, pl.DataFrame]: """Group data by date, converting to UTC and validating dates.""" - daily_batches: dict[str, list[dict[str, Any]]] = {} + daily_batches: Final[dict[str, list[dict[str, Any]]]] = {} # Ensure we have the required columns if "time/usage_start" not in data.columns: self.console.print("[red]Error: Missing 'time/usage_start' column for date grouping[/red]") return {} - timestamp_str: Optional[str] = None + timestamp_str: str | None = None for row in data.iter_rows(named=True): try: # Parse the timestamp and convert to UTC @@ -153,22 +153,22 @@ class CloudZeroStreamer: if batch_data.is_empty(): return - headers = { + headers: Final = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json", } # Use the correct API endpoint format from documentation - url = f"{self.base_url}/v2/connections/billing/anycost/{self.connection_id}/billing_drops" + url: Final = f"{self.base_url}/v2/connections/billing/anycost/{self.connection_id}/billing_drops" # Prepare the batch payload according to AnyCost API format - payload = self._prepare_batch_payload(batch_date, batch_data, operation) + payload: Final = self._prepare_batch_payload(batch_date, batch_data, operation) try: with httpx.Client(timeout=30.0) as client: self.console.print(f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]") - response = client.post(url, headers=headers, json=payload) + response: Final = client.post(url, headers=headers, json=payload) response.raise_for_status() self.console.print( @@ -188,28 +188,28 @@ class CloudZeroStreamer: """Prepare batch payload according to CloudZero AnyCost API format.""" # Convert batch_date to month for the API (YYYY-MM format) try: - date_obj = datetime.strptime(batch_date, "%Y-%m-%d") + date_obj: Final = datetime.strptime(batch_date, "%Y-%m-%d") month_str = date_obj.strftime("%Y-%m") except ValueError: # Fallback to current month month_str = datetime.now().strftime("%Y-%m") # Convert DataFrame rows to API format - data_records = [] + data_records: Final = [] for row in batch_data.iter_rows(named=True): record = self._convert_cbf_to_api_format(row) if record: data_records.append(record) - payload = {"month": month_str, "operation": operation, "data": data_records} + payload: Final = {"month": month_str, "operation": operation, "data": data_records} return payload - def _convert_cbf_to_api_format(self, row: dict[str, Any]) -> Optional[dict[str, Any]]: + def _convert_cbf_to_api_format(self, row: dict[str, Any]) -> dict[str, Any] | None: """Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them.""" try: # CloudZero expects CBF format field names directly, not converted names - api_record = {} + api_record: Final = {} # Copy all CBF fields, converting numeric values to strings as required by CloudZero for key, value in row.items(): @@ -241,7 +241,7 @@ class CloudZeroStreamer: return datetime.now(timezone.utc).isoformat() try: - dt = self._parse_and_convert_timestamp(timestamp_str) + dt: Final = self._parse_and_convert_timestamp(timestamp_str) return dt.isoformat().replace("+00:00", "Z") except Exception: # Fallback to current time in UTC diff --git a/litellm/integrations/cloudzero/database.py b/litellm/integrations/cloudzero/database.py index 71929398103..b050ee8e1ed 100644 --- a/litellm/integrations/cloudzero/database.py +++ b/litellm/integrations/cloudzero/database.py @@ -19,7 +19,7 @@ """Database connection and data extraction for LiteLLM.""" from datetime import datetime -from typing import Any, Optional, List +from typing import Any, Final import polars as pl @@ -39,12 +39,12 @@ class LiteLLMDatabase: async def get_usage_data( self, - limit: Optional[int] = None, - start_time_utc: Optional[datetime] = None, - end_time_utc: Optional[datetime] = None, + limit: int | None = None, + start_time_utc: datetime | None = None, + end_time_utc: datetime | None = None, ) -> pl.DataFrame: """Retrieve usage data from LiteLLM daily user spend table.""" - client = self._ensure_prisma_client() + client: Final = self._ensure_prisma_client() # Query to get user spend data with team information. Use parameter binding to # avoid SQL injection from user-supplied timestamps or limits. @@ -80,7 +80,7 @@ class LiteLLMDatabase: ORDER BY dus.date DESC, dus.created_at DESC """ - params: List[Any] = [ + params: Final[list[Any]] = [ start_time_utc, end_time_utc, ] @@ -93,9 +93,9 @@ class LiteLLMDatabase: query += " LIMIT $3" try: - db_response = await client.db.query_raw(query, *params) + db_response: Final = await client.db.query_raw(query, *params) # Convert the response to polars DataFrame with full schema inference # This prevents schema mismatch errors when data types vary across rows return pl.DataFrame(db_response, infer_schema_length=None) except Exception as e: - raise Exception(f"Error retrieving usage data: {str(e)}") + raise Exception(f"Error retrieving usage data: {e}") diff --git a/litellm/integrations/cloudzero/transform.py b/litellm/integrations/cloudzero/transform.py index c72001aee1a..f0d4d67fc22 100644 --- a/litellm/integrations/cloudzero/transform.py +++ b/litellm/integrations/cloudzero/transform.py @@ -19,7 +19,7 @@ """Transform LiteLLM data to CloudZero AnyCost CBF format.""" from datetime import datetime -from typing import Any, Optional +from typing import Any, Final import polars as pl @@ -40,7 +40,7 @@ class CBFTransformer: return pl.DataFrame() # Filter out records with zero successful_requests first - original_count = len(data) + original_count: Final = len(data) if "successful_requests" in data.columns: filtered_data = data.filter(pl.col("successful_requests") > 0) zero_requests_dropped = original_count - len(filtered_data) @@ -48,9 +48,9 @@ class CBFTransformer: filtered_data = data zero_requests_dropped = 0 - cbf_data = [] + cbf_data: Final = [] czrn_dropped_count = 0 - filtered_count = len(filtered_data) + filtered_count: Final = len(filtered_data) for row in filtered_data.iter_rows(named=True): try: @@ -65,7 +65,7 @@ class CBFTransformer: # Print summary of dropped records if any from rich.console import Console - console = Console() + console: Final = Console() if zero_requests_dropped > 0: console.print( @@ -86,35 +86,35 @@ class CBFTransformer: """Create a single CBF record from LiteLLM daily spend row.""" # Parse date (daily spend tables use date strings like '2025-04-19') - usage_date = self._parse_date(row.get("date")) + usage_date: Final = self._parse_date(row.get("date")) # Calculate total tokens - prompt_tokens = int(row.get("prompt_tokens", 0)) - completion_tokens = int(row.get("completion_tokens", 0)) - total_tokens = prompt_tokens + completion_tokens + prompt_tokens: Final = int(row.get("prompt_tokens", 0)) + completion_tokens: Final = int(row.get("completion_tokens", 0)) + total_tokens: Final = prompt_tokens + completion_tokens # Create CloudZero Resource Name (CZRN) as resource_id - resource_id = self.czrn_generator.create_from_litellm_data(row) + resource_id: Final = self.czrn_generator.create_from_litellm_data(row) # Build dimensions for CloudZero - model = str(row.get("model", "")) - api_key_hash = str(row.get("api_key", ""))[:8] # First 8 chars for identification + model: Final = str(row.get("model", "")) + api_key_hash: Final = str(row.get("api_key", ""))[:8] # First 8 chars for identification # Handle team information with fallbacks - team_id = row.get("team_id") - team_alias = row.get("team_alias") - user_email = row.get("user_email") + team_id: Final = row.get("team_id") + team_alias: Final = row.get("team_alias") + user_email: Final = row.get("user_email") # Use team_alias if available, otherwise team_id, otherwise fallback to 'unknown' - entity_id = str(team_alias) if team_alias else (str(team_id) if team_id else "unknown") + entity_id: Final = str(team_alias) if team_alias else (str(team_id) if team_id else "unknown") # Get alias fields if they exist - api_key_alias = row.get("api_key_alias") - organization_alias = row.get("organization_alias") - project_alias = row.get("project_alias") - user_alias = row.get("user_alias") + api_key_alias: Final = row.get("api_key_alias") + organization_alias: Final = row.get("organization_alias") + project_alias: Final = row.get("project_alias") + user_alias: Final = row.get("user_alias") - dimensions = { + dimensions: Final = { "entity_type": CZEntityType.TEAM.value, "entity_id": entity_id, "team_alias": str(team_alias) if team_alias else "unknown", @@ -135,7 +135,7 @@ class CBFTransformer: } # Extract CZRN components to populate corresponding CBF columns - czrn_components = self.czrn_generator.extract_components(resource_id) + czrn_components: Final = self.czrn_generator.extract_components(resource_id) ( service_type, provider, @@ -146,10 +146,10 @@ class CBFTransformer: ) = czrn_components # Build resource/account as concat of api_key_alias and api_key_prefix - resource_account = f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash + resource_account: Final = f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash # CloudZero CBF format with proper column names - cbf_record = { + cbf_record: Final = { # Required CBF fields "time/usage_start": ( usage_date.isoformat() if usage_date else None @@ -187,7 +187,7 @@ class CBFTransformer: return CBFRecord(cbf_record) - def _parse_date(self, date_str) -> Optional[datetime]: + def _parse_date(self, date_str) -> datetime | None: """Parse date string from daily spend tables (e.g., '2025-04-19').""" if date_str is None: return None diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index 759b2be3a84..f142f1b88a7 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -9,21 +9,21 @@ captured stdout back through the typed agentic loop plan. import json import time import uuid -from typing import Any, Literal, TypedDict, cast +from typing import Any, Final, Literal, TypedDict, cast -import litellm from pydantic import ValidationError +import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.types.integrations.code_interpreter_interception import ( CodeInterpreterInterceptionConfig, ) from litellm.types.integrations.custom_logger import ( - AgenticLoopPlan, - AgenticLoopRequestPatch, CHAT_COMPLETION_AGENTIC_SURFACE, NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + AgenticLoopPlan, + AgenticLoopRequestPatch, is_interception_internal_key, ) from litellm.types.llms.openai import ( @@ -37,14 +37,14 @@ from litellm.types.utils import ( ModelResponse, ) -LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution" -_INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active" -_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" -_SESSION_SCOPED_KEY = "_code_interpreter_interception_session_scoped" -_CONVERTED_STREAM_KEY = "_code_interpreter_interception_converted_stream" -_LITELLM_METADATA_KEY = "litellm_metadata" -_CACHE_TTL_SECONDS = 15 * 60 -_SESSION_SCOPED_PER_IDENTITY_CAP = 10 +LITELLM_CODE_EXECUTION_TOOL_NAME: Final = "litellm_code_execution" +_INTERCEPTION_ACTIVE_KEY: Final = "_code_interpreter_interception_active" +_SANDBOX_KEY: Final = "_code_interpreter_interception_sandbox_key" +_SESSION_SCOPED_KEY: Final = "_code_interpreter_interception_session_scoped" +_CONVERTED_STREAM_KEY: Final = "_code_interpreter_interception_converted_stream" +_LITELLM_METADATA_KEY: Final = "litellm_metadata" +_CACHE_TTL_SECONDS: Final = 15 * 60 +_SESSION_SCOPED_PER_IDENTITY_CAP: Final = 10 class CodeExecutionToolCall(TypedDict, total=False): @@ -200,16 +200,16 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if self.enabled_providers is not None and self._resolve_provider(kwargs) not in self.enabled_providers: return None - tools = kwargs.get("tools") + tools: Final = kwargs.get("tools") if not isinstance(tools, list): return None if not any(isinstance(tool, dict) and tool.get("type") == "code_interpreter" for tool in tools): return None kwargs[_INTERCEPTION_ACTIVE_KEY] = True - session_id = _extract_session_id(kwargs) + session_id: Final = _extract_session_id(kwargs) if session_id: - identity = _extract_identity(kwargs) + identity: Final = _extract_identity(kwargs) kwargs[_SANDBOX_KEY] = f"{identity}:{session_id}" if identity else session_id kwargs[_SESSION_SCOPED_KEY] = True else: @@ -219,7 +219,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): kwargs[_CONVERTED_STREAM_KEY] = True self._write_interception_metadata(kwargs) - function_tool = self._get_function_tool(call_type=call_type) + function_tool: Final = self._get_function_tool(call_type=call_type) kwargs["tools"] = [ (function_tool if isinstance(tool, dict) and tool.get("type") == "code_interpreter" else tool) for tool in tools @@ -230,10 +230,10 @@ class CodeInterpreterInterceptionLogger(CustomLogger): @staticmethod def _strip_interception_metadata(kwargs: dict[str, Any]) -> None: - metadata = kwargs.get(_LITELLM_METADATA_KEY) + metadata: Final = kwargs.get(_LITELLM_METADATA_KEY) if not isinstance(metadata, dict): return - filtered_metadata = { + filtered_metadata: Final = { key: value for key, value in metadata.items() if not is_interception_internal_key(key) @@ -264,7 +264,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): } def _get_function_tool(self, call_type: CallTypes | None) -> CodeExecutionFunctionTool: - description = "Execute python code in a sandbox and return stdout." + description: Final = "Execute python code in a sandbox and return stdout." if call_type in (CallTypes.completion, CallTypes.acompletion): return { "type": "function", @@ -299,7 +299,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool: if not isinstance(tool_choice, dict): return False - function = tool_choice.get("function") + function: Final = tool_choice.get("function") return ( tool_choice.get("type") == "code_interpreter" or tool_choice.get("name") == "code_interpreter" @@ -308,10 +308,10 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None: - provider = kwargs.get("custom_llm_provider") + provider: Final = kwargs.get("custom_llm_provider") if provider: return provider - model = kwargs.get("model") + model: Final = kwargs.get("model") if not isinstance(model, str): return None try: @@ -336,7 +336,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if self.enabled_providers is not None and custom_llm_provider not in self.enabled_providers: return False, {} - tool_calls = ( + tool_calls: Final = ( self._extract_chat_completion_code_execution_tool_calls(response=response) if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE else self._extract_code_execution_tool_calls(response=response) @@ -368,16 +368,16 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) await self._prune_expired_cache() - tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) - sandbox_key = kwargs.get(_SANDBOX_KEY) - is_session = bool(kwargs.get(_SESSION_SCOPED_KEY)) - identity = _extract_identity(kwargs) if is_session else None + tool_calls: Final = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) + sandbox_key: Final = kwargs.get(_SANDBOX_KEY) + is_session: Final = bool(kwargs.get(_SESSION_SCOPED_KEY)) + identity: Final = _extract_identity(kwargs) if is_session else None container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity) try: - container_id = cast(str | None, getattr(container, "id", None)) - input_list = self._normalize_messages(messages) - code_interpreter_calls: list[CodeInterpreterCall] = [] + container_id: Final = cast(str | None, getattr(container, "id", None)) + input_list: Final = self._normalize_messages(messages) + code_interpreter_calls: Final[list[CodeInterpreterCall]] = [] for tool_call in tool_calls: arguments = tool_call.get("arguments", "") code = self._parse_code(arguments) @@ -411,8 +411,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): await self._delete_container_for_cache_key(sandbox_key) raise - optional_params = anthropic_messages_optional_request_params - request_patch = AgenticLoopRequestPatch( + optional_params: Final = anthropic_messages_optional_request_params + request_patch: Final = AgenticLoopRequestPatch( model=model, messages=input_list, tools=self._get_followup_tools( @@ -443,15 +443,15 @@ class CodeInterpreterInterceptionLogger(CustomLogger): kwargs: dict[str, object], ) -> AgenticLoopPlan: await self._prune_expired_cache() - tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) - sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY)) - is_session = bool(kwargs.get(_SESSION_SCOPED_KEY)) - identity = _extract_identity(cast(dict[str, Any], kwargs)) if is_session else None + tool_calls: Final = cast(list[CodeExecutionToolCall], tools.get("tool_calls", [])) + sandbox_key: Final = cast(str | None, kwargs.get(_SANDBOX_KEY)) + is_session: Final = bool(kwargs.get(_SESSION_SCOPED_KEY)) + identity: Final = _extract_identity(cast(dict[str, Any], kwargs)) if is_session else None container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity) try: - container_id = cast(str | None, getattr(container, "id", None)) - tool_results = [ + container_id: Final = cast(str | None, getattr(container, "id", None)) + tool_results: Final = [ await self._build_chat_completion_tool_result( container=container, params=params, @@ -463,10 +463,10 @@ class CodeInterpreterInterceptionLogger(CustomLogger): except Exception: await self._delete_container_for_cache_key(sandbox_key) raise - tool_messages = [result[0] for result in tool_results] - code_interpreter_calls = [result[1] for result in tool_results] + tool_messages: Final = [result[0] for result in tool_results] + code_interpreter_calls: Final = [result[1] for result in tool_results] - request_patch = AgenticLoopRequestPatch( + request_patch: Final = AgenticLoopRequestPatch( model=model, messages=list(messages) + [self._build_chat_completion_assistant_message(tool_calls)] + tool_messages, tools=self._get_followup_tools( @@ -496,10 +496,10 @@ class CodeInterpreterInterceptionLogger(CustomLogger): tool_call: CodeExecutionToolCall, container_id: str | None, ) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]: - arguments = tool_call.get("arguments", "") - code = self._parse_code(arguments) - stdout = await self._run_tool_call(container=container, params=params, arguments=arguments) - tool_call_id = tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex + arguments: Final = tool_call.get("arguments", "") + code: Final = self._parse_code(arguments) + stdout: Final = await self._run_tool_call(container=container, params=params, arguments=arguments) + tool_call_id: Final = tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex return ( { "role": "tool", @@ -517,7 +517,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) async def async_agentic_loop_cleanup_hook(self, plan: AgenticLoopPlan, kwargs: dict) -> None: - metadata = plan.metadata or {} if plan else {} + metadata: Final = plan.metadata or {} if plan else {} if metadata.get("is_session_scoped"): return await self._delete_container_for_cache_key(metadata.get("sandbox_key")) @@ -544,33 +544,33 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ] def _get_followup_optional_params(self, optional_params: dict[str, object]) -> dict[str, object]: - drop_tool_choice = self._tool_choice_targets_code_interpreter(optional_params.get("tool_choice")) + drop_tool_choice: Final = self._tool_choice_targets_code_interpreter(optional_params.get("tool_choice")) return { k: v for k, v in optional_params.items() if k != "tools" and not (k == "tool_choice" and drop_tool_choice) } async def async_post_agentic_loop_response_hook(self, response: Any, plan: AgenticLoopPlan, kwargs: dict) -> Any: - metadata = plan.metadata or {} if plan else {} + metadata: Final = plan.metadata or {} if plan else {} if not metadata.get("is_session_scoped"): await self._delete_container_for_cache_key(metadata.get("sandbox_key")) - calls = metadata.get("code_interpreter_calls") + calls: Final = metadata.get("code_interpreter_calls") if not calls: return response - is_dict = isinstance(response, dict) - output = response.get("output") if is_dict else getattr(response, "output", None) + is_dict: Final = isinstance(response, dict) + output: Final = response.get("output") if is_dict else getattr(response, "output", None) if not isinstance(output, list): return response def _item_type(item: Any) -> Any: return item.get("type") if isinstance(item, dict) else getattr(item, "type", None) - insert_at = next( + insert_at: Final = next( (i for i, item in enumerate(output) if _item_type(item) == "message"), len(output), ) - new_output = output[:insert_at] + list(calls) + output[insert_at:] + new_output: Final = output[:insert_at] + list(calls) + output[insert_at:] if is_dict: response["output"] = new_output else: @@ -586,14 +586,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger): async def _run_tool_call(self, container: Any, params: dict[str, Any] | None, arguments: str) -> str: try: - code = json.loads(arguments).get("code", "") if arguments else "" + code: Final = json.loads(arguments).get("code", "") if arguments else "" except (json.JSONDecodeError, TypeError): return "[invalid tool arguments: could not parse code]" - result = await self._run_code(container=container, params=params, code=code) + result: Final = await self._run_code(container=container, params=params, code=code) if getattr(result, "error", None): - error = result.error - message = error.get("value") or error.get("name") if isinstance(error, dict) else str(error) + error: Final = result.error + message: Final = error.get("value") or error.get("name") if isinstance(error, dict) else str(error) return f"[execution error] {message}" return getattr(result, "stdout", "") or "" @@ -603,7 +603,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): identity: str | None = None, ) -> tuple[Any, dict[str, Any] | None]: if cache_key: - cached = self._container_cache.get(cache_key) + cached: Final = self._container_cache.get(cache_key) if cached is not None: self._container_cache[cache_key] = (cached[0], cached[1], time.time(), cached[3]) return cached[0], cached[1] @@ -616,7 +616,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): return container, params async def _evict_lru_session_if_over_cap(self, identity: str) -> None: - identity_entries = [(k, v) for k, v in self._container_cache.items() if v[3] == identity] + identity_entries: Final = [(k, v) for k, v in self._container_cache.items() if v[3] == identity] if len(identity_entries) < _SESSION_SCOPED_PER_IDENTITY_CAP: return lru_key, lru_entry = min(identity_entries, key=lambda item: item[1][2]) @@ -627,14 +627,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger): if self.sandbox_config is not None: return await self.sandbox_config.acreate_sandbox(), None - params = _resolve_sandbox_tool(self.sandbox_tool_name) + params: Final = _resolve_sandbox_tool(self.sandbox_tool_name) if params is None: raise ValueError( "CodeInterpreterInterception: no sandbox available. Provide a " "sandbox_config or configure a sandbox tool resolvable via " "sandbox_tool_name." ) - container = await litellm.acreate_sandbox( + container: Final = await litellm.acreate_sandbox( provider=params["sandbox_provider"], api_key=params.get("api_key"), api_base=params.get("api_base"), @@ -672,7 +672,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): async def _delete_container_for_cache_key(self, cache_key: str | None) -> None: if not cache_key: return - cached = self._container_cache.pop(cache_key, None) + cached: Final = self._container_cache.pop(cache_key, None) if cached is None: return await self._delete_container(container=cached[0], params=cached[1]) @@ -705,14 +705,14 @@ class CodeInterpreterInterceptionLogger(CustomLogger): def _extract_chat_completion_code_execution_tool_calls( self, response: ModelResponse | dict[str, Any] ) -> list[CodeExecutionToolCall]: - model_response = self._to_model_response(response) + model_response: Final = self._to_model_response(response) if model_response is None: return [] - choices = model_response.choices or [] + choices: Final = model_response.choices or [] if not choices: return [] - message = choices[0].message - tool_calls = message.tool_calls or [] + message: Final = choices[0].message + tool_calls: Final = message.tool_calls or [] return [ normalized @@ -783,8 +783,8 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) async def _prune_expired_cache(self) -> None: - now = time.time() - expired = [ + now: Final = time.time() + expired: Final = [ (cache_key, container, params) for cache_key, (container, params, last_accessed, *_) in self._container_cache.items() if now - last_accessed > _CACHE_TTL_SECONDS diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index f0a696aa1e1..ff2b1197c5f 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid -from typing import Any, Dict, List, Optional, Tuple, cast +from typing import Any, Final, cast from litellm._logging import verbose_logger from litellm.compression import compress @@ -22,8 +22,8 @@ from litellm.types.integrations.custom_logger import ( ) from litellm.types.utils import CallTypes -LITELLM_CONTENT_RETRIEVE_TOOL_NAME = "litellm_content_retrieve" -_CACHE_TTL_SECONDS = 15 * 60 +LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve" +_CACHE_TTL_SECONDS: Final = 15 * 60 def _compression_savings_from_counts( @@ -54,7 +54,7 @@ def _record_compression_savings(kwargs: dict[str, object], savings: CompressionS to the same object; replacing it would orphan writes made through those references. """ - existing = kwargs.get("litellm_metadata") + existing: Final = kwargs.get("litellm_metadata") if isinstance(existing, dict): existing["compression_savings"] = savings return @@ -76,9 +76,9 @@ class CompressionInterceptionLogger(CustomLogger): self, enabled: bool = True, compression_trigger: int = 200_000, - compression_target: Optional[int] = None, - embedding_model: Optional[str] = None, - embedding_model_params: Optional[Dict[str, Any]] = None, + compression_target: int | None = None, + embedding_model: str | None = None, + embedding_model_params: dict[str, Any] | None = None, ): super().__init__() self.enabled = enabled @@ -86,7 +86,7 @@ class CompressionInterceptionLogger(CustomLogger): self.compression_target = compression_target self.embedding_model = embedding_model self.embedding_model_params = embedding_model_params - self._compression_cache_by_call_id: Dict[str, Tuple[Dict[str, str], float]] = {} + self._compression_cache_by_call_id: dict[str, tuple[dict[str, str], float]] = {} @classmethod def from_config_yaml(cls, config: CompressionInterceptionConfig) -> "CompressionInterceptionLogger": @@ -100,8 +100,8 @@ class CompressionInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( - litellm_settings: Dict[str, Any], - callback_specific_params: Dict[str, Any], + litellm_settings: dict[str, Any], + callback_specific_params: dict[str, Any], ) -> "CompressionInterceptionLogger": compression_params: CompressionInterceptionConfig = {} if "compression_interception_params" in litellm_settings: @@ -115,9 +115,7 @@ class CompressionInterceptionLogger(CustomLogger): ) return CompressionInterceptionLogger.from_config_yaml(compression_params) - async def async_pre_call_deployment_hook( - self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] - ) -> Optional[dict]: + async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: if not self.enabled: return None if call_type is not None and call_type != CallTypes.anthropic_messages: @@ -125,8 +123,8 @@ class CompressionInterceptionLogger(CustomLogger): if int(kwargs.get("_agentic_loop_depth", 0) or 0) > 0: return None - messages = kwargs.get("messages") - model = kwargs.get("model") + messages: Final = kwargs.get("messages") + model: Final = kwargs.get("model") if not isinstance(messages, list) or not isinstance(model, str): return None @@ -135,7 +133,7 @@ class CompressionInterceptionLogger(CustomLogger): self._prune_expired_cache() - compressed = compress( # type: ignore + compressed: Final = compress( # type: ignore messages=messages, model=model, call_type=CallTypes.anthropic_messages, @@ -145,9 +143,9 @@ class CompressionInterceptionLogger(CustomLogger): embedding_model_params=self.embedding_model_params, ) - cache = cast(Dict[str, str], compressed.get("cache", {})) - skip_reason = cast(Optional[str], compressed.get("compression_skipped_reason")) - compressed_tools = cast(List[Dict[str, Any]], compressed.get("tools", [])) + cache: Final = cast(dict[str, str], compressed.get("cache", {})) + skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason")) + compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", [])) # Only mutate kwargs when compression actually produced a result. # If compression was a no-op (below trigger, invalid tool sequence, etc.), @@ -158,15 +156,15 @@ class CompressionInterceptionLogger(CustomLogger): kwargs["messages"] = compressed["messages"] if compressed_tools: kwargs["tools"] = self._merge_tools( - existing_tools=cast(Optional[List[Dict[str, Any]]], kwargs.get("tools")), + existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")), compressed_tools=compressed_tools, ) - call_id = cast(Optional[str], kwargs.get("litellm_call_id")) + call_id = cast(str | None, kwargs.get("litellm_call_id")) if not call_id: call_id = str(uuid.uuid4()) kwargs["litellm_call_id"] = call_id self._compression_cache_by_call_id[call_id] = (cache, time.time()) - savings = _compression_savings_from_counts( + savings: Final = _compression_savings_from_counts( original_tokens=compressed.get("original_tokens"), compressed_tokens=compressed.get("compressed_tokens"), ) @@ -193,12 +191,12 @@ class CompressionInterceptionLogger(CustomLogger): self, response: Any, model: str, - messages: List[Dict], - tools: Optional[List[Dict]], + messages: list[dict], + tools: list[dict] | None, stream: bool, custom_llm_provider: str, - kwargs: Dict, - ) -> Tuple[bool, Dict]: + kwargs: dict, + ) -> tuple[bool, dict]: if not self.enabled: return False, {} if not self._has_retrieval_tool(tools): @@ -216,25 +214,25 @@ class CompressionInterceptionLogger(CustomLogger): async def async_build_agentic_loop_plan( self, - tools: Dict, + tools: dict, model: str, - messages: List[Dict], + messages: list[dict], response: Any, anthropic_messages_provider_config: Any, - anthropic_messages_optional_request_params: Dict, + anthropic_messages_optional_request_params: dict, logging_obj: Any, stream: bool, - kwargs: Dict, + kwargs: dict, ) -> AgenticLoopPlan: self._prune_expired_cache() - tool_calls = cast(List[Dict[str, Any]], tools.get("tool_calls", [])) - thinking_blocks = cast(List[Dict[str, Any]], tools.get("thinking_blocks", [])) + tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", [])) + thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", [])) - call_id = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) - cache = self._get_cache(call_id=call_id) - retrieval_results = [self._resolve_retrieval_content(tc, cache) for tc in tool_calls] + call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs) + cache: Final = self._get_cache(call_id=call_id) + retrieval_results: Final = [self._resolve_retrieval_content(tc, cache) for tc in tool_calls] - assistant_message = { + assistant_message: Final = { "role": "assistant", "content": thinking_blocks + [ @@ -247,7 +245,7 @@ class CompressionInterceptionLogger(CustomLogger): for tc in tool_calls ], } - user_message = { + user_message: Final = { "role": "user", "content": [ { @@ -258,22 +256,22 @@ class CompressionInterceptionLogger(CustomLogger): for i in range(len(tool_calls)) ], } - follow_up_messages = messages + [assistant_message, user_message] + follow_up_messages: Final = messages + [assistant_message, user_message] - max_tokens = cast( - Optional[int], + max_tokens: Final = cast( + int | None, anthropic_messages_optional_request_params.get("max_tokens") or kwargs.get("max_tokens"), ) - optional_params_without_max_tokens = { + optional_params_without_max_tokens: Final = { k: v for k, v in anthropic_messages_optional_request_params.items() if k != "max_tokens" } full_model_name = model if logging_obj is not None: - agentic_params = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = cast(str, agentic_params.get("model", model)) - request_patch = AgenticLoopRequestPatch( + request_patch: Final = AgenticLoopRequestPatch( model=full_model_name, messages=follow_up_messages, max_tokens=max_tokens, @@ -288,7 +286,7 @@ class CompressionInterceptionLogger(CustomLogger): ) def _prune_expired_cache(self) -> None: - now = time.time() + now: Final = time.time() self._compression_cache_by_call_id = { call_id: (cache, created_at) for call_id, ( @@ -298,24 +296,24 @@ class CompressionInterceptionLogger(CustomLogger): if now - created_at <= _CACHE_TTL_SECONDS } - def _get_cache(self, call_id: Optional[str]) -> Dict[str, str]: + def _get_cache(self, call_id: str | None) -> dict[str, str]: if not call_id: return {} - cache_entry = self._compression_cache_by_call_id.get(call_id) + cache_entry: Final = self._compression_cache_by_call_id.get(call_id) if cache_entry is None: return {} return cache_entry[0] - def _resolve_call_id(self, logging_obj: Any, kwargs: Dict[str, Any]) -> Optional[str]: + def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None: if logging_obj is not None: - logging_call_id = getattr(logging_obj, "litellm_call_id", None) + logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) if isinstance(logging_call_id, str) and logging_call_id: return logging_call_id - kwargs_call_id = kwargs.get("litellm_call_id") - return cast(Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None) + kwargs_call_id: Final = kwargs.get("litellm_call_id") + return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None) - def _resolve_retrieval_content(self, tool_call: Dict[str, Any], cache: Dict[str, str]) -> str: - raw_input = tool_call.get("input", {}) + def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str: + raw_input: Final = tool_call.get("input", {}) key = "" if isinstance(raw_input, dict): key = str(raw_input.get("key", "") or "") @@ -325,7 +323,7 @@ class CompressionInterceptionLogger(CustomLogger): return cache[key] return f"[compressed content key '{key}' not found]" - def _extract_retrieval_tool_calls(self, response: Any) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: if isinstance(response, dict): content = response.get("content", []) else: @@ -334,8 +332,8 @@ class CompressionInterceptionLogger(CustomLogger): if not isinstance(content, list): return [], [] - tool_calls: List[Dict[str, Any]] = [] - thinking_blocks: List[Dict[str, Any]] = [] + tool_calls: Final[list[dict[str, Any]]] = [] + thinking_blocks: Final[list[dict[str, Any]]] = [] for block in content: if isinstance(block, dict): @@ -382,8 +380,8 @@ class CompressionInterceptionLogger(CustomLogger): return tool_calls, thinking_blocks - def _prepare_followup_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: - internal_keys = {"litellm_logging_obj"} + def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + internal_keys: Final = {"litellm_logging_obj"} return { k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys } @@ -404,10 +402,10 @@ class CompressionInterceptionLogger(CustomLogger): def _merge_tools( self, - existing_tools: Optional[List[Dict[str, Any]]], - compressed_tools: List[Dict[str, Any]], - ) -> List[Dict[str, Any]]: - merged = list(existing_tools or []) + existing_tools: list[dict[str, Any]] | None, + compressed_tools: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + merged: Final = list(existing_tools or []) if self._has_retrieval_tool(merged): return merged merged.extend(compressed_tools) diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index aded12fa399..c9e24913900 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -6,7 +6,7 @@ Use this if you want your logs to be stored in memory and flushed periodically. import asyncio import time -from typing import List, Optional +from typing import Final import litellm from litellm._logging import verbose_logger @@ -25,10 +25,10 @@ class CustomBatchLogger(CustomLogger): def __init__( self, - flush_lock: Optional[asyncio.Lock] = None, - batch_size: Optional[int] = None, - flush_interval: Optional[int] = None, - max_queue_size: Optional[int] = None, + flush_lock: asyncio.Lock | None = None, + batch_size: int | None = None, + flush_interval: int | None = None, + max_queue_size: int | None = None, **kwargs, ) -> None: """ @@ -36,7 +36,7 @@ class CustomBatchLogger(CustomLogger): flush_lock (Optional[asyncio.Lock], optional): Lock to use when flushing the queue. Defaults to None. Only used for custom loggers that do batching max_queue_size (Optional[int], optional): Maximum number of events to retain in ``log_queue``. When the limit is exceeded (e.g. because the send destination is unreachable and events are preserved for retry), the oldest events are dropped. Defaults to ``DEFAULT_MAX_QUEUE_SIZE``. """ - self.log_queue: List = [] + self.log_queue: list = [] self.flush_interval = flush_interval or litellm.DEFAULT_FLUSH_INTERVAL_SECONDS self.batch_size: int = batch_size or litellm.DEFAULT_BATCH_SIZE self.last_flush_time = time.time() @@ -48,7 +48,7 @@ class CustomBatchLogger(CustomLogger): async def periodic_flush(self): while True: await asyncio.sleep(self.flush_interval) - verbose_logger.debug(f"CustomLogger periodic flush after {self.flush_interval} seconds") + verbose_logger.debug("CustomLogger periodic flush after %s seconds", self.flush_interval) await self.flush_queue() async def flush_queue(self): @@ -57,7 +57,7 @@ class CustomBatchLogger(CustomLogger): async with self.flush_lock: if self.log_queue: - log_queue_length = len(self.log_queue) + log_queue_length: Final = len(self.log_queue) verbose_logger.debug("CustomLogger: Flushing batch of %s events", len(self.log_queue)) try: await self.async_send_batch() @@ -74,7 +74,7 @@ class CustomBatchLogger(CustomLogger): # Guard against unbounded queue growth if the destination # is persistently unreachable. Drop the oldest events # beyond ``max_queue_size``. - overflow = len(self.log_queue) - self.max_queue_size + overflow: Final = len(self.log_queue) - self.max_queue_size if overflow > 0: del self.log_queue[:overflow] verbose_logger.warning( diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9c7bbbd3b4c..a80b3ff5364 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,27 +3,16 @@ import hashlib import os import secrets from datetime import datetime -from typing import ( - TYPE_CHECKING, - Any, - ClassVar, - Dict, - List, - Literal, - Optional, - Type, - Union, - get_args, -) +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args from litellm._logging import verbose_logger +from litellm.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, redact_nested_match_and_regex_keys, ) -from litellm.caching import DualCache -from litellm.integrations.custom_logger import CustomLogger from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import ( DynamicGuardrailParams, @@ -49,7 +38,7 @@ except ImportError: if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -dc = DualCache() +dc: Final = DualCache() from litellm.constants import ( @@ -67,11 +56,11 @@ from litellm.exceptions import ( # honors markers carrying this token, so a caller cannot forge the metadata # field to suppress a guardrail on the direct-SDK path that never reaches the # proxy's metadata sanitizer. -_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16) +_PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) -_GUARDRAIL_BLOCK_STATUS_CODES = frozenset({400, 403, 422}) +_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) -_guardrail_self_recorded: contextvars.ContextVar[bool] = contextvars.ContextVar( +_guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( "litellm_guardrail_self_recorded", default=False ) @@ -83,25 +72,25 @@ def _strict_guardrail_modes_enabled() -> bool: for guardrails whose supported_event_hooks list newly includes their configured mode: log the mismatch and continue instead of raising at boot. """ - raw = os.environ.get("LITELLM_STRICT_GUARDRAIL_MODES") + raw: Final = os.environ.get("LITELLM_STRICT_GUARDRAIL_MODES") if raw is None: return True - parsed = str_to_bool(raw) + parsed: Final = str_to_bool(raw) return True if parsed is None else parsed -def get_session_id_from_request_data(request_data: Dict[str, Any]) -> Optional[str]: +def get_session_id_from_request_data(request_data: dict[str, Any]) -> str | None: """Extract session_id from request data (litellm_session_id or metadata).""" session_id = request_data.get("litellm_session_id") if session_id: return str(session_id) - metadata = request_data.get("metadata") or {} + metadata: Final = request_data.get("metadata") or {} session_id = metadata.get("session_id") if session_id: return str(session_id) - litellm_metadata = request_data.get("litellm_metadata") or {} + litellm_metadata: Final = request_data.get("litellm_metadata") or {} session_id = litellm_metadata.get("session_id") if session_id: return str(session_id) @@ -117,18 +106,18 @@ class CustomGuardrail(CustomLogger): def __init__( self, - guardrail_name: Optional[str] = None, - supported_event_hooks: Optional[List[GuardrailEventHooks]] = None, - event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = None, + guardrail_name: str | None = None, + supported_event_hooks: list[GuardrailEventHooks] | None = None, + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = None, default_on: bool = False, mask_request_content: bool = False, mask_response_content: bool = False, - violation_message_template: Optional[str] = None, - end_session_after_n_fails: Optional[int] = None, - on_violation: Optional[str] = None, - realtime_violation_message: Optional[str] = None, - on_sensitive_data: Optional[str] = None, - sensitive_data_route_to_model: Optional[str] = None, + violation_message_template: str | None = None, + end_session_after_n_fails: int | None = None, + on_violation: str | None = None, + realtime_violation_message: str | None = None, + on_sensitive_data: str | None = None, + sensitive_data_route_to_model: str | None = None, sticky_session_routing: bool = True, run_in_parallel: bool = False, only_scan_new_messages: bool = False, @@ -156,16 +145,16 @@ class CustomGuardrail(CustomLogger): """ self.guardrail_name = guardrail_name self.supported_event_hooks = supported_event_hooks - self.event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = event_hook + self.event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None = event_hook self.default_on: bool = default_on self.mask_request_content: bool = mask_request_content self.mask_response_content: bool = mask_response_content - self.violation_message_template: Optional[str] = violation_message_template - self.end_session_after_n_fails: Optional[int] = end_session_after_n_fails - self.on_violation: Optional[str] = on_violation - self.realtime_violation_message: Optional[str] = realtime_violation_message - self.on_sensitive_data: Optional[str] = on_sensitive_data - self.sensitive_data_route_to_model: Optional[str] = sensitive_data_route_to_model + self.violation_message_template: str | None = violation_message_template + self.end_session_after_n_fails: int | None = end_session_after_n_fails + self.on_violation: str | None = on_violation + self.realtime_violation_message: str | None = realtime_violation_message + self.on_sensitive_data: str | None = on_sensitive_data + self.sensitive_data_route_to_model: str | None = sensitive_data_route_to_model self.sticky_session_routing: bool = sticky_session_routing self.run_in_parallel: bool = run_in_parallel self.only_scan_new_messages: bool = only_scan_new_messages @@ -185,13 +174,13 @@ class CustomGuardrail(CustomLogger): ) super().__init__(**kwargs) - def render_violation_message(self, default: str, context: Optional[Dict[str, Any]] = None) -> str: + def render_violation_message(self, default: str, context: dict[str, Any] | None = None) -> str: """Return a custom violation message if template is configured.""" if not self.violation_message_template: return default - format_context: Dict[str, Any] = {"default_message": default} + format_context: Final[dict[str, Any]] = {"default_message": default} if context: format_context.update(context) try: @@ -207,8 +196,8 @@ class CustomGuardrail(CustomLogger): def raise_passthrough_exception( self, violation_message: str, - request_data: Dict[str, Any], - detection_info: Optional[Dict[str, Any]] = None, + request_data: dict[str, Any], + detection_info: dict[str, Any] | None = None, ) -> None: """ Raise a passthrough exception for guardrail violations. @@ -238,7 +227,7 @@ class CustomGuardrail(CustomLogger): detection_info=detection_info ) """ - model = request_data.get("model", "unknown") + model: Final = request_data.get("model", "unknown") raise ModifyResponseException( message=violation_message, @@ -251,8 +240,8 @@ class CustomGuardrail(CustomLogger): def raise_sensitive_data_route_exception( self, route_to_model: str, - request_data: Dict[str, Any], - detection_info: Optional[Dict[str, Any]] = None, + request_data: dict[str, Any], + detection_info: dict[str, Any] | None = None, ) -> None: """ Raise an exception to reroute the request to a different model. @@ -274,7 +263,7 @@ class CustomGuardrail(CustomLogger): Raises: SensitiveDataRouteException: Always raises to trigger rerouting """ - session_id = self._get_session_id_from_request_data(request_data) + session_id: Final = self._get_session_id_from_request_data(request_data) if not session_id: raise ValueError( "Cannot route sensitive data without a session_id. " @@ -289,7 +278,7 @@ class CustomGuardrail(CustomLogger): sticky_session_routing=self.sticky_session_routing, ) - def _get_session_id_from_request_data(self, request_data: Dict[str, Any]) -> Optional[str]: + def _get_session_id_from_request_data(self, request_data: dict[str, Any]) -> str | None: """Extract session_id from request data.""" return get_session_id_from_request_data(request_data) @@ -329,7 +318,7 @@ class CustomGuardrail(CustomLogger): ) return None - session_id = get_session_id_from_request_data(request_data) + session_id: Final = get_session_id_from_request_data(request_data) if not session_id: verbose_logger.debug( "Guardrail %s: only_scan_new_messages enabled but request has no session id; scanning full context.", @@ -338,7 +327,7 @@ class CustomGuardrail(CustomLogger): return None try: - cached: object = await cache.async_get_cache(key=self._scanned_texts_cache_key(session_id)) + cached: Final[object] = await cache.async_get_cache(key=self._scanned_texts_cache_key(session_id)) except Exception as e: # noqa: BLE001 # cache is best-effort; any failure must fall back to a full scan verbose_logger.warning( "Guardrail %s: failed to read scanned-message cache (%s); scanning full context.", @@ -347,7 +336,7 @@ class CustomGuardrail(CustomLogger): ) return None - seen: set[str] = {str(h) for h in cached} if isinstance(cached, list) else set() + seen: Final[set[str]] = {str(h) for h in cached} if isinstance(cached, list) else set() return [text for text in texts if self._scanned_text_hash(text) not in seen] async def mark_texts_scanned( @@ -365,16 +354,16 @@ class CustomGuardrail(CustomLogger): return if self.mask_request_content or self.mask_response_content: return - session_id = get_session_id_from_request_data(request_data) + session_id: Final = get_session_id_from_request_data(request_data) if not session_id: return - cache_key = self._scanned_texts_cache_key(session_id) - current_hashes = [self._scanned_text_hash(text) for text in texts] + cache_key: Final = self._scanned_texts_cache_key(session_id) + current_hashes: Final = [self._scanned_text_hash(text) for text in texts] try: - existing: object = await cache.async_get_cache(key=cache_key) - existing_hashes: list[str] = [str(h) for h in existing] if isinstance(existing, list) else [] - merged: list[str] = list(dict.fromkeys(existing_hashes + current_hashes)) + existing: Final[object] = await cache.async_get_cache(key=cache_key) + existing_hashes: Final[list[str]] = [str(h) for h in existing] if isinstance(existing, list) else [] + merged: Final[list[str]] = list(dict.fromkeys(existing_hashes + current_hashes)) await cache.async_set_cache( key=cache_key, value=merged, @@ -396,8 +385,8 @@ class CustomGuardrail(CustomLogger): def handle_sensitive_data_detection( self, - request_data: Dict[str, Any], - detection_info: Optional[Dict[str, Any]] = None, + request_data: dict[str, Any], + detection_info: dict[str, Any] | None = None, ) -> None: """ Handle sensitive data detection based on guardrail configuration. @@ -439,7 +428,7 @@ class CustomGuardrail(CustomLogger): ) @staticmethod - def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: + def get_config_model() -> type["GuardrailConfigModel"] | None: """ Returns the config model for the guardrail @@ -448,7 +437,7 @@ class CustomGuardrail(CustomLogger): return None @classmethod - def get_supported_event_hooks(cls) -> Optional[List[GuardrailEventHooks]]: + def get_supported_event_hooks(cls) -> list[GuardrailEventHooks] | None: """ Returns the event hooks this guardrail supports, for the UI to render. @@ -461,12 +450,12 @@ class CustomGuardrail(CustomLogger): def _validate_event_hook( self, - event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]], - supported_event_hooks: List[GuardrailEventHooks], + event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None, + supported_event_hooks: list[GuardrailEventHooks], ) -> None: def _validate_event_hook_list_is_in_supported_event_hooks( - event_hook: Union[List[GuardrailEventHooks], List[str]], - supported_event_hooks: List[GuardrailEventHooks], + event_hook: list[GuardrailEventHooks] | list[str], + supported_event_hooks: list[GuardrailEventHooks], ) -> None: for hook in event_hook: if isinstance(hook, str): @@ -481,7 +470,7 @@ class CustomGuardrail(CustomLogger): if isinstance(event_hook, list): _validate_event_hook_list_is_in_supported_event_hooks(event_hook, supported_event_hooks) elif isinstance(event_hook, Mode): - tag_values_flat: list = [] + tag_values_flat: Final[list] = [] for v in event_hook.tags.values(): if isinstance(v, list): tag_values_flat.extend(v) @@ -517,7 +506,7 @@ class CustomGuardrail(CustomLogger): key_meta = meta.get("user_api_key_metadata") or key_meta return {**team_meta, **key_meta} - def get_disable_global_guardrail(self, data: dict) -> Optional[bool]: + def get_disable_global_guardrail(self, data: dict) -> bool | None: """ Returns True if the global guardrail should be disabled. @@ -526,13 +515,13 @@ class CustomGuardrail(CustomLogger): """ return self._get_admin_metadata(data).get("disable_global_guardrails", False) - def get_opted_out_global_guardrails_from_metadata(self, data: dict) -> List[str]: + def get_opted_out_global_guardrails_from_metadata(self, data: dict) -> list[str]: """ Returns the list of global guardrail names the team/key has opted out of. Reads from admin-configured key/team metadata only. """ - value = self._get_admin_metadata(data).get("opted_out_global_guardrails") + value: Final = self._get_admin_metadata(data).get("opted_out_global_guardrails") return value if isinstance(value, list) else [] def _is_valid_response_type(self, result: Any) -> bool: @@ -548,7 +537,7 @@ class CustomGuardrail(CustomLogger): try: # Try isinstance check on valid types that support it - response_types = get_args(LLMResponseTypes) + response_types: Final = get_args(LLMResponseTypes) return isinstance(result, response_types) except TypeError as e: # TypedDict types don't support isinstance checks @@ -557,7 +546,7 @@ class CustomGuardrail(CustomLogger): return True raise - def get_guardrail_from_metadata(self, data: dict) -> Union[List[str], List[Dict[str, DynamicGuardrailParams]]]: + def get_guardrail_from_metadata(self, data: dict) -> list[str] | list[dict[str, DynamicGuardrailParams]]: """ Returns the guardrail(s) to be run from the metadata or root """ @@ -578,7 +567,7 @@ class CustomGuardrail(CustomLogger): def _guardrail_is_in_requested_guardrails( self, - requested_guardrails: Union[List[str], List[Dict[str, DynamicGuardrailParams]]], + requested_guardrails: list[str] | list[dict[str, DynamicGuardrailParams]], ) -> bool: for _guardrail in requested_guardrails: if isinstance(_guardrail, dict): @@ -590,13 +579,13 @@ class CustomGuardrail(CustomLogger): return False - def _pre_call_marker(self) -> Optional[str]: - name = self.guardrail_name + def _pre_call_marker(self) -> str | None: + name: Final = self.guardrail_name if not name: return None return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}" - def mark_pre_call_hook_ran(self, data: Dict[str, Any]) -> None: + def mark_pre_call_hook_ran(self, data: dict[str, Any]) -> None: """ Record that this guardrail's ``async_pre_call_hook`` already ran for this request, so the deployment-level hook does not run it a second time. @@ -606,7 +595,7 @@ class CustomGuardrail(CustomLogger): top-level request kwargs, which would otherwise re-trigger the same hook from ``async_pre_call_deployment_hook``. """ - marker = self._pre_call_marker() + marker: Final = self._pre_call_marker() if marker is None: return for meta_key in ("metadata", "litellm_metadata"): @@ -621,8 +610,8 @@ class CustomGuardrail(CustomLogger): return data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]} - def _pre_call_hook_already_ran(self, data: Dict[str, Any]) -> bool: - marker = self._pre_call_marker() + def _pre_call_hook_already_ran(self, data: dict[str, Any]) -> bool: + marker: Final = self._pre_call_marker() if marker is None: return False for meta_key in ("metadata", "litellm_metadata"): @@ -649,13 +638,11 @@ class CustomGuardrail(CustomLogger): ) from e return unified_guardrail - async def async_pre_call_deployment_hook( - self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] - ) -> Optional[dict]: + async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: from litellm.proxy._types import UserAPIKeyAuth # should run guardrail - litellm_guardrails = kwargs.get("guardrails") + litellm_guardrails: Final = kwargs.get("guardrails") if litellm_guardrails is None or not isinstance(litellm_guardrails, list): return kwargs @@ -667,10 +654,10 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - target = self._deployment_pre_call_target() + target: Final = self._deployment_pre_call_target() if target is not self: kwargs["guardrail_to_apply"] = self - result = await target.async_pre_call_hook( + result: Final = await target.async_pre_call_hook( user_api_key_dict=UserAPIKeyAuth( user_id=kwargs.get("user_api_key_user_id"), team_id=kwargs.get("user_api_key_team_id"), @@ -684,7 +671,7 @@ class CustomGuardrail(CustomLogger): ) if result is not None and isinstance(result, dict): - result_messages = result.get("messages") + result_messages: Final = result.get("messages") if result_messages is not None: # update for any pii / masking logic kwargs["messages"] = result_messages @@ -694,15 +681,15 @@ class CustomGuardrail(CustomLogger): self, request_data: dict, response: LLMResponseTypes, - call_type: Optional[CallTypes], - ) -> Optional[LLMResponseTypes]: + call_type: CallTypes | None, + ) -> LLMResponseTypes | None: """ Allow modifying / reviewing the response just after it's received from the deployment. """ from litellm.proxy._types import UserAPIKeyAuth # should run guardrail - litellm_guardrails = request_data.get("guardrails") + litellm_guardrails: Final = request_data.get("guardrails") if litellm_guardrails is None or not isinstance(litellm_guardrails, list): return response @@ -710,7 +697,7 @@ class CustomGuardrail(CustomLogger): return response # CHECK IF GUARDRAIL REJECTS THE REQUEST - result = await self.async_post_call_success_hook( + result: Final = await self.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth( user_id=request_data.get("user_api_key_user_id"), team_id=request_data.get("user_api_key_team_id"), @@ -735,9 +722,9 @@ class CustomGuardrail(CustomLogger): """ Returns True if the guardrail should be run on the event_type """ - requested_guardrails = self.get_guardrail_from_metadata(data) - disable_global_guardrail = self.get_disable_global_guardrail(data) - opted_out_global_guardrails = self.get_opted_out_global_guardrails_from_metadata(data) + requested_guardrails: Final = self.get_guardrail_from_metadata(data) + disable_global_guardrail: Final = self.get_disable_global_guardrail(data) + opted_out_global_guardrails: Final = self.get_opted_out_global_guardrails_from_metadata(data) verbose_logger.debug( "inside should_run_guardrail for guardrail=%s event_type= %s guardrail_supported_event_hooks= %s requested_guardrails= %s self.default_on= %s", self.guardrail_name, @@ -815,7 +802,7 @@ class CustomGuardrail(CustomLogger): elif event_type.value == tag_value: return True if self.event_hook.default: - default_list = ( + default_list: Final = ( self.event_hook.default if isinstance(self.event_hook.default, list) else [self.event_hook.default] ) return event_type.value in default_list @@ -840,7 +827,7 @@ class CustomGuardrail(CustomLogger): Args: request_data: The original `request_data` passed to LiteLLM Proxy """ - requested_guardrails = self.get_guardrail_from_metadata(request_data) + requested_guardrails: Final = self.get_guardrail_from_metadata(request_data) # Look for the guardrail configuration matching self.guardrail_name for guardrail in requested_guardrails: @@ -870,23 +857,23 @@ class CustomGuardrail(CustomLogger): if premium_user is not True: verbose_logger.warning( - f"Trying to use premium guardrail without premium user {CommonProxyErrors.not_premium_user.value}" + "Trying to use premium guardrail without premium user %s", CommonProxyErrors.not_premium_user.value ) return False return True def add_standard_logging_guardrail_information_to_request_data( self, - guardrail_json_response: Union[Exception, str, dict, List[dict]], + guardrail_json_response: Exception | str | dict | list[dict], request_data: dict, guardrail_status: GuardrailStatus, - start_time: Optional[float] = None, - end_time: Optional[float] = None, - duration: Optional[float] = None, - masked_entity_count: Optional[Dict[str, int]] = None, - guardrail_provider: Optional[str] = None, - event_type: Optional[GuardrailEventHooks] = None, - tracing_detail: Optional[GuardrailTracingDetail] = None, + start_time: float | None = None, + end_time: float | None = None, + duration: float | None = None, + masked_entity_count: dict[str, int] | None = None, + guardrail_provider: str | None = None, + event_type: GuardrailEventHooks | None = None, + tracing_detail: GuardrailTracingDetail | None = None, ) -> None: """ Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc. @@ -901,7 +888,7 @@ class CustomGuardrail(CustomLogger): from litellm.types.utils import GuardrailMode # Use event_type if provided, otherwise fall back to self.event_hook - guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]] + guardrail_mode: GuardrailEventHooks | GuardrailMode | list[GuardrailEventHooks] if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): @@ -938,7 +925,7 @@ class CustomGuardrail(CustomLogger): clean_guardrail_response = mask_credentials_in_payload(clean_guardrail_response) - slg = StandardLoggingGuardrailInformation( + slg: Final = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, guardrail_mode=guardrail_mode, @@ -952,8 +939,8 @@ class CustomGuardrail(CustomLogger): ) def _append_guardrail_info(container: dict) -> None: - key = "standard_logging_guardrail_information" - existing = container.get(key) + key: Final = "standard_logging_guardrail_information" + existing: Final = container.get(key) if existing is None: container[key] = [slg] elif isinstance(existing, list): @@ -1009,13 +996,13 @@ class CustomGuardrail(CustomLogger): def _process_response( self, - response: Optional[Dict], + response: dict | None, request_data: dict, - start_time: Optional[float] = None, - end_time: Optional[float] = None, - duration: Optional[float] = None, - event_type: Optional[GuardrailEventHooks] = None, - original_inputs: Optional[Dict] = None, + start_time: float | None = None, + end_time: float | None = None, + duration: float | None = None, + event_type: GuardrailEventHooks | None = None, + original_inputs: dict | None = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -1023,7 +1010,7 @@ class CustomGuardrail(CustomLogger): This gets logged on downsteam Langfuse, DataDog, etc. """ # Convert None to empty dict to satisfy type requirements - guardrail_response: Union[Dict[str, Any], str] = {} if response is None else response + guardrail_response: dict[str, Any] | str = {} if response is None else response # For apply_guardrail functions in custom_code_guardrail scenario, # simplify the logged response to "allow", "deny", or "mask" @@ -1034,7 +1021,7 @@ class CustomGuardrail(CustomLogger): else: guardrail_response = "allow" - verbose_logger.debug(f"Guardrail response: {response}") + verbose_logger.debug("Guardrail response: %s", response) self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, @@ -1090,22 +1077,22 @@ class CustomGuardrail(CustomLogger): self, e: Exception, request_data: dict, - start_time: Optional[float] = None, - end_time: Optional[float] = None, - duration: Optional[float] = None, - event_type: Optional[GuardrailEventHooks] = None, + start_time: float | None = None, + end_time: float | None = None, + duration: float | None = None, + event_type: GuardrailEventHooks | None = None, ): """ Add StandardLoggingGuardrailInformation to the request data This gets logged on downsteam Langfuse, DataDog, etc. """ - guardrail_status: GuardrailStatus = ( + guardrail_status: Final[GuardrailStatus] = ( "guardrail_intervened" if self._is_guardrail_intervention(e) else "guardrail_failed_to_respond" ) # For custom_code_guardrail scenario, log as "deny" instead of full exception # Check if this is from custom_code_guardrail by checking the class name - guardrail_response: Union[Exception, str] = e + guardrail_response: Exception | str = e if "CustomCodeGuardrail" in self.__class__.__name__: guardrail_response = "deny" @@ -1120,14 +1107,14 @@ class CustomGuardrail(CustomLogger): ) raise e - def _inputs_were_modified(self, original_inputs: Dict, response: Dict) -> bool: + def _inputs_were_modified(self, original_inputs: dict, response: dict) -> bool: """ Compare original inputs with response to determine if content was modified. Returns True if the inputs were modified (mask scenario), False otherwise (allow scenario). """ # Get all keys from both dictionaries - all_keys = set(original_inputs.keys()) | set(response.keys()) + all_keys: Final = set(original_inputs.keys()) | set(response.keys()) # Compare each key's value for key in all_keys: @@ -1165,8 +1152,8 @@ class CustomGuardrail(CustomLogger): setattr(self, key, value) def get_guardrails_messages_for_call_type( - self, call_type: CallTypes, data: Optional[dict] = None - ) -> Optional[List[AllMessageValues]]: + self, call_type: CallTypes, data: dict | None = None + ) -> list[AllMessageValues] | None: """ Returns the messages for the given call type and data """ @@ -1182,6 +1169,7 @@ class CustomGuardrail(CustomLogger): call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value or call_type == CallTypes.anthropic_messages.value + or call_type == CallTypes.call_mcp_tool.value ): return data.get("messages") @@ -1196,15 +1184,15 @@ class CustomGuardrail(CustomLogger): LiteLLMCompletionResponsesConfig, ) - input_data = data.get("input") + input_data: Final = data.get("input") if input_data is None: return None - messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=input_data, responses_api_request=data, ) - return cast(List[AllMessageValues], messages) + return cast(list[AllMessageValues], messages) return None @@ -1214,7 +1202,7 @@ def _append_slg_to_litellm_params(lp: object, entries: list) -> None: return if lp.get("metadata") is None: lp["metadata"] = {} - existing = lp["metadata"].setdefault("standard_logging_guardrail_information", []) + existing: Final = lp["metadata"].setdefault("standard_logging_guardrail_information", []) for entry in entries: if entry not in existing: existing.append(entry) @@ -1233,12 +1221,12 @@ def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object) """ if logging_obj is None: return - meta_src = request_data.get(get_metadata_variable_name_from_kwargs(request_data)) or {} - slg_info = meta_src.get("standard_logging_guardrail_information") + meta_src: Final = request_data.get(get_metadata_variable_name_from_kwargs(request_data)) or {} + slg_info: Final = meta_src.get("standard_logging_guardrail_information") if not slg_info: return - entries: list = slg_info if isinstance(slg_info, list) else [slg_info] - mcd = getattr(logging_obj, "model_call_details", None) or {} + entries: Final[list] = slg_info if isinstance(slg_info, list) else [slg_info] + mcd: Final = getattr(logging_obj, "model_call_details", None) or {} _append_slg_to_litellm_params(getattr(logging_obj, "litellm_params", None), entries) _append_slg_to_litellm_params(mcd.get("litellm_params"), entries) @@ -1280,7 +1268,7 @@ def log_guardrail_information(func): def _infer_event_type_from_function_name( func_name: str, - ) -> Optional[GuardrailEventHooks]: + ) -> GuardrailEventHooks | None: """Infer the actual event type from the function name""" if func_name == "async_pre_call_hook": return GuardrailEventHooks.pre_call @@ -1295,20 +1283,20 @@ def log_guardrail_information(func): @functools.wraps(func) async def async_wrapper(*args, **kwargs): - start_time = datetime.now() # Move start_time inside the wrapper - self: CustomGuardrail = args[0] - request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} - event_type = _infer_event_type_from_function_name(func.__name__) + start_time: Final = datetime.now() # Move start_time inside the wrapper + self: Final[CustomGuardrail] = args[0] + request_data: Final[dict] = kwargs.get("data") or kwargs.get("request_data") or {} + event_type: Final = _infer_event_type_from_function_name(func.__name__) # Store original inputs for comparison (for apply_guardrail functions) original_inputs = None if func.__name__ == "apply_guardrail" and "inputs" in kwargs: original_inputs = kwargs.get("inputs") - logging_obj = kwargs.get("logging_obj") - self_recorded_token = _guardrail_self_recorded.set(False) + logging_obj: Final = kwargs.get("logging_obj") + self_recorded_token: Final = _guardrail_self_recorded.set(False) try: - response = await func(*args, **kwargs) + response: Final = await func(*args, **kwargs) if self.records_own_guardrail_information or _guardrail_self_recorded.get(): return response return self._process_response( @@ -1337,20 +1325,20 @@ def log_guardrail_information(func): @functools.wraps(func) def sync_wrapper(*args, **kwargs): - start_time = datetime.now() # Move start_time inside the wrapper - self: CustomGuardrail = args[0] - request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} - event_type = _infer_event_type_from_function_name(func.__name__) + start_time: Final = datetime.now() # Move start_time inside the wrapper + self: Final[CustomGuardrail] = args[0] + request_data: Final[dict] = kwargs.get("data") or kwargs.get("request_data") or {} + event_type: Final = _infer_event_type_from_function_name(func.__name__) # Store original inputs for comparison (for apply_guardrail functions) original_inputs = None if func.__name__ == "apply_guardrail" and "inputs" in kwargs: original_inputs = kwargs.get("inputs") - logging_obj = kwargs.get("logging_obj") - self_recorded_token = _guardrail_self_recorded.set(False) + logging_obj: Final = kwargs.get("logging_obj") + self_recorded_token: Final = _guardrail_self_recorded.set(False) try: - response = func(*args, **kwargs) + response: Final = func(*args, **kwargs) if self.records_own_guardrail_information or _guardrail_self_recorded.get(): return response return self._process_response( diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 8b831b55da3..29ef04af123 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -2,25 +2,17 @@ # On success, logs events to Promptlayer import re import traceback -from typing import ( - TYPE_CHECKING, - Any, - AsyncGenerator, - Dict, - List, - Optional, - Tuple, - Union, -) +from collections.abc import AsyncGenerator +from typing import TYPE_CHECKING, Any, Final, Optional, Union from pydantic import BaseModel from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER from litellm.types.integrations.argilla import ArgillaItem +from litellm.types.integrations.custom_logger import AgenticLoopPlan from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.integrations.custom_logger import AgenticLoopPlan from litellm.types.utils import ( AdapterCompletionStreamWrapper, CallTypes, @@ -55,12 +47,12 @@ else: MCPPostCallResponseObject = Any MCPPreCallRequestObject = Any MCPPreCallResponseObject = Any - MCPDuringCallRequestObject = Any - MCPDuringCallResponseObject = Any + MCPDuringCallRequestObject: Final = Any + MCPDuringCallResponseObject: Final = Any PreRoutingHookResponse = Any -_BASE64_INLINE_PATTERN = re.compile( +_BASE64_INLINE_PATTERN: Final = re.compile( r"data:(?:application|image|audio|video)/[a-zA-Z0-9.+-]+;base64,[A-Za-z0-9+/=\s]+", re.MULTILINE, ) @@ -82,10 +74,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ self.message_logging = message_logging self.turn_off_message_logging = turn_off_message_logging - pass @staticmethod - def get_callback_env_vars(callback_name: Optional[str] = None) -> List[str]: + def get_callback_env_vars(callback_name: str | None = None) -> list[str]: """ Return the environment variables associated with a given callback name as defined in the proxy callback registry. @@ -99,24 +90,24 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if callback_name is None: return [] - normalized_name = callback_name.lower() + normalized_name: Final = callback_name.lower() - alias_map = { + alias_map: Final = { "langfuse_otel": "langfuse", } - lookup_name = alias_map.get(normalized_name, normalized_name) + lookup_name: Final = alias_map.get(normalized_name, normalized_name) try: from litellm.proxy._types import AllCallbacks except Exception: return [] - callbacks = AllCallbacks() - callback_info = getattr(callbacks, lookup_name, None) + callbacks: Final = AllCallbacks() + callback_info: Final = getattr(callbacks, lookup_name, None) if callback_info is None: return [] - params = getattr(callback_info, "litellm_callback_params", None) + params: Final = getattr(callback_info, "litellm_callback_params", None) if not params: return [] @@ -145,7 +136,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_pre_api_call(self, model, messages, kwargs): pass - async def async_pre_request_hook(self, model: str, messages: List, kwargs: Dict) -> Optional[Dict]: + async def async_pre_request_hook(self, model: str, messages: list, kwargs: dict) -> dict | None: """ Hook called before making the API request to allow modifying request parameters. @@ -169,7 +160,6 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return kwargs ``` """ - pass async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): pass @@ -179,26 +169,25 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"): """Called when an audit log is created. Override in subclasses to handle.""" - pass #### PROMPT MANAGEMENT HOOKS #### async def async_get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, litellm_logging_obj: LiteLLMLoggingObj, - prompt_spec: Optional[PromptSpec] = None, - tools: Optional[List[Dict]] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + tools: list[dict] | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Returns: - model: str - the model to use (can be pulled from prompt management tool) @@ -210,17 +199,17 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac def get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Returns: - model: str - the model to use (can be pulled from prompt management tool) @@ -237,11 +226,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_pre_routing_hook( self, model: str, - request_kwargs: Dict, - messages: Optional[List[Dict[str, Any]]] = None, - input: Optional[Union[str, List]] = None, - specific_deployment: Optional[bool] = False, - ) -> Optional[PreRoutingHookResponse]: + request_kwargs: dict, + messages: list[dict[str, Any]] | None = None, + input: str | list | None = None, + specific_deployment: bool | None = False, + ) -> PreRoutingHookResponse | None: """ This hook is called before the routing decision is made. @@ -252,16 +241,14 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_filter_deployments( self, model: str, - healthy_deployments: List, - messages: Optional[List[AllMessageValues]], - request_kwargs: Optional[dict] = None, - parent_otel_span: Optional[Span] = None, - ) -> List[dict]: + healthy_deployments: list, + messages: list[AllMessageValues] | None, + request_kwargs: dict | None = None, + parent_otel_span: Span | None = None, + ) -> list[dict]: return healthy_deployments - async def async_pre_call_deployment_hook( - self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] - ) -> Optional[dict]: + async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None: """ Allow modifying the request just before it's sent to the deployment. @@ -269,41 +256,38 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac Used in managed_files.py """ + + async def async_pre_call_check(self, deployment: dict, parent_otel_span: Span | None) -> dict | None: pass - async def async_pre_call_check(self, deployment: dict, parent_otel_span: Optional[Span]) -> Optional[dict]: - pass - - def pre_call_check(self, deployment: dict) -> Optional[dict]: + def pre_call_check(self, deployment: dict) -> dict | None: pass async def async_post_call_success_deployment_hook( self, request_data: dict, response: LLMResponseTypes, - call_type: Optional[CallTypes], - ) -> Optional[LLMResponseTypes]: + call_type: CallTypes | None, + ) -> LLMResponseTypes | None: """ Allow modifying / reviewing the response just after it's received from the deployment. """ - pass async def async_post_call_streaming_deployment_hook( self, request_data: dict, response_chunk: Any, - call_type: Optional[CallTypes], - ) -> Optional[Any]: + call_type: CallTypes | None, + ) -> Any | None: """ Allow modifying streaming chunks just before they're returned to the user. This is called for each streaming chunk in the response. """ - pass #### Fallback Events - router/proxy only #### async def log_model_group_rate_limit_error( - self, exception: Exception, original_model_group: Optional[str], kwargs: dict + self, exception: Exception, original_model_group: str | None, kwargs: dict ): pass @@ -315,33 +299,30 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac #### ADAPTERS #### Allow calling 100+ LLMs in custom format - https://github.com/BerriAI/litellm/pulls - def translate_completion_input_params(self, kwargs) -> Optional[ChatCompletionRequest]: + def translate_completion_input_params(self, kwargs) -> ChatCompletionRequest | None: """ Translates the input params, from the provider's native format to the litellm.completion() format. """ - pass - def translate_completion_output_params(self, response: ModelResponse) -> Optional[BaseModel]: + def translate_completion_output_params(self, response: ModelResponse) -> BaseModel | None: """ Translates the output params, from the OpenAI format to the custom format. """ - pass def translate_completion_output_params_streaming( self, completion_stream: Any - ) -> Optional[AdapterCompletionStreamWrapper]: + ) -> AdapterCompletionStreamWrapper | None: """ Translates the streaming chunk, from the OpenAI format to the custom format. """ - pass ### DATASET HOOKS #### - currently only used for Argilla async def async_dataset_hook( self, logged_item: ArgillaItem, - standard_logging_payload: Optional[StandardLoggingPayload], - ) -> Optional[ArgillaItem]: + standard_logging_payload: StandardLoggingPayload | None, + ) -> ArgillaItem | None: """ - Decide if the result should be logged to Argilla. - Modify the result before logging to Argilla. @@ -360,9 +341,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac cache: "DualCache", data: dict, call_type: CallTypesLiteral, - ) -> Optional[ - Union[Exception, str, dict] - ]: # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm + ) -> ( + Exception | str | dict | None + ): # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm pass async def async_post_call_response_headers_hook( @@ -370,9 +351,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac data: dict, user_api_key_dict: UserAPIKeyAuth, response: Any, - request_headers: Optional[Dict[str, str]] = None, - litellm_call_info: Optional[Dict[str, Any]] = None, - ) -> Optional[Dict[str, str]]: + request_headers: dict[str, str] | None = None, + litellm_call_info: dict[str, Any] | None = None, + ) -> dict[str, str] | None: """ Called after an LLM API call (success or failure) to allow injecting custom HTTP response headers. @@ -398,7 +379,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac request_data: dict, original_exception: Exception, user_api_key_dict: UserAPIKeyAuth, - traceback_str: Optional[str] = None, + traceback_str: str | None = None, ) -> Optional["HTTPException"]: """ Called after an LLM API call fails. Can return or raise HTTPException to transform error responses. @@ -413,7 +394,6 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac - Optional[HTTPException]: Return an HTTPException to transform the error response sent to the client. Return None to use the original exception. """ - pass async def async_post_call_success_hook( self, @@ -423,11 +403,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> Any: pass - async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: + async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result - def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> Tuple[dict, Any]: + def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]: """For masking logged request/response. Return a modified version of the request/result.""" return kwargs, result @@ -493,7 +473,6 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) except Exception: print_verbose(f"Custom Logger Error - {traceback.format_exc()}") - pass async def async_log_event(self, kwargs, response_obj, start_time, end_time, print_verbose, callback_func): # Method definition @@ -507,7 +486,6 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) except Exception: print_verbose(f"Custom Logger Error - {traceback.format_exc()}") - pass ######################################################### # MCP TOOL CALL HOOKS @@ -515,7 +493,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_post_mcp_tool_call_hook( self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time - ) -> Optional[MCPPostCallResponseObject]: + ) -> MCPPostCallResponseObject | None: """ This log gets called after the MCP tool call is made. @@ -537,12 +515,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac self, response: Any, model: str, - messages: List[Dict], - tools: Optional[List[Dict]], + messages: list[dict], + tools: list[dict] | None, stream: bool, custom_llm_provider: str, - kwargs: Dict, - ) -> Tuple[bool, Dict]: + kwargs: dict, + ) -> tuple[bool, dict]: """ Hook to determine if agentic loop should be executed. @@ -593,15 +571,15 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_run_agentic_loop( self, - tools: Dict, + tools: dict, model: str, - messages: List[Dict], + messages: list[dict], response: Any, anthropic_messages_provider_config: Any, - anthropic_messages_optional_request_params: Dict, + anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, - kwargs: Dict, + kwargs: dict, ) -> Any: """ Hook to execute agentic loop based on context from should_run hook. @@ -659,19 +637,18 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return final_response """ - pass async def async_build_agentic_loop_plan( self, - tools: Dict, + tools: dict, model: str, - messages: List[Dict], + messages: list[dict], response: Any, anthropic_messages_provider_config: Any, - anthropic_messages_optional_request_params: Dict, + anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, - kwargs: Dict, + kwargs: dict, ) -> AgenticLoopPlan: """ Build a typed rerun plan for Anthropic Messages agentic loops. @@ -685,7 +662,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac self, response: Any, plan: AgenticLoopPlan, - kwargs: Dict, + kwargs: dict, ) -> Any: """ Post-process the response returned by the agentic-loop follow-up call. @@ -718,18 +695,18 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac Default does nothing. """ - return None + return async def async_should_run_chat_completion_agentic_loop( self, response: Any, model: str, - messages: List[Dict], - tools: Optional[List[Dict]], + messages: list[dict], + tools: list[dict] | None, stream: bool, custom_llm_provider: str, - kwargs: Dict, - ) -> Tuple[bool, Dict]: + kwargs: dict, + ) -> tuple[bool, dict]: """ Hook to determine if chat completion agentic loop should be executed. """ @@ -737,30 +714,29 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_run_chat_completion_agentic_loop( self, - tools: Dict, + tools: dict, model: str, - messages: List[Dict], + messages: list[dict], response: Any, - optional_params: Dict, + optional_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, - kwargs: Dict, + kwargs: dict, ) -> Any: """ Hook to execute chat completion agentic loop based on context from should_run hook. """ - pass async def async_build_chat_completion_agentic_loop_plan( self, - tools: Dict, + tools: dict, model: str, - messages: List[Dict], + messages: list[dict], response: Any, - optional_params: Dict, + optional_params: dict, logging_obj: "LiteLLMLoggingObj", stream: bool, - kwargs: Dict, + kwargs: dict, ) -> AgenticLoopPlan: """ Build a typed rerun plan for chat-completions agentic loops. @@ -780,10 +756,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac This function truncates the error string and the message content if they exceed a certain length. """ - MAX_STR_LENGTH = 10_000 + MAX_STR_LENGTH: Final = 10_000 # Truncate fields that might exceed max length - fields_to_truncate = ["error_str", "messages", "response"] + fields_to_truncate: Final = ["error_str", "messages", "response"] for field in fields_to_truncate: self._truncate_field( standard_logging_object=standard_logging_object, @@ -807,9 +783,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac - Converting to string and then truncating the logged content catches this 2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user """ - field_value = standard_logging_object.get(field_name) # type: ignore + field_value: Final = standard_logging_object.get(field_name) # type: ignore if field_value: - str_value = str(field_value) + str_value: Final = str(field_value) if len(str_value) > max_length: standard_logging_object[field_name] = self._truncate_text( # type: ignore text=str_value, max_length=max_length @@ -823,7 +799,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac else text ) - def _select_metadata_field(self, request_kwargs: Optional[Dict] = None) -> Optional[str]: + def _select_metadata_field(self, request_kwargs: dict | None = None) -> str | None: """ Select the metadata field to use for logging @@ -838,7 +814,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return LITELLM_METADATA_FIELD return OLD_LITELLM_METADATA_FIELD - def redact_standard_logging_payload_from_model_call_details(self, model_call_details: Dict) -> Dict: + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. @@ -850,13 +826,13 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac This is useful for logging payloads that contain sensitive information. """ - import litellm from copy import copy + import litellm from litellm import Choices, Message, ModelResponse - turn_off_message_logging: bool = getattr(self, "turn_off_message_logging", False) - excluded_fields: Optional[List[str]] = getattr(litellm, "standard_logging_payload_excluded_fields", None) + turn_off_message_logging: Final[bool] = getattr(self, "turn_off_message_logging", False) + excluded_fields: Final[list[str] | None] = getattr(litellm, "standard_logging_payload_excluded_fields", None) # Early return if no processing needed if turn_off_message_logging is False and not excluded_fields: @@ -864,13 +840,13 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac # Only make a shallow copy of the top-level dict to avoid deepcopy issues # with complex objects like AuthenticationError that may be present - model_call_details_copy = copy(model_call_details) - standard_logging_object = model_call_details.get("standard_logging_object") + model_call_details_copy: Final = copy(model_call_details) + standard_logging_object: Final = model_call_details.get("standard_logging_object") if standard_logging_object is None: return model_call_details_copy # Make a copy of just the standard_logging_object to avoid modifying the original - standard_logging_object_copy = copy(standard_logging_object) + standard_logging_object_copy: Final = copy(standard_logging_object) # Handle excluded fields - remove them entirely from the payload if excluded_fields: @@ -880,19 +856,19 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac # Handle turn_off_message_logging - redact messages and responses (if not already excluded) if turn_off_message_logging: - redacted_str = "redacted-by-litellm" + redacted_str: Final = "redacted-by-litellm" if "messages" not in (excluded_fields or []) and standard_logging_object_copy.get("messages") is not None: standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()] if "response" not in (excluded_fields or []) and standard_logging_object_copy.get("response") is not None: - response = standard_logging_object_copy["response"] + response: Final = standard_logging_object_copy["response"] # Check if this is a ResponsesAPIResponse (has "output" field) if isinstance(response, dict) and "output" in response: # Make a copy to avoid modifying the original from copy import deepcopy - response_copy = deepcopy(response) + response_copy: Final = deepcopy(response) # Redact content in output array if isinstance(response_copy.get("output"), list): for output_item in response_copy["output"]: @@ -905,8 +881,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac standard_logging_object_copy["response"] = response_copy else: # Standard ModelResponse format - model_response = ModelResponse(choices=[Choices(message=Message(content=redacted_str))]) - model_response_dict = model_response.model_dump() + model_response: Final = ModelResponse(choices=[Choices(message=Message(content=redacted_str))]) + model_response_dict: Final = model_response.model_dump() standard_logging_object_copy["response"] = model_response_dict model_call_details_copy["standard_logging_object"] = standard_logging_object_copy @@ -915,11 +891,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def get_proxy_server_request_from_cold_storage_with_object_key( self, object_key: str, - ) -> Optional[dict]: + ) -> dict | None: """ Get the proxy server request from cold storage using the object key directly. """ - pass def handle_callback_failure(self, callback_name: str): """ @@ -931,23 +906,23 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac import litellm from litellm._logging import verbose_logger - all_callbacks = litellm.logging_callback_manager._get_all_callbacks() + all_callbacks: Final = litellm.logging_callback_manager._get_all_callbacks() for callback_obj in all_callbacks: if hasattr(callback_obj, "increment_callback_logging_failure"): - verbose_logger.debug(f"Incrementing callback failure metric for {callback_name}") + verbose_logger.debug("Incrementing callback failure metric for %s", callback_name) callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore return verbose_logger.debug( - f"No callback with increment_callback_logging_failure method found for {callback_name}. " - "Ensure 'prometheus' is in your callbacks config." + "No callback with increment_callback_logging_failure method found for %s. Ensure 'prometheus' is in your callbacks config.", + callback_name, ) except Exception as e: from litellm._logging import verbose_logger - verbose_logger.debug(f"Error in handle_callback_failure for {callback_name}: {str(e)}") + verbose_logger.debug("Error in handle_callback_failure for %s: %s", callback_name, e) async def _strip_base64_from_messages( self, @@ -964,9 +939,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Any = payload.get("messages", []) - messages: List[Any] = raw_messages if isinstance(raw_messages, list) else [] - verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages") + raw_messages: Final[Any] = payload.get("messages", []) + messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: payload["messages"] = self._process_messages(messages=messages, max_depth=max_depth) @@ -978,7 +953,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if isinstance(content, list): total_items += len(content) - verbose_logger.debug(f"[CustomLogger] Completed base64 strip; retained {total_items} content items") + verbose_logger.debug("[CustomLogger] Completed base64 strip; retained %s content items", total_items) return payload def _strip_base64_from_messages_sync( @@ -996,9 +971,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac • Keep untyped or text content. • Recursively redact inline base64 blobs in *any* string field, at any depth. """ - raw_messages: Any = payload.get("messages", []) - messages: List[Any] = raw_messages if isinstance(raw_messages, list) else [] - verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages") + raw_messages: Final[Any] = payload.get("messages", []) + messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else [] + verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages)) if messages: payload["messages"] = self._process_messages(messages=messages, max_depth=max_depth) @@ -1010,7 +985,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if isinstance(content, list): total_items += len(content) - verbose_logger.debug(f"[CustomLogger] Completed base64 strip; retained {total_items} content items") + verbose_logger.debug("[CustomLogger] Completed base64 strip; retained %s content items", total_items) return payload def _redact_base64( @@ -1021,12 +996,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac ) -> Any: """Recursively redact inline base64 from any nested structure with a max recursion depth limit.""" if depth > max_depth: - verbose_logger.warning(f"[CustomLogger] Max recursion depth {max_depth} reached while redacting base64") + verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth) return "[MAX_DEPTH_REACHED]" if isinstance(value, str): if _BASE64_INLINE_PATTERN.search(value): - verbose_logger.debug(f"[CustomLogger] Redacted inline base64 string: {value[:40]}...") + verbose_logger.debug("[CustomLogger] Redacted inline base64 string: %s...", value[:40]) return _BASE64_INLINE_PATTERN.sub("[BASE64_REDACTED]", value) return value @@ -1044,21 +1019,21 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return True if "file" in content: return False - ctype = content.get("type") + ctype: Final = content.get("type") return not (isinstance(ctype, str) and ctype != "text") def _process_messages( self, - messages: List[Any], + messages: list[Any], max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER, - ) -> List[Dict[str, Any]]: - filtered_messages: List[Dict[str, Any]] = [] + ) -> list[dict[str, Any]]: + filtered_messages: Final[list[dict[str, Any]]] = [] for msg in messages: if not isinstance(msg, dict): continue contents: Any = msg.get("content") if isinstance(contents, list): - cleaned: List[Any] = [] + cleaned: list[Any] = [] for c in contents: if self._should_keep_content(content=c): cleaned.append(self._redact_base64(value=c, max_depth=max_depth)) diff --git a/litellm/integrations/custom_prompt_management.py b/litellm/integrations/custom_prompt_management.py index fbca1867793..7078416b7a8 100644 --- a/litellm/integrations/custom_prompt_management.py +++ b/litellm/integrations/custom_prompt_management.py @@ -1,5 +1,3 @@ -from typing import List, Optional, Tuple - from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prompt_management_base import ( PromptManagementBase, @@ -13,8 +11,8 @@ from litellm.types.utils import StandardCallbackDynamicParams class CustomPromptManagement(CustomLogger, PromptManagementBase): def __init__( self, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, **kwargs, ): self.ignore_prompt_manager_model = ignore_prompt_manager_model @@ -23,17 +21,17 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase): def get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Returns: - model: str - the model to use (can be pulled from prompt management tool) @@ -48,30 +46,30 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase): def should_run_prompt_management( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], + prompt_id: str | None, + prompt_spec: PromptSpec | None, dynamic_callback_params: StandardCallbackDynamicParams, ) -> bool: return True def _compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_spec: PromptSpec | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: raise NotImplementedError("Custom prompt management does not support compile prompt helper") async def async_compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: raise NotImplementedError("Custom prompt management does not support async compile prompt helper") diff --git a/litellm/integrations/custom_secret_manager.py b/litellm/integrations/custom_secret_manager.py index a1bb7b00d92..e59842d409a 100644 --- a/litellm/integrations/custom_secret_manager.py +++ b/litellm/integrations/custom_secret_manager.py @@ -37,7 +37,7 @@ Usage: """ from abc import abstractmethod -from typing import Any, Dict, Optional, Union +from typing import Any import httpx @@ -87,7 +87,7 @@ class CustomSecretManager(BaseSecretManager): def __init__( self, - secret_manager_name: Optional[str] = None, + secret_manager_name: str | None = None, **kwargs, ): """ @@ -106,9 +106,9 @@ class CustomSecretManager(BaseSecretManager): async def async_read_secret( self, secret_name: str, - optional_params: Optional[dict] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Optional[str]: + optional_params: dict | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: """ Asynchronously read a secret from your custom secret manager. @@ -123,15 +123,14 @@ class CustomSecretManager(BaseSecretManager): Raises: Exception: If there's an error reading the secret """ - pass @abstractmethod def sync_read_secret( self, secret_name: str, - optional_params: Optional[dict] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - ) -> Optional[str]: + optional_params: dict | None = None, + timeout: float | httpx.Timeout | None = None, + ) -> str | None: """ Synchronously read a secret from your custom secret manager. @@ -146,17 +145,16 @@ class CustomSecretManager(BaseSecretManager): Raises: Exception: If there's an error reading the secret """ - pass async def async_write_secret( self, secret_name: str, secret_value: str, - description: Optional[str] = None, - optional_params: Optional[dict] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, - tags: Optional[Union[dict, list]] = None, - ) -> Dict[str, Any]: + description: str | None = None, + optional_params: dict | None = None, + timeout: float | httpx.Timeout | None = None, + tags: dict | list | None = None, + ) -> dict[str, Any]: """ Asynchronously write a secret to your custom secret manager. @@ -185,9 +183,9 @@ class CustomSecretManager(BaseSecretManager): async def async_delete_secret( self, secret_name: str, - recovery_window_in_days: Optional[int] = 7, - optional_params: Optional[dict] = None, - timeout: Optional[Union[float, httpx.Timeout]] = None, + recovery_window_in_days: int | None = 7, + optional_params: dict | None = None, + timeout: float | httpx.Timeout | None = None, ) -> dict: """ Asynchronously delete a secret from your custom secret manager. @@ -227,7 +225,7 @@ class CustomSecretManager(BaseSecretManager): verbose_logger.debug("No environment validation configured for custom secret manager") return True - async def async_health_check(self, timeout: Optional[Union[float, httpx.Timeout]] = None) -> bool: + async def async_health_check(self, timeout: float | httpx.Timeout | None = None) -> bool: """ Perform a health check on your secret manager. @@ -239,7 +237,7 @@ class CustomSecretManager(BaseSecretManager): Returns: True if the secret manager is healthy, False otherwise """ - verbose_logger.debug(f"Health check not implemented for {self.secret_manager_name}") + verbose_logger.debug("Health check not implemented for %s", self.secret_manager_name) return True def __repr__(self) -> str: diff --git a/litellm/integrations/custom_sso_handler.py b/litellm/integrations/custom_sso_handler.py index 202e488e0e4..345a9051c06 100644 --- a/litellm/integrations/custom_sso_handler.py +++ b/litellm/integrations/custom_sso_handler.py @@ -1,3 +1,5 @@ +from typing import Final + from fastapi import Request from fastapi_sso.sso.base import OpenID @@ -29,7 +31,7 @@ class CustomSSOLoginHandler(CustomLogger): feature_name="Custom UI SSO", ) - request_headers_dict = dict(request.headers) + request_headers_dict: Final = dict(request.headers) return OpenID( id=request_headers_dict.get("x-litellm-user-id"), email=request_headers_dict.get("x-litellm-user-email"), diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 20239d831cc..fd4faeed41a 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -18,8 +18,9 @@ import datetime import os import time import traceback +from collections.abc import Sequence from datetime import datetime as datetimeObj -from typing import Any, Dict, List, Optional, Sequence, Union +from typing import Any, Final import httpx from httpx import Response @@ -28,16 +29,16 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_batch_logger import CustomBatchLogger -from litellm.integrations.datadog.datadog_mock_client import ( - should_use_datadog_mock, - create_mock_datadog_client, -) from litellm.integrations.datadog.datadog_handler import ( + get_datadog_base_url_from_env, get_datadog_hostname, get_datadog_service, get_datadog_source, get_datadog_tags, - get_datadog_base_url_from_env, +) +from litellm.integrations.datadog.datadog_mock_client import ( + create_mock_datadog_client, + should_use_datadog_mock, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.llms.custom_httpx.http_handler import ( @@ -51,10 +52,10 @@ from litellm.types.integrations.datadog import ( DD_ERRORS, DD_MAX_BATCH_SIZE, DD_MAX_PAYLOAD_SIZE_BYTES, - DataDogStatus, DatadogInitParams, DatadogPayload, DatadogProxyFailureHookJsonMessage, + DataDogStatus, ) from litellm.types.services import ServiceLoggerPayload, ServiceTypes from litellm.types.utils import StandardLoggingPayload @@ -65,17 +66,17 @@ from ..additional_logging_utils import AdditionalLoggingUtils # specify what ServiceTypes are logged as success events to DD. (We don't want to spam DD traces with large number of service types) -DD_LOGGED_SUCCESS_SERVICE_TYPES = [ +DD_LOGGED_SUCCESS_SERVICE_TYPES: Final = [ ServiceTypes.RESET_BUDGET_JOB, ] def _resolve_dd_batch_size() -> int: - raw = os.getenv("DD_BATCH_SIZE") + raw: Final = os.getenv("DD_BATCH_SIZE") if raw is None: return DD_MAX_BATCH_SIZE try: - value = int(raw) + value: Final = int(raw) except ValueError: verbose_logger.warning( "Datadog: ignoring invalid DD_BATCH_SIZE=%r, using %s", @@ -93,10 +94,10 @@ class DataDogLogger( # Class variables or attributes def __init__( self, - dd_api_key: Optional[str] = None, - dd_site: Optional[str] = None, - dd_agent_host: Optional[str] = None, - dd_agent_port: Optional[str] = None, + dd_api_key: str | None = None, + dd_site: str | None = None, + dd_agent_host: str | None = None, + dd_agent_port: str | None = None, allow_env_credentials: bool = True, **kwargs, ): @@ -135,14 +136,14 @@ class DataDogLogger( ######################################################### # Handle datadog_params set as litellm.datadog_params ######################################################### - dict_datadog_params = self._get_datadog_params() + dict_datadog_params: Final = self._get_datadog_params() kwargs.update(dict_datadog_params) self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Configure DataDog endpoint (Agent or Direct API) # Prefer explicit kwargs, then fall back to env vars - resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST") + resolved_agent_host: Final = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST") if resolved_agent_host: self._configure_dd_agent( dd_agent_host=resolved_agent_host, @@ -158,7 +159,7 @@ class DataDogLogger( ) # Optional override for testing - dd_base_url = get_datadog_base_url_from_env() + dd_base_url: Final = get_datadog_base_url_from_env() if dd_base_url: self.intake_url = f"{dd_base_url}/api/v2/logs" self.sync_client = _get_httpx_client() @@ -170,20 +171,20 @@ class DataDogLogger( batch_size=_resolve_dd_batch_size(), ) except Exception as e: - verbose_logger.exception(f"Datadog: Got exception on init Datadog client {str(e)}") + verbose_logger.exception("Datadog: Got exception on init Datadog client %s", e) raise e - def _get_datadog_params(self) -> Dict: + def _get_datadog_params(self) -> dict: """ Get the datadog_params from litellm.datadog_params These are params specific to initializing the DataDogLogger e.g. turn_off_message_logging """ - dict_datadog_params: Dict = {} + dict_datadog_params: dict = {} if litellm.datadog_params is not None: if isinstance(litellm.datadog_params, DatadogInitParams): dict_datadog_params = litellm.datadog_params.model_dump() - elif isinstance(litellm.datadog_params, Dict): + elif isinstance(litellm.datadog_params, dict): # only allow params that are of DatadogInitParams dict_datadog_params = DatadogInitParams(**litellm.datadog_params).model_dump() return dict_datadog_params @@ -191,8 +192,8 @@ class DataDogLogger( def _configure_dd_agent( self, dd_agent_host: str, - dd_agent_port: Optional[str] = None, - dd_api_key: Optional[str] = None, + dd_agent_port: str | None = None, + dd_api_key: str | None = None, allow_env_credentials: bool = True, ) -> None: """ @@ -204,17 +205,17 @@ class DataDogLogger( dd_api_key: Datadog API key. Falls back to DD_API_KEY env var when allow_env_credentials is True. Optional when using agent. allow_env_credentials: When False, never read the API key from DD_API_KEY env var. """ - resolved_port = dd_agent_port or os.getenv("LITELLM_DD_AGENT_PORT", "10518") # default port for logs + resolved_port: Final = dd_agent_port or os.getenv("LITELLM_DD_AGENT_PORT", "10518") # default port for logs self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs" self.DD_API_KEY = dd_api_key or ( os.getenv("DD_API_KEY") if allow_env_credentials else None ) # Optional when using agent - verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}") + verbose_logger.debug("Datadog: Using DD Agent at %s", self.intake_url) def _configure_dd_direct_api( self, - dd_api_key: Optional[str] = None, - dd_site: Optional[str] = None, + dd_api_key: str | None = None, + dd_site: str | None = None, allow_env_credentials: bool = True, ) -> None: """ @@ -228,8 +229,8 @@ class DataDogLogger( Raises: Exception: If required credentials are not provided via args or env vars """ - resolved_api_key = dd_api_key or (os.getenv("DD_API_KEY") if allow_env_credentials else None) - resolved_site = dd_site or os.getenv("DD_SITE") + resolved_api_key: Final = dd_api_key or (os.getenv("DD_API_KEY") if allow_env_credentials else None) + resolved_site: Final = dd_site or os.getenv("DD_SITE") if resolved_api_key is None: raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>") @@ -256,8 +257,7 @@ class DataDogLogger( await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: - verbose_logger.exception(f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}") - pass + verbose_logger.exception("Datadog Layer Error - %s\n%s", e, traceback.format_exc()) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: @@ -265,16 +265,15 @@ class DataDogLogger( await self._log_async_event(kwargs, response_obj, start_time, end_time) except Exception as e: - verbose_logger.exception(f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}") - pass + verbose_logger.exception("Datadog Layer Error - %s\n%s", e, traceback.format_exc()) async def async_post_call_failure_hook( self, request_data: dict, original_exception: Exception, user_api_key_dict: Any, - traceback_str: Optional[str] = None, - ) -> Optional[Any]: + traceback_str: str | None = None, + ) -> Any | None: """ Log proxy-level failures (e.g. 401 auth, DB connection errors) to Datadog. @@ -288,23 +287,23 @@ class DataDogLogger( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - error_information = StandardLoggingPayloadSetup.get_error_information( + error_information: Final = StandardLoggingPayloadSetup.get_error_information( original_exception=original_exception, traceback_str=traceback_str, ) - _code = error_information.get("error_code") or "" - status_code: Optional[int] = None + _code: Final = error_information.get("error_code") or "" + status_code: int | None = None if _code and str(_code).strip().isdigit(): status_code = int(_code) # Use project-standard sanitized user context when running in proxy - user_context: Dict[str, Any] = {} + user_context: dict[str, Any] = {} try: from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, ) - _meta = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + _meta: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( user_api_key_dict=user_api_key_dict ) user_context = dict(_meta) if isinstance(_meta, dict) else _meta @@ -319,7 +318,7 @@ class DataDogLogger( if hasattr(user_api_key_dict, "end_user_id"): user_context["end_user_id"] = getattr(user_api_key_dict, "end_user_id", None) - message_payload: DatadogProxyFailureHookJsonMessage = { + message_payload: Final[DatadogProxyFailureHookJsonMessage] = { "exception": error_information.get("error_message") or str(original_exception), "error_class": error_information.get("error_class") or original_exception.__class__.__name__, "status_code": status_code, @@ -327,7 +326,7 @@ class DataDogLogger( "user_api_key_dict": user_context, } - dd_payload = DatadogPayload( + dd_payload: Final = DatadogPayload( ddsource=get_datadog_source(), ddtags=",".join(get_datadog_tags()), hostname=get_datadog_hostname(), @@ -341,7 +340,7 @@ class DataDogLogger( if len(self.log_queue) >= self.batch_size: await self.flush_queue() except Exception as e: - verbose_logger.exception(f"Datadog: async_post_call_failure_hook - {str(e)}\n{traceback.format_exc()}") + verbose_logger.exception("Datadog: async_post_call_failure_hook - %s\n%s", e, traceback.format_exc()) return None async def async_send_batch(self): @@ -359,7 +358,7 @@ class DataDogLogger( verbose_logger.exception("Datadog: log_queue does not exist") return - batch_to_send = self.log_queue[:] + batch_to_send: Final = self.log_queue[:] self.log_queue = [] try: @@ -372,18 +371,18 @@ class DataDogLogger( if self.is_mock_mode: verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") - undelivered = await self._send_with_413_split(batch_to_send) + undelivered: Final = await self._send_with_413_split(batch_to_send) if undelivered: self.log_queue = undelivered + self.log_queue if self.is_mock_mode: - verbose_logger.debug(f"[DATADOG MOCK] Batch of {len(batch_to_send)} events successfully mocked") + verbose_logger.debug("[DATADOG MOCK] Batch of %s events successfully mocked", len(batch_to_send)) except Exception as e: self.log_queue = batch_to_send + self.log_queue - verbose_logger.exception(f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}") + verbose_logger.exception("Datadog Error sending batch API - %s\n%s", e, traceback.format_exc()) - async def _send_with_413_split(self, batch: List) -> List: + async def _send_with_413_split(self, batch: list) -> list: """ Send a batch, halving any sub-batch that exceeds Datadog's intake limits before sending, and halving again on a 413 (payload too large) response, since Datadog @@ -396,7 +395,7 @@ class DataDogLogger( that could not be delivered because of a non-413 (transient) error, so the caller re-queues only those and never the events already accepted by Datadog. """ - pending: List[List] = [batch] + pending: Final[list[list]] = [batch] while pending: chunk = pending.pop() if not chunk: @@ -412,7 +411,7 @@ class DataDogLogger( if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413: response = e.response else: - verbose_logger.exception(f"Datadog Error sending batch API - {str(e)}") + verbose_logger.exception("Datadog Error sending batch API - %s", e) return self._undelivered(chunk, pending) if response.status_code == 413: @@ -441,7 +440,7 @@ class DataDogLogger( return [] @staticmethod - def _undelivered(chunk: List, pending: List[List]) -> List: + def _undelivered(chunk: list, pending: list[list]) -> list: return chunk + [event for remaining in reversed(pending) for event in remaining] @staticmethod @@ -456,7 +455,7 @@ class DataDogLogger( if len(chunk) > DD_MAX_BATCH_SIZE: return True - payload_size_bytes = len(safe_dumps(chunk).encode("utf-8")) + payload_size_bytes: Final = len(safe_dumps(chunk).encode("utf-8")) return payload_size_bytes > DD_MAX_PAYLOAD_SIZE_BYTES async def flush_queue(self): @@ -494,12 +493,12 @@ class DataDogLogger( ) # Build headers - headers = {} + headers: Final = {} # Add API key if available (required for direct API, optional for agent) if self.DD_API_KEY: headers["DD-API-KEY"] = self.DD_API_KEY - response = self.sync_client.post( + response: Final = self.sync_client.post( url=self.intake_url, json=dd_payload, # type: ignore headers=headers, @@ -516,12 +515,10 @@ class DataDogLogger( ) except Exception as e: - verbose_logger.exception(f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}") - pass - pass + verbose_logger.exception("Datadog Layer Error - %s\n%s", e, traceback.format_exc()) async def _log_async_event(self, kwargs, response_obj, start_time, end_time): - dd_payload = self.create_datadog_logging_payload( + dd_payload: Final = self.create_datadog_logging_payload( kwargs=kwargs, response_obj=response_obj, start_time=start_time, @@ -529,7 +526,7 @@ class DataDogLogger( ) self.log_queue.append(dd_payload) - verbose_logger.debug(f"Datadog, event added to queue. Will flush in {self.flush_interval} seconds...") + verbose_logger.debug("Datadog, event added to queue. Will flush in %s seconds...", self.flush_interval) if len(self.log_queue) >= self.batch_size: await self.flush_queue() @@ -541,9 +538,9 @@ class DataDogLogger( ) -> DatadogPayload: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - json_payload = safe_dumps(standard_logging_object) + json_payload: Final = safe_dumps(standard_logging_object) verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload) - dd_payload = DatadogPayload( + dd_payload: Final = DatadogPayload( ddsource=get_datadog_source(), ddtags=",".join(get_datadog_tags(standard_logging_object=standard_logging_object)), hostname=get_datadog_hostname(), @@ -556,7 +553,7 @@ class DataDogLogger( def create_datadog_logging_payload( self, - kwargs: Union[dict, Any], + kwargs: dict | Any, response_obj: Any, start_time: datetime.datetime, end_time: datetime.datetime, @@ -574,7 +571,7 @@ class DataDogLogger( DatadogPayload: defined in types.py """ - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) + standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: raise ValueError("standard_logging_object not found in kwargs") @@ -585,13 +582,13 @@ class DataDogLogger( # Build the initial payload self.truncate_standard_logging_payload_content(standard_logging_object) - dd_payload = self._create_datadog_logging_payload_helper( + dd_payload: Final = self._create_datadog_logging_payload_helper( standard_logging_object=standard_logging_object, status=status, ) return dd_payload - async def async_send_compressed_data(self, data: List) -> Response: + async def async_send_compressed_data(self, data: list) -> Response: """ Async helper to send compressed data to datadog self.intake_url @@ -605,10 +602,10 @@ class DataDogLogger( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - compressed_data = gzip.compress(safe_dumps(data).encode("utf-8")) + compressed_data: Final = gzip.compress(safe_dumps(data).encode("utf-8")) # Build headers - headers = { + headers: Final = { "Content-Encoding": "gzip", "Content-Type": "application/json", } @@ -617,7 +614,7 @@ class DataDogLogger( if self.DD_API_KEY: headers["DD-API-KEY"] = self.DD_API_KEY - response = await self.async_client.post( + response: Final = await self.async_client.post( url=self.intake_url, data=compressed_data, # type: ignore headers=headers, @@ -627,11 +624,11 @@ class DataDogLogger( async def async_service_failure_hook( self, payload: ServiceLoggerPayload, - error: Optional[str] = "", - parent_otel_span: Optional[Any] = None, - start_time: Optional[Union[datetimeObj, float]] = None, - end_time: Optional[Union[float, datetimeObj]] = None, - event_metadata: Optional[dict] = None, + error: str | None = "", + parent_otel_span: Any | None = None, + start_time: datetimeObj | float | None = None, + end_time: float | datetimeObj | None = None, + event_metadata: dict | None = None, ): """ Logs failures from Redis, Postgres (Adjacent systems), as 'WARNING' on DataDog @@ -639,12 +636,12 @@ class DataDogLogger( - example - Redis is failing / erroring, will be logged on DataDog """ try: - _payload_dict = payload.model_dump() + _payload_dict: Final = payload.model_dump() _payload_dict.update(event_metadata or {}) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - _dd_message_str = safe_dumps(_payload_dict) - _dd_payload = DatadogPayload( + _dd_message_str: Final = safe_dumps(_payload_dict) + _dd_payload: Final = DatadogPayload( ddsource=get_datadog_source(), ddtags=",".join(get_datadog_tags()), hostname=get_datadog_hostname(), @@ -656,17 +653,16 @@ class DataDogLogger( self.log_queue.append(_dd_payload) except Exception as e: - verbose_logger.exception(f"Datadog: Logger - Exception in async_service_failure_hook: {e}") - pass + verbose_logger.exception("Datadog: Logger - Exception in async_service_failure_hook: %s", e) async def async_service_success_hook( self, payload: ServiceLoggerPayload, - error: Optional[str] = "", - parent_otel_span: Optional[Any] = None, - start_time: Optional[Union[datetimeObj, float]] = None, - end_time: Optional[Union[float, datetimeObj]] = None, - event_metadata: Optional[dict] = None, + error: str | None = "", + parent_otel_span: Any | None = None, + start_time: datetimeObj | float | None = None, + end_time: float | datetimeObj | None = None, + event_metadata: dict | None = None, ): """ Logs success from Redis, Postgres (Adjacent systems), as 'INFO' on DataDog @@ -678,13 +674,13 @@ class DataDogLogger( if payload.service not in DD_LOGGED_SUCCESS_SERVICE_TYPES: return - _payload_dict = payload.model_dump() + _payload_dict: Final = payload.model_dump() _payload_dict.update(event_metadata or {}) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - _dd_message_str = safe_dumps(_payload_dict) - _dd_payload = DatadogPayload( + _dd_message_str: Final = safe_dumps(_payload_dict) + _dd_payload: Final = DatadogPayload( ddsource=get_datadog_source(), ddtags=",".join(get_datadog_tags()), hostname=get_datadog_hostname(), @@ -696,11 +692,11 @@ class DataDogLogger( self.log_queue.append(_dd_payload) except Exception as e: - verbose_logger.exception(f"Datadog: Logger - Exception in async_service_failure_hook: {e}") + verbose_logger.exception("Datadog: Logger - Exception in async_service_failure_hook: %s", e) def _create_v0_logging_payload( self, - kwargs: Union[dict, Any], + kwargs: dict | Any, response_obj: Any, start_time: datetime.datetime, end_time: datetime.datetime, @@ -712,14 +708,14 @@ class DataDogLogger( (Not Recommended) If you want this to get logged set `litellm.datadog_use_v1 = True` """ - litellm_params = kwargs.get("litellm_params", {}) - metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None - messages = kwargs.get("messages") - optional_params = kwargs.get("optional_params", {}) - call_type = kwargs.get("call_type", "litellm.completion") - cache_hit = kwargs.get("cache_hit", False) + litellm_params: Final = kwargs.get("litellm_params", {}) + metadata: Final = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None + messages: Final = kwargs.get("messages") + optional_params: Final = kwargs.get("optional_params", {}) + call_type: Final = kwargs.get("call_type", "litellm.completion") + cache_hit: Final = kwargs.get("cache_hit", False) usage = response_obj["usage"] - id = response_obj.get("id", str(uuid.uuid4())) + id: Final = response_obj.get("id", str(uuid.uuid4())) usage = dict(usage) try: response_time = (end_time - start_time).total_seconds() * 1000 @@ -734,7 +730,7 @@ class DataDogLogger( # Clean Metadata before logging - never log raw metadata # the raw metadata can contain circular references which leads to infinite recursion # we clean out all extra litellm metadata params before logging - clean_metadata = {} + clean_metadata: Final = {} if isinstance(metadata, dict): for key, value in metadata.items(): # clean litellm metadata before logging @@ -748,7 +744,7 @@ class DataDogLogger( clean_metadata[key] = value # Build the initial payload - payload = { + payload: Final = { "id": id, "call_type": call_type, "cache_hit": cache_hit, @@ -767,11 +763,11 @@ class DataDogLogger( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - json_payload = safe_dumps(payload) + json_payload: Final = safe_dumps(payload) verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload) - dd_payload = DatadogPayload( + dd_payload: Final = DatadogPayload( ddsource=get_datadog_source(), ddtags=",".join(get_datadog_tags()), hostname=get_datadog_hostname(), @@ -788,38 +784,38 @@ class DataDogLogger( """Attach Datadog APM trace context if one is active.""" try: - trace_context = self._get_active_trace_context() + trace_context: Final = self._get_active_trace_context() if trace_context is None: return dd_payload["dd.trace_id"] = trace_context["trace_id"] - span_id = trace_context.get("span_id") + span_id: Final = trace_context.get("span_id") if span_id is not None: dd_payload["dd.span_id"] = span_id except Exception: verbose_logger.exception("Datadog: Failed to attach trace context to payload") - def _get_active_trace_context(self) -> Optional[Dict[str, str]]: + def _get_active_trace_context(self) -> dict[str, str] | None: try: current_span = None - current_span_fn = getattr(tracer, "current_span", None) + current_span_fn: Final = getattr(tracer, "current_span", None) if callable(current_span_fn): current_span = current_span_fn() if current_span is None: - current_root_span_fn = getattr(tracer, "current_root_span", None) + current_root_span_fn: Final = getattr(tracer, "current_root_span", None) if callable(current_root_span_fn): current_span = current_root_span_fn() if current_span is None: return None - trace_id = getattr(current_span, "trace_id", None) + trace_id: Final = getattr(current_span, "trace_id", None) if trace_id is None: return None - span_id = getattr(current_span, "span_id", None) - trace_context: Dict[str, str] = {"trace_id": str(trace_id)} + span_id: Final = getattr(current_span, "span_id", None) + trace_context: Final[dict[str, str]] = {"trace_id": str(trace_id)} if span_id is not None: trace_context["span_id"] = str(span_id) return trace_context @@ -835,13 +831,13 @@ class DataDogLogger( create_dummy_standard_logging_payload, ) - standard_logging_object = create_dummy_standard_logging_payload() - dd_payload = self._create_datadog_logging_payload_helper( + standard_logging_object: Final = create_dummy_standard_logging_payload() + dd_payload: Final = self._create_datadog_logging_payload_helper( standard_logging_object=standard_logging_object, status=DataDogStatus.INFO, ) - log_queue = [dd_payload] - response = await self.async_send_compressed_data(log_queue) + log_queue: Final = [dd_payload] + response: Final = await self.async_send_compressed_data(log_queue) try: response.raise_for_status() return IntegrationHealthCheckStatus( @@ -862,7 +858,7 @@ class DataDogLogger( async def get_request_response_payload( self, request_id: str, - start_time_utc: Optional[datetimeObj], - end_time_utc: Optional[datetimeObj], - ) -> Optional[dict]: + start_time_utc: datetimeObj | None, + end_time_utc: datetimeObj | None, + ) -> dict | None: pass diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 714a50eb2f2..b30700e98f2 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -2,7 +2,7 @@ import asyncio import os import time from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple, cast +from typing import Any, Final, cast from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -27,7 +27,7 @@ from litellm.types.utils import StandardLoggingPayload # request_tags / metadata cannot overwrite these, even when the key is # allowlisted via cost_tag_keys, because that would let an authenticated caller # spoof cost attribution (e.g. request_tags=["team:victim-team"]). -_RESERVED_TAG_KEYS: frozenset = frozenset( +_RESERVED_TAG_KEYS: Final[frozenset] = frozenset( { "env", "service", @@ -44,8 +44,8 @@ _RESERVED_TAG_KEYS: frozenset = frozenset( class DatadogCostManagementLogger(CustomBatchLogger): - def __init__(self, cost_tag_keys: Optional[List[str]] = None, **kwargs): - self.cost_tag_keys: List[str] = list(cost_tag_keys) if cost_tag_keys else [] + def __init__(self, cost_tag_keys: list[str] | None = None, **kwargs): + self.cost_tag_keys: list[str] = list(cost_tag_keys) if cost_tag_keys else [] self.dd_api_key = os.getenv("DD_API_KEY") self.dd_app_key = os.getenv("DD_APP_KEY") self.dd_site = os.getenv("DD_SITE", "datadoghq.com") @@ -71,7 +71,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) + standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: return @@ -84,17 +84,17 @@ class DatadogCostManagementLogger(CustomBatchLogger): await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"Datadog Cost Management: Error in async_log_success_event: {str(e)}") + verbose_logger.exception("Datadog Cost Management: Error in async_log_success_event: %s", e) async def async_send_batch(self): if not self.log_queue: return - batch_to_send = self.log_queue[:] + batch_to_send: Final = self.log_queue[:] self.log_queue = [] try: - aggregated_entries = self._aggregate_costs(batch_to_send) + aggregated_entries: Final = self._aggregate_costs(batch_to_send) if not aggregated_entries: verbose_logger.debug( "Datadog Cost Management: batch produced no aggregable entries; dropping %d log(s) from queue.", @@ -104,14 +104,14 @@ class DatadogCostManagementLogger(CustomBatchLogger): await self._upload_to_datadog(aggregated_entries) except Exception as e: self.log_queue = batch_to_send + self.log_queue - verbose_logger.exception(f"Datadog Cost Management: Error in async_send_batch: {str(e)}") + verbose_logger.exception("Datadog Cost Management: Error in async_send_batch: %s", e) - def _aggregate_costs(self, logs: List[StandardLoggingPayload]) -> List[DatadogFOCUSCostEntry]: + def _aggregate_costs(self, logs: list[StandardLoggingPayload]) -> list[DatadogFOCUSCostEntry]: """ Aggregates costs by Provider, Model, and Date. Returns a list of DatadogFOCUSCostEntry. """ - aggregator: Dict[Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry] = {} + aggregator: Final[dict[tuple[str, str, str, tuple[tuple[str, str], ...]], DatadogFOCUSCostEntry]] = {} for log in logs: try: @@ -159,13 +159,13 @@ class DatadogCostManagementLogger(CustomBatchLogger): aggregator[key]["BilledCost"] += cost except Exception as e: - verbose_logger.warning(f"Error processing log for cost aggregation: {e}") + verbose_logger.warning("Error processing log for cost aggregation: %s", e) continue return list(aggregator.values()) - def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]: - tags: Dict[str, str] = { + def _extract_tags(self, log: StandardLoggingPayload) -> dict[str, str]: + tags: Final[dict[str, str]] = { "env": get_datadog_env(), "service": get_datadog_service(), "host": get_datadog_hostname(), @@ -180,12 +180,12 @@ class DatadogCostManagementLogger(CustomBatchLogger): # cast because StandardLoggingMetadata is a TypedDict; we iterate it # as a generic mapping below. - metadata: Dict[str, Any] = cast(Dict[str, Any], log.get("metadata") or {}) + metadata: Final[dict[str, Any]] = cast(dict[str, Any], log.get("metadata") or {}) # Backwards-compat: team/user/model_group preserved regardless of allowlist. if metadata.get("user_api_key_alias"): tags["user"] = str(metadata["user_api_key_alias"]) - team_tag = ( + team_tag: Final = ( metadata.get("user_api_key_team_alias") or metadata.get("team_alias") or metadata.get("user_api_key_team_id") @@ -200,7 +200,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): # Reserved keys are hard-blocked here regardless of allowlist membership — # see _RESERVED_TAG_KEYS for the rationale. if self.cost_tag_keys: - allow = set(self.cost_tag_keys) + allow: Final = set(self.cost_tag_keys) for rt in log.get("request_tags") or []: if not isinstance(rt, str) or ":" not in rt: continue @@ -220,7 +220,7 @@ class DatadogCostManagementLogger(CustomBatchLogger): return tags @staticmethod - def _set_custom_tag(tags: Dict[str, str], key: str, value: str) -> None: + def _set_custom_tag(tags: dict[str, str], key: str, value: str) -> None: if key in _RESERVED_TAG_KEYS: verbose_logger.debug( "Datadog Cost Management: dropping user-supplied tag %r=%r — " @@ -232,27 +232,27 @@ class DatadogCostManagementLogger(CustomBatchLogger): tags[key] = value @staticmethod - def _add_tag(tags: Dict[str, str], key: str, value: Any) -> None: + def _add_tag(tags: dict[str, str], key: str, value: Any) -> None: if value: tags[key] = str(value) - async def _upload_to_datadog(self, payload: List[Dict]): + async def _upload_to_datadog(self, payload: list[dict]): if not self.dd_api_key or not self.dd_app_key: return - headers = { + headers: Final = { "Content-Type": "application/json", "DD-API-KEY": self.dd_api_key, "DD-APPLICATION-KEY": self.dd_app_key, } # The API endpoint expects a list of objects directly in the body (file content behavior) - data_json = safe_dumps(payload) + data_json: Final = safe_dumps(payload) - response = await self.async_client.put(self.upload_url, content=data_json, headers=headers) + response: Final = await self.async_client.put(self.upload_url, content=data_json, headers=headers) response.raise_for_status() verbose_logger.debug( - f"Datadog Cost Management: Uploaded {len(payload)} cost entries. Status: {response.status_code}" + "Datadog Cost Management: Uploaded %s cost entries. Status: %s", len(payload), response.status_code ) diff --git a/litellm/integrations/datadog/datadog_handler.py b/litellm/integrations/datadog/datadog_handler.py index b6bb2b57037..2450382a192 100644 --- a/litellm/integrations/datadog/datadog_handler.py +++ b/litellm/integrations/datadog/datadog_handler.py @@ -3,7 +3,7 @@ from __future__ import annotations import os -from typing import List, Optional +from typing import Final from litellm.types.utils import StandardLoggingPayload @@ -20,7 +20,7 @@ def get_datadog_hostname() -> str: return os.getenv("HOSTNAME", "") -def get_datadog_base_url_from_env() -> Optional[str]: +def get_datadog_base_url_from_env() -> str | None: """ Get base URL override from common DD_BASE_URL env var. This is useful for testing or custom endpoints. @@ -37,8 +37,8 @@ def get_datadog_pod_name() -> str: def get_datadog_tags( - standard_logging_object: Optional[StandardLoggingPayload] = None, -) -> List[str]: + standard_logging_object: StandardLoggingPayload | None = None, +) -> list[str]: """Build Datadog tags as a list of individual tag strings. Returns a list of "key:value" strings suitable for Datadog LLM Observability @@ -46,7 +46,7 @@ def get_datadog_tags( comma: ",".join(get_datadog_tags(...)). """ - base_tags = { + base_tags: Final = { "env": get_datadog_env(), "service": get_datadog_service(), "version": os.getenv("DD_VERSION", "unknown"), @@ -54,15 +54,15 @@ def get_datadog_tags( "POD_NAME": get_datadog_pod_name(), } - tags: List[str] = [f"{k}:{v}" for k, v in base_tags.items()] + tags: Final[list[str]] = [f"{k}:{v}" for k, v in base_tags.items()] if standard_logging_object: - request_tags = standard_logging_object.get("request_tags", []) or [] + request_tags: Final = standard_logging_object.get("request_tags", []) or [] tags.extend(f"request_tag:{tag}" for tag in request_tags) # Add Team Tag - metadata = standard_logging_object.get("metadata", {}) or {} - team_tag = ( + metadata: Final = standard_logging_object.get("metadata", {}) or {} + team_tag: Final = ( metadata.get("user_api_key_team_alias") or metadata.get("team_alias") or metadata.get("user_api_key_team_id") diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 1078f05165a..704f0323e95 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -9,23 +9,23 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp import asyncio import json import os -from litellm._uuid import uuid from datetime import datetime -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Final, Literal import httpx import litellm from litellm._logging import verbose_logger +from litellm._uuid import uuid from litellm.integrations.custom_batch_logger import CustomBatchLogger -from litellm.integrations.datadog.datadog_mock_client import ( - should_use_datadog_mock, - create_mock_datadog_client, -) from litellm.integrations.datadog.datadog_handler import ( + get_datadog_base_url_from_env, get_datadog_service, get_datadog_tags, - get_datadog_base_url_from_env, +) +from litellm.integrations.datadog.datadog_mock_client import ( + create_mock_datadog_client, + should_use_datadog_mock, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -58,7 +58,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Configure DataDog endpoint (Agent or Direct API) # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST # Check for agent mode FIRST - agent mode doesn't require DD_API_KEY or DD_SITE - dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") + dd_agent_host: Final = os.getenv("LITELLM_DD_AGENT_HOST") self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) self.DD_API_KEY = os.getenv("DD_API_KEY") @@ -74,22 +74,22 @@ class DataDogLLMObsLogger(CustomBatchLogger): self._configure_dd_direct_api() # Optional override for testing - dd_base_url = get_datadog_base_url_from_env() + dd_base_url: Final = get_datadog_base_url_from_env() if dd_base_url: self.intake_url = f"{dd_base_url}/api/intake/llm-obs/v1/trace/spans" asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() - self.log_queue: List[LLMObsPayload] = [] + self.log_queue: list[LLMObsPayload] = [] ######################################################### # Handle datadog_llm_observability_params set as litellm.datadog_llm_observability_params ######################################################### - dict_datadog_llm_obs_params = self._get_datadog_llm_obs_params() + dict_datadog_llm_obs_params: Final = self._get_datadog_llm_obs_params() kwargs.update(dict_datadog_llm_obs_params) CustomBatchLogger.__init__(self, **kwargs, flush_lock=self.flush_lock) except Exception as e: - verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}") + verbose_logger.exception("DataDogLLMObs: Error initializing - %s", e) raise e def _configure_dd_agent(self, dd_agent_host: str): @@ -100,10 +100,10 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Reference: https://docs.datadoghq.com/llm_observability/setup/sdk/#agent-setup # Use specific port for LLM Obs (Trace Agent) to avoid conflict with Logs Agent (10518) - agent_port = os.getenv("LITELLM_DD_LLM_OBS_PORT", "8126") + agent_port: Final = os.getenv("LITELLM_DD_LLM_OBS_PORT", "8126") self.DD_SITE = "localhost" # Not used for URL construction in agent mode self.intake_url = f"http://{dd_agent_host}:{agent_port}/api/intake/llm-obs/v1/trace/spans" - verbose_logger.debug(f"DataDogLLMObs: Using DD Agent at {self.intake_url}") + verbose_logger.debug("DataDogLLMObs: Using DD Agent at %s", self.intake_url) def _configure_dd_direct_api(self): """ @@ -118,17 +118,17 @@ class DataDogLLMObsLogger(CustomBatchLogger): self.intake_url = f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans" - def _get_datadog_llm_obs_params(self) -> Dict: + def _get_datadog_llm_obs_params(self) -> dict: """ Get the datadog_llm_observability_params from litellm.datadog_llm_observability_params These are params specific to initializing the DataDogLLMObsLogger e.g. turn_off_message_logging """ - dict_datadog_llm_obs_params: Dict = {} + dict_datadog_llm_obs_params: dict = {} if litellm.datadog_llm_observability_params is not None: if isinstance(litellm.datadog_llm_observability_params, DatadogLLMObsInitParams): dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump() - elif isinstance(litellm.datadog_llm_observability_params, Dict): + elif isinstance(litellm.datadog_llm_observability_params, dict): # only allow params that are of DatadogLLMObsInitParams dict_datadog_llm_obs_params = DatadogLLMObsInitParams( **litellm.datadog_llm_observability_params @@ -137,40 +137,40 @@ class DataDogLLMObsLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug(f"DataDogLLMObs: Logging success event for model {kwargs.get('model', 'unknown')}") - payload = self.create_llm_obs_payload(kwargs, start_time, end_time) - verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}") + verbose_logger.debug("DataDogLLMObs: Logging success event for model %s", kwargs.get("model", "unknown")) + payload: Final = self.create_llm_obs_payload(kwargs, start_time, end_time) + verbose_logger.debug("DataDogLLMObs: Payload: %s", payload) self.log_queue.append(payload) if len(self.log_queue) >= self.batch_size: await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"DataDogLLMObs: Error logging success event - {str(e)}") + verbose_logger.exception("DataDogLLMObs: Error logging success event - %s", e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: - verbose_logger.debug(f"DataDogLLMObs: Logging failure event for model {kwargs.get('model', 'unknown')}") - payload = self.create_llm_obs_payload(kwargs, start_time, end_time) - verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}") + verbose_logger.debug("DataDogLLMObs: Logging failure event for model %s", kwargs.get("model", "unknown")) + payload: Final = self.create_llm_obs_payload(kwargs, start_time, end_time) + verbose_logger.debug("DataDogLLMObs: Payload: %s", payload) self.log_queue.append(payload) if len(self.log_queue) >= self.batch_size: await self.async_send_batch() except Exception as e: - verbose_logger.exception(f"DataDogLLMObs: Error logging failure event - {str(e)}") + verbose_logger.exception("DataDogLLMObs: Error logging failure event - %s", e) async def async_send_batch(self): try: if not self.log_queue: return - verbose_logger.debug(f"DataDogLLMObs: Flushing {len(self.log_queue)} events") + verbose_logger.debug("DataDogLLMObs: Flushing %s events", len(self.log_queue)) if self.is_mock_mode: verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted") # Prepare the payload - payload = { + payload: Final = { "data": DDIntakePayload( type="span", attributes=DDSpanAttributes( @@ -189,13 +189,13 @@ class DataDogLLMObsLogger(CustomBatchLogger): except Exception as debug_error: verbose_logger.debug("payload serialization failed: %s", str(debug_error)) - json_payload = safe_dumps(payload) + json_payload: Final = safe_dumps(payload) - headers = {"Content-Type": "application/json"} + headers: Final = {"Content-Type": "application/json"} if self.DD_API_KEY: headers["DD-API-KEY"] = self.DD_API_KEY - response = await self.async_client.post( + response: Final = await self.async_client.post( url=self.intake_url, content=json_payload, headers=headers, @@ -207,40 +207,40 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) if self.is_mock_mode: - verbose_logger.debug(f"[DATADOG MOCK] Batch of {len(self.log_queue)} events successfully mocked") + verbose_logger.debug("[DATADOG MOCK] Batch of %s events successfully mocked", len(self.log_queue)) else: - verbose_logger.debug(f"DataDogLLMObs: Successfully sent batch - status_code: {response.status_code}") + verbose_logger.debug("DataDogLLMObs: Successfully sent batch - status_code: %s", response.status_code) self.log_queue.clear() except httpx.HTTPStatusError as e: - verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {e.response.text}") + verbose_logger.exception("DataDogLLMObs: Error sending batch - %s", e.response.text) except Exception as e: - verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {str(e)}") + verbose_logger.exception("DataDogLLMObs: Error sending batch - %s", e) - def create_llm_obs_payload(self, kwargs: Dict, start_time: datetime, end_time: datetime) -> LLMObsPayload: - standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object") + def create_llm_obs_payload(self, kwargs: dict, start_time: datetime, end_time: datetime) -> LLMObsPayload: + standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") if standard_logging_payload is None: raise Exception("DataDogLLMObs: standard_logging_object is not set") messages = standard_logging_payload["messages"] messages = self._ensure_string_content(messages=messages) - metadata = kwargs.get("litellm_params", {}).get("metadata", {}) + metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {}) - input_meta = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages)) - output_meta = OutputMeta( + input_meta: Final = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages)) + output_meta: Final = OutputMeta( messages=self._get_response_messages( standard_logging_payload=standard_logging_payload, call_type=standard_logging_payload.get("call_type"), ) ) - error_info = self._assemble_error_info(standard_logging_payload) + error_info: Final = self._assemble_error_info(standard_logging_payload) - metadata_parent_id: Optional[str] = None + metadata_parent_id: str | None = None if isinstance(metadata, dict): metadata_parent_id = metadata.get("parent_id") - meta = Meta( + meta: Final = Meta( kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id), input=input_meta, output=output_meta, @@ -249,7 +249,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) # Calculate metrics (you may need to adjust these based on available data) - metrics = LLMMetrics( + metrics: Final = LLMMetrics( input_tokens=float(standard_logging_payload.get("prompt_tokens", 0)), output_tokens=float(standard_logging_payload.get("completion_tokens", 0)), total_tokens=float(standard_logging_payload.get("total_tokens", 0)), @@ -257,7 +257,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload), ) - payload: LLMObsPayload = LLMObsPayload( + payload: Final[LLMObsPayload] = LLMObsPayload( parent_id=metadata_parent_id if metadata_parent_id else "undefined", trace_id=standard_logging_payload.get("trace_id", str(uuid.uuid4())), span_id=metadata.get("span_id", str(uuid.uuid4())), @@ -270,36 +270,36 @@ class DataDogLLMObsLogger(CustomBatchLogger): tags=get_datadog_tags(standard_logging_object=standard_logging_payload), ) - apm_trace_id = self._get_apm_trace_id() + apm_trace_id: Final = self._get_apm_trace_id() if apm_trace_id is not None: payload["apm_id"] = apm_trace_id return payload - def _get_apm_trace_id(self) -> Optional[str]: + def _get_apm_trace_id(self) -> str | None: """Retrieve the current APM trace ID if available.""" try: - current_span_fn = getattr(tracer, "current_span", None) + current_span_fn: Final = getattr(tracer, "current_span", None) if callable(current_span_fn): - current_span = current_span_fn() + current_span: Final = current_span_fn() if current_span is not None: - trace_id = getattr(current_span, "trace_id", None) + trace_id: Final = getattr(current_span, "trace_id", None) if trace_id is not None: return str(trace_id) except Exception: pass return None - def _assemble_error_info(self, standard_logging_payload: StandardLoggingPayload) -> Optional[DDLLMObsError]: + def _assemble_error_info(self, standard_logging_payload: StandardLoggingPayload) -> DDLLMObsError | None: """ Assemble error information for failure cases according to DD LLM Obs API spec """ # Handle error information for failure cases according to DD LLM Obs API spec - error_info: Optional[DDLLMObsError] = None + error_info: DDLLMObsError | None = None if standard_logging_payload.get("status") == "failure": # Try to get structured error information first - error_information: Optional[StandardLoggingPayloadErrorInformation] = standard_logging_payload.get( + error_information: Final[StandardLoggingPayloadErrorInformation | None] = standard_logging_payload.get( "error_information" ) @@ -321,9 +321,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): For non streaming calls, CompletionStartTime is time we get the response back """ - start_time: Optional[float] = standard_logging_payload.get("startTime") - completion_start_time: Optional[float] = standard_logging_payload.get("completionStartTime") - end_time: Optional[float] = standard_logging_payload.get("endTime") + start_time: Final[float | None] = standard_logging_payload.get("startTime") + completion_start_time: Final[float | None] = standard_logging_payload.get("completionStartTime") + end_time: Final[float | None] = standard_logging_payload.get("endTime") if completion_start_time is not None and start_time is not None: return completion_start_time - start_time @@ -333,8 +333,8 @@ class DataDogLLMObsLogger(CustomBatchLogger): return 0.0 def _get_response_messages( - self, standard_logging_payload: StandardLoggingPayload, call_type: Optional[str] - ) -> List[Any]: + self, standard_logging_payload: StandardLoggingPayload, call_type: str | None + ) -> list[Any]: """ Get the messages from the response object @@ -372,7 +372,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): try: # Safely extract message from response_obj, handle failure cases if isinstance(response_obj, dict) and "choices" in response_obj: - choices = response_obj["choices"] + choices: Final = response_obj["choices"] if choices and len(choices) > 0 and "message" in choices[0]: return [choices[0]["message"]] return [] @@ -382,7 +382,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): return [] def _get_datadog_span_kind( - self, call_type: Optional[str], parent_id: Optional[str] = None + self, call_type: str | None, parent_id: str | None = None ) -> Literal["llm", "tool", "task", "embedding", "retrieval"]: """ Map liteLLM call_type to appropriate DataDog LLM Observability span kind. @@ -484,7 +484,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): # Default fallback for unknown or passthrough operations return "llm" - def _ensure_string_content(self, messages: Optional[Union[str, List[Any], Dict[Any, Any]]]) -> List[Any]: + def _ensure_string_content(self, messages: str | list[Any] | dict[Any, Any] | None) -> list[Any]: if messages is None: return [] if isinstance(messages, str): @@ -495,11 +495,11 @@ class DataDogLLMObsLogger(CustomBatchLogger): return [str(messages.get("content", ""))] return [] - def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> Dict[str, Any]: + def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]: """ Fields to track in DD LLM Observability metadata from litellm standard logging payload """ - _metadata: Dict[str, Any] = { + _metadata: Final[dict[str, Any]] = { "model_name": standard_logging_payload.get("model", "unknown"), "model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"), "id": standard_logging_payload.get("id", "unknown"), @@ -514,20 +514,20 @@ class DataDogLLMObsLogger(CustomBatchLogger): ######################################################### # Add latency metrics to metadata ######################################################### - latency_metrics = self._get_latency_metrics(standard_logging_payload) + latency_metrics: Final = self._get_latency_metrics(standard_logging_payload) _metadata.update({"latency_metrics": dict(latency_metrics)}) ######################################################### # Add spend metrics to metadata ######################################################### - spend_metrics = self._get_spend_metrics(standard_logging_payload) + spend_metrics: Final = self._get_spend_metrics(standard_logging_payload) _metadata.update({"spend_metrics": dict(spend_metrics)}) ## extract tool calls and add to metadata - tool_call_metadata = self._extract_tool_call_metadata(standard_logging_payload) + tool_call_metadata: Final = self._extract_tool_call_metadata(standard_logging_payload) _metadata.update(tool_call_metadata) - _standard_logging_metadata: dict = dict(standard_logging_payload.get("metadata", {})) or {} + _standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {} _metadata.update(_standard_logging_metadata) return _metadata @@ -535,27 +535,27 @@ class DataDogLLMObsLogger(CustomBatchLogger): """ Get the latency metrics from the standard logging payload """ - latency_metrics: DDLLMObsLatencyMetrics = DDLLMObsLatencyMetrics() + latency_metrics: Final[DDLLMObsLatencyMetrics] = DDLLMObsLatencyMetrics() # Add latency metrics to metadata # Time to first token (convert from seconds to milliseconds for consistency) - time_to_first_token_seconds = self._get_time_to_first_token_seconds(standard_logging_payload) + time_to_first_token_seconds: Final = self._get_time_to_first_token_seconds(standard_logging_payload) if time_to_first_token_seconds > 0: latency_metrics["time_to_first_token_ms"] = time_to_first_token_seconds * 1000 # LiteLLM overhead time - hidden_params = standard_logging_payload.get("hidden_params", {}) - litellm_overhead_ms = hidden_params.get("litellm_overhead_time_ms") + hidden_params: Final = standard_logging_payload.get("hidden_params", {}) + litellm_overhead_ms: Final = hidden_params.get("litellm_overhead_time_ms") if litellm_overhead_ms is not None: latency_metrics["litellm_overhead_time_ms"] = litellm_overhead_ms # Guardrail overhead latency - guardrail_info: Optional[list[StandardLoggingGuardrailInformation]] = standard_logging_payload.get( + guardrail_info: Final[list[StandardLoggingGuardrailInformation] | None] = standard_logging_payload.get( "guardrail_information" ) if guardrail_info is not None: total_duration = 0.0 for info in guardrail_info: - _guardrail_duration_seconds: Optional[float] = info.get("duration") + _guardrail_duration_seconds: float | None = info.get("duration") if _guardrail_duration_seconds is not None: total_duration += float(_guardrail_duration_seconds) @@ -581,7 +581,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): return True # Fallback to model_parameters.stream for original request parameters - model_params = standard_logging_payload.get("model_parameters", {}) + model_params: Final = standard_logging_payload.get("model_parameters", {}) if isinstance(model_params, dict): stream_value = model_params.get("stream") if stream_value is True: @@ -594,29 +594,29 @@ class DataDogLLMObsLogger(CustomBatchLogger): """ Get the spend metrics from the standard logging payload """ - spend_metrics: DDLLMObsSpendMetrics = DDLLMObsSpendMetrics() + spend_metrics: Final[DDLLMObsSpendMetrics] = DDLLMObsSpendMetrics() # send response cost spend_metrics["response_cost"] = standard_logging_payload.get("response_cost", 0.0) # Get budget information from metadata - metadata = standard_logging_payload.get("metadata", {}) + metadata: Final = standard_logging_payload.get("metadata", {}) # API key max budget - user_api_key_max_budget = metadata.get("user_api_key_max_budget") + user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget") if user_api_key_max_budget is not None: spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget) # API key spend - user_api_key_spend = metadata.get("user_api_key_spend") + user_api_key_spend: Final = metadata.get("user_api_key_spend") if user_api_key_spend is not None: try: spend_metrics["user_api_key_spend"] = float(user_api_key_spend) except (ValueError, TypeError): - verbose_logger.debug(f"Invalid user_api_key_spend value: {user_api_key_spend}") + verbose_logger.debug("Invalid user_api_key_spend value: %s", user_api_key_spend) # API key budget reset datetime - user_api_key_budget_reset_at = metadata.get("user_api_key_budget_reset_at") + user_api_key_budget_reset_at: Final = metadata.get("user_api_key_budget_reset_at") if user_api_key_budget_reset_at is not None: try: from datetime import datetime, timezone @@ -640,21 +640,21 @@ class DataDogLLMObsLogger(CustomBatchLogger): spend_metrics["user_api_key_budget_reset_at"] = iso_string # Debug logging to verify the conversion - verbose_logger.debug(f"Converted budget_reset_at to ISO format: {iso_string}") + verbose_logger.debug("Converted budget_reset_at to ISO format: %s", iso_string) except Exception as e: - verbose_logger.debug(f"Error processing budget reset datetime: {e}") - verbose_logger.debug(f"Original value: {user_api_key_budget_reset_at}") + verbose_logger.debug("Error processing budget reset datetime: %s", e) + verbose_logger.debug("Original value: %s", user_api_key_budget_reset_at) return spend_metrics - def _process_input_messages_preserving_tool_calls(self, messages: List[Any]) -> List[Dict[str, Any]]: + def _process_input_messages_preserving_tool_calls(self, messages: list[Any]) -> list[dict[str, Any]]: """ Process input messages while preserving tool_calls and tool message types. This bypasses the lossy string conversion when tool calls are present, allowing complex nested tool_calls objects to be preserved for Datadog. """ - processed = [] + processed: Final = [] for msg in messages: if isinstance(msg, dict): # Preserve messages with tool_calls or tool role as-is @@ -671,13 +671,13 @@ class DataDogLLMObsLogger(CustomBatchLogger): return processed @staticmethod - def _tool_calls_kv_pair(tool_calls: List[Dict[str, Any]]) -> Dict[str, Any]: + def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, Any]: """ Extract tool call information into key-value pairs for Datadog metadata. Similar to OpenTelemetry's implementation but adapted for Datadog's format. """ - kv_pairs: Dict[str, Any] = {} + kv_pairs: Final[dict[str, Any]] = {} for idx, tool_call in enumerate(tool_calls): try: # Extract tool call ID @@ -707,20 +707,20 @@ class DataDogLLMObsLogger(CustomBatchLogger): kv_pairs[f"tool_calls.{idx}.function.arguments"] = json.dumps(function_arguments) except (KeyError, TypeError, ValueError) as e: - verbose_logger.debug(f"DataDogLLMObs: Error processing tool call {idx}: {str(e)}") + verbose_logger.debug("DataDogLLMObs: Error processing tool call %s: %s", idx, e) continue return kv_pairs - def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> Dict[str, Any]: + def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]: """ Extract tool call information from both input messages and response for Datadog metadata. """ - tool_call_metadata: Dict[str, Any] = {} + tool_call_metadata: Final[dict[str, Any]] = {} try: # Extract tool calls from input messages - messages = standard_logging_payload.get("messages", []) + messages: Final = standard_logging_payload.get("messages", []) if messages and isinstance(messages, list): for message in messages: if isinstance(message, dict) and "tool_calls" in message: @@ -732,9 +732,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): tool_call_metadata[f"input_{key}"] = value # Extract tool calls from response - response_obj = standard_logging_payload.get("response") + response_obj: Final = standard_logging_payload.get("response") if response_obj and isinstance(response_obj, dict): - choices = response_obj.get("choices", []) + choices: Final = response_obj.get("choices", []) for choice in choices: if isinstance(choice, dict): message = choice.get("message") @@ -747,6 +747,6 @@ class DataDogLLMObsLogger(CustomBatchLogger): tool_call_metadata[f"output_{key}"] = value except Exception as e: - verbose_logger.debug(f"DataDogLLMObs: Error extracting tool call metadata: {str(e)}") + verbose_logger.debug("DataDogLLMObs: Error extracting tool call metadata: %s", e) return tool_call_metadata diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index b1e4bc73e77..37421126985 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -3,7 +3,7 @@ import gzip import os import time from datetime import datetime -from typing import List, Optional, Union +from typing import Final from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -60,13 +60,13 @@ class DatadogMetricsLogger(CustomBatchLogger): def _extract_tags( self, log: StandardLoggingPayload, - status_code: Optional[Union[str, int]] = None, - ) -> List[str]: + status_code: str | int | None = None, + ) -> list[str]: """ Builds the list of tags for a Datadog metric point """ # Base tags - tags = [ + tags: Final = [ f"env:{get_datadog_env()}", f"service:{get_datadog_service()}", f"version:{os.getenv('DD_VERSION', 'unknown')}", @@ -88,8 +88,8 @@ class DatadogMetricsLogger(CustomBatchLogger): tags.append(f"status_code:{status_code}") # Extract team tag - metadata = log.get("metadata", {}) or {} - team_tag = ( + metadata: Final = log.get("metadata", {}) or {} + team_tag: Final = ( metadata.get("user_api_key_team_alias") or metadata.get("team_alias") # type: ignore or metadata.get("user_api_key_team_id") @@ -105,22 +105,22 @@ class DatadogMetricsLogger(CustomBatchLogger): self, log: StandardLoggingPayload, kwargs: dict, - status_code: Union[str, int] = "200", + status_code: str | int = "200", ): """ Extracts latencies and appends Datadog metric series to the queue """ - tags = self._extract_tags(log, status_code=status_code) + tags: Final = self._extract_tags(log, status_code=status_code) # We record metrics with the end_time as the timestamp for the point - end_time_dt = kwargs.get("end_time") or datetime.now() - timestamp = int(end_time_dt.timestamp()) + end_time_dt: Final = kwargs.get("end_time") or datetime.now() + timestamp: Final = int(end_time_dt.timestamp()) # 1. Total Request Latency Metric (End to End) - start_time_dt = kwargs.get("start_time") + start_time_dt: Final = kwargs.get("start_time") if start_time_dt and end_time_dt: - total_duration = (end_time_dt - start_time_dt).total_seconds() - series_total_latency: DatadogMetricSeries = { + total_duration: Final = (end_time_dt - start_time_dt).total_seconds() + series_total_latency: Final[DatadogMetricSeries] = { "metric": "litellm.request.total_latency", "type": 3, # gauge "points": [{"timestamp": timestamp, "value": total_duration}], @@ -129,10 +129,10 @@ class DatadogMetricsLogger(CustomBatchLogger): self.log_queue.append(series_total_latency) # 2. LLM API Latency Metric (Provider alone) - api_call_start_time = kwargs.get("api_call_start_time") + api_call_start_time: Final = kwargs.get("api_call_start_time") if api_call_start_time and end_time_dt: - llm_api_duration = (end_time_dt - api_call_start_time).total_seconds() - series_llm_latency: DatadogMetricSeries = { + llm_api_duration: Final = (end_time_dt - api_call_start_time).total_seconds() + series_llm_latency: Final[DatadogMetricSeries] = { "metric": "litellm.llm_api.latency", "type": 3, # gauge "points": [{"timestamp": timestamp, "value": llm_api_duration}], @@ -141,11 +141,11 @@ class DatadogMetricsLogger(CustomBatchLogger): self.log_queue.append(series_llm_latency) # 3. LiteLLM Overhead Latency Metric (total - llm_api time) - hidden_params = log.get("hidden_params", {}) or {} - litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms") + hidden_params: Final = log.get("hidden_params", {}) or {} + litellm_overhead_time_ms: Final = hidden_params.get("litellm_overhead_time_ms") if litellm_overhead_time_ms is not None: - overhead_tags = self._extract_tags(log) # no status_code on latency metric - series_overhead: DatadogMetricSeries = { + overhead_tags: Final = self._extract_tags(log) # no status_code on latency metric + series_overhead: Final[DatadogMetricSeries] = { "metric": "litellm.overhead.latency", "type": 3, # gauge "points": [ @@ -159,7 +159,7 @@ class DatadogMetricsLogger(CustomBatchLogger): self.log_queue.append(series_overhead) # 4. Request Count / Status Code - series_count: DatadogMetricSeries = { + series_count: Final[DatadogMetricSeries] = { "metric": "litellm.llm_api.request_count", "type": 1, # count "points": [{"timestamp": timestamp, "value": 1.0}], @@ -170,7 +170,7 @@ class DatadogMetricsLogger(CustomBatchLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) + standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: return @@ -181,19 +181,19 @@ class DatadogMetricsLogger(CustomBatchLogger): await self.flush_queue() except Exception as e: - verbose_logger.exception(f"Datadog Metrics: Error in async_log_success_event: {str(e)}") + verbose_logger.exception("Datadog Metrics: Error in async_log_success_event: %s", e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): try: - standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None) + standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: return # Extract status code from error information status_code = "500" # default - error_information = standard_logging_object.get("error_information", {}) or {} - error_code = error_information.get("error_code") # type: ignore + error_information: Final = standard_logging_object.get("error_information", {}) or {} + error_code: Final = error_information.get("error_code") # type: ignore if error_code is not None: status_code = str(error_code) @@ -203,26 +203,26 @@ class DatadogMetricsLogger(CustomBatchLogger): await self.flush_queue() except Exception as e: - verbose_logger.exception(f"Datadog Metrics: Error in async_log_failure_event: {str(e)}") + verbose_logger.exception("Datadog Metrics: Error in async_log_failure_event: %s", e) async def async_send_batch(self): if not self.log_queue: return - batch = self.log_queue.copy() - payload_data: DatadogMetricsPayload = {"series": batch} + batch: Final = self.log_queue.copy() + payload_data: Final[DatadogMetricsPayload] = {"series": batch} try: await self._upload_to_datadog(payload_data) except Exception as e: - verbose_logger.exception(f"Datadog Metrics: Error in async_send_batch: {str(e)}") + verbose_logger.exception("Datadog Metrics: Error in async_send_batch: %s", e) raise async def _upload_to_datadog(self, payload: DatadogMetricsPayload): if not self.dd_api_key: return - headers = { + headers: Final = { "Content-Type": "application/json", "DD-API-KEY": self.dd_api_key, } @@ -230,11 +230,11 @@ class DatadogMetricsLogger(CustomBatchLogger): if self.dd_app_key: headers["DD-APPLICATION-KEY"] = self.dd_app_key - json_data = safe_dumps(payload) - compressed_data = gzip.compress(json_data.encode("utf-8")) + json_data: Final = safe_dumps(payload) + compressed_data: Final = gzip.compress(json_data.encode("utf-8")) headers["Content-Encoding"] = "gzip" - response = await self.async_client.post( + response: Final = await self.async_client.post( self.upload_url, content=compressed_data, headers=headers, # type: ignore @@ -243,7 +243,7 @@ class DatadogMetricsLogger(CustomBatchLogger): response.raise_for_status() verbose_logger.debug( - f"Datadog Metrics: Uploaded {len(payload['series'])} metric points. Status: {response.status_code}" + "Datadog Metrics: Uploaded %s metric points. Status: %s", len(payload["series"]), response.status_code ) async def async_health_check(self) -> IntegrationHealthCheckStatus: @@ -252,18 +252,18 @@ class DatadogMetricsLogger(CustomBatchLogger): """ try: # Send a test metric point to Datadog - test_metric_point: DatadogMetricPoint = { + test_metric_point: Final[DatadogMetricPoint] = { "timestamp": int(time.time()), "value": 1.0, } - test_metric_series: DatadogMetricSeries = { + test_metric_series: Final[DatadogMetricSeries] = { "metric": "litellm.health_check", "type": 3, # Gauge "points": [test_metric_point], "tags": ["env:health_check"], } - payload_data: DatadogMetricsPayload = {"series": [test_metric_series]} + payload_data: Final[DatadogMetricsPayload] = {"series": [test_metric_series]} await self._upload_to_datadog(payload_data) @@ -280,7 +280,7 @@ class DatadogMetricsLogger(CustomBatchLogger): async def get_request_response_payload( self, request_id: str, - start_time_utc: Optional[datetime], - end_time_utc: Optional[datetime], - ) -> Optional[dict]: + start_time_utc: datetime | None, + end_time_utc: datetime | None, + ) -> dict | None: pass diff --git a/litellm/integrations/datadog/datadog_mock_client.py b/litellm/integrations/datadog/datadog_mock_client.py index c50cdc6a019..a90ffbc6512 100644 --- a/litellm/integrations/datadog/datadog_mock_client.py +++ b/litellm/integrations/datadog/datadog_mock_client.py @@ -8,13 +8,15 @@ Usage: Set DATADOG_MOCK=true in environment variables or config to enable mock mode. """ +from typing import Final + from litellm.integrations.mock_client_factory import ( MockClientConfig, create_mock_client_factory, ) # Create mock client using factory -_config = MockClientConfig( +_config: Final = MockClientConfig( name="DATADOG", env_var="DATADOG_MOCK", default_latency_ms=100, diff --git a/litellm/integrations/datadog/datadog_team_handler.py b/litellm/integrations/datadog/datadog_team_handler.py index cae954f753c..dd848108968 100644 --- a/litellm/integrations/datadog/datadog_team_handler.py +++ b/litellm/integrations/datadog/datadog_team_handler.py @@ -5,7 +5,7 @@ Used to get the DataDogLogger for a given request. Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler. """ -from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict +from typing import TYPE_CHECKING, Any, Final, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams @@ -19,10 +19,10 @@ else: class DatadogLoggingConfig(TypedDict): - dd_api_key: Optional[str] - dd_site: Optional[str] - dd_agent_host: Optional[str] - dd_agent_port: Optional[str] + dd_api_key: str | None + dd_site: str | None + dd_agent_host: str | None + dd_agent_port: str | None class DataDogHandler: @@ -42,10 +42,10 @@ class DataDogHandler: The global (env-var based) DataDogLogger is managed separately by _init_custom_logger_compatible_class via _in_memory_loggers. """ - _credentials = DataDogHandler.get_dynamic_datadog_logging_config( + _credentials: Final = DataDogHandler.get_dynamic_datadog_logging_config( standard_callback_dynamic_params=standard_callback_dynamic_params, ) - credentials_dict = dict(_credentials) + credentials_dict: Final = dict(_credentials) # check if datadog logger is already cached temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache( @@ -63,7 +63,7 @@ class DataDogHandler: @staticmethod def _create_datadog_logger_from_credentials( - credentials: Dict, + credentials: dict, in_memory_dynamic_logger_cache: DynamicLoggingCache, ) -> DataDogLogger: """ @@ -71,8 +71,8 @@ class DataDogHandler: """ # When the destination is caller-supplied (dd_agent_host/dd_site), never fall back to the # proxy's DD_API_KEY env var, otherwise it would be sent to a team-controlled host. - allow_env_credentials = credentials.get("dd_agent_host") is None and credentials.get("dd_site") is None - datadog_logger = DataDogLogger( + allow_env_credentials: Final = credentials.get("dd_agent_host") is None and credentials.get("dd_site") is None + datadog_logger: Final = DataDogLogger( dd_api_key=credentials.get("dd_api_key"), dd_site=credentials.get("dd_site"), dd_agent_host=credentials.get("dd_agent_host"), diff --git a/litellm/integrations/deepeval/api.py b/litellm/integrations/deepeval/api.py index fccc5970433..60639fe2941 100644 --- a/litellm/integrations/deepeval/api.py +++ b/litellm/integrations/deepeval/api.py @@ -1,14 +1,17 @@ # duplicate -> https://github.com/confident-ai/deepeval/blob/main/deepeval/confident/api.py import logging -import httpx from enum import Enum +from typing import Final + +import httpx + from litellm._logging import verbose_logger -DEEPEVAL_BASE_URL = "https://deepeval.confident-ai.com" -DEEPEVAL_BASE_URL_EU = "https://eu.deepeval.confident-ai.com" -API_BASE_URL = "https://api.confident-ai.com" -API_BASE_URL_EU = "https://eu.api.confident-ai.com" -retryable_exceptions = httpx.HTTPError +DEEPEVAL_BASE_URL: Final = "https://deepeval.confident-ai.com" +DEEPEVAL_BASE_URL_EU: Final = "https://eu.deepeval.confident-ai.com" +API_BASE_URL: Final = "https://api.confident-ai.com" +API_BASE_URL_EU: Final = "https://eu.api.confident-ai.com" +retryable_exceptions: Final = httpx.HTTPError from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, @@ -18,12 +21,12 @@ from litellm.llms.custom_httpx.http_handler import ( def log_retry_error(details): - exception = details.get("exception") - tries = details.get("tries") + exception: Final = details.get("exception") + tries: Final = details.get("tries") if exception: - logging.error(f"Confident AI Error: {exception}. Retrying: {tries} time(s)...") + logging.error("Confident AI Error: %s. Retrying: %s time(s)...", exception, tries) else: - logging.error(f"Retrying: {tries} time(s)...") + logging.error("Retrying: %s time(s)...", tries) class HttpMethods(Enum): @@ -76,8 +79,8 @@ class Api: raise e def send_request(self, method: HttpMethods, endpoint: Endpoints, body=None, params=None): - url = f"{self.base_api_url}{endpoint.value}" - res = self._http_request( + url: Final = f"{self.base_api_url}{endpoint.value}" + res: Final = self._http_request( method=method.value, url=url, headers=self._headers, @@ -98,7 +101,7 @@ class Api: if method != HttpMethods.POST: raise Exception("Only POST requests are supported") - url = f"{self.base_api_url}{endpoint.value}" + url: Final = f"{self.base_api_url}{endpoint.value}" try: await self.async_http_handler.post( url=url, diff --git a/litellm/integrations/deepeval/deepeval.py b/litellm/integrations/deepeval/deepeval.py index 90c1d8eedce..71d149e4014 100644 --- a/litellm/integrations/deepeval/deepeval.py +++ b/litellm/integrations/deepeval/deepeval.py @@ -1,4 +1,7 @@ import os +from typing import Final + +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.deepeval.api import Api, Endpoints, HttpMethods @@ -12,7 +15,6 @@ from litellm.integrations.deepeval.utils import ( to_zod_compatible_iso, validate_environment, ) -from litellm._logging import verbose_logger # This file includes the custom callbacks for LiteLLM Proxy @@ -21,7 +23,7 @@ class DeepEvalLogger(CustomLogger): """Logs litellm traces to DeepEval's platform.""" def __init__(self, *args, **kwargs): - api_key = os.getenv("CONFIDENT_API_KEY") + api_key: Final = os.getenv("CONFIDENT_API_KEY") self.litellm_environment = os.getenv("LITELM_ENVIRONMENT", "development") validate_environment(self.litellm_environment) if not api_key: @@ -46,17 +48,17 @@ class DeepEvalLogger(CustomLogger): await self._async_event_handler(kwargs, response_obj, start_time, end_time, is_success=True) def _prepare_trace_api(self, kwargs, response_obj, start_time, end_time, is_success): - _start_time = to_zod_compatible_iso(start_time) - _end_time = to_zod_compatible_iso(end_time) - _standard_logging_object = kwargs.get("standard_logging_object", {}) - base_api_span = self._create_base_api_span( + _start_time: Final = to_zod_compatible_iso(start_time) + _end_time: Final = to_zod_compatible_iso(end_time) + _standard_logging_object: Final = kwargs.get("standard_logging_object", {}) + base_api_span: Final = self._create_base_api_span( kwargs, standard_logging_object=_standard_logging_object, start_time=_start_time, end_time=_end_time, is_success=is_success, ) - trace_api = self._create_trace_api( + trace_api: Final = self._create_trace_api( base_api_span, standard_logging_object=_standard_logging_object, start_time=_start_time, @@ -74,9 +76,9 @@ class DeepEvalLogger(CustomLogger): return body def _sync_event_handler(self, kwargs, response_obj, start_time, end_time, is_success): - body = self._prepare_trace_api(kwargs, response_obj, start_time, end_time, is_success) + body: Final = self._prepare_trace_api(kwargs, response_obj, start_time, end_time, is_success) try: - response = self.api.send_request( + response: Final = self.api.send_request( method=HttpMethods.POST, endpoint=Endpoints.TRACING_ENDPOINT, body=body, @@ -86,8 +88,8 @@ class DeepEvalLogger(CustomLogger): verbose_logger.debug("DeepEvalLogger: sync_log_failure_event: Api response %s", response) async def _async_event_handler(self, kwargs, response_obj, start_time, end_time, is_success): - body = self._prepare_trace_api(kwargs, response_obj, start_time, end_time, is_success) - response = await self.api.a_send_request( + body: Final = self._prepare_trace_api(kwargs, response_obj, start_time, end_time, is_success) + response: Final = await self.api.a_send_request( method=HttpMethods.POST, endpoint=Endpoints.TRACING_ENDPOINT, body=body, @@ -97,7 +99,7 @@ class DeepEvalLogger(CustomLogger): def _create_base_api_span(self, kwargs, standard_logging_object, start_time, end_time, is_success): # extract usage - usage = standard_logging_object.get("response", {}).get("usage", {}) + usage: Final = standard_logging_object.get("response", {}).get("usage", {}) if is_success: output = ( standard_logging_object.get("response", {}) diff --git a/litellm/integrations/deepeval/types.py b/litellm/integrations/deepeval/types.py index afaf4436db9..c86d01d0468 100644 --- a/litellm/integrations/deepeval/types.py +++ b/litellm/integrations/deepeval/types.py @@ -1,7 +1,8 @@ # Duplicate -> https://github.com/confident-ai/deepeval/blob/main/deepeval/tracing/api.py from enum import Enum -from typing import Any, ClassVar, Dict, List, Optional, Union, Literal -from pydantic import BaseModel, Field, ConfigDict +from typing import Any, ClassVar, Literal + +from pydantic import BaseModel, ConfigDict, Field class SpanApiType(Enum): @@ -24,37 +25,37 @@ class BaseApiSpan(BaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(use_enum_values=True) uuid: str - name: Optional[str] = None + name: str | None = None status: TraceSpanApiStatus type: SpanApiType trace_uuid: str = Field(alias="traceUuid") - parent_uuid: Optional[str] = Field(None, alias="parentUuid") + parent_uuid: str | None = Field(None, alias="parentUuid") start_time: str = Field(alias="startTime") end_time: str = Field(alias="endTime") - input: Optional[Union[Dict, list, str]] = None - output: Optional[Union[Dict, list, str]] = None - error: Optional[str] = None + input: dict | list | str | None = None + output: dict | list | str | None = None + error: str | None = None # llm - model: Optional[str] = None - input_token_count: Optional[int] = Field(None, alias="inputTokenCount") - output_token_count: Optional[int] = Field(None, alias="outputTokenCount") - cost_per_input_token: Optional[float] = Field(None, alias="costPerInputToken") - cost_per_output_token: Optional[float] = Field(None, alias="costPerOutputToken") + model: str | None = None + input_token_count: int | None = Field(None, alias="inputTokenCount") + output_token_count: int | None = Field(None, alias="outputTokenCount") + cost_per_input_token: float | None = Field(None, alias="costPerInputToken") + cost_per_output_token: float | None = Field(None, alias="costPerOutputToken") class TraceApi(BaseModel): uuid: str - base_spans: List[BaseApiSpan] = Field(alias="baseSpans") - agent_spans: List[BaseApiSpan] = Field(alias="agentSpans") - llm_spans: List[BaseApiSpan] = Field(alias="llmSpans") - retriever_spans: List[BaseApiSpan] = Field(alias="retrieverSpans") - tool_spans: List[BaseApiSpan] = Field(alias="toolSpans") + base_spans: list[BaseApiSpan] = Field(alias="baseSpans") + agent_spans: list[BaseApiSpan] = Field(alias="agentSpans") + llm_spans: list[BaseApiSpan] = Field(alias="llmSpans") + retriever_spans: list[BaseApiSpan] = Field(alias="retrieverSpans") + tool_spans: list[BaseApiSpan] = Field(alias="toolSpans") start_time: str = Field(alias="startTime") end_time: str = Field(alias="endTime") - metadata: Optional[Dict[str, Any]] = Field(None) - tags: Optional[List[str]] = Field(None) - environment: Optional[str] = Field(None) + metadata: dict[str, Any] | None = Field(None) + tags: list[str] | None = Field(None) + environment: str | None = Field(None) class Environment(Enum): diff --git a/litellm/integrations/deepeval/utils.py b/litellm/integrations/deepeval/utils.py index 3df9aceb241..92e2600a3f1 100644 --- a/litellm/integrations/deepeval/utils.py +++ b/litellm/integrations/deepeval/utils.py @@ -1,4 +1,6 @@ from datetime import datetime, timezone +from typing import Final + from litellm.integrations.deepeval.types import Environment @@ -8,5 +10,5 @@ def to_zod_compatible_iso(dt: datetime) -> str: def validate_environment(environment: str): if environment not in [env.value for env in Environment]: - valid_values = ", ".join(f'"{env.value}"' for env in Environment) + valid_values: Final = ", ".join(f'"{env.value}"' for env in Environment) raise ValueError(f"Invalid environment: {environment}. Please use one of the following instead: {valid_values}") diff --git a/litellm/integrations/dotprompt/__init__.py b/litellm/integrations/dotprompt/__init__.py index 8432d50e32b..578b7c63871 100644 --- a/litellm/integrations/dotprompt/__init__.py +++ b/litellm/integrations/dotprompt/__init__.py @@ -1,17 +1,18 @@ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Final, Optional if TYPE_CHECKING: - from .prompt_manager import PromptManager, PromptTemplate - from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec from litellm.integrations.custom_prompt_management import CustomPromptManagement + from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec + + from .prompt_manager import PromptManager, PromptTemplate from litellm.types.prompts.init_prompts import SupportedPromptIntegrations from .dotprompt_manager import DotpromptManager # Global instances -global_prompt_directory: Optional[str] = None -global_prompt_manager: Optional["PromptManager"] = None +global_prompt_directory: Final[str | None] = None +global_prompt_manager: Final[Optional["PromptManager"]] = None def set_global_prompt_directory(directory: str) -> None: @@ -35,7 +36,7 @@ def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict: from .prompt_manager import PromptManager # Parse the dotprompt content to extract frontmatter and content - temp_manager = PromptManager() + temp_manager: Final = PromptManager() metadata, content = temp_manager._parse_frontmatter(dotprompt_content) # Convert to prompt_data format @@ -46,23 +47,23 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom """ Initialize a prompt from a .prompt file. """ - prompt_directory = getattr(litellm_params, "prompt_directory", None) + prompt_directory: Final = getattr(litellm_params, "prompt_directory", None) prompt_data = getattr(litellm_params, "prompt_data", None) - prompt_id = getattr(litellm_params, "prompt_id", None) + prompt_id: Final = getattr(litellm_params, "prompt_id", None) if prompt_directory: raise ValueError( "Cannot set prompt_directory when working with prompt_initializer. Needs to be a specific dotprompt file" ) - prompt_file = getattr(litellm_params, "prompt_file", None) + prompt_file: Final = getattr(litellm_params, "prompt_file", None) # Handle dotprompt_content from database - dotprompt_content = getattr(litellm_params, "dotprompt_content", None) + dotprompt_content: Final = getattr(litellm_params, "dotprompt_content", None) if dotprompt_content and not prompt_data and not prompt_file: prompt_data = _get_prompt_data_from_dotprompt_content(dotprompt_content) try: - dot_prompt_manager = DotpromptManager( + dot_prompt_manager: Final = DotpromptManager( prompt_directory=prompt_directory, prompt_data=prompt_data, prompt_file=prompt_file, @@ -74,16 +75,16 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom raise e -prompt_initializer_registry = { +prompt_initializer_registry: Final = { SupportedPromptIntegrations.DOT_PROMPT.value: prompt_initializer, } # Export public API __all__ = [ - "PromptManager", "DotpromptManager", + "PromptManager", "PromptTemplate", - "set_global_prompt_directory", "global_prompt_directory", "global_prompt_manager", + "set_global_prompt_directory", ] diff --git a/litellm/integrations/dotprompt/dotprompt_manager.py b/litellm/integrations/dotprompt/dotprompt_manager.py index 3ba9efd68b7..bedeb803c27 100644 --- a/litellm/integrations/dotprompt/dotprompt_manager.py +++ b/litellm/integrations/dotprompt/dotprompt_manager.py @@ -4,7 +4,7 @@ Builds on top of PromptManagementBase to provide .prompt file support. """ import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Final from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.integrations.prompt_management_base import PromptManagementClient @@ -42,10 +42,10 @@ class DotpromptManager(CustomPromptManagement): def __init__( self, - prompt_directory: Optional[str] = None, - prompt_file: Optional[str] = None, - prompt_data: Optional[Union[dict, str]] = None, - prompt_id: Optional[str] = None, + prompt_directory: str | None = None, + prompt_file: str | None = None, + prompt_data: dict | str | None = None, + prompt_id: str | None = None, ): import litellm @@ -56,7 +56,7 @@ class DotpromptManager(CustomPromptManagement): else: self.prompt_data = prompt_data or {} - self._prompt_manager: Optional[PromptManager] = None + self._prompt_manager: PromptManager | None = None self.prompt_file = prompt_file self.prompt_id = prompt_id @@ -84,8 +84,8 @@ class DotpromptManager(CustomPromptManagement): def should_run_prompt_management( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], + prompt_id: str | None, + prompt_spec: PromptSpec | None, dynamic_callback_params: StandardCallbackDynamicParams, ) -> bool: """ @@ -103,12 +103,12 @@ class DotpromptManager(CustomPromptManagement): def _compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_spec: Optional[PromptSpec], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_spec: PromptSpec | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """ Compile a .prompt file into a PromptManagementClient structure. @@ -125,26 +125,26 @@ class DotpromptManager(CustomPromptManagement): try: # Get the prompt template (versioned or base) - template = self.prompt_manager.get_prompt(prompt_id=prompt_id, version=prompt_version) + template: Final = self.prompt_manager.get_prompt(prompt_id=prompt_id, version=prompt_version) if template is None: - version_str = f" (version {prompt_version})" if prompt_version else "" + version_str: Final = f" (version {prompt_version})" if prompt_version else "" raise ValueError(f"Prompt '{prompt_id}'{version_str} not found in prompt directory") # Render the template with variables (pass version for proper lookup) - rendered_content = self.prompt_manager.render( + rendered_content: Final = self.prompt_manager.render( prompt_id=prompt_id, prompt_variables=prompt_variables, version=prompt_version, ) # Convert rendered content to chat messages - messages = self._convert_to_messages(rendered_content) + messages: Final = self._convert_to_messages(rendered_content) # Extract model from metadata (if specified) - template_model = template.model + template_model: Final = template.model # Extract optional parameters from metadata - optional_params = self._extract_optional_params(template) + optional_params: Final = self._extract_optional_params(template) return PromptManagementClient( prompt_id=prompt_id, @@ -159,12 +159,12 @@ class DotpromptManager(CustomPromptManagement): async def async_compile_prompt_helper( self, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, ) -> PromptManagementClient: """ Async version of compile prompt helper. Since dotprompt operations are synchronous, @@ -185,17 +185,17 @@ class DotpromptManager(CustomPromptManagement): def get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, - prompt_spec: Optional[PromptSpec] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: from litellm.integrations.prompt_management_base import PromptManagementBase return PromptManagementBase.get_chat_completion_prompt( @@ -214,19 +214,19 @@ class DotpromptManager(CustomPromptManagement): async def async_get_chat_completion_prompt( self, model: str, - messages: List[AllMessageValues], + messages: list[AllMessageValues], non_default_params: dict, - prompt_id: Optional[str], - prompt_variables: Optional[dict], + prompt_id: str | None, + prompt_variables: dict | None, dynamic_callback_params: StandardCallbackDynamicParams, litellm_logging_obj: LiteLLMLoggingObj, - prompt_spec: Optional[PromptSpec] = None, - tools: Optional[List[Dict]] = None, - prompt_label: Optional[str] = None, - prompt_version: Optional[int] = None, - ignore_prompt_manager_model: Optional[bool] = False, - ignore_prompt_manager_optional_params: Optional[bool] = False, - ) -> Tuple[str, List[AllMessageValues], dict]: + prompt_spec: PromptSpec | None = None, + tools: list[dict] | None = None, + prompt_label: str | None = None, + prompt_version: int | None = None, + ignore_prompt_manager_model: bool | None = False, + ignore_prompt_manager_optional_params: bool | None = False, + ) -> tuple[str, list[AllMessageValues], dict]: """ Async version - delegates to PromptManagementBase async implementation. """ @@ -249,7 +249,7 @@ class DotpromptManager(CustomPromptManagement): ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params, ) - def _convert_to_messages(self, rendered_content: str) -> List[AllMessageValues]: + def _convert_to_messages(self, rendered_content: str) -> list[AllMessageValues]: """ Convert rendered prompt content to chat messages. @@ -259,14 +259,14 @@ class DotpromptManager(CustomPromptManagement): 3. Already formatted as a single message """ # Clean up the content - content = rendered_content.strip() + content: Final = rendered_content.strip() # Try to parse role-based format (System: ..., User: ..., etc.) - messages = [] + messages: Final = [] current_role = None current_content = [] - lines = content.split("\n") + lines: Final = content.split("\n") for line in lines: line = line.strip() @@ -298,7 +298,7 @@ class DotpromptManager(CustomPromptManagement): # Add the last message if current_role and current_content: - content_text = "\n".join(current_content).strip() + content_text: Final = "\n".join(current_content).strip() if content_text: # Only add if there's actual content messages.append(self._create_message(current_role, content_text)) @@ -321,7 +321,7 @@ class DotpromptManager(CustomPromptManagement): Includes parameters like temperature, max_tokens, etc. """ - optional_params = {} + optional_params: Final = {} # Extract common parameters from metadata if template.optional_params is not None: @@ -339,20 +339,20 @@ class DotpromptManager(CustomPromptManagement): if self._prompt_manager: self._prompt_manager.reload_prompts() - def add_prompt_from_json(self, prompt_id: str, json_data: Dict[str, Any]) -> None: + def add_prompt_from_json(self, prompt_id: str, json_data: dict[str, Any]) -> None: """Add a prompt from JSON data.""" - content = json_data.get("content", "") - metadata = json_data.get("metadata", {}) + content: Final = json_data.get("content", "") + metadata: Final = json_data.get("metadata", {}) self.prompt_manager.add_prompt(prompt_id, content, metadata) - def load_prompts_from_json(self, prompts_data: Dict[str, Dict[str, Any]]) -> None: + def load_prompts_from_json(self, prompts_data: dict[str, dict[str, Any]]) -> None: """Load multiple prompts from JSON data.""" self.prompt_manager.load_prompts_from_json_data(prompts_data) - def get_prompts_as_json(self) -> Dict[str, Dict[str, Any]]: + def get_prompts_as_json(self) -> dict[str, dict[str, Any]]: """Get all prompts in JSON format.""" return self.prompt_manager.get_all_prompts_as_json() - def convert_prompt_file_to_json(self, file_path: str) -> Dict[str, Any]: + def convert_prompt_file_to_json(self, file_path: str) -> dict[str, Any]: """Convert a .prompt file to JSON format.""" return self.prompt_manager.prompt_file_to_json(file_path) diff --git a/litellm/integrations/dotprompt/prompt_manager.py b/litellm/integrations/dotprompt/prompt_manager.py index dd198ba1272..70bad2f7290 100644 --- a/litellm/integrations/dotprompt/prompt_manager.py +++ b/litellm/integrations/dotprompt/prompt_manager.py @@ -4,7 +4,7 @@ Based on Google's GenAI Kit dotprompt implementation: https://google.github.io/d import re from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Final import yaml from jinja2 import DictLoader, select_autoescape @@ -17,15 +17,15 @@ class PromptTemplate: def __init__( self, content: str, - metadata: Optional[Dict[str, Any]] = None, - template_id: Optional[str] = None, + metadata: dict[str, Any] | None = None, + template_id: str | None = None, ): self.content = content self.metadata = metadata or {} self.template_id = template_id # Extract common metadata fields - restricted_keys = ["model", "input", "output"] + restricted_keys: Final = ["model", "input", "output"] self.model = self.metadata.get("model") self.input_schema = self.metadata.get("input", {}).get("schema", {}) self.output_format = self.metadata.get("output", {}).get("format") @@ -52,13 +52,13 @@ class PromptManager: def __init__( self, - prompt_id: Optional[str] = None, - prompt_directory: Optional[str] = None, - prompt_data: Optional[Dict[str, Dict[str, Any]]] = None, - prompt_file: Optional[str] = None, + prompt_id: str | None = None, + prompt_directory: str | None = None, + prompt_data: dict[str, dict[str, Any]] | None = None, + prompt_file: str | None = None, ): self.prompt_directory = Path(prompt_directory) if prompt_directory else None - self.prompts: Dict[str, PromptTemplate] = {} + self.prompts: dict[str, PromptTemplate] = {} self.prompt_file = prompt_file # Sandboxed env: templates can come from user input via /prompts/test, # so we must block access to unsafe Python attributes and mutation of @@ -83,7 +83,7 @@ class PromptManager: if not prompt_id: raise ValueError("prompt_id is required when prompt_file is provided") - template = self._load_prompt_file(self.prompt_file, prompt_id) + template: Final = self._load_prompt_file(self.prompt_file, prompt_id) self.prompts[prompt_id] = template # Load prompts from JSON data if provided @@ -95,7 +95,7 @@ class PromptManager: if not self.prompt_directory or not self.prompt_directory.exists(): raise ValueError(f"Prompt directory does not exist: {self.prompt_directory}") - prompt_files = list(self.prompt_directory.glob("*.prompt")) + prompt_files: Final = list(self.prompt_directory.glob("*.prompt")) for prompt_file in prompt_files: try: @@ -107,7 +107,7 @@ class PromptManager: # Optional: print(f"Error loading prompt file {prompt_file}") pass - def _load_prompts_from_json(self, prompt_data: Dict[str, Dict[str, Any]], prompt_id: Optional[str] = None) -> None: + def _load_prompts_from_json(self, prompt_data: dict[str, dict[str, Any]], prompt_id: str | None = None) -> None: """Load prompts from JSON data structure. Expected format: @@ -143,12 +143,12 @@ class PromptManager: # Optional: print(f"Error loading prompt from JSON: {prompt_id}") pass - def _load_prompt_file(self, file_path: Union[str, Path], prompt_id: str) -> PromptTemplate: + def _load_prompt_file(self, file_path: str | Path, prompt_id: str) -> PromptTemplate: """Load and parse a single .prompt file.""" if isinstance(file_path, str): file_path = Path(file_path) - content = file_path.read_text(encoding="utf-8") + content: Final = file_path.read_text(encoding="utf-8") # Split frontmatter and content frontmatter, template_content = self._parse_frontmatter(content) @@ -159,14 +159,14 @@ class PromptManager: template_id=prompt_id, ) - def _parse_frontmatter(self, content: str) -> Tuple[Dict[str, Any], str]: + def _parse_frontmatter(self, content: str) -> tuple[dict[str, Any], str]: """Parse YAML frontmatter from prompt content.""" # Match YAML frontmatter between --- delimiters - frontmatter_pattern = r"^---\s*\n(.*?)\n---\s*\n(.*)$" - match = re.match(frontmatter_pattern, content, re.DOTALL) + frontmatter_pattern: Final = r"^---\s*\n(.*?)\n---\s*\n(.*)$" + match: Final = re.match(frontmatter_pattern, content, re.DOTALL) if match: - frontmatter_yaml = match.group(1) + frontmatter_yaml: Final = match.group(1) template_content = match.group(2) try: @@ -183,8 +183,8 @@ class PromptManager: def render( self, prompt_id: str, - prompt_variables: Optional[Dict[str, Any]] = None, - version: Optional[int] = None, + prompt_variables: dict[str, Any] | None = None, + version: int | None = None, ) -> str: """ Render a prompt template with the given variables. @@ -202,14 +202,14 @@ class PromptManager: ValueError: If template rendering fails """ # Get the template (versioned or base) - template = self.get_prompt(prompt_id=prompt_id, version=version) + template: Final = self.get_prompt(prompt_id=prompt_id, version=version) if template is None: - available_prompts = list(self.prompts.keys()) - version_str = f" (version {version})" if version else "" + available_prompts: Final = list(self.prompts.keys()) + version_str: Final = f" (version {version})" if version else "" raise KeyError(f"Prompt '{prompt_id}'{version_str} not found. Available prompts: {available_prompts}") - variables = prompt_variables or {} + variables: Final = prompt_variables or {} # Validate input variables against schema if defined if template.input_schema: @@ -217,13 +217,13 @@ class PromptManager: try: # Create Jinja2 template and render - jinja_template = self.jinja_env.from_string(template.content) - rendered = jinja_template.render(**variables) + jinja_template: Final = self.jinja_env.from_string(template.content) + rendered: Final = jinja_template.render(**variables) return rendered except Exception as e: raise ValueError(f"Error rendering template '{prompt_id}': {e}") - def _validate_input(self, variables: Dict[str, Any], schema: Dict[str, Any]) -> None: + def _validate_input(self, variables: dict[str, Any], schema: dict[str, Any]) -> None: """Basic validation of input variables against schema.""" for field_name, field_type in schema.items(): if field_name in variables: @@ -236,9 +236,9 @@ class PromptManager: f"expected {getattr(expected_type, '__name__', str(expected_type))}, got {type(value).__name__}" ) - def _get_python_type(self, schema_type: str) -> Union[type, tuple]: + def _get_python_type(self, schema_type: str) -> type | tuple: """Convert schema type string to Python type.""" - type_mapping: Dict[str, Union[type, tuple]] = { + type_mapping: Final[dict[str, type | tuple]] = { "string": str, "str": str, "number": (int, float), @@ -255,7 +255,7 @@ class PromptManager: return type_mapping.get(schema_type.lower(), str) # type: ignore - def get_prompt(self, prompt_id: str, version: Optional[int] = None) -> Optional[PromptTemplate]: + def get_prompt(self, prompt_id: str, version: int | None = None) -> PromptTemplate | None: """ Get a prompt template by ID and optional version. @@ -268,20 +268,20 @@ class PromptManager: """ if version is not None: # Try versioned prompt first: prompt_id.v{version} - versioned_id = f"{prompt_id}.v{version}" + versioned_id: Final = f"{prompt_id}.v{version}" if versioned_id in self.prompts: return self.prompts[versioned_id] # Fall back to base prompt_id return self.prompts.get(prompt_id) - def list_prompts(self) -> List[str]: + def list_prompts(self) -> list[str]: """Get a list of all available prompt IDs.""" return list(self.prompts.keys()) - def get_prompt_metadata(self, prompt_id: str) -> Optional[Dict[str, Any]]: + def get_prompt_metadata(self, prompt_id: str) -> dict[str, Any] | None: """Get metadata for a specific prompt.""" - template = self.prompts.get(prompt_id) + template: Final = self.prompts.get(prompt_id) return template.metadata if template else None def reload_prompts(self) -> None: @@ -290,12 +290,12 @@ class PromptManager: if self.prompt_directory: self._load_prompts() - def add_prompt(self, prompt_id: str, content: str, metadata: Optional[Dict[str, Any]] = None) -> None: + def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, Any] | None = None) -> None: """Add a prompt template programmatically.""" - template = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id) + template: Final = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id) self.prompts[prompt_id] = template - def prompt_file_to_json(self, file_path: Union[str, Path]) -> Dict[str, Any]: + def prompt_file_to_json(self, file_path: str | Path) -> dict[str, Any]: """Convert a .prompt file to JSON format. Args: @@ -305,14 +305,14 @@ class PromptManager: Dictionary with 'content' and 'metadata' keys """ file_path = Path(file_path) - content = file_path.read_text(encoding="utf-8") + content: Final = file_path.read_text(encoding="utf-8") # Parse frontmatter and content frontmatter, template_content = self._parse_frontmatter(content) return {"content": template_content.strip(), "metadata": frontmatter} - def json_to_prompt_file(self, prompt_data: Dict[str, Any]) -> str: + def json_to_prompt_file(self, prompt_data: dict[str, Any]) -> str: """Convert JSON prompt data to .prompt file format. Args: @@ -321,8 +321,8 @@ class PromptManager: Returns: String content in .prompt file format """ - content = prompt_data.get("content", "") - metadata = prompt_data.get("metadata", {}) + content: Final = prompt_data.get("content", "") + metadata: Final = prompt_data.get("metadata", {}) if not metadata: # No metadata, return just the content @@ -331,17 +331,17 @@ class PromptManager: # Convert metadata to YAML frontmatter import yaml - frontmatter_yaml = yaml.dump(metadata, default_flow_style=False) + frontmatter_yaml: Final = yaml.dump(metadata, default_flow_style=False) return f"---\n{frontmatter_yaml}---\n{content}" - def get_all_prompts_as_json(self) -> Dict[str, Dict[str, Any]]: + def get_all_prompts_as_json(self) -> dict[str, dict[str, Any]]: """Get all loaded prompts in JSON format. Returns: Dictionary mapping prompt_id to prompt data """ - result = {} + result: Final = {} for prompt_id, template in self.prompts.items(): result[prompt_id] = { "content": template.content, @@ -349,6 +349,6 @@ class PromptManager: } return result - def load_prompts_from_json_data(self, prompt_data: Dict[str, Dict[str, Any]]) -> None: + def load_prompts_from_json_data(self, prompt_data: dict[str, dict[str, Any]]) -> None: """Load additional prompts from JSON data (merges with existing prompts).""" self._load_prompts_from_json(prompt_data) diff --git a/litellm/integrations/dynamodb.py b/litellm/integrations/dynamodb.py index ab76fa3c8bd..38f5924a233 100644 --- a/litellm/integrations/dynamodb.py +++ b/litellm/integrations/dynamodb.py @@ -3,10 +3,10 @@ import os import traceback -from litellm._uuid import uuid -from typing import Any +from typing import Any, Final import litellm +from litellm._uuid import uuid class DyanmoDBLogger: @@ -32,16 +32,16 @@ class DyanmoDBLogger: # construct payload to send to DynamoDB # follows the same params as langfuse.py - litellm_params = kwargs.get("litellm_params", {}) - metadata = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None - messages = kwargs.get("messages") - optional_params = kwargs.get("optional_params", {}) - call_type = kwargs.get("call_type", "litellm.completion") - usage = response_obj["usage"] - id = response_obj.get("id", str(uuid.uuid4())) + litellm_params: Final = kwargs.get("litellm_params", {}) + metadata: Final = litellm_params.get("metadata", {}) or {} # if litellm_params['metadata'] == None + messages: Final = kwargs.get("messages") + optional_params: Final = kwargs.get("optional_params", {}) + call_type: Final = kwargs.get("call_type", "litellm.completion") + usage: Final = response_obj["usage"] + id: Final = response_obj.get("id", str(uuid.uuid4())) # Build the initial payload - payload = { + payload: Final = { "id": id, "call_type": call_type, "startTime": start_time, @@ -66,14 +66,13 @@ class DyanmoDBLogger: print_verbose(f"\nDynamoDB Logger - Logging payload = {payload}") # put data in dyanmo DB - table = self.dynamodb.Table(self.table_name) + table: Final = self.dynamodb.Table(self.table_name) # Assuming log_data is a dictionary with log information - response = table.put_item(Item=payload) + response: Final = table.put_item(Item=payload) - print_verbose(f"Response from DynamoDB:{str(response)}") + print_verbose(f"Response from DynamoDB:{response}") print_verbose(f"DynamoDB Layer Logging - final response object: {response_obj}") return response except Exception: print_verbose(f"DynamoDB Layer Error - {traceback.format_exc()}") - pass diff --git a/litellm/integrations/email_alerting.py b/litellm/integrations/email_alerting.py index 35d63a691f9..351896425bb 100644 --- a/litellm/integrations/email_alerting.py +++ b/litellm/integrations/email_alerting.py @@ -3,18 +3,18 @@ Functions for sending Email Alerts """ import os -from typing import List, Optional +from typing import Final from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.proxy._types import WebhookEvent from litellm.repositories.team_repository import TeamRepository # we use this for the email header, please send a test email if you change this. verify it looks good on email -LITELLM_LOGO_URL = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" -LITELLM_SUPPORT_CONTACT = "support@berri.ai" +LITELLM_LOGO_URL: Final = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png" +LITELLM_SUPPORT_CONTACT: Final = "support@berri.ai" -async def get_all_team_member_emails(team_id: Optional[str] = None) -> list: +async def get_all_team_member_emails(team_id: str | None = None) -> list: verbose_logger.debug("Email Alerting: Getting all team members for team_id=%s", team_id) if team_id is None: return [] @@ -23,7 +23,7 @@ async def get_all_team_member_emails(team_id: Optional[str] = None) -> list: if prisma_client is None: raise Exception("Not connected to DB!") - team_row = await TeamRepository(prisma_client).table.find_unique( + team_row: Final = await TeamRepository(prisma_client).table.find_unique( where={ "team_id": team_id, } @@ -32,33 +32,33 @@ async def get_all_team_member_emails(team_id: Optional[str] = None) -> list: if team_row is None: return [] - _team_members = team_row.members_with_roles + _team_members: Final = team_row.members_with_roles verbose_logger.debug( "Email Alerting: Got team members for team_id=%s Team Members: %s", team_id, _team_members, ) - _team_member_user_ids: List[str] = [] + _team_member_user_ids: Final[list[str]] = [] for member in _team_members: if member and isinstance(member, dict): _user_id = member.get("user_id") if _user_id and isinstance(_user_id, str): _team_member_user_ids.append(_user_id) - sql_query = """ + sql_query: Final = """ SELECT user_email FROM "LiteLLM_UserTable" WHERE user_id = ANY($1::TEXT[]); """ - _result = await prisma_client.db.query_raw(sql_query, _team_member_user_ids) + _result: Final = await prisma_client.db.query_raw(sql_query, _team_member_user_ids) verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result) if _result is None: return [] - emails = [] + emails: Final = [] for user in _result: if user and isinstance(user, dict) and user.get("user_email", None) is not None: emails.append(user.get("user_email")) @@ -72,8 +72,8 @@ async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool: """ from litellm.proxy.utils import send_email - _team_id = webhook_event.team_id - team_alias = webhook_event.team_alias + _team_id: Final = webhook_event.team_id + team_alias: Final = webhook_event.team_alias verbose_logger.debug("Email Alerting: Sending Team Budget Alert for team=%s", team_alias) email_logo_url = os.getenv("SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", None)) @@ -87,12 +87,12 @@ async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool: email_logo_url = LITELLM_LOGO_URL if email_support_contact is None: email_support_contact = LITELLM_SUPPORT_CONTACT - recipient_emails = await get_all_team_member_emails(_team_id) - recipient_emails_str: str = ",".join(recipient_emails) + recipient_emails: Final = await get_all_team_member_emails(_team_id) + recipient_emails_str: Final[str] = ",".join(recipient_emails) verbose_logger.debug("Email Alerting: Sending team budget alert to %s", recipient_emails_str) - event_name = webhook_event.event_message - max_budget = webhook_event.max_budget + event_name: Final = webhook_event.event_message + max_budget: Final = webhook_event.max_budget email_html_content = "Alert from LiteLLM Server" if recipient_emails_str is None: @@ -116,7 +116,7 @@ async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool: The LiteLLM team
""" - email_event = { + email_event: Final = { "to": recipient_emails_str, "subject": f"LiteLLM {event_name} for Team {team_alias}", "html": email_html_content, diff --git a/litellm/integrations/email_templates/email_footer.py b/litellm/integrations/email_templates/email_footer.py index feb692354a0..950496a5c07 100644 --- a/litellm/integrations/email_templates/email_footer.py +++ b/litellm/integrations/email_templates/email_footer.py @@ -1,4 +1,6 @@ -EMAIL_FOOTER = """ +from typing import Final + +EMAIL_FOOTER: Final = """