diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 7ae79aa666f..cfa0390e836 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,6 +1,9 @@ /ui/ @yuneng-berri @ryan-crabbe-berri /litellm/proxy/_experimental/out/ @yuneng-berri @ryan-crabbe-berri +/ui/Dockerfile +/ui/nginx.conf /ui/litellm-dashboard/src/lib/http/schema.d.ts +/ui/litellm-dashboard/tsconfig.tsbuildinfo /model_prices_and_context_window.json @mateo-berri /litellm/model_prices_and_context_window_backup.json @mateo-berri /litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri diff --git a/.github/scripts/e2e_egress_sentinel.py b/.github/scripts/e2e_egress_sentinel.py new file mode 100755 index 00000000000..b40c72d1fd7 --- /dev/null +++ b/.github/scripts/e2e_egress_sentinel.py @@ -0,0 +1,198 @@ +"""Prove an e2e replay run makes zero outbound provider calls, by counting them. + +`serve` pins each provider host (`--host`) to a local sink address in the hosts +file and binds a counting listener on that address, so any connection the proxy +or the record/replay edge opens to a real provider is redirected to the sink, +recorded as one line in `--hits-file`, and never leaves the box. The record and +replay edge only ever dials `127.0.0.1:` (a different host than the +pinned provider names), so in a clean replay the sink sees nothing; a single hit +means a provider call escaped the bundle. `assert-empty` turns that hit file into +the pass/fail check. + +Stdlib only, so CI runs it under the system interpreter as root (binding :443 and +editing the hosts file both need root); `--sink-address`, `--port`, and +`--hosts-file` are injectable so it runs unprivileged against a temp hosts file on +a high port under test. +""" + +# ruff: noqa: T201 # CLI script: its stdout/stderr progress and results are the interface +from __future__ import annotations + +import argparse +import json +import os +import signal +import socket +import sys +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from types import FrameType +from typing import Final + +_BLOCK_BEGIN: Final = "# BEGIN e2e-egress-sentinel" +_BLOCK_END: Final = "# END e2e-egress-sentinel" + + +@dataclass(frozen=True, slots=True) +class ServeConfig: + hosts: tuple[str, ...] + sink_address: str + ports: tuple[int, ...] + hits_file: Path + hosts_file: Path + ready_file: Path | None + pid_file: Path | None + + +def _pin_block(sink_address: str, hosts: tuple[str, ...]) -> str: + lines = "\n".join(f"{sink_address}\t{host}" for host in hosts) + return f"\n{_BLOCK_BEGIN}\n{lines}\n{_BLOCK_END}\n" + + +def _install_pins(hosts_file: Path, sink_address: str, hosts: tuple[str, ...]) -> bytes: + original = hosts_file.read_bytes() if hosts_file.exists() else b"" + hosts_file.write_bytes(original + _pin_block(sink_address, hosts).encode()) + return original + + +def _restore_pins(hosts_file: Path, original: bytes) -> None: + hosts_file.write_bytes(original) + + +def _bind(sink_address: str, port: int) -> socket.socket: + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind((sink_address, port)) + listener.listen(128) + return listener + + +@dataclass(frozen=True, slots=True) +class _HitLog: + path: Path + _lock: threading.Lock + + def record(self, *, port: int, peer: tuple[str, int]) -> None: + entry = json.dumps({"ts": time.time(), "port": port, "peer": list(peer)}) + with self._lock: + with self.path.open("a", encoding="utf-8") as handle: + handle.write(entry + "\n") + + +def _serve_socket(listener: socket.socket, port: int, hits: _HitLog, stop: threading.Event) -> None: + while not stop.is_set(): + try: + conn, peer = listener.accept() + except OSError: + return + hits.record(port=port, peer=(peer[0], peer[1])) + try: + conn.close() + except OSError: + pass + + +def serve(config: ServeConfig) -> int: + config.hits_file.write_text("", encoding="utf-8") + original_hosts = _install_pins(config.hosts_file, config.sink_address, config.hosts) + try: + listeners = tuple(_bind(config.sink_address, port) for port in config.ports) + except OSError as exc: + _restore_pins(config.hosts_file, original_hosts) + print(f"egress sentinel could not bind a sink: {exc}", file=sys.stderr) + return 1 + + stop = threading.Event() + hits = _HitLog(path=config.hits_file, _lock=threading.Lock()) + threads = tuple( + threading.Thread(target=_serve_socket, args=(listener, port, hits, stop), daemon=True) + for listener, port in zip(listeners, config.ports) + ) + for thread in threads: + thread.start() + + def _handle(_signum: int, _frame: FrameType | None) -> None: + stop.set() + for listener in listeners: + try: + listener.close() + except OSError: + pass + + signal.signal(signal.SIGTERM, _handle) + signal.signal(signal.SIGINT, _handle) + + if config.pid_file is not None: + config.pid_file.write_text(str(os.getpid()), encoding="utf-8") + if config.ready_file is not None: + config.ready_file.write_text("ready", encoding="utf-8") + print( + f"egress sentinel up: pinned {', '.join(config.hosts)} to {config.sink_address} " + f"on port(s) {', '.join(str(p) for p in config.ports)}", + flush=True, + ) + + stop.wait() + _restore_pins(config.hosts_file, original_hosts) + if config.ready_file is not None and config.ready_file.exists(): + config.ready_file.unlink() + if config.pid_file is not None and config.pid_file.exists(): + config.pid_file.unlink() + return 0 + + +def assert_empty(hits_file: Path) -> int: + if not hits_file.exists(): + print(f"egress sentinel recorded no provider calls ({hits_file} absent): zero egress") + return 0 + hits = [line for line in hits_file.read_text(encoding="utf-8").splitlines() if line.strip()] + if not hits: + print("egress sentinel recorded no provider calls: zero egress") + return 0 + print(f"egress sentinel recorded {len(hits)} provider call(s); replay was not hermetic:", file=sys.stderr) + for line in hits: + print(f" {line}", file=sys.stderr) + return 1 + + +def _serve_from_args(args: argparse.Namespace) -> int: + config = ServeConfig( + hosts=tuple(args.host), + sink_address=args.sink_address, + ports=tuple(args.port), + hits_file=Path(args.hits_file), + hosts_file=Path(args.hosts_file), + ready_file=Path(args.ready_file) if args.ready_file else None, + pid_file=Path(args.pid_file) if args.pid_file else None, + ) + return serve(config) + + +def main(argv: tuple[str, ...]) -> int: + parser = argparse.ArgumentParser(description="count outbound provider calls during an e2e replay") + sub = parser.add_subparsers(dest="command", required=True) + + serve_parser = sub.add_parser("serve", help="pin provider hosts and count connection attempts") + serve_parser.add_argument("--host", action="append", required=True, help="provider host to pin and watch") + serve_parser.add_argument("--sink-address", default="127.0.0.1") + serve_parser.add_argument("--port", action="append", type=int, default=None) + serve_parser.add_argument("--hits-file", required=True) + serve_parser.add_argument("--hosts-file", default="/etc/hosts") + serve_parser.add_argument("--ready-file", default=None) + serve_parser.add_argument("--pid-file", default=None) + + assert_parser = sub.add_parser("assert-empty", help="exit non-zero if any provider call was recorded") + assert_parser.add_argument("--hits-file", required=True) + + args = parser.parse_args(argv) + if args.command == "serve": + if args.port is None: + args.port = [443] + return _serve_from_args(args) + return assert_empty(Path(args.hits_file)) + + +if __name__ == "__main__": + raise SystemExit(main(tuple(sys.argv[1:]))) diff --git a/.github/scripts/e2e_fetch_fixture_bundle.sh b/.github/scripts/e2e_fetch_fixture_bundle.sh new file mode 100755 index 00000000000..b76ced74b3b --- /dev/null +++ b/.github/scripts/e2e_fetch_fixture_bundle.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO="${1:-${GITHUB_REPOSITORY:?REPO required}}" +ARTIFACT_NAME="${2:-e2e-fixtures-bundle}" +BASE_BRANCH="${3:?base branch required}" +DEST_DIR="${4:?destination bundle dir required}" + +: "${GH_TOKEN:?GH_TOKEN required to query and download artifacts}" + +WORKDIR="$(mktemp -d)" +trap 'rm -rf "${WORKDIR}"' EXIT + +echo "resolving newest non-expired '${ARTIFACT_NAME}' artifact on ${REPO}@${BASE_BRANCH}" + +SELECTED="$( + gh api "repos/${REPO}/actions/artifacts" -X GET -f per_page=100 --paginate \ + --jq ".artifacts[] | select(.name == \"${ARTIFACT_NAME}\" and .expired == false and .workflow_run.head_branch == \"${BASE_BRANCH}\") | {id, digest, created_at, run_id: .workflow_run.id, run_number: .workflow_run.run_number}" \ + | jq -s 'sort_by(.created_at) | reverse | .[0] // empty' +)" + +if [[ -z "${SELECTED}" ]]; then + echo "no usable '${ARTIFACT_NAME}' artifact on ${BASE_BRANCH}: the last record run produced none (a red Saturday), so there is nothing fresh to replay; failing loudly instead of replaying a stale bundle" >&2 + exit 1 +fi + +RUN_ID="$(echo "${SELECTED}" | jq -r '.run_id')" +RUN_NUMBER="$(echo "${SELECTED}" | jq -r '.run_number')" +ARTIFACT_ID="$(echo "${SELECTED}" | jq -r '.id')" +GH_DIGEST="$(echo "${SELECTED}" | jq -r '.digest // "unknown"')" +CREATED_AT="$(echo "${SELECTED}" | jq -r '.created_at')" + +echo "pinned bundle: run #${RUN_NUMBER} (run_id=${RUN_ID}, artifact_id=${ARTIFACT_ID}), recorded ${CREATED_AT}, github digest ${GH_DIGEST}" + +gh run download "${RUN_ID}" --repo "${REPO}" -n "${ARTIFACT_NAME}" -D "${WORKDIR}" + +TARBALL="$(find "${WORKDIR}" -name '*.tar.gz' -type f | head -n 1)" +if [[ -z "${TARBALL}" ]]; then + echo "downloaded artifact contained no tarball" >&2 + exit 1 +fi +SIDECAR="${TARBALL}.sha256" +if [[ ! -f "${SIDECAR}" ]]; then + echo "downloaded artifact has no ${SIDECAR}: cannot verify the bundle digest" >&2 + exit 1 +fi + +echo "verifying bundle against its recorded sha256 digest" +( cd "$(dirname "${TARBALL}")" && sha256sum -c "$(basename "${SIDECAR}")" ) + +mkdir -p "${DEST_DIR}" +tar xzf "${TARBALL}" -C "${DEST_DIR}" + +echo "extracted bundle into ${DEST_DIR}" +python3 -c "import json,sys; m=json.load(open(sys.argv[1])); print(' recorded_at', m['recorded_at'], 'harness', m['harness_version'], 'format_version', m['format_version'])" "${DEST_DIR}/manifest.json" diff --git a/.github/scripts/e2e_pack_fixture_bundle.sh b/.github/scripts/e2e_pack_fixture_bundle.sh new file mode 100755 index 00000000000..2105447cc65 --- /dev/null +++ b/.github/scripts/e2e_pack_fixture_bundle.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 2 ]]; then + echo "usage: $0 " >&2 + exit 2 +fi + +BUNDLE_DIR="$1" +OUT_TARBALL="$2" + +MANIFEST="${BUNDLE_DIR}/manifest.json" +if [[ ! -f "${MANIFEST}" ]]; then + echo "no ${MANIFEST}: refusing to publish a bundle with no manifest (record produced nothing)" >&2 + exit 1 +fi + +echo "packing fixture bundle from ${BUNDLE_DIR}" +python3 -c "import json,sys; m=json.load(open(sys.argv[1])); print(' format_version', m['format_version'], 'recorded_at', m['recorded_at'], 'harness', m['harness_version'])" "${MANIFEST}" + +TEST_DIRS=$(find "${BUNDLE_DIR}" -mindepth 1 -maxdepth 1 -type d | wc -l | tr -d ' ') +if [[ "${TEST_DIRS}" -eq 0 ]]; then + echo "bundle at ${BUNDLE_DIR} has a manifest but no recorded interactions; refusing to publish an empty bundle" >&2 + exit 1 +fi +echo " ${TEST_DIRS} recorded test director(ies)" + +mkdir -p "$(dirname "${OUT_TARBALL}")" +tar czf "${OUT_TARBALL}" -C "${BUNDLE_DIR}" . + +OUT_DIR="$(cd "$(dirname "${OUT_TARBALL}")" && pwd)" +OUT_BASE="$(basename "${OUT_TARBALL}")" +( cd "${OUT_DIR}" && sha256sum "${OUT_BASE}" > "${OUT_BASE}.sha256" ) + +echo "wrote ${OUT_TARBALL} ($(du -h "${OUT_TARBALL}" | cut -f1)) and ${OUT_BASE}.sha256" +cat "${OUT_DIR}/${OUT_BASE}.sha256" diff --git a/.github/workflows/e2e_record_replay.yml b/.github/workflows/e2e_record_replay.yml new file mode 100644 index 00000000000..ca52f0b4d81 --- /dev/null +++ b/.github/workflows/e2e_record_replay.yml @@ -0,0 +1,237 @@ +name: "E2E Record and Replay" + +on: + schedule: + - cron: "0 8 * * 6" + - cron: "0 8 * * 1-5" + workflow_dispatch: + inputs: + mode: + description: "record (hits real providers and publishes a fresh bundle) or replay (bundle only, zero provider egress)" + type: choice + options: + - record + - replay + default: record + +permissions: + contents: read + +jobs: + record: + name: "Record the e2e suite against real providers" + if: >- + (github.event_name != 'schedule' || github.repository == 'BerriAI/litellm') && + (github.event.schedule == '0 8 * * 6' || + (github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'record')) + runs-on: ubuntu-latest + timeout-minutes: 45 + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-e2e-record-replay + LITELLM_LOCAL_MODEL_COST_MAP: "True" + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Start the proxy + run: | + nohup uv run --no-sync litellm --config tests/e2e/gateway/record_replay_ci_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Record the replayable e2e lane + env: + E2E_FIXTURE_MODE: record + run: | + uv run --no-sync pytest tests/e2e -m replayable --reruns 0 -v --tb=short -rA + + - name: Pack the fixture bundle + run: | + .github/scripts/e2e_pack_fixture_bundle.sh tests/e2e/.fixtures "${RUNNER_TEMP}/bundle/e2e-fixtures.tar.gz" + + - name: Publish the fixture bundle + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: e2e-fixtures-bundle + path: | + ${{ runner.temp }}/bundle/e2e-fixtures.tar.gz + ${{ runner.temp }}/bundle/e2e-fixtures.tar.gz.sha256 + if-no-files-found: error + retention-days: 30 + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log + + replay: + name: "Replay the e2e suite from the pinned bundle with zero egress" + if: >- + (github.event_name != 'schedule' || github.repository == 'BerriAI/litellm') && + (github.event.schedule == '0 8 * * 1-5' || + (github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'replay')) + runs-on: ubuntu-latest + timeout-minutes: 45 + permissions: + contents: read + actions: read + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U llmproxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + LITELLM_MASTER_KEY: sk-e2e-record-replay + LITELLM_LOCAL_MODEL_COST_MAP: "True" + GH_TOKEN: ${{ github.token }} + OPENAI_API_KEY: sk-replay-must-never-reach-a-provider + ANTHROPIC_API_KEY: sk-ant-replay-must-never-reach-a-provider + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Set up uv + uses: ./.github/actions/setup-uv-with-retries + with: + version: "0.10.9" + + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + - name: Install dependencies + run: | + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + + - name: Generate Prisma client + run: | + uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma + + - name: Fetch the pinned fixture bundle by digest + env: + BASE_BRANCH: ${{ github.ref_name }} + run: | + .github/scripts/e2e_fetch_fixture_bundle.sh \ + "${GITHUB_REPOSITORY}" \ + e2e-fixtures-bundle \ + "${BASE_BRANCH}" \ + tests/e2e/.fixtures + + - name: Start the proxy + run: | + nohup uv run --no-sync litellm --config tests/e2e/gateway/record_replay_ci_config.yml --port 4000 > proxy.log 2>&1 & + for _ in $(seq 1 90); do + if curl -fs http://localhost:4000/health/liveliness > /dev/null; then + exit 0 + fi + sleep 2 + done + echo "proxy never became live" + tail -n 100 proxy.log + exit 1 + + - name: Start the egress sentinel + run: | + # shellcheck disable=SC2024 # the log redirect is deliberately the runner user's, so a later non-sudo cat can read it + sudo python3 .github/scripts/e2e_egress_sentinel.py serve \ + --host api.openai.com \ + --host api.anthropic.com \ + --hits-file "${RUNNER_TEMP}/egress-hits.jsonl" \ + --ready-file "${RUNNER_TEMP}/egress-ready" \ + --pid-file "${RUNNER_TEMP}/egress.pid" \ + > "${RUNNER_TEMP}/egress-sentinel.log" 2>&1 & + for _ in $(seq 1 30); do + if [[ -f "${RUNNER_TEMP}/egress-ready" ]]; then + cat "${RUNNER_TEMP}/egress-sentinel.log" + exit 0 + fi + sleep 1 + done + echo "egress sentinel never became ready" + cat "${RUNNER_TEMP}/egress-sentinel.log" + exit 1 + + - name: Replay the replayable e2e lane + env: + E2E_FIXTURE_MODE: replay + run: | + uv run --no-sync pytest tests/e2e -m replayable --reruns 0 -v --tb=short -rA + + - name: Stop the egress sentinel and assert zero provider egress + if: always() + run: | + if [[ -f "${RUNNER_TEMP}/egress.pid" ]]; then + sudo kill -TERM "$(cat "${RUNNER_TEMP}/egress.pid")" 2>/dev/null || true + sleep 2 + fi + python3 .github/scripts/e2e_egress_sentinel.py assert-empty --hits-file "${RUNNER_TEMP}/egress-hits.jsonl" + + - name: Show proxy log on failure + if: failure() + run: tail -n 300 proxy.log diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index d798df4c3a4..bb04563c1a8 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -23,6 +23,8 @@ on: - tests/proxy_migration_tests/** - uv.lock - ui/litellm-dashboard/package-lock.json + - ui/Dockerfile + - ui/nginx.conf - .github/workflows/image-scan.yml schedule: - cron: "41 6 * * *" @@ -185,6 +187,35 @@ jobs: python -m pip install "pytest==9.0.3" python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v + ui-image: + name: ui-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 UI image + run: docker build -f ui/Dockerfile -t litellm-ui-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify the UI serves offline as an arbitrary uid with a read-only root fs + env: + LITELLM_IMAGE: litellm-ui-scan:${{ github.sha }} + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_ui_image_serves_offline.py -v + backend-image: name: backend-image runs-on: ubuntu-latest diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index 2a832d1956e..c112bf2bb22 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -131,6 +131,9 @@ jobs: - name: check_e2e_no_raw_requests run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py + - name: check_migrations_no_data_rewrites + run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py + - name: memory_test run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index a7c67f2b35d..2dfca3d308f 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -211,7 +211,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: proxy-extras artifact-name: proxy-extras @@ -219,7 +219,7 @@ jobs: workers: 2 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: enterprise-package artifact-name: enterprise-package @@ -227,7 +227,7 @@ jobs: workers: 4 reruns: 2 timeout-minutes: 20 - job-timeout-minutes: 55 + job-timeout-minutes: 60 - shard: responses-caching-types artifact-name: responses-caching-types diff --git a/CLAUDE.md b/CLAUDE.md index b3383b4a895..b7bd06954d0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -79,6 +79,8 @@ Do not put names of customers or customer company names in code, PR descriptions CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI +Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: ` + Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): - Composition over inheritance diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 664e1669834..19c28716025 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19955 + "limit": 19949 }, "reportArgumentType": { "limit": 2566 @@ -54,7 +54,7 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5663 + "limit": 5661 }, "reportMissingTypeArgument": { "limit": 15555 @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1822 + "limit": 1810 }, "reportRedeclaration": { "limit": 8 @@ -105,10 +105,10 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 39011 + "limit": 39009 }, "reportUnknownParameterType": { - "limit": 19885 + "limit": 19883 }, "reportUnknownVariableType": { "limit": 30569 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 998d225a317..06ba72d84b3 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -428,9 +428,11 @@ ui: 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. + # writes its pid, cache, and proxy temp files under /tmp, so it boots as + # any (arbitrary, non-root) uid; `securityContext.readOnlyRootFilesystem: + # true` here needs an emptyDir volume mounted over /tmp. Images before + # the /tmp move instead need emptyDirs over /var/cache/nginx and /run to + # run as a non-root uid at all. podLabels: {} podSecurityContext: {} securityContext: {} diff --git a/litellm/__init__.py b/litellm/__init__.py index e95b553c5d4..a520d284906 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -199,6 +199,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = ( None # Fields to exclude from StandardLoggingPayload before callbacks receive it ) log_raw_request_response: bool = False +log_client_error_tracebacks: bool = False request_correlation_in_logs: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False @@ -1628,6 +1629,9 @@ if TYPE_CHECKING: AmazonMantleMessagesConfig as AmazonMantleMessagesConfig, ) from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig + from .llms.together_ai.chat.transformation import ( + TogetherAIChatConfig as TogetherAIChatConfig, + ) from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig as VertexGeminiConfig, @@ -1801,6 +1805,9 @@ if TYPE_CHECKING: from .llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig, ) + from .llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig as VertexAIInteractionsConfig, + ) from .llms.openai.chat.o_series_transformation import ( OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config, diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 89c72acc06d..1c833256598 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -177,6 +177,7 @@ LLM_CONFIG_NAMES: Final = ( "AmazonAnthropicClaudeMessagesConfig", "AmazonMantleMessagesConfig", "TogetherAIConfig", + "TogetherAIChatConfig", "NLPCloudConfig", "VertexGeminiConfig", "GoogleAIStudioGeminiConfig", @@ -242,6 +243,7 @@ LLM_CONFIG_NAMES: Final = ( "OpenRouterResponsesAPIConfig", "BedrockMantleResponsesAPIConfig", "GoogleAIStudioInteractionsConfig", + "VertexAIInteractionsConfig", "OpenAIOSeriesConfig", "AnthropicSkillsConfig", "BaseSkillsAPIConfig", @@ -740,6 +742,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "AmazonMantleMessagesConfig", ), "TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"), + "TogetherAIChatConfig": ( + ".llms.together_ai.chat.transformation", + "TogetherAIChatConfig", + ), "NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"), "VertexGeminiConfig": ( ".llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini", @@ -977,6 +983,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.gemini.interactions.transformation", "GoogleAIStudioInteractionsConfig", ), + "VertexAIInteractionsConfig": ( + ".llms.vertex_ai.interactions.transformation", + "VertexAIInteractionsConfig", + ), "OpenAIOSeriesConfig": ( ".llms.openai.chat.o_series_transformation", "OpenAIOSeriesConfig", diff --git a/litellm/a2a_protocol/card_resolver.py b/litellm/a2a_protocol/card_resolver.py index d14d892256b..25f2e1a9a0d 100644 --- a/litellm/a2a_protocol/card_resolver.py +++ b/litellm/a2a_protocol/card_resolver.py @@ -4,6 +4,7 @@ Custom A2A Card Resolver for LiteLLM. Extends the A2A SDK's card resolver to support multiple well-known paths. """ +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger @@ -48,6 +49,43 @@ def is_localhost_or_internal_url(url: str | None) -> bool: return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS) +_CANONICAL_PROTOCOL_BINDINGS: Final = MappingProxyType( + { + "jsonrpc": "JSONRPC", + "http+json": "HTTP+JSON", + "grpc": "GRPC", + } +) + +_LEGACY_PROTOCOL_VERSION: Final = "0.3" + + +def normalize_agent_card_interfaces(agent_card: "AgentCard") -> "AgentCard": + """ + Canonicalize the supported interfaces of spec-adjacent agent cards. + + Some A2A servers (e.g. LangGraph Platform) serve agent cards with lowercase + bindings like "jsonrpc", but a2a-sdk's ClientFactory matches bindings + case-sensitively against its uppercase TransportProtocol constants and fails + with "no compatible transports found." for spec-adjacent casings. + + The same servers also speak the A2A 0.3 JSON dialect ("kind"-discriminated + payloads) while declaring protocolVersion "1.0", which a2a-sdk's strict v1 + proto parsing rejects. A mis-cased binding fingerprints such a server, so its + declared version is downgraded to 0.3 to route the SDK's ClientFactory onto + its v0.3 compat transport, which speaks that dialect. + """ + normalized: Final = type(agent_card)() + normalized.CopyFrom(agent_card) + for interface in normalized.supported_interfaces: + canonical: str | None = _CANONICAL_PROTOCOL_BINDINGS.get(interface.protocol_binding.lower()) + if canonical is None or canonical == interface.protocol_binding: + continue + interface.protocol_binding = canonical + interface.protocol_version = _LEGACY_PROTOCOL_VERSION + return normalized + + def get_agent_card_url(agent_card: "AgentCard") -> str | None: """Return the agent endpoint URL from the resolved SDK card.""" url: Final = getattr(agent_card, "url", None) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 1c6ebf0b95c..56b8089b0af 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -73,6 +73,7 @@ except ImportError: from litellm.a2a_protocol.card_resolver import ( LiteLLMA2ACardResolver, get_agent_card_url, + normalize_agent_card_interfaces, ) from litellm.a2a_protocol.exception_mapping_utils import ( handle_a2a_localhost_retry, @@ -782,13 +783,17 @@ async def create_a2a_client( if extra_headers: verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys())) + resolver: Final = A2ACardResolver(httpx_client=httpx_client, base_url=base_url) + agent_card: Final = normalize_agent_card_interfaces( + await resolver.get_agent_card(http_kwargs={"headers": extra_headers} if extra_headers else None) + ) + a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall] - base_url, + agent_card, client_config=ClientConfig( # pyright: ignore[reportOptionalCall] httpx_client=httpx_client, streaming=streaming, ), - resolver_http_kwargs={"headers": extra_headers} if extra_headers else None, ) # Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse # the configured httpx client and this agent's headers without excavating @@ -799,9 +804,7 @@ async def create_a2a_client( if extra_headers else None ) - agent_card: Final = getattr(a2a_client, "_card", None) - if agent_card is not None: - a2a_client._litellm_agent_card = agent_card + a2a_client._litellm_agent_card = agent_card verbose_logger.info("A2A client created for %s", base_url) diff --git a/litellm/caching/redis_cluster_node_isolation.py b/litellm/caching/redis_cluster_node_isolation.py index 8b0c120e80c..ae8c78709d9 100644 --- a/litellm/caching/redis_cluster_node_isolation.py +++ b/litellm/caching/redis_cluster_node_isolation.py @@ -18,6 +18,14 @@ already does when one of its pooled connections errors), leaving every other nod connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-covered, retry-exhaustion) is unchanged from upstream, since those already carry real evidence the topology changed. + +redis-py 8.x fixed this upstream with gentler machinery than this override's +``node.disconnect()`` (which also kills connections other coroutines are mid-operation +on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per +killed connection): it marks in-use connections for reconnect only after their current +operation completes, disconnects only the idle pooled ones, and defers reinitialization +to the outer retry loop. When the installed ``ClusterNode`` has that per-connection +recovery API, the factory returns the base ``RedisCluster`` unmodified. """ import asyncio @@ -72,8 +80,16 @@ class _ClusterAttrs(Protocol): _VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"}) -def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: - """Builds the ``RedisCluster`` subclass with the per-node isolation fix. +def get_litellm_async_redis_cluster_class( + cluster_node_class: type | None = None, +) -> type["_AsyncRedisClusterType"]: + """Returns the base ``RedisCluster`` when the installed redis-py already recovers a + node-level connection error per-connection (8.x+), else builds the ``RedisCluster`` + subclass with the per-node isolation fix for older versions whose upstream branch + tears down the whole cluster client. + + ``cluster_node_class`` exists for dependency injection in tests; production callers + leave it unset and the installed ``ClusterNode`` is used. Imported lazily because this module is reachable from a base ``import litellm`` while redis is not a base dependency. Cheap to call repeatedly: the underlying redis @@ -81,7 +97,10 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: """ import redis from redis.asyncio.cluster import ( - RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin + ClusterNode as _AsyncClusterNode, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin + ) + from redis.asyncio.cluster import ( + RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # same stale-stub gap as the import above ) from redis.cluster import get_node_name from redis.commands import READ_COMMANDS @@ -98,6 +117,15 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: from redis.exceptions import ConnectionError as _RedisConnectionError from redis.exceptions import TimeoutError as _RedisTimeoutError + node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode + if hasattr(node_class, "update_active_connections_for_reconnect"): + verbose_logger.debug( + "redis-py %s recovers a node-level connection error per-connection upstream; " + "using the base RedisCluster without litellm's node-isolation override.", + redis.__version__, + ) + return _BaseAsyncRedisCluster + if redis.__version__ not in _VERIFIED_REDIS_VERSIONS: verbose_logger.warning( "redis-py %s is not in the set this cluster-teardown-storm fix was verified " diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 6103b1bf484..17815976b4a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,7 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args from openai.types.responses.custom_tool_param import CustomToolParam from openai.types.responses.response_input_param import ( @@ -21,6 +21,9 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + responses_reasoning_item_from_thinking_blocks, +) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, @@ -32,6 +35,7 @@ from litellm.responses.sse_output_recovery import ( ) from litellm.responses.utils import normalize_responses_api_stream_options from litellm.types.llms.openai import ( + REASONING_EFFORT, ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionToolCallChunk, @@ -85,6 +89,22 @@ def _get_reasoning_items( return [] +def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: # mutable-ok: API message payload + """Reasoning input items for an assistant message. + + Stored reasoning items win because they carry an id the Responses API minted; thinking + blocks are the fallback for turns that arrived over another API surface. + """ + items: Final = _get_reasoning_items(msg) + stored: Final = [_reasoning_item_to_response_input(item) for item in items] # mutable-ok: API message payload + if stored: + return stored + raw_blocks: Final = msg.get("thinking_blocks") or () + blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json + from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks) + return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload + + def _build_reasoning_item( item_id: str, encrypted_content: str | None, @@ -372,8 +392,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) ) 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)) + input_items.extend(_reasoning_input_items(msg)) + if content: + input_items.append( + { # mutable-ok: API message payload + "type": "message", + "role": "assistant", + "content": self._convert_content_to_responses_format(content, "assistant"), + } + ) for tool_call in tool_calls: function = tool_call.get("function") custom = tool_call.get("custom") @@ -400,15 +427,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): raise ValueError(f"tool call not supported: {tool_call}") elif content is not None: if role == "assistant": - for r_item in _get_reasoning_items(msg): - input_items.append(_reasoning_item_to_response_input(r_item)) + input_items.extend(_reasoning_input_items(msg)) input_items.append( - { + { # mutable-ok: API message payload "type": "message", "role": role, "content": self._convert_content_to_responses_format(content, cast(str, role)), } ) + elif role == "assistant": + input_items.extend(_reasoning_input_items(msg)) return input_items, instructions @@ -1086,22 +1114,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" ) - # If string is passed, map with optional summary based on flag/env var - if reasoning_effort == "none": - return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") - elif reasoning_effort == "high": - return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high") - elif reasoning_effort == "xhigh": - return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") - elif reasoning_effort == "medium": + if reasoning_effort in get_args(REASONING_EFFORT): return ( - Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium") - ) - elif reasoning_effort == "low": - return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low") - elif reasoning_effort == "minimal": - return ( - Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal") + Reasoning(effort=reasoning_effort, summary="detailed") + if auto_summary_enabled + else Reasoning(effort=reasoning_effort) ) return None diff --git a/litellm/constants.py b/litellm/constants.py index aaaddd063e7..765bbfe1e54 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -48,6 +48,7 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = ( # Data URIs exceeding this are replaced with a size placeholder. # Set to 0 to disable truncation. MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)) +REDACTED_BY_LITELLM: Final = "redacted-by-litellm" MAX_STRING_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", 4096) @@ -146,6 +147,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [ "x-litellm-adaptive-router-model", "x-litellm-applied-guardrails", "x-litellm-guardrail-scan-id", + "x-litellm-cache-key", ] # Gemini model-specific minimal thinking budget constants @@ -749,6 +751,7 @@ openai_compatible_endpoints: Final[list] = [ "api.groq.com/openai/v1", "https://integrate.api.nvidia.com/v1", "api.deepseek.com/v1", + "api.together.ai/v1", "api.together.xyz/v1", "app.empower.dev/api/v1", "https://api.friendli.ai/serverless/v1", @@ -1563,6 +1566,19 @@ STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, int(os.getenv("STALE_OBJECT_CLEA # installations with large numbers of stale managed objects). _batch_polling_env: Final = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() PROXY_BATCH_POLLING_ENABLED: Final = _batch_polling_env == "true" +BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS: Final = float( + os.getenv("BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS", "5") +) +BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS: Final = float( + os.getenv("BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS", "60") +) +BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS: Final = float( + os.getenv("BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", "3600") +) +_background_interaction_cost_polling_env: Final = os.getenv( + "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true" +).lower() +BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_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) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 8f7cd09d364..6536941a094 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + InteractionsUsageObjectTransformation, TranscriptionUsageObjectTransformation, ) from litellm.litellm_core_utils.llm_cost_calc.utils import ( @@ -150,6 +151,7 @@ _VIDEO_CALL_TYPES: Final = frozenset( } ) + _SPEECH_CALL_TYPES: Final = frozenset( { CallTypes.speech.value, @@ -912,6 +914,8 @@ def _get_usage_object( usage_obj, ) ) + elif isinstance(usage_obj, dict) and InteractionsUsageObjectTransformation.is_interactions_usage_object(usage_obj): + return InteractionsUsageObjectTransformation.transform_interactions_usage_object(usage_obj) elif isinstance(usage_obj, dict): return Usage(**usage_obj) elif isinstance(usage_obj, BaseModel): @@ -1288,6 +1292,10 @@ def completion_cost( ) if tr_usage is not None: _usage = tr_usage.model_dump() + elif InteractionsUsageObjectTransformation.is_interactions_usage_object(_usage): + _usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + _usage + ).model_dump() else: _usage = _usage @@ -1372,23 +1380,36 @@ def completion_cost( 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: - _metadata = _litellm_params.get("metadata", {}) or {} - _video_model_info = _metadata.get("model_info", None) + _video_model_info = next( + ( + model_info + for _metadata_key in ("metadata", "litellm_metadata") + if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info")) + is not None + ), + None, + ) usage_obj = getattr(completion_response, "usage", None) duration_seconds: float | None = None video_resolution: str | None = None + provider_reported_cost: float | None = None if completion_response is not None and usage_obj: # Handle both dict and Pydantic Usage object if isinstance(usage_obj, dict): duration_seconds = usage_obj.get("duration_seconds", None) _vr = usage_obj.get("video_resolution", None) + provider_reported_cost = usage_obj.get("provider_reported_cost_usd", None) else: duration_seconds = getattr(usage_obj, "duration_seconds", None) _vr = getattr(usage_obj, "video_resolution", None) + provider_reported_cost = getattr(usage_obj, "provider_reported_cost_usd", None) if _vr is not None: video_resolution = str(_vr).strip().lower() + if _video_model_info is None and provider_reported_cost is not None: + return float(provider_reported_cost) + if duration_seconds is not None: # Calculate cost based on video duration using video-specific cost calculation from litellm.llms.openai.cost_calculation import ( diff --git a/litellm/images/main.py b/litellm/images/main.py index ae4818b1967..1688087c2da 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -19,6 +19,7 @@ from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT from litellm.exceptions import LiteLLMUnknownProvider from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_request_utils import flatten_form_field_values from litellm.litellm_core_utils.mock_functions import mock_image_generation from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -422,24 +423,32 @@ def image_generation( aimg_generation=aimg_generation, ) elif custom_llm_provider == "azure_ai": - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, + ) api_base = AzureFoundryModelInfo.get_api_base(api_base) api_key = AzureFoundryModelInfo.get_api_key(api_key) if extra_headers is not None: optional_params["extra_headers"] = extra_headers - default_headers = { + caller_header_names = frozenset(name.lower() for name in headers) + caller_set_auth = "api-key" in caller_header_names or "authorization" in caller_header_names + auth_headers = ( + headers + if caller_set_auth + else get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params_dict, + api_key_header="api-key", + ) + ) + request_headers: Final = { "Content-Type": "application/json", + **auth_headers, + **headers, } - # Only add api-key header if api_key is not None - # Azure AD authentication will use Authorization header instead - if api_key is not None: - default_headers["api-key"] = api_key - - for k, v in default_headers.items(): - if k not in headers: - headers[k] = v model_response = azure_chat_completions.image_generation( model=model, @@ -455,7 +464,7 @@ def image_generation( api_version=api_version, aimg_generation=aimg_generation, client=client, - headers=headers, + headers=request_headers, litellm_params=litellm_params_dict, ) elif ( @@ -846,6 +855,18 @@ def image_edit( additional_drop_params=kwargs.get("additional_drop_params"), ) + if ( + custom_llm_provider == "openai" + or custom_llm_provider == "azure" + or custom_llm_provider in litellm.openai_compatible_providers + ): + image_edit_request_params.update( + flatten_form_field_values( + non_default_params, + extra_body if isinstance(extra_body, dict) else None, + ) + ) + # Pre Call logging litellm_logging_obj.update_from_kwargs( kwargs=kwargs, @@ -995,6 +1016,9 @@ async def aimage_edit( response_format=response_format, size=size, user=user, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, timeout=timeout, custom_llm_provider=custom_llm_provider, **kwargs, diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 89f1a30c143..9607eccef52 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -168,17 +168,20 @@ class LangsmithLogger(CustomBatchLogger): return outputs def _ensure_required_ids(self, data: dict, run_id: str | None): + resolved_id: Final = run_id or str(uuid.uuid4()) if "id" not in data or data["id"] is None: - run_id = str(uuid.uuid4()) - data["id"] = run_id + data["id"] = resolved_id - if "trace_id" not in data or data["trace_id"] is None: - if run_id is not None and isinstance(run_id, str): - data["trace_id"] = run_id + # LangSmith rejects the whole ingest batch unless a root run's trace_id + # equals the run id embedded in the first segment of dotted_order + posts_as_root: Final = ("parent_run_id" not in data or data["parent_run_id"] is None) and ( + "dotted_order" not in data or data["dotted_order"] is None + ) + if posts_as_root or "trace_id" not in data or data["trace_id"] is None: + data["trace_id"] = resolved_id if "dotted_order" not in data or data["dotted_order"] is None: - if run_id is not None and isinstance(run_id, str): - data["dotted_order"] = self.make_dot_order(run_id=run_id) + data["dotted_order"] = self.make_dot_order(run_id=resolved_id) def _prepare_log_data( self, @@ -193,6 +196,11 @@ class LangsmithLogger(CustomBatchLogger): metadata = _litellm_params.get("metadata", {}) or {} fields: Final = self._extract_metadata_fields(metadata, credentials) + # the proxy header fan-out mirrors one value into both keys, and LangSmith + # rejects the whole ingest batch when run-body session_id is not an + # existing tracer-session uuid + if fields["session_id"] == fields["trace_id"]: + fields["session_id"] = None verbose_logger.debug( "Langsmith Logging - project_name: %s, run_name %s", fields["project_name"], fields["run_name"] ) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index d52bcda525f..10b4c0dd433 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -11,6 +11,7 @@ import time from collections.abc import Mapping from datetime import datetime from typing import Final, cast +from urllib.parse import quote import litellm from litellm._logging import print_verbose, verbose_logger @@ -206,6 +207,23 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): params.get("s3_sse_kms_key_id") or s3_sse_kms_key_id, ) + def _build_object_url(self, s3_object_key: str) -> str: + """ + Build the exact URL that is both signed and sent, with the key percent-encoded once. + + S3SigV4Auth signs the path verbatim while S3 canonicalizes the received path with reserved + characters encoded, so an unencoded `=`, `+`, `&`, `#`, `?`, `%` or space in the key makes + the two signatures disagree (403 SignatureDoesNotMatch). + """ + encoded_key: Final = quote(s3_object_key, safe="/") + if self.s3_endpoint_url and self.s3_bucket_name: + if self.s3_use_virtual_hosted_style: + endpoint_host: Final = self.s3_endpoint_url.replace("https://", "").replace("http://", "") + protocol: Final = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" + return f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{encoded_key}" + return f"{self.s3_endpoint_url}/{self.s3_bucket_name}/{encoded_key}" + return f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{encoded_key}" + def _sse_headers(self) -> Mapping[str, str]: candidates: Final = { "x-amz-server-side-encryption": self.s3_server_side_encryption, @@ -292,7 +310,6 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): import base64 import hashlib - import requests from botocore.auth import S3SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: @@ -316,18 +333,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): verbose_logger.debug("s3_v2 logger - uploading data to s3 - %s", batch_logging_element.s3_object_key) verbose_logger.debug("s3_v2 logger - s3_verify setting: %s", self.s3_verify) - # Prepare the URL - url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}" - - if self.s3_endpoint_url and self.s3_bucket_name: - if self.s3_use_virtual_hosted_style: - # Virtual-hosted-style: bucket.endpoint/key - endpoint_host: Final = self.s3_endpoint_url.replace("https://", "").replace("http://", "") - protocol: Final = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" - url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" - else: - # Path-style: endpoint/bucket/key - url = self.s3_endpoint_url + "/" + self.s3_bucket_name + "/" + batch_logging_element.s3_object_key + url: Final = self._build_object_url(batch_logging_element.s3_object_key) # Convert JSON to string json_string: Final = safe_dumps(batch_logging_element.payload) @@ -348,29 +354,19 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", **self._sse_headers(), } - req: Final = requests.Request("PUT", url, data=json_string, headers=headers) - prepped: Final = req.prepare() # Sign the request - aws_request: Final = AWSRequest( - method=prepped.method, - url=prepped.url, - data=prepped.body, - headers=prepped.headers, - ) + aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) - # Use prepared URL so path segments match SigV4 canonical request (e.g. %20 for spaces). - request_url: Final = prepped.url or url - # Make the request with retry for transient S3 errors (500/503) max_retries: Final = 3 for attempt in range(max_retries): - response = await self.async_httpx_client.put(request_url, data=json_string, headers=signed_headers) + response = await self.async_httpx_client.put(url, data=json_string, headers=signed_headers) if response.status_code in (500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( @@ -478,7 +474,6 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): import base64 import hashlib - import requests from botocore.auth import S3SigV4Auth from botocore.awsrequest import AWSRequest from botocore.credentials import Credentials @@ -493,18 +488,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): aws_region_name=self.s3_region_name, ) - # Prepare the URL - url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}" - - if self.s3_endpoint_url and self.s3_bucket_name: - if self.s3_use_virtual_hosted_style: - # Virtual-hosted-style: bucket.endpoint/key - endpoint_host: Final = self.s3_endpoint_url.replace("https://", "").replace("http://", "") - protocol: Final = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" - url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}" - else: - # Path-style: endpoint/bucket/key - url = self.s3_endpoint_url + "/" + self.s3_bucket_name + "/" + batch_logging_element.s3_object_key + url: Final = self._build_object_url(batch_logging_element.s3_object_key) # Convert JSON to string json_string: Final = safe_dumps(batch_logging_element.payload) @@ -525,32 +509,22 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", **self._sse_headers(), } - req: Final = requests.Request("PUT", url, data=json_string, headers=headers) - prepped: Final = req.prepare() # Sign the request - aws_request: Final = AWSRequest( - method=prepped.method, - url=prepped.url, - data=prepped.body, - headers=prepped.headers, - ) + aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers) aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name) S3SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) - # Use prepared URL so path segments match SigV4 canonical request (e.g. %20 for spaces). - request_url: Final = prepped.url or url - httpx_client: Final = _get_httpx_client( params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None) ) # Make the request with retry for transient S3 errors (500/503) max_retries: Final = 3 for attempt in range(max_retries): - response = httpx_client.put(request_url, data=json_string, headers=signed_headers) + response = httpx_client.put(url, data=json_string, headers=signed_headers) if response.status_code in (500, 503) and attempt < max_retries - 1: wait_time = 2**attempt # 1s, 2s verbose_logger.warning( @@ -582,7 +556,6 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): try: import hashlib - import requests from botocore.auth import S3SigV4Auth from botocore.awsrequest import AWSRequest except ImportError: @@ -607,18 +580,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): verbose_logger.debug("s3_v2 logger - downloading data from s3 - %s", s3_object_key) - # Prepare the URL - url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{s3_object_key}" - - if self.s3_endpoint_url and self.s3_bucket_name: - if self.s3_use_virtual_hosted_style: - # Virtual-hosted-style: bucket.endpoint/key - endpoint_host: Final = self.s3_endpoint_url.replace("https://", "").replace("http://", "") - protocol: Final = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" - url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{s3_object_key}" - else: - # Path-style: endpoint/bucket/key - url = self.s3_endpoint_url + "/" + self.s3_bucket_name + "/" + s3_object_key + url: Final = self._build_object_url(s3_object_key) # Prepare the request for GET operation # For GET requests, we need x-amz-content-sha256 with hash of empty string @@ -626,22 +588,15 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): headers: Final = { "x-amz-content-sha256": empty_string_hash, } - req: Final = requests.Request("GET", url, headers=headers) - prepped: Final = req.prepare() # Sign the request - aws_request: Final = AWSRequest( - method=prepped.method, - url=prepped.url, - headers=prepped.headers, - ) + aws_request: Final = AWSRequest(method="GET", url=url, headers=headers) S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request) # Prepare the signed headers signed_headers: Final = dict(aws_request.headers.items()) - request_url: Final = prepped.url or url - response: Final = await self.async_httpx_client.get(request_url, headers=signed_headers) + response: Final = await self.async_httpx_client.get(url, headers=signed_headers) if response.status_code != 200: verbose_logger.exception("S3 object not found, saw response=", response.text) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py new file mode 100644 index 00000000000..51325354e7d --- /dev/null +++ b/litellm/interactions/background_cost_polling.py @@ -0,0 +1,313 @@ +""" +Cost tracking for background interactions. + +A create request with ``background=true`` returns ``in_progress`` with no +usage block, and GET polls are deliberately never billed (billing them would +double-charge every poll; the GET response also does not echo ``background``, +so a poll cannot be told apart from a re-fetch of an already-billed +interaction). The create call is therefore the only place that can own +billing: it schedules a poll task that fetches the interaction until it +reaches a terminal status and logs the final usage as a single success event +attributed to the original request. + +``requires_action`` is terminal for the interaction it names. The API has no +operation that resumes one: a caller answers a tool request by creating a new +interaction whose ``previous_interaction_id`` points at it, and that new +interaction bills itself. The paused interaction keeps the tokens it already +spent producing the tool request, so it is billed and settled where it stops +rather than polled until the timeout, which would both lose that usage and +hold its budget reservation open for the whole timeout window. + +Deleting an interaction makes every subsequent poll fail, which would let a +caller retrieve the completed output themselves and then delete it before the +poll task settles, leaving the work unbilled and the budget reservation +refunded at the poll timeout. ``adelete`` therefore settles any pending poll +for the interaction before dispatching the delete: it fetches the current +state with the create's credentials, bills it if it is terminal with usage, +and releases the reservation otherwise. A settlement gate on the create's +logging object makes the poll task and the delete path mutually exclusive, so +the interaction is billed exactly once no matter who settles first. +""" + +import asyncio +from collections.abc import Awaitable, Callable, Iterator, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, TypeAlias + +from litellm._logging import verbose_logger +from litellm.constants import ( + BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS, + BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS, + BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS, + BACKGROUND_INTERACTION_COST_POLLING_ENABLED, +) +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.types.interactions import InteractionsAPIResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_TERMINAL_STATUSES: Final = frozenset( + {"completed", "failed", "cancelled", "incomplete", "budget_exceeded", "requires_action"} +) + +_POLLABLE_STATUSES: Final = frozenset({"in_progress", "queued"}) + +_STATUSES_THAT_PRODUCED_OUTPUT: Final = frozenset({"completed", "requires_action"}) + + +@dataclass(frozen=True, slots=True) +class BackgroundInteractionPollContext: + interaction_id: str + custom_llm_provider: str + logging_obj: "LiteLLMLoggingObj" + api_key: str | None = None + api_base: str | None = None + initial_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS + max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS + timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS + + +FetchInteraction: TypeAlias = Callable[[BackgroundInteractionPollContext], Awaitable[InteractionsAPIResponse]] + + +async def _fetch_interaction(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: + from litellm.interactions import aget + + return await aget( + interaction_id=context.interaction_id, + custom_llm_provider=context.custom_llm_provider, + api_key=context.api_key, + api_base=context.api_base, + **{ + "no-log": True + }, # mutable-ok: "no-log" is not a valid identifier, so it can only be passed through a mapping + ) + + +def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[float]: + elapsed = 0.0 + interval = initial + while interval > 0 and elapsed + interval <= timeout: + yield interval + elapsed += interval + interval = min(interval * 2, maximum) + + +_SETTLED_KEY = "background_interaction_settled" + + +def _is_settled(logging_obj: "LiteLLMLoggingObj") -> bool: + return logging_obj.model_call_details.get(_SETTLED_KEY) is True + + +def _claim_settlement(logging_obj: "LiteLLMLoggingObj") -> bool: + """ + Exactly-once gate between the poll task and the delete-time settlement: + both run on the same event loop and neither awaits between reading and + setting the flag, so whichever claims first owns billing or release. + """ + if _is_settled(logging_obj): + return False + logging_obj.model_call_details[_SETTLED_KEY] = True # rebind-ok: both settlers must see the same settlement flag + return True + + +async def poll_and_log_background_interaction_cost( + context: BackgroundInteractionPollContext, + fetch_interaction: FetchInteraction = _fetch_interaction, +) -> None: + last_seen_status: str | None = None + for interval in _poll_intervals( + initial=context.initial_interval_seconds, + maximum=context.max_interval_seconds, + timeout=context.timeout_seconds, + ): + await asyncio.sleep(interval) + if _is_settled(context.logging_obj): + return + try: + response = await fetch_interaction(context) + except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop + verbose_logger.debug( + "Background interaction cost poll for %s failed, will retry: %s", + context.interaction_id, + e, + ) + continue + last_seen_status = response.status + if response.status not in _TERMINAL_STATUSES: + continue + if not _claim_settlement(context.logging_obj): + return + if response.usage is not None: + await _bill_settled_interaction(logging_obj=context.logging_obj, response=response) + else: + await _release_open_budget_reservation(logging_obj=context.logging_obj) + return + if not _claim_settlement(context.logging_obj): + return + if last_seen_status is not None and last_seen_status not in _POLLABLE_STATUSES: + verbose_logger.error( + "Gave up cost polling for background interaction %s after %ss: its last status %r is in neither " + "the pollable nor the terminal set, so this proxy never learned how to settle it and its usage " + "will not be tracked", + context.interaction_id, + context.timeout_seconds, + last_seen_status, + ) + else: + verbose_logger.warning( + "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", + context.interaction_id, + context.timeout_seconds, + ) + await _release_open_budget_reservation(logging_obj=context.logging_obj) + + +async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None: + """ + The proxy keeps the pre-call budget reservation open for an in-progress + background interaction so concurrent creates cannot stack past the budget. + The completion success event reconciles it to the actual cost; when the + interaction terminates without billable usage (or polling gives up, or it + is deleted before settling), no such event fires, so whoever claims the + settlement must release the reservation here or the spend counters stay + pinned at the estimated cost. + """ + metadata = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + budget_reservation = metadata.get("user_api_key_budget_reservation") + if not isinstance(budget_reservation, dict): + return + + from litellm.proxy.spend_tracking.budget_reservation import release_budget_reservation + + try: + await release_budget_reservation(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # a failed release must not crash the poll task; counters expire via TTL + verbose_logger.exception("Failed to release budget reservation for an unbilled background interaction") + + +async def _bill_settled_interaction(logging_obj: "LiteLLMLoggingObj", response: InteractionsAPIResponse) -> None: + """ + Claiming the settlement makes the claimer solely responsible for the + reservation, and no one retries a claim that is already set. A billing + failure here must therefore release the reservation on its way out, or it + stays pinned at the estimated cost until the whole poll times out. + """ + try: + await logging_obj.async_log_background_interaction_completion(result=response) + except Exception: + await _release_open_budget_reservation(logging_obj=logging_obj) + raise + + +def is_pollable_background_interaction(response: InteractionsAPIResponse) -> bool: + """ + The single gate deciding whether a create's response gets a poll task. + The proxy's success callback defers releasing the budget reservation for + exactly these responses, on the promise that a poll task will settle them, + so a response one site accepts and the other refuses strands its + reservation on the spend counters with nothing left to reconcile it. + + ``queued`` belongs here alongside ``in_progress``. It is the API's + not-started-yet state, so it reaches a terminal status the same way and + needs polling for the same reason: nothing else in the proxy ever bills a + create that came back without usage, so a status missing from both this + set and ``_TERMINAL_STATUSES`` is billed nowhere and alerts nobody. + """ + return response.status in _POLLABLE_STATUSES and bool(response.id) + + +def missing_usage_is_expected(response: InteractionsAPIResponse) -> bool: + """ + Whether a response arriving with no usage block is a normal outcome rather + than lost billing data. An interaction that is still running, or that + stopped at ``failed``, ``cancelled``, ``incomplete`` or ``budget_exceeded``, + has nothing to charge for and should not raise a cost-tracking alarm. + + ``completed`` and ``requires_action`` both mean the model produced output, + so a usage block is always expected with them. If one arrives without it + the charge for real work has been lost, which is precisely what the + proxy's cost-tracking alert exists to surface. + """ + return response.status not in _STATUSES_THAT_PRODUCED_OUTPUT + + +@dataclass(frozen=True, slots=True) +class _ActiveBackgroundPoll: + task: "asyncio.Task[None]" + context: BackgroundInteractionPollContext + + +_ACTIVE_POLLS: dict[str, _ActiveBackgroundPoll] = {} # mutable-ok: asyncio needs strong refs to running poll tasks + + +def _discard_poll(interaction_id: str, task: "asyncio.Task[None]") -> None: + entry = _ACTIVE_POLLS.get(interaction_id) + if entry is not None and entry.task is task: + del _ACTIVE_POLLS[interaction_id] + + +def maybe_schedule_background_interaction_cost_polling( + response: object, + create_kwargs: Mapping[str, object], + custom_llm_provider: str, +) -> "asyncio.Task[None] | None": + from litellm.litellm_core_utils.litellm_logging import Logging + + if not BACKGROUND_INTERACTION_COST_POLLING_ENABLED: + return None + if not isinstance(response, InteractionsAPIResponse): + return None + if not is_pollable_background_interaction(response): + return None + logging_obj = create_kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return None + try: + asyncio.get_running_loop() + except RuntimeError: + return None + api_key = create_kwargs.get("api_key") + api_base = create_kwargs.get("api_base") + context = BackgroundInteractionPollContext( + interaction_id=response.id, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + api_key=api_key if isinstance(api_key, str) else None, + api_base=api_base if isinstance(api_base, str) else None, + ) + task = asyncio.create_task(poll_and_log_background_interaction_cost(context)) + _ACTIVE_POLLS[context.interaction_id] = _ActiveBackgroundPoll(task=task, context=context) + task.add_done_callback( + lambda finished, interaction_id=context.interaction_id: _discard_poll(interaction_id, finished) + ) + return task + + +async def maybe_settle_background_interaction_before_delete( + interaction_id: str, + fetch_interaction: FetchInteraction = _fetch_interaction, +) -> None: + entry = _ACTIVE_POLLS.get(interaction_id) + if entry is None: + return + context = entry.context + try: + response = await fetch_interaction(context) + except Exception as e: # noqa: BLE001 # unfetchable pre-delete state settles by releasing the reservation + verbose_logger.debug( + "Could not fetch background interaction %s before delete, releasing its reservation: %s", + interaction_id, + e, + ) + if _claim_settlement(context.logging_obj): + await _release_open_budget_reservation(logging_obj=context.logging_obj) + return + if not _claim_settlement(context.logging_obj): + return + if response.status in _TERMINAL_STATUSES and response.usage is not None: + await _bill_settled_interaction(logging_obj=context.logging_obj, response=response) + return + await _release_open_budget_reservation(logging_obj=context.logging_obj) diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index 3e8c381fdf7..a2c3d510fae 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -40,6 +40,10 @@ from typing import Any, Final import httpx import litellm +from litellm.interactions.background_cost_polling import ( + maybe_schedule_background_interaction_cost_polling, + maybe_settle_background_interaction_before_delete, +) from litellm.interactions.http_handler import interactions_http_handler from litellm.interactions.utils import ( InteractionsAPIRequestUtils, @@ -171,6 +175,12 @@ async def acreate( else: response = init_response + maybe_schedule_background_interaction_cost_polling( + response=response, + create_kwargs=kwargs, + custom_llm_provider=custom_llm_provider, + ) + return response except Exception as e: raise litellm.exception_type( @@ -462,6 +472,8 @@ async def adelete( loop: Final = asyncio.get_event_loop() kwargs["adelete_interaction"] = True + await maybe_settle_background_interaction_before_delete(interaction_id=interaction_id) + func: Final = partial( delete, interaction_id=interaction_id, diff --git a/litellm/interactions/utils.py b/litellm/interactions/utils.py index 8a1e8836894..3895a85061d 100644 --- a/litellm/interactions/utils.py +++ b/litellm/interactions/utils.py @@ -47,6 +47,13 @@ def get_provider_interactions_api_config( return GoogleAIStudioInteractionsConfig() + if provider in (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value): + from litellm.llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig, + ) + + return VertexAIInteractionsConfig() + return None diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index de1092bc02f..b71ef4c8e06 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -58,6 +58,26 @@ def safe_divide( return numerator / denominator +def is_expected_client_error(exception: BaseException | None) -> bool: + """ + True when the exception maps to an HTTP 4xx status. + + ProxyException stores the status on .code (as a str), HTTPException and + litellm exceptions on .status_code. + """ + if exception is None: + return False + code: Final[object] = getattr(exception, "code", None) + status_code: Final[object] = code if code is not None else getattr(exception, "status_code", None) + if status_code is None or isinstance(status_code, bool): + return False + try: + status: Final = int(str(status_code)) + except ValueError: + return False + return 400 <= status < 500 + + def coerce_token_limit(value: object) -> int | None: """ Coerce a max_input_tokens / max_output_tokens value to an int, treating a diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 4a25eb218c0..4ee726b67de 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -2301,6 +2301,7 @@ def exception_type( or custom_llm_provider == "custom_openai" or custom_llm_provider in litellm.openai_compatible_providers or custom_llm_provider == "mistral" + or custom_llm_provider == "runwayml" ): _map_openai_exception( model=model, diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index e674fc37673..005e94ebe82 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -272,6 +272,14 @@ def get_llm_provider( elif endpoint == "api.deepseek.com/v1": custom_llm_provider = "deepseek" dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY") + elif endpoint == "api.together.ai/v1" or endpoint == "api.together.xyz/v1": + custom_llm_provider = "together_ai" + dynamic_api_key = api_key or ( + get_secret_str("TOGETHER_API_KEY") + or get_secret_str("TOGETHER_AI_API_KEY") + or get_secret_str("TOGETHERAI_API_KEY") + or get_secret_str("TOGETHER_AI_TOKEN") + ) elif endpoint == "ollama.com": custom_llm_provider = "ollama" dynamic_api_key = get_secret_str("OLLAMA_API_KEY") @@ -707,7 +715,7 @@ def _get_openai_compatible_provider_info( dynamic_api_key, ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "together_ai": - api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1" + api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.ai/v1" dynamic_api_key = api_key or ( get_secret_str("TOGETHER_API_KEY") or get_secret_str("TOGETHER_AI_API_KEY") diff --git a/litellm/litellm_core_utils/get_supported_openai_params.py b/litellm/litellm_core_utils/get_supported_openai_params.py index 72f36661f4c..7a16ffe4d85 100644 --- a/litellm/litellm_core_utils/get_supported_openai_params.py +++ b/litellm/litellm_core_utils/get_supported_openai_params.py @@ -172,7 +172,7 @@ def get_supported_openai_params( if request_type == "embeddings": return litellm.JinaAIEmbeddingConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "together_ai": - return litellm.TogetherAIConfig().get_supported_openai_params(model=model) + return litellm.TogetherAIChatConfig().get_supported_openai_params(model=model) elif custom_llm_provider == "databricks": if request_type == "chat_completion": return litellm.DatabricksConfig().get_supported_openai_params(model=model) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3ad4c187b6d..626af7530a4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -62,7 +62,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger -from litellm.litellm_core_utils.core_helpers import reconstruct_model_name +from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( cost_breakdown_with_guardrail, @@ -71,6 +71,9 @@ from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) +from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + InteractionsUsageObjectTransformation, +) from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.redact_messages import ( @@ -83,6 +86,10 @@ from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.containers.main import ContainerObject +from litellm.types.interactions import ( + InteractionsAPIResponse, + InteractionsAPIStreamingResponse, +) from litellm.types.llms.openai import ( AllMessageValues, Batch, @@ -2145,6 +2152,11 @@ class Logging(LiteLLMLoggingBaseClass): or isinstance(logging_result, OpenAIModerationResponse) or isinstance(logging_result, OCRResponse) # OCR or isinstance(logging_result, SearchResponse) # Search API + or ( + isinstance(logging_result, InteractionsAPIResponse) + and logging_result.usage is not None + and self._is_interactions_create_call_type() + ) or isinstance(logging_result, dict) and logging_result.get("object") == "vector_store.search_results.page" or isinstance(logging_result, dict) @@ -2157,6 +2169,87 @@ class Logging(LiteLLMLoggingBaseClass): return True return False + def _is_interactions_create_call_type(self) -> bool: + """ + Only interaction creation is billable. GET polls, deletes, and cancels + also return an ``InteractionsAPIResponse`` (with usage once completed), + so recognizing those would write spend on every poll of a background + interaction. The proxy sets ``call_type`` from its route_type + (``create_interaction``/``acreate_interaction``); the SDK sets it from + the decorated function name (``create``/``acreate``). + + Recognition additionally requires a usage block (checked at the call + site): a ``background=true`` create returns ``in_progress`` without + usage, and billing it would write a $0 spend log under the interaction + id that collides with the row the background poll task writes once the + interaction completes (see + ``litellm.interactions.background_cost_polling``). + """ + return self.call_type in ( + CallTypes.create_interaction.value, + CallTypes.acreate_interaction.value, + "create", + "acreate", + ) + + async def async_log_background_interaction_completion( + self, + result: InteractionsAPIResponse, + ) -> None: + """ + Log the terminal result of a background interaction as a fresh success + event. The create request already ran success logging for its + ``in_progress`` response (no usage, so no cost was tracked); clearing + the dedup flags lets the completed result flow through cost calculation + and spend tracking exactly once, spanning create to completion. + + The poll fetched this body through its own client call, which priced it + against a throwaway logging object holding none of this request's + deployment context: no ``model_info``, no router ``model_id``, no + deployment ``litellm_params``. Keeping that price would bill a + custom-priced deployment at the wrong rate, and it would also satisfy + the "already calculated" shortcut and skip repricing here, leaving the + cost breakdown at the zeros the usage-less create stamped and writing + those zeros to the spend log. Dropping it makes this event price the + settled body itself, against the deployment that served the create. + + The same throwaway call stamped the deployment identity that travels + with the price, so ``model_id`` and ``litellm_model_name`` go with it. + Left in place they overwrite the create's real deployment with the + poll's empty one in the payload every logging integration reads. + """ + settled_hidden_params: Final = getattr(result, "_hidden_params", None) + if isinstance(settled_hidden_params, dict): + for poll_scoped_key in ("response_cost", "model_id", "litellm_model_name"): + settled_hidden_params.pop(poll_scoped_key, None) + self._reset_success_emission_dedupe() + await self.async_success_handler(result=result) + + def _reset_success_emission_dedupe(self) -> None: + """ + Success callbacks dedupe per request, because the sync and async + handlers both fire on some paths and would otherwise report one call + twice. A settled background interaction is a genuinely second success + event on the same request, so every such marker has to be cleared or + the completion, the only event that carries usage and cost, is + discarded as a duplicate of the in-progress create. + """ + self.model_call_details.pop("has_logged_async_success", None) + litellm_params = self.model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + metadata = litellm_params.get("metadata") + if not isinstance(metadata, dict): + return + otel_internal = metadata.get("_otel_internal") + if not isinstance(otel_internal, dict): + return + spans_logged = otel_internal.get("spans_logged") + if not isinstance(spans_logged, dict): + return + for scope in [key for key in spans_logged if isinstance(key, tuple) and key[-1:] == ("success",)]: + del spans_logged[scope] + def _flush_passthrough_collected_chunks_helper( self, raw_bytes: list[bytes], @@ -2282,7 +2375,9 @@ class Logging(LiteLLMLoggingBaseClass): is_sync_request: Final = self._is_sync_litellm_request(litellm_params) try: ## BUILD COMPLETE STREAMED RESPONSE - complete_streaming_response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None = None + complete_streaming_response: ( + ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None + ) = None if "complete_streaming_response" in self.model_call_details: return # break out of this. complete_streaming_response = self._get_assembled_streaming_response( @@ -2768,14 +2863,14 @@ class Logging(LiteLLMLoggingBaseClass): ## BUILD COMPLETE STREAMED RESPONSE if "async_complete_streaming_response" in self.model_call_details: return # break out of this. - complete_streaming_response: Final[ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None] = ( - self._get_assembled_streaming_response( - result=result, - start_time=start_time, - end_time=end_time, - is_async=True, - streaming_chunks=self.streaming_chunks, - ) + complete_streaming_response: Final[ + ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None + ] = self._get_assembled_streaming_response( + result=result, + start_time=start_time, + end_time=end_time, + is_async=True, + streaming_chunks=self.streaming_chunks, ) if complete_streaming_response is not None: @@ -3029,6 +3124,13 @@ class Logging(LiteLLMLoggingBaseClass): if not hasattr(self, "model_call_details"): self.model_call_details = {} + if ( + self.model_call_details.get("log_event_type") == "failed_api_call" + and self.model_call_details.get("exception") is exception + and self.model_call_details.get("standard_logging_object") is not None + ): + return start_time, self.model_call_details["end_time"] + self.model_call_details["log_event_type"] = "failed_api_call" self.model_call_details["exception"] = exception self.model_call_details["traceback_exception"] = ( @@ -3558,7 +3660,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time: datetime.datetime, is_async: bool, streaming_chunks: list[object], - ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | None: + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse | None: if self.stream is not True: return None if isinstance(result, ModelResponse) or isinstance(result, TextCompletionResponse): @@ -3583,9 +3685,40 @@ class Logging(LiteLLMLoggingBaseClass): ), ) return result.response + elif isinstance(result, InteractionsAPIStreamingResponse): + return self._assemble_completed_interaction_response(result) else: return None + @staticmethod + def _assemble_completed_interaction_response( + result: InteractionsAPIStreamingResponse, + ) -> InteractionsAPIResponse | None: + """ + The Interactions API streaming iterator hands the terminal event to the + success handlers: the new schema (Api-Revision: 2026-05-20) emits + ``interaction.completed`` carrying the full interaction object, the + legacy schema (2026-05-07) emits a chunk with ``status="completed"`` + and usage on the chunk itself. Build the equivalent non-streaming + response so cost calculation and spend tracking see one shape. + """ + if result.event_type == "interaction.completed" and result.interaction is not None: + return InteractionsAPIResponse(**result.interaction) + if result.status == "completed": + return InteractionsAPIResponse( + **result.model_dump( + exclude={ # mutable-ok: pydantic types exclude as set[str], which a frozenset does not satisfy + "event_type", + "delta", + "index", + "step", + "interaction_id", + "interaction", + } + ) + ) + return None + def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelResponse: """ Handles logging for Anthropic messages responses. @@ -5092,6 +5225,8 @@ class StandardLoggingPayloadSetup: elif isinstance(usage, dict): if ResponseAPILoggingUtils._is_response_api_usage(usage): return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(usage) + if InteractionsUsageObjectTransformation.is_interactions_usage_object(usage): + return InteractionsUsageObjectTransformation.transform_interactions_usage_object(usage) return Usage(**usage) raise ValueError(f"usage is required, got={usage} of type {type(usage)}") @@ -5118,6 +5253,8 @@ class StandardLoggingPayloadSetup: if isinstance(_raw, dict): if ResponseAPILoggingUtils._is_response_api_usage(_raw): return ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(_raw).model_dump() + if InteractionsUsageObjectTransformation.is_interactions_usage_object(_raw): + return InteractionsUsageObjectTransformation.transform_interactions_usage_object(_raw).model_dump() return _raw if isinstance(_raw, Usage): return _raw.model_dump() @@ -5325,9 +5462,10 @@ class StandardLoggingPayloadSetup: error_class: Final[str] = str(original_exception.__class__.__name__) if original_exception else "" _llm_provider_in_exception: Final = getattr(original_exception, "llm_provider", "") - # Get traceback information (first 100 lines) traceback_info = traceback_str or "" - if original_exception: + if original_exception and ( + litellm.log_client_error_tracebacks or not is_expected_client_error(original_exception) + ): tb: Final[TracebackType | None] = getattr(original_exception, "__traceback__", None) if tb: tb_lines: Final = traceback.format_tb(tb) @@ -5800,11 +5938,15 @@ def get_standard_logging_object_payload( response_model_name = final_response_obj.get("model") # For Azure Model Router, preserve the actual model in the top-level standard - # logging payload only when the user has opted in. + # logging payload. + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + requested_model: Final = kwargs.get("model") - if ( - isinstance(requested_model, str) - and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower()) + stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) + if stamped_selected_model is not None: + model_name = stamped_selected_model + elif ( + AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params) and isinstance(response_model_name, str) and response_model_name ): diff --git a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py index 210ac72cd8a..df436ef7611 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py +++ b/litellm/litellm_core_utils/llm_cost_calc/usage_object_transformation.py @@ -1,6 +1,9 @@ +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import Any from litellm.types.utils import ( + CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, TranscriptionUsageDurationObject, TranscriptionUsageTokensObject, @@ -34,3 +37,127 @@ class TranscriptionUsageObjectTransformation: ), ) return None + + +_INTERACTIONS_MODALITY_FIELDS: Mapping[str, str] = MappingProxyType( + { + "text": "text_tokens", + "audio": "audio_tokens", + "image": "image_tokens", + "video": "video_tokens", + "document": "text_tokens", + } +) + + +def _modality_field(entry: Mapping[str, Any]) -> str | None: + return _INTERACTIONS_MODALITY_FIELDS.get(str(entry.get("modality", "")).lower()) + + +def _token_count(value: object) -> int: + return value if isinstance(value, int) else 0 + + +def _modality_token_sums(entries: Sequence[Mapping[str, Any]]) -> Mapping[str, int]: + fields = frozenset(field for entry in entries if (field := _modality_field(entry)) is not None) + return MappingProxyType( + { + field: sum(_token_count(entry.get("tokens")) for entry in entries if _modality_field(entry) == field) + for field in fields + } + ) + + +def _google_search_query_count(usage_object: Mapping[str, Any]) -> int: + return sum( + _token_count(entry.get("count")) + for entry in tuple(usage_object.get("grounding_tool_count") or ()) + if isinstance(entry, Mapping) and entry.get("type") == "google_search" # pyright: ignore[reportUnnecessaryIsInstance] # provider JSON, not the empty tuple inferred from `or ()` + ) + + +def _subtract_cached_from_input( + input_sums: Mapping[str, int], + cached_sums: Mapping[str, int], + total_cached_tokens: int, +) -> Mapping[str, int]: + if cached_sums: + return MappingProxyType( + {field: max(0, tokens - cached_sums.get(field, 0)) for field, tokens in input_sums.items()} + ) + if total_cached_tokens and "text_tokens" in input_sums: + return MappingProxyType( + { + **input_sums, + "text_tokens": max(0, input_sums["text_tokens"] - total_cached_tokens), + } + ) + return input_sums + + +class InteractionsUsageObjectTransformation: + """ + Maps the Google Interactions API usage block (total_input_tokens, + output_tokens_by_modality, ...) into LiteLLM's chat-format ``Usage`` so the + generic cost calculator and spend tracking can bill it. + """ + + @staticmethod + def is_interactions_usage_object(usage_object: object) -> bool: + if not isinstance(usage_object, dict): + return False + if "prompt_tokens" in usage_object or "input_tokens" in usage_object: + return False + return "total_input_tokens" in usage_object or "total_output_tokens" in usage_object + + @staticmethod + def transform_interactions_usage_object(usage_object: Mapping[str, Any]) -> Usage: + input_entries = tuple(usage_object.get("input_tokens_by_modality") or ()) + tuple( + usage_object.get("tool_use_tokens_by_modality") or () + ) + cached_sums = _modality_token_sums(tuple(usage_object.get("cached_tokens_by_modality") or ())) + output_sums = _modality_token_sums(tuple(usage_object.get("output_tokens_by_modality") or ())) + + total_cached_tokens = _token_count(usage_object.get("total_cached_tokens")) + input_sums = _subtract_cached_from_input( + input_sums=_modality_token_sums(input_entries), + cached_sums=cached_sums, + total_cached_tokens=total_cached_tokens, + ) + + reasoning_tokens = _token_count(usage_object.get("total_reasoning_tokens")) or _token_count( + usage_object.get("total_thought_tokens") + ) + prompt_tokens = _token_count(usage_object.get("total_input_tokens")) + _token_count( + usage_object.get("total_tool_use_tokens") + ) + completion_tokens = _token_count(usage_object.get("total_output_tokens")) + reasoning_tokens + total_tokens = _token_count(usage_object.get("total_tokens")) or (prompt_tokens + completion_tokens) + + web_search_requests = _google_search_query_count(usage_object) + prompt_tokens_details = ( + PromptTokensDetailsWrapper( + cached_tokens=total_cached_tokens or None, + web_search_requests=web_search_requests or None, + **input_sums, + ) + if input_sums or total_cached_tokens or web_search_requests + else None + ) + completion_tokens_details = ( + CompletionTokensDetailsWrapper( + reasoning_tokens=reasoning_tokens or None, + **output_sums, + ) + if output_sums or reasoning_tokens + else None + ) + + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + prompt_tokens_details=prompt_tokens_details, + completion_tokens_details=completion_tokens_details, + cache_read_input_tokens=total_cached_tokens or None, + ) diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index b4e27b129fe..c833d57b6a9 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -1,8 +1,88 @@ +from collections.abc import Mapping from typing import Final import litellm +def _form_field_value(value: object) -> str: + if value is True: + return "true" + if value is False: + return "false" + return str(value) + + +def _flatten_form_field(key: str, value: object) -> tuple[tuple[str, str], ...]: + if isinstance(value, Mapping): + return tuple( + item for subkey, subvalue in value.items() for item in _flatten_form_field(f"{key}[{subkey}]", subvalue) + ) + if isinstance(value, (list, tuple)): + return tuple(item for entry in value for item in _flatten_form_field(f"{key}[]", entry)) + if value is None: + return () + serialized: Final = _form_field_value(value) + if not serialized: + return () + return ((key, serialized),) + + +def _is_form_scalar(value: object) -> bool: + return value is not None and not isinstance(value, (Mapping, list, tuple)) + + +def _flatten_form_data_field(key: str, value: object) -> tuple[tuple[str, str | tuple[str, ...]], ...]: + if isinstance(value, Mapping): + return tuple( + item + for subkey, subvalue in value.items() + for item in _flatten_form_data_field(f"{key}[{subkey}]", subvalue) + ) + if isinstance(value, (list, tuple)): + if all(_is_form_scalar(entry) for entry in value): + serialized_fields: Final = tuple(field for entry in value if (field := _form_field_value(entry))) + return ((key, serialized_fields),) if serialized_fields else () + return tuple(item for entry in value for item in _flatten_form_data_field(f"{key}[]", entry)) + if value is None: + return () + serialized: Final = _form_field_value(value) + if not serialized: + return () + return ((key, serialized),) + + +def flatten_form_field_values(*sources: Mapping[str, object] | None) -> tuple[tuple[str, str | tuple[str, ...]], ...]: + """ + Flatten JSON-shaped bodies into ``(name, value)`` form fields for a ``dict``-backed + multipart body, applying ``sources`` in order so a later source wins on a key collision + under ``dict.update``. Nested objects become ``key[subkey]`` fields the way the OpenAI SDK + serializes them, so provider params reach a multipart request without handing the httpx + encoder a nested value it rejects with ``Invalid type for value``. A scalar list becomes a + single field carrying a tuple value, which httpx emits as one repeated part per element, so + every element survives instead of collapsing to the last under ``dict.update``. + """ + return tuple( + pair + for source in sources + if source is not None + for top_key, top_value in source.items() + for pair in _flatten_form_data_field(top_key, top_value) + ) + + +def serialize_multipart_form_fields(data: Mapping[str, object]) -> tuple[tuple[str, tuple[None, str]], ...]: + """ + Encode a JSON-shaped body as OpenAI-SDK-style multipart file-tuples so a file-less + request is still sent as multipart/form-data, working around httpx downgrading a + file-less ``data=`` payload to application/x-www-form-urlencoded. + """ + return tuple( + (key, (None, serialized)) + for top_key, top_value in data.items() + for key, serialized in _flatten_form_field(top_key, top_value) + ) + + def _ensure_extra_body_is_safe(extra_body: dict | None) -> dict | None: """ Ensure that the extra_body sent in the request is safe, otherwise users will see this error diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 41d2af27eeb..77792671f6f 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -5,7 +5,7 @@ import asyncio import atexit import contextvars import logging -from collections.abc import Coroutine +from collections.abc import Coroutine, Iterator from typing import Final from typing_extensions import TypedDict @@ -61,6 +61,19 @@ class LoggingWorker: # Register cleanup handler to flush remaining events on exit atexit.register(self._flush_on_exit) + @staticmethod + def _drain_pending(queue: "asyncio.Queue[LoggingTask]") -> tuple[LoggingTask, ...]: + """Pop every task still queued, without awaiting them, so they can be moved to another queue.""" + + def _pop_until_empty() -> Iterator[LoggingTask]: + while True: + try: + yield queue.get_nowait() + except asyncio.QueueEmpty: + return + + return tuple(_pop_until_empty()) + def _ensure_queue(self) -> None: """Initialize the queue if it doesn't exist or if event loop has changed.""" try: @@ -69,14 +82,27 @@ class LoggingWorker: # No running loop, can't initialize return - # Check if we need to reinitialize due to event loop change + # The queue, semaphore and worker task are all bound to the loop that created them. On a + # loop change we hand the still-pending tasks to a fresh queue instead of dropping them, + # so queued spend-logging coroutines are not silently discarded (and never left un-awaited). if self._queue is not None and self._bound_loop is not current_loop: - verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker") - # Clear old state - these are bound to the old loop - self._queue = None + carried_over: Final = self._drain_pending(self._queue) + new_queue: Final[asyncio.Queue[LoggingTask]] = asyncio.Queue(maxsize=self.max_queue_size) + for carried_task in carried_over: + new_queue.put_nowait(carried_task) + if carried_over: + verbose_logger.warning( + "LoggingWorker: event loop changed; carried %d pending logging task(s) onto the new loop", + len(carried_over), + ) + else: + verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker") self._sem = None self._worker_task = None self._running_tasks.clear() + self._queue = new_queue + self._bound_loop = current_loop + return if self._queue is None: self._queue = asyncio.Queue(maxsize=self.max_queue_size) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 2db5776047b..72ea85dfa33 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -28,8 +28,12 @@ from litellm.types.llms.openai import ( ChatCompletionAssistantMessage, ChatCompletionFileObject, ChatCompletionImageObject, + ChatCompletionReasoningItem, + ChatCompletionReasoningSummaryTextBlock, + ChatCompletionRedactedThinkingBlock, ChatCompletionResponseMessage, ChatCompletionTextObject, + ChatCompletionThinkingBlock, ChatCompletionToolParam, ChatCompletionUserMessage, ) @@ -466,6 +470,8 @@ def update_messages_with_model_file_ids( from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, convert_b64_uid_to_unified_uid, + get_original_file_id, + is_model_embedded_id, ) for message in messages: @@ -504,6 +510,11 @@ def update_messages_with_model_file_ids( unified_file_id = convert_b64_uid_to_unified_uid(file_id) if "llm_output_file_id," in unified_file_id: provider_file_id = unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + if not provider_file_id and is_model_embedded_id(file_id): + # `litellm:;model,` encoding from the + # x-litellm-model upload path. Strip the wrapper + # so the provider sees its own ID. + provider_file_id = get_original_file_id(file_id) file_object_file_field["file_id"] = provider_file_id or file_id if format: file_object_file_field["format"] = format @@ -531,6 +542,8 @@ def update_responses_input_with_model_file_ids( from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, convert_b64_uid_to_unified_uid, + get_original_file_id, + is_model_embedded_id, ) if isinstance(input, str): @@ -574,6 +587,13 @@ def update_responses_input_with_model_file_ids( updated_content_item = content_item.copy() updated_content_item["file_id"] = provider_file_id updated_content.append(updated_content_item) + elif is_model_embedded_id(file_id): + # `litellm:;model,` encoding from the + # x-litellm-model upload path. Strip the wrapper + # so the provider sees its own ID. + updated_content_item = content_item.copy() + updated_content_item["file_id"] = get_original_file_id(file_id) + updated_content.append(updated_content_item) else: # Not a managed file, keep as-is updated_content.append(content_item) @@ -1549,6 +1569,44 @@ def _extract_reasoning_content(message: dict) -> tuple[str | None, str | None]: return None, message_content +def _readable_thinking_text( + block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, +) -> str: + """The text a chat model can read back, empty for redacted blocks and malformed ones.""" + if block.get("type") != "thinking": + return "" + thinking: Final = cast(ChatCompletionThinkingBlock, block).get("thinking") # cast-ok: narrowed by the type tag + return str(thinking or "") + + +def reasoning_content_from_thinking_blocks( + thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], +) -> str: + """Flatten Anthropic thinking blocks into the `reasoning_content` string chat models expect. + + Redacted blocks carry no readable text, so they contribute nothing. + """ + return "\n".join(text for block in thinking_blocks if (text := _readable_thinking_text(block))) + + +def responses_reasoning_item_from_thinking_blocks( + thinking_blocks: Iterable[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], +) -> ChatCompletionReasoningItem | None: + """Build a Responses API `reasoning` input item from Anthropic thinking blocks. + + The item carries no `id`: the Responses API rejects an empty one and 404s on any id it + did not mint itself, while an item without an id is always accepted. + """ + summary: Final[list[ChatCompletionReasoningSummaryTextBlock]] = [ # mutable-ok: API message payload + ChatCompletionReasoningSummaryTextBlock(type="summary_text", text=text) + for block in thinking_blocks + if (text := _readable_thinking_text(block)) + ] + if not summary: + return None + return ChatCompletionReasoningItem(type="reasoning", summary=summary) + + def _parse_content_for_reasoning( message_text: str | None, ) -> tuple[str | None, str | None]: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b676077ab0e..86cfbf70255 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -16,6 +16,7 @@ import litellm.types import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import REDACTED_BY_LITELLM from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client from litellm.types.files import get_file_extension_from_mime_type @@ -642,49 +643,6 @@ def claude_2_1_pt( return prompt -### TOGETHER AI - - -def get_model_info(token, model): - try: - headers: Final = {"Authorization": f"Bearer {token}"} - client: Final = HTTPHandler(concurrent_limit=1) - response: Final = client.get("https://api.together.xyz/models/info", headers=headers) - if response.status_code == 200: - model_info: Final = response.json() - for m in model_info: - if m["name"].lower().strip() == model.strip(): - return m["config"].get("prompt_format", None), m["config"].get("chat_template", None) - return None, None - else: - return None, None - except Exception: # safely fail a prompt template request - return None, None - - -## OLD TOGETHER AI FLOW -# def format_prompt_togetherai(messages, prompt_format, chat_template): -# if prompt_format is None: -# return default_pt(messages) - -# human_prompt, assistant_prompt = prompt_format.split("{prompt}") - -# if chat_template is not None: -# prompt = hf_chat_template( -# model=None, messages=messages, chat_template=chat_template -# ) -# elif prompt_format is not None: -# prompt = custom_prompt( -# role_dict={}, -# messages=messages, -# initial_prompt_value=human_prompt, -# final_prompt_value=assistant_prompt, -# ) -# else: -# prompt = default_pt(messages) -# return prompt - - ### IBM Granite @@ -5383,12 +5341,13 @@ def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> return raw if not isinstance(raw, str): return {} + normalized_raw: Final = "{}" if raw == REDACTED_BY_LITELLM else raw from litellm.litellm_core_utils.prompt_templates.common_utils import ( parse_tool_call_arguments, ) try: - parsed: Final = parse_tool_call_arguments(raw, tool_name=tool_name, context=context) + parsed: Final = parse_tool_call_arguments(normalized_raw, tool_name=tool_name, context=context) except ValueError as e: verbose_logger.warning("Failed to parse tool call arguments: %s", e) return {} diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 0d590e1ceba..be187b091b8 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -13,6 +13,7 @@ import inspect from typing import TYPE_CHECKING, Any, Final import litellm +from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -84,29 +85,29 @@ def _redact_tool_calls(tool_calls) -> None: for tool_call in tool_calls: function = getattr(tool_call, "function", None) if function is not None and hasattr(function, "arguments"): - function.arguments = "redacted-by-litellm" + function.arguments = REDACTED_BY_LITELLM def _redact_function_call(function_call) -> None: """Redact legacy assistant function_call arguments.""" if function_call is not None and hasattr(function_call, "arguments"): - function_call.arguments = "redacted-by-litellm" + function_call.arguments = REDACTED_BY_LITELLM def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" if isinstance(choice, litellm.Choices): - choice.message.content = "redacted-by-litellm" + choice.message.content = REDACTED_BY_LITELLM if hasattr(choice.message, "reasoning_content"): - choice.message.reasoning_content = "redacted-by-litellm" + choice.message.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.message, "thinking_blocks"): choice.message.thinking_blocks = None _redact_tool_calls(getattr(choice.message, "tool_calls", None)) _redact_function_call(getattr(choice.message, "function_call", None)) elif isinstance(choice, litellm.utils.StreamingChoices): - choice.delta.content = "redacted-by-litellm" + choice.delta.content = REDACTED_BY_LITELLM if hasattr(choice.delta, "reasoning_content"): - choice.delta.reasoning_content = "redacted-by-litellm" + choice.delta.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.delta, "thinking_blocks"): choice.delta.thinking_blocks = None _redact_tool_calls(getattr(choice.delta, "tool_calls", None)) @@ -117,22 +118,22 @@ def _redact_responses_api_output(output_items): """Helper to redact ResponsesAPIResponse output items.""" for output_item in output_items: if hasattr(output_item, "text"): - output_item.text = "redacted-by-litellm" + output_item.text = REDACTED_BY_LITELLM if hasattr(output_item, "content") and isinstance(output_item.content, list): for content_part in output_item.content: if hasattr(content_part, "text"): - content_part.text = "redacted-by-litellm" + content_part.text = REDACTED_BY_LITELLM # Redact reasoning items in output array if hasattr(output_item, "type") and output_item.type == "reasoning": if hasattr(output_item, "summary") and isinstance(output_item.summary, list): for summary_item in output_item.summary: if hasattr(summary_item, "text"): - summary_item.text = "redacted-by-litellm" + summary_item.text = REDACTED_BY_LITELLM if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"): - output_item.arguments = "redacted-by-litellm" + output_item.arguments = REDACTED_BY_LITELLM def _redact_responses_api_output_dict(output_items, redacted_str: str): @@ -164,7 +165,7 @@ def _redact_standard_logging_object(model_call_details: dict): if standard_logging_object is None: return - redacted_str: Final = "redacted-by-litellm" + redacted_str: Final = REDACTED_BY_LITELLM if standard_logging_object.get("messages") is not None: standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}] @@ -235,7 +236,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons copy via redact_streaming_responses_for_custom_logger instead. """ # Redact model_call_details - model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}] + model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}] model_call_details["prompt"] = "" model_call_details["input"] = "" _redact_standard_logging_object(model_call_details) @@ -256,7 +257,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons or hasattr(result, "__anext__") # async generator ): # async iterator # For async objects, return a simple redacted response without deepcopy - return {"text": "redacted-by-litellm"} + return {"text": REDACTED_BY_LITELLM} if not ( isinstance(result, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse)) @@ -273,11 +274,11 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: - _redact_model_response_dict_choices(_result["choices"], "redacted-by-litellm") + _redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM) redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): - _redact_responses_api_output_dict(_result["output"], "redacted-by-litellm") + _redact_responses_api_output_dict(_result["output"], REDACTED_BY_LITELLM) elif isinstance(_result, litellm.ResponsesAPIResponse): if hasattr(_result, "output"): _redact_responses_api_output(_result.output) @@ -288,7 +289,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons if hasattr(_result, "data") and _result.data is not None: _result.data = [] else: - return {"text": "redacted-by-litellm"} + return {"text": REDACTED_BY_LITELLM} return _result diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index ef278c8f723..26380ad0af8 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1,7 +1,8 @@ import json import re import time -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, NoReturn, cast import httpx @@ -121,6 +122,32 @@ else: # response side. _ANTHROPIC_TOOL_NAME_INVALID_CHARS: Final = re.compile(r"[^a-zA-Z0-9_-]") _ANTHROPIC_TOOL_NAME_MAX_LEN: Final = 128 + +_ENUM_TYPE_CHECKS: Final[Mapping[str, Callable[[Any], bool]]] = MappingProxyType( + { + "null": lambda v: v is None, + "boolean": lambda v: isinstance(v, bool), + "integer": lambda v: isinstance(v, int) and not isinstance(v, bool), + "number": lambda v: isinstance(v, (int, float)) and not isinstance(v, bool), + "string": lambda v: isinstance(v, str), + "array": lambda v: isinstance(v, list), + "object": lambda v: isinstance(v, dict), + } +) + + +def _enum_conflicts_with_declared_type(schema: Mapping[str, Any]) -> bool: + """Whether ``schema``'s ``enum`` cannot match its declared ``type``.""" + enum_values: Final = schema.get("enum") + declared_type: Final = schema.get("type") + if not isinstance(enum_values, list) or declared_type is None: + return False + if isinstance(declared_type, list): + return True + check: Final = _ENUM_TYPE_CHECKS.get(declared_type) + return check is not None and not all(check(value) for value in enum_values) + + # Single, internal-only key on ``litellm_params`` used to thread the per- # request reverse map (sanitized -> original) from request build to response # parsing. ``litellm_params`` is never serialized to a provider; ``optional_ @@ -565,9 +592,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): else: result["description"] = constraint_note + drops_conflicting_type: Final = _enum_conflicts_with_declared_type(schema) + for key, value in schema.items(): if key in unsupported_fields: continue + if key == "type" and drops_conflicting_type: + continue if key == "description" and "description" in result: # Already handled above continue @@ -1184,8 +1215,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if reasoning_effort is None or reasoning_effort == "none": return None if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider): + # without display, Anthropic defaults adaptive thinking to + # display="omitted" and returns a blank thinking block return AnthropicThinkingParam( type="adaptive", + display="summarized", ) elif reasoning_effort == "low": return AnthropicThinkingParam( @@ -2113,7 +2147,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: + def thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: details: Final = usage_object.get("output_tokens_details") if not isinstance(details, Mapping): return None @@ -2145,7 +2179,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): reported_thinking_tokens: Final = ( iteration_thinking_tokens if iteration_thinking_tokens is not None - else self._thinking_tokens_from_usage(usage_object) + else self.thinking_tokens_from_usage(usage_object) ) if reported_thinking_tokens is not None: capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens) @@ -2168,7 +2202,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: per_iteration: Final = tuple( - self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None + self.thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None for iteration in iterations ) reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 3297aa95715..53cc3464761 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -38,6 +38,21 @@ DROP_DISABLED_THINKING_WARNING: Final = ( "thinking blocks, and those thinking tokens are billed as output tokens." ) +# Anthropic error `type` (both the JSON error body and SSE `event: error` +# payloads use this field) mapped to the HTTP status code it corresponds to. +ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = MappingProxyType( + { + "invalid_request_error": 400, + "authentication_error": 401, + "permission_error": 403, + "not_found_error": 404, + "rate_limit_error": 429, + "api_error": 500, + "overloaded_error": 503, + "timeout_error": 504, + } +) + _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") @@ -440,6 +455,16 @@ class AnthropicModelInfo(BaseLLMModelInfo): """ return AnthropicModelInfo._supports_model_capability(model, "thinking_always_on", custom_llm_provider) + @staticmethod + def _supports_legacy_thinking(model: str, custom_llm_provider: str) -> bool: + """Whether ``model`` is an adaptive-thinking model that still accepts legacy + ``thinking.type=enabled`` with ``budget_tokens`` (the Claude 4.6 family). + The model cost map is authoritative: an explicit ``supports_legacy_thinking`` + entry resolved under ``custom_llm_provider``, or a ``fallback_generalizations`` + rule for unmapped 4.6 ids. Absent flag means the model rejects the legacy shape. + """ + return AnthropicModelInfo._supports_model_capability(model, "supports_legacy_thinking", custom_llm_provider) + @staticmethod def maybe_drop_disabled_thinking( model: str, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 34c2d837127..7c89da81fe6 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -64,6 +64,7 @@ from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingCho from litellm.litellm_core_utils.prompt_templates.common_utils import ( parse_tool_call_arguments, + reasoning_content_from_thinking_blocks, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -592,6 +593,9 @@ class LiteLLMAnthropicMessagesAdapter: assistant_message["tool_calls"] = tool_calls if len(thinking_blocks) > 0: assistant_message["thinking_blocks"] = thinking_blocks + reasoning_content = reasoning_content_from_thinking_blocks(thinking_blocks) + if reasoning_content: + assistant_message["reasoning_content"] = reasoning_content new_messages.append(assistant_message) return new_messages diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 922769dbbfd..0a12bc3135f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime from typing import Any, Final, Protocol, runtime_checkable @@ -11,9 +11,11 @@ from typing_extensions import TypedDict from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.utils import GenericStreamingChunk, ModelResponseStream @@ -33,26 +35,239 @@ def _is_message_stop_chunk(chunk: object) -> bool: return False -def _is_provider_error_chunk(chunk: object) -> bool: +def is_anthropic_ping_chunk(chunk: object) -> bool: + """ + Whether a chunk is a pure ``ping`` keepalive frame. It carries no content + and can recur indefinitely on a slow-starting or idle connection, so a + mid-stream fallback wrapper drops it outright while still deciding + whether to commit to the primary stream, rather than buffering it. + + A physical transport chunk that coalesces a ping with any other SSE + event (``message_start``, ``content_block_delta``, ``event: error``, ...) + is NOT a pure ping - dropping it whole would discard those events - so + only a chunk whose every ``event:`` line is ``event: ping`` qualifies. + """ if isinstance(chunk, dict): - return chunk.get("type") == "error" + return chunk.get("type") == "ping" if isinstance(chunk, (bytes, bytearray)): - return any(line == b"event: error" for line in chunk.splitlines()) + event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:")) + return bool(event_lines) and all(line == b"event: ping" for line in event_lines) return False +def is_anthropic_content_delta_chunk(chunk: object) -> bool: + """ + Whether a chunk carries actual assistant-generated output (a + ``content_block_delta`` frame), as opposed to a lifecycle/bookkeeping + frame (``message_start``, ``content_block_start``/``stop``, + ``message_delta``, ``message_stop``, ``ping``) that carries nothing + worth preserving before an invisible mid-stream fallback retry. + """ + if isinstance(chunk, dict): + return chunk.get("type") == "content_block_delta" + if isinstance(chunk, (bytes, bytearray)): + return any(line == b"event: content_block_delta" for line in chunk.splitlines()) + return False + + +def _decoded_sse_data_line(line: bytes) -> object | None: + if not line.startswith(b"data:"): + return None + try: + return json.loads(line[len(b"data:") :].strip()) + except (ValueError, TypeError): + return None + + +def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: + if isinstance(chunk, dict): + return chunk if chunk.get("type") == "error" else None + if isinstance(chunk, (bytes, bytearray)): + decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines()) + return next( + ( + candidate + for candidate in decoded_lines + if isinstance(candidate, dict) and candidate.get("type") == "error" + ), + None, + ) + return None + + +def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None: + """Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None.""" + payload: Final = _anthropic_error_event_payload(chunk) + error_body: Final = payload.get("error") if payload is not None else None + return error_body if isinstance(error_body, dict) else None + + +def _is_provider_error_chunk(chunk: object) -> bool: + return _anthropic_error_body(chunk) is not None + + +def parse_anthropic_error_event(chunk: object) -> tuple[str, str, int] | None: + """ + Extract ``(error_type, message, http_status_code)`` from an Anthropic SSE + ``event: error`` chunk (raw bytes or an already-decoded dict), or None if + ``chunk`` is not an error event. + + The status code is looked up via ANTHROPIC_ERROR_STATUS_CODE_MAP, + defaulting to 500 for an error ``type`` Anthropic hasn't documented yet. + """ + error_body: Final = _anthropic_error_body(chunk) + if error_body is None: + return None + error_type: Final = error_body.get("type") + if not isinstance(error_type, str): + return None + message: Final = error_body.get("message") + return ( + error_type, + message if isinstance(message, str) else error_type, + ANTHROPIC_ERROR_STATUS_CODE_MAP.get(error_type, 500), + ) + + def _is_terminal_stream_chunk(chunk: object) -> bool: return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk) +def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes: + return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode() + + def _incomplete_stream_error_sse_event() -> bytes: - payload: Final = json.dumps( - { - "type": "error", - "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}, - } + return _sse_event( # mutable-ok: one-shot JSON payload, never mutated after construction + "error", + {"type": "error", "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}}, + ) + + +def _anthropic_content_block_start_and_deltas( + block: Mapping[str, object], +) -> tuple[Mapping[str, object], tuple[Mapping[str, object], ...]]: + """ + ``(content_block_start.content_block, content_block_delta.delta events)`` + for one Anthropic response content block. A thinking block emits both a + thinking_delta and a trailing signature_delta - a real Anthropic stream + does the same, and dropping the signature makes any replay of that + assistant message (a follow-up turn, a tool-use continuation) fail + Anthropic's thinking-signature verification. redacted_thinking has no + delta at all - it is sent complete in content_block_start. + """ + match block.get("type"): + case "tool_use": + return ( + { # mutable-ok: one-shot payload + "id": block.get("id"), + "name": block.get("name"), + "input": {}, # mutable-ok: one-shot payload + "type": "tool_use", + }, + ( + { # mutable-ok: one-shot payload + "partial_json": json.dumps(block.get("input") or {}), # mutable-ok: one-shot payload + "type": "input_json_delta", + }, + ), + ) + case "thinking": + signature: Final = block.get("signature") + signature_deltas: Final = ( + ({"signature": signature, "type": "signature_delta"},) # mutable-ok: one-shot payload + if isinstance(signature, str) and signature + else () + ) + return ( + {"thinking": "", "signature": "", "type": "thinking"}, # mutable-ok: one-shot payload + ( + {"thinking": block.get("thinking") or "", "type": "thinking_delta"}, # mutable-ok: one-shot payload + *signature_deltas, + ), + ) + case "redacted_thinking": + return ({"type": "redacted_thinking", "data": block.get("data")}, ()) # mutable-ok: one-shot JSON payload + case _: + return ( + {"type": "text", "text": ""}, # mutable-ok: one-shot JSON payload + ({"type": "text_delta", "text": block.get("text") or ""},), # mutable-ok: one-shot JSON payload + ) + + +def anthropic_messages_response_as_sse_events(response: AnthropicMessagesResponse) -> tuple[bytes, ...]: + """ + Render a complete (non-streaming) AnthropicMessagesResponse as the SSE + event sequence a real streaming request would have produced. + + A mid-stream fallback can resolve to a non-streaming response even + though the client asked to stream (e.g. an agentic tool-use loop that + intercepts and returns a complete message) - yielding that dict directly + into a `/v1/messages` SSE byte stream would produce a malformed + response, so it's synthesized into the message_start/content_block_*/ + message_delta/message_stop lifecycle a real stream would have sent. + """ + content_blocks: Final = response.get("content") or () + content_events: Final = ( + event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block) + ) + # A real message_start always carries a null stop_reason/stop_sequence and + # a zero output_tokens - those are only known once generation finishes, so + # copying the completed response's final values here would let a client + # treat the message as already finished, or double-count output tokens. + message_start_usage: Final = { # mutable-ok: one-shot JSON payload + **(response.get("usage") or {}), + "output_tokens": 0, + } + message_start_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction + "type": "message_start", + "message": { # mutable-ok: one-shot JSON payload + **response, + "content": [], # mutable-ok: one-shot JSON payload + "stop_reason": None, + "stop_sequence": None, + "usage": message_start_usage, + }, + } + message_delta_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction + "type": "message_delta", + "delta": { # mutable-ok: one-shot JSON payload + "stop_reason": response.get("stop_reason"), + "stop_sequence": response.get("stop_sequence"), + }, + "usage": response.get("usage") or {}, # mutable-ok: one-shot JSON payload + } + return ( + _sse_event("message_start", message_start_payload), + *content_events, + _sse_event("message_delta", message_delta_payload), + _sse_event("message_stop", {"type": "message_stop"}), # mutable-ok: one-shot JSON payload + ) + + +def _anthropic_content_block_events(index: int, block: Mapping[str, object]) -> tuple[bytes, ...]: + start_block, deltas = _anthropic_content_block_start_and_deltas(block) + start_payload: Final = { # mutable-ok: one-shot payload + "type": "content_block_start", + "index": index, + "content_block": start_block, + } + stop_payload: Final = { # mutable-ok: one-shot payload + "type": "content_block_stop", + "index": index, + } + delta_events: Final = tuple( + _sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": index, "delta": delta}, # mutable-ok: one-shot payload + ) + for delta in deltas + ) + return ( + _sse_event("content_block_start", start_payload), + *delta_events, + _sse_event("content_block_stop", stop_payload), ) - return f"event: error\ndata: {payload}\n\n".encode() class AnthropicMessagesStreamHiddenParams(TypedDict): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index adabfa2d62d..032bf0130ce 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -379,13 +379,19 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): def _translate_legacy_thinking_for_adaptive_model( model: str, optional_params: dict, custom_llm_provider: str ) -> None: - """Translate legacy ``thinking.type=enabled`` to adaptive for 4.6/4.7. - Caller-provided ``output_config.effort`` is never overridden. + """Translate legacy ``thinking.type=enabled`` to adaptive for the + adaptive-thinking models that reject it (4.7+ and the 5 families). + Models flagged ``supports_legacy_thinking`` (the 4.6 family) accept the + legacy shape natively, so it is forwarded verbatim and the caller's + ``budget_tokens`` cap keeps applying. Caller-provided + ``output_config.effort`` is never overridden. """ from litellm.llms.anthropic.chat.transformation import AnthropicConfig if not AnthropicModelInfo._is_adaptive_thinking_model(model, custom_llm_provider): return + if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider): + return thinking: Final = optional_params.get("thinking") if not isinstance(thinking, dict) or thinking.get("type") != "enabled": return diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index e2ad9c9c6d3..292d2622c7f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -152,7 +152,10 @@ class AnthropicResponsesStreamWrapper: if block_idx < 0: if not delta: return - block_idx = self._open_block(item_id, {"type": "thinking", "thinking": ""}) + block_idx = self._open_block( + item_id, + {"type": "thinking", "thinking": "", "signature": ""}, # mutable-ok: API message payload + ) self._chunk_queue.append( { "type": "content_block_delta", diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 25d729d8606..6d47d0de19f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -6,12 +6,14 @@ path used for OpenAI and Azure models. """ import json -from collections.abc import Iterable +from collections.abc import Iterable, Mapping +from itertools import groupby from typing import Any, Final, cast from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + responses_reasoning_item_from_thinking_blocks, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.reasoning_effort_utils import ( @@ -36,7 +38,11 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, AnthropicUsage, ) -from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.llms.openai import ( + ChatCompletionThinkingBlock, + ResponseAPIUsage, + ResponsesAPIResponse, +) class LiteLLMAnthropicToResponsesAPIAdapter: @@ -100,6 +106,58 @@ class LiteLLMAnthropicToResponsesAPIAdapter: if isinstance(block, dict) and block.get("type") == "text" and (text := block.get("text")) # pyright: ignore[reportUnnecessaryIsInstance] # untrusted client payload ] + @staticmethod + def _summary_part_text(part: object) -> str: + if isinstance(part, Mapping): + mapping: Final = cast(Mapping[str, Any], part) # cast-ok: summary parts are untyped provider json + return str(mapping.get("text") or "") + return str(getattr(part, "text", None) or "") + + @classmethod + def _thinking_blocks_from_reasoning_item( + cls, + summary: Iterable[object], + ) -> tuple[dict[str, Any], ...]: # mutable-ok: API message payload + """Anthropic thinking blocks for one Responses reasoning item. + + The signature stays empty: only Anthropic can sign a thinking block, and a stand-in + value would be replayed as a real one and rejected by every backend that verifies it. + """ + return tuple( + AnthropicResponseContentBlockThinking( + type="thinking", + thinking=text, + signature=None, + ).model_dump() + for part in summary + if (text := cls._summary_part_text(part)) + ) + + @staticmethod + def _assistant_block_group_key(indexed_block: tuple[int, Mapping[str, Any]]) -> str: + """Group a run of consecutive thinking blocks together; keep every other block alone.""" + index, block = indexed_block + return "thinking" if block.get("type") == "thinking" else f"block:{index}" + + @classmethod + def _assistant_group_to_input_item( + cls, group: tuple[Mapping[str, Any], ...] + ) -> dict[str, Any] | None: # mutable-ok: API message payload + first: Final = group[0] + btype: Final = first.get("type") + if btype == "thinking": + blocks: Final = cast(tuple[ChatCompletionThinkingBlock, ...], group) # cast-ok: untrusted client payload + reasoning_item: Final = responses_reasoning_item_from_thinking_blocks(blocks) + return None if reasoning_item is None else dict(reasoning_item) # mutable-ok: API message payload + if btype == "tool_use": + return { # mutable-ok: API message payload + "type": "function_call", + "call_id": first.get("id", ""), + "name": first.get("name", ""), + "arguments": json.dumps(first.get("input", {})), # mutable-ok: API message payload + } + return None + def translate_messages_to_responses_input( self, messages: list[AllAnthropicPassThroughMessageValues], @@ -113,6 +171,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: user image -> message(role=user, input_image) user tool_result -> function_call_output assistant text -> message(role=assistant, output_text) + assistant thinking -> reasoning assistant tool_use -> function_call """ input_items: Final[list[dict[str, Any]]] = [] @@ -233,27 +292,17 @@ class LiteLLMAnthropicToResponsesAPIAdapter: } ) elif isinstance(content, list): - asst_parts: list[dict[str, Any]] = [] - for block in content: - if not isinstance(block, dict): - continue - btype = block.get("type") - if btype == "text": - asst_parts.append({"type": "output_text", "text": block.get("text", "")}) - elif btype == "tool_use": - # tool_use becomes a top-level function_call item - input_items.append( - { - "type": "function_call", - "call_id": block.get("id", ""), - "name": block.get("name", ""), - "arguments": json.dumps(block.get("input", {})), - } - ) - elif btype == "thinking": - thinking_text = block.get("thinking", "") - if thinking_text: - asst_parts.append({"type": "output_text", "text": thinking_text}) + blocks = tuple(block for block in content if isinstance(block, dict)) + input_items.extend( + item + for _, group in groupby(enumerate(blocks), key=self._assistant_block_group_key) + if (item := self._assistant_group_to_input_item(tuple(block for _, block in group))) is not None + ) + asst_parts: list[dict[str, Any]] = [ # mutable-ok: API message payload + {"type": "output_text", "text": block.get("text", "")} # mutable-ok: API message payload + for block in blocks + if block.get("type") == "text" + ] if asst_parts: input_items.append( { @@ -514,16 +563,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: for item in response.output: if isinstance(item, ResponseReasoningItem): - for summary in item.summary: - text = getattr(summary, "text", "") - if text: - content.append( - AnthropicResponseContentBlockThinking( - type="thinking", - thinking=text, - signature=None, - ).model_dump() - ) + content.extend(self._thinking_blocks_from_reasoning_item(item.summary)) elif isinstance(item, ResponseOutputMessage): for part in item.content: @@ -555,6 +595,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content.append( AnthropicResponseContentBlockText(type="text", text=part.get("text", "")).model_dump() ) + elif item_type == "reasoning": + content.extend( + self._thinking_blocks_from_reasoning_item( + cast(Iterable[object], item.get("summary") or ()), # cast-ok: untyped provider json + ) + ) elif item_type == "function_call": try: input_data = json.loads(item.get("arguments", "{}")) diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index 0c62418708f..5fdf2ceff7f 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -22,19 +22,7 @@ from litellm.types.llms.openai import ( from litellm.types.utils import CallTypes, LlmProviders, ModelResponse from ..chat.transformation import AnthropicConfig -from ..common_utils import AnthropicModelInfo - -# Map Anthropic error types to HTTP status codes -ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = { - "invalid_request_error": 400, - "authentication_error": 401, - "permission_error": 403, - "not_found_error": 404, - "rate_limit_error": 429, - "api_error": 500, - "overloaded_error": 503, - "timeout_error": 504, -} +from ..common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP, AnthropicModelInfo class AnthropicFilesHandler: diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index b77ba2f9460..2c34851d275 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -3,6 +3,7 @@ import hashlib import json import os from collections.abc import Callable, Mapping +from functools import lru_cache from typing import Any, Final, Literal, NamedTuple, cast import httpx @@ -75,6 +76,24 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict: return {**llm_response_headers, **openai_headers} +@lru_cache(maxsize=128) +def _cached_entra_id_token_provider( + tenant_id: str, + client_id: str, + client_secret: str, + scope: str, +) -> Callable[[], str]: + """Build (once per credential set) a bearer token provider backed by a `ClientSecretCredential`. + + The credential caches the access token internally and only talks to Entra ID when it is close + to expiry, so reusing the provider keeps one AAD round trip per token lifetime instead of one + per request. + """ + from azure.identity import ClientSecretCredential, get_bearer_token_provider + + return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) + + def get_azure_ad_token_from_entra_id( tenant_id: str, client_id: str, @@ -93,8 +112,6 @@ def get_azure_ad_token_from_entra_id( Returns: callable that returns a bearer token. """ - from azure.identity import ClientSecretCredential, get_bearer_token_provider - verbose_logger.debug("Getting Azure AD Token from Entra ID") if tenant_id.startswith("os.environ/"): @@ -120,9 +137,13 @@ def get_azure_ad_token_from_entra_id( ) if _tenant_id is None or _client_id is None or _client_secret is None: raise ValueError("tenant_id, client_id, and client_secret must be provided") - credential: Final = ClientSecretCredential(_tenant_id, _client_id, _client_secret) - token_provider: Final = get_bearer_token_provider(credential, scope) + token_provider: Final = _cached_entra_id_token_provider( + tenant_id=_tenant_id, + client_id=_client_id, + client_secret=_client_secret, + scope=scope, + ) verbose_logger.debug("token_provider %s", token_provider) diff --git a/litellm/llms/azure/search/__init__.py b/litellm/llms/azure/search/__init__.py new file mode 100644 index 00000000000..2414ba2b1e8 --- /dev/null +++ b/litellm/llms/azure/search/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.azure.search.transformation import BingGroundingSearchConfig + +__all__ = ("BingGroundingSearchConfig",) diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py new file mode 100644 index 00000000000..0754c9b1fda --- /dev/null +++ b/litellm/llms/azure/search/transformation.py @@ -0,0 +1,442 @@ +""" +Calls the Microsoft Foundry Responses API with the `bing_grounding` or `web_search` +tool to search the web (Grounding with Bing Search). + +Microsoft docs: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/how-to/tools/bing-grounding + +Setup: + 1. Set BING_GROUNDING_PROJECT_ENDPOINT to the Foundry project endpoint, e.g. + https://.services.ai.azure.com/api/projects/ + 2. Set BING_GROUNDING_MODEL to a model deployment in that project (e.g. gpt-4.1); + it runs the grounded search and its tokens are billed on that deployment + 3. Optional: set BING_GROUNDING_CONNECTION_ID to a Grounding with Bing Search + project connection id to use the `bing_grounding` tool; without it the + project's built-in `web_search` tool is used + 4. Auth: pass api_key (an Azure API key, sent in the api-key header), or set + BING_GROUNDING_TOKEN to an Entra bearer token for scope + https://ai.azure.com/.default, or configure azure-identity (AZURE_CLIENT_ID / + AZURE_CLIENT_SECRET / AZURE_TENANT_ID, managed identity, or any + DefaultAzureCredential source) and the token is minted automatically + +Usage: + response = litellm.search( + query="latest AI developments", + search_provider="bing_grounding", + max_results=5, + ) +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal + +import httpx +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + SearchResponse, + SearchResult, +) +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_DOCS_URL: Final = "https://learn.microsoft.com/en-us/azure/ai-foundry/agents/how-to/tools/bing-grounding" + +PROJECT_ENDPOINT_ENV: Final = "BING_GROUNDING_PROJECT_ENDPOINT" +MODEL_ENV: Final = "BING_GROUNDING_MODEL" +CONNECTION_ID_ENV: Final = "BING_GROUNDING_CONNECTION_ID" +TOKEN_ENV: Final = "BING_GROUNDING_TOKEN" + +ENTRA_SCOPE: Final = "https://ai.azure.com/.default" + +_RESPONSES_PATH: Final = "/openai/v1/responses" +_SNIPPET_FALLBACK_LENGTH: Final = 300 +_UPSTREAM_ERROR_STATUS: Final = 502 +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" + + +class _Annotation(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + url: str | None = None + title: str | None = None + start_index: int | None = None + end_index: int | None = None + + +class _ContentPart(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + text: str = "" + annotations: tuple[_Annotation, ...] = () + + +class _OutputItem(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + content: tuple[_ContentPart, ...] = () + + +class _ErrorBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + message: str | None = None + + +class _IncompleteDetails(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + reason: str | None = None + + +class _ResponsesEnvelope(BaseModel): + """A Foundry Responses API body. `output` is required: a body without it is not a + Responses API response and must not be reported as a successful empty search. + + A 200 body can still carry `status` `failed` or `incomplete`; those are surfaced as + errors rather than reported as a successful empty search.""" + + model_config = ConfigDict(extra="ignore", frozen=True) + + output: tuple[_OutputItem, ...] + status: str | None = None + error: _ErrorBody | None = None + incomplete_details: _IncompleteDetails | None = None + + +class _ErrorEnvelope(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + error: _ErrorBody | None = None + + +def _unwrap_error_detail(error_message: str) -> str: + """ + Surface the human-readable message inside Foundry's error envelope. + + Tool failures nest a second JSON document as a string inside `error.message` + (observed live for `bing_grounding` connection errors), so the unwrap runs twice. + Falls back to the raw body for anything else. + """ + try: + envelope: Final = _ErrorEnvelope.model_validate_json(error_message) + except ValidationError: + return error_message + message: Final = envelope.error.message if envelope.error else None + if message is None: + return error_message + try: + nested: Final = _ErrorBody.model_validate_json(message) + except ValidationError: + return message + return nested.message or message + + +def _snippet(text: str, annotation: _Annotation) -> str: + """ + The text a citation supports, not the citation marker itself. + + A url_citation's start/end indices span the inline marker ("([host](url))"), + which follows the claim it backs, so the snippet is the marker's own line up + to where the marker starts. + """ + start: Final = annotation.start_index + marker_start: Final = start if start is not None and 0 <= start <= len(text) else len(text) + claim: Final = text[:marker_start].rsplit("\n", 1)[-1].strip() + if claim: + return claim[-_SNIPPET_FALLBACK_LENGTH:] + return text[:_SNIPPET_FALLBACK_LENGTH] + + +def _citation_results(envelope: _ResponsesEnvelope) -> tuple[SearchResult, ...]: + """One result per cited URL: first occurrence wins, order preserved as answered.""" + cited: Final = tuple( + SearchResult( + title=annotation.title or "", + url=annotation.url or "", + snippet=_snippet(part.text, annotation), + date=None, + last_updated=None, + ) + for item in envelope.output + if item.type == "message" + for part in item.content + if part.type == "output_text" + for annotation in part.annotations + if annotation.type == "url_citation" and annotation.url + ) + first_by_url: Final = MappingProxyType({result.url: result for result in reversed(cited)}) + return tuple(first_by_url[url] for url in dict.fromkeys(result.url for result in cited)) + + +def _valid_max_results(max_results: object) -> int | None: + """A positive-int `max_results`, else None. Rejects bools, an `int` subclass, and + non-positive values so neither the request-side `count` nor the response-side cap + forwards a value the other would silently ignore. + """ + if isinstance(max_results, bool) or not isinstance(max_results, int): + return None + return max_results if max_results > 0 else None + + +def _requested_max_results(response_kwargs: Mapping[str, object]) -> int | None: + """The unified `max_results` cap the caller asked for, if any. + + The built-in web_search tool has no server-side result-count knob, so the cap is + enforced here after the fact; connection mode also honors it as a hard ceiling on + top of the tool's `count` hint. + """ + optional_params: Final = response_kwargs.get("optional_params") + if not isinstance(optional_params, Mapping): + return None + return _valid_max_results(optional_params.get("max_results")) + + +def _capped(results: tuple[SearchResult, ...], max_results: int | None) -> tuple[SearchResult, ...]: + return results[:max_results] if max_results is not None else results + + +class _SearchConfiguration(BaseModel): + model_config = ConfigDict(frozen=True) + + project_connection_id: str + count: int | None = None + + +class _BingGroundingParams(BaseModel): + model_config = ConfigDict(frozen=True) + + search_configurations: tuple[_SearchConfiguration, ...] + + +class _BingGroundingTool(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["bing_grounding"] = "bing_grounding" + bing_grounding: _BingGroundingParams + + +class _UserLocation(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["approximate"] = "approximate" + country: str + + +class _WebSearchTool(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["web_search"] = "web_search" + user_location: _UserLocation | None = None + + +class _ResponsesRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + model: str + input: str + tools: tuple[_BingGroundingTool | _WebSearchTool, ...] + + +def _search_tool(optional_params: Mapping[str, object]) -> _BingGroundingTool | _WebSearchTool: + connection_id: Final = get_secret_str(CONNECTION_ID_ENV) + max_results: Final = optional_params.get("max_results") + country: Final = optional_params.get("country") + if connection_id: + configuration: Final = _SearchConfiguration( + project_connection_id=connection_id, + count=_valid_max_results(max_results), + ) + return _BingGroundingTool(bing_grounding=_BingGroundingParams(search_configurations=(configuration,))) + location: Final = _UserLocation(country=country.upper()) if isinstance(country, str) else None + return _WebSearchTool(user_location=location) + + +def _default_entra_token_minter() -> str: + from litellm.secret_managers.get_azure_ad_token_provider import get_azure_ad_token_provider + + return get_azure_ad_token_provider(azure_scope=ENTRA_SCOPE)() + + +class BingGroundingSearchConfig(BaseSearchConfig): + def __init__(self, entra_token_minter: Callable[[], str] | None = None) -> None: + super().__init__() + self._entra_token_minter = entra_token_minter + + @staticmethod + def ui_friendly_name() -> str: + return "Grounding with Bing Search" + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseSearchConfig.validate_environment signature + api_key: str | None = None, + api_base: str | None = None, + **kwargs: object, # kwargs-ok: BaseSearchConfig.validate_environment signature + ) -> dict[str, str]: # mutable-ok: the http handler passes this straight to httpx as headers + """ + Validate environment and return headers. + + Returns a new dict rather than mutating ``headers``: the http handler calls this + a second time after ``litellm/search/main.py`` already did, so it has to be idempotent. + """ + return { # mutable-ok: httpx requires a plain dict of headers + **headers, + **self._auth_header(api_key, api_base), + "Content-Type": "application/json", + } + + def _auth_header(self, api_key: str | None, api_base: str | None) -> Mapping[str, str]: + """ + A caller-supplied ``api_key`` is an Azure API key and rides the ``api-key`` header; + an Entra bearer token (``BING_GROUNDING_TOKEN`` or one minted via azure-identity) + rides ``Authorization: Bearer``. Foundry rejects the wrong scheme for each. + """ + if api_key: + return MappingProxyType({"api-key": api_key}) + token: Final = self.resolve_server_api_key( + caller_api_key=None, + caller_api_base=api_base, + key_env_vars=(TOKEN_ENV,), + base_env_var=PROJECT_ENDPOINT_ENV, + default_api_base=None, + ) or self._mint_entra_token(api_base) + return MappingProxyType({"Authorization": f"Bearer {token}"}) + + def _mint_entra_token(self, caller_api_base: str | None) -> str: + self._assert_trusted_api_base_for_server_credential( + caller_api_base, None, PROJECT_ENDPOINT_ENV, "Azure AD token" + ) + minter: Final = self._entra_token_minter or _default_entra_token_minter + try: + return minter() + except Exception as e: + raise ValueError( + f"Grounding with Bing Search: no credential available. Pass api_key, set {TOKEN_ENV} " + f"to an Entra bearer token, or configure azure-identity (AZURE_CLIENT_ID / " + f"AZURE_CLIENT_SECRET / AZURE_TENANT_ID or any DefaultAzureCredential source) " + f"for scope {ENTRA_SCOPE}. Underlying error: {e}" + ) from e + + def get_complete_url( + self, + api_base: str | None, + optional_params: dict[str, object], # mutable-ok: BaseSearchConfig.get_complete_url signature + data: dict[str, object] | list[dict[str, object]] | None = None, # mutable-ok: base signature + **kwargs: object, # kwargs-ok: BaseSearchConfig.get_complete_url signature + ) -> str: + resolved_base: Final = api_base or get_secret_str(PROJECT_ENDPOINT_ENV) + if not resolved_base: + raise ValueError( + f"{PROJECT_ENDPOINT_ENV} is not set. Set it to your Microsoft Foundry project " + f"endpoint, e.g. https://.services.ai.azure.com/api/projects/." + ) + trimmed: Final = resolved_base.rstrip("/") + if trimmed.endswith(_RESPONSES_PATH): + return trimmed + return f"{trimmed}{_RESPONSES_PATH}" + + def transform_search_request( + self, + query: str | list[str], # mutable-ok: BaseSearchConfig.transform_search_request signature + optional_params: dict[str, object], # mutable-ok: base signature + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_request signature + ) -> dict[str, object]: # mutable-ok: the http handler passes this straight to httpx as the JSON body + """ + Transform Search request to the Foundry Responses API format. + + The unified params map as far as the API allows: + - max_results -> the bing_grounding search configuration's `count`; the built-in + web_search tool has no result-count knob, so that mode instead caps the returned + results after the fact (see transform_search_response) + - country -> web_search's approximate `user_location` (bing_grounding's `market` + wants a full locale like en-US, which a bare country code cannot fill) + - search_domain_filter, max_tokens_per_page -> no API equivalent, dropped + """ + model: Final = get_secret_str(MODEL_ENV) + if not model: + raise ValueError( + f"{MODEL_ENV} is not set. Set it to a model deployment in the Foundry project " + f"that runs the grounded search, e.g. gpt-4.1." + ) + request: Final = _ResponsesRequest( + model=model, + input=" ".join(query) if isinstance(query, list) else query, + tools=(_search_tool(optional_params),), + ) + return request.model_dump(mode="json", exclude_none=True) + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_response signature + ) -> SearchResponse: + try: + parsed: Final = _ResponsesEnvelope.model_validate_json(raw_response.content) + except ValidationError as e: + raise self.get_error_class( + error_message=f"response does not match the Foundry Responses API schema: {e}", + status_code=raw_response.status_code, + headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature + ) + if parsed.status == "failed": + detail: Final = ( + parsed.error.message if parsed.error and parsed.error.message else "the grounded search failed" + ) + raise self._upstream_error(detail, raw_response) + results: Final = _capped(_citation_results(parsed), _requested_max_results(kwargs)) + if not results and parsed.status == "incomplete": + reason: Final = ( + parsed.incomplete_details.reason + if parsed.incomplete_details and parsed.incomplete_details.reason + else "unknown reason" + ) + raise self._upstream_error(f"the grounded search was incomplete: {reason}", raw_response) + return self._priced(results) + + def _upstream_error(self, detail: str, raw_response: httpx.Response) -> Exception: + return self.get_error_class( + error_message=detail, + status_code=_UPSTREAM_ERROR_STATUS, + headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature + ) + + def _priced(self, results: tuple[SearchResult, ...]) -> SearchResponse: + """web_search mode runs no paid Grounding with Bing transaction, so it must not + inherit the connection-mode ``bing_grounding/search`` price; zero its per-query + cost while leaving connection mode to the cost map.""" + response: Final = SearchResponse( + results=list(results), # mutable-ok: SearchResponse.results is list[SearchResult] + object="search", + ) + if get_secret_str(CONNECTION_ID_ENV): + return response + response._hidden_params[ + "additional_headers" + ] = { # mutable-ok: response_cost_calculator writes into _hidden_params + _RESPONSE_COST_HEADER: 0.0 + } + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str], # mutable-ok: BaseSearchConfig.get_error_class signature + ) -> Exception: + detail: Final = _unwrap_error_detail(error_message).rstrip(". ") + return BaseLLMException( + status_code=status_code, + message=f"Grounding with Bing Search: {detail}. See {_DOCS_URL} for details.", + headers=headers, + ) diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index 1a924088390..61cbc213b11 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -65,15 +65,24 @@ class AzureModelRouterConfig(AzureAIStudioConfig): Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07) and returns it with the azure_ai/ prefix for proper display and cost tracking. + + Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs, + response restamping) can read it instead of guessing the route from the model string. """ - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + AzureFoundryModelInfo, + ) + from litellm.router_utils.add_retry_fallback_headers import ( + get_hidden_params_dict, + ) # Get base model for the parent call (strips routing prefixes for API compatibility) base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model) # Call parent transform_response first - this will extract the actual model # from the raw response (e.g., "gpt-5-nano-2025-08-07") - model_response = super().transform_response( + transformed_response: Final = super().transform_response( model=base_model, raw_response=raw_response, model_response=model_response, @@ -86,7 +95,15 @@ class AzureModelRouterConfig(AzureAIStudioConfig): api_key=api_key, json_mode=json_mode, ) - return model_response + selected_model: Final = transformed_response.model + if selected_model: + # Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a + # class-level dict, so an in-place write can bleed into unrelated responses. + transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter # mutable-ok: ModelResponse requires _hidden_params to be a plain dict + **get_hidden_params_dict(transformed_response), + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model, + } + return transformed_response def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None: """ diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index bc8ea31ea8c..9e7161120cc 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -30,7 +30,12 @@ class AzureFoundryErrorStrings(str, enum.Enum): SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'" -NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ("thinking_blocks", "provider_specific_fields", "cache_control") +NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( + "thinking_blocks", + "reasoning_content", + "provider_specific_fields", + "cache_control", +) class AzureAIStudioConfig(OpenAIConfig): @@ -173,7 +178,8 @@ class AzureAIStudioConfig(OpenAIConfig): """ - Azure AI Studio doesn't support content as a list. This handles: 1. Strips message fields that are not part of the OpenAI chat-completions - schema (thinking_blocks, provider_specific_fields, cache_control). + schema (thinking_blocks, reasoning_content, provider_specific_fields, + cache_control). Azure AI Foundry backends set additionalProperties=false and reject these with "Extra inputs are not permitted", which breaks multi-turn Anthropic-format clients that echo thinking blocks back as history. diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index d25a8fd6561..26a90157455 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,9 +1,57 @@ +from collections.abc import Mapping from typing import Final, Literal import litellm from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues +from litellm.types.router import GenericLiteLLMParams + +AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"] + + +def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None: + """ + Resolve an Entra ID / OAuth access token for an Azure AI Foundry deployment. + + Accepts the same credential set as the `azure` provider: service principal + (`tenant_id` / `client_id` / `client_secret`), a pre-fetched `azure_ad_token`, an OIDC + federated token, username/password, or `DefaultAzureCredential` / managed identity. + """ + from litellm.llms.azure.common_utils import get_azure_ad_token + + params = GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams() + + return get_azure_ad_token(params) + + +def get_azure_ai_auth_headers( + api_key: str | None, + litellm_params: Mapping[str, object] | None = None, + api_key_header: AzureAIApiKeyHeader = "Authorization", + api_key_env_var: str = "AZURE_AI_API_KEY", +) -> Mapping[str, str]: + """ + Build the auth headers for an Azure AI Foundry route. + + Prefers the API key when one is configured, and otherwise falls back to Entra ID / OAuth, + sending the access token as a bearer token. + """ + if api_key: + return {api_key_header: f"Bearer {api_key}" if api_key_header == "Authorization" else api_key} + + azure_ad_token = get_azure_ai_entra_token(litellm_params=litellm_params) + if azure_ad_token: + return {"Authorization": f"Bearer {azure_ad_token}"} + + raise ValueError( + f"Missing Azure AI credentials - set an API key (`api_key` or {api_key_env_var}), or Entra ID / OAuth " + "credentials (`tenant_id` + `client_id` + `client_secret`, `azure_ad_token`, an OIDC token, or a managed " + "identity with `litellm.enable_azure_ad_token_refresh = True`)" + ) + + +AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model" class AzureFoundryModelInfo(BaseLLMModelInfo): @@ -37,13 +85,48 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): return "model_router" return "default" + @staticmethod + def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None: + """The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``. + + Reading this beats re-deriving the route from a model string: the stamp is set on the + code path that was actually taken, so it holds no matter what the caller named the model. + """ + if not hidden_params: + return None + selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY) + if isinstance(selected, str) and selected: + return selected + return None + + @staticmethod + def is_model_router_call( + model: str | None = None, + hidden_params: Mapping[str, object] | None = None, + ) -> bool: + """Whether a request went down the Azure Model Router route. + + Prefers the response stamp, then the deployment's litellm model path, and only then the + caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router + name heuristic lives in exactly one place. + """ + if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None: + return True + deployment_model: Final = ( + hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None + ) + return any( + isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router" + for candidate in (deployment_model, model) + ) + @staticmethod def get_api_base(api_base: str | None = None) -> str | None: return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE") @staticmethod def get_api_key(api_key: str | None = None) -> str | None: - return api_key or litellm.api_key or litellm.openai_key or get_secret_str("AZURE_AI_API_KEY") + return api_key or litellm.api_key or get_secret_str("AZURE_AI_API_KEY") @property def api_version(self, api_version: str | None = None) -> str | None: diff --git a/litellm/llms/azure_ai/image_edit/flux2_transformation.py b/litellm/llms/azure_ai/image_edit/flux2_transformation.py index 3aac08ddcaf..a09a80985b7 100644 --- a/litellm/llms/azure_ai/image_edit/flux2_transformation.py +++ b/litellm/llms/azure_ai/image_edit/flux2_transformation.py @@ -5,7 +5,10 @@ from typing import Any, Final from httpx._types import RequestFiles import litellm -from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, +) from litellm.llms.azure_ai.image_generation.flux_transformation import ( AzureFoundryFluxImageGenerationConfig, ) @@ -71,16 +74,13 @@ class AzureFoundryFlux2ImageEditConfig(OpenAIImageEditConfig): """ Validate Azure AI Foundry environment and set up authentication """ - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - if not api_key: - raise ValueError( - f"Azure AI API key is required for model {model}. Set AZURE_AI_API_KEY environment variable or pass api_key parameter." - ) - headers.update( { - "Api-Key": api_key, + **get_azure_ai_auth_headers( + api_key=AzureFoundryModelInfo.get_api_key(api_key), + litellm_params=litellm_params, + api_key_header="Api-Key", + ), "Content-Type": "application/json", } ) diff --git a/litellm/llms/azure_ai/image_edit/mai_transformation.py b/litellm/llms/azure_ai/image_edit/mai_transformation.py index 73bd9957b8f..e639c20292b 100644 --- a/litellm/llms/azure_ai/image_edit/mai_transformation.py +++ b/litellm/llms/azure_ai/image_edit/mai_transformation.py @@ -3,7 +3,10 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles -from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, +) from litellm.llms.azure_ai.image_generation.mai_transformation import ( AzureFoundryMAIImageGenerationConfig, ) @@ -91,15 +94,13 @@ class AzureFoundryMAIImageEditConfig(OpenAIImageEditConfig): litellm_params: dict | None = None, api_base: str | None = None, ) -> dict: - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - if not api_key: - raise ValueError( - f"Azure AI API key is required for model {model}. " - "Set AZURE_AI_API_KEY environment variable or pass api_key parameter." + headers.update( + get_azure_ai_auth_headers( + api_key=AzureFoundryModelInfo.get_api_key(api_key), + litellm_params=litellm_params, + api_key_header="api-key", ) - - headers.update({"api-key": api_key}) + ) return headers def get_complete_url( diff --git a/litellm/llms/azure_ai/image_edit/transformation.py b/litellm/llms/azure_ai/image_edit/transformation.py index efe3d8b2b88..1c626458df4 100644 --- a/litellm/llms/azure_ai/image_edit/transformation.py +++ b/litellm/llms/azure_ai/image_edit/transformation.py @@ -3,7 +3,10 @@ from typing import Final import httpx import litellm -from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo +from litellm.llms.azure_ai.common_utils import ( + AzureFoundryModelInfo, + get_azure_ai_auth_headers, +) from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.utils import _add_path_to_api_base @@ -30,19 +33,14 @@ class AzureFoundryFluxImageEditConfig(OpenAIImageEditConfig): ) -> dict: """ Validate Azure AI Foundry environment and set up authentication - Uses Api-Key header format + Uses the Api-Key header format, or an Entra ID / OAuth bearer token when no key is set """ - api_key = AzureFoundryModelInfo.get_api_key(api_key) - - if not api_key: - raise ValueError( - f"Azure AI API key is required for model {model}. Set AZURE_AI_API_KEY environment variable or pass api_key parameter." - ) - headers.update( - { - "Api-Key": api_key, # Azure AI Foundry uses Api-Key header format - } + get_azure_ai_auth_headers( + api_key=AzureFoundryModelInfo.get_api_key(api_key), + litellm_params=litellm_params, + api_key_header="Api-Key", + ) ) return headers diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index e7b94b3812b..a0e427eab9a 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -26,6 +26,7 @@ from litellm.constants import ( ) from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin, encode_url_path_segment +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.base_llm.ocr.transformation import ( OCR_REQUEST_FORMAT_PARAM, BaseOCRConfig, @@ -236,17 +237,13 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): """ Validate environment and return headers for Azure Document Intelligence. - Authentication uses Ocp-Apim-Subscription-Key header. + Authentication uses the Ocp-Apim-Subscription-Key header, or an Entra ID / OAuth bearer + token when no subscription key is set. """ # Get API key from environment if not provided if api_key is None: api_key = get_secret_str(AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR) - if api_key is None: - raise ValueError( - "Missing Azure Document Intelligence API Key - Set AZURE_DOCUMENT_INTELLIGENCE_API_KEY environment variable or pass api_key parameter" - ) - # Validate API base/endpoint is provided if api_base is None: api_base = get_secret_str("AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT") @@ -257,7 +254,12 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): ) headers = { - "Ocp-Apim-Subscription-Key": api_key, + **get_azure_ai_auth_headers( + api_key=api_key, + litellm_params=litellm_params, + api_key_header="Ocp-Apim-Subscription-Key", + api_key_env_var=AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR, + ), "Content-Type": "application/json", **headers, } diff --git a/litellm/llms/azure_ai/ocr/transformation.py b/litellm/llms/azure_ai/ocr/transformation.py index 24f96868eb3..dff6af71c99 100644 --- a/litellm/llms/azure_ai/ocr/transformation.py +++ b/litellm/llms/azure_ai/ocr/transformation.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, convert_url_to_base64, ) +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.base_llm.ocr.transformation import DocumentType, OCRRequestData from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.secret_managers.main import get_secret_str @@ -47,17 +48,12 @@ class AzureAIOCRConfig(MistralOCRConfig): """ Validate environment and return headers for Azure AI OCR. - Azure AI uses Bearer token authentication with AZURE_AI_API_KEY. + Authenticates with AZURE_AI_API_KEY, or with an Entra ID / OAuth token when no key is set. """ # Get API key from environment if not provided if api_key is None: api_key = get_secret_str(AZURE_AI_OCR_API_KEY_ENV_VAR) - if api_key is None: - raise ValueError( - "Missing Azure AI API Key - A call is being made to Azure AI but no key is set either in the environment variables or via params" - ) - # Validate API base is provided if api_base is None: api_base = get_secret_str("AZURE_AI_API_BASE") @@ -68,7 +64,7 @@ class AzureAIOCRConfig(MistralOCRConfig): ) headers = { - "Authorization": f"Bearer {api_key}", + **get_azure_ai_auth_headers(api_key=api_key, litellm_params=litellm_params), "Content-Type": "application/json", **headers, } diff --git a/litellm/llms/azure_ai/rerank/transformation.py b/litellm/llms/azure_ai/rerank/transformation.py index 1212d2c1689..64372c53f09 100644 --- a/litellm/llms/azure_ai/rerank/transformation.py +++ b/litellm/llms/azure_ai/rerank/transformation.py @@ -2,12 +2,14 @@ Translate between Cohere's `/rerank` format and Azure AI's `/rerank` format. """ +from collections.abc import Mapping from typing import Final import httpx import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.utils import RerankResponse @@ -64,15 +66,13 @@ class AzureAIRerankConfig(CohereRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("AZURE_AI_API_KEY") or litellm.azure_key - if api_key is None: - raise ValueError("Azure AI API key is required. Please set 'AZURE_AI_API_KEY' or 'litellm.azure_key'") - default_headers: Final = { - "Authorization": f"Bearer {api_key}", + **get_azure_ai_auth_headers(api_key=api_key, litellm_params=litellm_params), "accept": "application/json", "content-type": "application/json", } diff --git a/litellm/llms/base_llm/rerank/transformation.py b/litellm/llms/base_llm/rerank/transformation.py index 3a946fb4af4..5d2f92b5e82 100644 --- a/litellm/llms/base_llm/rerank/transformation.py +++ b/litellm/llms/base_llm/rerank/transformation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -24,6 +25,7 @@ class BaseRerankConfig(ABC): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: pass diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index 1aea3cafe33..f725b295d0f 100644 --- a/litellm/llms/base_llm/videos/transformation.py +++ b/litellm/llms/base_llm/videos/transformation.py @@ -1,9 +1,10 @@ import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx -from httpx._types import RequestFiles +from httpx._types import FileContent, RequestFiles from litellm.types.responses.main import * from litellm.types.router import GenericLiteLLMParams @@ -91,6 +92,14 @@ class BaseVideoConfig(ABC): raise ValueError("api_base is required") return api_base + def use_multipart_form_data(self) -> bool: + """ + Whether video create requests without files must still be sent as + multipart/form-data (the encoding the OpenAI SDK always uses for + /videos), instead of falling back to JSON. + """ + return False + @abstractmethod def transform_video_create_request( self, @@ -332,14 +341,18 @@ class BaseVideoConfig(ABC): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, + video_file: FileContent | None = None, extra_body: dict[str, Any] | None = None, prefetched_source_data: dict[str, Any] | None = None, - ) -> tuple[str, dict]: + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: """ - Transform the video edit request into a URL and JSON data. + Transform the video edit request into a URL plus either JSON data or + multipart form fields and files. Returns: - Tuple[str, Dict]: (url, data) for the POST request + tuple[str, Mapping[str, object], RequestFiles | None]: (url, data, + files). When files is None the handler sends data as JSON; otherwise + data holds the form fields and files holds the uploaded source video. """ raise NotImplementedError("video edit is not supported for this provider") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b437e25d24b..767677cbcbf 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1617,6 +1617,8 @@ class AmazonConverseConfig(BaseConfig): } if additional_request_params: data["additionalModelRequestFields"] = additional_request_params + if "thinking" in additional_request_params: + data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",) if system_content_blocks: data["system"] = system_content_blocks @@ -1801,6 +1803,17 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list + @staticmethod + def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None: + """Converse omits thinking tokens from its usage block; they only arrive under + ``additionalModelResponseFields`` when ``/usage/output_tokens_details`` is requested.""" + if not isinstance(additional_fields, Mapping): + return None + usage: Final = additional_fields.get("usage") + if not isinstance(usage, Mapping): + return None + return AnthropicConfig.thinking_tokens_from_usage(usage) + @staticmethod def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool: """Converse-family models report camelCase token counts, not Anthropic's snake_case.""" @@ -1842,6 +1855,7 @@ class AmazonConverseConfig(BaseConfig): usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, thinking_ran: bool = False, + provider_reasoning_tokens: int | None = None, ) -> Usage: input_tokens = usage["inputTokens"] output_tokens: Final = usage["outputTokens"] @@ -1862,9 +1876,14 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens: Final = ( + estimated_reasoning_tokens: Final = ( token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 ) + reasoning_tokens: Final = ( + min(max(0, provider_reasoning_tokens), output_tokens) + if provider_reasoning_tokens is not None + else estimated_reasoning_tokens + ) completion_tokens_details: Final = ( CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens, @@ -2272,6 +2291,9 @@ class AmazonConverseConfig(BaseConfig): completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), thinking_ran=reasoningContentBlocks is not None, + provider_reasoning_tokens=self.thinking_tokens_from_additional_fields( + completion_response.get("additionalModelResponseFields") + ), ) ## HANDLE TOOL CALLS diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index ce89c6c23e2..fc34e403beb 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -331,6 +331,7 @@ class AWSEventStreamDecoder: self.json_mode = json_mode self._current_tool_name: str | None = None self._thinking_ran = False + self._provider_reasoning_tokens: int | None = None def check_empty_tool_call_args(self) -> bool: """ @@ -559,14 +560,22 @@ class AWSEventStreamDecoder: tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) + self._provider_reasoning_tokens = AmazonConverseConfig.thinking_tokens_from_additional_fields( + chunk_data.get("additionalModelResponseFields") + ) elif "usage" in chunk_data: usage = converse_config.transform_usage( chunk_data.get("usage", {}), thinking_ran=self._thinking_ran, + provider_reasoning_tokens=self._provider_reasoning_tokens, ) if thinking_blocks: self._thinking_ran = True + carries_message_content: Final = any( + key in chunk_data for key in ("start", "delta", "contentBlockIndex", "stopReason", "trace") + ) + model_response_provider_specific_fields: Final = {} if "trace" in chunk_data: trace: Final = chunk_data.get("trace") @@ -577,8 +586,8 @@ class AWSEventStreamDecoder: finish_reason=finish_reason, index=0, # Always 0 - Bedrock never returns multiple choices delta=Delta( - content=text, - role="assistant", + content=text if carries_message_content else None, + role="assistant" if carries_message_content else None, tool_calls=[tool_use] if tool_use else None, provider_specific_fields=(provider_specific_fields if provider_specific_fields else None), thinking_blocks=thinking_blocks, diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 0ce2e6f60d3..d0a3c37ffb3 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, cast from httpx import Response @@ -93,6 +94,9 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD endpoint_url, ) + def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None: + return None + def sign_request( self, headers: dict, @@ -109,6 +113,7 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD request_data=request_data or {}, api_base=api_base, model=model, + api_key=self.get_bedrock_bearer_token(optional_params), ) def logging_non_streaming_response( diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 889361cd808..d877fbb4e09 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -13,6 +13,7 @@ global state. """ import re +from collections.abc import Mapping from typing import Final from botocore.exceptions import ( @@ -31,30 +32,39 @@ BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE) +def resolve_mantle_bearer_token(api_key: str | None) -> str | None: + return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + + +def resolve_mantle_region(params: Mapping[str, object]) -> str: + region: Final = params.get("aws_region_name") + if isinstance(region, str) and region: + BaseAWSLLM._validate_aws_region_name(region) + return region + api_base: Final = params.get("api_base") + base: Final = (api_base if isinstance(api_base, str) else None) or get_secret_str("BEDROCK_MANTLE_API_BASE") + if base: + match: Final = MANTLE_HOST_RE.match(base.rstrip("/")) + if match: + return match.group(1) + return ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION_NAME") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + + class BedrockMantleAuthMixin: _aws_signer: BaseAWSLLM @staticmethod def _resolve_bearer_token(api_key: str | None) -> str | None: - return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + return resolve_mantle_bearer_token(api_key) @staticmethod def _resolve_region(params: dict) -> str: - region: Final = params.get("aws_region_name") - if region: - BaseAWSLLM._validate_aws_region_name(region) - return region - base: Final = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE") - if base: - match: Final = MANTLE_HOST_RE.match(base.rstrip("/")) - if match: - return match.group(1) - return ( - get_secret_str("BEDROCK_MANTLE_REGION") - or get_secret_str("AWS_REGION_NAME") - or get_secret_str("AWS_REGION") - or BEDROCK_MANTLE_DEFAULT_REGION - ) + return resolve_mantle_region(params) def sign_request( self, diff --git a/litellm/llms/bedrock_mantle/passthrough/transformation.py b/litellm/llms/bedrock_mantle/passthrough/transformation.py new file mode 100644 index 00000000000..e6b831efa57 --- /dev/null +++ b/litellm/llms/bedrock_mantle/passthrough/transformation.py @@ -0,0 +1,71 @@ +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final, Literal, Optional + +from httpx import Response + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.llms.bedrock_mantle.common_utils import ( + MANTLE_HOST_RE, + resolve_mantle_bearer_token, + resolve_mantle_region, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.types.utils import CostResponseTypes + + +class BedrockMantlePassthroughConfig(BedrockPassthroughConfig): + """Native Bedrock runtime passthrough (InvokeModel, Converse) for deployments declared as bedrock_mantle. + + The Mantle host only serves the OpenAI-compatible surface, so a Mantle api_base lends its region and the + request itself goes to bedrock-runtime, signed with the deployment's Bearer token or SigV4 credentials. + """ + + def _get_aws_region_name( + self, + optional_params: Mapping[str, object], + model: str | None = None, + model_id: str | None = None, + ) -> str: + return resolve_mantle_region(optional_params) + + def get_runtime_endpoint( + self, + api_base: str | None, + aws_bedrock_runtime_endpoint: str | None, + aws_region_name: str, + endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime", + ) -> tuple[str, str]: + is_mantle_host: Final = api_base is not None and MANTLE_HOST_RE.match(api_base.rstrip("/")) is not None + return super().get_runtime_endpoint( + api_base=None if is_mantle_host else api_base, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + aws_region_name=aws_region_name, + endpoint_type=endpoint_type, + ) + + def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None: + api_key: Final = litellm_params.get("api_key") + return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None) + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: dict, # mutable-ok: mirrors the inherited BedrockPassthroughConfig signature + logging_obj: Logging, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + is_converse: Final = "invoke" not in endpoint and "converse" in endpoint + shape_provider: Final = LlmProviders.BEDROCK.value if is_converse else custom_llm_provider + return super().logging_non_streaming_response( + model=model, + custom_llm_provider=shape_provider, + httpx_response=httpx_response, + request_data=request_data, + logging_obj=logging_obj, + endpoint=endpoint, + ) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 92da5835b2d..3e5dd4ff87d 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -15,8 +15,12 @@ role / access key / profile / web identity), signed via the shared BaseAWSLLM._sign_request after the request body is finalized. """ +import json +from collections.abc import Mapping from typing import Any, Final +from typing_extensions import ReadOnly, TypedDict + import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -50,6 +54,33 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" +_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message" +_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction" +_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call" + + +class _RewrittenOutputTextBlock(TypedDict): + type: ReadOnly[str] + text: ReadOnly[str] + + +class _RewrittenAssistantMessageItem(TypedDict): + type: ReadOnly[str] + role: ReadOnly[str] + content: ReadOnly[tuple[_RewrittenOutputTextBlock, ...]] + + +class _RewrittenCompactionItem(TypedDict): + type: ReadOnly[str] + encrypted_content: ReadOnly[str] + + +class _RewrittenFunctionCallItem(TypedDict): + type: ReadOnly[str] + call_id: ReadOnly[str] + name: ReadOnly[str] + arguments: ReadOnly[str] + class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig): def __init__( @@ -155,6 +186,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI headers: dict, ) -> dict: remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input) + normalized_input: Final = self._normalize_codex_input_items(remaining_input) request_params: Final = ( { **response_api_optional_request_params, @@ -168,7 +200,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) return super().transform_responses_api_request( model=model, - input=remaining_input, + input=normalized_input, response_api_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, @@ -210,6 +242,91 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) return remaining_input, cls._filter_unsupported_tools(hoisted_tools) + @staticmethod + def _agent_message_text(item: "Mapping[str, Any]") -> str: + content: Final = item.get("content") + if not isinstance(content, list): + return "" + return "".join( + str(block.get("text") or block.get("encrypted_content") or "") + for block in content + if isinstance(block, dict) + ) + + @classmethod + def _normalize_agent_message_item(cls, item: "Mapping[str, Any]") -> "_RewrittenAssistantMessageItem | None": + text: Final = cls._agent_message_text(item) + if not text: + return None + rewritten: Final[_RewrittenAssistantMessageItem] = { + "type": "message", + "role": "assistant", + "content": ({"type": "output_text", "text": text},), + } + return rewritten + + @staticmethod + def _normalize_context_compaction_item(item: "Mapping[str, Any]") -> "_RewrittenCompactionItem | None": + encrypted_content: Final = item.get("encrypted_content") + if not isinstance(encrypted_content, str) or not encrypted_content: + return None + rewritten: Final[_RewrittenCompactionItem] = {"type": "compaction", "encrypted_content": encrypted_content} + return rewritten + + @staticmethod + def _normalize_local_shell_call_item(item: "Mapping[str, Any]") -> "_RewrittenFunctionCallItem | None": + call_id: Final = item.get("call_id") + if not isinstance(call_id, str) or not call_id: + return None + action: Final = item.get("action") + rewritten: Final[_RewrittenFunctionCallItem] = { + "type": "function_call", + "call_id": call_id, + "name": "local_shell", + "arguments": json.dumps(action) if isinstance(action, dict) else "{}", + } + return rewritten + + @classmethod + def _normalize_codex_input_item(cls, item: object) -> "tuple[object, str | None]": + """Returns (normalized item or None to drop it, original type when rewritten).""" + if not isinstance(item, dict): + return item, None + item_type: Final = item.get("type") + if item_type == _CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: + return cls._normalize_agent_message_item(item), item_type + if item_type == _CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: + return cls._normalize_context_compaction_item(item), item_type + if item_type == _CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: + return cls._normalize_local_shell_call_item(item), item_type + return item, None + + @classmethod + def _normalize_codex_input_items( + cls, + input: "str | ResponseInputParam", + ) -> "str | ResponseInputParam": + """Rewrite Codex history item types Mantle rejects with 400 "Invalid + 'input': value did not match any expected variant" into supported + equivalents. `agent_message` (Codex multi-agent traffic; its + encrypted_content slot carries the plaintext payload when the model + never issued encrypted args) becomes an assistant message, + `context_compaction` becomes the `compaction` spelling Mantle accepts, + and `local_shell_call` becomes the function_call its recorded + function_call_output already pairs with. + """ + if not isinstance(input, list): + return input + normalized: Final = tuple(cls._normalize_codex_input_item(item) for item in input) + rewritten_types: Final = sorted(frozenset(item_type for _, item_type in normalized if item_type is not None)) + if rewritten_types: + verbose_logger.warning( + "Bedrock Mantle Responses API: rewrote Codex input item type(s) %s that Mantle rejects.", + rewritten_types, + ) + kept: Final = [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list + return kept # pyright: ignore[reportReturnType] # Codex passthrough items sit outside the OpenAI input union + def map_openai_params( self, response_api_optional_params: ResponsesAPIOptionalRequestParams, diff --git a/litellm/llms/cerebras/chat.py b/litellm/llms/cerebras/chat.py index 8827b0afd87..c3aa26ade35 100644 --- a/litellm/llms/cerebras/chat.py +++ b/litellm/llms/cerebras/chat.py @@ -68,6 +68,8 @@ class CerebrasConfig(OpenAIGPTConfig): "tool_choice", "tools", "user", + "max_retries", + "extra_headers", ] # Only add reasoning_effort for models that support it diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index 76386252b79..a8e755406d8 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -81,6 +82,7 @@ class CohereRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("COHERE_API_KEY") or get_secret_str("CO_API_KEY") or litellm.cohere_key diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 52f30e31641..777ab576de2 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,7 +9,7 @@ import threading import time from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy -from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional, TypeAlias, TypedDict import certifi import httpx @@ -933,11 +933,83 @@ class AsyncHTTPHandler: response.raise_for_status() return response + # Strong references to finalizer-scheduled client-close tasks. A bare + # create_task() result may be garbage-collected before it runs, leaving + # the underlying aiohttp session unclosed ("Unclosed client session"). + # Mirrors LiteLLMAiohttpTransport._background_close_tasks. + _finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes + + @classmethod + def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None: + cls._finalizer_close_tasks.discard(task) + if task.cancelled(): + return + exc: Final = task.exception() + if exc is not None: + verbose_logger.debug("Error closing client at finalization: %s", exc) + + def _aiohttp_session_bound_elsewhere(self, loop: asyncio.AbstractEventLoop) -> bool: + """True when the wrapped aiohttp session is bound to a loop other than + ``loop`` — awaiting ``aclose()`` here would touch that loop's internals.""" + from litellm.llms.custom_httpx.aiohttp_transport import ( + LiteLLMAiohttpTransport, + ) + + transport: Final = getattr(self._client, "_transport", None) + if not isinstance(transport, LiteLLMAiohttpTransport): + return False + session: Final = transport.client + if not isinstance(session, ClientSession) or session.closed: + return False + return getattr(session, "_loop", None) is not loop + + def _dispose_wrapped_aiohttp_session(self) -> None: + """Dispose the wrapped aiohttp session when ``aclose()`` cannot run here. + + Finalization either has no running loop, or a loop the session is not + bound to. Delegating to the transport's lifecycle-aware disposal picks + the safe path per session state (async close on its own loop, threadsafe + handoff to a loop running elsewhere, or the synchronous connector + teardown that flips the flags ``ClientSession.__del__`` checks), so no + "Unclosed client session" / "Unclosed connector" warnings fire at + garbage collection. + """ + from litellm.llms.custom_httpx.aiohttp_transport import ( + LiteLLMAiohttpTransport, + ) + + transport: Final = getattr(self._client, "_transport", None) + if not isinstance(transport, LiteLLMAiohttpTransport): + return + # A shared session (e.g. the proxy's) is never this handler's to close. + if not getattr(transport, "_owns_session", False): + return + session: Final = transport.client + if isinstance(session, ClientSession) and not session.closed: + transport._close_recycled_session(session) # pyright: ignore[reportPrivateUsage] # deliberate reuse of the transport's lifecycle-aware disposal; an async close can never run in this context + def __del__(self) -> None: try: if not _handler_may_close_client(sys.getrefcount(self._client), self._owns_client): return - asyncio.get_running_loop().create_task(self._client.aclose()) + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + # No running loop at finalization time (worker threads after + # their loop closed, interpreter/worker shutdown, GC in a + # sync context). An async close can never run here. + self._dispose_wrapped_aiohttp_session() + return + if self._aiohttp_session_bound_elsewhere(loop): + # GC ran on a live loop (e.g. the app's) but the session + # belongs to another, possibly dead, loop — awaiting aclose() + # here is the cross-loop path the transport refuses. + self._dispose_wrapped_aiohttp_session() + return + task: Final = loop.create_task(self._client.aclose()) + cls: Final = type(self) + cls._finalizer_close_tasks.add(task) + task.add_done_callback(cls._on_finalizer_close_done) except Exception: pass diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ed079197513..b7ddb55ae89 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, Type from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx +from httpx._types import FileContent from openai.types.file_deleted import FileDeleted import litellm @@ -24,6 +25,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( validated_max_agentic_loops, ) from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import realtime_error_event, websocket_close_reason from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -1108,6 +1110,7 @@ class BaseLLMHTTPHandler: headers=headers or {}, model=model, optional_params=optional_rerank_params, + litellm_params=litellm_params, ) api_base = provider_config.get_complete_url( @@ -1844,6 +1847,7 @@ class BaseLLMHTTPHandler: return provider_config.transform_search_response( raw_response=response, logging_obj=logging_obj, + optional_params=optional_params, ) async def async_search( @@ -1942,6 +1946,7 @@ class BaseLLMHTTPHandler: return provider_config.transform_search_response( raw_response=response, logging_obj=logging_obj, + optional_params=optional_params, ) async def _async_post_anthropic_messages_with_http_error_retry( @@ -7050,9 +7055,7 @@ class BaseLLMHTTPHandler: ) try: - # Use JSON when no files, otherwise use form data with files if files and len(files) > 0: - # Use multipart/form-data when files are present response = sync_httpx_client.post( url=api_base, headers=headers, @@ -7060,9 +7063,14 @@ class BaseLLMHTTPHandler: files=files, timeout=timeout, ) - + elif video_generation_provider_config.use_multipart_form_data(): + response = sync_httpx_client.post( # rebind-ok: one of three mutually-exclusive branches + url=api_base, + headers=headers, + files=serialize_multipart_form_fields(data), + timeout=timeout, + ) else: - # Use JSON content type for POST requests without files response = sync_httpx_client.post( url=api_base, headers=headers, @@ -7154,20 +7162,26 @@ class BaseLLMHTTPHandler: ) try: - # Use JSON when no files, otherwise use form data with files - if files is None or len(files) == 0: + if files and len(files) > 0: response = await async_httpx_client.post( url=api_base, headers=headers, - json=data, + data=data, + files=files, + timeout=timeout, + ) + elif video_generation_provider_config.use_multipart_form_data(): + response = await async_httpx_client.post( # rebind-ok: one of three mutually-exclusive branches + url=api_base, + headers=headers, + files=serialize_multipart_form_fields(data), timeout=timeout, ) else: response = await async_httpx_client.post( url=api_base, headers=headers, - data=data, - files=files, + json=data, timeout=timeout, ) @@ -7827,6 +7841,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, + video_file: FileContent | None = None, extra_headers: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, timeout: float | None = None, @@ -7838,6 +7853,7 @@ class BaseLLMHTTPHandler: return self.async_video_edit_handler( prompt=prompt, video_id=video_id, + video_file=video_file, video_provider_config=video_provider_config, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, @@ -7891,9 +7907,10 @@ class BaseLLMHTTPHandler: prefetched_source_data = prefetch_resp.json() try: - url, data = video_provider_config.transform_video_edit_request( + url, data, files = video_provider_config.transform_video_edit_request( prompt=prompt, video_id=video_id, + video_file=video_file, api_base=api_base, litellm_params=litellm_params, headers=headers, @@ -7912,11 +7929,10 @@ class BaseLLMHTTPHandler: }, ) - response: Final = sync_httpx_client.post( - url=url, - headers=headers, - json=data, - timeout=timeout, + response: Final = ( + sync_httpx_client.post(url=url, headers=headers, data=data, files=files, timeout=timeout) + if files + else sync_httpx_client.post(url=url, headers=headers, json=data, timeout=timeout) ) response.raise_for_status() return video_provider_config.transform_video_edit_response( @@ -7936,6 +7952,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, + video_file: FileContent | None = None, extra_headers: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, timeout: float | None = None, @@ -7987,9 +8004,10 @@ class BaseLLMHTTPHandler: prefetched_source_data = prefetch_resp.json() try: - url, data = video_provider_config.transform_video_edit_request( + url, data, files = video_provider_config.transform_video_edit_request( prompt=prompt, video_id=video_id, + video_file=video_file, api_base=api_base, litellm_params=litellm_params, headers=headers, @@ -8008,11 +8026,10 @@ class BaseLLMHTTPHandler: }, ) - response: Final = await async_httpx_client.post( - url=url, - headers=headers, - json=data, - timeout=timeout, + response: Final = await ( + async_httpx_client.post(url=url, headers=headers, data=data, files=files, timeout=timeout) + if files + else async_httpx_client.post(url=url, headers=headers, json=data, timeout=timeout) ) response.raise_for_status() return video_provider_config.transform_video_edit_response( diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 3c7801d4d3c..98be4e4f2e7 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -22,6 +22,7 @@ as supported only for gte-rerank-v2 / qwen3-vl-rerank. Docs - https://help.aliyun.com/zh/model-studio/text-rerank-api """ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -85,6 +86,7 @@ class DashScopeRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("DASHSCOPE_API_KEY") diff --git a/litellm/llms/databricks/cost_calculator.py b/litellm/llms/databricks/cost_calculator.py index 05647883ebf..64166e6fc11 100644 --- a/litellm/llms/databricks/cost_calculator.py +++ b/litellm/llms/databricks/cost_calculator.py @@ -3,10 +3,31 @@ Helper util for handling databricks-specific cost calculation - e.g.: handling 'dbrx-instruct-*' """ +from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import Usage -from litellm.utils import get_model_info + +_LEGACY_ENDPOINT_NAMES: Final = MappingProxyType( + { + "dbrx-instruct": "databricks-dbrx-instruct", + "meta-llama-3.1-70b-instruct": "databricks-meta-llama-3-1-70b-instruct", + "meta-llama-3.1-405b-instruct": "databricks-meta-llama-3-1-405b-instruct", + "mixtral-8x7b-instruct-v0.1": "databricks-mixtral-8x7b-instruct", + "bge-large-en": "databricks-bge-large-en", + "gte-large-en": "databricks-gte-large-en", + "llama-2-70b-chat": "databricks-llama-2-70b-chat", + } +) + + +def _registry_key(model: str) -> str: + name: Final = model.removeprefix("databricks/") + return next( + (key for prefix, key in _LEGACY_ENDPOINT_NAMES.items() if name.startswith(prefix)), + name, + ) def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: @@ -20,36 +41,8 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ - base_model = model - if model.startswith("databricks/dbrx-instruct") or model.startswith("dbrx-instruct"): - base_model = "databricks-dbrx-instruct" - elif model.startswith("databricks/meta-llama-3.1-70b-instruct") or model.startswith("meta-llama-3.1-70b-instruct"): - base_model = "databricks-meta-llama-3-1-70b-instruct" - elif model.startswith("databricks/meta-llama-3.1-405b-instruct") or model.startswith( - "meta-llama-3.1-405b-instruct" - ): - base_model = "databricks-meta-llama-3-1-405b-instruct" - elif ( - model.startswith("databricks/mixtral-8x7b-instruct-v0.1") - or model.startswith("mixtral-8x7b-instruct-v0.1") - or model.startswith("databricks/mixtral-8x7b-instruct-v0.1") - or model.startswith("mixtral-8x7b-instruct-v0.1") - ): - base_model = "databricks-mixtral-8x7b-instruct" - elif model.startswith("databricks/bge-large-en") or model.startswith("bge-large-en"): - base_model = "databricks-bge-large-en" - elif model.startswith("databricks/gte-large-en") or model.startswith("gte-large-en"): - base_model = "databricks-gte-large-en" - elif model.startswith("databricks/llama-2-70b-chat") or model.startswith("llama-2-70b-chat"): - base_model = "databricks-llama-2-70b-chat" - ## GET MODEL INFO - model_info: Final = get_model_info(model=base_model, custom_llm_provider="databricks") - - ## CALCULATE INPUT COST - - prompt_cost: Final[float] = usage["prompt_tokens"] * model_info["input_cost_per_token"] - - ## CALCULATE OUTPUT COST - completion_cost: Final = usage["completion_tokens"] * model_info["output_cost_per_token"] - - return prompt_cost, completion_cost + return generic_cost_per_token( + model=_registry_key(model), + usage=usage, + custom_llm_provider="databricks", + ) diff --git a/litellm/llms/deepinfra/rerank/transformation.py b/litellm/llms/deepinfra/rerank/transformation.py index 366b82e1dcf..e52c56af82b 100644 --- a/litellm/llms/deepinfra/rerank/transformation.py +++ b/litellm/llms/deepinfra/rerank/transformation.py @@ -2,6 +2,7 @@ Translate between Cohere's `/rerank` format and Deepinfra's `/rerank` format. """ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -67,6 +68,7 @@ class DeepinfraRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("DEEPINFRA_API_KEY") diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index e64237da978..4e9731ef485 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -504,6 +504,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): m = cast(dict, message) m.pop("provider_specific_fields", None) m.pop("thinking_blocks", None) + m.pop("reasoning_content", None) return messages diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index fde4f55e75b..8ef2c9acccb 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -4,6 +4,7 @@ Fireworks AI Rerank API transformation Reference: https://docs.fireworks.ai/inference-api-reference/rerank """ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -102,6 +103,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: api_key = self._get_api_key(api_key) if api_key is None: diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index db75be3b6a2..6a1fc144c42 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -566,6 +566,7 @@ class GeminiVideoConfig(BaseVideoConfig): api_base, litellm_params, headers, + video_file=None, extra_body=None, prefetched_source_data=None, ): diff --git a/litellm/llms/hosted_vllm/chat/transformation.py b/litellm/llms/hosted_vllm/chat/transformation.py index 46a2320b655..29dc485732f 100644 --- a/litellm/llms/hosted_vllm/chat/transformation.py +++ b/litellm/llms/hosted_vllm/chat/transformation.py @@ -164,12 +164,13 @@ class HostedVLLMChatConfig(OpenAIGPTConfig): """ Support translating: - video files from file_id or file_data to video_url - - thinking_blocks on assistant messages are removed, and content lists - are converted to strings for vLLM compatibility + - thinking_blocks and reasoning_content on assistant messages are removed, + and content lists are converted to strings for vLLM compatibility """ for message in messages: if message["role"] == "assistant": message.pop("thinking_blocks", None) + message.pop("reasoning_content", None) existing_content = message.get("content") if isinstance(existing_content, list): text_parts = [] diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 74c13b450f5..0e8fa294f5d 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -2,6 +2,7 @@ Transformation logic for Hosted VLLM rerank """ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -107,6 +108,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key" diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index d56a76c933f..334c60ee848 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -1,4 +1,5 @@ import os +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final import httpx @@ -123,6 +124,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, api_base: str | None = None, ) -> dict: # Get API credentials diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index ebcbf1b5a07..2526eb3b6a4 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -4,6 +4,7 @@ Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank` Why separate file? Make it easy to see how transformation works """ +from collections.abc import Mapping from typing import Final import httpx @@ -46,6 +47,7 @@ class InfinityRerankConfig(CohereRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("INFINITY_API_KEY") or get_secret_str("INFINITY_API_KEY") or litellm.infinity_key diff --git a/litellm/llms/jina_ai/rerank/transformation.py b/litellm/llms/jina_ai/rerank/transformation.py index 25607443292..199599d6b9c 100644 --- a/litellm/llms/jina_ai/rerank/transformation.py +++ b/litellm/llms/jina_ai/rerank/transformation.py @@ -6,6 +6,7 @@ Why separate file? Make it easy to see how transformation works Docs - https://jina.ai/reranker """ +from collections.abc import Mapping from typing import Any, Final from httpx import URL, Response @@ -139,6 +140,7 @@ class JinaAIRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: raise ValueError("api_key is required. Set via `api_key` parameter or `JINA_API_KEY` environment variable.") diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index bb07f9ec74f..93e00dad9a1 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import Any, Final, Literal import httpx @@ -152,6 +153,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: """ Validate that the Nvidia NIM API key is present. diff --git a/litellm/llms/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index ffa3de0d5c6..3a65e4a9426 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -16,7 +16,7 @@ def _normalize_reasoning_effort_for_chat_completion( ) -> str | None: """Convert reasoning_effort to the string format expected by OpenAI chat completion API. - The chat completion API expects a simple string: 'none', 'low', 'medium', 'high', or 'xhigh'. + The chat completion API expects an effort string such as 'low' or 'high'. Config/deployments may pass the Responses API format: {'effort': 'high', 'summary': 'detailed'}. """ if value is None: diff --git a/litellm/llms/openai/cost_calculation.py b/litellm/llms/openai/cost_calculation.py index 0352d246c09..115b2e27983 100644 --- a/litellm/llms/openai/cost_calculation.py +++ b/litellm/llms/openai/cost_calculation.py @@ -134,14 +134,7 @@ def cost_per_second(model: str, custom_llm_provider: str | None, duration: float def _video_resolution_to_cost_field_suffix(resolution: str) -> str | None: - """ - Map usage resolution to a safe suffix for ``output_cost_per_second_`` keys. - - Note: Currently only ``output_cost_per_second_1080p`` is explicitly declared in - ModelInfo (types/utils.py). Other resolution tiers (e.g., 720p, 4k) can be added - to model_prices_and_context_window.json but are not exposed via get_model_info() - until added to the ModelInfo TypedDict. - """ + """Map usage resolution to a safe suffix for ``output_cost_per_second_`` keys.""" r: Final = resolution.strip().lower() if not r: return None diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 50b466ae996..f1b6dcb330a 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -1,5 +1,7 @@ import mimetypes +from collections.abc import Mapping from io import BufferedReader, BytesIO +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import quote @@ -101,6 +103,9 @@ class OpenAIVideoConfig(BaseVideoConfig): return f"{api_base.rstrip('/')}/videos" + def use_multipart_form_data(self) -> bool: + return True + def transform_video_create_request( self, model: str, @@ -499,15 +504,26 @@ class OpenAIVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, + video_file: FileContent | None = None, extra_body: dict[str, object] | None = None, prefetched_source_data: dict[str, object] | None = None, - ) -> tuple[str, dict]: - original_video_id: Final = extract_original_video_id(video_id) + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: url: Final = f"{api_base.rstrip('/')}/edits" + + if video_file is not None: + files: Final[RequestFiles] = (self._video_file_tuple(video_file, "video"),) + form_data: Final = ( + MappingProxyType({"prompt": prompt, **extra_body}) + if extra_body + else MappingProxyType({"prompt": prompt}) + ) + return url, form_data, files + + original_video_id: Final = extract_original_video_id(video_id) data: Final[dict[str, object]] = {"prompt": prompt, "video": {"id": original_video_id}} if extra_body: data.update(extra_body) - return url, data + return url, data, None def transform_video_edit_response( self, @@ -567,21 +583,22 @@ class OpenAIVideoConfig(BaseVideoConfig): else: files_list.append((field_name, ("input_reference.png", image, image_content_type))) + def _video_file_tuple(self, video: FileContent, field_name: str) -> tuple[str, FileTypes]: + """ + Build a multipart field tuple for a video upload with proper video MIME + type detection: these paths must send video/mp4, not image/* content types. + """ + filename: Final = getattr(video, "name", None) or "input_video.mp4" + content_type: Final = self._get_video_content_type(video=video, filename=filename) + return (field_name, (filename, video, content_type)) + def _add_video_to_files( self, files_list: list[tuple[str, FileTypes]], video: FileContent, field_name: str, ) -> None: - """ - Add a video to files with proper video MIME type detection. - - This path is used by POST /videos/characters and must send video/mp4, - not image/* content types. - """ - filename: Final = getattr(video, "name", None) or "input_video.mp4" - content_type: Final = self._get_video_content_type(video=video, filename=filename) - files_list.append((field_name, (filename, video, content_type))) + files_list.append(self._video_file_tuple(video, field_name)) def _get_video_content_type(self, video: FileContent, filename: str) -> str: guessed_content_type, _ = mimetypes.guess_type(filename) diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index b8e57fa7cc0..8dd77db08c0 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -1,5 +1,6 @@ from collections.abc import Mapping, Sequence from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal import httpx @@ -33,6 +34,10 @@ else: LiteLLMLoggingObj = Any +class RunwayMLError(BaseLLMException): + pass + + class _RunwayTaskResponse(TypedDict, total=False): id: ReadOnly[str] status: ReadOnly[str] @@ -41,7 +46,8 @@ class _RunwayTaskResponse(TypedDict, total=False): output: ReadOnly[Sequence[str] | str] failureCode: ReadOnly[str] failure: ReadOnly[str] - progress: ReadOnly[int] + progress: ReadOnly[float] + estimatedCost: ReadOnly[Mapping[str, float]] class _VideoObjectData(TypedDict, extra_items=object): @@ -56,12 +62,54 @@ def _parse_runway_task_response(raw_response: httpx.Response) -> _RunwayTaskResp return response_data +_USD_PER_CREDIT: Final = 0.01 + +_RESOLUTION_AREA_TIERS: Final[tuple[tuple[int, str], ...]] = ( + (600_000, "480p"), + (1_500_000, "720p"), + (4_000_000, "1080p"), +) + + +def _ratio_to_resolution(ratio: object) -> str | None: + if not isinstance(ratio, str) or ":" not in ratio: + return None + width_str, _, height_str = ratio.partition(":") + if not (width_str.isdigit() and height_str.isdigit()): + return None + area: Final = int(width_str) * int(height_str) + return next((label for threshold, label in _RESOLUTION_AREA_TIERS if area < threshold), "4k") + + +def _duration_seconds(seconds: str | None) -> float | None: + if not seconds: + return None + try: + return float(seconds) + except ValueError: + return None + + +def _estimated_cost_usd(response_data: _RunwayTaskResponse) -> float | None: + estimated_cost: Final = response_data.get("estimatedCost") + if not isinstance(estimated_cost, Mapping): + return None + credits: Final = estimated_cost.get("credits") + if not isinstance(credits, (int, float)): + return None + return float(credits) * _USD_PER_CREDIT + + +def _progress_percent(progress: float) -> int: + return min(100, max(0, round(float(progress) * 100))) + + class RunwayMLVideoConfig(BaseVideoConfig): """ Configuration class for RunwayML video generation. RunwayML uses a task-based API where: - 1. POST /v1/image_to_video creates a task + 1. POST /v1/text_to_video, /v1/image_to_video, or /v1/video_to_video creates a task 2. The task returns immediately with a task ID 3. Client must poll or wait for task completion """ @@ -195,31 +243,36 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ Transform the video creation request for RunwayML API. - RunwayML expects: - { - "model": "gen4_turbo", - "promptImage": "https://... or data:image/...", - "promptText": "description", - "ratio": "1280:720", - "duration": 5 - } + RunwayML has three generation endpoints discriminated by which input is + present, and each request body rejects unknown fields: + - /text_to_video: promptText only (rejects promptImage) + - /image_to_video: promptImage (+ optional promptText) + - /video_to_video: promptVideo or videoUri (rejects promptImage) """ - # Build the request data + merged_params: Final = MappingProxyType( + { + "model": model, + "promptText": prompt, + **video_create_optional_request_params, + } + ) + + endpoint: Final = self._select_generation_endpoint(merged_params) + request_data: Final[dict[str, object]] = { - "model": model, - "promptText": prompt, + key: value for key, value in merged_params.items() if endpoint == "image_to_video" or key != "promptImage" } - # Add mapped parameters - request_data.update(video_create_optional_request_params) - - # RunwayML uses JSON body, no files multipart files_list: Final[RequestFiles] = [] - # Append the specific endpoint for video generation - full_api_base: Final = f"{api_base}/image_to_video" + return request_data, files_list, f"{api_base}/{endpoint}" - return request_data, files_list, full_api_base + def _select_generation_endpoint(self, request_data: Mapping[str, object]) -> str: + if request_data.get("promptVideo") is not None or request_data.get("videoUri") is not None: + return "video_to_video" + if request_data.get("promptImage") is not None: + return "image_to_video" + return "text_to_video" def transform_video_create_response( self, @@ -285,13 +338,15 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) # Add usage data for cost tracking - usage_data: Final = {} - if video_obj and hasattr(video_obj, "seconds") and video_obj.seconds: - try: - usage_data["duration_seconds"] = float(video_obj.seconds) - except (ValueError, TypeError): - pass - video_obj.usage = usage_data + video_obj.usage = { + key: value + for key, value in ( + ("duration_seconds", _duration_seconds(video_obj.seconds)), + ("video_resolution", _ratio_to_resolution(request_data.get("ratio") if request_data else None)), + ("provider_reported_cost_usd", _estimated_cost_usd(response_data)), + ) + if value is not None + } return video_obj @@ -581,8 +636,9 @@ class RunwayMLVideoConfig(BaseVideoConfig): if "completedAt" in response_data: video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) - if "progress" in response_data: - video_data["progress"] = response_data["progress"] + progress_value: Final = response_data.get("progress") + if progress_value is not None: + video_data["progress"] = _progress_percent(progress_value) if "failureCode" in response_data or "failure" in response_data: video_data["error"] = { @@ -616,6 +672,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): api_base, litellm_params, headers, + video_file=None, extra_body=None, prefetched_source_data=None, ): @@ -646,9 +703,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): raise NotImplementedError("video extension is not supported for RunwayML") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - from ...base_llm.chat.transformation import BaseLLMException - - raise BaseLLMException( + return RunwayMLError( status_code=status_code, message=error_message, headers=headers, diff --git a/litellm/llms/together_ai/chat.py b/litellm/llms/together_ai/chat.py deleted file mode 100644 index 58d47e45faa..00000000000 --- a/litellm/llms/together_ai/chat.py +++ /dev/null @@ -1,58 +0,0 @@ -""" -Support for OpenAI's `/v1/chat/completions` endpoint. - -Calls done in OpenAI/openai.py as TogetherAI is openai-compatible. - -Docs: https://docs.together.ai/reference/completions-1 -""" - -from typing import Final - -from litellm._logging import verbose_logger -from litellm.utils import supports_function_calling - -from ..openai.chat.gpt_transformation import OpenAIGPTConfig - - -class TogetherAIConfig(OpenAIGPTConfig): - def get_supported_openai_params(self, model: str) -> list: - """ - Only some together models support response_format / tool calling - - Docs: https://docs.together.ai/docs/json-mode - """ - # Use supports_function_calling() — which reads _get_model_info_helper - # directly — instead of get_model_info(). get_model_info() calls - # get_supported_openai_params() as its first step, which routes back - # into this method for together_ai models, creating a recursion that - # only terminates when Python's recursion limit or the "not mapped" - # exception in _get_model_info_helper is hit (~332 deep calls). - supports_fc: bool | None = None - try: - supports_fc = supports_function_calling(model, custom_llm_provider="together_ai") - except Exception as e: - verbose_logger.debug("Error getting supported openai params: %s", e) - - optional_params: Final = super().get_supported_openai_params(model) - if supports_fc is not True: - verbose_logger.debug( - "Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling" - ) - optional_params.remove("tools") - optional_params.remove("tool_choice") - optional_params.remove("function_call") - optional_params.remove("response_format") - return optional_params - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) - - if "response_format" in mapped_openai_params and mapped_openai_params["response_format"] == {"type": "text"}: - mapped_openai_params.pop("response_format") - return mapped_openai_params diff --git a/litellm/llms/together_ai/chat/__init__.py b/litellm/llms/together_ai/chat/__init__.py new file mode 100644 index 00000000000..f260d9126d7 --- /dev/null +++ b/litellm/llms/together_ai/chat/__init__.py @@ -0,0 +1,3 @@ +from .transformation import TogetherAIChatConfig as TogetherAIChatConfig + +TogetherAIConfig = TogetherAIChatConfig diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py new file mode 100644 index 00000000000..88fd79f2366 --- /dev/null +++ b/litellm/llms/together_ai/chat/transformation.py @@ -0,0 +1,48 @@ +""" +Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/completions`. + +Docs: https://docs.together.ai/docs/chat-overview +""" + +from types import MappingProxyType +from typing import Final + +from litellm._logging import verbose_logger +from litellm.utils import supports_function_calling + +from ...openai.chat.gpt_transformation import OpenAIGPTConfig + +FUNCTION_CALLING_ONLY_PARAMS: Final = ("tools", "tool_choice", "function_call", "response_format") +PLAIN_TEXT_RESPONSE_FORMAT: Final = MappingProxyType({"type": "text"}) + + +class TogetherAIChatConfig(OpenAIGPTConfig): + def get_supported_openai_params(self, model: str) -> list: + supports_fc: bool | None = None + try: + supports_fc = supports_function_calling(model, custom_llm_provider="together_ai") + except Exception as e: + verbose_logger.debug("Error getting supported openai params: %s", e) + + supported_params: Final = super().get_supported_openai_params(model) + if supports_fc is True: + return supported_params + verbose_logger.debug( + "Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling" + ) + return [ # mutable-ok: the inherited contract returns a plain list; building fresh avoids mutating the base class's value + param for param in supported_params if param not in FUNCTION_CALLING_ONLY_PARAMS + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params) + + if mapped_openai_params.get("response_format") == PLAIN_TEXT_RESPONSE_FORMAT: + mapped_openai_params.pop("response_format") + return mapped_openai_params diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py index 10246451a9d..b8079e52c97 100644 --- a/litellm/llms/together_ai/rerank/handler.py +++ b/litellm/llms/together_ai/rerank/handler.py @@ -16,11 +16,16 @@ from litellm.llms.together_ai.rerank.transformation import TogetherAIRerankConfi from litellm.types.rerank import RerankRequest, RerankResponse +def _rerank_url(api_base: str) -> str: + return f"{api_base.rstrip('/')}/rerank" + + class TogetherAIRerank(BaseLLM): def rerank( self, model: str, api_key: str, + api_base: str, query: str, documents: list[str | dict[str, Any]], top_n: int | None = None, @@ -46,10 +51,10 @@ class TogetherAIRerank(BaseLLM): raise ValueError("TogetherAI does not support max_chunks_per_doc") if _is_async: - return self.async_rerank(request_data_dict, api_key) # Call async method + return self.async_rerank(request_data_dict, api_key, api_base) response: Final = client.post( - "https://api.together.xyz/v1/rerank", + _rerank_url(api_base), headers={ "accept": "application/json", "content-type": "application/json", @@ -69,11 +74,12 @@ class TogetherAIRerank(BaseLLM): self, request_data_dict: dict[str, Any], api_key: str, + api_base: str, ) -> RerankResponse: client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.TOGETHER_AI) # Use async client response: Final = await client.post( - "https://api.together.xyz/v1/rerank", + _rerank_url(api_base), headers={ "accept": "application/json", "content-type": "application/json", diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py new file mode 100644 index 00000000000..0764a8bea62 --- /dev/null +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -0,0 +1,149 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig +from litellm.llms.vertex_ai.common_utils import validate_vertex_location +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +VERTEX_INTERACTIONS_API_VERSION: Final = "v1beta1" +VERTEX_INTERACTIONS_DEFAULT_LOCATION: Final = "global" + + +@dataclass(frozen=True, slots=True) +class VertexInteractionsTarget: + base_url: str + project_id: str + location: str + + @property + def collection_url(self) -> str: + return ( + f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}" + f"/projects/{self.project_id}/locations/{self.location}/interactions" + ) + + def interaction_url(self, interaction_id: str) -> str: + encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") + return f"{self.collection_url}/{encoded_interaction_id}" + + +class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig): + def __init__( + self, + mint_access_token: Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]] | None = None, + ) -> None: + super().__init__() + self._mint_access_token: Final[Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]]] = ( + mint_access_token or self._mint_access_token_with_vertex_base + ) + + def _mint_access_token_with_vertex_base( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + return self._ensure_access_token( + credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai" + ) + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.VERTEX_AI + + @property + def api_version(self) -> str: + return VERTEX_INTERACTIONS_API_VERSION + + def get_default_vertex_location(self) -> str: + return VERTEX_INTERACTIONS_DEFAULT_LOCATION + + def _mint(self, litellm_params: GenericLiteLLMParams) -> tuple[str, str]: + raw_params: Final = litellm_params.model_dump() + return self._mint_access_token( + self.safe_get_vertex_ai_credentials(raw_params), + self.safe_get_vertex_ai_project(raw_params), + ) + + def _target(self, api_base: str | None, litellm_params: GenericLiteLLMParams) -> VertexInteractionsTarget: + _, project_id = self._mint(litellm_params) + if not project_id: + raise ValueError( + "Vertex AI project is required. Set vertex_project, litellm.vertex_project, or VERTEXAI_PROJECT" + ) + location: Final = validate_vertex_location( + self.explicit_vertex_ai_location(litellm_params.model_dump()) or VERTEX_INTERACTIONS_DEFAULT_LOCATION + ) + return VertexInteractionsTarget( + base_url=self.get_api_base(api_base or None, location), + project_id=project_id, + location=location, + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict: # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers + access_token, _ = self._mint(litellm_params or GenericLiteLLMParams()) + return { # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers + "Content-Type": "application/json", + "Authorization": f"Bearer {access_token}", + **headers, + } + + def get_complete_url( + self, + api_base: str | None, + model: str | None, + agent: str | None = None, + litellm_params: Mapping[str, object] | None = None, + stream: bool | None = None, + ) -> str: + params: Final = ( + GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams() + ) + collection_url: Final = self._target(api_base, params).collection_url + return f"{collection_url}?alt=sse" if stream else collection_url + + def _interaction_by_id_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + url_suffix: str = "", + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + target: Final = self._target(api_base or None, litellm_params) + return f"{target.interaction_url(interaction_id)}{url_suffix}", {} # mutable-ok: same base contract + + def transform_get_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params) + + def transform_delete_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params) + + def transform_cancel_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel") diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 36876f67afc..2ec4f2da79b 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -4,6 +4,7 @@ Translates from Cohere's `/v1/rerank` input format to Vertex AI Discovery Engine Why separate file? Make it easy to see how transformation works """ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -74,14 +75,15 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: """ Validate and set up authentication for Vertex AI Discovery Engine API """ # Get credentials and project info from optional_params (which contains vertex_credentials, etc.) - litellm_params: Final = optional_params.copy() if optional_params else {} - vertex_credentials: Final = self.safe_get_vertex_ai_credentials(litellm_params) - vertex_project: Final = self.safe_get_vertex_ai_project(litellm_params) + vertex_params: Final = optional_params.copy() if optional_params else {} + vertex_credentials: Final = self.safe_get_vertex_ai_credentials(vertex_params) + vertex_project: Final = self.safe_get_vertex_ai_project(vertex_params) # Get access token using the base class method access_token, project_id = self._ensure_access_token( diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 16e72e3062d..25db3a673a1 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -7,11 +7,11 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer import base64 import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, TypedDict, cast import httpx -from httpx._types import RequestFiles +from httpx._types import FileContent, RequestFiles from typing_extensions import ReadOnly from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS @@ -677,9 +677,10 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, + video_file: FileContent | None = None, extra_body: dict[str, object] | None = None, prefetched_source_data: dict[str, Any] | None = None, - ) -> tuple[str, dict]: + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: """ Build a predictLongRunning edit request from the pre-fetched source video. @@ -727,7 +728,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): request_data["parameters"] = vertex_params edit_url: Final = f"{api_base.rstrip('/')}/{model}:predictLongRunning" - return edit_url, request_data + return edit_url, request_data, None def transform_video_edit_response( self, diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 497b2f62a97..ee330c92f1a 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -4,6 +4,7 @@ Transformation logic for Voyage AI's /v1/rerank endpoint. Docs - https://docs.voyageai.com/docs/reranker """ +from collections.abc import Mapping from typing import Any, Final import httpx @@ -137,6 +138,7 @@ class VoyageRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: if api_key is None: api_key = get_secret_str("VOYAGE_API_KEY") or get_secret_str("VOYAGE_AI_API_KEY") diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 32b96db2817..293880b188d 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -5,6 +5,7 @@ Docs - https://cloud.ibm.com/apidocs/watsonx-ai#text-rerank """ import uuid +from collections.abc import Mapping from typing import Any, Final, cast import httpx @@ -60,6 +61,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): model: str, api_key: str | None = None, optional_params: dict | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict: optional_params = optional_params or {} diff --git a/litellm/main.py b/litellm/main.py index 2cf53833c5a..98f92e50599 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -24,6 +24,7 @@ from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy from functools import partial +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args from litellm._logging import _redact_string @@ -416,7 +417,7 @@ async def acompletion( logprobs: bool | None = None, top_logprobs: int | None = None, deployment_id=None, - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, safety_identifier: str | None = None, service_tier: str | None = None, @@ -602,7 +603,7 @@ async def acompletion( _, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, - api_base=base_url, + api_base=kwargs.get("api_base") or base_url, ) fallbacks = fallbacks or litellm.model_fallbacks @@ -1811,6 +1812,56 @@ def _complete_fireworks_ai( return response +def _complete_together_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: + acompletion: Final = ctx.acompletion + api_base: Final = ctx.api_base + api_key: Final = ctx.api_key + client: Final = _dispatch_client_http(ctx) + custom_llm_provider: Final = ctx.custom_llm_provider + headers: Final = ctx.headers + litellm_params: Final = ctx.litellm_params + logging: Final = ctx.logging + messages: Final = ctx.messages + model: Final = ctx.model + model_response: Final = ctx.model_response + optional_params: Final = ctx.optional_params + provider_config: Final = ctx.provider_config + shared_session: Final = ctx.shared_session + stream: Final = ctx.stream + timeout: Final = ctx.timeout + + try: + response: Final = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=_get_encoding(), + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args=MappingProxyType({"headers": headers}), + ) + raise + + return response + + def _complete_heroku(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base @@ -4920,7 +4971,7 @@ def completion( logit_bias: dict | None = None, user: str | None = None, # openai v1.0+ new params - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, response_format: dict | type[BaseModel] | None = None, seed: int | None = None, @@ -5600,6 +5651,8 @@ def completion( elif custom_llm_provider == "fireworks_ai": ## COMPLETION CALL response = _complete_fireworks_ai(_dispatch_ctx) + elif custom_llm_provider == "together_ai": + response = _complete_together_ai(_dispatch_ctx) elif custom_llm_provider == "heroku": response = _complete_heroku(_dispatch_ctx) @@ -5649,7 +5702,6 @@ def completion( or custom_llm_provider == "volcengine" or custom_llm_provider == "anyscale" or custom_llm_provider == "openai" - or custom_llm_provider == "together_ai" or custom_llm_provider == "nebius" or custom_llm_provider == "wandb" or custom_llm_provider == "clarifai" @@ -5699,14 +5751,6 @@ def completion( response = _complete_openrouter(_dispatch_ctx) elif custom_llm_provider == "vercel_ai_gateway": response = _complete_vercel_ai_gateway(_dispatch_ctx) - elif ( - custom_llm_provider == "together_ai" - or ("togethercomputer" in model) - or (model in litellm.together_ai_models) - ): - """ - Deprecated. We now do together ai calls via the openai client - https://docs.together.ai/docs/openai-api-compatibility - """ elif custom_llm_provider == "palm": raise ValueError( "Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en" @@ -6828,6 +6872,8 @@ def embedding( aembedding=aembedding, ) elif custom_llm_provider == "azure_ai": + from litellm.llms.azure_ai.common_utils import get_azure_ai_entra_token + api_base = ( api_base # for deepinfra/perplexity/anyscale/groq/friendliai we check in get_llm_provider and pass in the api base from there or litellm.api_base @@ -6837,8 +6883,8 @@ def embedding( api_key = ( api_key or litellm.api_key # for deepinfra/perplexity/anyscale/friendliai we check in get_llm_provider and pass in the api key from there - or litellm.openai_key or get_secret_str("AZURE_AI_API_KEY") + or get_azure_ai_entra_token(litellm_params=litellm_params_dict) ) ## EMBEDDING CALL @@ -8566,7 +8612,7 @@ def stream_chunk_builder( if len(chunks) == 0: return None ## Route to the text completion logic - first_chunk_with_choices: Final = next((c for c in chunks if c["choices"]), None) + first_chunk_with_choices: Final = next((c for c in chunks if c.get("choices")), None) if first_chunk_with_choices is not None and isinstance( first_chunk_with_choices["choices"][0], litellm.utils.TextChoices ): # route to the text completion logic @@ -8581,7 +8627,7 @@ def stream_chunk_builder( simple_content_parts: Final[list[str]] = [] is_simple_text_stream = True for chunk in chunks: - if len(chunk["choices"]) == 0: + if not chunk.get("choices"): continue choice = chunk["choices"][0] @@ -8647,7 +8693,7 @@ def stream_chunk_builder( tool_call_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "tool_calls" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["tool_calls"] is not None ] @@ -8661,7 +8707,7 @@ def stream_chunk_builder( function_call_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "function_call" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["function_call"] is not None ] @@ -8674,7 +8720,7 @@ def stream_chunk_builder( content_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "content" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["content"] is not None ] @@ -8685,7 +8731,7 @@ def stream_chunk_builder( thinking_blocks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "thinking_blocks" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["thinking_blocks"] is not None ] @@ -8698,7 +8744,7 @@ def stream_chunk_builder( reasoning_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "reasoning_content" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["reasoning_content"] is not None ] @@ -8711,7 +8757,7 @@ def stream_chunk_builder( annotation_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "annotations" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["annotations"] is not None ] @@ -8728,7 +8774,7 @@ def stream_chunk_builder( audio_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "audio" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["audio"] is not None ] @@ -8742,7 +8788,7 @@ def stream_chunk_builder( image_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "images" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["images"] is not None ] @@ -8759,7 +8805,7 @@ def stream_chunk_builder( provider_specific_chunks: Final = [ chunk for chunk in chunks - if len(chunk["choices"]) > 0 + if chunk.get("choices") and "provider_specific_fields" in chunk["choices"][0]["delta"] and chunk["choices"][0]["delta"]["provider_specific_fields"] is not None ] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3af7d9e5019..9ca8d9e1bac 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1019,6 +1019,7 @@ }, "anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1053,6 +1054,7 @@ }, "global.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1087,6 +1089,7 @@ }, "us.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1121,6 +1124,7 @@ }, "eu.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1155,6 +1159,7 @@ }, "au.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -2233,6 +2238,7 @@ }, "anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -2266,6 +2272,7 @@ }, "global.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -2299,6 +2306,7 @@ }, "us.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2332,6 +2340,7 @@ }, "eu.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2365,6 +2374,7 @@ }, "au.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2398,6 +2408,7 @@ }, "jp.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2950,6 +2961,7 @@ "azure_ai/claude-opus-4-6": { "deprecation_date": "2027-02-02", "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, "litellm_provider": "azure_ai", @@ -3181,6 +3193,7 @@ "azure_ai/claude-sonnet-4-6": { "deprecation_date": "2027-02-10", "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12489,6 +12502,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -12698,6 +12712,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12735,6 +12750,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -14551,6 +14567,8 @@ ] }, "databricks/databricks-bge-large-en": { + "cache_creation_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.0003e-07, "input_dbu_cost_per_token": 1.429e-06, "litellm_provider": "databricks", @@ -14566,6 +14584,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-3-7-sonnet": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14581,10 +14601,41 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-fable-5": { + "cache_creation_input_token_cost": 1.250004e-05, + "cache_read_input_token_cost": 1.00002e-06, + "input_cost_per_token": 1.000006e-05, + "input_dbu_cost_per_token": 0.000142858, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.000002e-05, + "output_dbu_cost_per_token": 0.000714286, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": false, + "thinking_always_on": true + }, "databricks/databricks-claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14600,10 +14651,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4": { + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14619,10 +14673,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-1": { + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14638,10 +14695,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14657,11 +14717,14 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_output_config": true }, "databricks/databricks-claude-opus-4-6": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14677,10 +14740,94 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_legacy_thinking": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-opus-4-7": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 2048, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-5": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-claude-sonnet-4": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14696,10 +14843,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-1": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14715,10 +14865,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14734,10 +14887,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-6": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14753,10 +14909,41 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_legacy_thinking": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-sonnet-5": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, + "input_cost_per_token": 2.99999e-06, + "input_dbu_cost_per_token": 4.2857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + }, + "mode": "chat", + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.000214286, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-gemini-2-5-flash": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 3.0002e-08, "input_cost_per_token": 3.0001999999999996e-07, "input_dbu_cost_per_token": 4.285999999999999e-06, "litellm_provider": "databricks", @@ -14771,9 +14958,12 @@ "output_dbu_cost_per_token": 3.5714e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-2-5-pro": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14788,9 +14978,12 @@ "output_dbu_cost_per_token": 0.000142857, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-flash-lite": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.122e-08, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -14805,9 +14998,12 @@ "output_dbu_cost_per_token": 2.6786e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14822,9 +15018,12 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-flash": { + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.251e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -14839,9 +15038,12 @@ "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14856,9 +15058,12 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -14874,6 +15079,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-5": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14886,9 +15093,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14901,9 +15111,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-max": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14916,9 +15129,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-mini": { + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -14931,9 +15147,12 @@ "mode": "chat", "output_cost_per_token": 1.99997e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2": { + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14946,9 +15165,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2-codex": { + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14961,9 +15183,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-3-codex": { + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14976,9 +15201,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14991,9 +15219,12 @@ "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-mini": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15006,9 +15237,12 @@ "mode": "chat", "output_cost_per_token": 4.50002e-06, "output_dbu_cost_per_token": 6.4286e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-nano": { + "cache_creation_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost": 2.002e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15021,9 +15255,12 @@ "mode": "chat", "output_cost_per_token": 1.24999e-06, "output_dbu_cost_per_token": 1.7857e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-mini": { + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15036,9 +15273,12 @@ "mode": "chat", "output_cost_per_token": 1.9999700000000004e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-nano": { + "cache_creation_input_token_cost": 4.998e-08, + "cache_read_input_token_cost": 4.97e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", @@ -15051,9 +15291,12 @@ "mode": "chat", "output_cost_per_token": 3.9998000000000007e-07, "output_dbu_cost_per_token": 5.714000000000001e-06, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15069,6 +15312,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-oss-20b": { + "cache_creation_input_token_cost": 7e-08, + "cache_read_input_token_cost": 7e-08, "input_cost_per_token": 7e-08, "input_dbu_cost_per_token": 1e-06, "litellm_provider": "databricks", @@ -15084,6 +15329,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gte-large-en": { + "cache_creation_input_token_cost": 1.2999e-07, + "cache_read_input_token_cost": 1.2999e-07, "input_cost_per_token": 1.2999000000000001e-07, "input_dbu_cost_per_token": 1.857e-06, "litellm_provider": "databricks", @@ -15099,6 +15346,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-llama-2-70b-chat": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15115,6 +15364,8 @@ "supports_tool_choice": true }, "databricks/databricks-llama-4-maverick": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15131,6 +15382,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-405b-instruct": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 5.00003e-06, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -15147,6 +15400,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-8b-instruct": { + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15162,6 +15417,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-meta-llama-3-3-70b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15178,6 +15435,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-70b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15194,6 +15453,8 @@ "supports_tool_choice": true }, "databricks/databricks-mixtral-8x7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15210,6 +15471,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-30b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15226,6 +15489,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -16890,6 +17155,14 @@ "notes": "Web Search on Amazon Bedrock AgentCore, billed by AWS on the gateway" } }, + "bing_grounding/search": { + "input_cost_per_query": 0.035, + "litellm_provider": "bing_grounding", + "mode": "search", + "metadata": { + "notes": "Grounding with Bing Search (G1 SKU): $35 per 1,000 transactions. Tokens for the Foundry model deployment that runs the grounded search are billed separately on that deployment." + } + }, "tinyfish/search": { "input_cost_per_query": 0.0, "litellm_provider": "tinyfish", @@ -23180,6 +23453,7 @@ }, "github_copilot/claude-opus-4.6-fast": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "litellm_provider": "github_copilot", "max_input_tokens": 128000, "max_output_tokens": 16000, @@ -33563,6 +33837,7 @@ }, "openrouter/anthropic/claude-sonnet-4.6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, @@ -33607,6 +33882,7 @@ }, "openrouter/anthropic/claude-opus-4.6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, @@ -35681,6 +35957,7 @@ }, "perplexity/anthropic/claude-opus-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, @@ -37609,6 +37886,7 @@ "output_cost_per_token": 1e-07 }, "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37625,6 +37903,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 262000, @@ -37637,6 +37916,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37649,6 +37929,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 40000, @@ -37660,6 +37941,7 @@ "supports_tool_choice": false }, "together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37672,11 +37954,15 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 3e-06, "litellm_provider": "together_ai", "max_input_tokens": 128000, "max_output_tokens": 20480, "max_tokens": 20480, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 7e-06, "supports_function_calling": true, @@ -37685,6 +37971,7 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1-0528-tput": { + "deprecation_date": "2026-02-03", "input_cost_per_token": 5.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -37702,6 +37989,9 @@ "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.25e-06, "supports_function_calling": true, @@ -37710,9 +38000,13 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-V3.1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_tokens": 16384, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.7e-06, "source": "https://www.together.ai/models/deepseek-v3-1", @@ -37724,6 +38018,7 @@ "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37732,16 +38027,21 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": { - "input_cost_per_token": 8.8e-07, + "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 1.04e-06, + "source": "https://docs.together.ai/docs/serverless-models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 0, "litellm_provider": "together_ai", "mode": "chat", @@ -37752,6 +38052,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-03-31", "input_cost_per_token": 2.7e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37762,6 +38063,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37772,6 +38074,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 3.5e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -37782,6 +38085,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37792,6 +38096,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37802,6 +38107,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-7B-Instruct-v0.1": { + "deprecation_date": "2025-11-13", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37810,6 +38116,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { + "deprecation_date": "2026-04-02", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37817,6 +38124,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37829,6 +38137,9 @@ "together_ai/moonshotai/Kimi-K2-Instruct": { "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-instruct", @@ -37872,6 +38183,7 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.5-Air-FP8": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -37889,6 +38201,9 @@ "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2.2e-06, "source": "https://www.together.ai/models/glm-4-6", @@ -37898,11 +38213,15 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.7": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://www.together.ai/models/glm-4-7", @@ -37912,11 +38231,15 @@ "supports_tool_choice": true }, "together_ai/moonshotai/Kimi-K2.5": { + "deprecation_date": "2026-05-21", "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 2.8e-06, "source": "https://www.together.ai/models/kimi-k2-5", @@ -37926,9 +38249,13 @@ "supports_reasoning": true }, "together_ai/moonshotai/Kimi-K2-Instruct-0905": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-0905", @@ -37937,9 +38264,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.7-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct", @@ -37949,9 +38280,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.6-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking", @@ -37961,6 +38296,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3.5-397B-A17B": { + "deprecation_date": "2026-06-29", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -37972,6 +38308,292 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/MiniMaxAI/MiniMax-M3": { + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Prism-ML/Ternary-Bonsai-27B": { + "input_cost_per_token": 0.0, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.5-9B": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_reasoning": true + }, + "together_ai/Qwen/Qwen3.7-Max": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.7-Plus": { + "input_cost_per_token": 3.2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.28e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.8-2.4T-A95B": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1010000, + "max_output_tokens": 1010000, + "max_tokens": 1010000, + "mode": "chat", + "output_cost_per_token": 6.25e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/arize-ai/qwen-2-1.5b-instruct": { + "input_cost_per_token": 1e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro": { + "input_cost_per_token": 1.74e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 512000, + "max_output_tokens": 512000, + "max_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 3.48e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/google/gemma-3n-E4B-it": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/google/gemma-4-31B-it": { + "input_cost_per_token": 3.9e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 9.7e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/intfloat/multilingual-e5-large-instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 514, + "max_tokens": 514, + "mode": "embedding", + "output_cost_per_token": 2e-08, + "output_vector_size": 1024, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-llama/Llama-Guard-4-12B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-models/Muse-Glimmer-30B": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/moonshotai/Kimi-K2.7-Code": { + "input_cost_per_token": 9.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/moonshotai/Kimi-K3": { + "input_cost_per_token": 3e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/nvidia/nemotron-3-ultra-550b-a55b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 512288, + "max_output_tokens": 512288, + "max_tokens": 512288, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/pearl-ai/gemma-4-31b-it": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8.6e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/thinkingmachines/Inkling": { + "input_cost_per_token": 1e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/thinkingmachines/Inkling-Small": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/zai-org/GLM-5.2": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048575, + "max_output_tokens": 1048575, + "max_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", @@ -39056,6 +39678,7 @@ }, "vercel_ai_gateway/anthropic/claude-opus-4.6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, @@ -40315,6 +40938,7 @@ "deprecation_date": "2027-02-05", "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -40347,6 +40971,7 @@ "deprecation_date": "2027-02-05", "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -40712,6 +41337,7 @@ "vertex_ai/claude-sonnet-4-6": { "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -43795,10 +44421,10 @@ "comment": "5 credits per second @ $0.01 per credit = $0.05 per second" } }, - "runwayml/gen4_aleph": { + "runwayml/gen4.5": { "litellm_provider": "runwayml", "mode": "video_generation", - "output_cost_per_video_per_second": 0.15, + "output_cost_per_second": 0.12, "source": "https://docs.dev.runwayml.com/guides/pricing/", "supported_modalities": [ "text", @@ -43808,13 +44434,136 @@ "video" ], "metadata": { - "comment": "15 credits per second @ $0.01 per credit = $0.15 per second" + "comment": "12 credits per second @ $0.01 per credit = $0.12 per second" } }, - "runwayml/gen3a_turbo": { + "runwayml/aleph2": { "litellm_provider": "runwayml", "mode": "video_generation", - "output_cost_per_video_per_second": 0.05, + "output_cost_per_second": 0.28, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "28 credits per second @ $0.01 per credit = $0.28 per second; 56 credit minimum per task not modeled" + } + }, + "runwayml/seedance2": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.36, + "output_cost_per_second_1080p": 0.4, + "output_cost_per_second_4k": 1.5, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "36 credits per second at 480p/720p, 40 at 1080p, 150 at 4K @ $0.01 per credit" + } + }, + "runwayml/seedance2_fast": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.29, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "29 credits per second at 480p/720p @ $0.01 per credit = $0.29 per second" + } + }, + "runwayml/seedance2_mini": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.16, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "16 credits per second @ $0.01 per credit = $0.16 per second; 64 credit minimum per task not modeled" + } + }, + "runwayml/seedance2_5": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.3, + "output_cost_per_second_480p": 0.2, + "output_cost_per_second_1080p": 0.68, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "Output: 20/30/68 credits per second at 480p/720p/1080p @ $0.01 per credit; input video billed additionally at 10/15/34 credits per input second and the 80 credit minimum per task are not modeled" + } + }, + "runwayml/hailuo3": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.15, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "10 credits per second at 768P, 15 at 2K (mapped to the 1080p tier) @ $0.01 per credit; 2 credits per reference image not modeled" + } + }, + "runwayml/gemini_omni_flash": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.1, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "10 credits per second @ $0.01 per credit = $0.10 per second" + } + }, + "runwayml/veo3.1": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.4, "source": "https://docs.dev.runwayml.com/guides/pricing/", "supported_modalities": [ "text", @@ -43824,7 +44573,23 @@ "video" ], "metadata": { - "comment": "5 credits per second @ $0.01 per credit = $0.05 per second" + "comment": "40 credits per second with audio, 20 without @ $0.01 per credit; priced at the with-audio rate" + } + }, + "runwayml/veo3.1_fast": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.15, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "15 credits per second with audio, 10 without @ $0.01 per credit; priced at the with-audio rate" } }, "runwayml/gen4_image": { @@ -48481,6 +49246,7 @@ "vertex_ai/claude-sonnet-4-6@default": { "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -48595,12 +49361,13 @@ "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -48627,12 +49394,13 @@ "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -48659,12 +49427,13 @@ "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49265,6 +50034,7 @@ }, "snowflake/claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "max_tokens": 16384, "max_input_tokens": 200000, "max_output_tokens": 16384, @@ -50255,6 +51025,14 @@ "supports_adaptive_thinking": true } }, + { + "name": "claude-legacy-thinking", + "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", + "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", + "model_info": { + "supports_legacy_thinking": true + } + }, { "name": "claude-always-on-thinking", "pattern": "claude-(?:fable|mythos)-", diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 8a2ee2a3af8..4b30afb2f98 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -199,7 +199,7 @@ def llm_passthrough_route( api_key=api_key, ) - litellm_params_dict: Final = get_litellm_params(**kwargs) + litellm_params_dict: Final = get_litellm_params(api_key=api_key, api_base=api_base, **kwargs) if client is None: from litellm.llms.custom_httpx.http_handler import ( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0840d37ffa1..628e569e1b8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -815,6 +815,7 @@ class LiteLLMRoutes(enum.Enum): "/team/member_add", "/team/member_delete", "/team/member_update", + "/team/{team_id}/member/{user_id}/reset_spend", "/team/permissions_list", "/team/permissions_update", "/team/daily/activity", @@ -1287,6 +1288,16 @@ class RegenerateKeyRequest(GenerateKeyRequest): class ResetSpendRequest(LiteLLMPydanticObjectBase): reset_to: float + @field_validator("reset_to", mode="before") + @classmethod + def reject_bool_reset_to(cls, v): + # bool is a subclass of int, so pydantic silently coerces True/False into + # 1.0/0.0 for a `float` field: a caller who accidentally sends a boolean + # would otherwise get an unintended spend reset instead of a 422. + if isinstance(v, bool): + raise ValueError("reset_to must be a number, not a boolean") # noqa: TRY004 # pydantic needs ValueError + return v + class KeyRequest(LiteLLMPydanticObjectBase): keys: list[str] | None = None diff --git a/litellm/proxy/analytics_endpoints/cache_activity.py b/litellm/proxy/analytics_endpoints/cache_activity.py index 751d6c15dfa..b87b8eac3ef 100644 --- a/litellm/proxy/analytics_endpoints/cache_activity.py +++ b/litellm/proxy/analytics_endpoints/cache_activity.py @@ -34,10 +34,18 @@ class CacheActivityFilterOptions(BaseModel): models: list[str] +class CacheActivityErrorBucket(BaseModel): + call_type: str + error_code: str + error_class: str + count: int + + class CacheActivityResponse(BaseModel): groups: list[CacheActivityGroup] totals: CacheActivityTotals filter_options: CacheActivityFilterOptions + error_breakdown: tuple[CacheActivityErrorBucket, ...] GROUPS_SQL: Final = """ @@ -65,6 +73,26 @@ GROUPS_SQL: Final = """ ORDER BY (COUNT(*)) DESC """ +ERROR_BREAKDOWN_SQL: Final = """ + SELECT + CASE WHEN sl."call_type" = '' THEN 'Unknown' ELSE sl."call_type" END AS call_type, + COALESCE(NULLIF(sl."metadata"->'error_information'->>'error_code', ''), 'Unknown') AS error_code, + COALESCE(NULLIF(sl."metadata"->'error_information'->>'error_class', ''), 'Unknown') AS error_class, + COUNT(*)::int AS count + FROM "LiteLLM_SpendLogs" sl + LEFT JOIN "LiteLLM_VerificationToken" vt ON sl."api_key" = vt."token" + WHERE + sl."status" = 'failure' + AND sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') + AND ($3::jsonb = '[]'::jsonb + OR COALESCE(vt."key_alias", 'Unnamed Key') IN (SELECT jsonb_array_elements_text($3::jsonb))) + AND ($4::jsonb = '[]'::jsonb + OR sl."model" IN (SELECT jsonb_array_elements_text($4::jsonb))) + GROUP BY 1, 2, 3 + ORDER BY (COUNT(*)) DESC +""" + KEY_ALIAS_OPTIONS_SQL: Final = """ SELECT DISTINCT COALESCE(vt."key_alias", 'Unnamed Key') AS key_alias FROM "LiteLLM_SpendLogs" sl @@ -95,6 +123,7 @@ class _ModelRow(BaseModel): _groups_adapter: Final = TypeAdapter(list[CacheActivityGroup]) +_error_buckets_adapter: Final = TypeAdapter(tuple[CacheActivityErrorBucket, ...]) _key_alias_rows_adapter: Final = TypeAdapter(list[_KeyAliasRow]) _model_rows_adapter: Final = TypeAdapter(list[_ModelRow]) @@ -120,10 +149,11 @@ async def get_cache_activity( key_aliases: Sequence[str], models: Sequence[str], ) -> CacheActivityResponse: - group_rows, key_alias_rows, model_rows = await asyncio.gather( - prisma_client.db.query_raw( - GROUPS_SQL, start_date, end_date, json.dumps(list(key_aliases)), json.dumps(list(models)) - ), + key_aliases_json: Final = json.dumps(list(key_aliases)) + models_json: Final = json.dumps(list(models)) + group_rows, error_rows, key_alias_rows, model_rows = await asyncio.gather( + prisma_client.db.query_raw(GROUPS_SQL, start_date, end_date, key_aliases_json, models_json), + prisma_client.db.query_raw(ERROR_BREAKDOWN_SQL, start_date, end_date, key_aliases_json, models_json), prisma_client.db.query_raw(KEY_ALIAS_OPTIONS_SQL, start_date, end_date), prisma_client.db.query_raw(MODEL_OPTIONS_SQL, start_date, end_date), ) @@ -135,4 +165,5 @@ async def get_cache_activity( key_aliases=[row.key_alias for row in _key_alias_rows_adapter.validate_python(key_alias_rows or [])], models=[row.model for row in _model_rows_adapter.validate_python(model_rows or [])], ), + error_breakdown=_error_buckets_adapter.validate_python(error_rows or []), ) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index 46ee9b0911d..09ba2c93ea0 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -6,21 +6,21 @@ Plugins are stored as metadata + git source references in LiteLLM database. Actual plugin files are hosted on GitHub/GitLab/Bitbucket. Endpoints: -/claude-code/marketplace.json - GET - List plugins for Claude Code discovery -/claude-code/plugins - POST - Register a new plugin (create-only) -/claude-code/plugins - GET - List plugins (admin) -/claude-code/plugins/{name} - GET - Get plugin details -/claude-code/plugins/{name} - PUT - Update an existing plugin -/claude-code/plugins/{name}/enable - POST - Enable a plugin -/claude-code/plugins/{name}/disable - POST - Disable a plugin -/claude-code/plugins/{name} - DELETE - Delete a plugin +/claude-code/marketplace.json - GET - List plugins for Claude Code discovery (unauthenticated) +/claude-code/plugins - POST - Register a new plugin (create-only, proxy admin only) +/claude-code/plugins - GET - List plugins (any authenticated key) +/claude-code/plugins/{name} - GET - Get plugin details (any authenticated key) +/claude-code/plugins/{name} - PUT - Update an existing plugin (proxy admin only) +/claude-code/plugins/{name}/enable - POST - Enable a plugin (proxy admin only) +/claude-code/plugins/{name}/disable - POST - Disable a plugin (proxy admin only) +/claude-code/plugins/{name} - DELETE - Delete a plugin (proxy admin only) """ import json import re from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Final, Protocol, TypedDict +from typing import Annotated, Final, Protocol, TypedDict from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse @@ -28,6 +28,7 @@ from fastapi.responses import JSONResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.repositories.table_repositories import ClaudeCodePluginRepository from litellm.types.proxy.claude_code_endpoints import ( ListPluginsResponse, @@ -221,6 +222,18 @@ def _name_conflict_error(name: str) -> HTTPException: ) +def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: + """Catalog mutations are restricted to proxy admins: marketplace.json is served + unauthenticated and any registered/updated entry is immediately installable by + every user, so a non-admin key must never be able to add or overwrite one. + """ + if not is_proxy_admin(user_api_key_dict): + raise HTTPException( + status_code=403, + detail={"error": "Only proxy admins may modify the Claude Code plugin marketplace."}, + ) + + @router.post( "/claude-code/plugins", tags=["Claude Code Marketplace"], @@ -242,6 +255,8 @@ async def register_plugin( the same name already exists it returns 409 Conflict; use PUT /claude-code/plugins/{plugin_name} to update an existing plugin. + Requires a proxy admin API key. + Parameters: - name: Plugin name (kebab-case) - source: Git source reference (github, url, or git-subdir format) @@ -271,6 +286,8 @@ async def register_plugin( from prisma.errors import UniqueViolationError try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() if not re.match(r"^[a-z0-9-]+$", request.name): @@ -468,6 +485,7 @@ async def get_plugin( async def update_plugin( plugin_name: str, request: UpdatePluginRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ Update an existing plugin in the LiteLLM marketplace. @@ -481,6 +499,8 @@ async def update_plugin( Returns 404 if no plugin with the given name exists; use POST /claude-code/plugins to create a new plugin. + Requires a proxy admin API key. + Parameters: - plugin_name: Name of the plugin to update (path parameter) - source: Git source reference (github, url, or git-subdir format) @@ -509,6 +529,8 @@ async def update_plugin( from prisma.errors import PrismaError try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() _validate_plugin_source(request.source) @@ -566,10 +588,14 @@ async def enable_plugin( """ Enable a disabled plugin. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to enable """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( @@ -611,10 +637,14 @@ async def disable_plugin( """ Disable a plugin without deleting it. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to disable """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( @@ -656,10 +686,14 @@ async def delete_plugin( """ Delete a plugin from the marketplace. + Requires a proxy admin API key. + Parameters: - plugin_name: The name of the plugin to delete """ try: + _require_proxy_admin(user_api_key_dict) + prisma_client: Final = await _get_prisma_client() plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e7b98b3cc7f..52857909621 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -71,7 +71,6 @@ from litellm.proxy.auth.budget_throttle import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, @@ -87,6 +86,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, tag_cache_key, tag_registry_cache_key, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -1129,7 +1130,8 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: Parameters: - user_route: str - the route the user is trying to call - - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. + - allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name + (e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*"). """ from starlette.routing import compile_path @@ -1139,7 +1141,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool: regex, _, _ = compile_path(template) if regex.match(user_route): return True - elif allowed_route == user_route: + elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route): return True return False @@ -1967,7 +1969,7 @@ async def get_team_membership( if user_id is None or team_id is None: return None - _key: Final = f"team_membership:{user_id}:{team_id}" + _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) # check if in cache cached_membership_obj: Final = await user_api_key_cache.async_get_cache( @@ -2402,6 +2404,116 @@ async def _cache_team_object( ) +async def invalidate_team_member_spend_state( + user_id: str, + team_id: str, + user_api_key_cache: UserApiKeyCache, + new_spend: float | None = None, +) -> None: + """ + Clear every cached read path for one team member's budget so a spend + reset or a raised cap takes effect on the next request instead of + waiting on the membership cache's TTL. + + Two independently-keyed cache entries hold the same LiteLLM_TeamMembership + row: user_api_key_auth.py's admission check writes ``{team_id}_{user_id}``, + while budget_reservation.py's pre-call reservation and auth_checks.py's own + get_team_membership() (used by _check_team_member_budget) both write + ``team_membership:{user_id}:{team_id}``. Both formats must be invalidated + explicitly; writing one does not refresh the other. All keys are also + broadcast (LIT-3803): each worker's own in-memory copy (membership object, + spend counter, or the counter's own short-TTL DB-floor marker) survives + eviction elsewhere until its TTL, so the handling worker alone clearing its + copy leaves every other worker still enforcing the pre-reset budget. + + ``new_spend`` is only passed by reset_team_member_spend_fn, which knows the + exact post-reset value: it is SET everywhere (matching /key/{key}/reset_spend's + own precedent) rather than deleted, so a worker's next read reflects it + directly instead of re-deriving it through a DB reseed. team_member_update + only changes the budget cap, not the tracked spend, so it passes no + new_spend; the live spend counter is untouched in that case (deleting it + would force a reseed from the DB's own spend column, which lags the live + counter via periodic batch writes, briefly under-enforcing the raised cap + against a spend value lower than what was actually tracked) and only the + membership caches carrying the new cap are invalidated. + + The floor marker (``spend_db_floor:``, proxy_server.py's + _authoritative_floor_spend) caches the pre-reset DB spend for + SPEND_DB_FLOOR_CACHE_TTL_SECONDS; left stale after a real reset, a request + landing on the pod that cached it can read that higher floor and raise the + counter right back above the just-reset spend. It is overwritten here with + the post-reset floor (not merely deleted) and _authoritative_floor_spend + re-checks the marker after its DB read, so a floor read already in flight + on this pod when the reset commits cannot clobber it with the pre-reset + value. Both keys are broadcast as SETs carrying new_spend, not deletes: + every subscriber (remote pods AND this pod's own, which receives its own + message) writes the post-reset value, so the self-delivered message cannot + erase the guard just written here. + + Raises HTTPException(503) if Redis still holds the stale pre-reset counter + after both the SET and the fallback DELETE fail: budget checks read Redis + first, so returning success would leave the old value authoritative for + every worker despite the DB write having committed. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + publish_auth_cache_invalidation, + ) + + if new_spend is not None: + from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache + + spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}" + spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}" + + spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60) + except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up + verbose_proxy_logger.warning( + "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next " + "read reseeds from the DB rather than keeping the stale pre-reset value authoritative", + spend_counter_key, + e, + ) + try: + await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key) + except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success + verbose_proxy_logger.warning( + "Failed to delete stale spend counter %s in Redis after a failed reset write", + spend_counter_key, + exc_info=True, + ) + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={ # mutable-ok: HTTPException.detail takes a dict + "error": "Spend was reset in the database, but Redis is unreachable and still " + "holds the pre-reset counter. Retry once Redis is reachable." + }, + ) from e + + spend_counter_cache.in_memory_cache.set_cache( + key=spend_db_floor_key, + value=new_spend, + ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ) + await publish_auth_cache_invalidation(cache_key=spend_counter_key, new_value=new_spend, ttl=60) + await publish_auth_cache_invalidation( + cache_key=spend_db_floor_key, + new_value=new_spend, + ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ) + + await evict_and_broadcast( + cache_keys=( + team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + ), + user_api_key_cache=user_api_key_cache, + ) + + async def delete_cache_team_object( team_id: str, team_alias: str | None, @@ -2629,20 +2741,9 @@ async def _get_team_object_from_user_api_key_cache( async def _get_team_object_from_cache( key: str, - proxy_logging_obj: ProxyLogging | None, user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, ) -> LiteLLM_TeamTableCachedObj | None: - ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ## - if proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache: - cached_raw: Final = await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( - key=key, parent_otel_span=parent_otel_span - ) - if cached_raw is not None: - from_internal: Final = CacheCodec.deserialize(cached_raw, LiteLLM_TeamTableCachedObj) - if from_internal is not None: - return from_internal - decoded: Final = await user_api_key_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span, @@ -2678,7 +2779,6 @@ async def get_team_object( if not check_db_only: cached_team_obj: Final = await _get_team_object_from_cache( key=key, - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) @@ -2841,7 +2941,6 @@ async def get_team_object_by_alias( cached_team_obj: Final = await _get_team_object_from_cache( key=cache_key, - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 233679126f8..a42187b3a44 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import EMPTY_MAPPING from litellm.integrations.otel.runtime import seed_request_identity +from litellm.litellm_core_utils.core_helpers import is_expected_client_error from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, @@ -109,7 +110,12 @@ class UserAPIKeyAuthExceptionHandler: request=request, use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True, ) - verbose_proxy_logger.exception( + log_fn: Final = ( + verbose_proxy_logger.error + if is_expected_client_error(e) and not litellm.log_client_error_tracebacks + else verbose_proxy_logger.exception + ) + log_fn( "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", e, requester_ip, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index d04a71535ef..9b1a6ba5aa7 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -608,7 +608,7 @@ def route_in_additonal_public_routes(current_route: str): # Check wildcard patterns for route_pattern in routes_defined: - if RouteChecks._route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): + if RouteChecks.route_matches_wildcard_pattern(route=current_route, pattern=route_pattern): return True return False diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index cea21ca088b..4dba2497bb9 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -181,7 +181,7 @@ class RouteChecks: # check if wildcard pattern is allowed for allowed_route in valid_token.allowed_routes: - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): return True if denied_auth_enforced_pass_through_route: @@ -329,7 +329,7 @@ class RouteChecks: route_allowed = True break - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route): route_allowed = True break @@ -397,7 +397,7 @@ class RouteChecks: return True # Check for wildcard patterns like "/containers/*" if RouteChecks._is_wildcard_pattern(pattern=openai_route): - if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=openai_route): + if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route): return True # Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent" @@ -517,7 +517,7 @@ class RouteChecks: return pattern.endswith("*") @staticmethod - def _route_matches_wildcard_pattern(route: str, pattern: str) -> bool: + def route_matches_wildcard_pattern(route: str, pattern: str) -> bool: """ Check if route matches the wildcard pattern @@ -594,7 +594,7 @@ class RouteChecks: # e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/* ######################################################### if any( - RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route) + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) for allowed_route in allowed_routes if RouteChecks._is_wildcard_pattern(pattern=allowed_route) ): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 658d176f6a7..28d76e6799c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -87,7 +87,10 @@ from litellm.proxy.common_utils.http_parsing_utils import ( populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_auth_cache_key, +) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import ( @@ -1970,8 +1973,10 @@ async def _user_api_key_auth_builder( # Check 3. Check if user is in their team budget if not skip_budget_checks and valid_token.team_member_spend is not None: - if prisma_client is not None: - _cache_key: Final = f"{valid_token.team_id}_{valid_token.user_id}" + _user_id: Final = valid_token.user_id + _team_id: Final = valid_token.team_id + if prisma_client is not None and _user_id is not None and _team_id is not None: + _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id) team_member_info = await user_api_key_cache.async_get_cache( key=_cache_key, @@ -1979,25 +1984,21 @@ async def _user_api_key_auth_builder( ) if team_member_info is None: # read from DB - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - - if _user_id is not None and _team_id is not None: - _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, - include={"litellm_budget_table": True}, + _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( + where={ + "user_id": _user_id, + "team_id": _team_id, + }, + include={"litellm_budget_table": True}, + ) + if _db_member is not None: + team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=team_member_info, + model_type=LiteLLM_TeamMembership, + ttl=5, ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) if team_member_info is not None and team_member_info.litellm_budget_table is not None: team_member_budget: Final = team_member_info.litellm_budget_table.max_budget @@ -2013,11 +2014,16 @@ async def _user_api_key_auth_builder( max_budget=team_member_budget, ) if team_member_spend > team_member_budget: + _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, + message=( + f"Budget has been exceeded! TeamMember={_entity_id} " + f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" + ), entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{valid_token.team_id}", + entity_id=_entity_id, ) # Check 3. If token is expired diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index dbbf9cb673e..6fd743deb44 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -33,7 +33,7 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -1138,24 +1138,25 @@ async def open_sse_before_first_byte( ) -def _is_azure_model_router_request(model: str) -> bool: +def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool: """ - Check if the requested model is an Azure Model Router. + Check if a request went down the Azure Model Router route. - Azure Model Router models follow the pattern: - - azure_ai/model_router/ - - azure_ai/model-router - - model_router/ - - model-router + ``model`` here is what the *client* sent, a model group alias with no ``model_router/`` + prefix, so matching on it alone only works when the operator happened to put "model-router" + in the alias. Where the response is in hand its stamp answers this outright, so callers + should pass ``hidden_params``. Args: model: The requested model name + hidden_params: ``_hidden_params`` from the response, when the caller has it Returns: bool: True if this is an Azure Model Router request """ - model_lower: Final = model.lower() - return "model-router" in model_lower or "model_router" in model_lower + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + + return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params) def _override_openai_response_model( @@ -1223,7 +1224,7 @@ def _override_openai_response_model( return # Check if this is an Azure Model Router request - if so, preserve the actual model used - if _is_azure_model_router_request(requested_model): + if _is_azure_model_router_request(requested_model, hidden_params): verbose_proxy_logger.debug( "%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.", log_context, @@ -1379,7 +1380,12 @@ def _log_llm_api_exception(e: Exception) -> None: "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled" ) return - verbose_proxy_logger.exception("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e) + log_fn: Final = ( + verbose_proxy_logger.error + if is_expected_client_error(e) and not litellm.log_client_error_tracebacks + else verbose_proxy_logger.exception + ) + log_fn("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e) async def _cancel_llm_call_on_client_disconnect( diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index acdc9728390..fb2ca6372c0 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -12,6 +12,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( ) if TYPE_CHECKING: + from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -30,15 +31,24 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @dataclass(frozen=True, slots=True) class _CacheInvalidationMessage: cache_key: str + new_value: float | None = None + ttl: float | None = None -def _cache_invalidation_message_json(cache_key: str) -> str: - return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key))) +def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: + message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) + return json.dumps({field: value for field, value in message.items() if value is not None}) -def _cache_key_from_message_data(data: object) -> str | None: +def _finite_number_or_none(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _message_from_data(data: object) -> _CacheInvalidationMessage | None: if isinstance(data, bytes): - data = data.decode("utf-8", errors="replace") + data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str if not isinstance(data, str): return None try: @@ -48,14 +58,28 @@ def _cache_key_from_message_data(data: object) -> str | None: if not isinstance(parsed, dict): return None cache_key: Final = parsed.get("cache_key") - return cache_key if isinstance(cache_key, str) else None + if not isinstance(cache_key, str): + return None + return _CacheInvalidationMessage( + cache_key=cache_key, + new_value=_finite_number_or_none(parsed.get("new_value")), + ttl=_finite_number_or_none(parsed.get("ttl")), + ) -async def publish_auth_cache_invalidation(cache_key: str) -> None: +async def publish_auth_cache_invalidation( + cache_key: str, new_value: float | None = None, ttl: float | None = None +) -> None: """ Best-effort broadcast so every worker drops its local in-memory copy of a mutated management object; without this, only the handling worker and Redis are evicted and other workers keep serving the stale object until its TTL. + + Passing ``new_value`` broadcasts a SET instead of a delete: every subscriber + (including the publishing worker's own, which receives its own message) + writes the value into its additional in-memory caches rather than deleting + the key. A spend reset uses this so the handler's self-delivered message + cannot erase the freshly-written post-reset counter or floor marker. """ redis_cache: Final = coordination_redis_cache() if redis_cache is None: @@ -68,7 +92,10 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None: cache_key, ) return - await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key)) + await client.publish( + auth_cache_invalidation_channel(redis_cache), + _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl), + ) except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) @@ -95,15 +122,17 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") + __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache") def __init__( self, redis_cache: "RedisCache", user_api_key_cache: "UserApiKeyCache", + additional_in_memory_caches: Sequence["InMemoryCache"] = (), ) -> None: self._redis_cache = redis_cache self._user_api_key_cache = user_api_key_cache + self._additional_in_memory_caches = tuple(additional_in_memory_caches) self._task: asyncio.Task[None] | None = None def start(self) -> None: @@ -160,12 +189,18 @@ class AuthCacheInvalidationSubscriber: def _apply_message(self, message: object) -> None: data: Final = message.get("data") if isinstance(message, dict) else None - cache_key: Final = _cache_key_from_message_data(data) - if cache_key is None: + parsed: Final = _message_from_data(data) + if parsed is None: + return + if parsed.new_value is not None: + for additional_cache in self._additional_in_memory_caches: + additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl) return in_memory_cache: Final = self._user_api_key_cache.in_memory_cache if in_memory_cache is not None: - in_memory_cache.delete_cache(cache_key) + in_memory_cache.delete_cache(parsed.cache_key) + for additional_cache in self._additional_in_memory_caches: + additional_cache.delete_cache(parsed.cache_key) @staticmethod async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 1e4344a71f4..96621b08ba1 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -274,10 +274,8 @@ async def get_form_data(request: Request) -> dict[str, Any]: Handles when OpenAI SDKs pass form keys as `timestamp_granularities[]="word"` instead of `timestamp_granularities=["word", "sentence"]` """ form: Final = await request.form() - form_data: Final = dict(form) parsed_form_data: Final[dict[str, Any]] = {} - for key, value in form_data.items(): - # OpenAI SDKs pass form keys as `timestamp_granularities[]="word"` instead of `timestamp_granularities=["word", "sentence"]` + for key, value in form.multi_items(): # not dict(form), which keeps only the last repeat if key.endswith("[]"): clean_key = key[:-2] parsed_form_data.setdefault(clean_key, []).append(value) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 93d51bdd461..b8df0105b7b 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -200,6 +200,21 @@ def end_user_restricted_registry_cache_key() -> str: return "end_user_restricted_registry" +def team_membership_auth_cache_key(team_id: str, user_id: str) -> str: + """Cache key one team member's ``LiteLLM_TeamMembership`` row is stored under for the admission check.""" + return f"{team_id}_{user_id}" + + +def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: + """Cache key the pre-call budget reservation stores the same ``LiteLLM_TeamMembership`` row under. + + Deliberately not unified with ``team_membership_auth_cache_key``: the two readers wrote independent + keys before this file existed, so a fix that invalidates one must invalidate both explicitly rather + than assume a single write is visible to both. + """ + return f"team_membership:{user_id}:{team_id}" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml b/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml new file mode 100644 index 00000000000..ac8a3db2d21 --- /dev/null +++ b/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml @@ -0,0 +1,30 @@ +# Web search via Microsoft Foundry (Grounding with Bing Search / the built-in +# web_search tool), called through the Foundry Responses API. +# +# Configure the provider with env vars (setup and pricing are in the LiteLLM docs; +# the code lives in litellm/llms/azure/search/transformation.py): +# BING_GROUNDING_PROJECT_ENDPOINT (required) the Foundry project endpoint +# BING_GROUNDING_MODEL (required) a model deployment in that project +# BING_GROUNDING_CONNECTION_ID (optional) a Grounding with Bing connection id; +# without it the built-in web_search tool is used +# BING_GROUNDING_TOKEN (optional) an Entra bearer token; without it (and +# without api_key) azure-identity mints one + +model_list: + - model_name: claude-sonnet + litellm_params: + model: bedrock/us.anthropic.claude-sonnet-5 + aws_region_name: us-east-1 + +search_tools: + - search_tool_name: bing-grounding-search + litellm_params: + search_provider: bing_grounding + # Optional: an Azure API key instead of BING_GROUNDING_TOKEN / azure-identity + # api_key: os.environ/AZURE_AI_API_KEY + +litellm_settings: + callbacks: ["websearch_interception"] + websearch_interception_params: + enabled_providers: ["bedrock"] + search_tool_name: bing-grounding-search diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index f9d408fb7de..d8fde8ca5dc 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -18,6 +18,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS, ) +from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model ILLEGAL_DISPLAY_PARAMS: Final = [ "messages", @@ -182,30 +183,17 @@ async def run_with_timeout(task, timeout): return {"error": "Timeout exceeded", "exception": timeout_exception} -def _is_semantic_auto_router_deployment(litellm_params: dict) -> bool: - """ - True for semantic auto_router deployments (auto_router/) that are not - sub-strategies (complexity_router, adaptive_router, quality_router). - - These are meta-routers that select among real LLM deployments at request time; - they have no LLM endpoint to health-check. - """ +def _is_strategy_router_deployment(litellm_params: dict) -> bool: + """True for strategy-router deployments.""" model: Final[object] = litellm_params.get("model", "") - if not isinstance(model, str): - return False - if not model.startswith("auto_router/"): - return False - for sub_strategy in ("complexity_router", "adaptive_router", "quality_router"): - if model.startswith(f"auto_router/{sub_strategy}"): - return False - return True + return isinstance(model, str) and classify_strategy_router_model(model) is not None async def _run_model_health_check(model: dict): litellm_params = model["litellm_params"] model_info: Final = model.get("model_info", {}) - if _is_semantic_auto_router_deployment(litellm_params): + if _is_strategy_router_deployment(litellm_params): return {} mode: Final = _resolve_health_check_mode( @@ -445,6 +433,9 @@ def _update_litellm_params_for_health_check(model_info: dict, litellm_params: di """ Update the litellm params for health check. + - merges `model_info.health_check_params` into the probe request, so a deployment whose provider + requires a payload field litellm does not synthesize (e.g. `mediaSource` for Bedrock TwelveLabs + Pegasus) can supply it. The dedicated knobs below are applied afterwards and win on conflict. - gets a short `messages` param for health check - adds a bounded `max_tokens` when the deployment is a chat-style mode (`chat`, `completion`, `responses`) or the operator explicitly opts in @@ -459,6 +450,16 @@ def _update_litellm_params_for_health_check(model_info: dict, litellm_params: di model_info, litellm_params, # any-ok: untyped router config dict ) + _health_check_params: Final = model_info.get("health_check_params", None) + if isinstance(_health_check_params, dict): + litellm_params.update(_health_check_params) + elif _health_check_params is not None: + logger.warning( + "health_check_params for model %s is a %s, expected a dict. Ignoring it.", + litellm_params.get("model"), + type(_health_check_params).__name__, + ) + litellm_params["messages"] = _get_random_llm_message() if _should_inject_health_check_max_tokens( model_info, diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 33894777bc3..c23cde6052e 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1888,6 +1888,8 @@ async def test_model_connection( # already resolved before reaching this endpoint; any remaining # reference must have come from the request body. _reject_os_environ_references(request_litellm_params) + if model_info: + _reject_os_environ_references(model_info) model_name: Final = request_litellm_params.get("model") # Look up model configuration from router if model name is provided @@ -1950,23 +1952,23 @@ async def test_model_connection( **request_litellm_params, } - ## Auth check - auth_model_info: Final = loaded_model_info if loaded_model_info is not None else model_info + resolved_model_info: Final = loaded_model_info if loaded_model_info is not None else model_info + litellm_params = _update_litellm_params_for_health_check( + model_info=resolved_model_info or {}, + litellm_params=litellm_params, + ) + + ## Auth check, on the final probe params so health_check_params cannot retarget it afterwards await ModelManagementAuthChecks.can_user_make_model_call( model_params=Deployment( model_name="test_model", litellm_params=LiteLLM_Params(**litellm_params), - model_info=auth_model_info, + model_info=resolved_model_info, ), user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, premium_user=premium_user, ) - # Include health_check_params if provided - litellm_params = _update_litellm_params_for_health_check( - model_info={}, - litellm_params=litellm_params, - ) mode = mode or litellm_params.pop("mode", None) result: Final = await run_with_timeout( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 99d0c94d11b..6abfca1d3a0 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -5,6 +5,7 @@ from typing import Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, @@ -318,6 +319,21 @@ class _ProxyDBLogger(CustomLogger): elif budget_reservation is not None: await _release_budget_reservation(budget_reservation=budget_reservation) else: + if _is_unbilled_interaction_response(completion_response): + if BACKGROUND_INTERACTION_COST_POLLING_ENABLED and _is_unbilled_in_progress_interaction( + completion_response + ): + verbose_proxy_logger.debug( + "Cost tracking deferred for in-progress background interaction; " + "the budget reservation stays open until the poll task logs the final usage" + ) + return + await _release_budget_reservation(budget_reservation=budget_reservation) + verbose_proxy_logger.debug( + "Released the budget reservation for an interaction create with no usage " + "that no poll task will settle" + ) + return await _release_budget_reservation(budget_reservation=budget_reservation) # Non-model call types (health checks, afile_delete) have no model or standard_logging_object. # Use .get() for "stream" to avoid KeyError on health checks. @@ -463,6 +479,24 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +def _is_unbilled_interaction_response(completion_response: object) -> bool: + from litellm.interactions.background_cost_polling import missing_usage_is_expected + from litellm.types.interactions import InteractionsAPIResponse + + if not isinstance(completion_response, InteractionsAPIResponse): + return False + return completion_response.usage is None and missing_usage_is_expected(completion_response) + + +def _is_unbilled_in_progress_interaction(completion_response: object) -> bool: + from litellm.interactions.background_cost_polling import is_pollable_background_interaction + from litellm.types.interactions import InteractionsAPIResponse + + if not isinstance(completion_response, InteractionsAPIResponse): + return False + return completion_response.usage is None and is_pollable_background_interaction(completion_response) + + def _should_track_cost_callback( user_api_key: str | None, user_id: str | None, diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 414beabe014..83caa92ede5 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,6 +1,7 @@ import asyncio import io import traceback +from collections.abc import Sequence from typing import Final import orjson @@ -33,10 +34,10 @@ async def uploadfile_to_bytesio(upload: UploadFile) -> io.BytesIO: async def batch_to_bytesio( - uploads: list[UploadFile] | None, + uploads: Sequence[UploadFile] | None, ) -> list[io.BytesIO] | None: """ - Convert a list of UploadFiles to a list of BytesIO buffers, or None. + Convert a sequence of UploadFiles to a list of BytesIO buffers, or None. """ if not uploads: return None diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 2271501d480..c803d87e436 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -210,7 +210,6 @@ async def _patch_team_caches_add_access_group( for team_id in team_ids: cached_team = await _get_team_object_from_cache( key=f"team_id:{team_id}", - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=None, ) @@ -240,7 +239,6 @@ async def _patch_team_caches_remove_access_group( for team_id in team_ids: cached_team = await _get_team_object_from_cache( key=f"team_id:{team_id}", - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=None, ) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 46aac82473c..1322f50d4af 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.team_repository import TeamRepository from litellm.router_strategy.complexity_router import ComplexityRouter +from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model from litellm.types.management_endpoints.auto_router_endpoints import ( SHADOW_EVAL_TURN_VALVE, AutoRouterBenchmarkGroup, @@ -510,6 +511,53 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow: ) +def _strategy_router_key(deployment: object) -> tuple[str, str] | None: + """``(model_name, kind)`` for a deployment whose routing the session rollup records. + + Kinds come from ``classify_strategy_router_model``, the same rule the Router registers a + deployment by, so this arm cannot disagree with the arm that stamped ``router_type`` onto + the session rows. Semantic auto-routers return None: they record no routing decision, so + they can never own a session row, and ``AutoRouterBenchmarkGroup.router_type`` has no + value for them. A permanent zero would read as "no traffic" rather than "not instrumented". + """ + if not isinstance(deployment, Mapping): + return None + litellm_params: Final = deployment.get("litellm_params") + router_name: Final = deployment.get("model_name") + if not (isinstance(litellm_params, Mapping) and isinstance(router_name, str) and router_name): + return None + model: Final = litellm_params.get("model") + if not isinstance(model, str): + return None + kind: Final = classify_strategy_router_model(model) + return None if kind is None or kind == "semantic" else (router_name, kind) + + +def _idle_router_groups( + llm_router: "Router | None", covered: frozenset[tuple[str, str]] +) -> tuple[AutoRouterBenchmarkGroup, ...]: + """Zeroed groups for configured strategy routers the window's traffic did not cover. + + The dashboard's router picker has to list a router the moment it is created rather than + once it has spent something, so the registry drives the list and the rollup only supplies + the measures. ``_summed_agg_row`` over no sessions is already the zero element of the + fold, so a group with every measure at zero costs one relabel rather than a literal that + would go stale the next time the response grows a field. + """ + if llm_router is None: + return () + zero: Final = _summed_agg_row(()) + idle: Final = frozenset( + key + for key in (_strategy_router_key(deployment) for deployment in llm_router.model_list or ()) + if key is not None and key not in covered + ) + return tuple( + _benchmark_group(zero.model_copy(update=MappingProxyType({"router_name": name, "router_type": kind}))) + for name, kind in sorted(idle) + ) + + @router.get( "/auto_router/benchmarks", tags=("auto router",), @@ -532,8 +580,13 @@ async def get_auto_router_benchmarks( overlaps it: its last turn is on or after start_date and its first turn is on or before end_date. Overall hit rate is over telemetry-bearing turns; each bucket's hit rate is over that bucket's turns. + + The rollup supplies the measures, never the list. Which routers appear comes from the + model registry, so one shows up as soon as it is configured and reads zero until it + serves traffic, and `routers_in_scope` counts those too rather than only the routers the + window recorded. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import llm_router, prisma_client _require_admin_viewer(user_api_key_dict, "view auto-router benchmarks across the deployment") if prisma_client is None: @@ -555,11 +608,14 @@ async def get_auto_router_benchmarks( (end_day + timedelta(days=1)).isoformat(), ) rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ()) - groups: Final = tuple(_benchmark_group(row) for row in rows) + groups: Final = ( + *(_benchmark_group(row) for row in rows), + *_idle_router_groups(llm_router, frozenset((row.router_name, row.router_type) for row in rows)), + ) return AutoRouterBenchmarksResponse( start_date=start_day.strftime("%Y-%m-%d"), end_date=end_day.strftime("%Y-%m-%d"), - routers_in_scope=len(rows), + routers_in_scope=len(groups), totals=_benchmark_totals(_summed_agg_row(rows)), groups=groups, ) diff --git a/litellm/proxy/management_endpoints/scim/scim_transformations.py b/litellm/proxy/management_endpoints/scim/scim_transformations.py index 2d95d0bea29..496be05b4b0 100644 --- a/litellm/proxy/management_endpoints/scim/scim_transformations.py +++ b/litellm/proxy/management_endpoints/scim/scim_transformations.py @@ -198,15 +198,8 @@ class ScimTransformations: @staticmethod def _get_scim_member_value(member: Member) -> str: - """ - Get the SCIM member value. Use user_email if available, otherwise use user_id. - SCIM member value should be the unique identifier for the user. - """ - if hasattr(member, "user_email") and member.user_email: - return member.user_email - elif hasattr(member, "user_id"): - return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE - return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE + """The member's SCIM resource id, which LiteLLM serves as user_id (RFC 7643 §8.7.1).""" + return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE @staticmethod def _get_scim_member_display(member: Member) -> str: diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index 7183e6cb402..6658963d024 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -176,6 +176,10 @@ class UserProvisionerHelpers: is persisted too, so re-upserting an existing email demotes a user who is no longer in the admin group instead of leaving the stale role. + IdPs like Entra manage membership exclusively through /Groups and never send + ``groups`` on POST /Users, so a request without teams means "unspecified", + not "remove from every team": existing memberships are preserved then. + Args: prisma_client: Database client new_user_request: New user request data @@ -194,7 +198,8 @@ class UserProvisionerHelpers: if not existing_user: return None - new_teams: Final = list(dict.fromkeys(new_user_request.teams or [])) + requested_teams: Final = list(dict.fromkeys(new_user_request.teams or [])) + new_teams: Final = requested_teams if requested_teams else list(existing_user.teams or []) if new_user_request.user_id != existing_user.user_id: verbose_proxy_logger.info( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 01254d5c064..49461d7841d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -16,7 +16,7 @@ import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import Annotated, Final, NamedTuple, Protocol, TypedDict, TypeVar, cast +from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -56,6 +56,7 @@ from litellm.proxy._types import ( PatchTeamRequest, ProxyErrorTypes, ProxyException, + ResetSpendRequest, SpecialManagementEndpointEnums, SpecialModelNames, SpecialProxyStrings, @@ -84,6 +85,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_membership, get_team_object, get_user_object, + invalidate_team_member_spend_state, ) from litellm.proxy.auth.auth_utils import ( enforce_batch_enqueued_token_limit_is_admin_only, @@ -3392,7 +3394,7 @@ async def team_member_update( Update team member budgets and team member role """ - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3491,6 +3493,12 @@ async def team_member_update( budget_patch=budget_patch, team_default_budget_id=team_default_budget_id, ) + if budget_patch: + await invalidate_team_member_spend_state( + user_id=received_user_id, + team_id=data.team_id, + user_api_key_cache=user_api_key_cache, + ) ### update team member role if data.role is not None: @@ -3527,6 +3535,125 @@ async def team_member_update( ) +def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAuth) -> None: + """ + _verify_team_access authorizes a team admin (or org admin) over their own + team, with no check that the target user_id differs from the caller. Left + unchecked, that admin could target their own LiteLLM_TeamMembership row and + repeatedly reset it to 0 right before it crosses their per-member cap, + consuming the shared team budget without the configured limit ever binding. + Only a proxy admin may reset an admin's own spend. + """ + if user_id == user_api_key_dict.user_id and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + _raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.") + + +def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn: + detail: Final = {"error": message} # mutable-ok: HTTPException.detail takes a dict + raise HTTPException(status_code=status_code, detail=detail) + + +def _validate_team_member_reset_spend_value( + reset_to: object, + membership: LiteLLM_TeamMembership, +) -> float: + if not isinstance(reset_to, (int, float)): + _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a float") + + reset_to_float: Final = float(reset_to) + if not math.isfinite(reset_to_float) or reset_to_float < 0: + _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a finite number >= 0") + + current_spend: Final = membership.spend or 0.0 + if reset_to_float > current_spend: + _raise_reset_spend_error( + status.HTTP_400_BAD_REQUEST, + f"reset_to ({reset_to_float}) must be <= current spend ({current_spend})", + ) + + max_budget: Final = membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None + if max_budget is not None and reset_to_float > max_budget: + _raise_reset_spend_error( + status.HTTP_400_BAD_REQUEST, + f"reset_to ({reset_to_float}) must be <= budget ({max_budget})", + ) + + return reset_to_float + + +@router.post( + "/team/{team_id}/member/{user_id}/reset_spend", + tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence + dependencies=(Depends(user_api_key_auth),), +) +@management_endpoint_wrapper +async def reset_team_member_spend_fn( + team_id: str, + user_id: str, + data: ResetSpendRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + """ + Reset a team member's tracked spend against their per-member budget. + + A member's spend is tracked separately from both their own personal + budget and the team's own budget (LiteLLM_TeamMembership.spend), so + neither /user/update nor /team/update can clear it: this is the only + endpoint that does. The cross-pod spend counter and cached membership + reads are invalidated so the reset takes effect on the member's next + request rather than waiting on the membership cache's TTL. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + _raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None") + + team_obj: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) + _check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict) + + membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument + } + _membership_row: Final = await _team_membership_db(prisma_client).find_unique( + where=membership_where, + include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument + ) + if _membership_row is None: + _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.") + membership: Final = LiteLLM_TeamMembership.model_validate(_membership_row.model_dump()) + + current_spend: Final = membership.spend or 0.0 + reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership) + + await _team_membership_db(prisma_client).update( + where=membership_where, + data={"spend": reset_to}, # mutable-ok: prisma client requires a plain dict data= argument + ) + + await invalidate_team_member_spend_state( + user_id=user_id, + team_id=team_id, + user_api_key_cache=user_api_key_cache, + new_spend=reset_to, + ) + + return { # mutable-ok: matches this router's established untyped-response-dict convention + "team_id": team_id, + "user_id": user_id, + "spend": reset_to, + "previous_spend": current_spend, + "max_budget": membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None, + } + + def _create_results_from_response( members: list[Member], response: TeamAddMemberResponse, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 7ce41c1d5b6..850b50227f4 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -9,7 +9,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc. import json import os import re -from collections.abc import Callable +from collections.abc import Callable, Mapping from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, cast @@ -29,7 +29,11 @@ from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.auth.user_api_key_auth import user_api_key_auth, user_api_key_auth_websocket +from litellm.proxy.auth.user_api_key_auth import ( + _get_bearer_token, + user_api_key_auth, + user_api_key_auth_websocket, +) from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, @@ -106,6 +110,36 @@ def is_passthrough_request_streaming(request_body: object) -> bool: return bool(request_body.get("stream", False)) +def get_passthrough_router_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, Any]: + """ + Build the request metadata carrying key-level spend attribution and the + pre-call budget reservation for a router-model passthrough request. + + Router-model passthrough branches call ``allm_passthrough_route`` directly, + bypassing ``add_litellm_data_to_request``. Without this metadata the cost + callback cannot attribute spend to the calling key and never releases the + budget reservation minted at auth time, so the shared spend counter drifts + up until the key falsely trips ``BudgetExceededError``. + + The payload rides the ``litellm_metadata`` bucket, not ``metadata``: the + router hop ``_ageneric_api_call_with_fallbacks`` canonicalises this call + type into ``litellm_metadata``, and the cost callback reads spend + attribution from that bucket while only backfilling ``user_api_key*`` keys + from ``metadata``. Passing ``metadata=`` would silently drop the secondary + attribution fields the helper sets (``agent_id``, + ``user_api_end_user_max_budget``) before the callback ever sees them. + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + request_data: Final = {"litellm_metadata": {}} # mutable-ok: builder + litellm mutate this in place + LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( + data=request_data, + user_api_key_dict=user_api_key_dict, + _metadata_variable_name="litellm_metadata", + ) + return request_data["litellm_metadata"] + + async def llm_passthrough_factory_proxy_route( custom_llm_provider: str, endpoint: str, @@ -346,6 +380,7 @@ async def vllm_proxy_route( params=None, headers=None, cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), ), ) @@ -1475,6 +1510,7 @@ async def azure_proxy_route( params=None, headers=None, cookies=None, + litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict), ) if is_streaming_request: @@ -1726,6 +1762,154 @@ def _override_vertex_params_from_router_credentials( return vertex_project, vertex_location +_CREDENTIALLESS_VERTEX_MISSING_CREDENTIAL_DETAIL: Final = ( + "No Vertex AI credential is configured on this proxy and the request carried no upstream " + "Google credential. The LiteLLM virtual key is not forwarded to Google. Configure a Vertex " + "credential (DEFAULT_VERTEXAI_PROJECT / DEFAULT_VERTEXAI_LOCATION / DEFAULT_VERTEXAI_CREDENTIALS, " + "or a model with use_in_pass_through: true), or send your own Google OAuth token in the " + "Authorization header." +) + + +def _normalize_credential_value(value: str) -> str: + """Reduce a header value to the bare token, matching how ``user_api_key_auth`` + reads a caller's key. + + Reuses the auth module's ``_get_bearer_token`` so the caller-key comparison + strips exactly the schemes authentication accepts (``Bearer`` / ``bearer`` / + ``Basic`` / ``AWS4-HMAC-SHA256`` credential), rather than re-deriving a + narrower normalization here. ``_get_bearer_token`` returns ``""`` for a value + with no recognized scheme prefix, so a bare token (or a real Google + credential that carries no scheme) falls back to its own value. + """ + return _get_bearer_token(value) or value + + +_VERTEX_UPSTREAM_CREDENTIAL_HEADERS: Final = frozenset({"authorization", "x-goog-api-key"}) +_HEADERS_NEVER_FORWARDED_TO_VERTEX: Final = frozenset({"content-length", "host"}) | ( + SpecialHeaders.litellm_credential_header_names() - _VERTEX_UPSTREAM_CREDENTIAL_HEADERS +) + + +_VERTEX_CALLER_KEY_HEADER_PRECEDENCE: Final = ( + SpecialHeaders.custom_litellm_api_key.value.lower(), + SpecialHeaders.openai_authorization.value.lower(), + SpecialHeaders.azure_authorization.value.lower(), + SpecialHeaders.anthropic_authorization.value.lower(), + SpecialHeaders.google_ai_studio_authorization.value.lower(), + SpecialHeaders.azure_apim_authorization.value.lower(), +) + +_MAPPED_ROUTE_CALLER_KEY_HEADER: Final = "litellm_user_api_key" + + +def _operator_configured_caller_key_header_names() -> tuple[tuple[str, ...], tuple[str, ...]]: + """Operator-configured caller-key header names, as (override, pass_through). + + ``user_api_key_auth`` accepts the caller's key from two runtime-configured + header sources beyond the built-in ones, at opposite ends of its precedence. + ``general_settings.litellm_key_header_name`` overrides every built-in source + (it replaces the resolved key after ``get_api_key`` runs), so it is highest + precedence. Each ``general_settings.pass_through_endpoints`` entry's + ``headers.litellm_user_api_key`` is checked last inside ``get_api_key``, so it + is lowest. Google never consumes either, so both are also dropped by name. + """ + from litellm.proxy.proxy_server import general_settings + + custom_key_header: Final = general_settings.get("litellm_key_header_name") + override: Final = (custom_key_header.lower(),) if isinstance(custom_key_header, str) else () + pass_through_endpoints: Final = general_settings.get("pass_through_endpoints") + endpoints: Final = pass_through_endpoints if isinstance(pass_through_endpoints, list) else () + pass_through: Final = tuple( + dict.fromkeys( + headers["litellm_user_api_key"].lower() + for endpoint in endpoints + if isinstance(endpoint, dict) + for headers in (endpoint.get("headers"),) + if isinstance(headers, dict) and isinstance(headers.get("litellm_user_api_key"), str) + ) + ) + return override, pass_through + + +def _authenticated_caller_key_values(request: Request) -> frozenset[str]: + """The value ``user_api_key_auth`` would accept as this caller's LiteLLM key. + + The Vertex route authenticates through ``Depends(user_api_key_auth)``, which + resolves the key by precedence, matched here exactly. The ``/vertex_ai`` route + is a mapped pass-through route, so a header literally named + ``litellm_user_api_key`` overrides every other source (``user_api_key_auth`` + applies it last), making it highest precedence. Then an operator + ``litellm_key_header_name``, then the built-in headers in ``get_api_key`` order, + then a ``pass_through_endpoints`` ``litellm_user_api_key`` header which + ``get_api_key`` checks last. Some of those headers (``Authorization``, + ``x-goog-api-key``) are also kept as genuine bring-your-own Google credentials, + so returning only the value that actually authenticated lets the filter strip + that value wherever it appears while leaving a real Google credential in place. + An empty set means no caller key was found, so nothing is value-stripped. + """ + incoming: Final = _safe_get_request_headers(request) + override_headers, pass_through_headers = _operator_configured_caller_key_header_names() + ordered_names: Final = ( + (_MAPPED_ROUTE_CALLER_KEY_HEADER,) + + override_headers + + _VERTEX_CALLER_KEY_HEADER_PRECEDENCE + + pass_through_headers + ) + present_values: Final = (incoming[name] for name in ordered_names if incoming.get(name)) + authenticated_key: Final = next( + (stripped for value in present_values if (stripped := _normalize_credential_value(value))), + "", + ) + return frozenset({authenticated_key}) if authenticated_key else frozenset() + + +def _forwarded_headers_for_credentialless_vertex_passthrough(request: Request) -> Mapping[str, str]: + """ + Header set to forward on the bring-your-own-credentials Vertex passthrough + branch, used when the proxy has no Vertex credential configured. + + No credential the proxy accepts for caller authentication is forwarded to + Google. ``user_api_key_auth`` reads the caller's key from every header in + ``SpecialHeaders.litellm_credential_header_names()``, and Vertex only ever + authenticates with an OAuth token in ``Authorization`` or an API key in + ``x-goog-api-key``. So the proxy-only auth headers Google never consumes + (everything in that set except those two, e.g. ``x-litellm-api-key`` / + ``api-key`` / ``x-api-key`` / ``Ocp-Apim-Subscription-Key``, plus the mapped + pass-through ``litellm_user_api_key`` header and any operator-configured + ``litellm_key_header_name`` / ``pass_through_endpoints`` key header) are dropped + by name. ``Authorization`` and ``x-goog-api-key`` may + instead carry a genuine bring-your-own Google credential, so they are kept + unless their value is the caller's authenticated LiteLLM key, which is dropped + by value (normalizing any ``Bearer`` / ``Basic`` / ``AWS4`` auth-scheme prefix + the same way authentication does). Because the value that authenticated is + resolved by the same precedence ``user_api_key_auth`` uses, a virtual key sent + only in ``x-goog-api-key`` (or in an operator-configured key header) is dropped + too, while a real Google key in ``x-goog-api-key`` alongside a virtual key in a + higher-precedence header is preserved. When neither a surviving + ``Authorization`` nor ``x-goog-api-key`` remains the request is rejected so the + virtual key cannot leak upstream. + """ + incoming: Final = _safe_get_request_headers(request) + caller_key_values: Final = _authenticated_caller_key_values(request) + override_headers, pass_through_headers = _operator_configured_caller_key_header_names() + never_forwarded: Final = ( + _HEADERS_NEVER_FORWARDED_TO_VERTEX.union((_MAPPED_ROUTE_CALLER_KEY_HEADER,)) + .union(override_headers) + .union(pass_through_headers) + ) + forwarded: Final = MappingProxyType( + { + name: value + for name, value in incoming.items() + if name not in never_forwarded and _normalize_credential_value(value) not in caller_key_values + } + ) + if "authorization" not in forwarded and "x-goog-api-key" not in forwarded: + raise HTTPException(status_code=401, detail=_CREDENTIALLESS_VERTEX_MISSING_CREDENTIAL_DETAIL) + return forwarded + + async def _prepare_vertex_auth_headers( request: Request, vertex_credentials: Any | None, @@ -1734,7 +1918,7 @@ async def _prepare_vertex_auth_headers( vertex_location: str | None, base_target_url: str | None, get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler, -) -> tuple[dict, str | None, bool, str | None, str | None]: +) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]: """ Prepare authentication headers for Vertex AI pass-through requests. @@ -1760,11 +1944,11 @@ async def _prepare_vertex_auth_headers( # Use headers from the incoming request if no vertex credentials are found if (vertex_credentials is None or vertex_credentials.vertex_project is None) and router_credentials is None: - headers = _safe_get_request_headers(request).copy() + headers = _forwarded_headers_for_credentialless_vertex_passthrough(request) headers_passed_through = True - verbose_proxy_logger.debug("default_vertex_config not set, incoming request headers %s", headers) - headers.pop("content-length", None) - headers.pop("host", None) + verbose_proxy_logger.debug( + "default_vertex_config not set, forwarding caller-provided headers %s", tuple(headers.keys()) + ) else: if router_credentials is not None: vertex_credentials_str = None @@ -1850,7 +2034,7 @@ async def _base_vertex_proxy_route( encoded_endpoint = httpx.URL(endpoint).path verbose_proxy_logger.debug("requested endpoint %s", endpoint) - headers: dict = {} + headers: Mapping[str, str] = {} api_key_to_use = get_litellm_virtual_key(request=request) user_api_key_dict = await user_api_key_auth( request=request, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 1915a853983..3d721dead4d 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -470,7 +470,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): ``items()`` collapses duplicate keys to the last value. Files go out as a list of ``(field_name, (filename, content, content_type))`` tuples and repeated non-file fields are grouped into list values, both of which httpx - encodes as separate multipart parts. + encodes as separate multipart parts. A form with no file parts is sent + entirely through ``files`` as ``(field_name, (None, value))`` tuples, + because httpx downgrades a file-less ``data=`` payload to + application/x-www-form-urlencoded. """ form_items: Final = (await request.form()).multi_items() @@ -500,6 +503,11 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): ) } + multipart_files: Final = ( + files if files else tuple((field_name, (None, field_value)) for field_name, field_value in non_file_items) + ) + multipart_data: Final = form_data_dict if files else None + # Remove content-type header - httpx will set it correctly with the new boundary # when it creates the multipart body from files/data parameters headers_copy: Final = headers.copy() @@ -512,8 +520,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): url, headers=headers_copy, params=requested_query_params, - files=files, - data=form_data_dict, + files=multipart_files, + data=multipart_data, ) return await async_client.send(req, stream=True) @@ -522,8 +530,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): url=url, headers=headers_copy, params=requested_query_params, - files=files, - data=form_data_dict, + files=multipart_files, + data=multipart_data, ) @staticmethod diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index f66dc4e7bbe..001e4115374 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -30,7 +30,7 @@ class PolicyMatcher: """ Check if a value matches any of the given patterns. - Uses the existing RouteChecks._route_matches_wildcard_pattern helper. + Uses the existing RouteChecks.route_matches_wildcard_pattern helper. Args: value: The value to check (e.g., team alias, key alias, model) @@ -45,7 +45,7 @@ class PolicyMatcher: for pattern in patterns: # Use existing wildcard pattern matching helper - if RouteChecks._route_matches_wildcard_pattern(route=value, pattern=pattern): + if RouteChecks.route_matches_wildcard_pattern(route=value, pattern=pattern): return True return False diff --git a/litellm/proxy/policy_engine/policy_resolve_endpoints.py b/litellm/proxy/policy_engine/policy_resolve_endpoints.py index 346586c1e5a..70b98933d0f 100644 --- a/litellm/proxy/policy_engine/policy_resolve_endpoints.py +++ b/litellm/proxy/policy_engine/policy_resolve_endpoints.py @@ -100,7 +100,7 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple: key_alias = key.key_alias or "" key_tags = _get_tags_from_metadata(key.metadata, getattr(key, "metadata_json", None)) if key_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in key_tags for pat in tag_patterns ): @@ -123,7 +123,7 @@ def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple: team_alias = team.team_alias or "" team_tags = _get_tags_from_metadata(team.metadata) if team_tags and any( - RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat) + RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat) for tag in team_tags for pat in tag_patterns ): @@ -152,7 +152,7 @@ async def _find_affected_by_team_patterns( for team in all_teams: team_alias = team.team_alias or "" if team_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns + RouteChecks.route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns ): if team_alias not in existing_teams: new_teams.append(team_alias) @@ -190,7 +190,7 @@ async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list for key in keys: key_alias = key.key_alias or "" if key_alias and any( - RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns + RouteChecks.route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns ): if key_alias not in existing_keys: affected.append(key_alias) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dced4e26b6..0abcdeaf3f6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2555,6 +2555,12 @@ async def _authoritative_floor_spend( if db_spend is None: return None + # a spend reset that committed during the DB read above wrote the post-reset + # floor to the marker; keep it over this read's now-stale pre-commit value + rechecked: Final = spend_counter_cache.in_memory_cache.get_cache(key=marker_key) + if rechecked is not None: + return float(rechecked) + spend_counter_cache.in_memory_cache.set_cache( key=marker_key, value=db_spend, @@ -6798,6 +6804,7 @@ class ProxyConfig: subscriber: Final = AuthCacheInvalidationSubscriber( redis_cache=redis_cache, user_api_key_cache=user_api_key_cache, + additional_in_memory_caches=(spend_counter_cache.in_memory_cache,), ) self.auth_cache_invalidation_subscriber = subscriber subscriber.start() diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 9e2b1c9d82d..4d62f1d6d71 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -27,6 +27,13 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, get_form_data, ) +from litellm.proxy.rag_endpoints.upload_security import ( + MAX_UPLOAD_SIZE_BYTES, + EicarTestMalwareScanner, + MalwareScanner, + RejectedUpload, + validate_upload, +) from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, ) @@ -287,8 +294,22 @@ async def _save_vector_store_to_db_from_rag_ingest( verbose_proxy_logger.exception("Failed to save vector store %s to database: %s", vector_store_id, db_error) +def _secure_uploaded_file( + file_data: tuple[str, bytes, str], + scanner: MalwareScanner, +) -> tuple[str, bytes, str]: + validation: Final = validate_upload(content=file_data[1], scanner=scanner) + if isinstance(validation, RejectedUpload): + raise HTTPException( + status_code=400, + detail={"error": validation.message, "reason": validation.reason.value}, + ) + return validation.safe_filename, file_data[1], validation.content_type + + async def parse_rag_ingest_request( request: Request, + scanner: MalwareScanner, ) -> tuple[dict[str, Any], tuple[str, bytes, str] | None, str | None, str | None]: """ Parse RAG ingest request. @@ -297,6 +318,11 @@ async def parse_rag_ingest_request( - Form: file + request JSON in form field - JSON body for URL-based ingestion + Uploaded file bytes are validated against the vector-store upload controls + (size limit, format allowlist with content inspection, archive rejection, + and the injected malware scanner) and given a server-generated filename + before they are returned. + Returns: Tuple of (ingest_options, file_data, file_url, file_id) """ @@ -315,7 +341,7 @@ async def parse_rag_ingest_request( # Get file file_obj = form_data.get("file") if file_obj is not None and hasattr(file_obj, "read"): - file_content = await file_obj.read() + file_content = await file_obj.read(MAX_UPLOAD_SIZE_BYTES + 1) file_data = (file_obj.filename, file_content, file_obj.content_type) # Parse JSON from 'request' form field (contains full request body as JSON) @@ -357,6 +383,10 @@ async def parse_rag_ingest_request( detail={"error": "Must provide file, file_url, or file_id"}, ) + secured_file_data: Final[tuple[str, bytes, str] | None] = ( + _secure_uploaded_file(file_data, scanner) if file_data is not None else None + ) + if "vector_store" not in ingest_options: raise HTTPException( status_code=400, @@ -398,7 +428,7 @@ async def parse_rag_ingest_request( }, ) - return ingest_options, file_data, file_url, file_id + return ingest_options, secured_file_data, file_url, file_id @router.post( @@ -461,7 +491,9 @@ async def rag_ingest( try: # Parse request - ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(request) + ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request( + request, scanner=EicarTestMalwareScanner() + ) # INTERNAL_USER_VIEW_ONLY can ingest to existing vector stores only if user_api_key_dict.user_role == LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value and not ingest_options.get( diff --git a/litellm/proxy/rag_endpoints/upload_security.py b/litellm/proxy/rag_endpoints/upload_security.py new file mode 100644 index 00000000000..f0318f2f709 --- /dev/null +++ b/litellm/proxy/rag_endpoints/upload_security.py @@ -0,0 +1,289 @@ +"""Security controls for vector-store file uploads. + +Content is classified by inspecting its actual bytes (magic signatures and a +strict UTF-8 decode), never by trusting the client-supplied filename or +content-type. Uploads are restricted to an allowlist of non-executable formats, +capped in size, screened for archives, and passed through a dependency-injected +malware scanner before they are accepted. Accepted uploads are given a +server-generated filename so the client-controlled name never reaches storage. +""" + +from __future__ import annotations + +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from types import MappingProxyType +from typing import Final, Protocol, TypeAlias, runtime_checkable + +from typing_extensions import assert_never + +MAX_UPLOAD_SIZE_BYTES: Final = 512 * 1024 * 1024 + +EICAR_TEST_SIGNATURE: Final = b"X5O!P%@AP[4\\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" + +_ARCHIVE_MAGIC_PREFIXES: Final[tuple[bytes, ...]] = ( + b"PK\x03\x04", + b"PK\x05\x06", + b"PK\x07\x08", + b"\x1f\x8b", + b"\xfd7zXZ\x00", + b"7z\xbc\xaf\x27\x1c", + b"Rar!\x1a\x07\x00", + b"Rar!\x1a\x07\x01\x00", + b"\x04\x22\x4d\x18", + b"\x28\xb5\x2f\xfd", +) + +_ARCHIVE_MAGIC_PREFIXES_ASCII_AMBIGUOUS: Final[tuple[bytes, ...]] = (b"BZh",) + +_EXECUTABLE_MAGIC_PREFIXES: Final[tuple[bytes, ...]] = ( + b"\x7fELF", + b"\xca\xfe\xba\xbe", + b"\xfe\xed\xfa\xce", + b"\xfe\xed\xfa\xcf", + b"\xce\xfa\xed\xfe", + b"\xcf\xfa\xed\xfe", + b"\x00asm", +) + +_EXECUTABLE_MAGIC_PREFIXES_ASCII_AMBIGUOUS: Final[tuple[bytes, ...]] = (b"MZ", b"dex\n") + +_TAR_USTAR_MAGIC: Final = b"ustar" +_TAR_USTAR_OFFSET: Final = 257 + +_UTF8_BOM: Final = b"\xef\xbb\xbf" + + +class DetectedFormat(str, Enum): + PDF = "pdf" + TEXT = "text" + + +class DisallowedKind(str, Enum): + ARCHIVE = "archive" + EXECUTABLE = "executable" + UNKNOWN_BINARY = "unknown_binary" + + +class RejectionReason(str, Enum): + EMPTY_FILE = "empty_file" + FILE_TOO_LARGE = "file_too_large" + ARCHIVE_NOT_ALLOWED = "archive_not_allowed" + EXECUTABLE_NOT_ALLOWED = "executable_not_allowed" + UNSUPPORTED_FORMAT = "unsupported_format" + MALWARE_DETECTED = "malware_detected" + MALWARE_SCAN_ERROR = "malware_scan_error" + + +class ScanVerdict(str, Enum): + CLEAN = "clean" + INFECTED = "infected" + ERROR = "error" + + +@dataclass(frozen=True, slots=True) +class ScanResult: + verdict: ScanVerdict + signature: str | None = None + + +@runtime_checkable +class MalwareScanner(Protocol): + def scan(self, content: bytes) -> ScanResult: ... + + +@dataclass(frozen=True, slots=True) +class EicarTestMalwareScanner: + """Placeholder scanner that only flags the EICAR anti-malware test file. + + It exists to prove the scan hook is wired end to end and to satisfy the + EICAR retest; it provides no real protection. Inject a scanner backed by a + real engine through the ``scanner`` parameter of :func:`validate_upload` to + screen production uploads. + """ + + def scan(self, content: bytes) -> ScanResult: + if EICAR_TEST_SIGNATURE in content: + return ScanResult(verdict=ScanVerdict.INFECTED, signature="EICAR-STANDARD-ANTIVIRUS-TEST-FILE") + return ScanResult(verdict=ScanVerdict.CLEAN) + + +@dataclass(frozen=True, slots=True) +class AllowedContent: + format: DetectedFormat + + +@dataclass(frozen=True, slots=True) +class DisallowedContent: + kind: DisallowedKind + + +ContentInspection: TypeAlias = AllowedContent | DisallowedContent + + +@dataclass(frozen=True, slots=True) +class SecuredUpload: + safe_filename: str + content_type: str + detected_format: DetectedFormat + size_bytes: int + + +@dataclass(frozen=True, slots=True) +class RejectedUpload: + reason: RejectionReason + message: str + + +UploadValidation: TypeAlias = SecuredUpload | RejectedUpload + +_SAFE_EXTENSION: Final[Mapping[DetectedFormat, str]] = MappingProxyType( + { + DetectedFormat.PDF: "pdf", + DetectedFormat.TEXT: "txt", + } +) + +_SAFE_CONTENT_TYPE: Final[Mapping[DetectedFormat, str]] = MappingProxyType( + { + DetectedFormat.PDF: "application/pdf", + DetectedFormat.TEXT: "text/plain", + } +) + + +def _starts_with_any(content: bytes, prefixes: tuple[bytes, ...]) -> bool: + return any(content.startswith(prefix) for prefix in prefixes) + + +def _is_archive(content: bytes) -> bool: + if _starts_with_any(content, _ARCHIVE_MAGIC_PREFIXES): + return True + tar_magic_end: Final = _TAR_USTAR_OFFSET + len(_TAR_USTAR_MAGIC) + if len(content) >= tar_magic_end and content[_TAR_USTAR_OFFSET:tar_magic_end] == _TAR_USTAR_MAGIC: + return True + return _starts_with_any(content, _ARCHIVE_MAGIC_PREFIXES_ASCII_AMBIGUOUS) and not _is_utf8_text(content) + + +def _is_utf8_text(content: bytes) -> bool: + if b"\x00" in content: + return False + try: + content.decode("utf-8") + except UnicodeDecodeError: + return False + return True + + +def _is_executable_binary(content: bytes) -> bool: + if _starts_with_any(content, _EXECUTABLE_MAGIC_PREFIXES): + return True + return _starts_with_any(content, _EXECUTABLE_MAGIC_PREFIXES_ASCII_AMBIGUOUS) and not _is_utf8_text(content) + + +def _looks_like_shebang(content: bytes) -> bool: + body: Final = content.removeprefix(_UTF8_BOM).lstrip() + return body.startswith(b"#!") + + +def inspect_content(content: bytes) -> ContentInspection: + if _looks_like_shebang(content): + return DisallowedContent(DisallowedKind.EXECUTABLE) + if content.startswith(b"%PDF-"): + return AllowedContent(DetectedFormat.PDF) + if _is_archive(content): + return DisallowedContent(DisallowedKind.ARCHIVE) + if _is_executable_binary(content): + return DisallowedContent(DisallowedKind.EXECUTABLE) + if _is_utf8_text(content): + return AllowedContent(DetectedFormat.TEXT) + return DisallowedContent(DisallowedKind.UNKNOWN_BINARY) + + +def generate_safe_filename(detected_format: DetectedFormat) -> str: + return f"{uuid.uuid4().hex}.{_SAFE_EXTENSION[detected_format]}" + + +def _reject_disallowed(kind: DisallowedKind) -> RejectedUpload: + match kind: + case DisallowedKind.ARCHIVE: + return RejectedUpload( + RejectionReason.ARCHIVE_NOT_ALLOWED, + "Archive uploads are not allowed.", + ) + case DisallowedKind.EXECUTABLE: + return RejectedUpload( + RejectionReason.EXECUTABLE_NOT_ALLOWED, + "Executable uploads are not allowed.", + ) + case DisallowedKind.UNKNOWN_BINARY: + return RejectedUpload( + RejectionReason.UNSUPPORTED_FORMAT, + "Only PDF and UTF-8 text documents are accepted.", + ) + assert_never(kind) + + +def _scan_rejection(content: bytes, scanner: MalwareScanner) -> RejectedUpload | None: + result: Final = scanner.scan(content) + match result.verdict: + case ScanVerdict.CLEAN: + return None + case ScanVerdict.INFECTED: + return RejectedUpload( + RejectionReason.MALWARE_DETECTED, + f"Uploaded file was flagged by malware scanning ({result.signature or 'unknown signature'}).", + ) + case ScanVerdict.ERROR: + return RejectedUpload( + RejectionReason.MALWARE_SCAN_ERROR, + "Malware scanning could not complete; upload rejected.", + ) + assert_never(result.verdict) + + +def validate_upload( + *, + content: bytes, + scanner: MalwareScanner, + max_size_bytes: int = MAX_UPLOAD_SIZE_BYTES, +) -> UploadValidation: + size: Final = len(content) + if size == 0: + return RejectedUpload(RejectionReason.EMPTY_FILE, "Uploaded file is empty.") + if size > max_size_bytes: + return RejectedUpload( + RejectionReason.FILE_TOO_LARGE, + f"Uploaded file is {size} bytes, exceeding the {max_size_bytes}-byte limit.", + ) + + inspection: Final = inspect_content(content) + if isinstance(inspection, DisallowedContent): + return _reject_disallowed(inspection.kind) + + scan_rejection: Final = _scan_rejection(content, scanner) + if scan_rejection is not None: + return scan_rejection + + return SecuredUpload( + safe_filename=generate_safe_filename(inspection.format), + content_type=_SAFE_CONTENT_TYPE[inspection.format], + detected_format=inspection.format, + size_bytes=size, + ) + + +def _sanitize_header_filename(filename: str) -> str: + stripped: Final = "".join(char for char in filename if char not in '"\\\r\n').strip() + return stripped or "download" + + +def safe_download_headers(filename: str) -> Mapping[str, str]: + return MappingProxyType( + { + "Content-Disposition": f'attachment; filename="{_sanitize_header_filename(filename)}"', + "X-Content-Type-Options": "nosniff", + } + ) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ce6c9330620..149f9b960a1 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -25,7 +25,11 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_cache_key, + tag_cache_key, + team_membership_reservation_cache_key, +) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -546,7 +550,9 @@ async def _get_team_member_budget_counter( if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None - membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}" + membership_cache_key: Final = team_membership_reservation_cache_key( + user_id=valid_token.user_id, team_id=team_object.team_id + ) cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key) team_membership: LiteLLM_TeamMembership | None = None if isinstance(cached_team_membership, LiteLLM_TeamMembership): diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b6f695db512..d931d712a92 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -444,7 +444,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs or None ) raw_model: Final = cast(str, kwargs.get("model") or "") - model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) + model_name: Final = ( + standard_logging_payload.get("model") if standard_logging_payload is not None else None + ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: Final[SpendLogsPayload] = SpendLogsPayload( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 86d954c0913..2bbdad9c487 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,7 +91,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert -from litellm.litellm_core_utils.core_helpers import coerce_token_limit +from litellm.litellm_core_utils.core_helpers import coerce_token_limit, is_expected_client_error from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -2575,20 +2575,36 @@ class ProxyLogging: api_key="", ) - # log the custom exception - await litellm_logging_obj.async_failure_handler( - exception=original_exception, - traceback_exception=traceback.format_exc(), + await self._dispatch_proxy_only_failure_handlers( + litellm_logging_obj=litellm_logging_obj, + original_exception=original_exception, ) - threading.Thread( - target=litellm_logging_obj.failure_handler, - args=( - original_exception, - traceback.format_exc(), - ), - daemon=True, - ).start() + @staticmethod + async def _dispatch_proxy_only_failure_handlers( + litellm_logging_obj: Logging, + original_exception: Exception | None, + ) -> None: + """Runs the async failure handler plus the threaded sync handler. Expected + client (4xx) errors skip traceback formatting unless + litellm.log_client_error_tracebacks is set.""" + include_traceback: Final = litellm.log_client_error_tracebacks or not is_expected_client_error( + original_exception + ) + traceback_str: Final = traceback.format_exc() if include_traceback else "" + await litellm_logging_obj.async_failure_handler( + exception=original_exception, + traceback_exception=traceback_str, + ) + + threading.Thread( + target=litellm_logging_obj.failure_handler, + args=( + original_exception, + traceback_str, + ), + daemon=True, + ).start() async def post_call_success_hook( self, diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index c9b89bcd390..957ed9fd0b9 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( handle_model_based_routing, prepare_data_with_credentials, ) +from litellm.proxy.rag_endpoints.upload_security import safe_download_headers from litellm.proxy.vector_store_endpoints.utils import ( assert_user_can_access_vector_store_id, is_allowed_to_call_vector_store_files_endpoint, @@ -885,6 +886,9 @@ async def vector_store_file_content( if original_managed_file_id: response = _replace_file_id_in_response(response, original_managed_file_id) + for header_name, header_value in safe_download_headers(file_id).items(): + fastapi_response.headers[header_name] = header_value + return response except Exception as e: # noqa: BLE001 raise await processor._handle_llm_api_exception( diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 6c6b004fd17..d985a546fa7 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -2,9 +2,9 @@ from typing import Any, Final -import orjson from fastapi import APIRouter, Depends, File, Form, Request, Response, UploadFile from fastapi.responses import ORJSONResponse +from starlette.datastructures import UploadFile as StarletteUploadFile from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth @@ -20,6 +20,7 @@ from litellm.proxy.video_endpoints.utils import ( encode_character_id_in_response, extract_model_from_target_model_names, get_custom_provider_from_data, + video_reference_to_id, ) from litellm.types.videos.utils import ( decode_character_id_with_provider, @@ -451,9 +452,7 @@ async def video_remix( version, ) - # Read request body - body: Final = await request.body() - data: Final = orjson.loads(body) + data: Final = await _read_request_body(request=request) data["video_id"] = video_id decoded: Final = decode_video_id_with_provider(video_id) @@ -760,15 +759,17 @@ async def video_edit( version, ) - body: Final = await request.body() - data: Final = orjson.loads(body) + data: Final = await _read_request_body(request=request) + uploaded_video: Final = data.pop("video", None) + if isinstance(uploaded_video, StarletteUploadFile): + video_files: Final = await batch_to_bytesio((uploaded_video,)) + if video_files: + data["video"] = video_files[0] + data["video_id"] = "" + else: + data["video_id"] = video_reference_to_id(uploaded_video) - # Extract video_id from nested video object - video_ref: Final = data.pop("video", {}) - video_id: Final = video_ref.get("id", "") if isinstance(video_ref, dict) else "" - data["video_id"] = video_id - - decoded: Final = decode_video_id_with_provider(video_id) + decoded: Final = decode_video_id_with_provider(data["video_id"]) provider_from_id: Final = decoded.get("custom_llm_provider") model_id_from_decoded: Final = decoded.get("model_id") @@ -860,15 +861,10 @@ async def video_extension( version, ) - body: Final = await request.body() - data: Final = orjson.loads(body) + data: Final = await _read_request_body(request=request) + data["video_id"] = video_reference_to_id(data.pop("video", None)) - # Extract video_id from nested video object - video_ref: Final = data.pop("video", {}) - video_id: Final = video_ref.get("id", "") if isinstance(video_ref, dict) else "" - data["video_id"] = video_id - - decoded: Final = decode_video_id_with_provider(video_id) + decoded: Final = decode_video_id_with_provider(data["video_id"]) provider_from_id: Final = decoded.get("custom_llm_provider") model_id_from_decoded: Final = decoded.get("model_id") diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index d6b398e3476..a38226cc253 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -13,6 +13,18 @@ def extract_model_from_target_model_names(target_model_names: Any) -> str | None return target_model_names[0] if target_model_names else None +def video_reference_to_id(video_ref: object) -> str: + if isinstance(video_ref, dict): + return video_ref.get("id", "") + if not isinstance(video_ref, str): + return "" + try: + parsed_ref: Final = orjson.loads(video_ref) + except orjson.JSONDecodeError: + return video_ref + return parsed_ref.get("id", "") if isinstance(parsed_ref, dict) else video_ref + + def get_custom_provider_from_data(data: dict[str, Any]) -> str | None: custom_llm_provider: Final = data.get("custom_llm_provider") if custom_llm_provider: diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 15a6f18a6bb..c8f7842aebf 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -277,6 +277,8 @@ def rerank( if api_key is None: raise ValueError("TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment") + api_base = dynamic_api_base or optional_params.api_base or litellm.api_base or "https://api.together.ai/v1" + response = together_rerank.rerank( model=model, query=query, @@ -286,6 +288,7 @@ def rerank( return_documents=return_documents, max_chunks_per_doc=max_chunks_per_doc, api_key=api_key, + api_base=api_base, _is_async=_is_async, ) elif _custom_llm_provider == litellm.LlmProviders.JINA_AI: diff --git a/litellm/router.py b/litellm/router.py index 045fd32847c..d07effd0d90 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8,6 +8,7 @@ # Thank you ! We ❤️ you! - Krrish & Ishaan import asyncio +import contextlib import copy import enum import hashlib @@ -20,7 +21,7 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping, Sequence from functools import lru_cache from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -80,6 +81,7 @@ from litellm.litellm_core_utils.request_timeout_resolver import ( from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, + mask_credentials_in_payload, mask_sensitive_structure, ) from litellm.llms.openai_like.json_loader import JSONProviderRegistry @@ -167,6 +169,11 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, ) +from litellm.router_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -242,6 +249,7 @@ from .router_utils.pattern_match_deployments import PatternMatchRouter if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.exceptions import MidStreamFallbackError from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, ) @@ -258,6 +266,9 @@ if TYPE_CHECKING: from litellm.router_strategy.quality_router.quality_router import ( QualityRouter, ) + from litellm.types.llms.anthropic_messages.anthropic_response import ( + AnthropicMessagesResponse, + ) from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject from litellm.types.llms.openai import ( ResponseAPIUsage, @@ -355,6 +366,101 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) return False +# Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks +# until real content commits the primary stream; a hostile or slow-starting +# upstream that never emits content or an error could otherwise grow that +# buffer without bound, so hitting this cap forces an early commit instead. +MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200 + + +def _anthropic_stream_should_drop_pre_content_ping(chunk: object, has_generated_content: bool) -> bool: + """A `ping` keepalive seen before any real content is dropped outright - it recurs indefinitely on a + slow-starting connection and carries nothing worth buffering toward a possible fallback.""" + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import is_anthropic_ping_chunk + + if has_generated_content: + return False + return is_anthropic_ping_chunk(chunk) + + +def _is_retriable_anthropic_status(status_code: int) -> bool: + return status_code == 429 or status_code >= 500 + + +def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error: "MidStreamFallbackError") -> bool: + """ + A MidStreamFallbackError raised directly by the source iterator (the + completion-bridge path's CustomStreamWrapper, e.g. on a transport drop) + carries its own pre_first_chunk bookkeeping - gated the same way a + detected SSE error event is, so a fallback is never appended after real + content already reached the client on either path. + """ + return has_generated_content or not error.is_pre_first_chunk + + +def _anthropic_stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: + """ + Whether `chunk` should make Router._aanthropic_messages_streaming_iterator + commit to the primary Anthropic stream (real content arrived, or the + pre-content buffer cap was hit) rather than keep buffering lifecycle + frames toward a possible fallback. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + is_anthropic_content_delta_chunk, + ) + + if has_generated_content: + return False + return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + + +class FallbackAwareAnthropicMessagesStream: + """ + Bare async generators can't carry the `_hidden_params` attribute the + proxy reads response headers off of (see + router_utils.add_retry_fallback_headers.get_hidden_params_dict), so this + thin wrapper carries it through from the source iterator - mirrors + AnthropicMessagesStreamingResponse. Used by + Router._aanthropic_messages_streaming_iterator. + """ + + def __init__(self, async_generator: AsyncGenerator[bytes, None], source_iterator: object) -> None: + self._async_generator = async_generator + self._hidden_params = dict( # mutable-ok: mutated in place by merge_fallback_hidden_params + getattr(source_iterator, "_hidden_params", None) or {} + ) + + def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream": + return self + + async def __anext__(self) -> bytes: + return await self._async_generator.__anext__() + + async def aclose(self) -> None: + await self._async_generator.aclose() + + def merge_fallback_hidden_params( + self, + fallback_hidden_params: Mapping[str, object], + fallback_headers: Mapping[str, object], + ) -> None: + """ + Raw bytes can't carry their own _hidden_params the way a + ModelResponseStream/ResponsesAPI event can, so a mid-stream + fallback's provider headers (e.g. Bedrock's x-amzn-requestid) are + merged onto the wrapper itself instead - mirrors + Router._apply_fallback_hidden_params_to_item's merge shape. + """ + existing_headers: Final = cast( # cast-ok: additional_headers is always a dict[str, object] when present + "dict[str, object]", self._hidden_params.get("additional_headers") or {} + ) + self._hidden_params = { # mutable-ok: matches _hidden_params' existing dict[str, object] shape + **self._hidden_params, + **fallback_hidden_params, + "additional_headers": {**existing_headers, **fallback_headers}, # mutable-ok: same shape + } + + class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) @@ -375,10 +481,17 @@ def _replay_live_router_model_cost() -> None: set_live_deployment_replay(_replay_live_router_model_cost) -# Kwargs that log_retry must not copy into a retry breadcrumb. The breadcrumbs reach spend -# logs and logging callbacks, and these carry either the request payload or router-internal -# walk state rather than anything that identifies the failed attempt. -RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(("messages", "original_function", "attempted_targets")) +# Kwargs that carry no signal about the failed attempt, so log_retry drops them from a +# breadcrumb entirely: the request payload and the router-internal walk state. Credentials are +# handled separately by mask_credentials_in_payload, which scrubs credential-named values from +# whatever kwargs remain rather than trying to enumerate every credential-bearing key here. +RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( + ( + "messages", + "original_function", + "attempted_targets", + ) +) class Router: @@ -4793,6 +4906,264 @@ class Router: ) return response + async def _aanthropic_messages_streaming_iterator( + self, + response: AsyncIterator[bytes], + initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain + ) -> AsyncIterator[bytes]: + """ + Wrap an anthropic_messages (/v1/messages) streaming response so a + mid-stream provider error triggers the Router's fallback chain + (parity with _acompletion_streaming_iterator for the + chat-completions path). See #24004. + + anthropic_messages goes through _ageneric_api_call_with_fallbacks + rather than _acompletion, so the returned byte iterator is never + wrapped by the chat-completions fallback handler. Two failure + shapes land here: + - the completion-bridge path (deployments with no native + /v1/messages endpoint, via + LiteLLMMessagesToCompletionTransformationHandler) already + raises MidStreamFallbackError out of its underlying + CustomStreamWrapper; this wrapper only needs to catch it. + - a native Anthropic/Bedrock passthrough never raises anything + for a provider SSE `event: error` frame (e.g. `overloaded_error`, + `internal_server_error`) - it is forwarded to the client as-is - + so this wrapper detects it via parse_anthropic_error_event and + raises MidStreamFallbackError itself. + + Only an error before any real content (a content_block_delta frame) + has reached the caller triggers a fallback attempt, mirroring the + restriction _acompletion_streaming_iterator applies: once generated + output has already reached the caller, retrying would start a + second, overlapping Anthropic message lifecycle on the same SSE + stream, so the error is left to propagate instead of being retried + invisibly. A non-retriable client error (4xx other than 429) is + never worth a fallback attempt either, so it is also left to + propagate. + + Lifecycle/bookkeeping frames (message_start, content_block_start, + ping, ...) do not by themselves disqualify a fallback attempt - + Anthropic routinely sends message_start before an overload error - + but they are BUFFERED rather than forwarded immediately, since + forwarding one and then appending a fallback attempt's own + message_start would produce two overlapping message lifecycles on + one SSE stream. Buffered frames are flushed, in order, the moment + real content arrives (the primary attempt has committed by then + anyway) or once the stream ends without ever producing content or + an error. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + aclose_if_supported, + parse_anthropic_error_event, + ) + + source_iterator: Final = response + + async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: + from litellm.exceptions import MidStreamFallbackError + + # Lifecycle/bookkeeping frames (message_start, content_block_start, + # ping, ...) are held back rather than forwarded immediately: + # Anthropic routinely sends message_start before an overload + # error, and once a byte reaches the client a fallback attempt + # can only append its OWN message_start, producing two + # overlapping message lifecycles on one SSE stream. Buffered + # frames are flushed the moment real content (content_block_delta) + # arrives - at that point the primary attempt has committed and a + # clean retry is no longer possible anyway - or once the primary + # stream ends without ever producing content. A `ping` keepalive + # is dropped outright rather than buffered, since it can recur + # indefinitely on a slow-starting connection and carries nothing + # worth preserving; hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + # forces the same early commit as real content arriving, so a + # hostile or pathological upstream can't grow the buffer forever. + has_generated_content = False # rebind-ok: set once real content is seen, or the buffer cap is hit + buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline + model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group + try: + async for chunk in source_iterator: + if _anthropic_stream_should_drop_pre_content_ping(chunk, has_generated_content): + continue + if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): + has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit + error_event = parse_anthropic_error_event(chunk) + retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over + not has_generated_content + and error_event is not None + and _is_retriable_anthropic_status(error_event[2]) + ) + if not has_generated_content and not retriable_pending_error and error_event is None: + buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) + continue + if retriable_pending_error: + assert error_event is not None # guard-ok: retriable_pending_error implies this + _error_type, message, status_code = error_event + raise MidStreamFallbackError( + message=message, + model=model, + llm_provider="anthropic", + original_exception=litellm.exceptions.APIError( + status_code=status_code, + message=message, + llm_provider="anthropic", + model=model, + ), + is_pre_first_chunk=True, + ) + for buffered_chunk in buffered_lifecycle_chunks: + yield buffered_chunk + buffered_lifecycle_chunks = () + yield chunk + for buffered_chunk in buffered_lifecycle_chunks: + yield buffered_chunk + except MidStreamFallbackError as e: + if _anthropic_stream_should_decline_fallback(has_generated_content, e): + for buffered_chunk in buffered_lifecycle_chunks: + yield buffered_chunk + if e.original_exception is not None: + raise e.original_exception from e + raise + async for item in self._aanthropic_messages_fallback_attempt(e, initial_kwargs, wrapper): + yield item + finally: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(source_iterator) + + # Referenced by stream_with_fallbacks via closure - assigned here, before + # the generator body ever runs, so the reference resolves fine despite + # being defined textually after the function that captures it. + wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator) + return wrapper + + async def _aanthropic_messages_fallback_attempt( + self, + e: "MidStreamFallbackError", + initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain + wrapper: "FallbackAwareAnthropicMessagesStream", + ) -> AsyncGenerator[bytes, None]: + """ + Re-enters the Router's fallback chain for a mid-stream + anthropic_messages error and yields whatever the fallback attempt + produces. Split out of _aanthropic_messages_streaming_iterator to + keep each function's cyclomatic complexity within the repo's C901 + budget. + """ + from litellm.exceptions import MidStreamFallbackError + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + aclose_if_supported, + anthropic_messages_response_as_sse_events, + ) + + fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted + try: + model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group + fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param + "fallbacks", self.fallbacks + ) + context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below + "context_window_fallbacks", self.context_window_fallbacks + ) + content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below + "content_policy_fallbacks", self.content_policy_fallbacks + ) + initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper + self._update_kwargs_before_fallbacks( + model=model_group, + kwargs=initial_kwargs, + metadata_variable_name="litellm_metadata", + ) + fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success + e=e, + disable_fallbacks=False, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + model_group=model_group, + args=(), + kwargs=initial_kwargs, + include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, + ) + fallback_hidden_params, fallback_headers = Router._prepare_fallback_hidden_params(fallback_response) + wrapper.merge_fallback_hidden_params(fallback_hidden_params, fallback_headers) + if hasattr(fallback_response, "__aiter__"): + async for fallback_item in fallback_response: + yield fallback_item + else: + # A fallback can resolve to a complete AnthropicMessagesResponse + # dict even for a streaming request (e.g. an agentic tool-use + # interception loop) - yielding it as-is would put a raw dict + # into a byte stream, so it's synthesized into the SSE + # lifecycle a real stream would have sent instead. + for event in anthropic_messages_response_as_sse_events( + cast("AnthropicMessagesResponse", fallback_response) # cast-ok: non-streaming shape by elimination + ): + yield event + except Exception as fallback_error: + verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error) + if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: + raise fallback_error.original_exception from fallback_error + raise + finally: + if fallback_response is not None: + with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): + await aclose_if_supported(fallback_response) + + async def _aanthropic_messages_with_streaming_fallbacks( + self, + original_function: Callable, + **kwargs: object, # kwargs-ok: forwarded verbatim to original_function, shape varies per call site + ) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]: + """ + _ageneric_api_call_with_fallbacks for anthropic_messages, with the + addition of mid-stream fallback handling (see + _aanthropic_messages_streaming_iterator). Parity with + _aresponses_with_streaming_fallbacks for the Responses API. + """ + from litellm.litellm_core_utils.core_helpers import safe_deep_copy + + # Snapshot the request kwargs before the primary attempt mutates them + # in place: _update_kwargs_with_deployment writes deployment-specific + # fields (deployment, model_info, api_base, tags, ...) into the + # SAME litellm_metadata/metadata dicts a shallow .copy() would still + # share, leaking primary-deployment metadata into the mid-stream + # fallback request. safe_deep_copy avoids deep-copying the full + # kwargs (which can hold non-deepcopyable logging handles/clients). + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry + if isinstance(fallback_kwargs.get("litellm_metadata"), dict): + fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) + if isinstance(fallback_kwargs.get("metadata"), dict): + fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"]) + fallback_kwargs["original_generic_function"] = original_function + + response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + + if kwargs.get("stream") and hasattr(response, "__aiter__"): + return await self._aanthropic_messages_streaming_iterator( + response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator + initial_kwargs=fallback_kwargs, + ) + return response + + async def _dispatch_generic_call_type( + self, + call_type: str, + original_function: Callable, + **kwargs: object, # kwargs-ok: forwarded verbatim to the per-call-type helper, shape varies per call site + ): + """ + factory_function's shared dispatch for call types with no + call-specific handling, except anthropic_messages: kept out of + factory_function's own async_wrapper (already at the repo's C901 + complexity ceiling) so routing its mid-stream fallback handling + (#24004) doesn't add another branch there. + """ + if call_type == "anthropic_messages": + return await self._aanthropic_messages_with_streaming_fallbacks( + original_function=original_function, **kwargs + ) + return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): """ Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router @@ -5979,7 +6350,8 @@ class Router: "aget_skill", "adelete_skill", ): - return await self._ageneric_api_call_with_fallbacks( + return await self._dispatch_generic_call_type( + call_type=call_type, original_function=original_function, **kwargs, ) @@ -7347,7 +7719,8 @@ class Router: if len(self.previous_models) > 3: self.previous_models.pop(0) - self.previous_models.append(previous_model) + scrubbed_previous_model: Final = mask_credentials_in_payload(previous_model) + self.previous_models.append(scrubbed_previous_model) kwargs[_metadata_var]["previous_models"] = self.previous_models return kwargs except Exception as e: @@ -8332,6 +8705,7 @@ class Router: ) = litellm.get_llm_provider( model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), + api_base=deployment.litellm_params.api_base, ) # done reading model["litellm_params"] # Check if provider is supported: either in enum or JSON-configured @@ -9439,6 +9813,8 @@ class Router: except Exception: model_info = None + deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict) + # get llm provider litellm_model, llm_provider = "", "" try: @@ -9481,6 +9857,7 @@ class Router: "model_group": user_facing_model_group_name, "providers": [llm_provider], **model_info, + "supported_reasoning_efforts": None, } ) else: @@ -9558,6 +9935,11 @@ class Router: if model_info.get("rpm", None) is not None and _deployment_rpm is None: _deployment_rpm = model_info.get("rpm") + model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( + model_group_info.supported_reasoning_efforts, + resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=deployment_is_mapped), + ) + if _deployment_tpm is not None: if total_tpm is None: total_tpm = 0 diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index cbaba69f696..087c1f7278d 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -19,7 +19,7 @@ import asyncio import random import re from collections.abc import Iterator, Mapping, Sequence -from itertools import accumulate, islice +from itertools import accumulate, islice, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast @@ -274,6 +274,8 @@ _REMINDER_CLOSE: Final = "" _DEFAULT_REMINDER_MARKERS: Final = ((_REMINDER_OPEN, _REMINDER_CLOSE),) _TRUNCATION_MARKER: Final = "..." +_TRUNCATION_HEAD_FRACTION: Final = 0.3 +_MIN_QUOTED_TURN_CHARS: Final = 120 _CJK_CHARACTER: Final = re.compile("[぀-ヿㇰ-ㇿ㐀-䶿一-鿿豈-﫿ｦ-ﾝ\U00020000-\U0003ffff]") @@ -552,8 +554,21 @@ def _matched_plan_mode_sentinel( def _truncate(text: str, limit: int) -> str: - """Cap text at limit characters, marking it so the classifier can tell the turn was cut short.""" - return text if len(text) <= limit else f"{text[:limit]}{_TRUNCATION_MARKER}" + """Cap text at limit characters, keeping both ends and eliding the middle. + + A chat turn states its ask at the end, so cutting the tail keeps the preamble and discards the + request the turn exists to make: a turn opening with an incident report and closing with "rewrite + the retry path and prove it cannot livelock" reached the classifier as the incident report alone. + Keeping both ends costs nothing at the same budget and is what the truncation literature finds + best for classifying long text, head+tail measuring above both head-only and tail-only in Sun et + al. 2019. The marker sits at the cut, so the turn reads as having its middle removed rather than + as trailing off mid-thought. + """ + if len(text) <= limit: + return text + head_chars: Final = max(int(limit * _TRUNCATION_HEAD_FRACTION), 0) + tail_chars: Final = max(limit - head_chars, 0) + return f"{text[:head_chars]}{_TRUNCATION_MARKER}{text[len(text) - tail_chars :]}" def _iter_context_turns_newest_first( @@ -579,11 +594,40 @@ def _iter_context_turns_newest_first( ) +def _turns_within_budget( + turns: Sequence[tuple[str, str]], + budget_chars: int, +) -> tuple[tuple[str, str], ...]: + """The newest-first turns that fit budget_chars, quoted whole wherever they fit. + + Bounding the block rather than every turn in it is what lets an ordinary conversation reach the + classifier intact: a per-turn cap cuts a 785 character turn even when the whole block would have + been 353 characters, which is three orders of magnitude below anything the classifier call is + near. Once the budget does run out the older turns are dropped entire rather than shortened, so + at most one turn is ever cut and the rest read as themselves. A remainder too small to carry a + sentence buys less signal than the ellipses it would arrive wrapped in, so that turn is dropped. + + The boundary turn is cut to leave room for the marker rather than to the remainder itself, so the + quoted block never exceeds budget_chars; the marker is part of what the budget buys, not an extra + charged on top of it. + """ + spent: Final = accumulate(len(text) for _, text in turns) + fitting: Final = tuple(takewhile(lambda pair: pair[1] <= budget_chars, zip(turns, spent))) + remaining: Final = budget_chars - (fitting[-1][1] if fitting else 0) + whole: Final = tuple(turn for turn, _ in fitting) + cut_to: Final = remaining - len(_TRUNCATION_MARKER) + if len(whole) == len(turns) or cut_to < _MIN_QUOTED_TURN_CHARS: + return whole + boundary_role, boundary_text = turns[len(whole)] + return (*whole, (boundary_role, _truncate(boundary_text, cut_to))) + + def _extract_prior_turns( messages: Sequence[Mapping[str, object]], current_ask: str | None, window_size: int, - per_turn_chars: int, + budget_chars: int, + per_turn_chars: int | None, include_assistant: bool, marker_pairs: tuple[tuple[str, str], ...] = _DEFAULT_REMINDER_MARKERS, ) -> tuple[tuple[str, str], ...]: @@ -598,19 +642,29 @@ def _extract_prior_turns( window_size counts turns of every eligible role, so with assistant turns included it is the last N of the conversation rather than the last N asks. A turn carrying only tool calls or thinking blocks flattens to empty text and is skipped, so it never spends a slot. + + Three bounds apply and the tightest wins: window_size caps how many turns, budget_chars caps the + block they form, and per_turn_chars optionally caps any single one of them before the block is + measured. They are separate because they answer separate questions, and only the block bound + tracks what the classifier call actually costs. """ if window_size <= 0 or not messages: return () - prior: Final = islice( - ( - turn - for turn in _iter_context_turns_newest_first(messages, include_assistant, marker_pairs) - if turn[1] != current_ask - ), - window_size, + prior: Final = tuple( + islice( + ( + turn + for turn in _iter_context_turns_newest_first(messages, include_assistant, marker_pairs) + if turn[1] != current_ask + ), + window_size, + ) ) - return tuple((role, _truncate(text, per_turn_chars)) for role, text in reversed(tuple(prior))) + clamped: Final = ( + prior if per_turn_chars is None else tuple((role, _truncate(text, per_turn_chars)) for role, text in prior) + ) + return tuple(reversed(_turns_within_budget(clamped, budget_chars))) def _decision_is_pinnable(decision: StandardLoggingRoutingDecision | None) -> bool: @@ -1349,6 +1403,7 @@ class ComplexityRouter(CustomLogger): messages, current_ask=prompt, window_size=self.config.classifier_context_window_size, + budget_chars=self.config.classifier_context_budget_chars, per_turn_chars=self.config.classifier_context_per_turn_chars, include_assistant=include_assistant, marker_pairs=self._reminder_markers, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index d3c4bd7938b..9907407d84d 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -49,7 +49,7 @@ TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = ( DEFAULT_TIER_DISTANCE_PENALTY: Final[float] = 0.5 DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE: Final[int] = 3 -DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS: Final[int] = 200 +DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS: Final[int] = 8000 class KeywordTierRule(BaseModel): @@ -645,12 +645,30 @@ class ComplexityRouterConfig(BaseModel): "classifier_type is 'llm'." ), ) - classifier_context_per_turn_chars: int = Field( - default=DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS, + classifier_context_budget_chars: int = Field( + default=DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, + ge=0, + description=( + "Maximum characters of prior-turn text quoted to the LLM classifier, across the whole " + "context window, per classification call. Turns are taken newest first and quoted whole " + "while they fit, so a conversation small enough to quote entirely is never cut; once the " + "budget runs out the older turns are dropped whole and only the turn straddling the " + "boundary is truncated, into whatever space is left. The current ask and the caller's " + "system prompt sit outside this budget and are always sent in full, as does the numbering " + "each quoted turn carries. A budget under 120 leaves no room to quote a turn and " + "suppresses the block; set classifier_context_window_size to 0 to turn context off " + "deliberately. Only applies when classifier_type is 'llm'." + ), + ) + classifier_context_per_turn_chars: int | None = Field( + default=None, gt=0, description=( - "Maximum character length for each prior turn's text in the classifier context window. " - "Turns exceeding this are truncated. Only applies when classifier_type is 'llm'." + "Optional cap on each individual prior turn's text, applied before " + "classifier_context_budget_chars bounds the block. Unset by default, so one long turn may " + "spend the whole budget, which is usually what a follow-up needs; set it when no single " + "turn should dominate the context the classifier sees. A capped turn keeps its opening " + "and its ending with the middle elided. Only applies when classifier_type is 'llm'." ), ) classifier_context_include_assistant_turns: bool = Field( @@ -662,9 +680,9 @@ class ComplexityRouterConfig(BaseModel): "word 'yes'. When enabled, classifier_context_window_size counts the last N turns of the " "conversation across both roles rather than the last N user turns, and assistant text is " "sent to the classifier model, which may be a different deployment or provider than the " - "routed completion model. Assistant replies share classifier_context_per_turn_chars with " - "user turns, so raise it if replies are truncated before the part that carries the " - "difficulty. Off by default because enabling it shifts tier decisions, and therefore " + "routed completion model. Assistant replies spend classifier_context_budget_chars " + "alongside user turns, so raise it if the oldest turns stop being quoted once replies " + "join the window. Off by default because enabling it shifts tier decisions, and therefore " "spend, for an already-deployed router. Only applies when classifier_type is 'llm'." ), ) diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py new file mode 100644 index 00000000000..3e4478e5e21 --- /dev/null +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -0,0 +1,146 @@ +"""Resolve which reasoning_effort values a deployment, and by intersection a model group, accepts. + +The model map's supports_*_reasoning_effort flags are the only signal, and each level's polarity +mirrors how a request path reads that same flag. medium and high are unconditional for a reasoning +model. minimal and low are opt-out: openai/chat/gpt_5_transformation.py refuses them only when the +map says false. xhigh and max are opt-in. none is opt-out everywhere except the azure gpt-5 family, +whose config raises UnsupportedParamsError without an explicit true. + +xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at +all: every entry carrying supports_max_reasoning_effort is Claude-family, and +anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort +path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an +explicit flag is the only signal that the tier is a real one rather than litellm rounding the level +to a budget, and a missing flag costs advisory metadata rather than a rejected request. + +A deployment the map describes with no effort flags at all resolves to None rather than to the +opt-out defaults. 689 of the map's 854 reasoning entries carry no flag, and the o-series, xai and +bedrock nova entries among them take neither none nor minimal, so composing a set out of the +defaults alone would advertise levels those providers reject. + +The advertisement order is the REASONING_EFFORT declaration order, which is presentation only. It +is not a strength scale and does not reconcile with bedrock's output_config ceiling order in +llms/bedrock/common_utils.py, which ranks max below xhigh while the thinking-budget constants rank +it above. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, get_args + +import litellm +from litellm.types.llms.openai import REASONING_EFFORT + +REASONING_EFFORT_ADVERTISEMENT_ORDER: Final = get_args(REASONING_EFFORT) +_EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) + +_EFFORT_FLAGS: Final = ( + ("none", "supports_none_reasoning_effort"), + ("minimal", "supports_minimal_reasoning_effort"), + ("low", "supports_low_reasoning_effort"), + ("xhigh", "supports_xhigh_reasoning_effort"), + ("max", "supports_max_reasoning_effort"), +) +_OPT_OUT_EFFORTS: Final = ("minimal", "low") +_OPT_IN_EFFORTS: Final = ("xhigh", "max") +_UNCONDITIONAL_EFFORTS: Final = frozenset(("medium", "high")) + + +def _bare_model_entry(model_info: Mapping[str, object]) -> Mapping[str, object]: + """The unprefixed twin of a provider-prefixed map entry, which is where the flags often live: + azure/gpt-5-mini carries none of them while gpt-5-mini carries all three. The request-path + gates resolve through the same twin (_supports_factory, #20885), so reading it here is what + keeps the advertisement and the gate on the same answer.""" + key: Final = model_info.get("key") + provider: Final = model_info.get("litellm_provider") + if not isinstance(key, str) or not isinstance(provider, str) or not key.startswith(f"{provider}/"): + return _EMPTY_ENTRY + entry: Final[Mapping[str, object] | None] = litellm.model_cost.get(key.removeprefix(f"{provider}/")) + return entry if entry is not None else _EMPTY_ENTRY + + +def _declared_effort_flags(model_info: Mapping[str, object]) -> Mapping[str, object]: + bare: Final = _bare_model_entry(model_info) + return MappingProxyType( + { + effort: model_info.get(flag) if model_info.get(flag) is not None else bare.get(flag) + for effort, flag in _EFFORT_FLAGS + } + ) + + +def _supports_none_reasoning_effort(model_info: Mapping[str, object], flag: object) -> bool: + """Opt-in only where a request path refuses the level. AzureOpenAIGPT5Config raises + UnsupportedParamsError on reasoning_effort='none' without an explicit true, and it is selected + only for the gpt-5 family, so every other azure deployment keeps the opt-out default.""" + if model_info.get("litellm_provider") != "azure": + return flag is not False + + from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + + key: Final = model_info.get("key") + if not isinstance(key, str) or not AzureOpenAIGPT5Config.is_model_gpt_5_model(key): + return flag is not False + return flag is True + + +def deployment_is_catalog_mapped( + resolved_model_info: Mapping[str, object] | None, + operator_model_info: Mapping[str, object], +) -> bool: + """Whether the model map described this deployment, as opposed to the operator describing it. + + Every deployment is registered in the cost map under its own id, so a mode the operator wrote + on an off-map deployment reads back here exactly like one the catalog supplied. Excluding it is + what stops such a deployment from claiming to be a known non-reasoning model and emptying the + levels its mapped siblings agree on. + """ + if resolved_model_info is None or resolved_model_info.get("mode") is None: + return False + return operator_model_info.get("mode") is None + + +def resolve_supported_reasoning_efforts( + model_info: Mapping[str, object], + *, + deployment_is_mapped: bool, +) -> tuple[str, ...] | None: + """None = nothing is known about this deployment, so it must not narrow its group; () = a known + model that accepts no effort level, which correctly empties the group. + + Telling those apart needs provenance the flattened ModelInfo does not carry. A deployment the + map does not describe arrives with supports_reasoning None, exactly like a mapped non-reasoning + model: 2273 of the map's 3165 entries omit the key rather than setting it false, so reading an + unset flag as () would let one custom deployment empty every level its mapped siblings agree + on. deployment_is_mapped is that provenance, and an operator who wants either answer for an + off-map deployment gets it by setting supports_reasoning explicitly. + """ + supports_reasoning: Final = model_info.get("supports_reasoning") + if supports_reasoning is not True: + return () if supports_reasoning is False or deployment_is_mapped else None + + flags: Final = _declared_effort_flags(model_info) + if all(value is None for value in flags.values()): + return None + + opt_out: Final = frozenset(effort for effort in _OPT_OUT_EFFORTS if flags[effort] is not False) + opt_in: Final = frozenset(effort for effort in _OPT_IN_EFFORTS if flags[effort] is True) + none_level: Final = ( + frozenset(("none",)) if _supports_none_reasoning_effort(model_info, flags["none"]) else frozenset() + ) + allowed: Final = opt_out | _UNCONDITIONAL_EFFORTS | opt_in | none_level + return tuple(effort for effort in REASONING_EFFORT_ADVERTISEMENT_ORDER if effort in allowed) + + +def intersect_supported_reasoning_efforts( + current: Sequence[str] | None, + resolved: Sequence[str] | None, +) -> tuple[str, ...] | None: + """Deployments without metadata (None) never narrow the group; an effort survives only when + every deployment with metadata accepts it, so the group offers nothing routing could reject.""" + if resolved is None: + return tuple(current) if current is not None else None + if current is None: + return tuple(resolved) + keep: Final = frozenset(current) & frozenset(resolved) + return tuple(effort for effort in REASONING_EFFORT_ADVERTISEMENT_ORDER if effort in keep) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index cc6eccbf3e0..d3b0f334163 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -685,6 +685,7 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai class AnthropicThinkingParam(TypedDict, total=False): type: ReadOnly[Literal["enabled", "adaptive", "disabled"]] budget_tokens: int + display: ReadOnly[Literal["summarized", "omitted"]] class ANTHROPIC_HOSTED_TOOLS(str, Enum): diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5665aa3277a..6ae2e31fe60 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from collections.abc import Sequence from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal @@ -396,7 +397,7 @@ class OutputConfigBlock(TypedDict, total=False): class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows additionalModelRequestFields: dict - additionalModelResponseFieldPaths: list[str] + additionalModelResponseFieldPaths: Sequence[str] inferenceConfig: InferenceConfig system: list[SystemContentBlock] toolConfig: ToolConfigBlock diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e7a3f825455..4a6c4a5bbb5 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1840,7 +1840,7 @@ ResponsesAPIStreamingResponse = Annotated[ ] -REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh"] +REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh", "max"] class OpenAIRealtimeStreamSession(TypedDict, total=False): diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index e2469d4c78f..a88ffeec6b5 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -158,9 +158,18 @@ class AutoRouterBenchmarksResponse(BaseModel): start_date: str = Field(description="Window start day, YYYY-MM-DD UTC, inclusive") end_date: str = Field(description="Window end day, YYYY-MM-DD UTC, inclusive") - routers_in_scope: int + routers_in_scope: int = Field( + description="How many groups this response carries. Every auto-router configured on the " + "proxy counts, whether or not it served anything in the window. To count only the routers " + "that did serve traffic, filter `groups` to the entries whose `sessions` is above zero" + ) totals: AutoRouterBenchmarkTotals - groups: tuple[AutoRouterBenchmarkGroup, ...] + groups: tuple[AutoRouterBenchmarkGroup, ...] = Field( + description="One entry per auto-router, listed from the model registry rather than from " + "the rollup, so a router appears as soon as it is configured and reads zero until it " + "serves traffic. Semantic auto-routers are absent: they record no routing decision, so no " + "session can ever be attributed to them" + ) ShadowEvalStatus: TypeAlias = Literal["running", "completed", "stopped"] diff --git a/litellm/types/router.py b/litellm/types/router.py index 99a4603ae49..d4c735387a5 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -10,7 +10,7 @@ from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hin import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator -from typing_extensions import Protocol, Required, TypedDict, runtime_checkable +from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable from litellm._uuid import uuid @@ -480,7 +480,9 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): output_cost_per_token: float | None input_cost_per_second: float | None output_cost_per_second: float | None + output_cost_per_second_480p: ReadOnly[float | None] output_cost_per_second_1080p: float | None + output_cost_per_second_4k: ReadOnly[float | None] num_retries: int | None ## MOCK RESPONSES ## mock_response: str | ModelResponse | Exception | None @@ -635,6 +637,7 @@ class ModelGroupInfo(BaseModel): supports_url_context: bool = Field(default=False) supports_reasoning: bool = Field(default=False) supports_function_calling: bool = Field(default=False) + supported_reasoning_efforts: tuple[str, ...] | None = Field(default=None) supported_openai_params: list[str] | None = Field(default=[]) configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a5beb7728e5..96402116a89 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -154,6 +154,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_web_search: bool | None supports_reasoning: bool | None supports_adaptive_thinking: bool | None + supports_legacy_thinking: ReadOnly[bool | None] thinking_always_on: ReadOnly[bool | None] supports_tool_search: bool | None supports_mid_conversation_system: bool | None @@ -277,6 +278,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_second_1080p: ( float | None ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) + output_cost_per_second_480p: ReadOnly[float | None] + output_cost_per_second_4k: ReadOnly[float | None] ocr_cost_per_page: float | None # for OCR models ocr_cost_per_credit: float | None # for OCR models priced by credit annotation_cost_per_page: float | None # for OCR models @@ -440,6 +443,12 @@ class CallTypes(str, Enum): query = "query" aquery = "aquery" + ######################################################### + # Google Interactions API Call Types + ######################################################### + create_interaction = "create_interaction" + acreate_interaction = "acreate_interaction" + ######################################################### # Container Call Types ######################################################### @@ -3332,6 +3341,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_second: float | None = None output_cost_per_second: float | None = None output_cost_per_second_1080p: float | None = None + output_cost_per_second_480p: float | None = None + output_cost_per_second_4k: float | None = None input_cost_per_pixel: float | None = None output_cost_per_pixel: float | None = None @@ -3845,6 +3856,7 @@ class SearchProviders(str, Enum): TINYFISH = "tinyfish" AGENTCORE = "agentcore" NIMBLE = "nimble" + BING_GROUNDING = "bing_grounding" # Create a set of all search provider values for quick lookup diff --git a/litellm/utils.py b/litellm/utils.py index e5ce7157e77..59a2de609d7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4130,7 +4130,7 @@ def get_optional_params( drop_params=(drop_params if drop_params is not None and isinstance(drop_params, bool) else False), ) elif custom_llm_provider == "together_ai": - optional_params = litellm.TogetherAIConfig().map_openai_params( + optional_params = litellm.TogetherAIChatConfig().map_openai_params( non_default_params=non_default_params, optional_params=optional_params, model=model, @@ -5726,6 +5726,8 @@ def _get_model_info_helper( ), output_cost_per_second=_model_info.get("output_cost_per_second", None), output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), + output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), + output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), @@ -5753,6 +5755,7 @@ def _get_model_info_helper( supports_url_context=_model_info.get("supports_url_context", None), supports_reasoning=_model_info.get("supports_reasoning", None), supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None), + supports_legacy_thinking=_model_info.get("supports_legacy_thinking", None), thinking_always_on=_model_info.get("thinking_always_on", None), supports_tool_search=_model_info.get("supports_tool_search", None), supports_mid_conversation_system=_model_info.get("supports_mid_conversation_system", None), @@ -6525,24 +6528,6 @@ def acreate(*args, **kwargs): ## Thin client to handle the acreate langchain ca return litellm.acompletion(*args, **kwargs) -def prompt_token_calculator(model, messages): - # use tiktoken or anthropic's tokenizer depending on the model - text: Final = " ".join(message["content"] for message in messages) - num_tokens = 0 - if "claude" in model: - try: - import anthropic - except Exception: - Exception("Anthropic import failed please run `pip install anthropic`") - from anthropic import AI_PROMPT, HUMAN_PROMPT, Anthropic - - anthropic_obj: Final = Anthropic() - num_tokens = anthropic_obj.count_tokens(text) - else: - num_tokens = len(_get_default_encoding().encode(text)) - return num_tokens - - def valid_model(model): try: # for a given model name, check if the user has the right permissions to access the model @@ -7913,7 +7898,7 @@ class ProviderConfigManager: LlmProviders.GALADRIEL: (lambda: litellm.GaladrielChatConfig(), False), LlmProviders.REPLICATE: (lambda: litellm.ReplicateConfig(), False), LlmProviders.HUGGINGFACE: (lambda: litellm.HuggingFaceChatConfig(), False), - LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIConfig(), False), + LlmProviders.TOGETHER_AI: (lambda: litellm.TogetherAIChatConfig(), False), LlmProviders.OPENROUTER: (lambda: litellm.OpenrouterConfig(), False), LlmProviders.VERCEL_AI_GATEWAY: ( lambda: litellm.VercelAIGatewayConfig(), @@ -8625,6 +8610,12 @@ class ProviderConfigManager: ) return BedrockPassthroughConfig() + elif LlmProviders.BEDROCK_MANTLE == provider: + from litellm.llms.bedrock_mantle.passthrough.transformation import ( + BedrockMantlePassthroughConfig, + ) + + return BedrockMantlePassthroughConfig() elif LlmProviders.VLLM == provider or LlmProviders.HOSTED_VLLM == provider: from litellm.llms.vllm.passthrough.transformation import ( VLLMPassthroughConfig, @@ -9111,6 +9102,7 @@ class ProviderConfigManager: from litellm.llms.apiserpent.search.transformation import ( APISerpentSearchConfig, ) + from litellm.llms.azure.search.transformation import BingGroundingSearchConfig from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig @@ -9152,6 +9144,7 @@ class ProviderConfigManager: SearchProviders.TINYFISH: TinyfishSearchConfig, SearchProviders.AGENTCORE: AgentCoreSearchConfig, SearchProviders.NIMBLE: NimbleSearchConfig, + SearchProviders.BING_GROUNDING: BingGroundingSearchConfig, } config_class: Final = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 978849ac006..445435a30fa 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -5,6 +5,8 @@ from collections.abc import Coroutine from functools import partial from typing import Final, Literal, overload +from httpx._types import FileContent + import litellm from litellm.constants import DEFAULT_VIDEO_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT @@ -1344,6 +1346,8 @@ async def avideo_edit( extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, + *, + video: FileContent | None = None, **kwargs, ) -> VideoObject: """ @@ -1359,6 +1363,7 @@ async def avideo_edit( video_edit, video_id=video_id, prompt=prompt, + video=video, timeout=timeout, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, @@ -1396,6 +1401,8 @@ def video_edit( extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, + *, + video: FileContent | None = None, **kwargs, ) -> VideoObject | Coroutine[object, object, VideoObject]: """ @@ -1444,6 +1451,7 @@ def video_edit( return base_llm_http_handler.video_edit_handler( prompt=prompt, video_id=video_id, + video_file=video, video_provider_config=provider_config, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3af7d9e5019..9ca8d9e1bac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1019,6 +1019,7 @@ }, "anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1053,6 +1054,7 @@ }, "global.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -1087,6 +1089,7 @@ }, "us.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1121,6 +1124,7 @@ }, "eu.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -1155,6 +1159,7 @@ }, "au.anthropic.claude-opus-4-6-v1": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.875e-06, "cache_creation_input_token_cost_above_1hr": 1.1e-05, "cache_read_input_token_cost": 5.5e-07, @@ -2233,6 +2238,7 @@ }, "anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -2266,6 +2272,7 @@ }, "global.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -2299,6 +2306,7 @@ }, "us.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2332,6 +2340,7 @@ }, "eu.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2365,6 +2374,7 @@ }, "au.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2398,6 +2408,7 @@ }, "jp.anthropic.claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 4.125e-06, "cache_creation_input_token_cost_above_1hr": 6.6e-06, "cache_read_input_token_cost": 3.3e-07, @@ -2950,6 +2961,7 @@ "azure_ai/claude-opus-4-6": { "deprecation_date": "2027-02-02", "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "input_cost_per_token": 5e-06, "output_cost_per_token": 2.5e-05, "litellm_provider": "azure_ai", @@ -3181,6 +3193,7 @@ "azure_ai/claude-sonnet-4-6": { "deprecation_date": "2027-02-10", "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12489,6 +12502,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "supports_assistant_prefill": true, "supports_computer_use": true, "supports_function_calling": true, @@ -12698,6 +12712,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -12735,6 +12750,7 @@ "search_context_size_medium": 0.01 }, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, @@ -14551,6 +14567,8 @@ ] }, "databricks/databricks-bge-large-en": { + "cache_creation_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.0003e-07, "input_dbu_cost_per_token": 1.429e-06, "litellm_provider": "databricks", @@ -14566,6 +14584,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-3-7-sonnet": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14581,10 +14601,41 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-fable-5": { + "cache_creation_input_token_cost": 1.250004e-05, + "cache_read_input_token_cost": 1.00002e-06, + "input_cost_per_token": 1.000006e-05, + "input_dbu_cost_per_token": 0.000142858, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.000002e-05, + "output_dbu_cost_per_token": 0.000714286, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": false, + "thinking_always_on": true + }, "databricks/databricks-claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -14600,10 +14651,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4": { + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14619,10 +14673,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-1": { + "cache_creation_input_token_cost": 1.874999e-05, + "cache_read_input_token_cost": 1.50003e-06, "input_cost_per_token": 1.5000020000000002e-05, "input_dbu_cost_per_token": 0.000214286, "litellm_provider": "databricks", @@ -14638,10 +14695,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14657,11 +14717,14 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_output_config": true }, "databricks/databricks-claude-opus-4-6": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -14677,10 +14740,94 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_legacy_thinking": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-opus-4-7": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 2048, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-4-8": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-claude-opus-5": { + "cache_creation_input_token_cost": 6.25002e-06, + "cache_read_input_token_cost": 5.0001e-07, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.000357143, + "prompt_cache_min_tokens": 512, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-claude-sonnet-4": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14696,10 +14843,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-1": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14715,10 +14865,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14734,10 +14887,13 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, "databricks/databricks-claude-sonnet-4-6": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, "input_cost_per_token": 2.9999900000000002e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", @@ -14753,10 +14909,41 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, + "supports_legacy_thinking": true, + "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-sonnet-5": { + "cache_creation_input_token_cost": 3.74997e-06, + "cache_read_input_token_cost": 3.0002e-07, + "input_cost_per_token": 2.99999e-06, + "input_dbu_cost_per_token": 4.2857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + }, + "mode": "chat", + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.000214286, + "prompt_cache_min_tokens": 1024, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true + }, "databricks/databricks-gemini-2-5-flash": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 3.0002e-08, "input_cost_per_token": 3.0001999999999996e-07, "input_dbu_cost_per_token": 4.285999999999999e-06, "litellm_provider": "databricks", @@ -14771,9 +14958,12 @@ "output_dbu_cost_per_token": 3.5714e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-2-5-pro": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.24999e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14788,9 +14978,12 @@ "output_dbu_cost_per_token": 0.000142857, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-flash-lite": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.122e-08, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -14805,9 +14998,12 @@ "output_dbu_cost_per_token": 2.6786e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-1-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14822,9 +15018,12 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-flash": { + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.251e-08, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -14839,9 +15038,12 @@ "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemini-3-pro": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14856,9 +15058,12 @@ "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, + "supports_prompt_caching": true, "supports_tool_choice": true }, "databricks/databricks-gemma-3-12b": { + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -14874,6 +15079,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-5": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14886,9 +15093,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14901,9 +15111,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-max": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_read_input_token_cost": 1.2502e-07, "input_cost_per_token": 1.24999e-06, "input_dbu_cost_per_token": 1.7857e-05, "litellm_provider": "databricks", @@ -14916,9 +15129,12 @@ "mode": "chat", "output_cost_per_token": 9.999990000000002e-06, "output_dbu_cost_per_token": 0.000142857, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-1-codex-mini": { + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -14931,9 +15147,12 @@ "mode": "chat", "output_cost_per_token": 1.99997e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2": { + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14946,9 +15165,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-2-codex": { + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14961,9 +15183,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-3-codex": { + "cache_creation_input_token_cost": 1.75e-06, + "cache_read_input_token_cost": 1.75e-07, "input_cost_per_token": 1.75e-06, "input_dbu_cost_per_token": 2.5e-05, "litellm_provider": "databricks", @@ -14976,9 +15201,12 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_dbu_cost_per_token": 0.0002, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.4997e-07, "input_cost_per_token": 2.49998e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", @@ -14991,9 +15219,12 @@ "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-mini": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, "input_cost_per_token": 7.4998e-07, "input_dbu_cost_per_token": 1.0714e-05, "litellm_provider": "databricks", @@ -15006,9 +15237,12 @@ "mode": "chat", "output_cost_per_token": 4.50002e-06, "output_dbu_cost_per_token": 6.4286e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-4-nano": { + "cache_creation_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost": 2.002e-08, "input_cost_per_token": 1.9999e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", @@ -15021,9 +15255,12 @@ "mode": "chat", "output_cost_per_token": 1.24999e-06, "output_dbu_cost_per_token": 1.7857e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-mini": { + "cache_creation_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost": 2.499e-08, "input_cost_per_token": 2.4997000000000006e-07, "input_dbu_cost_per_token": 3.571e-06, "litellm_provider": "databricks", @@ -15036,9 +15273,12 @@ "mode": "chat", "output_cost_per_token": 1.9999700000000004e-06, "output_dbu_cost_per_token": 2.8571e-05, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-5-nano": { + "cache_creation_input_token_cost": 4.998e-08, + "cache_read_input_token_cost": 4.97e-09, "input_cost_per_token": 4.998e-08, "input_dbu_cost_per_token": 7.14e-07, "litellm_provider": "databricks", @@ -15051,9 +15291,12 @@ "mode": "chat", "output_cost_per_token": 3.9998000000000007e-07, "output_dbu_cost_per_token": 5.714000000000001e-06, - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_prompt_caching": true }, "databricks/databricks-gpt-oss-120b": { + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15069,6 +15312,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gpt-oss-20b": { + "cache_creation_input_token_cost": 7e-08, + "cache_read_input_token_cost": 7e-08, "input_cost_per_token": 7e-08, "input_dbu_cost_per_token": 1e-06, "litellm_provider": "databricks", @@ -15084,6 +15329,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-gte-large-en": { + "cache_creation_input_token_cost": 1.2999e-07, + "cache_read_input_token_cost": 1.2999e-07, "input_cost_per_token": 1.2999000000000001e-07, "input_dbu_cost_per_token": 1.857e-06, "litellm_provider": "databricks", @@ -15099,6 +15346,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-llama-2-70b-chat": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15115,6 +15364,8 @@ "supports_tool_choice": true }, "databricks/databricks-llama-4-maverick": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15131,6 +15382,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-405b-instruct": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 5.00003e-06, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", @@ -15147,6 +15400,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-8b-instruct": { + "cache_creation_input_token_cost": 1.5001e-07, + "cache_read_input_token_cost": 1.5001e-07, "input_cost_per_token": 1.5000999999999998e-07, "input_dbu_cost_per_token": 2.1429999999999996e-06, "litellm_provider": "databricks", @@ -15162,6 +15417,8 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-meta-llama-3-3-70b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15178,6 +15435,8 @@ "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-70b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15194,6 +15453,8 @@ "supports_tool_choice": true }, "databricks/databricks-mixtral-8x7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -15210,6 +15471,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-30b-instruct": { + "cache_creation_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost": 1.00002e-06, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", @@ -15226,6 +15489,8 @@ "supports_tool_choice": true }, "databricks/databricks-mpt-7b-instruct": { + "cache_creation_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.0001e-07, "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", @@ -16890,6 +17155,14 @@ "notes": "Web Search on Amazon Bedrock AgentCore, billed by AWS on the gateway" } }, + "bing_grounding/search": { + "input_cost_per_query": 0.035, + "litellm_provider": "bing_grounding", + "mode": "search", + "metadata": { + "notes": "Grounding with Bing Search (G1 SKU): $35 per 1,000 transactions. Tokens for the Foundry model deployment that runs the grounded search are billed separately on that deployment." + } + }, "tinyfish/search": { "input_cost_per_query": 0.0, "litellm_provider": "tinyfish", @@ -23180,6 +23453,7 @@ }, "github_copilot/claude-opus-4.6-fast": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "litellm_provider": "github_copilot", "max_input_tokens": 128000, "max_output_tokens": 16000, @@ -33563,6 +33837,7 @@ }, "openrouter/anthropic/claude-sonnet-4.6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, @@ -33607,6 +33882,7 @@ }, "openrouter/anthropic/claude-opus-4.6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, @@ -35681,6 +35957,7 @@ }, "perplexity/anthropic/claude-opus-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "litellm_provider": "perplexity", "mode": "responses", "supports_web_search": true, @@ -37609,6 +37886,7 @@ "output_cost_per_token": 1e-07 }, "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37625,6 +37903,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 262000, @@ -37637,6 +37916,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37649,6 +37929,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 40000, @@ -37660,6 +37941,7 @@ "supports_tool_choice": false }, "together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": { + "deprecation_date": "2026-06-04", "input_cost_per_token": 2e-06, "litellm_provider": "together_ai", "max_input_tokens": 256000, @@ -37672,11 +37954,15 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 3e-06, "litellm_provider": "together_ai", "max_input_tokens": 128000, "max_output_tokens": 20480, "max_tokens": 20480, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 7e-06, "supports_function_calling": true, @@ -37685,6 +37971,7 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-R1-0528-tput": { + "deprecation_date": "2026-02-03", "input_cost_per_token": 5.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -37702,6 +37989,9 @@ "max_input_tokens": 65536, "max_output_tokens": 8192, "max_tokens": 8192, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.25e-06, "supports_function_calling": true, @@ -37710,9 +38000,13 @@ "supports_tool_choice": true }, "together_ai/deepseek-ai/DeepSeek-V3.1": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_tokens": 16384, + "metadata": { + "successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro" + }, "mode": "chat", "output_cost_per_token": 1.7e-06, "source": "https://www.together.ai/models/deepseek-v3-1", @@ -37724,6 +38018,7 @@ "max_output_tokens": 16384 }, "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37732,16 +38027,21 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": { - "input_cost_per_token": 8.8e-07, + "input_cost_per_token": 1.04e-06, "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 8.8e-07, + "output_cost_per_token": 1.04e-06, + "source": "https://docs.together.ai/docs/serverless-models", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true }, "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 0, "litellm_provider": "together_ai", "mode": "chat", @@ -37752,6 +38052,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": { + "deprecation_date": "2026-03-31", "input_cost_per_token": 2.7e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37762,6 +38063,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37772,6 +38074,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": { + "deprecation_date": "2026-02-06", "input_cost_per_token": 3.5e-06, "litellm_provider": "together_ai", "mode": "chat", @@ -37782,6 +38085,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 8.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37792,6 +38096,7 @@ "supports_tool_choice": true }, "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1.8e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37802,6 +38107,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-7B-Instruct-v0.1": { + "deprecation_date": "2025-11-13", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37810,6 +38116,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": { + "deprecation_date": "2026-04-02", "litellm_provider": "together_ai", "mode": "chat", "supports_function_calling": true, @@ -37817,6 +38124,7 @@ "supports_tool_choice": true }, "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "mode": "chat", @@ -37829,6 +38137,9 @@ "together_ai/moonshotai/Kimi-K2-Instruct": { "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-instruct", @@ -37872,6 +38183,7 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.5-Air-FP8": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 2e-07, "litellm_provider": "together_ai", "max_input_tokens": 128000, @@ -37889,6 +38201,9 @@ "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2.2e-06, "source": "https://www.together.ai/models/glm-4-6", @@ -37898,11 +38213,15 @@ "supports_tool_choice": true }, "together_ai/zai-org/GLM-4.7": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 4.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 200000, "max_output_tokens": 200000, "max_tokens": 200000, + "metadata": { + "successor": "together_ai/zai-org/GLM-5.2" + }, "mode": "chat", "output_cost_per_token": 2e-06, "source": "https://www.together.ai/models/glm-4-7", @@ -37912,11 +38231,15 @@ "supports_tool_choice": true }, "together_ai/moonshotai/Kimi-K2.5": { + "deprecation_date": "2026-05-21", "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 2.8e-06, "source": "https://www.together.ai/models/kimi-k2-5", @@ -37926,9 +38249,13 @@ "supports_reasoning": true }, "together_ai/moonshotai/Kimi-K2-Instruct-0905": { + "deprecation_date": "2026-03-06", "input_cost_per_token": 1e-06, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/moonshotai/Kimi-K3" + }, "mode": "chat", "output_cost_per_token": 3e-06, "source": "https://www.together.ai/models/kimi-k2-0905", @@ -37937,9 +38264,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": { + "deprecation_date": "2026-04-02", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.7-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct", @@ -37949,9 +38280,13 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, + "metadata": { + "successor": "together_ai/Qwen/Qwen3.6-Plus" + }, "mode": "chat", "output_cost_per_token": 1.5e-06, "source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking", @@ -37961,6 +38296,7 @@ "supports_tool_choice": true }, "together_ai/Qwen/Qwen3.5-397B-A17B": { + "deprecation_date": "2026-06-29", "input_cost_per_token": 6e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -37972,6 +38308,292 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "together_ai/MiniMaxAI/MiniMax-M3": { + "input_cost_per_token": 3e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Prism-ML/Ternary-Bonsai-27B": { + "input_cost_per_token": 0.0, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.5-9B": { + "input_cost_per_token": 1.7e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/Qwen/Qwen3.6-Plus": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_reasoning": true + }, + "together_ai/Qwen/Qwen3.7-Max": { + "input_cost_per_token": 1.25e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 3.75e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.7-Plus": { + "input_cost_per_token": 3.2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.28e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/Qwen/Qwen3.8-2.4T-A95B": { + "input_cost_per_token": 2.5e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1010000, + "max_output_tokens": 1010000, + "max_tokens": 1010000, + "mode": "chat", + "output_cost_per_token": 6.25e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/arize-ai/qwen-2-1.5b-instruct": { + "input_cost_per_token": 1e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": { + "input_cost_per_token": 1.4e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro": { + "input_cost_per_token": 1.74e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 512000, + "max_output_tokens": 512000, + "max_tokens": 512000, + "mode": "chat", + "output_cost_per_token": 3.48e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813": { + "input_cost_per_token": 1.32e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 3.96e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/google/gemma-3n-E4B-it": { + "input_cost_per_token": 6e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/google/gemma-4-31B-it": { + "input_cost_per_token": 3.9e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 9.7e-07, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/intfloat/multilingual-e5-large-instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "together_ai", + "max_input_tokens": 514, + "max_tokens": 514, + "mode": "embedding", + "output_cost_per_token": 2e-08, + "output_vector_size": 1024, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-llama/Llama-Guard-4-12B": { + "input_cost_per_token": 2e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 2e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/meta-models/Muse-Glimmer-30B": { + "input_cost_per_token": 3.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/moonshotai/Kimi-K2.7-Code": { + "input_cost_per_token": 9.5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/moonshotai/Kimi-K3": { + "input_cost_per_token": 3e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "together_ai/nvidia/nemotron-3-ultra-550b-a55b": { + "input_cost_per_token": 6e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 512288, + "max_output_tokens": 512288, + "max_tokens": 512288, + "mode": "chat", + "output_cost_per_token": 3.6e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/pearl-ai/gemma-4-31b-it": { + "input_cost_per_token": 2.8e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8.6e-07, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/thinkingmachines/Inkling": { + "input_cost_per_token": 1e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 4.05e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "together_ai/thinkingmachines/Inkling-Small": { + "input_cost_per_token": 5e-07, + "litellm_provider": "together_ai", + "max_input_tokens": 524288, + "max_output_tokens": 524288, + "max_tokens": 524288, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.together.ai/docs/serverless-models" + }, + "together_ai/zai-org/GLM-5.2": { + "input_cost_per_token": 1.4e-06, + "litellm_provider": "together_ai", + "max_input_tokens": 1048575, + "max_output_tokens": 1048575, + "max_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 4.4e-06, + "source": "https://docs.together.ai/docs/serverless-models", + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "tts-1": { "input_cost_per_character": 1.5e-05, "litellm_provider": "openai", @@ -39056,6 +39678,7 @@ }, "vercel_ai_gateway/anthropic/claude-opus-4.6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, @@ -40315,6 +40938,7 @@ "deprecation_date": "2027-02-05", "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -40347,6 +40971,7 @@ "deprecation_date": "2027-02-05", "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 6.25e-06, "cache_creation_input_token_cost_above_1hr": 1e-05, "cache_read_input_token_cost": 5e-07, @@ -40712,6 +41337,7 @@ "vertex_ai/claude-sonnet-4-6": { "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -43795,10 +44421,10 @@ "comment": "5 credits per second @ $0.01 per credit = $0.05 per second" } }, - "runwayml/gen4_aleph": { + "runwayml/gen4.5": { "litellm_provider": "runwayml", "mode": "video_generation", - "output_cost_per_video_per_second": 0.15, + "output_cost_per_second": 0.12, "source": "https://docs.dev.runwayml.com/guides/pricing/", "supported_modalities": [ "text", @@ -43808,13 +44434,136 @@ "video" ], "metadata": { - "comment": "15 credits per second @ $0.01 per credit = $0.15 per second" + "comment": "12 credits per second @ $0.01 per credit = $0.12 per second" } }, - "runwayml/gen3a_turbo": { + "runwayml/aleph2": { "litellm_provider": "runwayml", "mode": "video_generation", - "output_cost_per_video_per_second": 0.05, + "output_cost_per_second": 0.28, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "28 credits per second @ $0.01 per credit = $0.28 per second; 56 credit minimum per task not modeled" + } + }, + "runwayml/seedance2": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.36, + "output_cost_per_second_1080p": 0.4, + "output_cost_per_second_4k": 1.5, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "36 credits per second at 480p/720p, 40 at 1080p, 150 at 4K @ $0.01 per credit" + } + }, + "runwayml/seedance2_fast": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.29, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "29 credits per second at 480p/720p @ $0.01 per credit = $0.29 per second" + } + }, + "runwayml/seedance2_mini": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.16, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "16 credits per second @ $0.01 per credit = $0.16 per second; 64 credit minimum per task not modeled" + } + }, + "runwayml/seedance2_5": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.3, + "output_cost_per_second_480p": 0.2, + "output_cost_per_second_1080p": 0.68, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "Output: 20/30/68 credits per second at 480p/720p/1080p @ $0.01 per credit; input video billed additionally at 10/15/34 credits per input second and the 80 credit minimum per task are not modeled" + } + }, + "runwayml/hailuo3": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.1, + "output_cost_per_second_1080p": 0.15, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "10 credits per second at 768P, 15 at 2K (mapped to the 1080p tier) @ $0.01 per credit; 2 credits per reference image not modeled" + } + }, + "runwayml/gemini_omni_flash": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.1, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "10 credits per second @ $0.01 per credit = $0.10 per second" + } + }, + "runwayml/veo3.1": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.4, "source": "https://docs.dev.runwayml.com/guides/pricing/", "supported_modalities": [ "text", @@ -43824,7 +44573,23 @@ "video" ], "metadata": { - "comment": "5 credits per second @ $0.01 per credit = $0.05 per second" + "comment": "40 credits per second with audio, 20 without @ $0.01 per credit; priced at the with-audio rate" + } + }, + "runwayml/veo3.1_fast": { + "litellm_provider": "runwayml", + "mode": "video_generation", + "output_cost_per_second": 0.15, + "source": "https://docs.dev.runwayml.com/guides/pricing/", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ], + "metadata": { + "comment": "15 credits per second with audio, 10 without @ $0.01 per credit; priced at the with-audio rate" } }, "runwayml/gen4_image": { @@ -48481,6 +49246,7 @@ "vertex_ai/claude-sonnet-4-6@default": { "regional_endpoint_uplift_multiplier": 1.1, "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -48595,12 +49361,13 @@ "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -48627,12 +49394,13 @@ "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -48659,12 +49427,13 @@ "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49265,6 +50034,7 @@ }, "snowflake/claude-sonnet-4-6": { "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, "max_tokens": 16384, "max_input_tokens": 200000, "max_output_tokens": 16384, @@ -50255,6 +51025,14 @@ "supports_adaptive_thinking": true } }, + { + "name": "claude-legacy-thinking", + "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", + "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", + "model_info": { + "supports_legacy_thinking": true + } + }, { "name": "claude-always-on-thinking", "pattern": "claude-(?:fable|mythos)-", diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index f5560a20ab2..f7f60c7666d 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -428,6 +428,14 @@ "type": "number", "minimum": 0 }, + "output_cost_per_second_480p": { + "type": "number", + "minimum": 0 + }, + "output_cost_per_second_4k": { + "type": "number", + "minimum": 0 + }, "output_cost_per_token": { "type": "number", "minimum": 0, @@ -625,6 +633,9 @@ "supports_image_size": { "type": "boolean" }, + "supports_legacy_thinking": { + "type": "boolean" + }, "supports_low_reasoning_effort": { "type": "boolean" }, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index a990f7c3830..bbfbbbc1cfb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3020 + "limit": 3016 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2017 + "limit": 2015 }, "ANN202": { "limit": 852 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2920 + "limit": 2918 }, "C401": { "limit": 8 @@ -108,7 +108,7 @@ "limit": 3 }, "F401": { - "limit": 17 + "limit": 14 }, "LOG015": { "limit": 5 @@ -152,9 +152,6 @@ "PLW0127": { "limit": 57 }, - "PLW0133": { - "limit": 1 - }, "PLW0602": { "limit": 215 }, @@ -171,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 176 + "limit": 175 }, "RUF012": { "limit": 240 diff --git a/ruff-tests.toml b/ruff-tests.toml index e52e1a96d00..7d28d6d838a 100644 --- a/ruff-tests.toml +++ b/ruff-tests.toml @@ -40,6 +40,17 @@ # later binding makes the name local for the whole body, so the read raises # UnboundLocalError, and in an autouse fixture that takes every test in the # directory down with it +# F601 the same key literal twice in one dict. Python keeps the last value, so the +# first is dropped before the test ever runs, and a fixture that looks like it +# covers two cases covers one +# B023 a closure over a loop variable. Every closure sees the last iteration's value, +# so a per-case callback built in a loop checks the last case N times. Bind the +# value as a parameter instead +# B025 an `except` for a type an earlier `except` already catches. The second handler +# is unreachable, so the recovery or skip written there never happens +# F632 `is` against a literal. It compares identity, so it passes only where CPython +# happens to intern the value and stops meaning what it says the moment the +# value is built at runtime # # No target-version here on purpose: it resolves from requires-python (>=3.10), so # 3.11-only builtins like BaseExceptionGroup are correctly flagged in a tree that @@ -63,4 +74,8 @@ lint.select = [ "PLW0127", "RUF043", "F823", + "F601", + "B023", + "B025", + "F632", ] diff --git a/ruff.toml b/ruff.toml index 9b90910b355..44bdf9d8125 100644 --- a/ruff.toml +++ b/ruff.toml @@ -5,9 +5,9 @@ lint.ignore = ["F405", "E402", "F403"] lint.extend-select = [ "T20", "PGH004", "RUF008", "RUF009", "RUF100", "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", - "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", - "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", "UP012", - "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", + "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", + "RUF010", "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", + "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", ] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index ff2f3f817f9..3bd3c4d1d6c 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -14,6 +14,14 @@ longer signal it. ## [Unreleased] +### Added + +- **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it + +### Fixed + +- **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state + ### Changed - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/docs/resources/team.md b/terraform/provider/docs/resources/team.md index 68535309f10..65ab4bf82d4 100644 --- a/terraform/provider/docs/resources/team.md +++ b/terraform/provider/docs/resources/team.md @@ -24,11 +24,18 @@ resource "litellm_team" "advanced_team" { # Budget and rate limiting max_budget = 1000.0 + soft_budget = 800.0 budget_duration = "1mo" tpm_limit = 500000 rpm_limit = 5000 blocked = false + # Who gets paged when spend crosses soft_budget + soft_budget_alerting_emails = ["finops@example.com"] + + # Tags for spend tracking and tag-based routing + tags = ["team:ai-research", "environment:production"] + # Team member permissions team_member_permissions = [ "create_key", @@ -91,7 +98,9 @@ The following arguments are supported: * `models` - (Optional) List of model names that this team can access. -* `metadata` - (Optional) A map of metadata key-value pairs associated with the team. +* `metadata` - (Optional) A map of string metadata key-value pairs associated with the team. `tags` and `soft_budget_alerting_emails` are stored by the proxy under metadata but are managed through their own attributes below, not this map. + +* `tags` - (Optional) List of tags applied to the team, used for [spend tracking](https://docs.litellm.ai/docs/proxy/enterprise#tracking-spend-for-custom-tags) and [tag-based routing](https://docs.litellm.ai/docs/proxy/tag_routing). * `blocked` - (Optional) Whether the team is blocked from making requests. Default is `false`. @@ -101,6 +110,10 @@ The following arguments are supported: * `max_budget` - (Optional) Maximum budget allocated to the team. +* `soft_budget` - (Optional) Spend threshold at which the proxy sends a soft budget alert without blocking requests. + +* `soft_budget_alerting_emails` - (Optional) List of email addresses notified when the team's spend crosses `soft_budget`. + * `budget_duration` - (Optional) Duration for the budget cycle. Valid values are: * `daily` * `weekly` diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index 88e0dcd4811..2a167a1b5c4 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -53,6 +53,11 @@ func ResourceLiteLLMTeam() *schema.Resource { Type: schema.TypeFloat, Optional: true, }, + "soft_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Spend threshold that triggers a soft budget alert without blocking requests", + }, "budget_duration": { Type: schema.TypeString, Optional: true, @@ -72,6 +77,18 @@ func ResourceLiteLLMTeam() *schema.Resource { Elem: &schema.Schema{Type: schema.TypeString}, Description: "List of permissions granted to team members", }, + "tags": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Tags for spend tracking and tag-based routing", + }, + "soft_budget_alerting_emails": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Email addresses alerted when the team crosses soft_budget", + }, }, } } @@ -117,21 +134,20 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { return nil } - var teamResp TeamResponse - if err := json.NewDecoder(resp.Body).Decode(&teamResp); err != nil { + var infoResp TeamInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&infoResp); err != nil { return fmt.Errorf("error decoding team info response: %w", err) } + teamResp := infoResp.TeamInfo // Update the state with values from the response or fall back to the data passed in during creation d.Set("team_alias", GetStringValue(teamResp.TeamAlias, d.Get("team_alias").(string))) d.Set("organization_id", GetStringValue(teamResp.OrganizationID, d.Get("organization_id").(string))) - // Handle metadata separately as it's a map - if teamResp.Metadata != nil { - d.Set("metadata", teamResp.Metadata) - } else { - d.Set("metadata", d.Get("metadata")) - } + metadata, tags, alertEmails := splitTeamMetadata(teamResp.Metadata) + d.Set("metadata", metadata) + d.Set("tags", tags) + d.Set("soft_budget_alerting_emails", alertEmails) if teamResp.TPMLimit != nil { d.Set("tpm_limit", *teamResp.TPMLimit) @@ -142,6 +158,7 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { if teamResp.MaxBudget != nil { d.Set("max_budget", *teamResp.MaxBudget) } + d.Set("soft_budget", teamResp.SoftBudget) d.Set("budget_duration", GetStringValue(teamResp.BudgetDuration, d.Get("budget_duration").(string))) // Handle models separately as it's a list @@ -240,15 +257,77 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} "team_alias": d.Get("team_alias").(string), } - for _, key := range []string{"organization_id", "metadata", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions"} { + for _, key := range []string{"organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions"} { if v, ok := d.GetOk(key); ok { teamData[key] = v } } + if v, ok := d.GetOk("soft_budget"); ok { + teamData["soft_budget"] = v + } else if d.HasChange("soft_budget") { + teamData["soft_budget"] = nil + } + + if v, ok := d.GetOk("tags"); ok || d.HasChange("tags") { + teamData["tags"] = v + } + + if metadata := buildTeamMetadata(d); metadata != nil { + teamData["metadata"] = metadata + } + return teamData } +// /team/update replaces metadata wholesale, so the full map must go out whenever either half changed. +func buildTeamMetadata(d *schema.ResourceData) map[string]interface{} { + metadata := map[string]interface{}{} + for k, v := range d.Get("metadata").(map[string]interface{}) { + metadata[k] = v + } + if v, ok := d.GetOk("soft_budget_alerting_emails"); ok { + metadata["soft_budget_alerting_emails"] = v + } + if len(metadata) == 0 && !d.HasChange("metadata") && !d.HasChange("soft_budget_alerting_emails") { + return nil + } + return metadata +} + +func splitTeamMetadata(raw map[string]interface{}) (map[string]string, []string, []string) { + metadata := map[string]string{} + var tags, alertEmails []string + for k, v := range raw { + switch k { + case "tags": + tags = toStringSlice(v) + case "soft_budget_alerting_emails": + alertEmails = toStringSlice(v) + case "team_member_budget_id": + default: + if s, ok := v.(string); ok { + metadata[k] = s + } + } + } + return metadata, tags, alertEmails +} + +func toStringSlice(v interface{}) []string { + items, ok := v.([]interface{}) + if !ok { + return nil + } + out := make([]string, 0, len(items)) + for _, item := range items { + if s, ok := item.(string); ok { + out = append(out, s) + } + } + return out +} + func handleResponse(resp *http.Response, action string) error { if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go new file mode 100644 index 00000000000..1f74be4819d --- /dev/null +++ b/terraform/provider/litellm/resource_team_test.go @@ -0,0 +1,184 @@ +package litellm + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" +) + +func newTeamTestServer(t *testing.T, captured *map[string]interface{}, infoBody string) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case endpointTeamNew, endpointTeamUpdate: + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, captured) + w.Write([]byte(`{}`)) + case endpointTeamInfo: + w.Write([]byte(infoBody)) + case endpointTeamPermissionsList: + w.Write([]byte(`{"team_id":"team-1","team_member_permissions":[],"all_available_permissions":[]}`)) + default: + w.WriteHeader(http.StatusNotFound) + } + })) +} + +const teamInfoWithSoftBudget = `{ + "team_id": "team-1", + "team_info": { + "team_id": "team-1", + "team_alias": "insights", + "max_budget": 750.0, + "soft_budget": 600.0, + "models": ["claude-haiku-4-5"], + "metadata": { + "department": "customer-insights", + "tags": ["team:customer-insights", "environment:production"], + "soft_budget_alerting_emails": ["finops@example.com"], + "team_member_budget_id": "budget-1" + } + }, + "keys": [], + "team_memberships": [] +}` + +func TestTeamCreateSendsSoftBudgetTagsAndAlertEmails(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, ResourceLiteLLMTeam().Schema, map[string]interface{}{ + "team_alias": "insights", + "max_budget": 750.0, + "soft_budget": 600.0, + "tags": []interface{}{"team:customer-insights", "environment:production"}, + "soft_budget_alerting_emails": []interface{}{"finops@example.com"}, + "metadata": map[string]interface{}{"department": "customer-insights"}, + }) + + if err := resourceLiteLLMTeamCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if got := captured["soft_budget"]; got != 600.0 { + t.Fatalf("payload soft_budget = %v, want 600", got) + } + wantTags := []interface{}{"team:customer-insights", "environment:production"} + if got := captured["tags"]; !reflect.DeepEqual(got, wantTags) { + t.Fatalf("payload tags = %v, want %v", got, wantTags) + } + wantMetadata := map[string]interface{}{ + "department": "customer-insights", + "soft_budget_alerting_emails": []interface{}{"finops@example.com"}, + } + if got := captured["metadata"]; !reflect.DeepEqual(got, wantMetadata) { + t.Fatalf("payload metadata = %v, want %v", got, wantMetadata) + } +} + +func TestTeamReadMapsTeamInfoEnvelope(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, teamInfoWithSoftBudget) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, ResourceLiteLLMTeam().Schema, map[string]interface{}{}) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if got := d.Get("team_alias"); got != "insights" { + t.Fatalf("team_alias = %v, want insights", got) + } + if got := d.Get("soft_budget"); got != 600.0 { + t.Fatalf("soft_budget = %v, want 600", got) + } + if got := d.Get("max_budget"); got != 750.0 { + t.Fatalf("max_budget = %v, want 750", got) + } + wantTags := []interface{}{"team:customer-insights", "environment:production"} + if got := d.Get("tags"); !reflect.DeepEqual(got, wantTags) { + t.Fatalf("tags = %v, want %v", got, wantTags) + } + wantEmails := []interface{}{"finops@example.com"} + if got := d.Get("soft_budget_alerting_emails"); !reflect.DeepEqual(got, wantEmails) { + t.Fatalf("soft_budget_alerting_emails = %v, want %v", got, wantEmails) + } + wantMetadata := map[string]interface{}{"department": "customer-insights"} + if got := d.Get("metadata"); !reflect.DeepEqual(got, wantMetadata) { + t.Fatalf("metadata = %v, want %v (server-managed team_member_budget_id dropped)", got, wantMetadata) + } +} + +func TestTeamUpdateClearsRemovedTagsAndSoftBudget(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"insights"},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + res := ResourceLiteLLMTeam() + priorData := schema.TestResourceDataRaw(t, res.Schema, map[string]interface{}{ + "team_alias": "insights", + "soft_budget": 600.0, + "tags": []interface{}{"team:to-be-removed"}, + "soft_budget_alerting_emails": []interface{}{"ops@example.com"}, + "metadata": map[string]interface{}{"department": "eng"}, + }) + priorData.SetId("team-1") + prior := priorData.State() + config := terraform.NewResourceConfigRaw(map[string]interface{}{ + "team_alias": "insights", + "metadata": map[string]interface{}{"department": "eng"}, + }) + diff, err := res.Diff(context.Background(), prior, config, nil) + if err != nil { + t.Fatalf("diff failed: %v", err) + } + d, err := schema.InternalMap(res.Schema).Data(prior, diff) + if err != nil { + t.Fatalf("data failed: %v", err) + } + + if err := resourceLiteLLMTeamUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + + if got, ok := captured["soft_budget"]; !ok || got != nil { + t.Fatalf("payload soft_budget = %v (present=%v), want explicit null", got, ok) + } + if got := captured["tags"]; !reflect.DeepEqual(got, []interface{}{}) { + t.Fatalf("payload tags = %v, want []", got) + } + if got := captured["metadata"]; !reflect.DeepEqual(got, map[string]interface{}{"department": "eng"}) { + t.Fatalf("payload metadata = %v, want department only", got) + } +} + +func TestTeamReadClearsSoftBudgetWhenProxyReturnsNull(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id":"team-1","team_info":{"team_id":"team-1","team_alias":"insights","soft_budget":null},"keys":[],"team_memberships":[]}`) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, ResourceLiteLLMTeam().Schema, map[string]interface{}{ + "team_alias": "insights", + "soft_budget": 600.0, + }) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if got := d.Get("soft_budget"); got != 0.0 { + t.Fatalf("soft_budget = %v, want cleared after the proxy returned null", got) + } +} diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 069fe4b3e23..66d1f6a8ba9 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -33,6 +33,11 @@ type ModelRequest struct { Additional map[string]interface{} `json:"additional"` } +type TeamInfoResponse struct { + TeamID string `json:"team_id"` + TeamInfo TeamResponse `json:"team_info"` +} + // TeamResponse represents a response from the API containing team information. type TeamResponse struct { TeamID string `json:"team_id,omitempty"` @@ -42,6 +47,7 @@ type TeamResponse struct { TPMLimit *int `json:"tpm_limit,omitempty"` RPMLimit *int `json:"rpm_limit,omitempty"` MaxBudget *float64 `json:"max_budget,omitempty"` + SoftBudget *float64 `json:"soft_budget,omitempty"` BudgetDuration string `json:"budget_duration,omitempty"` Models []string `json:"models"` Blocked bool `json:"blocked,omitempty"` diff --git a/tests/code_coverage_tests/bedrock_pricing.py b/tests/code_coverage_tests/bedrock_pricing.py index 5984dd8b3a4..73a2000d2cf 100644 --- a/tests/code_coverage_tests/bedrock_pricing.py +++ b/tests/code_coverage_tests/bedrock_pricing.py @@ -93,7 +93,7 @@ def get_bedrock_pricing(url, providers): else: # General logic for other providers section = soup.find( - "h2", text=lambda t: t and provider.lower() in t.lower() + "h2", text=lambda t, needle=provider.lower(): t and needle in t.lower() ) if not section: pricing_data[provider] = "Provider section not found" diff --git a/tests/code_coverage_tests/check_migrations_no_data_rewrites.py b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py new file mode 100644 index 00000000000..e4375d6d8ba --- /dev/null +++ b/tests/code_coverage_tests/check_migrations_no_data_rewrites.py @@ -0,0 +1,962 @@ +#!/usr/bin/env python3 +"""Ban row-rewriting DML from Prisma migrations. + +Migrations run synchronously at proxy boot, before the process serves traffic, so +anything whose cost scales with existing table size turns into downtime. A single +`UPDATE` with no batching over a spend-log-sized table is minutes of unavailability +plus a doubled heap that plain autovacuum will not give back. + +What is banned is the row-rewriting DML behind that, not everything whose cost +scales that way. A non-concurrent `CREATE INDEX`, an `ALTER COLUMN ... TYPE` that is +not binary coercible, a volatile `DEFAULT` on a new column, a `CREATE TABLE ... AS +SELECT` or `SELECT ... INTO` filling a new table from an existing one, the rename +that pairs with one of those to swap a table out, and a `REFRESH MATERIALIZED VIEW` +all read the whole table and all pass. That is deliberate: a rule wide enough to +reach them fires on most ordinary migrations, and a marker everyone adds by reflex +stops carrying information. The outage this was written for was a backfill. + +Flagged, per statement, by its leading keyword: + + UPDATE rewrites every matching row, and `WHERE` does not bound the scan + DELETE same scan, and the dead tuples outlive the migration + MERGE both of the above in one statement + INSERT only when its rows come from a query rather than a literal `VALUES` + list. The query counts wherever it sits, since Postgres takes it + parenthesised, and `TABLE t` is one as much as a `SELECT` is. An + insert bounded by a `VALUES` list passes, written bare or in + parentheses, and so do the scalar subqueries in that list and the + `RETURNING` and `ON CONFLICT` clauses written after it, none of which + supply the rows. A `VALUES` reached through a subquery or joined to a + query by a set operation bounds nothing + WITH a CTE-led statement containing any of the above. An `INSERT` is read + against the part of the statement holding it, so a writable CTE + bounded by its own `VALUES` list is not handed the query the statement + ends with as the rows it copies + +Referential actions (`ON DELETE CASCADE`, `ON UPDATE CASCADE`) are schema, never a +statement's leading keyword, so they pass. + +A statement wrapped in `EXPLAIN` is judged on the statement itself, because the +`ANALYZE` form runs it rather than only planning it, and a rewrite left under one +rewrites the table on the way to printing its timings. Explaining a rewrite without +`ANALYZE` is flagged too: nothing here needs the plan of a statement it is being +told not to run at boot, and a marker is a cheap answer if one ever does. + +Statements inside dollar-quoted bodies are scanned too. `DO $$ ... $$` is this +repo's idiom for conditional DDL, so a body is where an `UPDATE` would otherwise +hide. A `CREATE FUNCTION` or `CREATE PROCEDURE` body is the exception, because +defining a routine only stores it: that body is read when the same migration names +the routine somewhere else, which is what defining a backfill and then running it +looks like, and left alone when nothing calls it. A routine whose name needed +quoting is read either way, since quoting is blanked at the call sites too and a +call written there could never be found. The SQL an `EXECUTE` runs is scanned the same way, since a rewrite reads the +same to Postgres whether it is spelled out or handed over as a string, and so is a +literal parked in a variable some `EXECUTE` in the same body then runs by name, +however it got there: an assignment with `:=`, the bare `=` PL/pgSQL takes as the +same operator, a query returning it through `INTO`, or a loop walking the query it +came out of. So is the body of a `DO` written in single quotes rather than dollar +quotes. A literal nothing runs is text, however much it reads like a statement, so +an error message naming a `DELETE` the application handles stays a message. + +Each literal is read on its own, so a keyword built by concatenating fragments that +do not contain it (`'UPD' || 'ATE ...'`) is not caught. Every fragment is scanned, +so a concatenation is caught wherever the keyword survives whole in one of them, +which covers `'UPDATE ' || quote_ident(t)` and the rest of the readable shapes. The +gap needs a keyword deliberately split down the middle, and this check is a guard +against a rewrite reaching a boot unnoticed, not a defence against someone hiding +one on purpose. + +Line numbers always count against the whole migration file, however deeply the +statement is nested, so a reported line points at the statement and the markers +below line up with the statements they exempt. + +Add a column and let the application populate it, or run the rewrite as an opt-in +batched job outside boot. When a rewrite is genuinely bounded and must ship inside +the migration, put `-- data-migration-ok: ` on the statement or on the line +above it, naming what bounds it. The reason is required. A marker sharing a line +with the statement it follows exempts that statement alone, so the next statement +down is still checked rather than picking the marker up as its own. A marker on an +`EXECUTE` or on the assignment feeding one covers the single-quoted SQL that +statement hands off, so it goes where the migration reads rather than inside the +string. A dollar-quoted payload is not a string to this check but a region read like +any other body, so a rewrite inside one takes its marker on the rewrite itself. That +placement is deliberate rather than an oversight: a marker covering a whole body +would let one written for a `DO` block silence a rewrite added to that block later. + +`GRANDFATHERED` freezes the violations that predate this check. Prisma records a +checksum for every applied migration and this repo treats applied files as +immutable, so those two cannot take an inline marker. The set is closed; a new +migration belongs nowhere in it. +""" + +from __future__ import annotations + +import re +import sys +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +MIGRATIONS_DIR = REPO_ROOT / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" + +GRANDFATHERED = frozenset( + { + "20260817000000_shadow_eval_multi_key", + "20260818224500_add_shadow_eval_stopped_by", + } +) + +MARKER = re.compile(r"--[ \t]*data-migration-ok:[ \t]*(\S.*?)[ \t]*$", re.MULTILINE) +DOLLAR_TAG = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") +FIRST_WORD = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +STATEMENT = re.compile(r"[^;]+") +RUN_BY_NAME = re.compile(r"\bEXECUTE\s+([A-Za-z_][A-Za-z0-9_]*)", re.IGNORECASE) +INTO_TARGETS = re.compile( + r"\bINTO\s+(?:STRICT\s+)?" + r"([A-Za-z_][A-Za-z0-9_]*(?:\s*,\s*[A-Za-z_][A-Za-z0-9_]*)*)", + re.IGNORECASE, +) +LOOP_TARGET = re.compile(r"\bFOR(?:EACH)?\s+([A-Za-z_][A-Za-z0-9_]*)\s+IN\b", re.IGNORECASE) +LOOP_HEADER = re.compile(r"\bFOR(?:EACH)?\b.*?\bLOOP\b", re.IGNORECASE | re.DOTALL) +WORD_OR_ASSIGN = re.compile(r"[A-Za-z_][A-Za-z0-9_]*|:=|(?!:=])=(?![=>])") +PRECEDING_WORD = re.compile(r"([A-Za-z_][A-Za-z0-9_]*)[^A-Za-z0-9_]*$") +QUALIFIER_GAP = re.compile(r"[\s.]*") +EXPLAIN_OPTIONS = re.compile(r"\bEXPLAIN\b(?:\s+(?:ANALYZE|ANALYSE|VERBOSE)\b)+", re.IGNORECASE) +DEFINES_A_ROUTINE = re.compile( + r"\bCREATE\b(?:\s+OR\s+REPLACE)?\s+(?:FUNCTION|PROCEDURE)\b", re.IGNORECASE +) +QUALIFIED_NAME = r"(?:\"[^\"]*\"|[A-Za-z_][A-Za-z0-9_$]*)" +ROUTINE_NAME = re.compile(rf"\s*(?:{QUALIFIED_NAME}\s*\.\s*)?({QUALIFIED_NAME})") +OPENS_A_CALL = re.compile(r"\s*\(") +NAMES_AN_INDEX = re.compile(r"\bCREATE\b.+\bINDEX\b", re.IGNORECASE | re.DOTALL) +INTRODUCES_A_RELATION = frozenset({"TABLE", "INTO", "REFERENCES", "EXISTS", "COPY"}) + +REWRITES_ROWS = frozenset({"UPDATE", "DELETE", "MERGE"}) + +JOINS_QUERIES = ("UNION", "INTERSECT", "EXCEPT") + +SET_OPERATION = re.compile(rf"\b(?:{'|'.join(JOINS_QUERIES)})\b", re.IGNORECASE) + +STATEMENT_KEYWORDS = REWRITES_ROWS | frozenset( + { + "INSERT", + "SELECT", + "WITH", + "ALTER", + "CREATE", + "DROP", + "TRUNCATE", + "COMMENT", + "GRANT", + "REVOKE", + "COPY", + "SET", + "PERFORM", + "RAISE", + "RETURN", + "EXECUTE", + "DO", + "CALL", + "REINDEX", + "REFRESH", + "VACUUM", + "ANALYZE", + } +) + +GUARDS_A_CONDITION = frozenset({"IF", "ELSIF", "ELSEIF", "CASE", "WHEN", "WHILE", "EXIT", "ASSERT"}) + +OPENS_A_BLOCK = frozenset({"BEGIN", "THEN", "ELSE", "LOOP"}) + +NEVER_A_VARIABLE = frozenset({"INTO", "USING"}) + +BIND_VALUES = re.compile(r"\bUSING\b", re.IGNORECASE) + +WRITES_ROWS = re.compile(r"\bINSERT\b", re.IGNORECASE) + +GUIDANCE = """ +Migrations apply at proxy boot, before it serves traffic, so a statement whose cost +scales with table size is downtime. Add the column and let the application backfill +it, or move the rewrite to a batched job outside boot. + +If the rewrite is genuinely bounded and has to ship in the migration, mark the +statement with the bound spelled out: + + -- data-migration-ok: + UPDATE ... +""" + + +@dataclass(frozen=True, slots=True) +class Violation: + migration: str + line: int + keyword: str + + def render(self) -> str: + location = f"{MIGRATIONS_DIR.relative_to(REPO_ROOT)}/{self.migration}/migration.sql" + return f"{location}:{self.line}: {self.keyword} rewrites existing rows at boot" + + +@dataclass(frozen=True, slots=True) +class Marker: + start: int + end: int + standalone: bool + + +@dataclass(frozen=True, slots=True) +class Markers: + sql: str + written: tuple[Marker, ...] + + def exempt(self, start: int, end: int) -> bool: + """Whether the statement spanning `start` to `end` carries a marker.""" + return any(self.speaks_for(marker, start, end) for marker in self.written) + + def speaks_for(self, marker: Marker, start: int, end: int) -> bool: + """Whether a marker is written against this statement. One alone on its line speaks for + the statement below it, which is how a marker written above a rewrite exempts it, and one + sharing its line with code speaks for the statement it follows. Either is matched by where + it sits rather than by the line it lands on, so a second statement sharing that line does + not inherit the exemption. A marker inside a statement speaks for it whichever kind it is, + which is how one on the opening line of a long statement still covers the whole of it.""" + if start <= marker.start < end: + return True + if marker.standalone: + return self.on_the_line_below(marker.end, start) + return self.only_separators(end, marker.start) + + def on_the_line_below(self, start: int, end: int) -> bool: + """Whether a marker on its own line is written directly above the statement, which means + one line break and nothing else that carries meaning. A blank line between the two leaves + the marker reading as a note about the file rather than a bound on what follows it.""" + return self.only_separators(start, end) and self.sql[start:end].count("\n") == 1 + + def only_separators(self, start: int, end: int) -> bool: + """Whether nothing but statement separators lie between two points, which is what makes a + marker and the statement it follows adjacent however they are laid out.""" + return start <= end and not self.sql[start:end].strip(" \t\r\n;") + + +def blank(text: str) -> str: + return "".join(character if character == "\n" else " " for character in text) + + +def undouble(literal: str) -> str: + """The SQL a single-quoted literal stands for, with each doubled quote read back as the one it + escapes. `mask` hands the literal on raw, `''` and all, so re-lexing it as SQL needs the escapes + resolved first: left doubled, the first quote of a pair opens an empty string and closes it on + the second, and a `--` or `/*` in what was a nested string is then bare and blanks the code + after it.""" + return literal.replace("''", "'") + + +def defuse_escapes(literal: str) -> str: + """The literal made safe to re-lex without moving anything: each doubled quote becomes a real + quote and a space, so a `--` or `/*` in a nested string stays inside its string the way + `undouble` achieves it, while the pair keeps its two characters. Every newline and every + character after a resolved escape then holds the offset it had in the document, so a rewrite + scanned out of the literal reports its true file line and lines up with the file's markers, + which `undouble` cannot promise because it shrinks the text as it collapses each pair.""" + return literal.replace("''", "' ") + + +def mask( + sql: str, +) -> tuple[str, tuple[tuple[int, int], ...], tuple[tuple[int, int], ...], tuple[tuple[int, int], ...]]: + """Blank comments and quoted text, keeping offsets, and locate the spans that can still + hold SQL: dollar-quoted bodies, and the single-quoted literals `EXECUTE` runs. Also locate + the double-quoted identifiers that open a call (`"backfill"(`), so a routine invoked through + one can be found by name even though the call is blanked here the way every other quoted run + of text is. Whether an identifier opens a call is read from the masked text rather than the + raw SQL, so a comment sitting between the name and its parenthesis, blanked to spaces here, is + skipped exactly as whitespace is. A double-quoted identifier that opens no call, a column, + index, or constraint name, is left out, so it never masquerades as a call to a like-named + routine, as is one whose parenthesis is a column list rather than an argument list, the table + of a `CREATE TABLE`, `INSERT INTO`, `REFERENCES`, `COPY`, or `CREATE INDEX`, which + `names_a_relation` reads from the word before the name.""" + chunks: list[str] = [] + bodies: list[tuple[int, int]] = [] + literals: list[tuple[int, int]] = [] + identifiers: list[tuple[int, int]] = [] + index = 0 + length = len(sql) + + while index < length: + pair = sql[index : index + 2] + + if pair == "--": + stop = sql.find("\n", index) + stop = length if stop == -1 else stop + chunks.append(blank(sql[index:stop])) + index = stop + continue + + if pair == "/*": + stop = skip_block_comment(sql, index) + chunks.append(blank(sql[index:stop])) + index = stop + continue + + character = sql[index] + + if character in "'\"": + stop = skip_quoted(sql, index, character) + if character == "'": + closed = sql[stop - 1 : stop] == character + literals.append((index + 1, max(index + 1, stop - 1 if closed else stop))) + else: + identifiers.append((index, stop)) + chunks.append(blank(sql[index:stop])) + index = stop + continue + + if character == "$": + tag = DOLLAR_TAG.match(sql, index) + if tag is not None: + closing = sql.find(tag.group(), tag.end()) + body_end = length if closing == -1 else closing + stop = length if closing == -1 else closing + len(tag.group()) + bodies.append((tag.end(), body_end)) + chunks.append(blank(sql[index:stop])) + index = stop + continue + + chunks.append(character) + index += 1 + + masked = "".join(chunks) + calls = tuple( + (start, end) + for start, end in identifiers + if OPENS_A_CALL.match(masked, end) and not names_a_relation(masked[:start]) + ) + return masked, tuple(bodies), tuple(literals), calls + + +def skip_block_comment(sql: str, start: int) -> int: + depth = 1 + index = start + 2 + while index < len(sql) and depth > 0: + pair = sql[index : index + 2] + if pair == "/*": + depth += 1 + index += 2 + elif pair == "*/": + depth -= 1 + index += 2 + else: + index += 1 + return index + + +def skip_quoted(sql: str, start: int, quote: str) -> int: + """One quoted run, up to and including its closing quote. A doubled quote is an escaped + quote sitting inside the run rather than the end of it. Closing on the first and reopening + on the second would mask the same span, which is why this looked like it needed no special + case, but the run is also handed on whole as one literal, and splitting it there offers the + tail of a string to be read as SQL in its own right.""" + index = start + 1 + while True: + stop = sql.find(quote, index) + if stop == -1: + return len(sql) + if sql[stop + 1 : stop + 2] == quote: + index = stop + 2 + continue + return stop + 1 + + +def strip_parens(statement: str) -> str: + """Blank parenthesised groups in place, so an `IF EXISTS (SELECT ...)` guard does not + stand in for the statement it guards.""" + chunks: list[str] = [] + depth = 0 + + for character in statement: + if character == "(": + depth += 1 + chunks.append(" ") + elif character == ")": + depth = max(depth - 1, 0) + chunks.append(" ") + elif depth > 0 and character != "\n": + chunks.append(" ") + else: + chunks.append(character) + + return "".join(chunks) + + +def strip_explain(statement: str) -> str: + """Blank an `EXPLAIN` written with bare options, since the `ANALYZE` among them would + otherwise stand in for the keyword of the statement being explained. That statement is + the one worth reading: `EXPLAIN ANALYZE` runs it rather than only planning it, so a + rewrite underneath rewrites the table for real. The parenthesised option list needs + nothing here, already being blanked as a group.""" + return EXPLAIN_OPTIONS.sub(lambda match: blank(match.group()), statement) + + +def leading_keyword(statement: str) -> re.Match[str] | None: + """The statement's own keyword, looking past what wraps it: a parenthesised guard, + PL/pgSQL block syntax such as `BEGIN`, `IF ... THEN` and `END`, and an `EXPLAIN`. + Offsets survive both strips, so the match still points into `statement` itself.""" + return next( + ( + word + for word in FIRST_WORD.finditer(strip_explain(strip_parens(statement))) + if word.group().upper() in STATEMENT_KEYWORDS + ), + None, + ) + + +def offending_keyword(statement: str) -> str | None: + word = leading_keyword(statement) + if word is None: + return None + + keyword = word.group().upper() + + if keyword in REWRITES_ROWS: + return keyword + + if keyword == "INSERT": + source = row_source_keyword(statement) + return None if source is None else f"INSERT ... {source}" + + if keyword == "WITH": + nested = next((name for name in sorted(REWRITES_ROWS) if contains(statement, name)), None) + if nested is not None: + return f"WITH ... {nested}" + if contains(statement, "INSERT"): + source = insert_row_source(statement) + if source is not None: + return f"WITH ... INSERT ... {source}" + + return None + + +def insert_row_source(statement: str) -> str | None: + """Which keyword supplies the rows to an `INSERT` written somewhere inside a `WITH` + statement. Only the parts that hold that insert are read, because a writable CTE sits + beside the query the statement ends with and reading the whole thing hands the insert + the outer `SELECT` as its row source: `WITH c AS (INSERT ... VALUES (1) RETURNING "x") + SELECT * FROM c` adds one literal row and copies nothing. A CTE keeps its insert in a + parenthesised group, and the statement's own insert, if it is the one writing, runs from + the keyword to the end, found in the text outside every parenthesis so a group's insert + is not counted twice.""" + inserts = [group for group in parenthesised_groups(statement) if contains(group, "INSERT")] + written = WRITES_ROWS.search(strip_parens(statement)) + if written is not None: + inserts.append(statement[written.start() :]) + sources = (row_source_keyword(insert) for insert in inserts) + return next((source for source in sources if source is not None), None) + + +def row_source_keyword(statement: str) -> str | None: + """Which keyword supplies an `INSERT` its rows, or `None` when a literal `VALUES` list + does. A query outside every parenthesis is the row source outright. Failing that, a + set operation at that same level joins several terms, and the insert is a rewrite when + any one of them is a query, so each term is read on its own rather than the statement + read whole. Failing that, a `VALUES` outside every parenthesis is itself the row source, + so the scalar subqueries and helper CTEs nested within that list do not make the insert + a rewrite. Failing all three, the rows come from a parenthesised group, which Postgres + accepts and which reading only the unparenthesised text would let through: + `INSERT INTO "t" ("a") (SELECT ...)` copies a whole table. Each group at that level is + read on its own terms until one of them supplies the rows, since the ones before it are + the column list and the ones after it are the conflict target and the rest of the clauses + an insert is allowed to carry. A wrapped `VALUES` list is the row source as much as a + wrapped query is, so it ends the search rather than being skipped over: reading past it + reaches a `RETURNING (SELECT ...)` or a `DO UPDATE SET "a" = (SELECT ...)` written after + it and calls that scalar subquery the rows the insert copies. The group is read on its + own terms before it is allowed to end the search, because a `VALUES` list joined to a + query by a set operation inside the group supplies every row the query does, and + stopping on the word `VALUES` alone would pass the whole copy.""" + outer = strip_parens(statement) + joined = row_source_in(outer) + if joined is not None: + return joined + if SET_OPERATION.search(outer): + sources = (row_source_keyword(term) for term in set_operation_terms(statement, outer)) + return next((source for source in sources if source is not None), None) + if contains(outer, "VALUES"): + return None + groups = list(parenthesised_groups(statement)) + if not groups: + return row_source_in(statement) + for group in groups: + source = row_source_keyword(group) + if source is not None: + return source + if contains(strip_parens(group), "VALUES"): + return None + return None + + +def set_operation_terms(statement: str, outer: str) -> Iterator[str]: + """The terms a top-level set operation joins. The operators are read from the text outside + every parenthesis, which `strip_parens` blanks in place rather than removing, so their + offsets are offsets into the statement itself and each term comes back from the original + text with its own parentheses intact. Reading them at that level is what keeps a set + operation written inside a `VALUES` list from cutting the list in half. An `ALL` or a + `DISTINCT` stays at the head of the term that follows, where it names no row source and + so reads as nothing.""" + edges = [0] + for operation in SET_OPERATION.finditer(outer): + edges += [operation.start(), operation.end()] + edges.append(len(statement)) + + for opens, closes in zip(edges[::2], edges[1::2]): + yield statement[opens:closes] + + +def parenthesised_groups(statement: str) -> Iterator[str]: + """What each group of parentheses closed at the statement's outermost level holds, in the + order they are written. One of them is where an `INSERT` keeps a row source it has + wrapped, since Postgres takes `INSERT INTO "t" ("a") (SELECT ...)` and `... (VALUES (1))` + alike, and reading a group on its own terms is what stops a scalar subquery nested inside + a wrapped `VALUES` list standing in for the rows.""" + depth = 0 + opens = None + + for index, character in enumerate(statement): + if character == "(": + if depth == 0: + opens = index + depth += 1 + elif character == ")": + depth = max(depth - 1, 0) + if depth == 0 and opens is not None: + yield statement[opens + 1 : index] + + +def row_source_in(text: str) -> str | None: + return next((word for word in ("SELECT", "TABLE") if contains(text, word)), None) + + +def hands_off_sql(statement: str, executed: frozenset[str]) -> bool: + """Whether a statement gives the server a string literal to run as SQL. `EXECUTE` runs one + outright, and so does `DO`, whose body is a string wherever it is not dollar-quoted. An + assignment parks one in a variable, which counts only when something further down runs + that variable by name, since a string the migration never executes is text.""" + if leads_with(statement, "EXECUTE") or leads_with(statement, "DO"): + return True + return bool(assigned_names(statement) & executed) + + +def assigned_names(statement: str) -> frozenset[str]: + """The candidate variable names a statement writes to. An assignment is read as every + word ahead of its operator, since a declaration carries its type and sometimes a leading + `DECLARE` alongside the name, and none of that is worth parsing when the only question + is which name is executed. A query assigns through the target list after its `INTO` + instead, and a loop through the variable it walks its query with, which is how a rewrite + reaches a variable with no operator appearing at all.""" + names = {word.lower() for word in assignment_reach(statement)} + + for targets in INTO_TARGETS.finditer(statement): + if names_a_table(statement[: targets.start()]): + continue + names.update(word.group().lower() for word in FIRST_WORD.finditer(targets.group(1))) + + names.update(loop.group(1).lower() for loop in LOOP_TARGET.finditer(statement)) + + return frozenset(names) + + +def names_a_table(before: str) -> bool: + """Whether the `INTO` this text runs up to introduces a table rather than a query's + target list. `INSERT INTO` is the one that does, and reading its table as somewhere a + string was parked would have an insert scanned for the SQL its own literals spell out. + An `INSERT` that really does assign reaches its `INTO` through a `RETURNING` list, so + the word immediately before is what separates the two.""" + word = PRECEDING_WORD.search(before) + return word is not None and word.group(1).upper() == "INSERT" + + +def names_a_relation(before: str) -> bool: + """Whether the parenthesised quoted identifier this text runs up to names a table with a + column list rather than opening a routine call. The two look alike, a name then a `(`, so + an uncalled routine sharing a name with a table would otherwise read as called. The word + immediately before tells most of them apart: `CREATE TABLE`, `INSERT INTO`, a foreign key's + `REFERENCES`, `CREATE TABLE IF NOT EXISTS`, and `COPY` each put a table there, and none can + precede a call. `ON` is the ambiguous one, since it introduces the table of a `CREATE INDEX` + but also a join condition that may itself be a call, so it counts only inside a statement + that creates an index, leaving `JOIN ... ON f()` and an index predicate's `WHERE f()` as + calls. A bare schema qualifier is read through: `INSERT INTO public."Foo"` parks the table's + introducing word a hop back behind `public.`, so any word ahead of the name that a dot follows, + touching or spaced as `public . "Foo"`, is the qualifier and the one before it decides. The + introducing word is settled before that, so a quoted schema, which blanks to spaces and leaves + `INTO` itself as the word ahead of the name however the dot is spaced, still reads as a relation, + while a genuine `SELECT public."f"()` reads through its qualifier to the `SELECT` and stays a call. + A word only introduces the name when nothing but whitespace and qualifier dots lies between them, + so a `(` in that gap keeps it from reaching across: a schema-qualified call inside a `CREATE INDEX` + expression, `ON "Foo" (public."f"(col))`, leaves `ON` behind the paren and the call stays a call.""" + word = PRECEDING_WORD.search(before) + if word is None: + return False + gap = before[word.end(1) :] + if QUALIFIER_GAP.fullmatch(gap): + keyword = word.group(1).upper() + if keyword in INTRODUCES_A_RELATION: + return True + if keyword == "ON": + return NAMES_AN_INDEX.search(before[before.rfind(";") + 1 :]) is not None + if "." in gap: + return names_a_relation(before[: word.start(1)]) + return False + + +def assignment_reach(statement: str) -> tuple[str, ...]: + """The words the statement's assignment is reached through, empty where it holds none. + PL/pgSQL spells the operator `:=` and takes a bare `=` as the same thing, so both count, + the second only where none of the words reached so far `marks_a_comparison`. The search + stops at the first operator that reads as an assignment, because a statement holds one + at most and everything after it is the expression being assigned, where an `=` only ever + compares: that is what keeps `ok := stmt = ''` from reading as a write to `stmt`. + What comes before can still be a comparison the assignment sits behind, as in + `IF n = 1 THEN stmt = ''`, and a word opening a block ends what it is reached + through, since nothing ahead of the `THEN` describes what follows it.""" + reached: list[str] = [] + compares = False + + for token in WORD_OR_ASSIGN.finditer(statement): + word = token.group().upper() + + if word == ":=": + return tuple(reached) + + if word == "=": + if not compares: + return tuple(reached) + continue + + if word in OPENS_A_BLOCK: + reached.clear() + compares = False + continue + + reached.append(word) + compares = compares or marks_a_comparison(word) + + return () + + +def marks_a_comparison(word: str) -> bool: + """Whether reaching a bare `=` through this word means the operator tests a variable + rather than writing one. These are all that tell the two apart: an assignment is reached + with a name and perhaps a type, while a comparison is reached either through a statement + carrying its own keyword or through a word that guards a condition.""" + return word in STATEMENT_KEYWORDS or word in GUARDS_A_CONDITION + + +def executed_names(masked: str) -> frozenset[str]: + """The variables handed to an `EXECUTE` by name. Reading these off the masked text keeps + an `EXECUTE` written inside a comment or a string from counting. Masking blanks a literal + in place rather than removing it, so `EXECUTE '...'` leaves whatever follows the literal + looking like the name being run. Only `INTO` and `USING` can sit there, since the syntax + allows nothing else between an `EXECUTE` and the semicolon ending it, and neither is ever + a variable, so both are dropped rather than left to collide with a query reaching one.""" + return frozenset( + match.group(1).lower() + for match in RUN_BY_NAME.finditer(masked) + if match.group(1).upper() not in NEVER_A_VARIABLE + ) + + +def leads_with(statement: str, keyword: str) -> bool: + word = leading_keyword(statement) + return word is not None and word.group().upper() == keyword + + +def contains(statement: str, keyword: str) -> bool: + return re.search(rf"\b{keyword}\b", statement, re.IGNORECASE) is not None + + +def read_markers(sql: str) -> Markers: + return Markers( + sql, + tuple( + Marker(match.start(), match.end(), alone_on_its_line(sql, match.start())) + for match in MARKER.finditer(sql) + ), + ) + + +def alone_on_its_line(sql: str, start: int) -> bool: + return not sql[sql.rfind("\n", 0, start) + 1 : start].strip() + + +def scan(sql: str, migration: str, markers: Markers) -> Iterator[Violation]: + yield from scan_region(sql, sql, migration, markers, 0) + + +def scan_region( + document: str, region: str, migration: str, markers: Markers, offset: int +) -> Iterator[Violation]: + """Violations in one region of `document`, whose text begins at `offset`. Positions are + always counted against the whole document, so a statement nested in a dollar-quoted body + reports its real file line and lines up with the markers read from that file. A single-quoted + literal that `DO` or `EXECUTE` runs as SQL has each doubled quote turned into a quote and a space + before it is scanned, so a `--` or `/*` in one of its nested strings blanks nothing and the + statement after it stays visible, and since that keeps every character on its offset, the + statement reports its true file line and lines up with the markers.""" + masked, bodies, literals, identifiers = mask(region) + executed = executed_names(masked) + runnable = executed_literals(masked, literals, executed) + + for match in STATEMENT.finditer(masked): + exempt = markers.exempt(offset + statement_start(match), offset + match.end()) + + for clause, base in clauses(match.group(), match.start()): + if hands_off_sql(clause, executed) and not exempt: + commands_end = base + bind_values_start(clause) + for start, end in literals: + if base <= start and end <= commands_end: + yield from scan_region( + document, + defuse_escapes(region[start:end]), + migration, + markers, + offset + start, + ) + + keyword = offending_keyword(clause) + if keyword is None or exempt: + continue + yield Violation(migration, line_of(document, offset + keyword_start(clause, base)), keyword) + + for body in bodies: + if not runs_when_applied(masked, region, bodies, runnable, identifiers, body): + continue + start, end = body + yield from scan_region(document, region[start:end], migration, markers, offset + start) + + +def executed_literals( + masked: str, literals: tuple[tuple[int, int], ...], executed: frozenset[str] +) -> tuple[tuple[int, int], ...]: + """The single-quoted literals a region runs as SQL, where a call to a routine the same + migration defines is as real as one written in the open. `DO '...'` runs its body and + `EXECUTE` runs the string it is handed, so a definition named inside one of those is called, + while a name in a message string or any literal nothing executes stays text. These are the + spans the direct scan already recurses into, read here so a call written in one is found when + the migration is searched for the routine's name.""" + return tuple( + (start, end) + for match in STATEMENT.finditer(masked) + for clause, base in clauses(match.group(), match.start()) + if hands_off_sql(clause, executed) + for start, end in literals + if base <= start and end <= base + bind_values_start(clause) + ) + + +def runs_when_applied( + masked: str, + region: str, + bodies: tuple[tuple[int, int], ...], + runnable: tuple[tuple[int, int], ...], + identifiers: tuple[tuple[int, int], ...], + body: tuple[int, int], +) -> bool: + """Whether a dollar-quoted body runs while the migration is being applied. A `DO` block runs + where it is written, and so does every other use of this quoting. A `CREATE FUNCTION` or a + `CREATE PROCEDURE` only stores its body, which runs when something calls the routine, so a + definition nothing calls rewrites no rows at boot and reporting it names a line that never + executes. Skipping every definition instead would let a migration define a backfill and then + run it unseen, which is the shape this check exists to catch, so the body is read whenever + the same migration names the routine anywhere outside the definition. The definition is + found in the masked text, where one written inside a comment has already been blanked, and + the name is read from the region at those same offsets, since masking blanks a quoted + identifier in place. A call written as a quoted identifier is blanked there too, and + `\"backfill\"()` is the same call as `backfill()` in Postgres, so the double-quoted call sites + are put back before the search and a routine invoked through one is found. A quoted name that + opens no call, a column or table sharing the routine's name, stays blanked and cannot be read + as a call it never makes. A definition whose + own name needs those quotes is read rather than trusted, since matching such a name once it is + put back in the open would be unreliable.""" + start, end = body + opens = masked.rfind(";", 0, start) + 1 + defined = DEFINES_A_ROUTINE.search(masked, opens, start) + if defined is None: + return True + named = ROUTINE_NAME.match(region, defined.end(), start) + if named is None or named.group(1).startswith('"'): + return True + restored = outside_definition(masked, region, bodies, runnable, identifiers, opens, end) + return contains(restored, re.escape(named.group(1))) + + +def outside_definition( + masked: str, + region: str, + bodies: tuple[tuple[int, int], ...], + runnable: tuple[tuple[int, int], ...], + identifiers: tuple[tuple[int, int], ...], + opens: int, + closes: int, +) -> str: + """The migration's text with one routine definition blanked out and every runnable body put + back: the dollar-quoted bodies and the single-quoted literals `DO` and `EXECUTE` run as SQL. + Masking blanks all of them alike, and a `DO` block, dollar-quoted or single-quoted, is the + ordinary way a migration runs a routine it has just defined, so a call written inside one has + to stay readable. Each comes back with its comments blanked, since a name written in a comment + is documentation rather than a call, while its string literals stay readable because `EXECUTE` + runs one as SQL and the call can be written inside it. A single-quoted payload is undoubled as + it goes back, so a `--` or `/*` in one of its nested strings blanks nothing and the call after + it stays visible, and it is padded to the span it fills so the later offsets still land. The + double-quoted call sites come back verbatim, so a routine invoked as `\"backfill\"()` reads as + the call it is, while a like-named identifier that opens no call was never collected and stays + blanked. The definition is blanked after they are restored, which takes its own body and + any identifier standing inside it with it, so a routine that names itself recursively does not + thereby count as called.""" + text = list(masked) + for start, end in bodies: + text[start:end] = without_comments(region[start:end]) + for start, end in runnable: + text[start:end] = without_comments(undouble(region[start:end])).ljust(end - start) + for start, end in identifiers: + text[start:end] = region[start:end] + text[opens:closes] = blank(region[opens:closes]) + return "".join(text) + + +def without_comments(sql: str) -> str: + """The text with its comments blanked in place and everything else kept, read with the same + lexing as `mask` so a `--` inside a string literal blanks nothing. A dollar-quoted body + nested within is read the same way on its own, which keeps a stray quote inside it from + reaching past its closing tag.""" + chunks: list[str] = [] + index = 0 + length = len(sql) + + while index < length: + pair = sql[index : index + 2] + + if pair == "--": + stop = sql.find("\n", index) + stop = length if stop == -1 else stop + chunks.append(blank(sql[index:stop])) + index = stop + continue + + if pair == "/*": + stop = skip_block_comment(sql, index) + chunks.append(blank(sql[index:stop])) + index = stop + continue + + character = sql[index] + + if character in "'\"": + stop = skip_quoted(sql, index, character) + chunks.append(sql[index:stop]) + index = stop + continue + + if character == "$": + tag = DOLLAR_TAG.match(sql, index) + if tag is not None: + closing = sql.find(tag.group(), tag.end()) + body_end = length if closing == -1 else closing + stop = length if closing == -1 else closing + len(tag.group()) + chunks.append(sql[index : tag.end()]) + chunks.append(without_comments(sql[tag.end() : body_end])) + chunks.append(sql[body_end:stop]) + index = stop + continue + + chunks.append(character) + index += 1 + + return "".join(chunks) + + +def clauses(statement: str, start: int) -> Iterator[tuple[str, int]]: + """The statements written inside one semicolon-delimited run, each with where it begins. A + `FOR ... LOOP` header takes no semicolon of its own, so the first statement of the loop body + is written into the same run, and reading the pair as one statement lets the header's row + source stand in as the keyword for both. That hides the statement the loop repeats, which is + the shape a row-by-row backfill takes. Splitting after each header, nested ones included, + reads the header and the body as the separate statements Postgres runs them as.""" + edges = (0, *(header.end() for header in LOOP_HEADER.finditer(statement)), len(statement)) + for opens, closes in zip(edges, edges[1:]): + if opens < closes: + yield statement[opens:closes], start + opens + + +def bind_values_start(statement: str) -> int: + """Where a statement stops handing commands to the server and starts listing bind values. + The expressions after `USING` are values substituted into the command, never commands in + their own right, so one that merely spells out a rewrite is not running it. Read off the + masked text, so a `USING` written inside the command string is not mistaken for this one, + and only once the parentheses have closed, so that the `USING` of a `JOIN` in a subquery + that helps build the command does not cut the command short and hide the rest of it.""" + for keyword in BIND_VALUES.finditer(statement): + preceding = statement[: keyword.start()] + if preceding.count("(") == preceding.count(")"): + return keyword.start() + return len(statement) + + +def statement_start(statement: re.Match[str]) -> int: + """Where the statement's own text begins, past the whitespace and blanked comments it picked + up from whatever sat between it and the statement before it, one of which can be a marker.""" + text = statement.group() + return statement.start() + len(text) - len(text.lstrip()) + + +def keyword_start(clause: str, base: int) -> int: + word = leading_keyword(clause) + return base + (0 if word is None else word.start()) + + +def line_of(sql: str, offset: int) -> int: + return sql.count("\n", 0, offset) + 1 + + +def scan_migration(directory: Path) -> tuple[Violation, ...]: + sql = (directory / "migration.sql").read_text(encoding="utf-8") + return tuple(scan(sql, directory.name, read_markers(sql))) + + +def stale_grandfathers(found: Mapping[str, tuple[Violation, ...]]) -> tuple[str, ...]: + clean = (name for name in GRANDFATHERED & found.keys() if not found[name]) + missing = GRANDFATHERED - found.keys() + return tuple(sorted((*clean, *missing))) + + +def main() -> int: + if not MIGRATIONS_DIR.is_dir(): + print(f"migrations directory not found: {MIGRATIONS_DIR}", file=sys.stderr) + return 2 + + directories = tuple(sorted(path for path in MIGRATIONS_DIR.iterdir() if (path / "migration.sql").is_file())) + found = {directory.name: scan_migration(directory) for directory in directories} + violations = tuple( + violation for name, results in found.items() if name not in GRANDFATHERED for violation in results + ) + + for violation in violations: + print(violation.render()) + + stale = stale_grandfathers(found) + for name in stale: + print(f"{name}: listed in GRANDFATHERED but no longer violates; remove it from the set") + + if violations: + print(f"\n{len(violations)} data-rewriting statement(s) in migrations.") + print(GUIDANCE) + + if violations or stale: + return 1 + + print(f"No data-rewriting statements in {len(directories)} migrations.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 5bd6326d8f2..b15a16ffc23 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -60,6 +60,8 @@ IGNORE_FUNCTIONS = [ "json_string_leaves", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); fails closed by raising at the cap so nothing goes unscanned. "with_json_string_leaves", # transitively bounded: only runs on a tree json_string_leaves already walked under the cap. "json_unrewritable_labels", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks. + "_flatten_form_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). + "_flatten_form_data_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). ] diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 15bd2c19ca9..cf912ddaa25 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -77,7 +77,7 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover The seam is `provider_edge.py`: `start_provider_edge` boots an in-process HTTP server (one shared instance per pytest process, `e2e_config.provider_edge_base` is the accessor) that mounts each supported provider under a path prefix (`EDGE_MOUNTS`: `/openai` -> `https://api.openai.com`, `/anthropic` -> `https://api.anthropic.com`). A test participates by registering its deployment with `api_base=provider_edge_base("openai")` plus the provider's path suffix; `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` is the reference. In live mode the accessor returns None and the deployment defaults to the real provider, so an edge-wired test runs in all three modes unchanged. Non-wired tests hit their providers live in every mode. The edge binds `E2E_PROVIDER_EDGE_BIND_HOST` (default 127.0.0.1) and advertises `E2E_PROVIDER_EDGE_ADVERTISE_HOST` in the api_base it hands out, for proxies running in containers -A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, `multipart/form-data` bodies store their ordinary fields plus a JSON list of the uploaded parts' `[field, filename, content-type]` triples and a digest of their content, so the per-request random boundary and the envelope never reach the key, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. `fixture_bundle.py` owns the format. Record serves the proxy the same filtered stored response replay will serve later, so the two modes are byte-identical from the proxy's side of the socket +A bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`) is a directory: `manifest.json` carries the record timestamp, harness git version, and format version, and each test gets a subdirectory holding one JSON file per provider call in call order (`0000-post-openai-v1-chat-completions.json`). Request headers are never stored (provider credentials never touch disk), non-JSON request bodies store a canonicalized sha256 digest instead of the bytes, `multipart/form-data` bodies store their ordinary fields plus a JSON list of the uploaded parts' `[field, filename, content-type]` triples and a digest of their content, so the per-request random boundary and the envelope never reach the key, and responses store status, filtered headers, and the verbatim body base64-encoded, which is part of why bundles are gitignored. Responses come in two shapes told apart by a `kind` tag: an ordinary one holding a single base64 body, and, for a response the provider streamed (`content-type: text/event-stream`), one holding its transfer chunks in order plus why the stream ended early if it did, so replay reproduces the split points the provider chose instead of one coalesced body. `fixture_bundle.py` owns the format, and `BUNDLE_FORMAT_VERSION` is checked on load, so a bundle recorded under older rules is refused by name rather than partially read. Record serves the proxy the same filtered stored response replay will serve later, chunk for chunk on a stream, so the two modes are byte-identical from the proxy's side of the socket Multipart identity is the fiddly corner, and the rules exist because each one had a collision behind it. A part counts as an upload when it carries a filename or declares its own content type, and everything else is an ordinary field. Field names get a `name[n]` suffix on repeats, with a literal `[` doubled first, so a form that repeats `purpose` never keys the same as one that literally sends `purpose[1]`. A field whose name reads as a credential is stored as ``, which stays key-preserving because the key is recomputed from the stored request rather than saved alongside it, so the live request carrying the real value still matches its redacted fixture. A field value that is not UTF-8 is stored as a base64 sha256 digest, base64 and not hex because the canonicalizer rewrites any 64-character hex run to `` and would fold every binary value onto one key. The uploaded parts contribute a JSON list rather than a `field:filename` string, so a separator inside a filename cannot impersonate a field boundary, and their byte length is stored for a reader's benefit but deliberately left out of the key, since the canonicalizer absorbs timestamp and id drift inside a file that changes its length @@ -87,16 +87,16 @@ A replayed response carries the recorded provider response id, and `LiteLLM_Spen The same id reuse reaches the managed-object tables. A replayed `/v1/files` or `/v1/batches` response carries the recorded provider object id, and `LiteLLM_ManagedObjectTable.model_object_id` is unique, so a unified batch create replayed against a database that still holds the record run's row fails on a Prisma unique-constraint violation, which surfaces as a 500, makes the router retry, and exhausts the recording. Replay the batches suite against a fresh database, or truncate `LiteLLM_ManagedObjectTable` and `LiteLLM_ManagedFileTable` before the run -Edge-wired today: `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` (the reference), `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic deployments in `llm_translation/test_messages_e2e.py` except the streaming test, and the OpenAI batch deployment behind `batches/` (`capabilities.openai_batch_params`). The mount base is not the same for both providers: OpenAI deployments register `f"{base}/v1"`, Anthropic deployments register `base` on its own, because litellm's Anthropic handler appends `/v1/messages` to `api_base` itself where the OpenAI handler appends only `/chat/completions`. Recording one suite locally is two runs against a proxy you already have up: +Edge-wired today: `quota_management/spend_tracking/test_provider_edge_spend_e2e.py` (the reference), `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic deployments in `llm_translation/test_messages_e2e.py` including the streaming test, and the OpenAI batch deployment behind `batches/` (`capabilities.openai_batch_params`). The mount base is not the same for both providers: OpenAI deployments register `f"{base}/v1"`, Anthropic deployments register `base` on its own, because litellm's Anthropic handler appends `/v1/messages` to `api_base` itself where the OpenAI handler appends only `/chat/completions`. Recording one suite locally is two runs against a proxy you already have up: ```bash E2E_FIXTURE_MODE=record E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 uv run pytest tests/e2e/llm_translation/test_chat_completions_contract_e2e.py E2E_FIXTURE_MODE=replay E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 uv run pytest tests/e2e/llm_translation/test_chat_completions_contract_e2e.py ``` -Point the proxy at bogus provider credentials for the replay run and it still has to pass: that is the whole proof that nothing left the process. Bundles are never committed. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and hard-fails after seven days, and publishing one for CI is LIT-5748 +Point the proxy at bogus provider credentials for the replay run and it still has to pass: that is the whole proof that nothing left the process. Bundles are never committed. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and hard-fails after seven days. CI records and replays this lane on a schedule in `.github/workflows/e2e_record_replay.yml`, publishing the bundle as a private `e2e-fixtures-bundle` artifact instead of committing it, selecting the tests with the `@pytest.mark.replayable` marker, and proving the bogus-credentials replay hermetic by counting provider egress with `.github/scripts/e2e_egress_sentinel.py` -Current limits: streaming chunk fidelity is LIT-5742 (a streamed response records as one buffered body), CI wiring is LIT-5748, Bedrock cannot be mounted (SigV4 signs the Host header, so a rewritten api_base fails signature verification), deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode +Current limits: Bedrock cannot be mounted (SigV4 signs the Host header, so a rewritten api_base fails signature verification), deployments baked into the proxy's config file cannot be edge-wired (only `/model/new` registrations can carry the edge api_base), and a file upload routed by `custom_llm_provider` through the proxy's `files_settings` block never passes a deployment at all, so the batches `model_param` and `provider_fallback` scenarios keep uploading live in every mode ## Typing diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 29778b06d7a..871d6b3904c 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -61,11 +61,15 @@ E2E_FIXTURE_MODE=record E2E_FIXTURE_DIR=/tmp/e2e-fixtures uv run pytest tests/e2 E2E_FIXTURE_MODE=replay E2E_FIXTURE_DIR=/tmp/e2e-fixtures uv run pytest tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py -v ``` -Bundles stay local. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and expires seven days after it was recorded, so record the suite you want before you replay it and never commit the result; publishing bundles for CI is LIT-5748 +Bundles stay local. `tests/e2e/.fixtures` is gitignored because a bundle holds verbatim provider response bodies and expires seven days after it was recorded, so record the suite you want before you replay it and never commit the result. CI keeps its bundle out of git too, as a private GitHub Actions artifact rather than a committed file, for the same reason + +In CI the `.github/workflows/e2e_record_replay.yml` lane runs record and replay on a schedule. A Saturday cron records the `replayable` marker's tests against the real providers and publishes the bundle as a private `e2e-fixtures-bundle` artifact carrying a SHA-256 sidecar; weekday crons pull that artifact by its pinned digest, verify the checksum before extracting, and replay it with provider credentials deliberately set to bogus values, so a run that ever reached a real provider would fail instead of passing. An egress sentinel (`.github/scripts/e2e_egress_sentinel.py`) pins the provider hostnames to a local sink for the whole replay job and counts every connection that reaches them, and the job asserts that count is zero, so hermeticity is proven by measurement rather than by an absent bill. A red Saturday publishes no bundle, so the next weekday finds nothing fresh and fails loudly rather than replaying a week-old recording, and the seven-day freshness gate hard-fails any bundle that has drifted too far from the live providers. Run the lane on demand from the Actions tab with the `mode` input: `record` re-records and republishes, `replay` replays the current bundle. A test joins the lane by carrying `@pytest.mark.replayable` on top of its edge wiring, so add that marker only to a test whose provider traffic actually replays with zero egress One sharp edge: a replayed response reuses the recorded provider response id, and that id is the primary key of `LiteLLM_SpendLogs`, so replaying against a database that still holds the record run's rows silently dedupes the spend writes and a spend assertion fails with zero rows. Run both commands above with `E2E_RESET_SPEND_LOGS=1` (and `DATABASE_URL` set in the pytest env) so each session truncates the spend log table after itself, or point replay at a fresh database -Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. The suites wired to the edge today are `quota_management/spend_tracking/test_provider_edge_spend_e2e.py`, `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the non-streaming Anthropic tests in `llm_translation/test_messages_e2e.py`, and the OpenAI batch deployment behind `batches/`. See `CLAUDE.md` in this directory for the bundle format, the edge design, and the current limits (streaming, Bedrock) +Another sharp edge, same root: record and replay derive every per-test token deterministically (the model name included, so a replay regenerates the exact requests the record run sent), which means an edge-wired deployment left in the database by an interrupted earlier run carries the same model name as the fresh one the current run registers. The proxy then holds two deployments under one model group and load-balances across both, and because the leftover's `api_base` points at the earlier run's edge process, which is gone, the calls that land on it fail with a connection error that reads like a transport bug rather than the stale row it is. Give each record or replay run a fresh database, or let a run finish so its own teardown deletes what it registered, and never reuse one long-lived proxy across back-to-back record/replay sessions. CI hands every job its own empty database and its own proxy, so it never sees this + +Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. The suites wired to the edge today are `quota_management/spend_tracking/test_provider_edge_spend_e2e.py`, `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic tests in `llm_translation/test_messages_e2e.py`, streamed and not, and the OpenAI batch deployment behind `batches/`. A streamed response replays as the chunk sequence the provider sent rather than one buffered body. See `CLAUDE.md` in this directory for the bundle format, the edge design, and the current limits (Bedrock). The scheduled CI record/replay lane is described above Tests marked `@pytest.mark.e2e` hard-fail when no proxy answers `/health/liveliness`, so a run that goes red with `No live proxy` at setup means the proxy isn't up; they never skip for a missing proxy, so an absent proxy can't be mistaken for a pass diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index f02d4eb4fe4..6d50cb436e2 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -81,6 +81,21 @@ File delete asserts `object=="file"` and `deleted==True`. | `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | | `conftest.py` | session-scoped batch deployment registration and teardown | | `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, per-backend content download, failure paths, second-hop routing, terminal state + cost | +| `test_managed_files_enforcement_e2e.py` | require_managed_files enforcement pins; deselected unless `E2E_MANAGED_FILES_STACK` is set (see below) | + +## require_managed_files enforcement (separate stack phase) + +`litellm_settings.require_managed_files` is a boot-time module global with no per-key +or runtime override, and turning it on 400s every upload that lacks +`target_model_names`, including the files_settings-routed `provider_fallback` +scenario above. So its pins cannot share a proxy with the rest of this suite: +`test_managed_files_enforcement_e2e.py` carries the `managed_files` marker, is +deselected unless `E2E_MANAGED_FILES_STACK` is set (the same pattern as the `weekly` +marker), and the PR gate runs it in a sequential phase after the main suite, against +the same ephemeral stack redeployed with the flag on. The pins: upload without +`target_model_names` is a 400, upload carrying a `model` param is a 400, a raw +provider file id on retrieve is a 400, and another user's managed unified id is a +403 while the owning user still retrieves it. ## Failure paths diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py index 73e8918e2ee..3b133fab680 100644 --- a/tests/e2e/batches/conftest.py +++ b/tests/e2e/batches/conftest.py @@ -12,12 +12,14 @@ the proxy config. from __future__ import annotations +import os from typing import Iterator import pytest from batch_client import BatchClient, build_client from capabilities import PROVIDERS +from e2e_config import MANAGED_FILES_OPT_IN_ENV from e2e_http import NoBody from proxy_client import ProxyClient @@ -29,6 +31,22 @@ def pytest_configure(config: pytest.Config) -> None: ) +def pytest_collection_modifyitems( + config: pytest.Config, items: list[pytest.Item] +) -> None: + if os.environ.get(MANAGED_FILES_OPT_IN_ENV): + return + deselected = [ + item for item in items if item.get_closest_marker("managed_files") is not None + ] + if not deselected: + return + config.hook.pytest_deselected(items=deselected) + items[:] = [ + item for item in items if item.get_closest_marker("managed_files") is None + ] + + @pytest.fixture(scope="session") def client(proxy: ProxyClient) -> BatchClient: return build_client(proxy) diff --git a/tests/e2e/batches/test_managed_files_enforcement_e2e.py b/tests/e2e/batches/test_managed_files_enforcement_e2e.py new file mode 100644 index 00000000000..7ad0b16adc3 --- /dev/null +++ b/tests/e2e/batches/test_managed_files_enforcement_e2e.py @@ -0,0 +1,118 @@ +"""Live e2e pins for litellm_settings.require_managed_files enforcement. + +require_managed_files is a boot-time module global, so these tests need a proxy +whose config enables it. The main ephemeral stack can never run with it on: the +flag would 400 every files_settings-routed upload in the rest of the suite. The +PR gate instead reconfigures the same stack sequentially after the main run and +executes only this file with E2E_MANAGED_FILES_STACK set; without that env every +test here is deselected (see conftest.py, mirroring the weekly marker). + +Pins: an upload without target_model_names is rejected 400, an upload that also +carries a model param is rejected 400, a raw provider file id is rejected 400 on +retrieve, and another user's managed unified file id is denied 403 while the +owning user still retrieves it. +""" + +from __future__ import annotations + +import json +from typing import Iterator + +import pytest + +from batch_client import BatchClient, FileObject +from capabilities import batch_model_name, is_managed_id, openai_batch_params +from e2e_config import unique_marker +from e2e_http import FileUploadForm, Result, UnknownApiError, unwrap +from lifecycle import ResourceManager + +pytestmark = [pytest.mark.e2e, pytest.mark.managed_files] + +UPLOAD_ROW = "llm.files.openai.require_managed_files_upload.nonstream.works" +ISOLATION_ROW = "llm.files.openai.require_managed_files_isolation.nonstream.works" + + +def batch_jsonl(model: str) -> bytes: + line = { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [{"role": "user", "content": "ping"}], + "max_tokens": 8, + }, + } + return (json.dumps(line) + "\n").encode() + + +def expect_api_error(result: Result[FileObject], status: int, needle: str) -> None: + match result: + case UnknownApiError(status_code=code, body=body) if code == status: + assert needle in body, f"expected {needle!r} in HTTP {status} body: {body[:300]}" + case _: + raise AssertionError(f"expected HTTP {status} containing {needle!r}, got: {result}") + + +@pytest.fixture(scope="module") +def managed_model(client: BatchClient) -> Iterator[str]: + model_name = batch_model_name("managed-files-openai") + model_id = client.create_model(model_name, openai_batch_params()) + yield model_name + client.delete_model(model_id) + + +@pytest.mark.covers(UPLOAD_ROW) +def test_upload_without_target_model_names_rejected( + client: BatchClient, scoped_key: str, managed_model: str +) -> None: + result = client.upload_file( + content=batch_jsonl(managed_model), + form=FileUploadForm(purpose="batch"), + key=scoped_key, + ) + expect_api_error(result, 400, "target_model_names is required") + + +@pytest.mark.covers(UPLOAD_ROW) +def test_upload_with_model_param_rejected( + client: BatchClient, scoped_key: str, managed_model: str +) -> None: + result = client.upload_file( + content=batch_jsonl(managed_model), + form=FileUploadForm(purpose="batch", target_model_names=managed_model), + model=managed_model, + key=scoped_key, + ) + expect_api_error(result, 400, "model is not allowed") + + +@pytest.mark.covers(ISOLATION_ROW) +def test_raw_provider_file_id_rejected(client: BatchClient, scoped_key: str) -> None: + result = client.retrieve_file("file-e2e-raw-provider-id", key=scoped_key) + expect_api_error(result, 400, "Raw provider file ids cannot be used") + + +@pytest.mark.covers(ISOLATION_ROW) +def test_cross_user_managed_id_denied_owner_allowed( + client: BatchClient, resources: ResourceManager, managed_model: str +) -> None: + run = unique_marker() + owner_key = resources.key(user_id=f"managed-files-owner-{run}") + other_key = resources.key(user_id=f"managed-files-other-{run}") + + uploaded = unwrap( + client.upload_file( + content=batch_jsonl(managed_model), + form=FileUploadForm(purpose="batch", target_model_names=managed_model), + key=owner_key, + ) + ) + resources.defer(lambda: client.delete_file(uploaded.id, key=owner_key)) + assert is_managed_id(uploaded.id), f"expected a managed unified file id, got {uploaded.id}" + + denied = client.retrieve_file(uploaded.id, key=other_key) + expect_api_error(denied, 403, "does not have access to this managed file") + + retrieved = unwrap(client.retrieve_file(uploaded.id, key=owner_key)) + assert retrieved.id == uploaded.id diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index dbe2d6e514e..e1b987cbfd9 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -43,6 +43,11 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "covers(cell_id, *, exercised_on=()): coverage-registry cell(s) this test covers", ) + config.addinivalue_line( + "markers", + "replayable: edge-wired test whose provider traffic replays from a fixture bundle, so it makes " + "zero provider calls in replay mode; the record/replay CI lane selects it with -m replayable", + ) config.addinivalue_line( "markers", "load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites", @@ -51,6 +56,10 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set", ) + config.addinivalue_line( + "markers", + "managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set", + ) def pytest_sessionstart(session: pytest.Session) -> None: diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index e6f08123b7c..47d296e61f3 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -45,6 +45,8 @@ - {id: llm.files.bedrock.upload.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "batches/capabilities.py:59", rationale: "Bedrock file upload to S3"} - {id: llm.files.gemini.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: gemini, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "Gemini Files API upload via proxy"} - {id: llm.files.hosted_vllm.upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: hosted_vllm, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "hosted_vllm OpenAI-compatible file upload"} +- {id: llm.files.openai.require_managed_files_upload.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "test_managed_files_enforcement_e2e.py / LIT-5902", rationale: "With require_managed_files enabled, an upload without target_model_names and an upload carrying a model param are both rejected 400; runs only in the sequential managed-files stack phase (E2E_MANAGED_FILES_STACK)"} +- {id: llm.files.openai.require_managed_files_isolation.nonstream.works, module: llm, tier: P1, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_managed_files_enforcement_e2e.py / LIT-5902", rationale: "With require_managed_files enabled, a raw provider file id is rejected 400 and another user's managed unified id is denied 403 while the owner still retrieves it; runs only in the managed-files stack phase"} - {id: llm.rerank.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rerank, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_rerank_e2e.py:29", rationale: "Cohere rerank, top_n + relevance_score"} - {id: llm.files.openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py", rationale: "GET /v1/files/{id}/content returns uploaded batch JSONL bytes"} - {id: llm.files.azure_openai.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: azure_openai, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on an Azure unified file returns the uploaded JSONL bytes verbatim"} diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 21a7a8c478a..21c5a338dc3 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -133,6 +133,7 @@ LOAD_MAX_SERIAL_LATENCY_SECONDS = float(os.environ.get("E2E_LOAD_MAX_SERIAL_LATE LOAD_MIN_CONCURRENCY_EFFICIENCY = float(os.environ.get("E2E_LOAD_MIN_CONCURRENCY_EFFICIENCY", "0.8")) WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY" +MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK" ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6")) ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6")) ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3")) diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 03f201e946e..bc76eb3ea7a 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -4,6 +4,11 @@ Enforced by tests/code_coverage_tests/check_e2e_no_raw_requests.py. Every reques body / query / header / response is a pydantic model; outcomes are a tagged union (``Result[R]``) so callers ``match`` on them instead of catching exceptions. +``forward`` relays one provider-bound request for the provider edge and buffers +the whole body; ``forward_stream`` relays the same request but hands back the +response head plus a lazy iterator over the upstream's own transfer chunks, which +is what lets a recording keep the split points a streamed response arrived on. + Named e2e_http (not http) so it does not shadow the stdlib ``http`` package that requests itself imports. """ @@ -12,7 +17,8 @@ from __future__ import annotations import time from collections.abc import Callable -from typing import Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast +from dataclasses import dataclass +from typing import Generator, Generic, Iterator, Literal, NewType, Protocol, TypeVar, cast import pytest import requests @@ -681,3 +687,88 @@ def forward( headers={name.lower(): value for name, value in resp.headers.items()}, body=resp.content, ) + + +@dataclass(frozen=True, slots=True) +class StreamChunk: + """One transfer chunk of a response body, exactly as the upstream framed it.""" + + data: bytes + + +@dataclass(frozen=True, slots=True) +class StreamTruncation: + """The body ended without its terminator, i.e. the upstream hung up mid-message. + Always the last step, and ``reason`` is the transport's own description of it.""" + + reason: str + + +type StreamStep = StreamChunk | StreamTruncation + + +@dataclass(frozen=True, slots=True) +class StreamHead: + """An upstream response whose head has arrived and whose body has not been read. + + A dataclass rather than a BaseModel because it owns a live socket: ``steps`` is + consumed once, in order, and closing it closes the underlying response.""" + + status_code: int + headers: dict[str, str] + steps: Generator[StreamStep, None, None] + + +def _stream_steps(resp: requests.Response) -> Generator[StreamStep, None, None]: + """The body as the upstream framed it, one step per transfer chunk. + + ``chunk_size=None`` is the whole point: urllib3 then returns exactly one piece + per wire chunk, so the provider's split points survive into the recording. Any + integer would re-slice the body into fixed-size pieces instead. Empty pieces are + dropped because a zero-length chunk is the terminator on the wire, and a failure + part way through becomes a final truncation step rather than an exception, since + the chunks already delivered are exactly what makes a mid-stream failure + different from a request that never streamed at all.""" + try: + for piece in cast("Iterator[bytes]", resp.iter_content(chunk_size=None)): + if piece: + yield StreamChunk(data=piece) + except requests.RequestException as exc: + yield StreamTruncation(reason=str(exc)) + finally: + resp.close() + + +def forward_stream( + method: str, + url: str, + *, + headers: dict[str, str], + body: bytes | None, + timeout: float = 60.0, +) -> StreamHead | NetworkError: + """Relay one provider-bound request for the provider edge and return as soon as + the response head arrives, with the body left unread behind ``StreamHead.steps``. + + Same contract as ``forward`` otherwise: no retries, no redirects, no schema. A + failure before the head arrives is still a ``NetworkError``; one raised while the + body streams arrives as the last step. With ``stream=True`` the timeout bounds + each socket read rather than the whole body, which is the right bound for a + stream and strictly more permissive for a long generation.""" + try: + resp = requests.request( + method, + url, + headers=headers, + data=body, + timeout=timeout, + allow_redirects=False, + stream=True, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return StreamHead( + status_code=resp.status_code, + headers={name.lower(): value for name, value in resp.headers.items()}, + steps=_stream_steps(resp), + ) diff --git a/tests/e2e/fixture_bundle.py b/tests/e2e/fixture_bundle.py index aa0ba100b6c..7c9dab1a687 100644 --- a/tests/e2e/fixture_bundle.py +++ b/tests/e2e/fixture_bundle.py @@ -6,15 +6,17 @@ per provider-bound interaction in call order. Bundles older than ``MAX_BUNDLE_AGE`` hard-fail replay at collection time (see conftest), so a green replay run can never certify against fixtures that have drifted more than a week from the live providers. Bump ``BUNDLE_FORMAT_VERSION`` whenever a change -moves recorded keys: a bundle recorded under the old rules then fails naming -both versions instead of quietly missing on every call. +moves recorded keys or changes the stored shape: a bundle recorded under the old +rules then fails naming both versions instead of quietly missing on every call. This module owns the format only. The provider-edge server that produces and -consumes it lives in provider_edge.py (LIT-5745) and the canonical match keys -it computes live in fixture_canonical.py (LIT-5741); streaming chunk fidelity -is a follow-up (LIT-5742). Every interaction file stores the full redacted -request because replay matches on its canonicalized content, and the response -as the raw HTTP status, filtered headers, and base64 body the provider sent. +consumes it lives in provider_edge.py (LIT-5745) and the canonical match keys it +computes live in fixture_canonical.py (LIT-5741). Every interaction file stores +the full redacted request because replay matches on its canonicalized content, +and a response in one of two shapes, told apart by their ``kind`` tag: an +ordinary ``RecordedHttpResponse`` holding one base64 body, or, for a response the +provider streamed, a ``RecordedStreamedResponse`` holding its transfer chunks in +order so replay reproduces the same split points (LIT-5742). """ from __future__ import annotations @@ -26,11 +28,11 @@ import subprocess from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Final +from typing import Annotated, Final, Literal -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, Field, JsonValue -BUNDLE_FORMAT_VERSION: Final = 3 +BUNDLE_FORMAT_VERSION: Final = 4 MAX_BUNDLE_AGE: Final = timedelta(days=7) MANIFEST_FILENAME: Final = "manifest.json" @@ -74,14 +76,38 @@ class RecordedHttpResponse(BaseModel): volatile entries (see provider_edge.py), and the body as base64 so binary payloads survive JSON.""" + kind: Literal["http"] = "http" status_code: int headers: dict[str, str] body_b64: str +class RecordedStreamedResponse(BaseModel): + """A response the provider streamed, kept chunk by chunk instead of buffered. + + ``chunks_b64`` holds one entry per upstream transfer chunk, in order, so replay + reproduces the split points the provider chose rather than one coalesced body. + ``truncated`` is None for a stream that reached its terminator and otherwise + says why it did not, prefixed by which side ended it (``upstream:`` for a + provider that hung up mid-stream, ``downstream:`` for a proxy that stopped + reading). Replay behaves the same for any truncation, delivering the recorded + chunks and then closing; the reason is there for whoever reads the bundle.""" + + kind: Literal["streamed"] = "streamed" + status_code: int + headers: dict[str, str] + chunks_b64: list[str] + truncated: str | None = None + + +type RecordedResponse = Annotated[ + RecordedHttpResponse | RecordedStreamedResponse, Field(discriminator="kind") +] + + class Interaction(BaseModel): request: RecordedRequest - response: RecordedHttpResponse + response: RecordedResponse def slugify(raw: str, *, limit: int = 60) -> str: @@ -128,7 +154,7 @@ class BundleRecorder: root: Path _ordinals: dict[str, int] = field(default_factory=dict) - def record(self, *, test_key: str, request: RecordedRequest, response: RecordedHttpResponse) -> None: + def record(self, *, test_key: str, request: RecordedRequest, response: RecordedResponse) -> None: slug = slug_for_test(test_key) ordinal = self._ordinals.get(slug, 0) self._ordinals[slug] = ordinal + 1 @@ -200,14 +226,27 @@ def _read_manifest(root: Path) -> Manifest | UnreadableBundle: return UnreadableBundle(reason=f"{MANIFEST_FILENAME} is invalid: {exc}") -def check_freshness(root: Path, *, now: datetime) -> BundleFreshness: +def _supported_manifest(root: Path) -> Manifest | UnreadableBundle: + """The manifest, refused when it was written under a different format version. + A bundle is atomic (record wipes and rewrites the whole directory and never + merges), so a foreign version is a hard reject rather than a partial read.""" manifest = _read_manifest(root) if isinstance(manifest, UnreadableBundle): return manifest if manifest.format_version != BUNDLE_FORMAT_VERSION: return UnreadableBundle( - reason=f"format_version {manifest.format_version} != supported {BUNDLE_FORMAT_VERSION}" + reason=( + f"format_version {manifest.format_version} != supported {BUNDLE_FORMAT_VERSION}; " + "re-record with E2E_FIXTURE_MODE=record" + ) ) + return manifest + + +def check_freshness(root: Path, *, now: datetime) -> BundleFreshness: + manifest = _supported_manifest(root) + if isinstance(manifest, UnreadableBundle): + return manifest recorded_at = ( manifest.recorded_at if manifest.recorded_at.tzinfo is not None @@ -231,7 +270,7 @@ class LoadedBundle: def load_bundle(root: Path) -> LoadedBundle | UnreadableBundle: - manifest = _read_manifest(root) + manifest = _supported_manifest(root) if isinstance(manifest, UnreadableBundle): return manifest interactions = { diff --git a/tests/e2e/gateway/record_replay_ci_config.yml b/tests/e2e/gateway/record_replay_ci_config.yml new file mode 100644 index 00000000000..08972969cf0 --- /dev/null +++ b/tests/e2e/gateway/record_replay_ci_config.yml @@ -0,0 +1,3 @@ +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY + store_model_in_db: true diff --git a/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py b/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py index 114beaae2fb..09b484eb120 100644 --- a/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py +++ b/tests/e2e/llm_translation/test_chat_completions_contract_e2e.py @@ -13,7 +13,7 @@ from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody from proxy_client import ProxyClient from pydantic import BaseModel -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.replayable] OPENAI_BACKEND = "openai/gpt-4o-mini" CHAT_PATH = "/chat/completions" diff --git a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py index 265cc202ff4..f951eb328f5 100644 --- a/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py +++ b/tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py @@ -40,6 +40,7 @@ def _openai_embeddings_params() -> LiteLLMParamsBody: class TestEmbeddingsEndpoint: + @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") def test_embeddings_returns_vector( self, endpoints_client: EndpointsClient, resources: ResourceManager @@ -109,6 +110,7 @@ class TestEmbeddingsEndpoint: f"embedding vector is all zeros: {result.body[:300]}" ) + @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.basic.nonstream.works") def test_array_input_returns_vectors( self, endpoints_client: EndpointsClient, resources: ResourceManager @@ -129,6 +131,7 @@ class TestEmbeddingsEndpoint: parsed = EmbeddingsResult.model_validate_json(result.body) assert len(parsed.data) == 3, f"expected 3 vectors: {result.body[:300]}" + @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.input_validation.nonstream.works") def test_missing_model_returns_client_error( self, endpoints_client: EndpointsClient, resources: ResourceManager @@ -141,6 +144,7 @@ class TestEmbeddingsEndpoint: ) assert_client_error(result, "embeddings missing model") + @pytest.mark.replayable @pytest.mark.covers("llm.embeddings.openai.input_validation.nonstream.works") def test_missing_input_returns_error( self, endpoints_client: EndpointsClient, resources: ResourceManager diff --git a/tests/e2e/llm_translation/test_messages_e2e.py b/tests/e2e/llm_translation/test_messages_e2e.py index 7f81a5e3946..e0bedd72eac 100644 --- a/tests/e2e/llm_translation/test_messages_e2e.py +++ b/tests/e2e/llm_translation/test_messages_e2e.py @@ -24,7 +24,7 @@ from models import ( ) from pydantic import BaseModel -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.replayable] class _OptionalMessagesBody(BaseModel): @@ -33,6 +33,25 @@ class _OptionalMessagesBody(BaseModel): max_tokens: int | None = None +class _MessagesEventDelta(BaseModel): + text: str = "" + + +class _MessagesEventUsage(BaseModel): + output_tokens: int | None = None + + +class _MessagesStreamEvent(BaseModel): + """One Anthropic SSE event, keeping only what the stream's shape is asserted on. + + ``delta.text`` is populated on ``content_block_delta`` and absent on the + ``message_delta`` that closes the turn, which is the event carrying ``usage``.""" + + type: str + delta: _MessagesEventDelta | None = None + usage: _MessagesEventUsage | None = None + + ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5" WEATHER_TOOL = AnthropicCustomTool( @@ -137,13 +156,14 @@ class TestAnthropicMessages: def test_messages_streams_completion( self, endpoints_client: EndpointsClient, resources: ResourceManager ) -> None: - """Stays on a live Anthropic deployment in every mode: the edge buffers a - streamed response into one body, so chunk fidelity waits on LIT-5742.""" - model, key = self._register( - endpoints_client, - resources, - LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY"), - ) + """Edge-wired like its non-streaming siblings, so record and replay both + carry the streamed response. + + Asserts the shape of the event sequence, not just that deltas and a stop + appeared somewhere in it: the answer arrives across several deltas, and the + usage event sits between the last of them and ``message_stop``. A replay that + coalesced the response into one buffered body could not satisfy either.""" + model, key = self._register(endpoints_client, resources) result = endpoints_client.proxy.messages_stream( key, @@ -151,18 +171,42 @@ class TestAnthropicMessages: model=model, max_tokens=64, stream=True, - messages=[ChatMessage(role="user", content="Count from one to three.")], + messages=[ChatMessage(role="user", content="Count from 1 to 20, one number per line.")], ), ) require_successful_call(result) assert result.is_streaming, f"response was not streamed: {result.headers}" assert not result.stream_error, f"stream errored: {result.stream_error}" assert result.stream_events, "stream produced no SSE events" - assert any("content_block_delta" in event for event in result.stream_events), ( - "stream carried no content deltas" + + events = [ + _MessagesStreamEvent.model_validate_json(event) for event in result.stream_events + ] + types = [event.type for event in events] + delta_positions = [ + index for index, event in enumerate(events) if event.type == "content_block_delta" + ] + assert len(delta_positions) >= 2, ( + f"stream carried {len(delta_positions)} content deltas, so it was not " + f"incremental: {types}" ) - assert any("message_stop" in event for event in result.stream_events), ( - "stream never reached message_stop" + text = "".join( + event.delta.text + for event in events + if event.type == "content_block_delta" and event.delta is not None + ) + assert text.strip(), f"content deltas assembled to no text: {result.stream_events[:5]}" + + usage_positions = [ + index + for index, event in enumerate(events) + if event.type == "message_delta" and event.usage is not None + ] + assert usage_positions, f"stream never reported usage: {types}" + assert "message_stop" in types, f"stream never reached message_stop: {types}" + stop_position = types.index("message_stop") + assert delta_positions[-1] < usage_positions[0] < stop_position, ( + f"usage did not land between the last content delta and message_stop: {types}" ) @pytest.mark.covers("llm.messages.anthropic.tool_use.nonstream.works") diff --git a/tests/e2e/provider_edge.py b/tests/e2e/provider_edge.py index 25a1e8043ed..ceb695ffcd6 100644 --- a/tests/e2e/provider_edge.py +++ b/tests/e2e/provider_edge.py @@ -19,10 +19,21 @@ headers must never touch disk. An unmatched replay call returns HTTP ``REPLAY_MISS_STATUS`` naming the closest recorded interaction, which the proxy relays as a provider error the failing test surfaces. +A response the provider streamed (one whose content type names +``text/event-stream``) is relayed and stored chunk by chunk instead of buffered +(LIT-5742): the edge reads one piece per upstream transfer chunk, writes each +one downstream in chunked framing as it arrives, and records the sequence, so +replay hands the proxy the same number of chunks split in the same places. A +provider that hangs up mid-stream is recorded as the chunks it did deliver plus +a truncation, and replays as those chunks followed by a connection close with no +terminator, which is the same incomplete chunked read the live failure produced +rather than a clean 502 that erases it. Everything else keeps the buffered +shape, byte for byte, framed with a content-length as before. + v1 limits: only the mounts in ``EDGE_MOUNTS`` (SigV4 providers like Bedrock -sign the Host header, so a forwarding edge breaks their signatures), streaming -fidelity is LIT-5742, and CI wiring is LIT-5748. Suites that do not wire the -edge keep hitting providers live in every mode. +sign the Host header, so a forwarding edge breaks their signatures), and CI +wiring is LIT-5748. Suites that do not wire the edge keep hitting providers +live in every mode. """ from __future__ import annotations @@ -34,24 +45,34 @@ import hashlib import re import threading from collections import deque -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from contextlib import closing from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from itertools import islice from pathlib import Path from types import MappingProxyType -from typing import Final, Literal, assert_never +from typing import Final, Generator, Literal, assert_never from urllib.parse import parse_qsl, urlsplit from pydantic import JsonValue, TypeAdapter -from e2e_http import NetworkError, RawResponse, forward +from e2e_http import ( + NetworkError, + StreamChunk, + StreamHead, + StreamStep, + StreamTruncation, + forward_stream, +) from fixture_bundle import ( BundleRecorder, Interaction, LoadedBundle, RecordedHttpResponse, RecordedRequest, + RecordedResponse, + RecordedStreamedResponse, UnreadableBundle, UnsafeBundleDir, interaction_filename, @@ -479,11 +500,28 @@ type EdgeBackend = RecordEdge | ReplayEdge @dataclass(frozen=True, slots=True) class EdgeReply: + """A whole response the edge already holds: written with a content-length.""" + status_code: int headers: dict[str, str] body: bytes +@dataclass(frozen=True, slots=True) +class EdgeStream: + """A response the edge relays chunk by chunk: written in chunked framing, one + transfer chunk per step, so the split points reach the proxy intact. Record and + replay both produce one of these, driven by different step sources, which is + what makes their framing identical by construction rather than by inspection.""" + + status_code: int + headers: dict[str, str] + steps: Generator[StreamStep, None, None] + + +type EdgeOutcome = EdgeReply | EdgeStream + + def _text_reply(status_code: int, message: str) -> EdgeReply: return EdgeReply( status_code=status_code, @@ -492,34 +530,87 @@ def _text_reply(status_code: int, message: str) -> EdgeReply: ) -def _reply_from_recorded(response: RecordedHttpResponse) -> EdgeReply: - return EdgeReply( - status_code=response.status_code, - headers=dict(response.headers), - body=base64.b64decode(response.body_b64), +def _recorded_steps( + chunks_b64: Sequence[str], truncated: str | None +) -> Generator[StreamStep, None, None]: + """Replay's step source: the recorded chunks in recorded order, as fast as the + socket takes them (inter-chunk delays are deliberately not reproduced), then the + recorded truncation if the stream ended without a terminator.""" + for chunk in chunks_b64: + yield StreamChunk(data=base64.b64decode(chunk)) + if truncated is not None: + yield StreamTruncation(reason=truncated) + + +def _recorded_outcome(response: RecordedResponse) -> EdgeOutcome: + match response: + case RecordedHttpResponse(status_code=status_code, headers=headers, body_b64=body_b64): + return EdgeReply( + status_code=status_code, + headers=dict(headers), + body=base64.b64decode(body_b64), + ) + case RecordedStreamedResponse( + status_code=status_code, headers=headers, chunks_b64=chunks_b64, truncated=truncated + ): + return EdgeStream( + status_code=status_code, + headers=dict(headers), + steps=_recorded_steps(chunks_b64, truncated), + ) + case _: + assert_never(response) + + +def _filtered_response_headers(headers: Mapping[str, str]) -> dict[str, str]: + """What the edge stores and serves: the provider's headers minus hop-by-hop and + volatile entries. Framing headers are in that set, so a stored header can never + contradict the framing the edge chooses when it serves the response.""" + return { + name: value for name, value in headers.items() if name not in _RESPONSE_DROPPED_HEADERS + } + + +def _network_error_response(message: str) -> RecordedHttpResponse: + return RecordedHttpResponse( + status_code=502, + headers={"content-type": "text/plain; charset=utf-8"}, + body_b64=base64.b64encode( + f"provider edge could not reach the provider: {message}".encode() + ).decode("ascii"), ) -def _recorded_response(outcome: RawResponse | NetworkError) -> RecordedHttpResponse: - match outcome: - case RawResponse(status_code=status_code, headers=headers, body=body): - return RecordedHttpResponse( - status_code=status_code, - headers={ - name: value - for name, value in headers.items() - if name not in _RESPONSE_DROPPED_HEADERS - }, - body_b64=base64.b64encode(body).decode("ascii"), - ) - case NetworkError(message=message): - return RecordedHttpResponse( - status_code=502, - headers={"content-type": "text/plain; charset=utf-8"}, - body_b64=base64.b64encode( - f"provider edge could not reach the provider: {message}".encode() - ).decode("ascii"), - ) +def _buffered_response( + status_code: int, headers: Mapping[str, str], body: bytes +) -> RecordedHttpResponse: + return RecordedHttpResponse( + status_code=status_code, + headers=_filtered_response_headers(headers), + body_b64=base64.b64encode(body).decode("ascii"), + ) + + +def _streamed_response( + status_code: int, headers: Mapping[str, str], chunks: Sequence[bytes], truncated: str | None +) -> RecordedStreamedResponse: + return RecordedStreamedResponse( + status_code=status_code, + headers=_filtered_response_headers(headers), + chunks_b64=[base64.b64encode(chunk).decode("ascii") for chunk in chunks], + truncated=truncated, + ) + + +def _is_streamed(headers: Mapping[str, str]) -> bool: + """Whether a response is one to relay incrementally, decided by content type. + + ``transfer-encoding: chunked`` would be the wrong signal: chunking is a + transport choice providers make freely for ordinary JSON, so keying off it would + move nearly every recording to the streamed shape for no gain. The content type + is the header that says "consume this as it arrives", and it is already how the + harness defines streaming everywhere else.""" + return "text/event-stream" in _header_value(headers, "content-type").lower() def _upstream_url(upstream_base: str, upstream_path: str, query: str) -> str: @@ -527,6 +618,75 @@ def _upstream_url(upstream_base: str, upstream_path: str, query: str) -> str: return f"{url}?{query}" if query else url +def _persist( + backend: RecordEdge, test_key: str, request: RecordedRequest, response: RecordedResponse +) -> None: + with backend.lock: + backend.recorder.record(test_key=test_key, request=request, response=response) + + +def _recording_steps( + backend: RecordEdge, test_key: str, request: RecordedRequest, head: StreamHead +) -> Generator[StreamStep, None, None]: + """Record mode's step source: hand each upstream chunk downstream and record it + only once that write has returned, then persist the whole sequence once, under + the lock. Relaying incrementally keeps record exercising the proxy's incremental + parser the way a live run does. + + A chunk is appended after its ``yield`` returns, so a downstream that hangs up + mid-relay records exactly the chunks it took and never the one whose write + raised. The ``except`` covers that downstream close and the proxy hanging up + mid-stream; either way the ``finally`` persists what arrived, marked truncated, + because recording a cut-short stream as a clean one would let a later replay + serve a well-terminated fraction of the response and pass a test that should + have gone red.""" + collected: list[bytes] = [] + truncated: str | None = None + try: + with closing(head.steps) as steps: + for step in steps: + match step: + case StreamChunk(): + pass + case StreamTruncation(reason=reason): + truncated = f"upstream: {reason}" + case _: + assert_never(step) + yield step + if isinstance(step, StreamChunk): + collected.append(step.data) + except GeneratorExit: + if truncated is None: + truncated = f"downstream: relay closed after {len(collected)} chunks" + raise + finally: + _persist( + backend, + test_key, + request, + _streamed_response(head.status_code, head.headers, collected, truncated), + ) + + +def _drain_to_response(head: StreamHead) -> RecordedHttpResponse: + """A response the detection rule did not call streamed: drain the same step + iterator, join the pieces, and store today's buffered shape byte for byte. A + truncation part way through degrades to the synthetic 502 exactly as the eager + read did, because storing half a JSON body under a content-length as though it + were whole would be a worse lie than failing.""" + pieces: list[bytes] = [] + with closing(head.steps) as steps: + for step in steps: + match step: + case StreamChunk(data=data): + pieces.append(data) + case StreamTruncation(reason=reason): + return _network_error_response(reason) + case _: + assert_never(step) + return _buffered_response(head.status_code, head.headers, b"".join(pieces)) + + def _handle_record( backend: RecordEdge, request: RecordedRequest, @@ -536,23 +696,37 @@ def _handle_record( headers: Mapping[str, str], body: bytes | None, timeout: float, -) -> EdgeReply: +) -> EdgeOutcome: + test_key: Final = current_test_key() forwarded: Final = { name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS } - outcome: Final = forward(method, url, headers=forwarded, body=body, timeout=timeout) - response: Final = _recorded_response(outcome) - with backend.lock: - backend.recorder.record(test_key=current_test_key(), request=request, response=response) - return _reply_from_recorded(response) + head: Final = forward_stream(method, url, headers=forwarded, body=body, timeout=timeout) + match head: + case NetworkError(message=message): + unreachable: Final = _network_error_response(message) + _persist(backend, test_key, request, unreachable) + return _recorded_outcome(unreachable) + case StreamHead() if _is_streamed(head.headers): + return EdgeStream( + status_code=head.status_code, + headers=_filtered_response_headers(head.headers), + steps=_recording_steps(backend, test_key, request, head), + ) + case StreamHead(): + buffered: Final = _drain_to_response(head) + _persist(backend, test_key, request, buffered) + return _recorded_outcome(buffered) + case _: + assert_never(head) -def _handle_replay(source: ReplaySource, request: RecordedRequest) -> EdgeReply: +def _handle_replay(source: ReplaySource, request: RecordedRequest) -> EdgeOutcome: try: interaction: Final = source.next_interaction(request) except ReplayMiss as miss: return _text_reply(REPLAY_MISS_STATUS, str(miss)) - return _reply_from_recorded(interaction.response) + return _recorded_outcome(interaction.response) def handle_edge_request( @@ -564,7 +738,7 @@ def handle_edge_request( body: bytes | None, *, timeout: float, -) -> EdgeReply: +) -> EdgeOutcome: """The edge's pure core, one HTTP exchange in and out: resolve the mount prefix, then record (forward + persist) or replay (serve from the bundle). Socket-free so unit tests exercise every branch without a server.""" @@ -618,7 +792,7 @@ class _EdgeHandler(BaseHTTPRequestHandler): assert isinstance(edge_server, _EdgeHTTPServer) length: Final = int(self.headers.get("content-length") or "0") body: Final = self.rfile.read(length) if length else None - reply: Final = handle_edge_request( + outcome: Final = handle_edge_request( edge_server.backend, edge_server.mounts, self.command, @@ -627,6 +801,15 @@ class _EdgeHandler(BaseHTTPRequestHandler): body, timeout=edge_server.forward_timeout, ) + match outcome: + case EdgeReply(): + self._write_reply(outcome) + case EdgeStream(): + self._write_stream(outcome) + case _: + assert_never(outcome) + + def _write_reply(self, reply: EdgeReply) -> None: self.send_response(reply.status_code) for name, value in reply.headers.items(): self.send_header(name, value) @@ -634,6 +817,32 @@ class _EdgeHandler(BaseHTTPRequestHandler): self.end_headers() self.wfile.write(reply.body) + def _write_stream(self, stream: EdgeStream) -> None: + """Write a streamed outcome in chunked framing, one transfer chunk per step. + + ``wbufsize`` is 0 on BaseHTTPRequestHandler, so ``wfile`` sends each write + straight down the socket and no flush is needed. A truncation step ends the + message without its terminator and closes the connection, which the stdlib + shuts down write-side first: the proxy sees a graceful close mid-message, + which is the incomplete chunked read a provider hanging up produces, and not + the reset that could discard the chunks already in flight.""" + self.send_response(stream.status_code) + for name, value in stream.headers.items(): + self.send_header(name, value) + self.send_header("transfer-encoding", "chunked") + self.end_headers() + with closing(stream.steps) as steps: + for step in steps: + match step: + case StreamChunk(data=data): + self.wfile.write(b"%x\r\n%s\r\n" % (len(data), data)) + case StreamTruncation(): + self.close_connection = True + return + case _: + assert_never(step) + self.wfile.write(b"0\r\n\r\n") + def log_message(self, format: str, *args: object) -> None: """Silence the per-request stderr line BaseHTTPRequestHandler emits.""" diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 8feb4505ce3..c3f8865f218 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -5,5 +5,7 @@ addopts = --strict-markers --strict-config --reruns 1 --only-rerun "kind='network'" --only-rerun "status_code=5[0-9][0-9]" markers = e2e: live test that requires a running proxy and real provider keys + replayable: edge-wired test whose provider traffic replays from a fixture bundle, so it makes zero provider calls in replay mode; the record/replay CI lane selects it with -m replayable load: heavy throughput/load test; collected last so it never perturbs latency-sensitive suites weekly: real-provider anomaly load test that spends real money; deselected unless E2E_WEEKLY_ANOMALY is set + managed_files: needs a proxy running with require_managed_files enabled; deselected unless E2E_MANAGED_FILES_STACK is set diff --git a/tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py index ced7c819d42..4931af4222d 100644 --- a/tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_provider_edge_spend_e2e.py @@ -18,7 +18,7 @@ from lifecycle import ResourceManager from models import LiteLLMParamsBody from spend_e2e_client import SpendClient, unique_marker, unwrap -pytestmark = pytest.mark.e2e +pytestmark = [pytest.mark.e2e, pytest.mark.replayable] @pytest.mark.covers("quota_management.spend_tracking.chat_completions.logs_cost") diff --git a/tests/e2e/test_fixture_bundle.py b/tests/e2e/test_fixture_bundle.py index b49ab565e39..c01d0b34fb5 100644 --- a/tests/e2e/test_fixture_bundle.py +++ b/tests/e2e/test_fixture_bundle.py @@ -22,6 +22,7 @@ from fixture_bundle import ( Manifest, RecordedHttpResponse, RecordedRequest, + RecordedStreamedResponse, StaleBundle, UnreadableBundle, UnsafeBundleDir, @@ -186,3 +187,45 @@ class TestRecordAndLoad: slug_for_test("suite/test_a.py::test_one"), slug_for_test("suite/test_b.py::test_two"), } + + def test_a_streamed_response_round_trips_through_the_bundle(self, tmp_path: Path) -> None: + """LIT-5742: the two response shapes share one file format and are told apart + by their ``kind`` tag, so a streamed recording comes back with its chunk list + intact rather than as a buffered response with an empty body.""" + root = tmp_path / "bundle" + recorder = prepared(root) + key = "suite/test_mod.py::test_streamed" + recorder.record( + test_key=key, + request=plain_request("/messages"), + response=RecordedStreamedResponse( + status_code=200, + headers={"content-type": "text/event-stream"}, + chunks_b64=["Zmly", "c3Q="], + truncated="upstream: hung up", + ), + ) + loaded = load_bundle(root) + assert isinstance(loaded, LoadedBundle) + (interaction,) = loaded.interactions[slug_for_test(key)] + response = interaction.response + assert isinstance(response, RecordedStreamedResponse) + assert response.chunks_b64 == ["Zmly", "c3Q="] + assert response.truncated == "upstream: hung up" + + def test_load_bundle_rejects_a_foreign_format_version(self, tmp_path: Path) -> None: + """A bundle is written atomically, so a manifest from another format version + means every response inside it may have a shape this code cannot read. Loading + has to refuse it by name, the way the freshness gate does, rather than parse + what it happens to understand.""" + root = tmp_path / "bundle" + prepared(root).record( + test_key="suite/test_mod.py::test_old", + request=plain_request("/chat"), + response=plain_response(), + ) + write_manifest(root, NOW, format_version=BUNDLE_FORMAT_VERSION - 1) + loaded = load_bundle(root) + assert isinstance(loaded, UnreadableBundle) + assert f"format_version {BUNDLE_FORMAT_VERSION - 1}" in loaded.reason + assert "E2E_FIXTURE_MODE=record" in loaded.reason diff --git a/tests/e2e/test_provider_edge.py b/tests/e2e/test_provider_edge.py index 14a9fd53393..8ab389ee43c 100644 --- a/tests/e2e/test_provider_edge.py +++ b/tests/e2e/test_provider_edge.py @@ -11,12 +11,20 @@ computed and closest recorded canonical keys (LIT-5741; the pure canonicalizer is pinned in test_fixture_canonical.py). Requests are made through ``e2e_http.forward`` so the whole HTTP surface of the edge is exercised; the pure ``handle_edge_request`` core is pinned socket-free alongside. + +Streaming fidelity (LIT-5742) is pinned at the transfer layer, because that is +the only layer where it is visible: a chunked provider sends a known list of +transfer chunks, one of which deliberately splits an SSE event mid-token, and a +raw-socket client reads the edge's own reply back as HTTP chunks. Counting SSE +events at the client would prove nothing, since a coalesced body carries the +same events as a chunk-per-event one. """ from __future__ import annotations import base64 import json +import socket import threading from collections.abc import Generator, Mapping from concurrent.futures import ThreadPoolExecutor @@ -28,7 +36,7 @@ from typing import Final import pytest from pydantic import TypeAdapter -from e2e_http import RawResponse, forward +from e2e_http import RawResponse, StreamChunk, forward from fixture_canonical import canonicalize from fixture_bundle import ( BundleRecorder, @@ -36,6 +44,7 @@ from fixture_bundle import ( LoadedBundle, RecordedHttpResponse, RecordedRequest, + RecordedStreamedResponse, load_bundle, prepare_bundle, slug_for_test, @@ -44,6 +53,8 @@ from fixture_mode import current_test_key from provider_edge import ( REPLAY_MISS_STATUS, EdgeBackend, + EdgeReply, + EdgeStream, ProviderEdge, RecordEdge, ReplayEdge, @@ -116,10 +127,175 @@ def fake_provider() -> Generator[_FakeProvider]: server.server_close() -def provider_url(server: _FakeProvider) -> str: +def provider_url(server: ThreadingHTTPServer) -> str: return f"http://127.0.0.1:{server.server_address[1]}" +STREAM_PATH = "/openai/v1/messages" +STREAM_BODY = json.dumps({"model": "claude", "stream": True}).encode() +MID_EVENT_HEAD = b'data: {"type":"content_bl' +MID_EVENT_TAIL = b'ock_delta","delta":{"text":" two"}}\n\n' +SSE_CHUNKS: tuple[bytes, ...] = ( + b'data: {"type":"content_block_delta","delta":{"text":"one"}}\n\n', + MID_EVENT_HEAD, + MID_EVENT_TAIL, + b'data: {"type":"message_delta","usage":{"output_tokens":7}}\n\n', + b"data: [DONE]\n\n", +) +JSON_CHUNKS: tuple[bytes, ...] = (b'{"echo":"one",', b'"chunked":true}') + + +class _ChunkedProvider(ThreadingHTTPServer): + """A provider that frames its response as a known list of transfer chunks, each + flushed on its own, and optionally hangs up part way through without writing the + terminating chunk. The chunk list is what the recording has to reproduce.""" + + daemon_threads = True + + def __init__( + self, + bind: tuple[str, int], + *, + chunks: tuple[bytes, ...], + content_type: str, + abort_after: int | None, + ) -> None: + super().__init__(bind, _ChunkedProviderHandler) + self.chunks = chunks + self.content_type = content_type + self.abort_after = abort_after + self.hits: list[str] = [] + + +class _ChunkedProviderHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self) -> None: + provider = self.server + assert isinstance(provider, _ChunkedProvider) + length = int(self.headers.get("content-length") or "0") + if length: + self.rfile.read(length) + provider.hits.append(f"{self.command} {self.path}") + self.send_response(200) + self.send_header("content-type", provider.content_type) + self.send_header("transfer-encoding", "chunked") + self.end_headers() + limit = len(provider.chunks) if provider.abort_after is None else provider.abort_after + for chunk in provider.chunks[:limit]: + self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) + self.wfile.flush() + if limit < len(provider.chunks): + self.close_connection = True + return + self.wfile.write(b"0\r\n\r\n") + self.wfile.flush() + + def log_message(self, format: str, *args: object) -> None: + """Silence the per-request stderr line BaseHTTPRequestHandler emits.""" + + +@contextmanager +def chunked_provider( + *, + chunks: tuple[bytes, ...] = SSE_CHUNKS, + content_type: str = "text/event-stream", + abort_after: int | None = None, +) -> Generator[_ChunkedProvider]: + server = _ChunkedProvider( + ("127.0.0.1", 0), chunks=chunks, content_type=content_type, abort_after=abort_after + ) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield server + finally: + server.shutdown() + server.server_close() + + +def response_header(head: str, name: str) -> str | None: + wanted = f"{name.lower()}:" + for line in head.splitlines()[1:]: + if line.lower().startswith(wanted): + return line.split(":", 1)[1].strip() + return None + + +def _read_chunked(sock: socket.socket, buffered: bytes) -> tuple[list[bytes], str]: + """A chunked body read back one entry per HTTP chunk, plus how the message ended. + + The framing is parsed rather than ``recv`` calls counted, because TCP is free to + coalesce two chunks into one segment or split one across two, so a read count + says nothing about how the sender framed the message.""" + chunks: list[bytes] = [] + try: + while True: + while b"\r\n" not in buffered: + piece = sock.recv(65536) + if not piece: + return chunks, "truncated" + buffered += piece + line, _, buffered = buffered.partition(b"\r\n") + size = int(line.split(b";")[0], 16) + if size == 0: + return chunks, "terminated" + while len(buffered) < size + 2: + piece = sock.recv(65536) + if not piece: + return chunks, "truncated" + buffered += piece + chunks.append(buffered[:size]) + buffered = buffered[size + 2 :] + except ConnectionResetError: + return chunks, "reset" + + +def _read_fixed(sock: socket.socket, buffered: bytes, length: int) -> tuple[list[bytes], str]: + while len(buffered) < length: + piece = sock.recv(65536) + if not piece: + return ([buffered] if buffered else []), "truncated" + buffered += piece + return ([buffered[:length]] if length else []), "terminated" + + +def raw_stream_post(port: int, path: str, body: bytes) -> tuple[str, list[bytes], str]: + """POST over a raw socket and read the reply at the transfer layer: the response + head, one entry per HTTP chunk (or the whole body for a content-length reply), + and how the message ended, ``terminated`` when its terminator arrived, + ``truncated`` on a graceful close before it, ``reset`` on an abortive one. + + ``call_edge`` goes through ``forward``, which buffers, so it cannot see any of + this; the streaming tests need the framing itself, so they read the socket.""" + sock = socket.create_connection(("127.0.0.1", port), timeout=15) + try: + sock.sendall( + ( + f"POST {path} HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\n" + f"content-type: application/json\r\ncontent-length: {len(body)}\r\n\r\n" + ).encode() + + body + ) + buffered = b"" + while b"\r\n\r\n" not in buffered: + piece = sock.recv(65536) + if not piece: + break + buffered += piece + head_bytes, _, rest = buffered.partition(b"\r\n\r\n") + head = head_bytes.decode("latin-1") + if (response_header(head, "transfer-encoding") or "").lower() == "chunked": + chunks, ending = _read_chunked(sock, rest) + else: + chunks, ending = _read_fixed( + sock, rest, int(response_header(head, "content-length") or 0) + ) + return head, chunks, ending + finally: + sock.close() + + @contextmanager def running_edge(backend: EdgeBackend, mounts: Mapping[str, str]) -> Generator[ProviderEdge]: running = start_provider_edge(backend, mounts=mounts, bind_host="127.0.0.1") @@ -787,6 +963,242 @@ class TestConcurrentReplay: assert source.leftover_error(current_test_key()) is None +def record_stream(root: Path, *, abort_after: int | None = None) -> tuple[str, list[bytes], str]: + with chunked_provider(abort_after=abort_after) as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + return raw_stream_post(edge.port, STREAM_PATH, STREAM_BODY) + + +def replay_stream(root: Path) -> tuple[str, list[bytes], str]: + with running_edge(ReplayEdge(source=replay_source(root)), REPLAY_MOUNTS) as edge: + return raw_stream_post(edge.port, STREAM_PATH, STREAM_BODY) + + +def only_recorded_response(root: Path) -> RecordedHttpResponse | RecordedStreamedResponse: + files = this_tests_files(root) + assert len(files) == 1, [file.name for file in files] + return Interaction.model_validate_json(files[0].read_text(encoding="utf-8")).response + + +def recorded_stream(root: Path) -> RecordedStreamedResponse: + response = only_recorded_response(root) + assert isinstance(response, RecordedStreamedResponse), response + return response + + +def stream_chunks(response: RecordedStreamedResponse) -> list[bytes]: + return [base64.b64decode(chunk) for chunk in response.chunks_b64] + + +class TestStreamingFidelity: + """LIT-5742: a streamed response records and replays as the chunk sequence the + provider actually sent, not as one coalesced body. The unit of fidelity is the + HTTP transfer chunk, so every assertion here is made at the transfer layer.""" + + def test_a_streamed_response_records_its_chunk_boundaries(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + record_stream(root) + + recorded = recorded_stream(root) + assert recorded.status_code == 200 + assert stream_chunks(recorded) == list(SSE_CHUNKS) + assert recorded.truncated is None + + def test_replay_reproduces_the_recorded_split_points(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + record_stream(root) + + head, chunks, ending = replay_stream(root) + assert head.startswith("HTTP/1.1 200 OK") + assert response_header(head, "transfer-encoding") == "chunked" + assert response_header(head, "content-type") == "text/event-stream" + assert len(chunks) > 1 + assert chunks == list(SSE_CHUNKS) + assert ending == "terminated" + + def test_record_mode_relays_the_stream_chunked_like_replay_will(self, tmp_path: Path) -> None: + """Record/replay parity at the framing level: what record serves the proxy + must be what replay serves it later, chunk for chunk.""" + root = tmp_path / "bundle" + recorded_head, recorded_chunks, recorded_ending = record_stream(root) + replayed_head, replayed_chunks, replayed_ending = replay_stream(root) + + assert response_header(recorded_head, "transfer-encoding") == "chunked" + assert recorded_chunks == list(SSE_CHUNKS) + assert recorded_chunks == replayed_chunks + assert recorded_ending == replayed_ending == "terminated" + assert response_header(recorded_head, "transfer-encoding") == response_header( + replayed_head, "transfer-encoding" + ) + + def test_a_chunk_split_inside_an_event_survives_replay(self, tmp_path: Path) -> None: + """The anti-tautology test. One provider chunk ends mid-token, so the two + halves of that SSE event must arrive as two chunks; an implementation that + joins the body and re-splits it on event boundaries cannot pass this.""" + root = tmp_path / "bundle" + record_stream(root) + + _, chunks, _ = replay_stream(root) + split_at = SSE_CHUNKS.index(MID_EVENT_HEAD) + assert chunks[split_at] == MID_EVENT_HEAD + assert chunks[split_at + 1] == MID_EVENT_TAIL + assert b"content_block_delta" not in chunks[split_at] + assert b"content_block_delta" in chunks[split_at] + chunks[split_at + 1] + + def test_the_usage_chunk_replays_in_its_recorded_position(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + record_stream(root) + recorded = stream_chunks(recorded_stream(root)) + + _, replayed, _ = replay_stream(root) + usage_positions = [ + index for index, chunk in enumerate(recorded) if b"output_tokens" in chunk + ] + assert usage_positions == [ + index for index, chunk in enumerate(replayed) if b"output_tokens" in chunk + ] + assert usage_positions == [len(replayed) - 2] + assert replayed[-1] == SSE_CHUNKS[-1] + + def test_a_mid_stream_upstream_failure_records_the_delivered_chunks_and_the_truncation( + self, tmp_path: Path + ) -> None: + """The provider delivers two chunks and hangs up. The deltas it did send are + the difference between a stream that died and a request that never streamed, + so they are recorded, and the recording says the stream never terminated.""" + root = tmp_path / "bundle" + head, chunks, ending = record_stream(root, abort_after=2) + + assert head.startswith("HTTP/1.1 200 OK") + assert chunks == list(SSE_CHUNKS[:2]) + assert ending == "truncated" + recorded = recorded_stream(root) + assert recorded.status_code == 200 + assert stream_chunks(recorded) == list(SSE_CHUNKS[:2]) + assert recorded.truncated is not None + assert recorded.truncated.startswith("upstream: ") + + def test_a_downstream_disconnect_mid_relay_records_only_the_delivered_chunks( + self, tmp_path: Path + ) -> None: + """The provider keeps sending, but the proxy the edge relays to hangs up after + two chunks. The chunk whose downstream write never landed must stay out of the + recording, or replay would hand back a byte the record run never delivered. + + Driven through the pure ``handle_edge_request`` core because a socket client + cannot force these tiny chunks to block mid-write, so closing the relay + generator is the faithful stand-in for the downstream write raising: it lands + the generator on the same suspended yield a broken pipe would.""" + root = tmp_path / "bundle" + with chunked_provider() as provider: + outcome = handle_edge_request( + record_backend(root), + {"openai": provider_url(provider)}, + "POST", + STREAM_PATH, + {"content-type": "application/json"}, + STREAM_BODY, + timeout=10.0, + ) + assert isinstance(outcome, EdgeStream) + steps = outcome.steps + first = next(steps) + second = next(steps) + assert isinstance(first, StreamChunk) and isinstance(second, StreamChunk) + assert (first.data, second.data) == (SSE_CHUNKS[0], SSE_CHUNKS[1]) + steps.close() + + recorded = recorded_stream(root) + assert recorded.status_code == 200 + assert stream_chunks(recorded) == [SSE_CHUNKS[0]] + assert recorded.truncated == "downstream: relay closed after 1 chunks" + + def test_a_truncated_recording_replays_as_a_truncated_stream(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + record_stream(root, abort_after=2) + + head, chunks, ending = replay_stream(root) + assert head.startswith("HTTP/1.1 200 OK") + assert response_header(head, "transfer-encoding") == "chunked" + assert chunks == list(SSE_CHUNKS[:2]) + assert ending == "truncated" + + def test_a_non_streamed_response_keeps_the_buffered_shape(self, tmp_path: Path) -> None: + """No-churn guard: an ordinary JSON response records and is framed exactly as + it was before streaming existed.""" + root = tmp_path / "bundle" + with fake_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + head, chunks, ending = raw_stream_post(edge.port, CHAT_PATH, chat_body("hi")) + + response = only_recorded_response(root) + assert isinstance(response, RecordedHttpResponse) + assert response_header(head, "transfer-encoding") is None + assert response_header(head, "content-length") is not None + assert ending == "terminated" + assert json_object(b"".join(chunks))["echo"] == chat_body("hi").decode() + + def test_a_chunked_non_sse_response_stays_buffered(self, tmp_path: Path) -> None: + """Detection keys off the content type, not the transfer encoding: providers + chunk ordinary JSON freely, and treating that as streamed would move nearly + every recording to the chunk-list shape for no gain.""" + root = tmp_path / "bundle" + with chunked_provider(chunks=JSON_CHUNKS, content_type="application/json") as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + head, chunks, _ = raw_stream_post(edge.port, CHAT_PATH, chat_body("hi")) + + response = only_recorded_response(root) + assert isinstance(response, RecordedHttpResponse) + assert base64.b64decode(response.body_b64) == b"".join(JSON_CHUNKS) + assert response_header(head, "transfer-encoding") is None + assert b"".join(chunks) == b"".join(JSON_CHUNKS) + + def test_replay_of_a_stream_makes_no_provider_connection(self, tmp_path: Path) -> None: + root = tmp_path / "bundle" + with chunked_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + raw_stream_post(edge.port, STREAM_PATH, STREAM_BODY) + hits_after_record = list(provider.hits) + with running_edge( + ReplayEdge(source=replay_source(root)), {"openai": provider_url(provider)} + ) as edge: + _, chunks, ending = raw_stream_post(edge.port, STREAM_PATH, STREAM_BODY) + assert provider.hits == hits_after_record == ["POST /v1/messages"] + assert chunks == list(SSE_CHUNKS) + assert ending == "terminated" + + def test_concurrent_streams_each_record_their_own_chunks(self, tmp_path: Path) -> None: + """The edge relays streams on concurrent threads and each one takes the + recorder lock once, at the end, so neither recording loses or borrows a chunk + from the other.""" + root = tmp_path / "bundle" + bodies = [ + json.dumps({"model": "claude", "stream": True, "n": index}).encode() + for index in range(2) + ] + barrier = threading.Barrier(len(bodies)) + with chunked_provider() as provider: + with running_edge(record_backend(root), {"openai": provider_url(provider)}) as edge: + + def consume(body: bytes) -> tuple[list[bytes], str]: + barrier.wait() + _, chunks, ending = raw_stream_post(edge.port, STREAM_PATH, body) + return chunks, ending + + with ThreadPoolExecutor(max_workers=len(bodies)) as executor: + served = list(executor.map(consume, bodies)) + + assert served == [(list(SSE_CHUNKS), "terminated")] * len(bodies) + files = this_tests_files(root) + assert len(files) == len(bodies) + for file in files: + response = Interaction.model_validate_json( + file.read_text(encoding="utf-8") + ).response + assert isinstance(response, RecordedStreamedResponse), response + assert stream_chunks(response) == list(SSE_CHUNKS) + + class TestHandleEdgeRequestPure: def test_unknown_mount_404s_naming_the_known_mounts(self, tmp_path: Path) -> None: root = tmp_path / "bundle" @@ -800,6 +1212,7 @@ class TestHandleEdgeRequestPure: b"{}", timeout=1.0, ) + assert isinstance(reply, EdgeReply) assert reply.status_code == 404 assert b"unknown provider mount 'bedrock'" in reply.body assert b"anthropic, openai" in reply.body @@ -824,6 +1237,7 @@ class TestHandleEdgeRequestPure: json.dumps({"prompt": "x"}).encode(), timeout=1.0, ) + assert isinstance(reply, EdgeReply) assert reply.status_code == 201 assert reply.body == b"ok" assert reply.headers == {"x-upstream": "fake"} diff --git a/tests/load_tests/test_langsmith_load_test.py b/tests/load_tests/test_langsmith_load_test.py index 84400d6974b..5eca0e339f5 100644 --- a/tests/load_tests/test_langsmith_load_test.py +++ b/tests/load_tests/test_langsmith_load_test.py @@ -64,11 +64,6 @@ def test_langsmith_logging_async(): except Exception as e: pytest.fail(f"An exception occurred - {e}") - except litellm.Timeout as e: - pass - except Exception as e: - pytest.fail(f"An exception occurred - {e}") - async def make_async_calls(metadata=None, **completion_kwargs): total_tasks = 300 diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 76ff23a9a1b..3d66064f5c0 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -4202,13 +4202,7 @@ def test_gemini_google_maps_tool_simple(): ) print(f"Response: {response.model_dump_json(indent=4)}") assert response.choices[0].message.content is not None - except (litellm.RateLimitError, litellm.InternalServerError): - # Transient Vertex-side failures (rate limiting, 500 INTERNAL from the - # Google Maps grounding backend) are not LiteLLM bugs — don't fail CI. - pass - except litellm.InternalServerError: - pytest.skip( - "Google Maps Platform returned a transient 500 (upstream flake); skipping." - ) + except (litellm.RateLimitError, litellm.InternalServerError) as e: + pytest.skip(f"Transient Vertex-side failure, not a LiteLLM bug: {e}") except Exception as e: pytest.fail(f"Error occurred: {e}") diff --git a/tests/logging_callback_tests/test_spend_logs.py b/tests/logging_callback_tests/test_spend_logs.py index feecfc9f4ab..15f073123a4 100644 --- a/tests/logging_callback_tests/test_spend_logs.py +++ b/tests/logging_callback_tests/test_spend_logs.py @@ -143,7 +143,6 @@ def test_spend_logs_payload(model_id: Optional[str]): "completion_start_time": datetime.datetime(2024, 6, 7, 12, 43, 30, 954146), "max_tokens": 10, "extra_body": {}, - "custom_llm_provider": "azure", "input": [ {"role": "system", "content": "you are a helpful assistant.\n"}, {"role": "user", "content": "bom dia"}, diff --git a/tests/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py index d1dc3ec7216..5823893afc0 100644 --- a/tests/mcp_tests/conftest.py +++ b/tests/mcp_tests/conftest.py @@ -7,6 +7,7 @@ import pytest import litellm import asyncio +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @pytest.fixture(scope="session") @@ -38,6 +39,8 @@ def setup_and_teardown(): yield # Teardown code (executes after the yield point) + # LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) loop.close() # Close the loop created earlier asyncio.set_event_loop(None) # Remove the reference to the loop diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 7ee745b311e..fc9f675f837 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -1,6 +1,9 @@ import os import pytest import asyncio +import subprocess +import sys +from pathlib import Path from typing import Optional from unittest.mock import AsyncMock, patch @@ -24,12 +27,20 @@ from mcp.types import Tool as MCPTool, CallToolResult, TextContent class TestMCPLogger(CustomLogger): def __init__(self): self.standard_logging_payload = None + self.mcp_tool_call_payloads = [] super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): print("success event") - self.standard_logging_payload = kwargs.get("standard_logging_object", None) - print(f"Captured standard_logging_payload: {self.standard_logging_payload}") + payload = kwargs.get("standard_logging_object", None) + self.standard_logging_payload = payload + # Async success events from other calls (e.g. a mocked acompletion whose + # log task is delivered late) race with the MCP event for the single + # last-writer slot; keep MCP tool calls in their own list so assertions + # are order-independent. + if payload is not None and payload.get("call_type") == "call_mcp_tool": + self.mcp_tool_call_payloads.append(payload) + print(f"Captured standard_logging_payload: {payload}") def _set_authorized_user(server_ids): @@ -138,7 +149,11 @@ async def test_mcp_cost_tracking(): # wait 1-2 seconds for logging to be processed await asyncio.sleep(2) - logged_standard_logging_payload = test_logger.standard_logging_payload + logged_standard_logging_payload = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print("logged_standard_logging_payload", logged_standard_logging_payload) # Add assertions @@ -277,7 +292,11 @@ async def test_mcp_cost_tracking_per_tool(): # wait for logging to be processed await asyncio.sleep(2) - logged_standard_logging_payload_1 = test_logger.standard_logging_payload + logged_standard_logging_payload_1 = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print( "logged_standard_logging_payload_1", logged_standard_logging_payload_1 ) @@ -290,6 +309,7 @@ async def test_mcp_cost_tracking_per_tool(): # Reset logger for second test test_logger.standard_logging_payload = None + test_logger.mcp_tool_call_payloads.clear() # Test 2: Call cheap_tool - should cost 0.1 response2 = await mcp_server_tool_call( @@ -300,7 +320,11 @@ async def test_mcp_cost_tracking_per_tool(): # wait for logging to be processed await asyncio.sleep(2) - logged_standard_logging_payload_2 = test_logger.standard_logging_payload + logged_standard_logging_payload_2 = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print( "logged_standard_logging_payload_2", logged_standard_logging_payload_2 ) @@ -329,16 +353,7 @@ async def test_mcp_cost_tracking_per_tool(): assert mock_client.call_tool.call_count == 2 -class MCPLoggerHook(CustomLogger): - def __init__(self): - self.standard_logging_payload = None - super().__init__() - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - print("success event") - self.standard_logging_payload = kwargs.get("standard_logging_object", None) - print(f"Captured standard_logging_payload: {self.standard_logging_payload}") - +class MCPLoggerHook(TestMCPLogger): async def async_post_mcp_tool_call_hook( self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time ) -> Optional[MCPPostCallResponseObject]: @@ -436,9 +451,55 @@ async def test_mcp_tool_call_hook(): await asyncio.sleep(2) # check logged standard logging payload - logged_standard_logging_payload = test_logger.standard_logging_payload + logged_standard_logging_payload = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print("logged_standard_logging_payload", logged_standard_logging_payload) assert ( logged_standard_logging_payload is not None ), "Standard logging payload should not be None" assert logged_standard_logging_payload["response_cost"] == 1.42 + + +_QUEUED_LOGGING_OUTLIVES_TEST = ''' +import time + +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + +ran_at = [] + + +async def _record_run(): + ran_at.append(time.monotonic()) + + +async def test_1_leaves_logging_queued_behind_a_stopped_worker(): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run()) + await GLOBAL_LOGGING_WORKER.stop() + assert ran_at == [] + + +async def test_2_starts_after_the_previous_tests_logging_ran(): + started_at = time.monotonic() + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run()) + await GLOBAL_LOGGING_WORKER.flush() + assert [t < started_at for t in ran_at] == [True, False] +''' + + +def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path): + """Regression: a logging coroutine queued by one test must not run inside a later test (it would log into that + test's callbacks, which is how test_mcp_tool_call_hook captured a gpt-4o-mini payload under xdist).""" + (tmp_path / "conftest.py").write_text((Path(__file__).parent / "conftest.py").read_text()) + (tmp_path / "pyproject.toml").write_text('[tool.pytest.ini_options]\nasyncio_mode = "auto"\n') + (tmp_path / "test_queued_logging.py").write_text(_QUEUED_LOGGING_OUTLIVES_TEST) + result = subprocess.run( + [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", "test_queued_logging.py"], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/proxy_behavior/management/test_team_member_reset_spend.py b/tests/proxy_behavior/management/test_team_member_reset_spend.py new file mode 100644 index 00000000000..ec2c78139fe --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_reset_spend.py @@ -0,0 +1,152 @@ +import uuid + +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_SEED_SPEND = 5.0 +_RESET_TO = 2.0 + + +# POST /team/{team_id}/member/{user_id}/reset_spend. The handler gate is +# _verify_team_access (proxy admin / team admin of this team / org admin of +# the team's org) — the same gate /team/member_update uses, so this mirrors +# that file's matrix exactly. +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None: + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + ) + elif shape == "beta": + await create_scratch_team(prisma, team_id, organization_id=world.org_b_id) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + await prisma.db.litellm_teammembership.create( + data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND} + ) + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_reset_spend_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + member_id = scratch.tag("member") + await _seed_target(prisma, world, shape, scratch.prefix, member_id) + caller = world.keys[actor] + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_spend", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"reset_to": _RESET_TO}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}} + ) + assert row is not None + if expected_status == 200: + assert row.spend == _RESET_TO + else: + assert row.spend == _SEED_SPEND, "denied but spend reset" + + +async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world): + resp = await proxy_client.post( + f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_spend_missing_membership_is_404( + proxy_client, prisma, scratch, world +): + """A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_spend_above_current_spend_is_400( + proxy_client, prisma, scratch, world +): + member_id = scratch.tag("member") + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + await prisma.db.litellm_teammembership.create( + data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0} + ) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 5.0}, + ) + assert resp.status_code == 400, resp.text + + +async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend( + proxy_client, prisma, scratch, world +): + """A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an + admin could repeatedly zero their own spend right before it crosses their per-member + cap, consuming the shared team budget without the configured limit ever binding.""" + team_admin = world.keys[Actor.TEAM_ADMIN] + await create_scratch_team( + prisma, + scratch.prefix, + organization_id=world.org_a_id, + admin_user_ids=[team_admin.user_id], + ) + await prisma.db.litellm_teammembership.create( + data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND} + ) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend", + headers={"Authorization": f"Bearer {team_admin.cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 403, resp.text + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}} + ) + assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset" diff --git a/tests/proxy_migration_tests/test_ui_image_serves_offline.py b/tests/proxy_migration_tests/test_ui_image_serves_offline.py new file mode 100644 index 00000000000..5ff68effd7f --- /dev/null +++ b/tests/proxy_migration_tests/test_ui_image_serves_offline.py @@ -0,0 +1,132 @@ +"""Image-level regression net for arbitrary-uid boot of the UI image. + +OpenShift ``restricted-v2`` ignores the image ``USER`` and assigns an +arbitrary uid in GID 0. The stock nginx base expects to start as root, so +its cache (``/var/cache/nginx``) and pid (``/run``) paths are root-owned +755 and the master process dies at startup with +``mkdir() "/var/cache/nginx/client_temp" failed (13: Permission denied)``. +The fix anchors everything nginx writes under ``/tmp`` in ``ui/nginx.conf``. + +Booting the image the way that deployment does, with a read-only root +filesystem and ``/tmp`` as the only writable mount, is what catches the +whole class: a boot as the default (root) uid passes even on the broken +config. + +Gated on LITELLM_IMAGE so it is skipped in the normal unit-test run and +exercised only where an image has been built (the image-scan workflow). +Requires a working docker CLI. +""" + +import os +import shutil +import subprocess +import time +import uuid +from collections.abc import Iterator + +import pytest + +IMAGE = os.getenv("LITELLM_IMAGE") +CURL_IMAGE = os.getenv("LITELLM_TEST_CURL_IMAGE", "curlimages/curl:8.11.1") +UI_PORT = os.getenv("LITELLM_UI_PORT", "3000") +ARBITRARY_UID = "1001200000:0" +STARTUP_TIMEOUT_SECONDS = int(os.getenv("LITELLM_UI_STARTUP_TIMEOUT", "60")) + +pytestmark = [ + pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"), + pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"), +] + + +def _docker(*args: str, check: bool = True) -> "subprocess.CompletedProcess[str]": + return subprocess.run(["docker", *args], capture_output=True, text=True, check=check) + + +@pytest.fixture() +def ui_container() -> Iterator[tuple[str, str]]: + """The UI container as an arbitrary uid in GID 0 on a network with no egress. + + ``--read-only`` with a tmpfs on ``/tmp`` mirrors the strictest supported + deployment: ``readOnlyRootFilesystem: true`` with an emptyDir on ``/tmp``. + A config that writes anywhere else fails here exactly like it does on + OpenShift. + """ + run_id = f"uiserve-{uuid.uuid4().hex[:8]}" + network = f"{run_id}-net" + container = f"{run_id}-ui" + + _docker("pull", "--quiet", CURL_IMAGE) + _docker("network", "create", "--internal", network) + try: + assert IMAGE is not None + _docker( + "run", "-d", "--name", container, "--network", network, + "--user", ARBITRARY_UID, + "--read-only", "--tmpfs", "/tmp", + IMAGE, + ) + yield network, container + finally: + _docker("logs", container, check=False) + _docker("rm", "-f", container, check=False) + _docker("network", "rm", network, check=False) + + +def _container_logs(container: str) -> str: + logs = _docker("logs", container, check=False) + return f"stdout:\n{logs.stdout}\nstderr:\n{logs.stderr}" + + +def _is_running(container: str) -> bool: + return bool( + _docker( + "ps", "--filter", f"name={container}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout.strip() + ) + + +def _probe(network: str, container: str, path: str) -> "subprocess.CompletedProcess[str]": + return _docker( + "run", "--rm", "--network", network, CURL_IMAGE, + "--silent", "--show-error", "--max-time", "10", + "--output", "/dev/null", "--write-out", "%{http_code}", + f"http://{container}:{UI_PORT}{path}", + check=False, + ) + + +def test_ui_serves_as_arbitrary_uid_read_only(ui_container: tuple[str, str]) -> None: + """nginx boots and serves as an arbitrary uid with a read-only root fs. + + On the pre-fix config nginx exits during startup with + ``mkdir() "/var/cache/nginx/client_temp" failed (13: Permission denied)`` + and the running-check below fails; it never reaches the probes. + """ + network, container = ui_container + + deadline = time.time() + STARTUP_TIMEOUT_SECONDS + healthz = None + while time.time() < deadline: + if not _is_running(container): + pytest.fail( + f"the UI container exited during startup as uid {ARBITRARY_UID} with a " + f"read-only root filesystem. nginx writes outside /tmp.\n" + f"{_container_logs(container)}" + ) + healthz = _probe(network, container, "/healthz") + if healthz.returncode == 0 and healthz.stdout.strip() == "200": + break + time.sleep(2) + + assert healthz is not None and healthz.stdout.strip() == "200", ( + f"/healthz never answered 200 within {STARTUP_TIMEOUT_SECONDS}s as uid " + f"{ARBITRARY_UID}.\n{_container_logs(container)}" + ) + + for path in ("/", "/ui", "/ui/login"): + page = _probe(network, container, path) + assert page.stdout.strip() == "200", ( + f"GET {path} returned {page.stdout.strip()!r} as uid {ARBITRARY_UID}.\n" + f"{_container_logs(container)}" + ) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index a3deeb46f6e..6115e627b26 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -3323,16 +3323,17 @@ async def test_team_access_groups(prisma_client): request._url = URL(url="/chat/completions") + def body_reader(requested_model: str): + async def return_body() -> bytes: + return f'{{"model": "{requested_model}"}}'.encode() + + return return_body + for model in ["gpt-4o", "gemini-pro-vision"]: # Expect these to pass - async def return_body(): - return_string = f'{{"model": "{model}"}}' - # return string as bytes - return return_string.encode() - request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - request.body = return_body + request.body = body_reader(model) # use generated key to auth in print( @@ -3342,14 +3343,9 @@ async def test_team_access_groups(prisma_client): for model in ["gpt-4", "gpt-4o-mini", "gemini-experimental"]: # Expect these to fail - async def return_body_2(): - return_string = f'{{"model": "{model}"}}' - # return string as bytes - return return_string.encode() - request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") - request.body = return_body_2 + request.body = body_reader(model) # use generated key to auth in print( diff --git a/tests/search_tests/test_bing_grounding_search.py b/tests/search_tests/test_bing_grounding_search.py new file mode 100644 index 00000000000..3d1737477a1 --- /dev/null +++ b/tests/search_tests/test_bing_grounding_search.py @@ -0,0 +1,199 @@ +""" +Tests for the Grounding with Bing Search (Microsoft Foundry) integration. +""" + +import json +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +import litellm +from tests.search_tests.base_search_unit_tests import BaseSearchTest + +PROJECT_ENDPOINT = "https://acct.services.ai.azure.com/api/projects/proj" + +_ANSWER_TEXT = ( + "LiteLLM is an open source LLM gateway ([github.com](https://github.com/BerriAI/litellm))\n" + "The docs live on docs.litellm.ai ([docs.litellm.ai](https://docs.litellm.ai/))" +) + + +def _annotation(marker: str, url: str, title: str) -> dict: + start = _ANSWER_TEXT.index(marker) + return { + "type": "url_citation", + "url": url, + "title": title, + "start_index": start, + "end_index": start + len(marker), + } + + +MOCK_BING_GROUNDING_RESPONSE = { + "id": "resp_mock", + "object": "response", + "status": "completed", + "model": "gpt-4.1", + "output": [ + {"type": "web_search_call", "status": "completed"}, + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": _ANSWER_TEXT, + "annotations": [ + _annotation( + "([github.com](https://github.com/BerriAI/litellm))", + "https://github.com/BerriAI/litellm", + "BerriAI/litellm - GitHub", + ), + _annotation( + "([docs.litellm.ai](https://docs.litellm.ai/))", + "https://docs.litellm.ai/", + "LiteLLM Docs", + ), + ], + } + ], + }, + ], + "usage": {"input_tokens": 100, "output_tokens": 50}, +} + + +def _mock_response(): + response = Mock() + response.status_code = 200 + response.headers = {} + response.content = json.dumps(MOCK_BING_GROUNDING_RESPONSE).encode() + return response + + +@pytest.mark.skip(reason="Local only tested search providers") +class TestBingGroundingSearch(BaseSearchTest): + """ + E2E tests for Grounding with Bing Search that make real API calls. + Inherits from BaseSearchTest to run standard search tests. + """ + + def get_search_provider(self) -> str: + return "bing_grounding" + + +class TestBingGroundingSearchTransformation: + """ + Full-stack tests through `litellm.search` / `litellm.asearch` with the HTTP layer mocked. + Transformation details are unit-tested in tests/test_litellm/llms/azure/search/. + """ + + @pytest.fixture(autouse=True) + def _server_env(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_PROJECT_ENDPOINT", PROJECT_ENDPOINT) + monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1") + monkeypatch.setenv("BING_GROUNDING_TOKEN", "test-entra-token") + monkeypatch.delenv("BING_GROUNDING_CONNECTION_ID", raising=False) + + def test_bing_grounding_search_request_and_response(self): + with patch( # test-quality-ok: litellm.search has no client injection seam + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + return_value=_mock_response(), + ) as mock_post: + response = litellm.search( + query="what is litellm", + search_provider="bing_grounding", + max_results=5, + country="us", + ) + + assert mock_post.called + call_kwargs = mock_post.call_args.kwargs + assert call_kwargs["url"] == f"{PROJECT_ENDPOINT}/openai/v1/responses" + assert call_kwargs["headers"]["Authorization"] == "Bearer test-entra-token" + + request_body = call_kwargs["json"] + assert request_body["model"] == "gpt-4.1" + assert request_body["input"] == "what is litellm" + assert request_body["tools"] == [ + {"type": "web_search", "user_location": {"type": "approximate", "country": "US"}} + ] + + assert response.object == "search" + assert len(response.results) == 2 + assert response.results[0].url == "https://github.com/BerriAI/litellm" + assert response.results[0].title == "BerriAI/litellm - GitHub" + assert response.results[0].snippet == "LiteLLM is an open source LLM gateway" + assert response.results[1].url == "https://docs.litellm.ai/" + assert response.results[1].snippet == "The docs live on docs.litellm.ai" + + def test_connection_mode_sends_the_bing_grounding_tool(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv( + "BING_GROUNDING_CONNECTION_ID", + "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.CognitiveServices" + "/accounts/acct/projects/proj/connections/bing-conn", + ) + with patch( # test-quality-ok: litellm.search has no client injection seam + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + return_value=_mock_response(), + ) as mock_post: + litellm.search( + query="what is litellm", + search_provider="bing_grounding", + max_results=3, + ) + + request_body = mock_post.call_args.kwargs["json"] + assert request_body["tools"] == [ + { + "type": "bing_grounding", + "bing_grounding": { + "search_configurations": [ + { + "project_connection_id": ( + "/subscriptions/sub/resourceGroups/rg/providers/Microsoft.CognitiveServices" + "/accounts/acct/projects/proj/connections/bing-conn" + ), + "count": 3, + } + ] + }, + } + ] + + @pytest.mark.asyncio + async def test_bing_grounding_asearch(self): + with patch( # test-quality-ok: litellm.asearch has no client injection seam + "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", + new=AsyncMock(return_value=_mock_response()), + ) as mock_post: + response = await litellm.asearch( + query="what is litellm", + search_provider="bing_grounding", + ) + + assert mock_post.call_args.kwargs["json"]["tools"] == [{"type": "web_search"}] + assert len(response.results) == 2 + + def test_web_search_mode_is_not_billed_the_g1_price(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + with patch( # test-quality-ok: litellm.search has no client injection seam + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + return_value=_mock_response(), + ): + response = litellm.search(query="pricing check", search_provider="bing_grounding") + + assert response._hidden_params["response_cost"] == 0.0 + + def test_connection_mode_tracks_the_g1_cost(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id") + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + with patch( # test-quality-ok: litellm.search has no client injection seam + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + return_value=_mock_response(), + ): + response = litellm.search(query="pricing check", search_provider="bing_grounding") + + assert response._hidden_params["response_cost"] == pytest.approx(0.035) diff --git a/tests/test_litellm/a2a_protocol/test_card_resolver.py b/tests/test_litellm/a2a_protocol/test_card_resolver.py index 053f28c940f..5cbfa51fa08 100644 --- a/tests/test_litellm/a2a_protocol/test_card_resolver.py +++ b/tests/test_litellm/a2a_protocol/test_card_resolver.py @@ -13,6 +13,7 @@ from litellm.a2a_protocol.card_resolver import ( LiteLLMA2ACardResolver, fix_agent_card_url, is_localhost_or_internal_url, + normalize_agent_card_interfaces, set_agent_card_url, ) @@ -114,3 +115,26 @@ def test_fix_agent_card_url_updates_interface_when_top_level_is_localhost(): assert result.url == "https://my-public-agent.example.com/" assert result.supported_interfaces[0].url == "https://my-public-agent.example.com/" + + +def test_normalize_agent_card_interfaces_downgrades_miscased_interfaces_to_the_0_3_dialect(): + pb2 = pytest.importorskip("a2a.types.a2a_pb2") + + card = pb2.AgentCard( + name="langgraph", + supported_interfaces=[ + pb2.AgentInterface(url="http://a/", protocol_binding="jsonrpc", protocol_version="1.0"), + pb2.AgentInterface(url="http://b/", protocol_binding="JSONRPC", protocol_version="1.0"), + pb2.AgentInterface(url="http://c/", protocol_binding="websocket", protocol_version="1.0"), + ], + ) + + normalized = normalize_agent_card_interfaces(card) + + assert [(i.protocol_binding, i.protocol_version) for i in normalized.supported_interfaces] == [ + ("JSONRPC", "0.3"), + ("JSONRPC", "1.0"), + ("websocket", "1.0"), + ] + assert card.supported_interfaces[0].protocol_binding == "jsonrpc" + assert card.supported_interfaces[0].protocol_version == "1.0" diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 08f6b9f25bb..8850a2eca6c 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -176,10 +176,62 @@ _AGENT_A_HEADERS = {"x-agent-token": "token-for-a", "x-tenant": "tenant-a"} _AGENT_B_HEADERS = {"x-agent-token": "token-for-b", "x-tenant": "tenant-b"} +_LANGGRAPH_TASK_REPLY = { + "jsonrpc": "2.0", + "id": "reply", + "result": { + "kind": "task", + "id": "run-1:task-1", + "contextId": "thread-1", + "history": [ + { + "kind": "message", + "role": "user", + "parts": [{"kind": "text", "text": "hi"}], + "messageId": "m-user", + "taskId": "run-1:task-1", + "contextId": "thread-1", + }, + { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": "langgraph echo: hi"}], + "messageId": "m-agent", + "taskId": "run-1:task-1", + "contextId": "thread-1", + }, + ], + "status": {"state": "completed", "timestamp": "2026-08-24T00:00:00+00:00"}, + "artifacts": [ + { + "artifactId": "art-1", + "name": "Assistant Response", + "parts": [{"kind": "text", "text": "langgraph echo: hi"}], + } + ], + }, +} + + +_LOWERCASE_BINDING_CARD = { + "name": "langgraph-agent", + "version": "1.0.0", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text/plain"], + "defaultOutputModes": ["text/plain"], + "skills": [], + "supportedInterfaces": [ + {"url": "http://127.0.0.1:9/", "protocolBinding": "jsonrpc", "protocolVersion": "1.0"} + ], +} + + class _RequestRecorder: """Records the headers httpx put on the wire, per outbound request.""" - def __init__(self): + def __init__(self, card=_AGENT_CARD, rpc_reply=_RPC_REPLY): + self.card = card + self.rpc_reply = rpc_reply self.card_requests = [] self.rpc_requests = [] self.client = None @@ -188,23 +240,23 @@ class _RequestRecorder: headers = {k.lower(): v for k, v in request.headers.items()} if request.method == "GET": self.card_requests.append(headers) - return httpx.Response(200, json=_AGENT_CARD) + return httpx.Response(200, json=self.card) self.rpc_requests.append(headers) - return httpx.Response(200, json=_RPC_REPLY) + return httpx.Response(200, json=self.rpc_reply) def _a2a_client_cache_key(timeout: float) -> str: return "async_httpx_client" + f"timeout_{timeout}" + httpxSpecialProvider.A2AProvider -async def _seed_shared_a2a_client() -> _RequestRecorder: +async def _seed_shared_a2a_client(card=_AGENT_CARD, rpc_reply=_RPC_REPLY) -> _RequestRecorder: """Put the one A2A client the cache will hand out behind a mock transport. Seeding has to happen on the test's own event loop, because the client cache keys on it. The injected client is a real httpx.AsyncClient, so the merge of per-request headers over client defaults, which is what these tests are about, stays real. """ - recorder = _RequestRecorder() + recorder = _RequestRecorder(card=card, rpc_reply=rpc_reply) handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT) owned_client = handler.client handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder)) @@ -311,6 +363,25 @@ async def test_streaming_send_carries_only_its_own_caller_headers(isolated_clien assert received["b"]["x-tenant"] == "tenant-b" +@pytest.mark.asyncio +async def test_lowercase_protocol_binding_card_round_trips_the_langgraph_dialect(isolated_client_cache): + """LangGraph Platform serves cards with protocolBinding "jsonrpc" and answers in the + A2A 0.3 JSON dialect ("kind"-discriminated) while declaring protocolVersion "1.0". + Without binding normalization client creation raises ValueError("no compatible + transports found."); without the version downgrade the SDK's strict v1 transport + rejects the reply with 'Message type "lf.a2a.v1.Task" has no field named "kind"'.""" + await _seed_shared_a2a_client(card=_LOWERCASE_BINDING_CARD, rpc_reply=_LANGGRAPH_TASK_REPLY) + + a2a_client = await create_a2a_client(base_url="http://127.0.0.1:9") + response = await _send_message(a2a_client, _send_request("lc")) + + assert type(response.root.result).__name__ == "Task" + assert response.root.result.artifacts[0].parts[0].root.text == "langgraph echo: hi" + interface = a2a_client._litellm_agent_card.supported_interfaces[0] + assert interface.protocol_binding == "JSONRPC" + assert interface.protocol_version == "0.3" + + @pytest.mark.asyncio async def test_agent_card_fetch_carries_the_callers_headers(isolated_client_cache): """Agent cards can sit behind the same auth as the agent, so the card fetch must stay diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py index f4cd3ab20ef..c16ceec8c31 100644 --- a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py +++ b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py @@ -28,6 +28,14 @@ if TYPE_CHECKING: from redis.asyncio.cluster import RedisCluster as _AsyncRedisClusterType +class _NodeClassWithPerConnectionRecovery: + def update_active_connections_for_reconnect(self) -> None: ... + + +class _NodeClassWithoutPerConnectionRecovery: + pass + + class _FakeClusterNode: def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None: self.name = name @@ -47,7 +55,9 @@ class _FakeNodesManager: def _build_cluster_instance() -> "_AsyncRedisClusterType": - cluster_cls = get_litellm_async_redis_cluster_class() + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithoutPerConnectionRecovery + ) instance = cluster_cls.__new__(cluster_cls) instance.RedisClusterRequestTTL = 1 instance.reinitialize_counter = 0 @@ -58,6 +68,33 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType": return instance +def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None: + """Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level + connection error per-connection, the factory must NOT install the copied override, + whose node.disconnect() also kills connections other coroutines are mid-operation on.""" + from redis.asyncio.cluster import RedisCluster + + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithPerConnectionRecovery + ) + + assert cluster_cls is RedisCluster + + +def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None: + """Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(), + so those versions must keep litellm's per-node isolation override.""" + from redis.asyncio.cluster import RedisCluster + + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithoutPerConnectionRecovery + ) + + assert cluster_cls is not RedisCluster + assert issubclass(cluster_cls, RedisCluster) + assert "_execute_command" in cluster_cls.__dict__ + + @pytest.mark.asyncio @pytest.mark.parametrize("error_cls", [RedisConnectionError, RedisTimeoutError]) async def test_node_level_error_resets_only_that_node_not_the_whole_client(error_cls: type[Exception]) -> None: diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 4ff92aaf87d..6ca48ce63b8 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -2,7 +2,7 @@ import datetime import json import os import unittest -from typing import TYPE_CHECKING, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple from unittest.mock import ANY, MagicMock, Mock, patch import httpx @@ -1585,10 +1585,16 @@ def test_map_reasoning_effort_adds_summary_detailed(monkeypatch): assert result_dict["summary"] == "custom_summary" print("✓ Dict input is passed through without modification") - # Test 5: None/unknown values return None - result_unknown = handler._map_reasoning_effort("unknown_value") - assert result_unknown is None - print("✓ Unknown reasoning_effort values return None") + # Test 5: every REASONING_EFFORT level reaches the provider, and anything else (a typo, an + # unshipped level, "default") is dropped so the request still succeeds at the provider default + from litellm.types.llms.openai import Reasoning + + for effort in ("max", "xhigh", "none"): + result_passthrough = handler._map_reasoning_effort(effort) + assert result_passthrough == Reasoning(effort=effort) + for dropped in ("ultra", "hgih", "unknown_value", "", "default"): + assert handler._map_reasoning_effort(dropped) is None + print("✓ Enumerated levels pass through and unknown ones are dropped") print( "✓ All reasoning_effort behaviors work correctly with flag/env var control" @@ -2438,6 +2444,32 @@ def test_map_optional_params_preserves_reasoning_summary(): assert responses_api_request["reasoning"]["summary"] == "detailed" +@pytest.mark.parametrize("reasoning_effort", ["max", "high"]) +def test_transform_request_bedrock_mantle_tools_keeps_reasoning_effort(monkeypatch, reasoning_effort): + """Regression for reasoning_effort=max being dropped on the chat -> Responses bridge (issue #38084).""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + monkeypatch.setattr(litellm, "reasoning_auto_summary", False) + monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False) + handler: Final = LiteLLMResponsesTransformationHandler() + + result: Final = handler.transform_request( + model="openai.gpt-5.6-sol", + messages=[{"role": "user", "content": "Say pong"}], + optional_params={ + "reasoning_effort": reasoning_effort, + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + }, + litellm_params={"custom_llm_provider": "bedrock_mantle"}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert result["reasoning"] == {"effort": reasoning_effort} + + def test_map_optional_params_tool_choice_chat_nested_to_responses_api(): """Chat tool_choice must become Responses ToolChoiceFunction (top-level name).""" from litellm.completion_extras.litellm_responses_transformation.transformation import ( @@ -3762,3 +3794,112 @@ def test_response_incomplete_stream_event_without_details_defaults_to_length(): result = iterator.chunk_parser(chunk) assert result.choices[0].finish_reason == "length" + + +def test_assistant_message_with_tool_calls_keeps_its_content(): + """Regression for https://github.com/BerriAI/litellm/issues/24985. + + An assistant turn that both answered and called a tool used to lose its whole message: + the branch handling tool_calls emitted the calls and dropped the text. + """ + handler = LiteLLMResponsesTransformationHandler() + messages = [ + {"role": "user", "content": "What is the weather in Denver?"}, + { + "role": "assistant", + "content": "Let me look that up.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Denver"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "88F"}, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + assistant_message = next( + item for item in input_items if item.get("type") == "message" and item.get("role") == "assistant" + ) + assert assistant_message["content"] == [{"type": "output_text", "text": "Let me look that up."}] + assert [item.get("type") for item in input_items] == [ + "message", + "message", + "function_call", + "function_call_output", + ] + + +def test_assistant_thinking_blocks_become_a_reasoning_input_item(): + """Thinking blocks are how an Anthropic-shaped turn carries reasoning into this bridge.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [ + {"role": "user", "content": "What is the weather in Denver?"}, + { + "role": "assistant", + "content": "Denver is sunny.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "August in Denver is dry.", "signature": "sig1"}, + {"type": "redacted_thinking", "data": "REDACTED"}, + ], + }, + {"role": "user", "content": "Why?"}, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + reasoning_item = next(item for item in input_items if item.get("type") == "reasoning") + assert reasoning_item["summary"] == [{"type": "summary_text", "text": "August in Denver is dry."}] + assert "id" not in reasoning_item + + +def test_thinking_only_assistant_turn_still_sends_its_reasoning(): + """An assistant turn can be pure reasoning, with no visible text and no tool call.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [ + {"role": "user", "content": "What is the weather in Denver?"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [ + {"type": "thinking", "thinking": "August in Denver is dry.", "signature": "sig1"} + ], + }, + {"role": "user", "content": "Why?"}, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + reasoning_items = [item for item in input_items if item.get("type") == "reasoning"] + assert len(reasoning_items) == 1 + assert reasoning_items[0]["summary"] == [{"type": "summary_text", "text": "August in Denver is dry."}] + + +def test_stored_reasoning_items_win_over_thinking_blocks(): + """A minted reasoning id beats a re-derived one, so the two must not both be sent.""" + handler = LiteLLMResponsesTransformationHandler() + messages = [ + { + "role": "assistant", + "content": "Denver is sunny.", + "reasoning_items": [ + { + "type": "reasoning", + "id": "rs_real", + "summary": [{"type": "summary_text", "text": "August in Denver is dry."}], + } + ], + "thinking_blocks": [ + {"type": "thinking", "thinking": "August in Denver is dry.", "signature": "rs_real"} + ], + }, + ] + + input_items, _ = handler.convert_chat_completion_messages_to_responses_api(messages) + + reasoning_items = [item for item in input_items if item.get("type") == "reasoning"] + assert len(reasoning_items) == 1 + assert reasoning_items[0]["id"] == "rs_real" diff --git a/tests/test_litellm/images/test_image_edit_extra_params.py b/tests/test_litellm/images/test_image_edit_extra_params.py new file mode 100644 index 00000000000..088faafa9f3 --- /dev/null +++ b/tests/test_litellm/images/test_image_edit_extra_params.py @@ -0,0 +1,147 @@ +""" +Regression tests for https://github.com/BerriAI/litellm/issues/36493 + +/v1/images/edits on the openai path silently dropped unknown provider params +(e.g. seed) and the extra_body escape hatch, unlike /v1/images/generations. +""" + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + +PNG_BYTES = b"\x89PNG\r\n\x1a\nfakepng" + + +def _capture_image_edit_request(captured): + def respond(request): + captured["content_type"] = request.headers.get("content-type") + captured["body"] = request.content + return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]}) + + return respond + + +def _multipart_text_fields(content_type: str, body: bytes) -> dict: + boundary = content_type.split("boundary=")[1].encode() + return { + part.split(b'name="')[1].split(b'"')[0].decode(): part.partition(b"\r\n\r\n")[2].rstrip(b"\r\n-").decode() + for part in body.split(b"--" + boundary) + if b'name="' in part and b"filename=" not in part + } + + +def test_image_edit_forwards_provider_params_and_extra_body(): + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_image_edit_request(captured)))) + + response = litellm.image_edit( + model="openai/gpt-image-1", + image=PNG_BYTES, + prompt="add a hat", + api_key="sk-test", + api_base="https://edit.example/v1", + client=client, + seed=42, + extra_body={"quality_level": "high"}, + ) + + assert captured["content_type"].startswith("multipart/form-data") + fields = _multipart_text_fields(captured["content_type"], captured["body"]) + assert fields["seed"] == "42" + assert fields["quality_level"] == "high" + assert "extra_body" not in fields + assert fields["model"] == "gpt-image-1" + assert fields["prompt"] == "add a hat" + assert b'name="image[]"' in captured["body"] + assert response.data + + +def test_image_edit_extra_body_takes_precedence_over_kwargs(): + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_image_edit_request(captured)))) + + litellm.image_edit( + model="openai/gpt-image-1", + image=PNG_BYTES, + prompt="add a hat", + api_key="sk-test", + api_base="https://edit.example/v1", + client=client, + seed=42, + extra_body={"seed": 7}, + ) + + assert _multipart_text_fields(captured["content_type"], captured["body"])["seed"] == "7" + + +def test_image_edit_flattens_nested_provider_params(): + """A nested value in extra_body (or a nested unknown kwarg) must be + serialized as OpenAI-SDK bracket form fields (key[subkey]) rather than + handed to the httpx multipart encoder, which raises 'Invalid type for + value. Expected primitive type' on a dict and 500s the request.""" + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_image_edit_request(captured)))) + + litellm.image_edit( + model="openai/gpt-image-1", + image=PNG_BYTES, + prompt="add a hat", + api_key="sk-test", + api_base="https://edit.example/v1", + client=client, + extra_body={"generation_config": {"steps": 30, "guidance": True}}, + ) + + fields = _multipart_text_fields(captured["content_type"], captured["body"]) + assert fields["generation_config[steps]"] == "30" + assert fields["generation_config[guidance]"] == "true" + assert "generation_config" not in fields + + +def test_image_edit_forwards_scalar_array_as_repeated_fields(): + """A list-valued provider param must reach the backend as one repeated part + per element, not collapse to its last element under dict.update.""" + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_image_edit_request(captured)))) + + litellm.image_edit( + model="openai/gpt-image-1", + image=PNG_BYTES, + prompt="add a hat", + api_key="sk-test", + api_base="https://edit.example/v1", + client=client, + loras=["style_a", "style_b", "style_c"], + ) + + body = captured["body"] + assert body.count(b'name="loras"') == 3 + assert b"style_a" in body and b"style_b" in body and b"style_c" in body + + +@pytest.mark.asyncio +async def test_aimage_edit_forwards_extra_body(): + """aimage_edit used to drop extra_headers/extra_query/extra_body when + building its partial, so they never reached image_edit.""" + captured = {} + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_capture_image_edit_request(captured))) + + response = await litellm.aimage_edit( + model="openai/gpt-image-1", + image=PNG_BYTES, + prompt="add a hat", + api_key="sk-test", + api_base="https://edit.example/v1", + client=client, + seed=42, + extra_body={"quality_level": "high"}, + ) + + fields = _multipart_text_fields(captured["content_type"], captured["body"]) + assert fields["seed"] == "42" + assert fields["quality_level"] == "high" + assert "extra_body" not in fields + assert response.data diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index b6e063a6d94..cf3318b32fe 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1599,6 +1599,14 @@ class TestEnableAnthropicPromptCaching: assert supports_prompt_caching(model=model, custom_llm_provider=provider) is True assert self._points(model=model, provider=provider) == [] + def test_databricks_claude_not_injected_despite_caching_support(self, monkeypatch, local_model_cost_map): + from litellm.utils import supports_prompt_caching + + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + model = "databricks/databricks-claude-sonnet-4-5" + assert supports_prompt_caching(model=model, custom_llm_provider="databricks") is True + assert self._points(model=model, provider="databricks") == [] + def test_model_without_caching_support_not_injected(self, monkeypatch): monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) assert self._points(model="anthropic.claude-3-5-sonnet-20240620-v1:0", provider="bedrock") == [] diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index d3393ac3d28..025aa86466c 100644 --- a/tests/test_litellm/integrations/test_langsmith_init.py +++ b/tests/test_litellm/integrations/test_langsmith_init.py @@ -432,3 +432,103 @@ class TestLangsmithRedactUserApiKeyInfo: ) assert data["inputs"]["metadata"]["user_api_key_hash"] == "abc123" + + +class TestLangsmithRootRunIdConsistency: + """Regression tests for LIT-5878 / #37269. + + A request that carries a session/trace header (e.g. x-claude-code-session-id) + fans the header value out into litellm metadata as both trace_id and + session_id. LangSmith then rejected the whole ingest batch twice over: + a root run whose trace_id does not match the run id embedded in dotted_order + (400), and a run-body session_id that does not reference an existing tracer + session (404, or 422 for non-UUID values). + """ + + def _prepare(self, request_metadata): + payload = { + "id": "slp-1", + "response": {"choices": []}, + "metadata": {}, + "startTime": 1.0, + "endTime": 2.0, + "request_tags": [], + "error_str": None, + "status": "success", + "response_cost": 0.0, + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + } + logger = LangsmithLogger( + langsmith_api_key="test-key", + langsmith_project="test-project", + ) + return logger._prepare_log_data( + kwargs={ + "litellm_params": {"metadata": request_metadata}, + "standard_logging_object": payload, + }, + response_obj=None, + start_time=1.0, + end_time=2.0, + credentials={ + "LANGSMITH_API_KEY": "test-key", + "LANGSMITH_PROJECT": "test-project", + "LANGSMITH_BASE_URL": "https://api.smith.langchain.com", + }, + ) + + def test_header_derived_ids_yield_self_consistent_root_run(self): + header_value = "ed29c3bb-44fa-4eec-9b7b-fecaa3e82d64" + data = self._prepare({"trace_id": header_value, "session_id": header_value}) + + assert data["trace_id"] == data["id"] + assert data["trace_id"] != header_value + assert data["dotted_order"].endswith(data["id"]) + assert len(data["dotted_order"]) == 22 + len(data["id"]) + assert "session_id" not in data + + def test_distinct_session_id_is_still_forwarded(self): + data = self._prepare({"session_id": "11111111-2222-3333-4444-555555555555"}) + + assert data["session_id"] == "11111111-2222-3333-4444-555555555555" + + def test_trace_id_only_root_run_is_overridden(self): + data = self._prepare({"trace_id": "ed29c3bb-44fa-4eec-9b7b-fecaa3e82d64"}) + + assert data["trace_id"] == data["id"] + assert data["trace_id"] != "ed29c3bb-44fa-4eec-9b7b-fecaa3e82d64" + assert data["dotted_order"].endswith(data["id"]) + + def test_root_run_without_caller_ids_is_self_consistent(self): + data = self._prepare({}) + + assert data["trace_id"] == data["id"] + assert data["dotted_order"].endswith(data["id"]) + + def test_child_run_keeps_caller_trace_id(self): + data = self._prepare( + { + "trace_id": "trace-1", + "parent_run_id": "parent-1", + "run_id": "child-1", + } + ) + + assert data["trace_id"] == "trace-1" + assert data["id"] == "child-1" + assert data["parent_run_id"] == "parent-1" + + def test_caller_supplied_dotted_order_and_trace_id_are_untouched(self): + dotted = "20260820T000000000000Ztrace-1.20260820T000001000000Zrun-1" + data = self._prepare( + { + "trace_id": "trace-1", + "run_id": "run-1", + "dotted_order": dotted, + } + ) + + assert data["trace_id"] == "trace-1" + assert data["dotted_order"] == dotted diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 933e41d17a0..0f307450b50 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1647,15 +1647,27 @@ def _signature_for(signer_cls, url: str, method: str, body: bytes | None, header return signer.signature(signer.string_to_sign(request, canonical_request), request) +def _as_s3_canonicalizes(url: str) -> str: + """ + The path S3 rebuilds from the wire path: percent-encode everything outside the unreserved + set, without normalizing or double-encoding. `=` becomes `%3D`, `%20` stays `%20`. + """ + from urllib.parse import quote, unquote, urlsplit, urlunsplit + + split = urlsplit(url) + return urlunsplit(split._replace(path=quote(unquote(split.path), safe="/~"))) + + def _assert_signed_for_s3_canonicalization(url: str, method: str, body: bytes | None, headers: dict[str, str]) -> None: """ S3 rebuilds the canonical request from the wire path with single percent-encoding, which botocore models as S3SigV4Auth; plain SigV4Auth double-encodes it (%2520 for a space) and S3 - answers 403 SignatureDoesNotMatch. Assert we signed the path the way S3 reads it. + answers 403 SignatureDoesNotMatch. Assert we sent an already-encoded path and signed it the + way S3 reads it. """ from botocore.auth import S3SigV4Auth, SigV4Auth - assert "%20" in url + assert url == _as_s3_canonicalizes(url) sent_signature = headers["Authorization"].split("Signature=")[1].strip() assert sent_signature == _signature_for(S3SigV4Auth, url, method, body, headers) assert sent_signature != _signature_for(SigV4Auth, url, method, body, headers) @@ -1744,3 +1756,103 @@ async def test_download_signs_object_key_with_space_the_way_s3_does(): body=None, headers=call.kwargs["headers"], ) + +_RESERVED_CHAR_KEYS = ( + "2026-08-21/time-05-29-36_resp_bGl0ZWxsbTpjdXN0b20=.json", + "session=logs/2026-08-21/time-05-29-36_abc.json", + "a+b/2026-08-21/time-05-29-36_abc.json", + "a&b/2026-08-21/time-05-29-36_abc.json", + "a#b/2026-08-21/time-05-29-36_abc.json", + "a?b/2026-08-21/time-05-29-36_abc.json", + "a%b/2026-08-21/time-05-29-36_abc.json", + _KEY_WITH_SPACE, +) + + +def _element_for(s3_object_key: str): + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + return s3BatchLoggingElement( + s3_object_key=s3_object_key, + payload={"test": "sigv4"}, + s3_object_download_filename="log.json", + ) + + +def _expected_wire_url(s3_object_key: str) -> str: + """The URL boto3 itself would put on the wire for this key.""" + from urllib.parse import quote + + return f"https://logs-bucket.s3.us-east-1.amazonaws.com/{quote(s3_object_key, safe='/')}" + + +@pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS) +@pytest.mark.asyncio +async def test_async_upload_percent_encodes_reserved_characters_in_object_key(s3_object_key): + from unittest.mock import AsyncMock, MagicMock + + logger = _logger_for_signing() + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.put.return_value = response + + await logger.async_upload_data_to_s3(_element_for(s3_object_key)) + + call = logger.async_httpx_client.put.call_args + assert call[0][0] == _expected_wire_url(s3_object_key) + _assert_signed_for_s3_canonicalization( + url=call[0][0], + method="PUT", + body=call.kwargs["data"].encode("utf-8"), + headers=call.kwargs["headers"], + ) + + +@pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS) +def test_sync_upload_percent_encodes_reserved_characters_in_object_key(s3_object_key): + from unittest.mock import MagicMock + + logger = _logger_for_signing() + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + mock_sync_client = MagicMock() + mock_sync_client.put.return_value = response + + with patch("litellm.integrations.s3_v2._get_httpx_client", return_value=mock_sync_client): + logger.upload_data_to_s3(_element_for(s3_object_key)) + + call = mock_sync_client.put.call_args + assert call[0][0] == _expected_wire_url(s3_object_key) + _assert_signed_for_s3_canonicalization( + url=call[0][0], + method="PUT", + body=call.kwargs["data"].encode("utf-8"), + headers=call.kwargs["headers"], + ) + + +@pytest.mark.parametrize("s3_object_key", _RESERVED_CHAR_KEYS) +@pytest.mark.asyncio +async def test_download_percent_encodes_reserved_characters_in_object_key(s3_object_key): + from unittest.mock import AsyncMock, MagicMock + + logger = _logger_for_signing() + response = MagicMock() + response.status_code = 200 + response.json = MagicMock(return_value={"downloaded": "data"}) + logger.async_httpx_client = AsyncMock() + logger.async_httpx_client.get.return_value = response + + assert await logger._download_object_from_s3(s3_object_key) == {"downloaded": "data"} + + call = logger.async_httpx_client.get.call_args + assert call[0][0] == _expected_wire_url(s3_object_key) + _assert_signed_for_s3_canonicalization( + url=call[0][0], + method="GET", + body=None, + headers=call.kwargs["headers"], + ) diff --git a/tests/test_litellm/interactions/test_background_cost_polling.py b/tests/test_litellm/interactions/test_background_cost_polling.py new file mode 100644 index 00000000000..97f09de1b52 --- /dev/null +++ b/tests/test_litellm/interactions/test_background_cost_polling.py @@ -0,0 +1,545 @@ +import asyncio +import time +from itertools import islice +from typing import Optional + +import pytest + +from litellm.interactions.background_cost_polling import ( + _SETTLED_KEY, + _poll_intervals, + BackgroundInteractionPollContext, + maybe_schedule_background_interaction_cost_polling, + maybe_settle_background_interaction_before_delete, + poll_and_log_background_interaction_cost, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.types.interactions import InteractionsAPIResponse + +USAGE_BLOCK = { + "total_tokens": 175, + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_cached_tokens": 0, + "total_output_tokens": 50, + "output_tokens_by_modality": [{"modality": "text", "tokens": 50}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 25, +} + + +def _logging_obj( + call_type: str = "acreate_interaction", + litellm_params: Optional[dict] = None, +) -> LitellmLogging: + logging_obj = LitellmLogging( + model="gemini-2.5-flash", + messages=[], + stream=False, + call_type=call_type, + start_time=time.time(), + litellm_call_id="bg-interactions-call-id", + function_id="bg-interactions-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params=litellm_params or {}, + optional_params={}, + model="gemini-2.5-flash", + custom_llm_provider="gemini", + input="hi", + ) + return logging_obj + + +def _reservation() -> dict: + return {"reserved_cost": 0.05, "entries": [], "finalized": False, "input_cost": 0.001} + + +def _logging_obj_with_reservation(reservation: dict) -> LitellmLogging: + return _logging_obj(litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}) + + +async def _raise_on_billing(result: InteractionsAPIResponse) -> None: + raise RuntimeError("cost calculation failed for a settled background interaction") + + +def _context(logging_obj: LitellmLogging, timeout_seconds: float = 1.0) -> BackgroundInteractionPollContext: + return BackgroundInteractionPollContext( + interaction_id="interactions/bg-abc", + custom_llm_provider="gemini", + logging_obj=logging_obj, + initial_interval_seconds=0.001, + max_interval_seconds=0.002, + timeout_seconds=timeout_seconds, + ) + + +def _response(status: str, with_usage: bool) -> InteractionsAPIResponse: + return InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-2.5-flash", + status=status, + steps=[], + usage=dict(USAGE_BLOCK) if with_usage else None, + ) + + +def _fetch_sequence(*responses): + remaining = list(responses) + calls = [] + + async def fetch(context): + calls.append(context.interaction_id) + item = remaining.pop(0) if len(remaining) > 1 else remaining[0] + if isinstance(item, Exception): + raise item + return item + + return fetch, calls + + +@pytest.mark.parametrize( + "initial, maximum", + [(0.0, 0.002), (0.001, 0.0), (-1.0, 0.002), (0.0, 0.0)], +) +def test_poll_intervals_stops_instead_of_looping_on_a_non_positive_interval(initial, maximum): + intervals = list(islice(_poll_intervals(initial=initial, maximum=maximum, timeout=3600.0), 10)) + + assert len(intervals) < 10 + assert all(interval > 0 for interval in intervals) + + +@pytest.mark.asyncio +async def test_poller_bills_once_when_interaction_completes(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + _response("in_progress", with_usage=False), + _response("completed", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 2 + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + + +@pytest.mark.asyncio +async def test_poller_bills_an_interaction_paused_for_a_tool_result(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + _response("in_progress", with_usage=False), + _response("requires_action", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 2 + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + + +@pytest.mark.asyncio +async def test_poller_does_not_pin_the_budget_for_an_interaction_paused_for_a_tool_result(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("requires_action", with_usage=True)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert logging_obj.model_call_details["response_cost"] > 0 + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_poller_stops_without_billing_on_terminal_status_without_usage(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(_response("failed", with_usage=False)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 1 + assert logging_obj.model_call_details.get("response_cost") is None + + +@pytest.mark.asyncio +async def test_poller_gives_up_after_timeout_without_billing(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(_response("in_progress", with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), + fetch_interaction=fetch, + ) + + assert len(calls) >= 2 + assert logging_obj.model_call_details.get("response_cost") is None + + +@pytest.mark.asyncio +async def test_poller_releases_budget_reservation_when_interaction_ends_without_usage(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("failed", with_usage=False)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_poller_releases_budget_reservation_on_timeout_give_up(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_poller_releases_budget_reservation_when_billing_raises(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + logging_obj.async_log_background_interaction_completion = _raise_on_billing + + with pytest.raises(RuntimeError): + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_poller_leaves_reservation_reconciliation_to_the_completion_event(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence( + _response("in_progress", with_usage=False), + _response("completed", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert logging_obj.model_call_details["response_cost"] > 0 + assert reservation["finalized"] is False + + +@pytest.mark.asyncio +async def test_poller_retries_after_fetch_error_and_still_bills(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + RuntimeError("transient network error"), + _response("completed", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 2 + assert logging_obj.model_call_details["response_cost"] > 0 + + +@pytest.mark.asyncio +async def test_schedule_creates_poll_task_for_in_progress_create(): + logging_obj = _logging_obj() + task = maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + ) + + assert isinstance(task, asyncio.Task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response,create_kwargs", + [ + (_response("completed", with_usage=True), {"litellm_logging_obj": "placeholder"}), + (_response("in_progress", with_usage=False), {}), + ("not a response", {"litellm_logging_obj": "placeholder"}), + ], +) +async def test_schedule_skips_non_pollable_results(response, create_kwargs): + if create_kwargs.get("litellm_logging_obj") == "placeholder": + create_kwargs = {"litellm_logging_obj": _logging_obj()} + + task = maybe_schedule_background_interaction_cost_polling( + response=response, + create_kwargs=create_kwargs, + custom_llm_provider="gemini", + ) + + assert task is None + + +def _register_poll(logging_obj: LitellmLogging, poll_fetch=None) -> asyncio.Task: + import litellm.interactions.background_cost_polling as bg + + if poll_fetch is None: + poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + context = _context(logging_obj) + task = asyncio.create_task(poll_and_log_background_interaction_cost(context, fetch_interaction=poll_fetch)) + bg._ACTIVE_POLLS[context.interaction_id] = bg._ActiveBackgroundPoll(task=task, context=context) + task.add_done_callback(lambda finished: bg._discard_poll(context.interaction_id, finished)) + return task + + +@pytest.mark.asyncio +async def test_delete_settlement_bills_an_interaction_paused_for_a_tool_result(): + logging_obj = _logging_obj() + task = _register_poll(logging_obj) + fetch, calls = _fetch_sequence(_response("requires_action", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert len(calls) == 1 + assert logging_obj.model_call_details["response_cost"] > 0 + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_bills_pending_background_interaction(): + logging_obj = _logging_obj() + task = _register_poll(logging_obj) + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert len(calls) == 1 + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_releases_reservation_when_still_in_progress(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + task = _register_poll(logging_obj) + fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + assert logging_obj.model_call_details.get("response_cost") is None + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_releases_reservation_when_prefetch_fails(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + task = _register_poll(logging_obj) + fetch, _ = _fetch_sequence(RuntimeError("interaction already deleted")) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + assert logging_obj.model_call_details.get("response_cost") is None + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_releases_reservation_when_billing_raises(): + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + task = _register_poll(logging_obj) + fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + logging_obj.async_log_background_interaction_completion = _raise_on_billing + + with pytest.raises(RuntimeError): + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is True + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_delete_settlement_ignores_interactions_without_pending_poll(): + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/never-polled", + fetch_interaction=fetch, + ) + + assert calls == [] + + +@pytest.mark.asyncio +async def test_delete_settlement_noop_after_poll_task_finished(): + logging_obj = _logging_obj() + poll_fetch, _ = _fetch_sequence(_response("completed", with_usage=True)) + task = _register_poll(logging_obj, poll_fetch=poll_fetch) + await asyncio.wait_for(task, timeout=5) + assert logging_obj.model_call_details["response_cost"] > 0 + + settle_fetch, settle_calls = _fetch_sequence(_response("completed", with_usage=True)) + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=settle_fetch, + ) + + assert settle_calls == [] + + +@pytest.mark.asyncio +async def test_delete_settlement_does_not_rebill_when_gate_already_claimed(): + logging_obj = _logging_obj() + logging_obj.model_call_details[_SETTLED_KEY] = True + task = _register_poll(logging_obj) + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + fetch_interaction=fetch, + ) + + assert len(calls) == 1 + assert logging_obj.model_call_details.get("response_cost") is None + await asyncio.wait_for(task, timeout=5) + + +@pytest.mark.asyncio +async def test_poller_exits_without_billing_once_settled_elsewhere(): + logging_obj = _logging_obj() + logging_obj.model_call_details[_SETTLED_KEY] = True + fetch, calls = _fetch_sequence(_response("completed", with_usage=True)) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert calls == [] + assert logging_obj.model_call_details.get("response_cost") is None + + +@pytest.mark.asyncio +async def test_schedule_respects_kill_switch(monkeypatch): + import litellm.interactions.background_cost_polling as module + + monkeypatch.setattr(module, "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", False) + + task = maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": _logging_obj()}, + custom_llm_provider="gemini", + ) + + assert task is None + + +def test_every_status_the_api_can_return_is_either_pollable_or_terminal(): + """ + The proxy bills a usage-less create in exactly two ways: it polls the + interaction until it settles, or it recognises the status as terminal and + settles immediately. A status in neither set is billed by nobody, alerts + nobody, and releases its budget reservation, which is the zero-spend bug + this whole module exists to fix. + + Pinned against the generated spec enum rather than a hand-written list, so + a status Google adds later breaks this test instead of silently shipping + another unbilled path. + """ + from litellm.interactions.background_cost_polling import _POLLABLE_STATUSES, _TERMINAL_STATUSES + from litellm.types.interactions.generated import Status1 + + spec_statuses = {member.value for member in Status1} + handled = _POLLABLE_STATUSES | _TERMINAL_STATUSES + + assert spec_statuses - handled == set() + assert handled - spec_statuses == set() + + +@pytest.mark.asyncio +async def test_schedule_creates_poll_task_for_queued_create(): + """ + ``queued`` is the API's not-started-yet state. It carries no usage, so the + create cannot bill it, and it is not terminal, so nothing settles it: + without a poll task it is never charged at all. + """ + logging_obj = _logging_obj() + task = maybe_schedule_background_interaction_cost_polling( + response=_response("queued", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + ) + + assert isinstance(task, asyncio.Task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +@pytest.mark.asyncio +async def test_poller_bills_an_interaction_that_started_out_queued(): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence( + _response("queued", with_usage=False), + _response("in_progress", with_usage=False), + _response("completed", with_usage=True), + ) + + await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert len(calls) == 3 + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + + +def test_poll_intervals_double_up_to_the_cap_and_stay_inside_the_timeout(): + """ + The degenerate cases are covered above; this pins the shape the proxy + actually ships, so an off-by-one in the doubling or in the remaining-budget + check cannot pass green. + """ + intervals = list(_poll_intervals(initial=5.0, maximum=60.0, timeout=3600.0)) + + assert intervals[:6] == [5.0, 10.0, 20.0, 40.0, 60.0, 60.0] + assert max(intervals) == 60.0 + assert sum(intervals) <= 3600.0 + assert sum(intervals) + 60.0 > 3600.0 + + +@pytest.mark.asyncio +async def test_giving_up_on_an_unrecognized_status_says_which_status_it_was(monkeypatch): + """ + A status outside both sets polls for the full timeout and then gives up. + The give-up line is the only trace it leaves, so it has to name the status + rather than reporting it as an interaction that was merely still running. + """ + import litellm.interactions.background_cost_polling as bg + + errors = [] + monkeypatch.setattr(bg.verbose_logger, "error", lambda *args, **kwargs: errors.append(args)) + + logging_obj = _logging_obj() + fetch, _ = _fetch_sequence(_response("halted_for_review", with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.01), fetch_interaction=fetch + ) + + assert len(errors) == 1 + assert "halted_for_review" in errors[0] diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index c8c36032793..6f513ce1bd4 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -478,10 +478,10 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_m ], ) def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model): - """Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K.""" + """Bedrock GPT-5.6 enforces a 1,050,000-token context window, billed at the long-context rates above 272K.""" model_cost_map = litellm.model_cost[model] - assert model_cost_map["max_input_tokens"] == 1000000 + assert model_cost_map["max_input_tokens"] == 1050000 cached_tokens = 100000 completion_tokens = 1000 diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py new file mode 100644 index 00000000000..2d8092959ce --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_usage_object_transformation.py @@ -0,0 +1,150 @@ +from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( + InteractionsUsageObjectTransformation, +) +from litellm.types.utils import Usage + +OMNI_VIDEO_USAGE = { + "total_tokens": 18247, + "total_input_tokens": 16, + "input_tokens_by_modality": [{"modality": "text", "tokens": 16}], + "total_cached_tokens": 0, + "total_output_tokens": 17937, + "output_tokens_by_modality": [{"modality": "video", "tokens": 17376}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 294, +} + + +def test_detects_interactions_usage_object(): + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(OMNI_VIDEO_USAGE) is True + + +def test_rejects_chat_and_responses_api_usage_objects(): + chat_usage = {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30} + responses_api_usage = {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30} + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(chat_usage) is False + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(responses_api_usage) is False + assert InteractionsUsageObjectTransformation.is_interactions_usage_object(None) is False + assert InteractionsUsageObjectTransformation.is_interactions_usage_object("usage") is False + + +def test_transforms_real_omni_video_usage_block(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(OMNI_VIDEO_USAGE) + + assert isinstance(usage, Usage) + assert usage.prompt_tokens == 16 + assert usage.completion_tokens == 17937 + 294 + assert usage.total_tokens == 18247 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 16 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.video_tokens == 17376 + assert usage.completion_tokens_details.reasoning_tokens == 294 + + +def test_transforms_reasoning_tokens_spec_field_name(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 10, + "total_output_tokens": 20, + "total_reasoning_tokens": 5, + } + ) + assert usage.completion_tokens == 25 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 + assert usage.total_tokens == 35 + + +def test_cached_tokens_subtracted_from_text_input(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 1000, + "input_tokens_by_modality": [{"modality": "text", "tokens": 1000}], + "total_cached_tokens": 400, + "total_output_tokens": 50, + } + ) + assert usage.prompt_tokens == 1000 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 600 + assert usage.prompt_tokens_details.cached_tokens == 400 + assert usage._cache_read_input_tokens == 400 + + +def test_cached_tokens_subtracted_per_modality_when_breakdown_present(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 1500, + "input_tokens_by_modality": [ + {"modality": "text", "tokens": 1000}, + {"modality": "audio", "tokens": 500}, + ], + "total_cached_tokens": 300, + "cached_tokens_by_modality": [{"modality": "audio", "tokens": 300}], + "total_output_tokens": 50, + } + ) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 1000 + assert usage.prompt_tokens_details.audio_tokens == 200 + assert usage.prompt_tokens_details.cached_tokens == 300 + + +def test_tool_use_tokens_billed_as_input(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_tool_use_tokens": 40, + "tool_use_tokens_by_modality": [{"modality": "text", "tokens": 40}], + "total_output_tokens": 10, + } + ) + assert usage.prompt_tokens == 140 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 140 + + +def test_google_search_grounding_count_maps_to_web_search_requests(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 103, + "input_tokens_by_modality": [{"modality": "text", "tokens": 103}], + "total_output_tokens": 226, + "total_thought_tokens": 351, + "grounding_tool_count": [ + {"type": "google_search", "count": 3}, + {"type": "url_context", "count": 2}, + ], + } + ) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.web_search_requests == 3 + + +def test_no_grounding_leaves_web_search_requests_unset(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 10, + "input_tokens_by_modality": [{"modality": "text", "tokens": 10}], + "total_output_tokens": 5, + } + ) + assert usage.prompt_tokens_details is not None + assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None + + +def test_document_modality_folds_into_text(): + usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object( + { + "total_input_tokens": 80, + "input_tokens_by_modality": [ + {"modality": "text", "tokens": 30}, + {"modality": "document", "tokens": 50}, + ], + "total_output_tokens": 10, + } + ) + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.text_tokens == 80 diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 44fa8fc8ae0..44f91c98d81 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -19,9 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( def test_get_format_from_file_id(): - unified_file_id = ( - "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf" - ) + unified_file_id = "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf" format = get_format_from_file_id(unified_file_id) @@ -48,9 +46,7 @@ def test_update_messages_with_model_file_ids(): model_file_id_mapping = {file_id: {"my_model_id": "provider_file_id"}} - updated_messages = update_messages_with_model_file_ids( - messages, model_id, model_file_id_mapping - ) + updated_messages = update_messages_with_model_file_ids(messages, model_id, model_file_id_mapping) assert updated_messages == [ { @@ -143,9 +139,7 @@ def test_add_system_prompt_to_messages_merge_with_first_system(): {"role": "system", "content": "Existing system prompt."}, {"role": "user", "content": "Hello"}, ] - result = add_system_prompt_to_messages( - messages, "You are helpful.", merge_with_first_system=True - ) + result = add_system_prompt_to_messages(messages, "You are helpful.", merge_with_first_system=True) assert result == [ {"role": "system", "content": "You are helpful.\n\nExisting system prompt."}, {"role": "user", "content": "Hello"}, @@ -155,9 +149,7 @@ def test_add_system_prompt_to_messages_merge_with_first_system(): def test_add_system_prompt_to_messages_merge_with_first_system_adds_new_when_no_system(): """When merge_with_first_system=True but no system message, adds new one at start.""" messages = [{"role": "user", "content": "Hello"}] - result = add_system_prompt_to_messages( - messages, "You are helpful.", merge_with_first_system=True - ) + result = add_system_prompt_to_messages(messages, "You are helpful.", merge_with_first_system=True) assert result == [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hello"}, @@ -492,14 +484,8 @@ def test_update_messages_with_model_file_ids_tolerates_non_dict_content_items(): messages_token_ids_batch = [{"role": "user", "content": [[15496, 995], [9906, 0]]}] # Both should pass through unchanged without raising. - assert ( - update_messages_with_model_file_ids(messages_token_ids, "model-A", {}) - == messages_token_ids - ) - assert ( - update_messages_with_model_file_ids(messages_token_ids_batch, "model-A", {}) - == messages_token_ids_batch - ) + assert update_messages_with_model_file_ids(messages_token_ids, "model-A", {}) == messages_token_ids + assert update_messages_with_model_file_ids(messages_token_ids_batch, "model-A", {}) == messages_token_ids_batch class TestExtractFileDataBareStr: @@ -645,9 +631,7 @@ class TestUnpackLegacyDefs: definitions = { f"L{i}": { "type": "object", - "properties": { - f"x{j}": {"$ref": f"#/definitions/L{i + 1}"} for j in range(fanout) - }, + "properties": {f"x{j}": {"$ref": f"#/definitions/L{i + 1}"} for j in range(fanout)}, } for i in range(depth) } @@ -712,9 +696,7 @@ class TestUnpackLegacyDefs: schema = { "type": "object", - "properties": { - f"r{i}": {"$ref": f"#/components/schemas/T{i}"} for i in range(50) - }, + "properties": {f"r{i}": {"$ref": f"#/components/schemas/T{i}"} for i in range(50)}, "components": { "schemas": { f"T{i}": { @@ -739,9 +721,7 @@ class TestTextCompletionPromptToMessages: text_completion_prompt_to_messages, ) - assert text_completion_prompt_to_messages("summarize this") == ( - {"role": "user", "content": "summarize this"}, - ) + assert text_completion_prompt_to_messages("summarize this") == ({"role": "user", "content": "summarize this"},) def test_list_of_strings_becomes_one_message_each(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -970,3 +950,80 @@ class TestCustomToolFormatShapeConversion: for weird in ({}, {"type": "grammar"}, {"type": "future_format", "x": 1}): assert convert_custom_tool_format_to_chat_shape(dict(weird)) in (weird, {"type": "grammar", "grammar": {}}) assert convert_custom_tool_format_to_responses_shape(dict(weird)) == weird + + +# --- x-litellm-model upload-path decoding (litellm #29830) ------------------- + + +def _xlitellm_encoded(raw_id: str, model: str) -> str: + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + return encode_file_id_with_model(raw_id, model) + + +def test_update_messages_with_model_file_ids_decodes_xlitellm_encoded_id(): + """x-litellm-model upload returns `file-;model,)>`. + Without decoding, the encoded id leaks to upstream OpenAI and errors as + 'Files [...] were not found'. Decode it back to raw provider id.""" + raw_id = "file-ExTuCawUqxEMjVFK6xwR9B" + encoded_id = _xlitellm_encoded(raw_id, "gpt-5.1") + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this."}, + {"type": "file", "file": {"file_id": encoded_id}}, + ], + } + ] + + updated = update_messages_with_model_file_ids(messages, "model-A", {}) + + assert updated[0]["content"][1]["file"]["file_id"] == raw_id + + +def test_update_responses_input_with_model_file_ids_decodes_xlitellm_encoded_id(): + """Same bug on /v1/responses path. Without decoding the encoded id (>64 + chars), OpenAI rejects with 'string too long. Expected ... maximum length + 64'.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + update_responses_input_with_model_file_ids, + ) + + raw_id = "file-ExTuCawUqxEMjVFK6xwR9B" + encoded_id = _xlitellm_encoded(raw_id, "gpt-5.1") + input_items = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize."}, + {"type": "input_file", "file_id": encoded_id}, + ], + } + ] + + updated = update_responses_input_with_model_file_ids(input_items) + + assert updated[0]["content"][1]["file_id"] == raw_id + + +def test_update_messages_xlitellm_decode_does_not_override_mapping(): + """If the call-site already resolved a provider id via the mapping, that + wins. The new decode fallback runs only when no mapping match.""" + raw_id = "file-ExTuCawUqxEMjVFK6xwR9B" + encoded_id = _xlitellm_encoded(raw_id, "gpt-5.1") + mapping = {encoded_id: {"model-A": "provider-explicit-id"}} + messages = [ + { + "role": "user", + "content": [ + {"type": "file", "file": {"file_id": encoded_id}}, + ], + } + ] + + updated = update_messages_with_model_file_ids(messages, "model-A", mapping) + + assert updated[0]["content"][0]["file"]["file_id"] == "provider-explicit-id" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 3a7e06d085a..6265779b90d 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1,6 +1,8 @@ import base64 import json +import logging import os +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -3309,6 +3311,66 @@ def test_get_tool_calls_from_response_include_all_choices_reads_every_choice(): assert names == ["tool_alpha", "tool_beta"] +def test_get_tool_calls_from_response_silences_redacted_arguments(caplog): + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) + + response: Final = { + "choices": [ + { + "message": { + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "Read", + "arguments": "redacted-by-litellm", + }, + } + ] + } + } + ] + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tool_calls: Final = get_tool_calls_from_response(response) + + assert tool_calls == [{"id": "call_1", "name": "Read", "arguments": {}}] + assert "Failed to parse tool call arguments" not in caplog.text + + +def test_get_tool_calls_from_response_warns_for_malformed_arguments(caplog): + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) + + response: Final = { + "choices": [ + { + "message": { + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "Read", + "arguments": "not-json", + }, + } + ] + } + } + ] + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tool_calls: Final = get_tool_calls_from_response(response) + + assert tool_calls == [{"id": "call_1", "name": "Read", "arguments": {}}] + assert "Failed to parse tool call arguments" in caplog.text + + def test_group_tool_exchanges_pairs_assistant_with_its_tool_rows(): from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index aa7f2990c13..33a33711698 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -255,3 +255,26 @@ class TestRedactNestedMatchAndRegexKeys: def test_passes_through_none_and_str(self): assert redact_nested_match_and_regex_keys(None) is None assert redact_nested_match_and_regex_keys("plain") == "plain" + + +class TestIsExpectedClientError: + def test_status_ranges(self): + from litellm.litellm_core_utils.core_helpers import is_expected_client_error + + class WithStatusCode(Exception): + def __init__(self, status_code): + self.status_code = status_code + + class WithCode(Exception): + def __init__(self, code): + self.code = code + + assert is_expected_client_error(WithStatusCode(400)) is True + assert is_expected_client_error(WithStatusCode(429)) is True + assert is_expected_client_error(WithStatusCode(499)) is True + assert is_expected_client_error(WithStatusCode(500)) is False + assert is_expected_client_error(WithStatusCode(399)) is False + assert is_expected_client_error(WithCode("403")) is True + assert is_expected_client_error(WithCode("invalid_request_error")) is False + assert is_expected_client_error(Exception("no status")) is False + assert is_expected_client_error(None) is False diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 599ad016827..6f7ea9da640 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -867,6 +867,7 @@ PROVIDERS_WITH_A_HANDLER = ( "openrouter", "perplexity", "replicate", + "runwayml", "sagemaker", "together_ai", "vertex_ai", @@ -956,6 +957,7 @@ PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW = ( "mistral", "openai", "perplexity", + "runwayml", "together_ai", "vertex_ai", "xai", @@ -971,6 +973,7 @@ PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK = ( "mistral", "openai", "perplexity", + "runwayml", "together_ai", "xai", ) diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index 882429fd7cd..0587628e2fe 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -378,7 +378,7 @@ def test_shipped_rules_flag_unmapped_fable_as_always_on_thinking(shipped_cost_ma "model,provider", [ ("claude-opus-4-9@20260101", "vertex_ai"), - ("databricks-claude-opus-5-1", "databricks"), + ("databricks-claude-haiku-5-1", "databricks"), ], ) def test_shipped_rules_are_provider_neutral_for_unmapped_ids(shipped_cost_map, model, provider): @@ -388,6 +388,8 @@ def test_shipped_rules_are_provider_neutral_for_unmapped_ids(shipped_cost_map, m assert info["supports_adaptive_thinking"] is True assert info["supports_mid_conversation_system"] is True assert info["supports_function_calling"] is True + assert not info.get("input_cost_per_token") + assert not info.get("output_cost_per_token") @pytest.mark.parametrize( diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py index bda7ab4afc6..6cacd119030 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py +++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py @@ -133,3 +133,54 @@ class TestGetLlmProviderRejectsAttackerSmuggledApiBase: assert provider == "groq" assert dynamic_api_key == "server-real-groq-key" + + +class TestTogetherApiBaseResolvesProvider: + """ + Regression for the Together host migration: both the current + ``api.together.ai`` host and the legacy ``api.together.xyz`` host must + resolve to ``together_ai`` when passed as ``api_base``. Before the fix + the endpoint list carried the legacy host but the provider-mapping + chain had no branch for it, so the match fell through with a None + provider and the deployment failed with "LLM Provider NOT provided". + """ + + @pytest.mark.parametrize( + "api_base", + [ + "https://api.together.ai/v1", + "https://api.together.xyz/v1", + ], + ) + def test_together_api_base_resolves_to_together_ai(self, api_base, monkeypatch): + monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env") + + model, provider, dynamic_api_key, returned_api_base = get_llm_provider( + model="some-model", + api_base=api_base, + ) + + assert provider == "together_ai" + assert dynamic_api_key == "together-key-from-env" + assert returned_api_base == api_base + assert model == "some-model" + + def test_explicit_api_key_beats_together_env_key(self, monkeypatch): + monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env") + + _, provider, dynamic_api_key, _ = get_llm_provider( + model="some-model", + api_base="https://api.together.ai/v1", + api_key="explicit-caller-key", + ) + + assert provider == "together_ai" + assert dynamic_api_key == "explicit-caller-key" + + def test_together_default_api_base_is_together_ai(self, monkeypatch): + monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False) + + _, provider, _, api_base = get_llm_provider(model="together_ai/some-model") + + assert provider == "together_ai" + assert api_base == "https://api.together.ai/v1" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1c706be51fa..f93daa61570 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - import time import httpx @@ -3867,6 +3866,90 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +# ── Azure Model Router selected-model attribution ──────────────────────────── + + +def _model_router_response(selected_model: str, stamp: bool): + """A ModelResponse as AzureModelRouterConfig hands it back, with or without the stamp.""" + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + from litellm.types.utils import ModelResponse + + response = ModelResponse(model=selected_model) + response._hidden_params = ( + {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} + ) + return response + + +def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): + """ + The selected model must win off the stamp, not off "model-router" appearing in the + requested model. An operator whose model group is named anything else was invisible + to the name check, so their logs and spend rows named the router instead. + """ + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "azure_ai/smart-pick", + "custom_llm_provider": "azure_ai", + "messages": [], + "litellm_params": {"metadata": {}}, + }, + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=True + ), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning" + + +def test_standard_logging_payload_keeps_requested_model_without_router_stamp( + logging_obj, +): + """ + Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp + is what redirects attribution rather than the response model winning unconditionally. + """ + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "azure_ai/smart-pick", + "custom_llm_provider": "azure_ai", + "messages": [], + "litellm_params": {"metadata": {}}, + }, + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=False + ), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["model"] == "azure_ai/smart-pick" + + def _make_dict_logging_obj(): """Build a Logging instance configured for a non-streaming dict result.""" obj = LitellmLogging( @@ -4546,6 +4629,323 @@ def test_zero_token_video_usage_preserves_duration_seconds(logging_obj): assert payload["completion_tokens"] == 0 +INTERACTIONS_USAGE_BLOCK = { + "total_tokens": 175, + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_cached_tokens": 0, + "total_output_tokens": 50, + "output_tokens_by_modality": [{"modality": "text", "tokens": 50}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 25, +} + + +def _interactions_logging_obj(stream: bool, call_type: str = "acreate"): + logging_obj = LitellmLogging( + model="gemini-2.5-flash", + messages=[], + stream=stream, + call_type=call_type, + start_time=time.time(), + litellm_call_id="interactions-call-id", + function_id="interactions-fn-id", + ) + logging_obj.update_environment_variables( + litellm_params={}, + optional_params={}, + model="gemini-2.5-flash", + custom_llm_provider="gemini", + input="hi", + ) + return logging_obj + + +@pytest.mark.parametrize("call_type", ["create", "acreate", "create_interaction", "acreate_interaction"]) +def test_interactions_response_is_recognized_for_logging(call_type): + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False, call_type=call_type) + response = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is True + + +@pytest.mark.parametrize("call_type", ["acreate", "acreate_interaction"]) +def test_in_progress_background_create_is_not_billed(call_type): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False, call_type=call_type) + response = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + + assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is False + + logging_obj._success_handler_helper_fn( + result=response, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + cache_hit=False, + ) + + assert logging_obj.model_call_details.get("response_cost") is None + assert logging_obj.model_call_details.get("standard_logging_object") is None + + +@pytest.mark.asyncio +async def test_background_interaction_completion_rebills_after_in_progress_success(): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False) + in_progress = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + await logging_obj.async_success_handler( + result=in_progress, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + ) + + assert logging_obj.model_call_details.get("response_cost") is None + assert logging_obj.should_run_logging(event_type="async_success") is False + + completed = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + await logging_obj.async_log_background_interaction_completion(result=completed) + + assert logging_obj.model_call_details["response_cost"] > 0 + assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175 + + +@pytest.mark.asyncio +async def test_background_interaction_completion_prices_the_settled_body_itself(): + """ + The poll fetches the settled body through its own client call, which + prices it against a throwaway logging object holding none of this + request's deployment context. Adopting that price would bill a + custom-priced deployment at the wrong rate, and it would also satisfy the + "already calculated" shortcut and skip repricing, leaving the breakdown at + the zeros the usage-less create stamped and writing those to the spend log. + """ + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False) + in_progress = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + await logging_obj.async_success_handler( + result=in_progress, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + ) + + completed = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + completed._hidden_params = {"response_cost": 99.0} + + await logging_obj.async_log_background_interaction_completion(result=completed) + + response_cost = logging_obj.model_call_details["response_cost"] + assert response_cost != 99.0 + assert response_cost > 0 + + cost_breakdown = logging_obj.model_call_details["standard_logging_object"]["cost_breakdown"] + assert cost_breakdown["total_cost"] == response_cost + assert cost_breakdown["input_cost"] > 0 + assert cost_breakdown["output_cost"] > 0 + + +@pytest.mark.asyncio +async def test_background_interaction_completion_lets_otel_emit_the_cost_span(): + """ + OTEL, and every integration that derives from it, dedupes span emission on + a marker kept in the request's own metadata. The in-progress create claims + that marker, so without clearing it the settled completion, the only event + carrying usage and cost, is discarded as a duplicate and every + OTEL-family backend shows the interaction as a span with no cost at all. + """ + import datetime as dt + + from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig + from litellm.types.interactions import InteractionsAPIResponse + + otel = OpenTelemetry(config=OpenTelemetryConfig(exporter="console")) + logging_obj = _interactions_logging_obj(stream=False) + in_progress = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress") + await logging_obj.async_success_handler( + result=in_progress, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + ) + + assert otel._emit_once(logging_obj.model_call_details, "success") is True + assert otel._emit_once(logging_obj.model_call_details, "success") is False + + completed = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + await logging_obj.async_log_background_interaction_completion(result=completed) + + assert otel._emit_once(logging_obj.model_call_details, "success") is True + + +@pytest.mark.parametrize( + "call_type", + ["aget", "get", "aget_interaction", "adelete_interaction", "acancel_interaction"], +) +def test_interactions_get_poll_is_not_billed(call_type): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False, call_type=call_type) + response = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + assert logging_obj._is_recognized_call_type_for_logging(logging_result=response) is False + + logging_obj._success_handler_helper_fn( + result=response, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + cache_hit=False, + ) + + assert logging_obj.model_call_details.get("response_cost") is None + assert logging_obj.model_call_details.get("standard_logging_object") is None + + +def test_non_streaming_interactions_success_sets_response_cost_and_usage(): + import datetime as dt + + from litellm.types.interactions import InteractionsAPIResponse + + logging_obj = _interactions_logging_obj(stream=False) + response = InteractionsAPIResponse( + id="interactions/abc", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + logging_obj._success_handler_helper_fn( + result=response, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + cache_hit=False, + ) + + assert logging_obj.model_call_details["response_cost"] > 0 + standard_logging_object = logging_obj.model_call_details["standard_logging_object"] + assert standard_logging_object["prompt_tokens"] == 100 + assert standard_logging_object["completion_tokens"] == 75 + assert standard_logging_object["total_tokens"] == 175 + assert standard_logging_object["response_cost"] == logging_obj.model_call_details["response_cost"] + + +def test_assembled_streaming_response_from_completed_interaction_event(): + import datetime as dt + + from litellm.types.interactions import ( + InteractionsAPIResponse, + InteractionsAPIStreamingResponse, + ) + + logging_obj = _interactions_logging_obj(stream=True) + completed_event = InteractionsAPIStreamingResponse( + event_type="interaction.completed", + interaction={ + "id": "interactions/abc", + "model": "gemini-2.5-flash", + "status": "completed", + "steps": [], + "usage": dict(INTERACTIONS_USAGE_BLOCK), + }, + ) + + assembled = logging_obj._get_assembled_streaming_response( + result=completed_event, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + is_async=True, + streaming_chunks=[], + ) + + assert isinstance(assembled, InteractionsAPIResponse) + assert assembled.usage == INTERACTIONS_USAGE_BLOCK + + in_progress_event = InteractionsAPIStreamingResponse(event_type="interaction.in_progress") + assert ( + logging_obj._get_assembled_streaming_response( + result=in_progress_event, + start_time=dt.datetime.now(), + end_time=dt.datetime.now(), + is_async=True, + streaming_chunks=[], + ) + is None + ) + + +def test_assembled_streaming_response_from_legacy_completed_chunk(): + from litellm.types.interactions import ( + InteractionsAPIResponse, + InteractionsAPIStreamingResponse, + ) + + legacy_chunk = InteractionsAPIStreamingResponse( + event_type="interaction.complete", + id="interactions/legacy", + model="gemini-2.5-flash", + status="completed", + outputs=[], + usage=dict(INTERACTIONS_USAGE_BLOCK), + ) + + assembled = LitellmLogging._assemble_completed_interaction_response(legacy_chunk) + + assert isinstance(assembled, InteractionsAPIResponse) + assert assembled.id == "interactions/legacy" + assert assembled.usage == INTERACTIONS_USAGE_BLOCK + + +def test_standard_logging_payload_maps_interactions_usage(): + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + usage = StandardLoggingPayloadSetup.get_usage_from_response_obj( + response_obj={"usage": dict(INTERACTIONS_USAGE_BLOCK)} + ) + + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 75 + assert usage.total_tokens == 175 + + def test_pre_call_does_not_pin_request_in_module_state(logging_obj): """ pre_call/post_call must not stash their locals (full messages, the Logging @@ -5278,3 +5678,62 @@ def test_get_custom_logger_compatible_class_finds_v2_newrelic(monkeypatch): logging_module._in_memory_loggers.clear() monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) is_otel_v2_enabled.cache_clear() + + +class _ClientError(Exception): + def __init__(self, status_code, message): + self.status_code = status_code + self.message = message + super().__init__(message) + + +def _raise_and_catch(exc): + try: + raise exc + except Exception as caught: + return caught + + +def test_get_error_information_skips_traceback_for_expected_4xx(monkeypatch): + """Regression for LIT-6043: expected client (4xx) errors must not pay for + traceback.format_tb on every rejected request unless + litellm.log_client_error_tracebacks is enabled.""" + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + client_exc = _raise_and_catch(_ClientError(status_code=403, message="team does not allow model")) + assert client_exc.__traceback__ is not None + result = StandardLoggingPayloadSetup.get_error_information(client_exc) + assert result["traceback"] == "" + + server_exc = _raise_and_catch(_ClientError(status_code=500, message="boom")) + result = StandardLoggingPayloadSetup.get_error_information(server_exc) + assert "test_litellm_logging" in result["traceback"] + + monkeypatch.setattr(litellm, "log_client_error_tracebacks", True) + result = StandardLoggingPayloadSetup.get_error_information(client_exc) + assert "test_litellm_logging" in result["traceback"] + + +def test_failure_handler_helper_fn_builds_payload_once_per_exception(): + """Regression for LIT-6043: async and sync failure handlers both call + _failure_handler_helper_fn for the same failed request; the standardized + payload must be built once, not once per handler.""" + obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6043-1", + function_id="f", + ) + exc = _raise_and_catch(_ClientError(status_code=400, message="invalid model")) + obj._failure_handler_helper_fn(exception=exc, traceback_exception="") + first_payload = obj.model_call_details["standard_logging_object"] + assert first_payload is not None + obj._failure_handler_helper_fn(exception=exc, traceback_exception="") + assert obj.model_call_details["standard_logging_object"] is first_payload + + other_exc = _raise_and_catch(_ClientError(status_code=429, message="rate limited")) + obj._failure_handler_helper_fn(exception=other_exc, traceback_exception="") + assert obj.model_call_details["standard_logging_object"] is not first_payload diff --git a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py new file mode 100644 index 00000000000..3a09702de45 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py @@ -0,0 +1,99 @@ +import httpx + +from litellm.litellm_core_utils.llm_request_utils import ( + flatten_form_field_values, + serialize_multipart_form_fields, +) + + +def _multipart_field_names(data: dict) -> list[str]: + request = httpx.Request( + "POST", + "http://backend/v1/images/edits", + data=data, + files=[("image[]", ("in.png", b"stub", "image/png"))], + ) + request.read() + body = request.content.decode("utf-8", "replace") + prefix = 'Content-Disposition: form-data; name="' + return [line[len(prefix) : line.index('"', len(prefix))] for line in body.splitlines() if line.startswith(prefix)] + + +def test_serialize_multipart_form_fields_flattens_like_the_openai_sdk(): + fields = serialize_multipart_form_fields( + { + "model": "sora-2", + "prompt": "a cat surfing", + "hd": True, + "watermark": False, + "seconds": 4, + "size": None, + "metadata": {"trace": {"id": "t1"}}, + "characters": [{"id": "char_1", "name": "Mia"}, "solo"], + } + ) + + assert fields == ( + ("model", (None, "sora-2")), + ("prompt", (None, "a cat surfing")), + ("hd", (None, "true")), + ("watermark", (None, "false")), + ("seconds", (None, "4")), + ("metadata[trace][id]", (None, "t1")), + ("characters[][id]", (None, "char_1")), + ("characters[][name]", (None, "Mia")), + ("characters[]", (None, "solo")), + ) + + +def test_serialize_multipart_form_fields_drops_empty_strings(): + assert serialize_multipart_form_fields({"prompt": "", "model": "sora-2"}) == (("model", (None, "sora-2")),) + + +def test_serialize_multipart_form_fields_empty_body(): + assert serialize_multipart_form_fields({}) == () + + +def test_flatten_form_field_values_flattens_nested_and_drops_empty(): + assert flatten_form_field_values( + { + "seed": 42, + "hd": True, + "size": None, + "prompt": "", + "generation_config": {"steps": 30, "guidance": True}, + } + ) == ( + ("seed", "42"), + ("hd", "true"), + ("generation_config[steps]", "30"), + ("generation_config[guidance]", "true"), + ) + + +def test_flatten_form_field_values_later_source_wins_on_collision(): + assert flatten_form_field_values({"seed": 1}, None, {"seed": 2}) == ( + ("seed", "1"), + ("seed", "2"), + ) + assert dict(flatten_form_field_values({"seed": 1}, {"seed": 2}))["seed"] == "2" + + +def test_flatten_form_field_values_keeps_scalar_lists_as_repeated_fields(): + assert flatten_form_field_values( + {"loras": ["a", "b", "c"], "generation_config": {"tags": [1, 2]}, "seed": 42} + ) == ( + ("loras", ("a", "b", "c")), + ("generation_config[tags]", ("1", "2")), + ("seed", "42"), + ) + + +def test_flatten_form_field_values_scalar_list_survives_update_into_multipart(): + request_params: dict = {"model": "my-edit-model"} + request_params.update(flatten_form_field_values({"loras": ["style_a", "style_b"]})) + + names = _multipart_field_names(request_params) + + assert names.count("loras") == 2 + assert names.count("model") == 1 diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 978f22ca2e4..80c9585dae9 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -413,3 +413,42 @@ class TestLoggingWorker: assert worker2._bound_loop is not None await worker2.stop() + + def test_event_loop_change_carries_pending_tasks_over(self): + """Regression (LIT-6028): a loop change must not silently drop queued coroutines. + + Before the fix ``_ensure_queue`` nulled ``self._queue`` on a loop change, discarding + every pending ``LoggingTask`` (each an un-awaited spend-logging coroutine). The tasks + must instead be moved onto the queue bound to the new loop and still execute there. + """ + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + executed: list[int] = [] + + async def spend_log(index: int) -> None: + executed.append(index) + + async def enqueue_on_first_loop() -> None: + worker._ensure_queue() + for i in range(5): + worker.enqueue(spend_log(i)) + assert worker._queue is not None + assert worker._queue.qsize() == 5 + + asyncio.run(enqueue_on_first_loop()) + + stale_queue = worker._queue + assert stale_queue is not None + + async def rebind_on_second_loop() -> None: + worker._ensure_queue() + assert worker._queue is not None + # A fresh queue bound to the new loop, holding every carried-over task (not dropped). + assert worker._queue is not stale_queue + assert worker._queue.qsize() == 5 + while not worker._queue.empty(): + task = worker._queue.get_nowait() + await task["context"].run(asyncio.create_task, task["coroutine"]) + + asyncio.run(rebind_on_second_loop()) + + assert sorted(executed) == [0, 1, 2, 3, 4] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 44e77506b3f..4b5b51cb4b8 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -989,6 +989,32 @@ def test_cost_field_in_usage_chunks(): assert usage.completion_tokens == 5 +def test_stream_chunk_builder_tolerates_trailing_chunk_without_choices(): + """Regression for https://github.com/BerriAI/litellm/issues/32051 + + The Responses-API bridge yields ModelResponseStream chunks with choices + followed by a trailing event object that has no ``choices`` key. Building + those chunks used to raise ``KeyError('choices')`` (surfaced as a 500 + APIError); it must now skip the choices-less chunk and assemble content. + """ + from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject + + content_chunks = [ + ModelResponseStream( + model="gpt-4o", + choices=[StreamingChoices(index=0, delta=Delta(content=part))], + ) + for part in ("Hello", " world") + ] + trailing_chunk = BaseLiteLLMOpenAIResponseObject() + assert "choices" not in trailing_chunk + + response = stream_chunk_builder(chunks=content_chunks + [trailing_chunk]) + + assert response is not None + assert response.choices[0].message.content == "Hello world" + + def test_anthropic_speed_and_geo_survive_stream_assembly(): """Anthropic prices fast mode and non-global regions with a multiplier read off ``usage.speed`` / ``usage.inference_geo``. Dropping them while reassembling a stream diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index e4dacc308dc..d0169963962 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -359,6 +359,48 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks(): assert result[1]["tool_calls"][0]["id"] == "toolu_01234" +def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): + """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. + + Without it Moonshot and DeepSeek fill in a single-space placeholder and the model gets + a blank where its own prior reasoning belongs. + """ + + anthropic_messages = [ + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "text", "text": "Which city is best for a picnic?"}], + ), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "thinking", "thinking": "Denver is dry in August.", "signature": "sig1"}, + {"type": "thinking", "thinking": "San Francisco is foggy.", "signature": "sig2"}, + {"type": "redacted_thinking", "data": "REDACTED"}, + {"type": "text", "text": "Denver."}, + ], + ), + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages) + + assert result[1]["reasoning_content"] == "Denver is dry in August.\nSan Francisco is foggy." + assert result[1]["content"] == "Denver." + + +def test_translate_anthropic_messages_to_openai_sets_no_reasoning_content_without_thinking(): + anthropic_messages = [ + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[{"type": "text", "text": "Denver."}], + ), + ] + + result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=anthropic_messages) + + assert "reasoning_content" not in result[0] + + def test_translate_anthropic_messages_to_openai_tool_message_placement(): """Test that tool result messages are placed before user messages in the conversation order.""" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py index f393a7b50b1..c9170efd18a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py @@ -2,7 +2,6 @@ import pytest -import litellm from litellm.constants import ( DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, @@ -17,7 +16,6 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran ) - @pytest.mark.parametrize( "reasoning_effort,expected_effort", [ @@ -44,7 +42,7 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} @@ -258,19 +256,22 @@ def test_reasoning_effort_in_supported_params(): "model", [ "claude-sonnet-4-6", - "bedrock/invoke/us.anthropic.claude-sonnet-4-6", - "vertex_ai/claude-sonnet-4-6", "claude-opus-4-6", + "claude-sonnet-4-6-20260219", + "bedrock/invoke/us.anthropic.claude-sonnet-4-6", "bedrock/invoke/us.anthropic.claude-opus-4-6-v1:0", + "vertex_ai/claude-sonnet-4-6", "vertex_ai/claude-opus-4-6", + "azure_ai/claude-sonnet-4-6", ], ) -def test_legacy_thinking_high_budget_clamps_to_high_when_xhigh_unsupported( - local_model_cost_map, model -): - """Claude Code sends ``thinking.budget_tokens=31999``; Sonnet 4.6 and Opus 4.6 - have no ``xhigh`` tier, so the translator must emit ``high`` rather than the - provider-invalid ``xhigh`` (regression for issue #29282).""" +def test_legacy_thinking_budget_preserved_verbatim_on_46(local_model_cost_map, model): + """Regression for the passthrough silently dropping a caller's hard thinking + budget: the 4.6 family accepts ``thinking.type=enabled`` with ``budget_tokens`` + natively, so rewriting it to ``thinking.type=adaptive`` + ``output_config.effort`` + (which carries no ceiling) let reasoning run past the requested cap. The legacy + shape must be forwarded verbatim, in every 4.6 id shape including unmapped dated + releases resolved by the ``claude-legacy-thinking`` fallback rule.""" config = AnthropicMessagesConfig() optional_params = { "max_tokens": 1024, @@ -285,8 +286,8 @@ def test_legacy_thinking_high_budget_clamps_to_high_when_xhigh_unsupported( headers={}, ) - assert result.get("thinking") == {"type": "adaptive"} - assert result.get("output_config") == {"effort": "high"} + assert result.get("thinking") == {"type": "enabled", "budget_tokens": 31999} + assert "output_config" not in result def test_legacy_thinking_high_budget_keeps_xhigh_when_supported(): @@ -343,11 +344,44 @@ def test_legacy_thinking_translates_to_adaptive_for_opus_48( assert result.get("output_config") == {"effort": "xhigh"} +@pytest.mark.parametrize( + "model,expected_effort", + [ + ("claude-sonnet-5", "xhigh"), + ("claude-opus-5", "xhigh"), + ("claude-newfamily-6", "high"), + ], +) +def test_legacy_thinking_translates_to_adaptive_for_5_and_future_models( + local_model_cost_map, model, expected_effort +): + """The 5 families reject ``thinking.type=enabled``, so the adaptive translation + stays the safe default for every adaptive model not flagged + ``supports_legacy_thinking``, unmapped future ids included. An unmapped id + cannot prove ``xhigh`` support, so its high-budget bucket clamps to ``high``.""" + config = AnthropicMessagesConfig() + optional_params = { + "max_tokens": 1024, + "thinking": {"type": "enabled", "budget_tokens": 31999}, + } + + result = config.transform_anthropic_messages_request( + model=model, + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive"} + assert result.get("output_config") == {"effort": expected_effort} + + @pytest.mark.parametrize( "budget_tokens,expected_effort", [ - (DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET * 2, "high"), - (DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, "high"), + (DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET * 2, "xhigh"), + (DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, "xhigh"), (DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, "high"), (DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET - 1, "medium"), (DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, "medium"), @@ -355,7 +389,9 @@ def test_legacy_thinking_translates_to_adaptive_for_opus_48( (1, "low"), ], ) -def test_legacy_thinking_budget_buckets_on_sonnet_46(budget_tokens, expected_effort): +def test_legacy_thinking_budget_buckets_on_opus_48( + local_model_cost_map, budget_tokens, expected_effort +): config = AnthropicMessagesConfig() optional_params = { "max_tokens": 1024, @@ -363,7 +399,7 @@ def test_legacy_thinking_budget_buckets_on_sonnet_46(budget_tokens, expected_eff } result = config.transform_anthropic_messages_request( - model="claude-sonnet-4-6", + model="claude-opus-4-8", messages=[{"role": "user", "content": "Hello"}], anthropic_messages_optional_request_params=optional_params, litellm_params={}, @@ -373,7 +409,29 @@ def test_legacy_thinking_budget_buckets_on_sonnet_46(budget_tokens, expected_eff assert result.get("output_config") == {"effort": expected_effort} -def test_legacy_thinking_does_not_override_explicit_output_config(): +def test_legacy_thinking_does_not_override_explicit_output_config(local_model_cost_map): + config = AnthropicMessagesConfig() + optional_params = { + "max_tokens": 1024, + "thinking": {"type": "enabled", "budget_tokens": 31999}, + "output_config": {"effort": "low"}, + } + + result = config.transform_anthropic_messages_request( + model="claude-opus-4-8", + messages=[{"role": "user", "content": "Hello"}], + anthropic_messages_optional_request_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert result.get("thinking") == {"type": "adaptive"} + assert result.get("output_config") == {"effort": "low"} + + +def test_legacy_thinking_with_explicit_output_config_untouched_on_46( + local_model_cost_map, +): config = AnthropicMessagesConfig() optional_params = { "max_tokens": 1024, @@ -389,6 +447,7 @@ def test_legacy_thinking_does_not_override_explicit_output_config(): headers={}, ) + assert result.get("thinking") == {"type": "enabled", "budget_tokens": 31999} assert result.get("output_config") == {"effort": "low"} diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index f33bb3dda8b..652c1f077a9 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -11,6 +11,10 @@ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterato BaseAnthropicMessagesStreamingIterator, _incomplete_stream_error_sse_event, _is_message_stop_chunk, + _is_provider_error_chunk, + anthropic_messages_response_as_sse_events, + is_anthropic_content_delta_chunk, + parse_anthropic_error_event, ) @@ -157,6 +161,96 @@ def test_is_message_stop_chunk_ignores_substring_in_payload(): assert _is_message_stop_chunk(delta_frame_with_substring) is False +def test_parse_anthropic_error_event_from_dict_chunk(): + """Regression for #24004: dict-shaped error chunks parse to + (type, message, status) so the Router can decide whether to fall back.""" + chunk = {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} + assert parse_anthropic_error_event(chunk) == ("overloaded_error", "Overloaded", 503) + assert _is_provider_error_chunk(chunk) is True + + +def test_parse_anthropic_error_event_from_sse_bytes(): + """Regression for #24004: a raw `event: error` SSE frame (what a native + Anthropic/Bedrock passthrough forwards verbatim today) must parse + identically to the dict shape so the Router can raise a fallback.""" + sse_chunk = ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "internal_server_error", "message": "boom"}}\n\n' + ) + assert parse_anthropic_error_event(sse_chunk) == ("internal_server_error", "boom", 500) + assert _is_provider_error_chunk(sse_chunk) is True + + +def test_parse_anthropic_error_event_defaults_status_for_unknown_type(): + chunk = {"type": "error", "error": {"type": "some_future_error_type", "message": "?"}} + assert parse_anthropic_error_event(chunk) == ("some_future_error_type", "?", 500) + + +def test_parse_anthropic_error_event_missing_message_falls_back_to_type(): + chunk = {"type": "error", "error": {"type": "overloaded_error"}} + assert parse_anthropic_error_event(chunk) == ("overloaded_error", "overloaded_error", 503) + + +def test_parse_anthropic_error_event_non_string_error_type_returns_none(): + """A malformed error body whose `type` field isn't a string (e.g. an + upstream bug sends null or a number) must not be treated as an error + event rather than crashing or forwarding a garbage error_type.""" + chunk = {"type": "error", "error": {"type": None, "message": "boom"}} + assert parse_anthropic_error_event(chunk) is None + + +def test_decoded_sse_data_line_swallows_invalid_json(): + """A `data:` line that isn't valid JSON (a malformed/truncated frame) + must not be treated as an error event or raise, just be ignored.""" + malformed_frame = b"event: error\ndata: {not valid json\n\n" + assert parse_anthropic_error_event(malformed_frame) is None + assert _is_provider_error_chunk(malformed_frame) is False + + +class TestIsAnthropicContentDeltaChunk: + def test_dict_content_block_delta(self): + assert is_anthropic_content_delta_chunk({"type": "content_block_delta"}) is True + + def test_dict_other_type(self): + assert is_anthropic_content_delta_chunk({"type": "message_start"}) is False + + def test_bytes_content_block_delta(self): + assert is_anthropic_content_delta_chunk(b"event: content_block_delta\ndata: {}\n\n") is True + + def test_bytes_other_event(self): + assert is_anthropic_content_delta_chunk(b"event: message_start\ndata: {}\n\n") is False + + def test_neither_dict_nor_bytes(self): + assert is_anthropic_content_delta_chunk("content_block_delta") is False + assert is_anthropic_content_delta_chunk(None) is False + + +@pytest.mark.parametrize( + "chunk", + [ + {"type": "content_block_delta", "delta": {"type": "text_delta", "text": "hi"}}, + b'event: content_block_delta\ndata: {"type": "content_block_delta"}\n\n', + b"raw-bytes", + "error", + None, + ], +) +def test_parse_anthropic_error_event_non_error_chunks_return_none(chunk): + assert parse_anthropic_error_event(chunk) is None + assert _is_provider_error_chunk(chunk) is False + + +def test_parse_anthropic_error_event_ignores_substring_in_payload(): + """A content_block_delta whose partial_json happens to contain the + literal string `"type": "error"` must not be misread as an error event.""" + delta_frame_with_substring = ( + b"event: content_block_delta\n" + b'data: {"type": "content_block_delta", "delta": ' + b'{"type": "input_json_delta", "partial_json": "\\"type\\": \\"error\\""}}\n\n' + ) + assert parse_anthropic_error_event(delta_frame_with_substring) is None + + @pytest.mark.asyncio async def test_async_sse_wrapper_emits_error_when_bytes_stream_only_mentions_message_stop_in_payload(): """ @@ -307,3 +401,117 @@ def test_incomplete_stream_error_sse_event_is_valid_anthropic_error(): "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}, } assert event.endswith("\n\n") + + +def _decode_sse_events(events: tuple[bytes, ...]) -> list[tuple[str, dict]]: + decoded = [] + for event in events: + assert isinstance(event, bytes) + lines = event.decode().split("\n") + assert lines[0].startswith("event: ") + decoded.append((lines[0].removeprefix("event: "), json.loads(lines[1].removeprefix("data: ")))) + return decoded + + +def test_anthropic_messages_response_as_sse_events_text_block(): + response = { + "id": "msg_1", + "model": "claude-haiku", + "role": "assistant", + "type": "message", + "stop_reason": "end_turn", + "stop_sequence": None, + "content": [{"type": "text", "text": "hello"}], + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + types = [event_type for event_type, _ in decoded] + assert types == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ] + # message_start must not carry generated content itself, matching a real + # streaming response - it arrives via the content_block_delta that follows. + assert decoded[0][1]["message"]["content"] == [] + assert decoded[0][1]["message"]["id"] == "msg_1" + # Bugbot regression: message_start must not carry the completed response's + # final stop_reason/stop_sequence/output_tokens - a real stream keeps those + # null/zero until message_delta, so a client could otherwise treat the + # message as already finished, or double-count output tokens. + assert decoded[0][1]["message"]["stop_reason"] is None + assert decoded[0][1]["message"]["stop_sequence"] is None + assert decoded[0][1]["message"]["usage"] == {"input_tokens": 3, "output_tokens": 0} + assert decoded[1][1]["content_block"] == {"type": "text", "text": ""} + assert decoded[2][1]["delta"] == {"type": "text_delta", "text": "hello"} + assert decoded[4][1]["delta"]["stop_reason"] == "end_turn" + assert decoded[4][1]["usage"] == {"input_tokens": 3, "output_tokens": 2} + + +def test_anthropic_messages_response_as_sse_events_tool_use_block(): + response = { + "id": "msg_2", + "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "NYC"}}], + "stop_reason": "tool_use", + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + content_block_start = dict(decoded)["content_block_start"] + assert content_block_start["content_block"] == { + "type": "tool_use", + "id": "toolu_1", + "name": "get_weather", + "input": {}, + } + content_block_delta = dict(decoded)["content_block_delta"] + assert json.loads(content_block_delta["delta"]["partial_json"]) == {"city": "NYC"} + assert content_block_delta["delta"]["type"] == "input_json_delta" + + +def test_anthropic_messages_response_as_sse_events_thinking_block_emits_signature_delta(): + """Bugbot regression: a thinking block's real `signature` must reach the + client via a trailing signature_delta, not be silently dropped - Anthropic + rejects a replayed assistant message (a follow-up turn, a tool-use + continuation) whose thinking block lacks its original signature.""" + response = { + "id": "msg_5", + "content": [{"type": "thinking", "thinking": "let me think", "signature": "sig-abc123"}], + "stop_reason": "end_turn", + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + deltas = [payload["delta"] for event_type, payload in decoded if event_type == "content_block_delta"] + assert deltas == [ + {"type": "thinking_delta", "thinking": "let me think"}, + {"type": "signature_delta", "signature": "sig-abc123"}, + ] + + +def test_anthropic_messages_response_as_sse_events_thinking_block_without_signature_omits_delta(): + response = { + "id": "msg_6", + "content": [{"type": "thinking", "thinking": "let me think", "signature": None}], + "stop_reason": "end_turn", + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + deltas = [payload["delta"] for event_type, payload in decoded if event_type == "content_block_delta"] + assert deltas == [{"type": "thinking_delta", "thinking": "let me think"}] + + +def test_anthropic_messages_response_as_sse_events_multiple_blocks_are_indexed(): + response = { + "id": "msg_3", + "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}], + } + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + starts = [payload for event_type, payload in decoded if event_type == "content_block_start"] + assert [s["index"] for s in starts] == [0, 1] + deltas = [payload for event_type, payload in decoded if event_type == "content_block_delta"] + assert [d["delta"]["text"] for d in deltas] == ["a", "b"] + + +def test_anthropic_messages_response_as_sse_events_no_content_blocks(): + response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"} + decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) + assert [event_type for event_type, _ in decoded] == ["message_start", "message_delta", "message_stop"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 8b591fcd7da..aebbed88c70 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -144,9 +144,15 @@ class TestReasoningItemWithoutSummaryText: ("content_block_delta", 1), ("content_block_stop", 1), ] - assert chunks[1]["content_block"] == {"type": "thinking", "thinking": ""} + assert chunks[1]["content_block"] == {"type": "thinking", "thinking": "", "signature": ""} assert "".join(c["delta"]["thinking"] for c in chunks[2:4]) == "Weighing options" + def test_the_reasoning_item_id_is_never_streamed_as_a_signature(self): + """A stand-in signature would be replayed as a real one, so none is ever sent.""" + chunks = _drain_async(self._gpt_turn(reasoning_summary_deltas=["Weighing options"])) + + assert not [c for c in chunks if c.get("delta", {}).get("type") == "signature_delta"] + class TestToolUseBlockClosedExactlyOnce: """Regression for https://github.com/BerriAI/litellm/issues/37273. diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 964f4b9f68b..8225e7cff39 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -486,8 +486,8 @@ class TestTranslateMessagesToResponsesInput: } ] - def test_assistant_thinking_block_becomes_output_text(self): - """Assistant thinking block text is included as output_text.""" + def test_assistant_thinking_block_becomes_reasoning_item(self): + """Assistant thinking block becomes a reasoning item, never visible assistant prose.""" messages = [ { "role": "assistant", @@ -495,7 +495,77 @@ class TestTranslateMessagesToResponsesInput: } ] result = _translate_messages(messages) - assert result[0]["content"] == [{"type": "output_text", "text": "Let me reason step by step."}] + assert result == [ + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "Let me reason step by step."}], + } + ] + + def test_reasoning_item_carries_no_id(self): + """A fabricated reasoning id 404s upstream, so the item must go out without one.""" + messages = [ + { + "role": "assistant", + "content": [{"type": "thinking", "thinking": "Private reasoning.", "signature": "rs_abc123"}], + } + ] + result = _translate_messages(messages) + assert "id" not in result[0] + + def test_consecutive_thinking_blocks_become_one_reasoning_item(self): + """Summary parts of one upstream reasoning item are regrouped into that item.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "First part."}, + {"type": "thinking", "thinking": "Second part."}, + ], + } + ] + result = _translate_messages(messages) + assert result == [ + { + "type": "reasoning", + "summary": [ + {"type": "summary_text", "text": "First part."}, + {"type": "summary_text", "text": "Second part."}, + ], + } + ] + + def test_a_tool_call_splits_the_reasoning_items_around_it(self): + """Thinking on either side of a tool call belongs to two different reasoning items.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "Before the call."}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Denver"}}, + {"type": "thinking", "thinking": "After the call."}, + ], + } + ] + result = _translate_messages(messages) + assert [item["type"] for item in result] == ["reasoning", "function_call", "reasoning"] + assert result[0]["summary"] == [{"type": "summary_text", "text": "Before the call."}] + assert result[2]["summary"] == [{"type": "summary_text", "text": "After the call."}] + + def test_thinking_and_text_stay_separate(self): + """The visible answer stays the only thing in the assistant message.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "The user wants Denver."}, + {"type": "text", "text": "Denver is the best pick."}, + ], + } + ] + result = _translate_messages(messages) + assert [item["type"] for item in result] == ["reasoning", "message"] + assert result[1]["content"] == [{"type": "output_text", "text": "Denver is the best pick."}] def test_assistant_empty_thinking_block_skipped(self): """Assistant thinking block with empty thinking text is skipped.""" @@ -1094,7 +1164,7 @@ def _make_function_call_item(call_id: str, name: str, arguments: str) -> MagicMo return item -def _make_reasoning_item(summaries: List[str]) -> MagicMock: +def _make_reasoning_item(summaries: List[str], item_id: str = "rs_test_1") -> MagicMock: """Build a mock ResponseReasoningItem.""" from openai.types.responses import ResponseReasoningItem # type: ignore[import] @@ -1105,6 +1175,7 @@ def _make_reasoning_item(summaries: List[str]) -> MagicMock: summary_mocks.append(s) item = MagicMock(spec=ResponseReasoningItem) + item.id = item_id item.summary = summary_mocks return item @@ -1178,6 +1249,53 @@ class TestTranslateResponse: result: Any = _ADAPTER.translate_response(response) assert result["content"] == [] + def test_null_summary_text_skipped_rather_than_stringified(self): + """A summary part whose text is null must not reach the client as the word "None".""" + response = _make_mock_response( + output=[ + { + "type": "reasoning", + "id": "rs_null_1", + "summary": [{"type": "summary_text", "text": None}], + } + ] + ) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [] + + def test_reasoning_item_id_never_becomes_a_thinking_signature(self): + """Only Anthropic can sign a thinking block, so a stand-in signature is never invented.""" + reasoning = _make_reasoning_item(["Part one.", "Part two."], item_id="rs_abc123") + response = _make_mock_response(output=[reasoning]) + result: Any = _ADAPTER.translate_response(response) + assert [block["signature"] for block in result["content"]] == [None, None] + + def test_dict_reasoning_item_becomes_thinking_block(self): + """A reasoning item arriving as a plain dict is kept, not dropped.""" + response = _make_mock_response( + output=[ + { + "type": "reasoning", + "id": "rs_dict_1", + "summary": [{"type": "summary_text", "text": "Weighing the options."}], + } + ] + ) + result: Any = _ADAPTER.translate_response(response) + assert result["content"] == [ + {"type": "thinking", "thinking": "Weighing the options.", "signature": None} + ] + + def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self): + """Replaying this turn to an Anthropic model must not send a signature it cannot verify.""" + from litellm.litellm_core_utils.prompt_templates.factory import ( + _drop_unsignable_thinking_blocks, + ) + + response = _make_mock_response(output=[_make_reasoning_item(["Part one."], item_id="rs_abc123")]) + result: Any = _ADAPTER.translate_response(response) + assert _drop_unsignable_thinking_blocks(result["content"]) == [] + def test_usage_mapped_correctly(self): """Input/output tokens from ResponseAPIUsage are mapped to AnthropicUsage.""" response = _make_mock_response( diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py index ef74249ca8e..288817dff07 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py @@ -5,6 +5,8 @@ Verifies that reasoning_effort=None returns None for all models, including Claude Opus 4.6. """ +import pytest + from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -35,6 +37,16 @@ class TestMapReasoningEffort: ) assert result["type"] == "adaptive" + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_adaptive_mapping_requests_summarized_display(self, effort): + """Regression LIT-5714: adaptive thinking without ``display`` makes Anthropic + return a blank thinking block, so reasoning_effort callers always got + ``reasoning_content: ""``.""" + result = AnthropicConfig._map_reasoning_effort( + reasoning_effort=effort, model="claude-opus-4-6", custom_llm_provider="anthropic" + ) + assert result["display"] == "summarized" + def test_other_model_low_returns_enabled_with_budget(self): result = AnthropicConfig._map_reasoning_effort( reasoning_effort="low", model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py b/tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py index 71c9cfe8f41..32423d6679f 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_schema_filter.py @@ -417,3 +417,94 @@ class TestFilterAnthropicOutputSchema: result = AnthropicConfig.filter_anthropic_output_schema(schema) assert result["additionalProperties"] is False + + def test_drops_union_type_alongside_enum(self): + """A union ``type`` can never match a single declared type. + + Anthropic rejects it with "Invalid schema: Enum value 'low' does not + match declared type '['string', 'null']'". ``enum`` is the tighter + constraint, so the conflicting ``type`` is dropped. + """ + schema = { + "type": "object", + "properties": { + "confidence": { + "enum": ["low", "medium", "high", None], + "type": ["string", "null"], + } + }, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "type" not in result["properties"]["confidence"] + assert result["properties"]["confidence"]["enum"] == [ + "low", + "medium", + "high", + None, + ] + + def test_drops_type_when_an_enum_value_does_not_match_it(self): + """``enum: ["x", None]`` with ``type: "string"`` is rejected too.""" + schema = { + "type": "object", + "properties": {"a": {"enum": ["x", None], "type": "string"}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "type" not in result["properties"]["a"] + + def test_preserves_type_when_every_enum_value_matches(self): + """The non-conflicting case must be left exactly as-is.""" + schema = { + "type": "object", + "properties": {"a": {"enum": ["x", "y"], "type": "string"}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["properties"]["a"]["type"] == "string" + assert result["properties"]["a"]["enum"] == ["x", "y"] + + def test_integer_enum_satisfies_number_type(self): + """JSON Schema ``number`` accepts integers, so this is not a conflict.""" + schema = { + "type": "object", + "properties": {"a": {"enum": [1, 2], "type": "number"}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert result["properties"]["a"]["type"] == "number" + + def test_bool_enum_does_not_satisfy_integer_type(self): + """``bool`` is a Python ``int`` subclass but is not a JSON integer.""" + schema = { + "type": "object", + "properties": {"a": {"enum": [True], "type": "integer"}}, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "type" not in result["properties"]["a"] + + def test_normalizes_enum_type_inside_array_items(self): + """Normalization applies at every recursion site, not just top level.""" + schema = { + "type": "object", + "properties": { + "rows": { + "type": "array", + "items": { + "type": "object", + "properties": {"c": {"enum": ["a", None], "type": ["string", "null"]}}, + }, + } + }, + } + + result = AnthropicConfig.filter_anthropic_output_schema(schema) + + assert "type" not in result["properties"]["rows"]["items"]["properties"]["c"] diff --git a/tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json b/tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json new file mode 100644 index 00000000000..2dba691d038 --- /dev/null +++ b/tests/test_litellm/llms/azure/search/foundry_responses_web_search_fixture.json @@ -0,0 +1,78 @@ +{ + "id": "resp_04bcf916b629f9be006a8c7c053e9881959a4efe7118cb3a34", + "object": "response", + "status": "completed", + "model": "gpt-4.1", + "output": [ + { + "type": "web_search_call", + "id": "ws_04bcf916b629f9be006a8c7c080e108195b3366ab048231129", + "response_id": "resp_04bcf916b629f9be006a8c7c053e9881959a4efe7118cb3a34", + "status": "completed", + "action": { + "type": "search", + "queries": [ + "latest LiteLLM release GitHub LiteLLM", + "LiteLLM github releases", + "LiteLLM repo LiteLLM" + ], + "query": "latest LiteLLM release GitHub" + } + }, + { + "type": "message", + "id": "msg_04bcf916b629f9be006a8c7c0c7b448195979bfafb3b28f1ba", + "response_id": "resp_04bcf916b629f9be006a8c7c053e9881959a4efe7118cb3a34", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "You asked: “What is the latest LiteLLM release on GitHub? Answer with the version number.”\n\nI verified the most recent release information from the official GitHub repository for LiteLLM (BerriAI/litellm) using multiple sources, including GitHub itself and secondary trackers:\n\n- On the GitHub **Releases** page for BerriAI/litellm, the topmost entry is **v1.99.0‑rc.1**, marked as a pre‑release, published “yesterday” (relative to today, August 24, 2026) ([github.com](https://github.com/BerriAI/litellm/releases)). This indicates that version **v1.99.0‑rc.1** is the most recent tag available.\n- An external release‑tracking site (ReleaseAlert) confirms: **Latest version: v1.99.0‑rc.1**, last published August 22, 2026 ([releasealert.dev](https://releasealert.dev/github/BerriAI/litellm)).\n- The GitHub API (via `releases/latest`) currently points to **v1.98.0** as the latest **stable** release, with published date August 23, 2026 ([api.github.com](https://api.github.com/repos/BerriAI/litellm/releases/latest)).\n\nTo summarize:\n\n- The absolute **latest** release tag on GitHub is **v1.99.0‑rc.1** (release candidate), published recently (August 22, 2026) ([github.com](https://github.com/BerriAI/litellm/releases)).\n- The most recent **stable** release is **v1.98.0**, published August 23, 2026 ([api.github.com](https://api.github.com/repos/BerriAI/litellm/releases/latest)).\n\nSince you asked for the “latest LiteLLM release on GitHub,” without specifying stable vs. pre‑release, the correct answer is:\n\n**v1.99.0‑rc.1**\n\nLet me know if you'd like details on what's new in that release, or if you'd prefer the latest stable version.", + "annotations": [ + { + "type": "url_citation", + "url": "https://github.com/BerriAI/litellm/releases", + "start_index": 456, + "end_index": 515, + "title": "Releases · BerriAI/litellm - GitHub" + }, + { + "type": "url_citation", + "url": "https://releasealert.dev/github/BerriAI/litellm", + "start_index": 722, + "end_index": 791, + "title": "BerriAI/litellm on GitHub | Release Alert" + }, + { + "type": "url_citation", + "url": "https://api.github.com/repos/BerriAI/litellm/releases/latest", + "start_index": 936, + "end_index": 1016, + "title": "api.github.com" + }, + { + "type": "url_citation", + "url": "https://github.com/BerriAI/litellm/releases", + "start_index": 1160, + "end_index": 1219, + "title": "Releases · BerriAI/litellm - GitHub" + }, + { + "type": "url_citation", + "url": "https://api.github.com/repos/BerriAI/litellm/releases/latest", + "start_index": 1300, + "end_index": 1380, + "title": "api.github.com" + } + ], + "logprobs": [] + } + ], + "status": "completed" + } + ], + "usage": { + "input_tokens": 15195, + "output_tokens": 467 + } +} diff --git a/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py new file mode 100644 index 00000000000..fdc6f7bc239 --- /dev/null +++ b/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py @@ -0,0 +1,380 @@ +import json +from pathlib import Path +from unittest.mock import Mock + +import pytest + +from litellm.llms.azure.search.transformation import BingGroundingSearchConfig + +REAL_FIXTURE = json.loads((Path(__file__).parent / "foundry_responses_web_search_fixture.json").read_text()) + +RESPONSES_URL = "https://acct.services.ai.azure.com/api/projects/proj/openai/v1/responses" + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch: pytest.MonkeyPatch): + for var in ( + "BING_GROUNDING_PROJECT_ENDPOINT", + "BING_GROUNDING_MODEL", + "BING_GROUNDING_CONNECTION_ID", + "BING_GROUNDING_TOKEN", + ): + monkeypatch.delenv(var, raising=False) + + +def _config(entra_token_minter=None) -> BingGroundingSearchConfig: + return BingGroundingSearchConfig(entra_token_minter=entra_token_minter) + + +def _resp(payload, status_code: int = 200): + r = Mock() + r.status_code = status_code + r.headers = {} + r.content = (payload if isinstance(payload, str) else json.dumps(payload)).encode() + return r + + +def _message_response(text: str, annotations: list) -> dict: + return { + "output": [ + {"type": "web_search_call", "status": "completed"}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": annotations}], + }, + ] + } + + +def _citation(url: str, title: str, start: int, end: int) -> dict: + return {"type": "url_citation", "url": url, "title": title, "start_index": start, "end_index": end} + + +def test_ui_friendly_name(): + assert _config().ui_friendly_name() == "Grounding with Bing Search" + + +def test_validate_environment_api_key_uses_api_key_header_not_bearer(): + headers = _config().validate_environment({}, api_key="azure-api-key") + assert headers["api-key"] == "azure-api-key" + assert "Authorization" not in headers + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_reads_env_token(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_TOKEN", "env-token") + headers = _config().validate_environment({}) + assert headers["Authorization"] == "Bearer env-token" + assert "api-key" not in headers + + +def test_validate_environment_falls_back_to_entra_minter(): + headers = _config(entra_token_minter=lambda: "entra-token").validate_environment({}) + assert headers["Authorization"] == "Bearer entra-token" + + +def test_validate_environment_api_key_beats_env_token(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_TOKEN", "env-token") + minter = Mock(return_value="entra-token") + headers = _config(entra_token_minter=minter).validate_environment({}, api_key="azure-api-key") + assert headers["api-key"] == "azure-api-key" + assert "Authorization" not in headers + minter.assert_not_called() + + +def test_validate_environment_env_token_beats_entra_minter(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_TOKEN", "env-token") + minter = Mock(return_value="entra-token") + assert _config(entra_token_minter=minter).validate_environment({})["Authorization"] == "Bearer env-token" + minter.assert_not_called() + + +def test_validate_environment_refuses_entra_token_for_caller_api_base(): + minter = Mock(return_value="entra-token") + with pytest.raises(ValueError, match="Refusing to send the server-configured"): + _config(entra_token_minter=minter).validate_environment({}, api_base="https://attacker.example.com") + minter.assert_not_called() + + +def test_validate_environment_entra_minter_failure_names_the_options(): + def failing_minter() -> str: + raise RuntimeError("no az login") + + with pytest.raises(ValueError, match="no credential available") as excinfo: + _config(entra_token_minter=failing_minter).validate_environment({}) + message = str(excinfo.value) + assert "BING_GROUNDING_TOKEN" in message + assert "https://ai.azure.com/.default" in message + assert "no az login" in message + + +def test_validate_environment_does_not_mutate_and_is_idempotent(): + config = _config() + caller_headers = {"X-Custom": "keep-me"} + + once = config.validate_environment(caller_headers, api_key="k") + twice = config.validate_environment(once, api_key="k") + + assert caller_headers == {"X-Custom": "keep-me"} + assert once == twice + assert once["X-Custom"] == "keep-me" + + +def test_get_complete_url_from_api_base(): + url = _config().get_complete_url("https://acct.services.ai.azure.com/api/projects/proj", {}) + assert url == RESPONSES_URL + + +def test_get_complete_url_reads_env_endpoint(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_PROJECT_ENDPOINT", "https://acct.services.ai.azure.com/api/projects/proj/") + assert _config().get_complete_url(None, {}) == RESPONSES_URL + + +def test_get_complete_url_missing_endpoint_raises(): + with pytest.raises(ValueError, match="BING_GROUNDING_PROJECT_ENDPOINT"): + _config().get_complete_url(None, {}) + + +@pytest.mark.parametrize( + "api_base", + [ + "https://acct.services.ai.azure.com/api/projects/proj", + "https://acct.services.ai.azure.com/api/projects/proj/", + "https://acct.services.ai.azure.com/api/projects/proj/openai/v1/responses", + "https://acct.services.ai.azure.com/api/projects/proj/openai/v1/responses/", + ], +) +def test_get_complete_url_appends_responses_path_exactly_once(api_base: str): + assert _config().get_complete_url(api_base, {}) == RESPONSES_URL + + +def test_transform_search_request_missing_model_raises(): + with pytest.raises(ValueError, match="BING_GROUNDING_MODEL"): + _config().transform_search_request("q", {}) + + +def test_transform_search_request_web_search_mode_exact_body(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1") + body = _config().transform_search_request("latest AI developments", {"max_results": 5}) + assert body == { + "model": "gpt-4.1", + "input": "latest AI developments", + "tools": [{"type": "web_search"}], + } + + +def test_transform_search_request_web_search_mode_maps_country(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1") + body = _config().transform_search_request("q", {"country": "us"}) + assert body["tools"] == [{"type": "web_search", "user_location": {"type": "approximate", "country": "US"}}] + + +def test_transform_search_request_connection_mode_exact_body(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1") + monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id") + body = _config().transform_search_request("q", {"max_results": 5}) + assert body == { + "model": "gpt-4.1", + "input": "q", + "tools": [ + { + "type": "bing_grounding", + "bing_grounding": {"search_configurations": [{"project_connection_id": "conn-id", "count": 5}]}, + } + ], + } + + +def test_transform_search_request_connection_mode_omits_count_without_max_results( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1") + monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id") + body = _config().transform_search_request("q", {}) + assert body["tools"][0]["bing_grounding"]["search_configurations"] == [{"project_connection_id": "conn-id"}] + + +@pytest.mark.parametrize("max_results", [True, False, 0, -1]) +def test_transform_search_request_connection_mode_omits_count_for_invalid_max_results( + monkeypatch: pytest.MonkeyPatch, max_results: object +): + monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1") + monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id") + body = _config().transform_search_request("q", {"max_results": max_results}) + assert body["tools"][0]["bing_grounding"]["search_configurations"] == [{"project_connection_id": "conn-id"}] + + +def test_transform_search_response_ignores_invalid_max_results_cap(): + annotations = [_citation(f"https://example.com/{i}", f"T{i}", 0, 5) for i in range(3)] + resp = _config().transform_search_response( + _resp(_message_response("claim", annotations)), logging_obj=Mock(), optional_params={"max_results": True} + ) + assert [r.url for r in resp.results] == [f"https://example.com/{i}" for i in range(3)] + + +def test_transform_search_request_joins_list_query(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1") + assert _config().transform_search_request(["foo", "bar"], {})["input"] == "foo bar" + + +def test_transform_search_response_real_fixture_dedupes_and_preserves_order(): + resp = _config().transform_search_response(_resp(REAL_FIXTURE), logging_obj=Mock()) + + assert resp.object == "search" + assert [r.url for r in resp.results] == [ + "https://github.com/BerriAI/litellm/releases", + "https://releasealert.dev/github/BerriAI/litellm", + "https://api.github.com/repos/BerriAI/litellm/releases/latest", + ] + assert resp.results[0].title == "Releases · BerriAI/litellm - GitHub" + assert resp.results[1].title == "BerriAI/litellm on GitHub | Release Alert" + + +def test_transform_search_response_real_fixture_snippets_are_the_cited_claims(): + resp = _config().transform_search_response(_resp(REAL_FIXTURE), logging_obj=Mock()) + + assert resp.results[0].snippet.startswith("- On the GitHub **Releases** page for BerriAI/litellm") + assert resp.results[1].snippet.startswith("- An external release") + assert resp.results[2].snippet.startswith("- The GitHub API (via `releases/latest`)") + for result in resp.results: + assert "url_citation" not in result.snippet + assert not result.snippet.startswith("([") + + +def test_transform_search_response_snippet_falls_back_to_text_head_for_leading_citation(): + text = "([example.com](https://example.com)) trailing prose" + payload = _message_response(text, [_citation("https://example.com", "Example", 0, 36)]) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert resp.results[0].snippet == text + + +def test_transform_search_response_snippet_without_indices_uses_last_line(): + payload = _message_response( + "first line\nthe claim on the last line", + [{"type": "url_citation", "url": "https://example.com", "title": "Example"}], + ) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert resp.results[0].snippet == "the claim on the last line" + + +def test_transform_search_response_ignores_non_citation_annotations(): + payload = _message_response("text", [{"type": "file_citation", "url": "https://example.com"}]) + assert _config().transform_search_response(_resp(payload), logging_obj=Mock()).results == [] + + +def test_transform_search_response_ignores_citation_without_url(): + payload = _message_response("text", [{"type": "url_citation", "title": "no url"}]) + assert _config().transform_search_response(_resp(payload), logging_obj=Mock()).results == [] + + +def test_transform_search_response_no_message_output(): + payload = {"output": [{"type": "web_search_call", "status": "completed"}]} + assert _config().transform_search_response(_resp(payload), logging_obj=Mock()).results == [] + + +@pytest.mark.parametrize( + "body", + [ + "502 Bad Gateway", + '{"output": "garbage"}', + '{"output": null}', + "{}", + ], +) +def test_transform_search_response_malformed_body_raises_instead_of_reporting_empty(body: str): + with pytest.raises(Exception, match="Grounding with Bing Search"): + _config().transform_search_response(_resp(body, status_code=502), logging_obj=Mock()) + + +def test_transform_search_response_caps_results_to_max_results(): + annotations = [_citation(f"https://example.com/{i}", f"T{i}", 0, 5) for i in range(5)] + resp = _config().transform_search_response( + _resp(_message_response("claim", annotations)), logging_obj=Mock(), optional_params={"max_results": 2} + ) + assert [r.url for r in resp.results] == ["https://example.com/0", "https://example.com/1"] + + +def test_transform_search_response_without_max_results_returns_all_citations(): + annotations = [_citation(f"https://example.com/{i}", f"T{i}", 0, 5) for i in range(4)] + resp = _config().transform_search_response(_resp(_message_response("c", annotations)), logging_obj=Mock()) + assert len(resp.results) == 4 + + +def test_transform_search_response_failed_status_raises_with_error_message(): + payload = {"output": [], "status": "failed", "error": {"message": "content was filtered"}} + with pytest.raises(Exception, match="content was filtered") as excinfo: + _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert excinfo.value.status_code == 502 + + +def test_transform_search_response_incomplete_with_no_results_raises_with_reason(): + payload = {"output": [], "status": "incomplete", "incomplete_details": {"reason": "max_output_tokens"}} + with pytest.raises(Exception, match="incomplete: max_output_tokens"): + _config().transform_search_response(_resp(payload), logging_obj=Mock()) + + +def test_transform_search_response_incomplete_with_partial_results_returns_them(): + payload = _message_response("claim", [_citation("https://example.com", "Example", 0, 5)]) + payload["status"] = "incomplete" + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert [r.url for r in resp.results] == ["https://example.com"] + + +def test_transform_search_response_web_search_mode_zeroes_per_query_cost(): + payload = _message_response("claim", [_citation("https://example.com", "Example", 0, 5)]) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert resp._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.0 + + +def test_transform_search_response_connection_mode_leaves_price_to_cost_map(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id") + payload = _message_response("claim", [_citation("https://example.com", "Example", 0, 5)]) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert "additional_headers" not in resp._hidden_params + + +def test_get_error_class_attributes_the_provider(): + error = _config().get_error_class(error_message="quota exceeded", status_code=429, headers={}) + assert error.status_code == 429 + assert "Grounding with Bing Search: quota exceeded" in str(error) + assert "learn.microsoft.com" in str(error) + + +def test_get_error_class_unwraps_the_nested_tool_error(): + nested_tool_error = json.dumps( + { + "error": "Tool_User_Error", + "message": ( + "The specified connection ID 'conn-id' in tool config input was not found " + "in the project or account connections." + ), + "code": "invalid_tool_input", + "tool": "bing_grounding", + } + ) + live_400_shape = json.dumps( + { + "error": { + "message": nested_tool_error, + "type": "invalid_request_error", + "param": None, + "code": "tool_user_error", + } + } + ) + error = _config().get_error_class(error_message=live_400_shape, status_code=400, headers={}) + assert ( + "Grounding with Bing Search: The specified connection ID 'conn-id' in tool config input " + "was not found in the project or account connections" in str(error) + ) + assert "Tool_User_Error" not in str(error) + + +def test_get_error_class_unwraps_a_plain_error_envelope(): + error = _config().get_error_class( + error_message='{"error":{"message":"The api key is invalid.","code":"401"}}', + status_code=401, + headers={}, + ) + assert "Grounding with Bing Search: The api key is invalid" in str(error) diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index f2c852e9509..f000abb4c9a 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -7,7 +7,12 @@ from unittest.mock import MagicMock, patch import pytest import litellm -from litellm.llms.azure.common_utils import BaseAzureLLM, get_azure_ad_token +from litellm.llms.azure.common_utils import ( + BaseAzureLLM, + _cached_entra_id_token_provider, + get_azure_ad_token, + get_azure_ad_token_from_entra_id, +) from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, ) @@ -413,6 +418,7 @@ def test_select_azure_base_url_called(setup_mocks): "avector_store_create", "avector_store_search", "acreate_skill", + "acreate_interaction", ] ], ) @@ -2001,6 +2007,62 @@ def test_azure_traditional_api_uses_azure_openai_client(): ), f"Expected AsyncAzureOpenAI client for api_version={api_version}" +class TestEntraIdTokenProviderCache: + def setup_method(self): + _cached_entra_id_token_provider.cache_clear() + + def teardown_method(self): + _cached_entra_id_token_provider.cache_clear() + + def test_reuses_credential_for_the_same_service_principal(self): + with ( + patch("azure.identity.ClientSecretCredential") as mock_credential, + patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"), + ): + first = get_azure_ad_token_from_entra_id( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://cognitiveservices.azure.com/.default", + ) + second = get_azure_ad_token_from_entra_id( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://cognitiveservices.azure.com/.default", + ) + + assert first is second + assert mock_credential.call_count == 1 + + @pytest.mark.parametrize( + "second_call_kwargs", + [ + {"tenant_id": "other-tenant"}, + {"client_id": "other-client"}, + {"client_secret": "other-secret"}, + {"scope": "https://ai.azure.com/.default"}, + ], + ) + def test_does_not_share_a_provider_across_credentials_or_scopes(self, second_call_kwargs): + base_kwargs = { + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "scope": "https://cognitiveservices.azure.com/.default", + } + + with ( + patch("azure.identity.ClientSecretCredential") as mock_credential, + patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"), + ): + first = get_azure_ad_token_from_entra_id(**base_kwargs) + second = get_azure_ad_token_from_entra_id(**{**base_kwargs, **second_call_kwargs}) + + assert first is not second + assert mock_credential.call_count == 2 + + def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch): """`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client. diff --git a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index 0fd9a381a5a..11a727c9635 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -201,6 +201,90 @@ def test_azure_model_router_response_shows_actual_model(): ) +def test_azure_model_router_stamps_selected_model_on_hidden_params(): + """ + The selected model must be stamped on _hidden_params, not left for downstream code to + re-derive by looking for "model-router" in the model string. Deployments whose alias + does not contain that text are invisible to the string check. + """ + from httpx import Response + + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + AzureFoundryModelInfo, + ) + from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj + from litellm.types.utils import ModelResponse + + raw_response_json = { + "id": "chatcmpl-test456", + "object": "chat.completion", + "created": 1234567890, + "model": "grok-4-1-fast-reasoning", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "pong"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + mock_response = MagicMock(spec=Response) + mock_response.json.return_value = raw_response_json + mock_response.text = json.dumps(raw_response_json) + mock_response.headers = {} + + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + logging_obj.model_call_details = {} + + result = AzureModelRouterConfig().transform_response( + model="smart-pick", + raw_response=mock_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "Reply with just pong"}], + optional_params={}, + litellm_params={"model": "azure_ai/model_router/smart-pick"}, + encoding=None, + api_key="test-key", + json_mode=False, + ) + + assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model + assert ( + result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] + == "azure_ai/grok-4-1-fast-reasoning" + ) + assert AzureFoundryModelInfo.get_model_router_selected_model( + result._hidden_params + ) == ("azure_ai/grok-4-1-fast-reasoning") + assert ( + AzureFoundryModelInfo.is_model_router_call( + model="smart-pick", hidden_params=result._hidden_params + ) + is True + ) + + +def test_azure_model_router_stamp_does_not_leak_across_responses(): + """ + ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written + as a fresh dict. Mutating in place would bleed the selected model into unrelated responses. + """ + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + from litellm.types.utils import ModelResponse + + untouched = ModelResponse() + + assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {}) + + def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): """ Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name. @@ -300,6 +384,7 @@ def test_azure_ai_strips_non_openai_spec_message_fields(): "cache_control": {"type": "ephemeral"}, } ], + "reasoning_content": "The user wants me to read a file.", "provider_specific_fields": {"thought_signature": "sig-top"}, "tool_calls": [ { @@ -327,6 +412,7 @@ def test_azure_ai_strips_non_openai_spec_message_fields(): transformed_messages = request["messages"] assert not _find_key_anywhere(transformed_messages, "thinking_blocks") + assert not _find_key_anywhere(transformed_messages, "reasoning_content") assert not _find_key_anywhere(transformed_messages, "provider_specific_fields") assert not _find_key_anywhere(transformed_messages, "cache_control") diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 53a432427d3..326edde743d 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -341,7 +341,7 @@ def test_messages_thinking_shape_follows_exact_azure_entry_flag(local_model_cost ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index 667552dcf60..b6cb7ea9b54 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py @@ -1,5 +1,9 @@ +import litellm +from litellm.llms.azure_ai.image_edit.flux2_transformation import ( + AzureFoundryFlux2ImageEditConfig, +) from litellm.llms.azure_ai.image_edit.transformation import ( AzureFoundryFluxImageEditConfig, ) @@ -27,3 +31,32 @@ def test_azure_ai_url_generation(): ) expected_url = f"{api_base}/openai/deployments/FLUX.1-Kontext-pro/images/edits?api-version=2025-04-01-preview" assert complete_url == expected_url + + +def test_azure_ai_validate_environment_with_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + config = AzureFoundryFluxImageEditConfig() + + headers = config.validate_environment( + {}, + "FLUX.1-Kontext-pro", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers == {"Authorization": "Bearer entra-token"} + + +def test_flux2_validate_environment_with_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + config = AzureFoundryFlux2ImageEditConfig() + + headers = config.validate_environment( + {}, + "flux.2-pro", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers["Authorization"] == "Bearer entra-token" + assert headers["Content-Type"] == "application/json" diff --git a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py index b948e46093a..284a912d9a4 100644 --- a/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure_ai/image_edit/test_mai_image_edit_transformation.py @@ -5,6 +5,7 @@ import httpx import pytest +import litellm from litellm.llms.azure_ai.image_edit import ( AzureFoundryMAIImageEditConfig, get_azure_ai_image_edit_config, @@ -166,3 +167,16 @@ class TestAzureMAIImageEdit: assert image_response.data[0].b64_json == "abc123" assert image_response.usage.output_tokens == 1024 assert image_response.usage.total_tokens == 1024 + + +def test_mai_validate_environment_with_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "api_key", None) + + headers = AzureFoundryMAIImageEditConfig().validate_environment( + headers={}, + model="MAI-Image-2.5", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers == {"Authorization": "Bearer entra-token"} diff --git a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py index ab497d06ca7..91bf665f18d 100644 --- a/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/azure_ai/rerank/test_azure_ai_rerank_transformation.py @@ -2,6 +2,7 @@ import pytest +import litellm from litellm.llms.azure_ai.rerank.transformation import AzureAIRerankConfig @@ -92,3 +93,26 @@ class TestAzureAIRerankConfigGetCompleteUrl: model=self.model, ) assert url == "https://my-resource.services.ai.azure.com/v1/rerank?r=1" + + +class TestAzureAIRerankConfigValidateEnvironment: + def test_uses_api_key_when_set(self): + headers = AzureAIRerankConfig().validate_environment( + headers={}, + model="azure_ai/cohere-rerank-v3-english", + api_key="my-key", + ) + + assert headers["Authorization"] == "Bearer my-key" + + def test_falls_back_to_entra_token(self, monkeypatch): + monkeypatch.delenv("AZURE_AI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "azure_key", None) + + headers = AzureAIRerankConfig().validate_environment( + headers={}, + model="azure_ai/cohere-rerank-v3-english", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers["Authorization"] == "Bearer entra-token" diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py new file mode 100644 index 00000000000..c55bb2c3c36 --- /dev/null +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_entra_auth.py @@ -0,0 +1,154 @@ +""" +Entra ID / OAuth auth for Azure AI Foundry routes. + +Every azure_ai route must authenticate with an Entra ID token when no API key is configured, +instead of requiring an API key. +""" + +from unittest.mock import patch + +import pytest + +import litellm +from litellm.llms.azure_ai.common_utils import get_azure_ai_auth_headers +from litellm.llms.azure_ai.ocr.transformation import AzureAIOCRConfig + +ENTRA_PARAMS = {"azure_ad_token": "entra-token"} + + +@pytest.fixture(autouse=True) +def clear_azure_env(monkeypatch): + for env_var in ( + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_SCOPE", + "OPENAI_API_KEY", + "AZURE_DOCUMENT_INTELLIGENCE_API_KEY", + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + + +def test_api_key_wins_over_entra_credentials(): + headers = get_azure_ai_auth_headers(api_key="my-key", litellm_params=ENTRA_PARAMS, api_key_header="Api-Key") + + assert headers == {"Api-Key": "my-key"} + + +def test_entra_token_used_when_no_api_key(): + headers = get_azure_ai_auth_headers(api_key=None, litellm_params=ENTRA_PARAMS, api_key_header="Api-Key") + + assert headers == {"Authorization": "Bearer entra-token"} + + +def test_service_principal_token_is_requested_with_the_configured_scope(): + with patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id") as mock_entra_id: # test-quality-ok: stubs the Entra token fetch to assert the SP credential+scope plumbing and the returned Bearer header; live SP path proven by the PR's Azure Foundry e2e QA + mock_entra_id.return_value = lambda: "sp-token" + + headers = get_azure_ai_auth_headers( + api_key=None, + litellm_params={ + "tenant_id": "tenant", + "client_id": "client", + "client_secret": "secret", + "azure_scope": "https://ai.azure.com/.default", + }, + ) + + mock_entra_id.assert_called_once_with( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://ai.azure.com/.default", + ) + assert headers == {"Authorization": "Bearer sp-token"} + + +def test_error_mentions_both_credential_types_when_nothing_is_configured(): + with pytest.raises(ValueError, match="AZURE_AI_API_KEY") as exc_info: + get_azure_ai_auth_headers(api_key=None, litellm_params={}) + + message = str(exc_info.value) + assert "AZURE_AI_API_KEY" in message + assert "client_secret" in message + + +def test_ocr_authenticates_with_entra_token(): + headers = AzureAIOCRConfig().validate_environment( + headers={}, + model="azure_ai/mistral-ocr", + api_base="https://my-resource.services.ai.azure.com", + litellm_params=ENTRA_PARAMS, + ) + + assert headers["Authorization"] == "Bearer entra-token" + + +def test_embedding_falls_back_to_entra_token_instead_of_openai_key(monkeypatch): # test-quality-ok: asserts the embedding handler is authed with the Entra token, not the OpenAI key fallback; live path proven by the PR's Azure Foundry e2e QA + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-key") + + with patch.object(litellm.main.azure_ai_embedding, "embedding") as mock_embedding: # test-quality-ok: no injection seam for the embedding handler through the public embedding() API; live path proven by the PR's Azure Foundry e2e QA + mock_embedding.return_value = litellm.EmbeddingResponse() + + litellm.embedding( + model="azure_ai/cohere-embed-v3-english", + input=["hello"], + api_base="https://my-resource.services.ai.azure.com", + azure_ad_token="entra-token", + ) + + assert mock_embedding.call_args.kwargs["api_key"] == "entra-token" + + +def test_image_generation_authenticates_with_entra_token(): + with patch.object(litellm.images.main.azure_chat_completions, "image_generation") as mock_image_generation: # test-quality-ok: asserts image_generation forwards the computed Entra bearer header; no injection seam through the public API; live path proven by the PR's Azure Foundry e2e QA + mock_image_generation.return_value = litellm.ImageResponse() + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + prompt="a red circle", + api_base="https://my-resource.services.ai.azure.com", + azure_ad_token="entra-token", + ) + + headers = mock_image_generation.call_args.kwargs["headers"] + assert headers["Authorization"] == "Bearer entra-token" + assert "api-key" not in headers + + +@pytest.mark.parametrize("header_name", ["Authorization", "authorization", "api-key", "API-KEY"]) +def test_image_generation_keeps_caller_supplied_auth_header(header_name): + with patch.object(litellm.images.main.azure_chat_completions, "image_generation") as mock_image_generation: # test-quality-ok: asserts a caller-supplied auth header is preserved over Entra; no injection seam through the public API; live path proven by the PR's Azure Foundry e2e QA + mock_image_generation.return_value = litellm.ImageResponse() + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + prompt="a red circle", + api_base="https://my-resource.services.ai.azure.com", + headers={header_name: "caller-credential"}, + ) + + headers = mock_image_generation.call_args.kwargs["headers"] + assert headers[header_name] == "caller-credential" + assert len(headers) == 2 + + +def test_image_generation_still_uses_api_key_header(): + with patch.object(litellm.images.main.azure_chat_completions, "image_generation") as mock_image_generation: # test-quality-ok: asserts the api-key header path still works alongside Entra; no injection seam through the public API; live path proven by the PR's Azure Foundry e2e QA + mock_image_generation.return_value = litellm.ImageResponse() + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + prompt="a red circle", + api_base="https://my-resource.services.ai.azure.com", + api_key="my-key", + ) + + headers = mock_image_generation.call_args.kwargs["headers"] + assert headers["api-key"] == "my-key" + assert "Authorization" not in headers diff --git a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py index 66f4f432eb8..ef4c78553f1 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_document_intelligence_ocr_transformation.py @@ -344,3 +344,30 @@ def test_get_complete_url_combines_pages_and_features(): assert "&pages=1,2,3" in url assert "&features=keyValuePairs,languages" in url + + +def test_validate_environment_uses_subscription_key(monkeypatch): + monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_API_KEY", raising=False) + + headers = AzureDocumentIntelligenceOCRConfig().validate_environment( + headers={}, + model="prebuilt-layout", + api_key="my-key", + api_base="https://example.cognitiveservices.azure.com", + ) + + assert headers["Ocp-Apim-Subscription-Key"] == "my-key" + + +def test_validate_environment_falls_back_to_entra_token(monkeypatch): + monkeypatch.delenv("AZURE_DOCUMENT_INTELLIGENCE_API_KEY", raising=False) + + headers = AzureDocumentIntelligenceOCRConfig().validate_environment( + headers={}, + model="prebuilt-layout", + api_base="https://example.cognitiveservices.azure.com", + litellm_params={"azure_ad_token": "entra-token"}, + ) + + assert headers["Authorization"] == "Bearer entra-token" + assert "Ocp-Apim-Subscription-Key" not in headers diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py index e4402bbec49..e6aad7688d1 100644 --- a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -16,6 +16,7 @@ import pytest import litellm from litellm.llms.apiserpent.search.transformation import APISerpentSearchConfig +from litellm.llms.azure.search.transformation import BingGroundingSearchConfig from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, _is_trusted_search_api_base, @@ -59,6 +60,7 @@ _BASE_ENV_VARS = ( "TINYFISH_API_BASE", "CRW_API_BASE", "NIMBLE_API_BASE", + "BING_GROUNDING_PROJECT_ENDPOINT", ) @@ -99,6 +101,7 @@ PROVIDERS: Tuple[ProviderSpec, ...] = ( (TinyfishSearchConfig, {"TINYFISH_API_KEY": "srv"}, "caller-key", {}), (FastCRWSearchConfig, {"CRW_API_KEY": "srv"}, "caller-key", {}), (NimbleSearchConfig, {"NIMBLE_API_KEY": "srv"}, "caller-key", {}), + (BingGroundingSearchConfig, {"BING_GROUNDING_TOKEN": "srv"}, "caller-key", {}), ) _IDS = tuple(spec[0].__name__ for spec in PROVIDERS) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 604f3414775..4d2c077b548 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -366,6 +366,96 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +def test_reasoning_effort_requests_summarized_display_converse(): + """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must + request the summarized display, otherwise the provider returns a blank thinking + block and reasoning_content is always empty.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high"}, + optional_params={}, + model="bedrock/converse/us.anthropic.claude-opus-4-7", + drop_params=False, + ) + + assert optional_params["thinking"]["type"] == "adaptive" + assert optional_params["thinking"]["display"] == "summarized" + + +def test_thinking_request_adds_output_tokens_details_response_path(): + """Regression LIT-5714: the Converse usage block has no thinking-token field, so + thinking requests must ask for ``/usage/output_tokens_details`` via + ``additionalModelResponseFieldPaths``.""" + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "high"}, + }, + litellm_params={}, + headers={}, + ) + + assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",) + + +def test_request_without_thinking_omits_response_field_paths(): + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={"maxTokens": 256}, + litellm_params={}, + headers={}, + ) + + assert "additionalModelResponseFieldPaths" not in result + + +def test_transform_usage_prefers_provider_reasoning_tokens(): + """Regression LIT-5714: provider-reported thinking tokens must win over the + token_counter estimate derived from visible reasoning text.""" + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + provider_reasoning_tokens=1033, + ) + + assert usage.completion_tokens_details.reasoning_tokens == 1033 + assert usage.completion_tokens_details.text_tokens == 3002 - 1033 + + +def test_transform_usage_falls_back_to_estimate_without_provider_tokens(): + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 300, "totalTokens": 340}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + ) + + assert usage.completion_tokens_details.reasoning_tokens > 0 + assert usage.completion_tokens_details.reasoning_tokens < 300 + + +def test_thinking_tokens_parsed_from_additional_model_response_fields(): + parsed = AmazonConverseConfig.thinking_tokens_from_additional_fields( + {"usage": {"output_tokens_details": {"thinking_tokens": 92}}} + ) + assert parsed == 92 + assert AmazonConverseConfig.thinking_tokens_from_additional_fields(None) is None + assert AmazonConverseConfig.thinking_tokens_from_additional_fields({"usage": {}}) is None + + @pytest.mark.parametrize( "model,effort,expected_effort", [ diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index e2892a6ccee..4bef59842f1 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -1,10 +1,12 @@ +import datetime from unittest.mock import AsyncMock, MagicMock import httpx import pytest - import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AWSEventStreamDecoder, make_call, @@ -206,6 +208,29 @@ def test_bedrock_converse_streaming_consistent_id(): ), "All chunk IDs must match the one captured from the messageStart event" +def test_converse_streaming_usage_uses_provider_thinking_tokens(): + """Regression LIT-5714: the messageStop event carries provider thinking tokens + under ``additionalModelResponseFields``; the usage chunk must report them instead + of a token_counter estimate.""" + chunks = [ + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"text": "thinking about it"}}, + }, + { + "stopReason": "end_turn", + "additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}}, + }, + {"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}}, + ] + + decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7") + parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks] + + usage = parsed[-1].usage + assert usage.completion_tokens_details.reasoning_tokens == 1033 + + @pytest.mark.asyncio async def test_make_call_does_not_rechunk_stream_by_default(): """Re-chunking the event stream into fixed 1024-byte blocks holds small @@ -292,6 +317,139 @@ def test_make_sync_call_honors_explicit_stream_chunk_size(): response.iter_bytes.assert_called_once_with(chunk_size=2048) +CONVERSE_MODEL = "anthropic.claude-sonnet-4-6" +CONVERSE_METADATA_EVENT = { + "usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, + "metrics": {"latencyMs": 100}, +} + + +def _converse_stream_wrapper(events): + async def bedrock_stream(): + decoder = AWSEventStreamDecoder(model=CONVERSE_MODEL) + for event in events: + yield decoder._chunk_parser(chunk_data=event) + + return CustomStreamWrapper( + completion_stream=bedrock_stream(), + model=CONVERSE_MODEL, + custom_llm_provider="bedrock", + logging_obj=LiteLLMLoggingObj( + model=CONVERSE_MODEL, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="1234", + function_id="1234", + ), + ) + + +@pytest.mark.parametrize( + "events, expected_finish_reason", + [ + pytest.param( + ( + {"role": "assistant"}, + {"contentBlockIndex": 0, "delta": {"text": "Hello"}}, + {"contentBlockIndex": 0, "delta": {"text": " world"}}, + {"contentBlockIndex": 0}, + {"stopReason": "end_turn"}, + CONVERSE_METADATA_EVENT, + ), + "stop", + id="text", + ), + pytest.param( + ( + {"role": "assistant"}, + {"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "t1", "name": "get_weather"}}}, + {"contentBlockIndex": 0, "delta": {"toolUse": {"input": '{"city": "SF"}'}}}, + {"contentBlockIndex": 0}, + {"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "t2", "name": "get_time"}}}, + {"contentBlockIndex": 1, "delta": {"toolUse": {"input": '{"tz": "PT"}'}}}, + {"contentBlockIndex": 1}, + {"stopReason": "tool_use"}, + CONVERSE_METADATA_EVENT, + ), + "tool_calls", + id="multiple_tool_calls", + ), + pytest.param( + ( + {"role": "assistant"}, + {"contentBlockIndex": 0, "start": {}}, + {"contentBlockIndex": 0, "delta": {"text": "Let me check."}}, + {"contentBlockIndex": 0}, + {"contentBlockIndex": 1, "start": {"toolUse": {"toolUseId": "t1", "name": "get_weather"}}}, + {"contentBlockIndex": 1, "delta": {"toolUse": {"input": '{"city": "SF"}'}}}, + {"contentBlockIndex": 1}, + {"stopReason": "tool_use"}, + CONVERSE_METADATA_EVENT, + ), + "tool_calls", + id="text_then_tool_call", + ), + pytest.param( + ( + {"role": "assistant"}, + {"contentBlockIndex": 0, "start": {}}, + {"contentBlockIndex": 0, "delta": {"reasoningContent": {"text": "thinking hard"}}}, + {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "sig123"}}}, + {"contentBlockIndex": 0}, + {"contentBlockIndex": 1, "start": {}}, + {"contentBlockIndex": 1, "delta": {"text": "Answer"}}, + {"contentBlockIndex": 1}, + {"stopReason": "end_turn"}, + CONVERSE_METADATA_EVENT, + ), + "stop", + id="reasoning_then_text", + ), + ], +) +@pytest.mark.asyncio +async def test_converse_stream_ends_on_finish_reason_chunk(events, expected_finish_reason): + """The usage-only metadata event Bedrock sends after messageStop must not reach the caller as an extra + assistant delta following the finish_reason chunk.""" + wrapper = _converse_stream_wrapper(events) + + chunks = [chunk async for chunk in wrapper] + + finish_reasons = [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] + assert finish_reasons == [expected_finish_reason] + assert chunks[-1].choices[0].finish_reason == expected_finish_reason, ( + f"stream must end on the finish_reason chunk, got trailing {chunks[-1].model_dump(exclude_none=True)}" + ) + roles = [choice.delta.role for chunk in chunks for choice in chunk.choices if choice.delta.role] + assert roles == ["assistant"] + assert any(getattr(chunk, "usage", None) is not None for chunk in wrapper.chunks) + + +@pytest.mark.asyncio +async def test_converse_stream_still_emits_guardrail_trace_after_finish_reason(): + """Guardrail metadata events carry a trace payload alongside usage; that chunk must still reach the caller + after the finish_reason chunk, as it did before the regression.""" + trace = {"guardrail": {"inputAssessment": {"g1": {}}}} + events = ( + {"role": "assistant"}, + {"contentBlockIndex": 0, "delta": {"text": "Hello"}}, + {"contentBlockIndex": 0}, + {"stopReason": "end_turn"}, + {**CONVERSE_METADATA_EVENT, "trace": trace}, + ) + wrapper = _converse_stream_wrapper(events) + + chunks = [chunk async for chunk in wrapper] + + finish_reasons = [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] + assert finish_reasons == ["stop"] + assert chunks[-1].provider_specific_fields == {"trace": trace} + assert chunks[-1].choices[0].delta.content == "" + assert chunks[-1].choices[0].delta.role == "assistant" + + def test_invoke_streaming_forwards_bedrock_response_headers(): response = MagicMock() response.status_code = 200 diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index d3c28302bf9..1e09afd6919 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1387,7 +1387,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} @@ -2935,7 +2935,7 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem(litellm.model_cost[model], "supports_adaptive_thinking", False) diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py new file mode 100644 index 00000000000..8c6eda605ca --- /dev/null +++ b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py @@ -0,0 +1,197 @@ +import json +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from botocore.credentials import Credentials + +from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.llms.bedrock_mantle.passthrough.transformation import BedrockMantlePassthroughConfig +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.passthrough.main import llm_passthrough_route +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws" +INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke" +CONVERSE_ENDPOINT = "model/us.openai.gpt-5.6-sol/converse" +REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64} + + +@pytest.fixture +def no_ambient_aws(monkeypatch): + for name in ( + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_KEY", + "BEDROCK_MANTLE_API_BASE", + "BEDROCK_MANTLE_REGION", + "AWS_BEDROCK_RUNTIME_ENDPOINT", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", + ): + monkeypatch.delenv(name, raising=False) + + +def test_bedrock_mantle_registers_its_own_bedrock_passthrough_config(): + config = ProviderConfigManager.get_provider_passthrough_config( + model="us.openai.gpt-5.6-sol", provider=LlmProviders.BEDROCK_MANTLE + ) + assert isinstance(config, BedrockMantlePassthroughConfig) + assert isinstance(config, BedrockPassthroughConfig) + + +def test_mantle_api_base_only_lends_its_region_to_the_runtime_url(no_ambient_aws): + url, base_url = BedrockMantlePassthroughConfig().get_complete_url( + api_base=MANTLE_API_BASE, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={"api_base": MANTLE_API_BASE}, + ) + assert str(url) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}" + assert base_url == "https://bedrock-runtime.us-east-2.amazonaws.com" + + +def test_explicit_region_and_non_mantle_api_base_are_kept(no_ambient_aws): + vpc_endpoint = "https://vpce-0123.bedrock-runtime.us-east-1.vpce.amazonaws.com" + url, base_url = BedrockMantlePassthroughConfig().get_complete_url( + api_base=vpc_endpoint, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={"api_base": vpc_endpoint, "aws_region_name": "us-east-1"}, + ) + assert str(url) == f"{vpc_endpoint}/{INVOKE_ENDPOINT}" + assert base_url == vpc_endpoint + + +def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws): + url, _ = BedrockMantlePassthroughConfig().get_complete_url( + api_base=None, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={}, + ) + assert str(url) == f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}" + + +@pytest.mark.parametrize( + ("litellm_params", "env", "expected_bearer"), + [ + ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"), + ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"), + ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"), + ], +) +def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer): + for name, value in env.items(): + monkeypatch.setenv(name, value) + headers, body = BedrockMantlePassthroughConfig().sign_request( + headers={}, + litellm_params=litellm_params, + request_data=REQUEST_BODY, + api_base=f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}", + model="us.openai.gpt-5.6-sol", + ) + assert headers["Authorization"] == f"Bearer {expected_bearer}" + assert body is not None + assert json.loads(body) == REQUEST_BODY + + +def test_sign_request_falls_back_to_sigv4_scoped_to_the_mantle_region(no_ambient_aws): + config = BedrockMantlePassthroughConfig() + with patch.object(config, "get_credentials", return_value=Credentials("AKIA", "secret")): + headers, body = config.sign_request( + headers={}, + litellm_params={"api_base": MANTLE_API_BASE}, + request_data=REQUEST_BODY, + api_base=f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}", + model="us.openai.gpt-5.6-sol", + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIA/") + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert body is not None + assert json.loads(body) == REQUEST_BODY + + +@pytest.mark.parametrize( + ("route_kwargs", "env", "expected_bearer"), + [ + ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"), + ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"), + ], +) +def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deployment( + no_ambient_aws, monkeypatch, route_kwargs, env, expected_bearer +): + for name, value in env.items(): + monkeypatch.setenv(name, value) + client = HTTPHandler() + with ( + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), + patch.object(client.client, "build_request", wraps=client.client.build_request) as build_request, + ): + response = llm_passthrough_route( + model="bedrock_mantle/us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + method="POST", + api_base=MANTLE_API_BASE, + json=dict(REQUEST_BODY), + client=client, + litellm_logging_obj=MagicMock(), + **route_kwargs, + ) + assert response.status_code == 200 + sent = build_request.call_args.kwargs + assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}" + assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}" + assert json.loads(sent["content"]) == REQUEST_BODY + + +def _logged_model_response(endpoint, body): + request = httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com/{endpoint}") + return BedrockMantlePassthroughConfig().logging_non_streaming_response( + model="us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock_mantle", + httpx_response=httpx.Response(200, json=body, request=request), + request_data={"messages": [{"role": "user", "content": [{"text": "say pong"}]}]}, + logging_obj=MagicMock(), + endpoint=endpoint, + ) + + +def test_converse_logging_parses_the_converse_response_shape(): + result = _logged_model_response( + CONVERSE_ENDPOINT, + { + "metrics": {"latencyMs": 800.0}, + "output": {"message": {"content": [{"text": "pong"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 8, "outputTokens": 5, "totalTokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5 + + +def test_invoke_logging_parses_the_openai_chat_response_shape(): + result = _logged_model_response( + INVOKE_ENDPOINT, + { + "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "pong", "role": "assistant"}}], + "created": 1787677792, + "id": "chatcmpl-regression", + "model": "us.openai.gpt-5.6-sol", + "object": "chat.completion", + "usage": {"completion_tokens": 5, "prompt_tokens": 8, "total_tokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5 diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 9e05d48a18f..00a319f99f0 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -8,7 +8,9 @@ gate, the URL construction for both paths, and the shared Bearer auth. """ import copy - +import json +import logging +from pathlib import Path import pytest from botocore.exceptions import ( @@ -623,6 +625,181 @@ class TestBedrockMantleCodexAdditionalTools: assert "additional_tools" in str(mock_debug.call_args) +class TestBedrockMantleCodexInputItemNormalization: + """Mantle 400s ("Invalid 'input': value did not match any expected variant") + on the Codex history item types agent_message, context_compaction, and + local_shell_call (verified against bedrock-mantle.us-east-1.api.aws with + openai.gpt-5.6-sol), so the config must rewrite them into supported + equivalents. agent_message is what every Codex multi-agent v2 session sends, + and its encrypted_content slot carries the verbatim plaintext payload when + the upstream model never issued encrypted args, so that slot must be + preserved, not dropped. Mantle also rejects assistant messages with + input_text content, so the rewrite must use output_text.""" + + _USER_MESSAGE = { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Continue."}], + } + + def _transform(self, input): + cfg = BedrockMantleResponsesAPIConfig() + return cfg.transform_responses_api_request( + model="openai.gpt-5.6-sol", + input=input, + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + def test_plaintext_agent_message_becomes_assistant_output_text_message(self): + body = self._transform( + input=[ + self._USER_MESSAGE, + { + "type": "agent_message", + "id": "amsg_1", + "author": "/root/arithmetic", + "recipient": "/root", + "content": [{"type": "input_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."}], + }, + ] + ) + assert body["input"] == [ + self._USER_MESSAGE, + { + "type": "message", + "role": "assistant", + "content": ({"type": "output_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."},), + }, + ] + + def test_agent_message_encrypted_content_payload_is_preserved(self): + body = self._transform( + input=[ + { + "type": "agent_message", + "author": "/root", + "recipient": "/root/arithmetic", + "content": [ + {"type": "input_text", "text": "Message Type: NEW_TASK\nPayload:\n"}, + {"type": "encrypted_content", "encrypted_content": "Answer the question 'what is 2+2'."}, + ], + }, + self._USER_MESSAGE, + ] + ) + assert body["input"][0] == { + "type": "message", + "role": "assistant", + "content": ( + { + "type": "output_text", + "text": "Message Type: NEW_TASK\nPayload:\nAnswer the question 'what is 2+2'.", + }, + ), + } + + def test_agent_message_without_any_text_is_dropped(self): + body = self._transform( + input=[ + {"type": "agent_message", "author": "/root", "recipient": "/root/a", "content": []}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_context_compaction_becomes_compaction_with_same_ciphertext(self): + body = self._transform( + input=[ + {"type": "context_compaction", "id": "cc_1", "encrypted_content": "smry_abc123"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [ + {"type": "compaction", "encrypted_content": "smry_abc123"}, + self._USER_MESSAGE, + ] + + def test_context_compaction_without_ciphertext_is_dropped(self): + body = self._transform( + input=[ + {"type": "context_compaction", "id": "cc_1"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_local_shell_call_becomes_function_call_keeping_call_id_pairing(self): + body = self._transform( + input=[ + { + "type": "local_shell_call", + "id": "lsh_1", + "call_id": "call_1", + "status": "completed", + "action": {"type": "exec", "command": ["echo", "hi"]}, + }, + {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [ + { + "type": "function_call", + "call_id": "call_1", + "name": "local_shell", + "arguments": '{"type": "exec", "command": ["echo", "hi"]}', + }, + {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"}, + self._USER_MESSAGE, + ] + + def test_local_shell_call_without_call_id_is_dropped(self): + body = self._transform( + input=[ + {"type": "local_shell_call", "status": "completed", "action": {"type": "exec", "command": ["ls"]}}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_mantle_supported_item_types_pass_through_untouched(self): + supported_items = [ + self._USER_MESSAGE, + {"type": "compaction", "encrypted_content": "smry_abc123"}, + {"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "call_2"}, + {"type": "function_call_output", "call_id": "call_2", "output": "ok"}, + {"type": "tool_search_call", "call_id": "call_3", "execution": "server", "arguments": {"query": "x"}}, + {"type": "tool_search_output", "call_id": "call_3", "status": "completed", "execution": "server", "tools": []}, + {"type": "compaction_trigger"}, + ] + body = self._transform(input=copy.deepcopy(supported_items)) + assert body["input"] == supported_items + + def test_string_input_passes_through(self): + body = self._transform(input="Say hi.") + assert body["input"] == "Say hi." + + def test_rewrite_is_logged_as_warning_naming_the_types(self, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body = self._transform( + input=[ + {"type": "agent_message", "author": "a", "recipient": "b", "content": [{"type": "input_text", "text": "hi"}]}, + self._USER_MESSAGE, + ] + ) + assert body["input"][0]["role"] == "assistant" + rewrite_warnings = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "rewrote Codex input item type" in record.getMessage() + ] + assert rewrite_warnings == [ + "Bedrock Mantle Responses API: rewrote Codex input item type(s) ['agent_message'] that Mantle rejects." + ] + + class TestBedrockMantleResponsesRegistry: def test_registry_returns_config_for_gpt_5_5(self, local_cost_map): # gpt-5.x advertises /v1/responses in supported_endpoints (capability) @@ -1523,7 +1700,7 @@ class TestBedrockMantleResponsesPricing: assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost) assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost) assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == 1000000 + assert info["max_input_tokens"] == 1050000 assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2) assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2) @@ -1565,3 +1742,42 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models + + +def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]: + repo_root = Path(__file__).resolve().parents[4] + paths = { + "root": repo_root / "model_prices_and_context_window.json", + "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json", + } + return json.loads(paths[map_name].read_text()) + + +class TestGpt56MantleRegistryEntries: + """Locks the gpt-5.6 frontier entries to Bedrock Mantle's live behavior. + + Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna + (oversize requests 400 with "prompt tokens (N) exceed model maximum + (1050000)", and a 1,030,590-token request completes), matching the OpenAI + Bedrock guide. mode must stay "responses": Mantle's native + /v1/chat/completions rejects function tools unless reasoning_effort is + "none", so chat traffic has to keep bridging to the Responses API + (see the responses_api_bridge tests above). + """ + + @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) + @pytest.mark.parametrize( + "key", + ( + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + ), + ) + def test_entry_matches_mantle_enforced_limits(self, map_name, key): + entry = _repo_cost_map(map_name)[key] + assert entry["max_input_tokens"] == 1050000 + assert entry["max_output_tokens"] == 128000 + assert entry["mode"] == "responses" + assert entry["use_openai_responses_path"] is True + assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] diff --git a/tests/test_litellm/llms/cerebras/__init__.py b/tests/test_litellm/llms/cerebras/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py new file mode 100644 index 00000000000..09718b1e6e0 --- /dev/null +++ b/tests/test_litellm/llms/cerebras/test_cerebras_chat_transformation.py @@ -0,0 +1,61 @@ +from litellm.llms.cerebras.chat import CerebrasConfig + + +def test_max_retries_in_supported_params() -> None: + config = CerebrasConfig() + params = config.get_supported_openai_params(model="llama-3.3-70b") + assert "max_retries" in params, ( + f"max_retries must be in CerebrasConfig.get_supported_openai_params(); got: {params!r}" + ) + + +def test_extra_headers_in_supported_params() -> None: + config = CerebrasConfig() + params = config.get_supported_openai_params(model="llama-3.3-70b") + assert "extra_headers" in params, ( + f"extra_headers must be in CerebrasConfig.get_supported_openai_params(); got: {params!r}" + ) + + +def test_core_openai_params_still_supported() -> None: + config = CerebrasConfig() + params = config.get_supported_openai_params(model="llama-3.3-70b") + for expected in ( + "max_tokens", + "max_completion_tokens", + "response_format", + "seed", + "stop", + "stream", + "temperature", + "top_p", + "tool_choice", + "tools", + "user", + ): + assert expected in params, f"{expected!r} unexpectedly missing from Cerebras supported params: {params!r}" + + +def test_map_openai_params_preserves_max_retries() -> None: + config = CerebrasConfig() + result = config.map_openai_params( + non_default_params={"max_retries": 0, "temperature": 0.7}, + optional_params={}, + model="llama-3.3-70b", + drop_params=False, + ) + assert result.get("max_retries") == 0, f"map_openai_params must preserve max_retries=0; got: {result!r}" + assert result.get("temperature") == 0.7 + + +def test_map_openai_params_preserves_max_retries_zero_falsy() -> None: + config = CerebrasConfig() + result = config.map_openai_params( + non_default_params={"max_retries": 0}, + optional_params={}, + model="llama-3.3-70b", + drop_params=False, + ) + assert "max_retries" in result and result["max_retries"] == 0, ( + f"max_retries=0 (falsy) must not be silently omitted; got: {result!r}" + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index f7f89cd1d8d..16d57437043 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -56,9 +56,7 @@ async def test_async_post_streaming_status_error_should_not_wait_forever_for_bod litellm_handler = AsyncHTTPHandler() await litellm_handler.client.aclose() - litellm_handler.client = httpx.AsyncClient( - transport=httpx.MockTransport(mock_handler) - ) + litellm_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler)) try: with pytest.raises(MaskedHTTPStatusError) as exc_info: await asyncio.wait_for( @@ -202,9 +200,7 @@ async def test_ssl_verification_with_aiohttp_transport(monkeypatch: pytest.Monke transport_connector = transport._get_valid_client_session().connector assert isinstance(transport_connector, TCPConnector) - aiohttp_session = aiohttp.ClientSession( - connector=aiohttp.TCPConnector(ssl=False) - ) + aiohttp_session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) try: aiohttp_connector = aiohttp_session.connector assert isinstance(aiohttp_connector, aiohttp.TCPConnector) @@ -378,7 +374,8 @@ async def test_get_async_httpx_client_with_shared_session(): # Test with shared session client = get_async_httpx_client( - llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session # type: ignore + llm_provider=LlmProviders.ANTHROPIC, + shared_session=mock_session, # type: ignore ) # Verify the client was created successfully @@ -397,9 +394,7 @@ async def test_get_async_httpx_client_without_shared_session(): from litellm.types.utils import LlmProviders # Test without shared session - client = get_async_httpx_client( - llm_provider=LlmProviders.ANTHROPIC, shared_session=None - ) + client = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC, shared_session=None) # Verify the client was created successfully assert client is not None @@ -476,11 +471,13 @@ async def test_session_reuse_integration(): # Create two clients with the same session client1 = get_async_httpx_client( - llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session # type: ignore + llm_provider=LlmProviders.ANTHROPIC, + shared_session=mock_session, # type: ignore ) client2 = get_async_httpx_client( - llm_provider=LlmProviders.OPENAI, shared_session=mock_session # type: ignore + llm_provider=LlmProviders.OPENAI, + shared_session=mock_session, # type: ignore ) # Both clients should be created successfully @@ -512,9 +509,7 @@ async def test_session_reuse_integration(): (None, None, None, False), # None value - skip configuration ], ) -def test_ssl_ecdh_curve( - env_curve, litellm_curve, expected_curve, should_call, monkeypatch -): +def test_ssl_ecdh_curve(env_curve, litellm_curve, expected_curve, should_call, monkeypatch): """Test SSL ECDH curve configuration with valid curves and precedence""" from litellm.llms.custom_httpx.http_handler import _ssl_context_cache @@ -717,9 +712,7 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout: _default_cached_client_timeout, ) - monkeypatch.setattr( - litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS - ) + monkeypatch.setattr(litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT @@ -734,9 +727,7 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout: assert resolved.read == 300.0 assert resolved.connect == 5.0 - def test_cached_async_client_built_with_explicit_request_timeout( - self, monkeypatch: pytest.MonkeyPatch - ): + def test_cached_async_client_built_with_explicit_request_timeout(self, monkeypatch: pytest.MonkeyPatch): from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.utils import LlmProviders @@ -1195,3 +1186,131 @@ async def test_aiohttp_session_never_replays_one_upstreams_cookie_to_another(): assert len(jar) == 0 assert dict(jar.filter_cookies(URL("https://upstream-a.example.com"))) == {} await session.close() + + +def _mint_session_on_dead_loop(handler: AsyncHTTPHandler) -> ClientSession: + """Create the transport's real ClientSession on a loop that then closes. + + This is the lifecycle of every client minted for a short-lived event loop + (the loop-id-keyed LLM client cache creates one handler per loop): the + session outlives its loop and can only ever be disposed loop-lessly. + """ + transport = handler.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport) + loop = asyncio.new_event_loop() + + async def _create() -> ClientSession: + return transport._get_valid_client_session() + + session = loop.run_until_complete(_create()) + loop.close() + return session + + +def test_finalizer_without_running_loop_closes_dead_loop_session(): + """A handler finalized with no running event loop must still dispose its + aiohttp session. + + The async close can never run in that context; without the synchronous + fallback the session and its connector are abandoned to GC and emit + "Unclosed client session" / "Unclosed connector" warnings.""" + handler = AsyncHTTPHandler(timeout=61.0) + session = _mint_session_on_dead_loop(handler) + assert not session.closed + + del handler + gc.collect() + + assert session.closed + + +@pytest.mark.asyncio +async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref(): + """With a running loop, finalization schedules an async close and must keep + a strong reference to the task until it completes — a bare create_task() + result may be collected before it runs, leaving the session unclosed.""" + handler = AsyncHTTPHandler(timeout=61.0) + transport = handler.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport) + session = transport._get_valid_client_session() + assert not session.closed + del transport + + baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) + del handler + gc.collect() + + scheduled = AsyncHTTPHandler._finalizer_close_tasks - baseline_tasks + assert len(scheduled) == 1 + + await asyncio.gather(*scheduled) + assert session.closed + assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled) + + +@pytest.mark.asyncio +async def test_sync_close_helper_respects_session_ownership(): + """The loop-less fallback closes only sessions the transport owns; a + shared session (e.g. the proxy's) must never be closed by a handler.""" + owned_handler = AsyncHTTPHandler(timeout=61.0) + owned_transport = owned_handler.client._transport + assert isinstance(owned_transport, LiteLLMAiohttpTransport) + owned_session = owned_transport._get_valid_client_session() + + baseline = set(LiteLLMAiohttpTransport._background_close_tasks) + owned_handler._dispose_wrapped_aiohttp_session() + scheduled = LiteLLMAiohttpTransport._background_close_tasks - baseline + await asyncio.gather(*scheduled) + assert owned_session.closed + + shared_session = ClientSession() + shared_handler = AsyncHTTPHandler(timeout=61.0, shared_session=shared_session) + shared_transport = shared_handler.client._transport + assert isinstance(shared_transport, LiteLLMAiohttpTransport) + assert shared_transport._owns_session is False + + shared_handler._dispose_wrapped_aiohttp_session() + assert not shared_session.closed + + await shared_session.close() + await shared_handler.close() + await owned_handler.close() + + +@pytest.mark.asyncio +async def test_finalizer_close_done_consumes_exception(): + """A failing finalizer close must have its exception retrieved by the done + callback, or asyncio emits "Task exception was never retrieved" at GC — + the same log noise the finalizer path exists to eliminate.""" + + async def failing_close() -> None: + raise RuntimeError("close failed") + + task = asyncio.get_running_loop().create_task(failing_close()) + AsyncHTTPHandler._finalizer_close_tasks.add(task) + await asyncio.sleep(0) + + AsyncHTTPHandler._on_finalizer_close_done(task) + assert task not in AsyncHTTPHandler._finalizer_close_tasks + + cancelled = asyncio.get_running_loop().create_task(asyncio.sleep(30)) + cancelled.cancel() + await asyncio.sleep(0) + AsyncHTTPHandler._on_finalizer_close_done(cancelled) + + +@pytest.mark.asyncio +async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_scheduling(): + """GC on a live loop (e.g. the app's) of a handler whose session belongs to + another, dead loop must not schedule aclose() here — that is the cross-loop + path the transport refuses — and must still dispose the session.""" + handler = AsyncHTTPHandler(timeout=61.0) + session = await asyncio.to_thread(_mint_session_on_dead_loop, handler) + assert not session.closed + + baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) + del handler + gc.collect() + + assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks + assert session.closed diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 9faa77d6dce..a93b14d45f3 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -26,6 +26,8 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _has_pre_call_deployment_hook, _rust_responses_websocket_enabled, ) +from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import TranscriptionResponse @@ -2524,3 +2526,132 @@ def test_only_callbacks_that_can_charge_a_frame_are_collected_for_ws_quota(monke monkeypatch.setattr(litellm, "callbacks", [plain, quota, decoy]) assert _collect_ws_project_quota_callbacks() == (quota,) + + +class _JSONBodyVideoConfig(OpenAIVideoConfig): + def use_multipart_form_data(self) -> bool: + return False + + +def _video_create_call_kwargs(config, **optional_params): + return { + "model": "sora-2", + "prompt": "a cat surfing", + "video_generation_provider_config": config, + "video_generation_optional_request_params": {"seconds": "4", **optional_params}, + "custom_llm_provider": "openai", + "litellm_params": GenericLiteLLMParams(api_key="sk-test", api_base="https://video.example/v1"), + "logging_obj": Mock(), + "timeout": 10.0, + } + + +def _capture_video_create_request(captured): + def respond(request): + captured["content_type"] = request.headers.get("content-type") + captured["body"] = request.content + return httpx.Response( + 200, + json={"id": "video_123", "object": "video", "status": "queued", "created_at": 1712697600, "model": "sora-2"}, + ) + + return respond + + +def _multipart_text_fields(content_type: str, body: bytes) -> dict: + boundary = content_type.split("boundary=")[1].encode() + return { + part.split(b'name="')[1].split(b'"')[0].decode(): part.partition(b"\r\n\r\n")[2].rstrip(b"\r\n-").decode() + for part in body.split(b"--" + boundary) + if b'name="' in part and b"filename=" not in part + } + + +def test_video_generation_without_file_sends_multipart_form_data(): + """Regression for #36493: the OpenAI SDK always sends /videos requests as + multipart/form-data, so OpenAI-compatible backends (SGLang Diffusion, + vLLM-Omni) reject the JSON body LiteLLM used to send when no + input_reference file was attached.""" + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_video_create_request(captured)))) + + result = BaseLLMHTTPHandler().video_generation_handler(client=client, **_video_create_call_kwargs(OpenAIVideoConfig())) + + assert captured["content_type"].startswith("multipart/form-data") + assert _multipart_text_fields(captured["content_type"], captured["body"]) == { + "model": "sora-2", + "prompt": "a cat surfing", + "seconds": "4", + } + assert result.status == "queued" + + +@pytest.mark.asyncio +async def test_async_video_generation_without_file_sends_multipart_form_data(): + captured = {} + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(_capture_video_create_request(captured))) + + result = await BaseLLMHTTPHandler().async_video_generation_handler( + client=client, **_video_create_call_kwargs(OpenAIVideoConfig()) + ) + + assert captured["content_type"].startswith("multipart/form-data") + assert _multipart_text_fields(captured["content_type"], captured["body"]) == { + "model": "sora-2", + "prompt": "a cat surfing", + "seconds": "4", + } + assert result.status == "queued" + + +def test_azure_video_generation_without_file_sends_multipart_form_data(): + """AzureVideoConfig subclasses OpenAIVideoConfig, so it inherits the + file-less multipart behavior. Azure's /openai/v1/videos surface is + OpenAI-SDK-compatible (the SDK sends multipart there too), so this is + intentional; lock it so the inherited flip can't silently regress to JSON.""" + assert AzureVideoConfig().use_multipart_form_data() is True + + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_video_create_request(captured)))) + + result = BaseLLMHTTPHandler().video_generation_handler(client=client, **_video_create_call_kwargs(AzureVideoConfig())) + + assert captured["content_type"].startswith("multipart/form-data") + assert _multipart_text_fields(captured["content_type"], captured["body"]) == { + "model": "sora-2", + "prompt": "a cat surfing", + "seconds": "4", + } + assert result.status == "queued" + + +def test_video_generation_json_provider_keeps_json_body(): + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_video_create_request(captured)))) + + result = BaseLLMHTTPHandler().video_generation_handler(client=client, **_video_create_call_kwargs(_JSONBodyVideoConfig())) + + assert captured["content_type"] == "application/json" + assert json.loads(captured["body"]) == {"model": "sora-2", "prompt": "a cat surfing", "seconds": "4"} + assert result.status == "queued" + + +def test_video_generation_with_input_reference_keeps_file_multipart(): + captured = {} + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_capture_video_create_request(captured)))) + + result = BaseLLMHTTPHandler().video_generation_handler( + client=client, + **_video_create_call_kwargs(OpenAIVideoConfig(), input_reference=b"\x89PNG\r\n\x1a\nfakepng"), + ) + + assert captured["content_type"].startswith("multipart/form-data") + assert b'name="input_reference"' in captured["body"] + assert b'filename="input_reference.png"' in captured["body"] + assert _multipart_text_fields(captured["content_type"], captured["body"]) == { + "model": "sora-2", + "prompt": "a cat surfing", + "seconds": "4", + } + assert result.status == "queued" diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py new file mode 100644 index 00000000000..21f047b753c --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -0,0 +1,264 @@ +import json +from decimal import Decimal +from pathlib import Path +from typing import Final + +import pytest + +import litellm +from litellm.llms.databricks.cost_calculator import cost_per_token +from litellm.types.utils import ModelInfo, Usage + +REPO_ROOT: Final = Path(__file__).parents[4] +MAIN_PRICES: Final = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PRICES: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" +NEW_MODELS: Final = ( + "databricks/databricks-claude-opus-4-7", + "databricks/databricks-claude-opus-4-8", + "databricks/databricks-claude-opus-5", + "databricks/databricks-claude-sonnet-5", + "databricks/databricks-claude-fable-5", +) + +DOLLARS_PER_DBU: Final = Decimal("0.070") +PRICE_FIELDS: Final = ( + "input_cost_per_token", + "output_cost_per_token", + "cache_creation_input_token_cost", + "cache_read_input_token_cost", +) +PUBLISHED_DBU_PER_MILLION: Final = { + "databricks/databricks-claude-fable-5": ("142.858", "714.286", "178.572", "14.286"), + "databricks/databricks-claude-opus-5": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-8": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-7": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-6": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-5": ("71.429", "357.143", "89.286", "7.143"), + "databricks/databricks-claude-opus-4-1": ("214.286", "1071.429", "267.857", "21.429"), + "databricks/databricks-claude-opus-4": ("214.286", "1071.429", "267.857", "21.429"), + "databricks/databricks-claude-sonnet-5": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4-6": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4-5": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4-1": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-4": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-3-7-sonnet": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-haiku-4-5": ("14.286", "71.429", "17.857", "1.429"), + "databricks/databricks-gpt-5": ("17.857", "142.857", "17.857", "1.786"), + "databricks/databricks-gpt-5-1": ("17.857", "142.857", "17.857", "1.786"), + "databricks/databricks-gpt-5-1-codex-max": ("17.857", "142.857", "17.857", "1.786"), + "databricks/databricks-gpt-5-1-codex-mini": ("3.571", "28.571", "3.571", "0.357"), + "databricks/databricks-gpt-5-mini": ("3.571", "28.571", "3.571", "0.357"), + "databricks/databricks-gpt-5-nano": ("0.714", "5.714", "0.714", "0.071"), + "databricks/databricks-gpt-5-2": ("25.000", "200.000", "25.000", "2.500"), + "databricks/databricks-gpt-5-2-codex": ("25.000", "200.000", "25.000", "2.500"), + "databricks/databricks-gpt-5-3-codex": ("25.000", "200.000", "25.000", "2.500"), + "databricks/databricks-gpt-5-4": ("35.714", "214.286", "35.714", "3.571"), + "databricks/databricks-gpt-5-4-mini": ("10.714", "64.286", "10.714", "1.071"), + "databricks/databricks-gpt-5-4-nano": ("2.857", "17.857", "2.857", "0.286"), + "databricks/databricks-gemini-3-1-pro": ("35.714", "214.286", "35.714", "3.571"), + "databricks/databricks-gemini-3-pro": ("35.714", "214.286", "35.714", "3.571"), + "databricks/databricks-gemini-3-flash": ("8.929", "53.571", "8.929", "0.893"), + "databricks/databricks-gemini-3-1-flash-lite": ("4.464", "26.786", "4.464", "0.446"), + "databricks/databricks-gemini-2-5-pro": ("22.321", "178.571", "22.321", "2.232"), + "databricks/databricks-gemini-2-5-flash": ("5.357", "44.643", "5.357", "0.536"), +} +PROMOTIONAL_DISCOUNT: Final = 0.80 +PROMOTION_EXPIRES: Final = "2027-01-31" +ENTRIES_STORING_PROMOTIONAL_RATE: Final = ( + "databricks/databricks-gemini-2-5-pro", + "databricks/databricks-gemini-2-5-flash", +) +ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION: Final = ( + "databricks/databricks-gemini-3-1-pro", + "databricks/databricks-gemini-3-pro", + "databricks/databricks-gemini-3-flash", + "databricks/databricks-gemini-3-1-flash-lite", +) +CACHE_FIELDS: Final = ("cache_creation_input_token_cost", "cache_read_input_token_cost") + + +def _model_info(model: str) -> ModelInfo: + return litellm.get_model_info(model=model, custom_llm_provider="databricks") + + +def _dollars_per_token(dbu_per_million: str) -> float: + return float(Decimal(dbu_per_million) * DOLLARS_PER_DBU / Decimal(10) ** 6) + + +@pytest.mark.parametrize( + "model", + [ + "databricks/databricks-claude-opus-4-8", + "databricks/databricks-claude-opus-5", + "databricks/databricks-claude-sonnet-5", + ], +) +def test_cached_tokens_bill_at_cache_rates(local_model_cost_map: None, model: str) -> None: + info: Final = _model_info(model) + usage: Final = Usage( + prompt_tokens=11000, + completion_tokens=500, + total_tokens=11500, + cache_creation_input_tokens=2000, + cache_read_input_tokens=8000, + ) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + assert prompt_cost == pytest.approx( + 1000 * info["input_cost_per_token"] + + 2000 * info["cache_creation_input_token_cost"] + + 8000 * info["cache_read_input_token_cost"] + ) + assert completion_cost == pytest.approx(500 * info["output_cost_per_token"]) + assert prompt_cost < 11000 * info["input_cost_per_token"] + + +def test_uncached_request_bills_every_prompt_token_at_the_input_rate(local_model_cost_map: None) -> None: + model: Final = "databricks/databricks-claude-sonnet-5" + info: Final = _model_info(model) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + assert prompt_cost == pytest.approx(1000 * info["input_cost_per_token"]) + assert completion_cost == pytest.approx(200 * info["output_cost_per_token"]) + + +def test_legacy_endpoint_names_still_resolve(local_model_cost_map: None) -> None: + info: Final = _model_info("databricks/databricks-mixtral-8x7b-instruct") + usage: Final = Usage(prompt_tokens=100, completion_tokens=100, total_tokens=200) + + prompt_cost, completion_cost = cost_per_token(model="databricks/mixtral-8x7b-instruct-v0.1", usage=usage) + + assert prompt_cost == pytest.approx(100 * info["input_cost_per_token"]) + assert completion_cost == pytest.approx(100 * info["output_cost_per_token"]) + + +@pytest.mark.parametrize("model", NEW_MODELS) +def test_new_models_price_at_published_dbu_rates(local_model_cost_map: None, model: str) -> None: + info: Final = _model_info(model) + + for field, dbu_per_million in zip(PRICE_FIELDS, PUBLISHED_DBU_PER_MILLION[model]): + assert info[field] == _dollars_per_token(dbu_per_million), field + + +@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(ENTRIES_STORING_PROMOTIONAL_RATE))) +def test_cache_rates_derive_from_published_cache_dbu(local_model_cost_map: None, model: str) -> None: + info: Final = _model_info(model) + cache_dbu_per_million: Final = PUBLISHED_DBU_PER_MILLION[model][2:] + + for field, dbu_per_million in zip(CACHE_FIELDS, cache_dbu_per_million): + assert info[field] == _dollars_per_token(dbu_per_million), field + + +@pytest.mark.parametrize("model", NEW_MODELS) +def test_new_models_carry_cache_pricing(local_model_cost_map: None, model: str) -> None: + info: Final = _model_info(model) + + assert info["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] > 0 + assert info["cache_creation_input_token_cost"] > info["input_cost_per_token"] + assert info["cache_read_input_token_cost"] < info["input_cost_per_token"] + assert info["supports_prompt_caching"] is True + + +def test_every_priced_databricks_model_declares_cache_rates(local_model_cost_map: None) -> None: + undeclared: Final = [ + model + for model, info in litellm.model_cost.items() + if model.startswith("databricks/") + and info.get("input_cost_per_token") is not None + and any(info.get(field) is None for field in CACHE_FIELDS) + ] + + assert undeclared == [] + + +def test_models_without_a_cache_discount_bill_cache_tokens_at_the_input_rate( + local_model_cost_map: None, +) -> None: + model: Final = "databricks/databricks-meta-llama-3-3-70b-instruct" + info: Final = _model_info(model) + usage: Final = Usage( + prompt_tokens=10000, + completion_tokens=100, + total_tokens=10100, + cache_read_input_tokens=8000, + ) + + prompt_cost, _ = cost_per_token(model=model, usage=usage) + + assert prompt_cost == pytest.approx(10000 * info["input_cost_per_token"]) + assert prompt_cost > 8000 * info["input_cost_per_token"] + + +def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_rate( + local_model_cost_map: None, +) -> None: + without_published_rates: Final = [ + model + for model, info in litellm.model_cost.items() + if model.startswith("databricks/") + and info.get("input_cost_per_token") + and model not in PUBLISHED_DBU_PER_MILLION + ] + + assert len(without_published_rates) == 14 + for model in without_published_rates: + info = _model_info(model) + for field in CACHE_FIELDS: + assert info[field] == pytest.approx(info["input_cost_per_token"]), (model, field) + + +@pytest.mark.parametrize("model", NEW_MODELS) +def test_backup_price_map_matches_main(model: str) -> None: + main_cost: Final = json.loads(MAIN_PRICES.read_text()) + backup_cost: Final = json.loads(BACKUP_PRICES.read_text()) + + assert model in main_cost + assert model in backup_cost + assert backup_cost[model] == main_cost[model] + + +def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: None) -> None: + sonnet_5: Final = _model_info("databricks/databricks-claude-sonnet-5") + sonnet_4_6: Final = _model_info("databricks/databricks-claude-sonnet-4-6") + + for field in PRICE_FIELDS: + assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field + + +@pytest.mark.parametrize("model", ENTRIES_STORING_PROMOTIONAL_RATE) +def test_entries_storing_the_promotional_rate_price_below_the_published_table( + local_model_cost_map: None, + model: str, +) -> None: + info: Final = _model_info(model) + input_dbu, output_dbu, _, _ = PUBLISHED_DBU_PER_MILLION[model] + expiry_hint: Final = f"the gemini promotion expires {PROMOTION_EXPIRES}, after which the list rate applies" + + assert info["input_cost_per_token"] == pytest.approx( + _dollars_per_token(input_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 + ), expiry_hint + assert info["output_cost_per_token"] == pytest.approx( + _dollars_per_token(output_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 + ), expiry_hint + assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) + assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"]) + + +@pytest.mark.parametrize("model", ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION) +def test_entries_storing_the_list_rate_bill_above_the_promotional_price( + local_model_cost_map: None, + model: str, +) -> None: + info: Final = _model_info(model) + input_dbu, _, _, _ = PUBLISHED_DBU_PER_MILLION[model] + list_rate: Final = _dollars_per_token(input_dbu) + + assert info["input_cost_per_token"] == pytest.approx(list_rate, rel=2e-4), ( + f"{model} moved off the list rate; if it now stores the discount that runs to " + f"{PROMOTION_EXPIRES}, move it into ENTRIES_STORING_PROMOTIONAL_RATE" + ) + assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index e728fc4bc40..63a749dab84 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -473,12 +473,14 @@ def test_transform_messages_helper_strips_thinking_blocks(): "thinking_blocks": [ {"type": "thinking", "thinking": "internal", "signature": ""} ], + "reasoning_content": "internal", }, ] out = config._transform_messages_helper( messages, model="accounts/fireworks/models/glm-5p1", litellm_params={} ) assert "thinking_blocks" not in out[1] + assert "reasoning_content" not in out[1] assert out[1]["content"] == "I can help." diff --git a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py index e316cd14dd4..82b05601a85 100644 --- a/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py +++ b/tests/test_litellm/llms/hosted_vllm/chat/test_hosted_vllm_chat_transformation.py @@ -200,6 +200,7 @@ def test_hosted_vllm_thinking_blocks_prepended_to_assistant_content(): "signature": "abc123", } ], + "reasoning_content": "Let me reason about this...", }, { "role": "user", @@ -218,6 +219,7 @@ def test_hosted_vllm_thinking_blocks_prepended_to_assistant_content(): assert isinstance(assistant_msg["content"], str) assert assistant_msg["content"] == "Here is my answer." assert "thinking_blocks" not in assistant_msg + assert "reasoning_content" not in assistant_msg def test_hosted_vllm_thinking_blocks_with_list_content(): diff --git a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py index acd69b94d02..eadc2bc9541 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_completion_transformation.py @@ -475,7 +475,7 @@ class TestOllamaTextCompletionResponseIterator: assert isinstance(result, ModelResponseStream) assert result.choices and result.choices[0].delta is not None assert result.choices[0].delta.content == None - assert getattr(result.choices[0].delta, "reasoning_content", None) is "" + assert getattr(result.choices[0].delta, "reasoning_content", None) == "" def test_chunk_parser_done_chunk(self): """Test that done chunks work correctly.""" diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index d279b119efe..a35b75a6106 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -1309,3 +1309,79 @@ def test_responses_gpt54_allow_temperature_effort_none( drop_params=False, ) assert params["temperature"] == 0.7 + + +@pytest.mark.parametrize("model", ["gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"]) +def test_gpt5_6_forwards_reasoning_effort_max_for_the_responses_bridge(config: OpenAIConfig, model: str): + """A chat request carrying tools or a reasoning summary is converted to /v1/responses further + down main.py, and that surface accepts max. This runs before litellm has decided to bridge, so + refusing max here would break the cursor thinking-max shape that works today. Plain chat + completions still answer max with a provider 400, and the capability list below is what keeps + the level out of the picker.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "max"}, + optional_params={}, + model=model, + drop_params=False, + ) + assert params["reasoning_effort"] == "max" + + +@pytest.mark.parametrize("model", ["gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"]) +def test_gpt5_6_never_advertises_reasoning_effort_max(model: str): + """/v1/chat/completions answers max with "Unsupported value: 'reasoning_effort' does not support + 'max' with this model. Supported values are: 'none', 'low', 'medium', 'high', and 'xhigh'", so no + gpt-5.6 entry asserts supports_max_reasoning_effort and the advertised set stops at xhigh.""" + from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + + resolved = resolve_supported_reasoning_efforts(litellm.get_model_info(model), deployment_is_mapped=True) + assert resolved is not None + assert "max" not in resolved + assert "xhigh" in resolved + + +def test_gpt5_6_keeps_reasoning_effort_max_on_the_responses_api( + responses_config: OpenAIResponsesAPIConfig, +): + """/v1/responses accepts max for gpt-5.6, and that is the surface the cursor thinking-max + variant resolves onto, so the responses path keeps carrying the level chat refuses.""" + params = responses_config.map_openai_params( + response_api_optional_params=ResponsesAPIOptionalRequestParams( + reasoning={"effort": "max"}, + ), + model="gpt-5.6", + drop_params=False, + ) + assert params["reasoning"] == {"effort": "max"} + + +def test_gpt5_forwards_levels_the_chat_gate_does_not_own(config: OpenAIConfig): + """Only xhigh is gated on this surface. max reaches the provider (or the responses bridge) and + is answered there, which is what happened before per-group capabilities existed.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "max"}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + assert params["reasoning_effort"] == "max" + + +def test_gpt5_rejects_xhigh_for_models_without_the_flag(config: OpenAIConfig): + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + + +def test_gpt5_drops_xhigh_when_requested(config: OpenAIConfig): + params = config.map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model="gpt-5.1", + drop_params=True, + ) + assert "reasoning_effort" not in params diff --git a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py index c15554a46a5..89f52261cba 100644 --- a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py +++ b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py @@ -1,3 +1,5 @@ +import io + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.utils import encode_character_id_with_provider @@ -68,3 +70,48 @@ def test_wrapped_character_id_is_decoded_then_encoded_as_path_segment(): == "https://api.openai.com/v1/videos/characters/..%2F..%2Fcharacters%3Fx%3D1%23frag" ) assert params == {} + + +def test_video_edit_request_forwards_uploaded_file_as_multipart(): + """An uploaded source video must leave as a multipart ``video`` file part, + not be dropped in favor of a JSON id reference.""" + config = OpenAIVideoConfig() + source = io.BytesIO(b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isomBODY") + source.name = "clip.mp4" + + url, data, files = config.transform_video_edit_request( + prompt="make it nighttime", + video_id="", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + video_file=source, + ) + + assert url == "https://api.openai.com/v1/videos/edits" + assert data == {"prompt": "make it nighttime"} + assert files is not None + field_names = [field for field, _ in files] + assert field_names == ["video"] + _, (filename, content, content_type) = files[0] + assert filename == "clip.mp4" + assert content is source + assert content_type == "video/mp4" + + +def test_video_edit_request_without_file_sends_json_id_reference(): + """The id-reference path must stay JSON (files is None) so existing + remix/edit-by-id callers keep working.""" + config = OpenAIVideoConfig() + + url, data, files = config.transform_video_edit_request( + prompt="brighter", + video_id="video_abc123", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.openai.com/v1/videos/edits" + assert data == {"prompt": "brighter", "video": {"id": "video_abc123"}} + assert files is None diff --git a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py index 24879ce83f9..afc8e7ec4a2 100644 --- a/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py +++ b/tests/test_litellm/llms/runwayml/videos/test_runway_video_transformation.py @@ -7,7 +7,12 @@ from unittest.mock import Mock import httpx import pytest -from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.runwayml.videos.transformation import ( + RunwayMLError, + RunwayMLVideoConfig, + _ratio_to_resolution, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoObject @@ -49,6 +54,158 @@ class TestRunwayMLVideoTransformation: # Validate URL has correct endpoint assert url == "https://api.dev.runwayml.com/v1/image_to_video" + def test_transform_video_create_request_text_to_video(self): + """A prompt-only request must hit /text_to_video, not /image_to_video.""" + data, files, url = self.config.transform_video_create_request( + model="veo3.1", + prompt="A serene mountain lake at sunrise", + api_base="https://api.dev.runwayml.com/v1", + video_create_optional_request_params={"duration": 8, "ratio": "1280:720"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.dev.runwayml.com/v1/text_to_video" + assert "promptImage" not in data + assert data["promptText"] == "A serene mountain lake at sunrise" + + def test_transform_video_create_request_video_to_video(self): + """A promptVideo request must hit /video_to_video with promptImage stripped.""" + data, files, url = self.config.transform_video_create_request( + model="aleph2", + prompt="Make it snow", + api_base="https://api.dev.runwayml.com/v1", + video_create_optional_request_params={ + "promptVideo": "https://example.com/source.mp4", + "promptImage": "https://example.com/reference.png", + "ratio": "1280:720", + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.dev.runwayml.com/v1/video_to_video" + assert data["promptVideo"] == "https://example.com/source.mp4" + assert "promptImage" not in data + + def test_transform_video_create_request_video_uri_routes_to_video_to_video(self): + _, _, url = self.config.transform_video_create_request( + model="aleph2", + prompt="Make it snow", + api_base="https://api.dev.runwayml.com/v1", + video_create_optional_request_params={"videoUri": "https://example.com/source.mp4"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.dev.runwayml.com/v1/video_to_video" + + def test_status_progress_fraction_scales_to_percent(self): + """Runway reports progress as a 0..1 float; VideoObject.progress is an int percent.""" + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "63fd0f13-f29d-4e58-99d3-1cb9efa14a5b", + "createdAt": "2025-11-11T21:48:50.448Z", + "status": "RUNNING", + "progress": 0.027, + } + + result = self.config.transform_video_status_retrieve_response( + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + custom_llm_provider="runwayml", + ) + + assert result.status == "in_progress" + assert result.progress == 3 + + def test_status_progress_null_leaves_progress_unset(self): + """Runway sends an explicit null progress for pending polls; scaling it must not crash.""" + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "63fd0f13-f29d-4e58-99d3-1cb9efa14a5b", + "createdAt": "2025-11-11T21:48:50.448Z", + "status": "PENDING", + "progress": None, + } + + result = self.config.transform_video_status_retrieve_response( + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + custom_llm_provider="runwayml", + ) + + assert result.status == "queued" + assert result.progress is None + + def test_get_error_class_returns_exception_instead_of_raising(self): + error = self.config.get_error_class( + error_message="Invalid API key", + status_code=401, + headers={}, + ) + + assert isinstance(error, RunwayMLError) + assert isinstance(error, BaseLLMException) + assert error.status_code == 401 + assert error.message == "Invalid API key" + + def test_create_response_usage_includes_resolution_and_provider_cost(self): + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "test-video-id-123", + "createdAt": "2025-11-11T21:48:50.448Z", + "status": "PENDING", + "estimatedCost": {"credits": 25.0}, + } + + video_obj = self.config.transform_video_create_response( + model="gen4_turbo", + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + custom_llm_provider="runwayml", + request_data={"model": "gen4_turbo", "ratio": "1280:720", "duration": 5}, + ) + + assert video_obj.usage == { + "duration_seconds": 5.0, + "video_resolution": "720p", + "provider_reported_cost_usd": 0.25, + } + + def test_create_response_usage_omits_unknown_fields(self): + mock_response = Mock(spec=httpx.Response) + mock_response.json.return_value = { + "id": "test-video-id-123", + "createdAt": "2025-11-11T21:48:50.448Z", + "status": "PENDING", + } + + video_obj = self.config.transform_video_create_response( + model="gen4_turbo", + raw_response=mock_response, + logging_obj=self.mock_logging_obj, + custom_llm_provider="runwayml", + request_data={"model": "gen4_turbo"}, + ) + + assert video_obj.usage == {} + + @pytest.mark.parametrize( + "ratio,expected", + [ + ("848:480", "480p"), + ("1280:720", "720p"), + ("1920:1080", "1080p"), + ("2560:1440", "1080p"), + ("3840:2160", "4k"), + (None, None), + ("banana", None), + ], + ) + def test_ratio_to_resolution_tiers(self, ratio, expected): + assert _ratio_to_resolution(ratio) == expected + def test_transform_video_status_with_timestamp_handling(self): """Test status retrieval handles RunwayML's ISO 8601 timestamps correctly.""" from litellm.types.videos.utils import encode_video_id_with_provider diff --git a/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py new file mode 100644 index 00000000000..6216d3bf225 --- /dev/null +++ b/tests/test_litellm/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -0,0 +1,232 @@ +import json +from unittest.mock import MagicMock + +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, +) +from litellm.llms.together_ai.chat.transformation import TogetherAIChatConfig +from litellm.types.utils import LlmProviders, ModelResponse + +TOOL_CALLING_MODEL = "openai/gpt-oss-20b" +REASONING_MODEL = "deepseek-ai/DeepSeek-V3.1" +PLAIN_MODEL = "Qwen/Qwen3-235B-A22B-fp8-tput" +UNMAPPED_MODEL = "MiniMaxAI/MiniMax-M3" + +FUNCTION_CALLING_PARAMS = ("tools", "tool_choice", "function_call", "response_format") + + +@pytest.fixture(autouse=True) +def force_local_model_cost(monkeypatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) + + +def test_supported_params_tool_calling_model(): + supported = TogetherAIChatConfig().get_supported_openai_params(model=TOOL_CALLING_MODEL) + + for param in FUNCTION_CALLING_PARAMS: + assert param in supported + + +def test_supported_params_plain_model(): + supported = TogetherAIChatConfig().get_supported_openai_params(model=PLAIN_MODEL) + + for param in FUNCTION_CALLING_PARAMS: + assert param not in supported + assert "temperature" in supported + assert "max_tokens" in supported + + +def test_supported_params_unmapped_model_treated_as_plain(): + supported = TogetherAIChatConfig().get_supported_openai_params(model=UNMAPPED_MODEL) + + for param in FUNCTION_CALLING_PARAMS: + assert param not in supported + assert "stream" in supported + + +def test_map_openai_params_tool_calling_model_passes_tools(): + tools = [{"type": "function", "function": {"name": "get_weather", "parameters": {}}}] + + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"tools": tools, "tool_choice": "auto"}, + optional_params={}, + model=TOOL_CALLING_MODEL, + drop_params=False, + ) + + assert mapped["tools"] == tools + assert mapped["tool_choice"] == "auto" + + +def test_map_openai_params_reasoning_model_passes_sampling_params(): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"temperature": 0.2, "max_tokens": 512}, + optional_params={}, + model=REASONING_MODEL, + drop_params=False, + ) + + assert mapped["temperature"] == 0.2 + assert mapped["max_tokens"] == 512 + + +def test_map_openai_params_drops_text_response_format(): + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"response_format": {"type": "text"}, "temperature": 0.5}, + optional_params={}, + model=REASONING_MODEL, + drop_params=False, + ) + + assert "response_format" not in mapped + assert mapped["temperature"] == 0.5 + + +def test_map_openai_params_keeps_json_response_format(): + response_format = {"type": "json_object"} + + mapped = TogetherAIChatConfig().map_openai_params( + non_default_params={"response_format": response_format}, + optional_params={}, + model=TOOL_CALLING_MODEL, + drop_params=False, + ) + + assert mapped["response_format"] == response_format + + +def _transform_response(message: dict) -> ModelResponse: + raw_response_json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1234567890, + "model": REASONING_MODEL, + "choices": [{"index": 0, "message": message, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = raw_response_json + mock_response.text = json.dumps(raw_response_json) + mock_response.headers = {} + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + logging_obj.model_call_details = {} + + return TogetherAIChatConfig().transform_response( + model=REASONING_MODEL, + raw_response=mock_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "What is 2+2?"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key="test-key", + json_mode=False, + ) + + +def test_transform_response_maps_reasoning_to_reasoning_content(): + result = _transform_response( + {"role": "assistant", "content": "4", "reasoning": "2+2 equals 4"} + ) + + assert result.choices[0].message.content == "4" + assert result.choices[0].message.reasoning_content == "2+2 equals 4" + + +def test_transform_response_preserves_reasoning_content_field(): + result = _transform_response( + {"role": "assistant", "content": "4", "reasoning_content": "adding 2 and 2"} + ) + + assert result.choices[0].message.reasoning_content == "adding 2 and 2" + + +def test_streaming_chunk_maps_delta_reasoning_to_reasoning_content(): + iterator = TogetherAIChatConfig().get_model_response_iterator( + streaming_response=iter(()), sync_stream=True + ) + assert isinstance(iterator, OpenAIChatCompletionStreamingHandler) + + parsed = iterator.chunk_parser( + { + "id": "chunk-1", + "created": 1234567890, + "model": REASONING_MODEL, + "choices": [{"index": 0, "delta": {"reasoning": "thinking about 2+2"}}], + } + ) + + assert parsed.choices[0]["delta"]["reasoning_content"] == "thinking about 2+2" + + +def test_together_ai_config_alias_points_at_chat_config(): + assert litellm.TogetherAIConfig is litellm.TogetherAIChatConfig + config = litellm.TogetherAIConfig(max_tokens=10) + assert isinstance(config, TogetherAIChatConfig) + + +def test_provider_config_manager_returns_together_chat_config(): + from litellm.utils import ProviderConfigManager + + config = ProviderConfigManager.get_provider_chat_config( + model=REASONING_MODEL, provider=LlmProviders.TOGETHER_AI + ) + + assert isinstance(config, TogetherAIChatConfig) + + +def test_completion_routes_through_together_chat_config(): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + captured_requests = [] + + def respond(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-together", + "object": "chat.completion", + "created": 1234567890, + "model": REASONING_MODEL, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "4", + "reasoning": "2+2 equals 4", + }, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + response = litellm.completion( + model=f"together_ai/{REASONING_MODEL}", + messages=[{"role": "user", "content": "What is 2+2?"}], + api_key="fake-key", + client=client, + ) + + request = captured_requests[0] + assert str(request.url) == "https://api.together.ai/v1/chat/completions" + assert request.headers["authorization"] == "Bearer fake-key" + assert json.loads(request.content)["model"] == REASONING_MODEL + assert response.choices[0].message.content == "4" + assert response.choices[0].message.reasoning_content == "2+2 equals 4" diff --git a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py new file mode 100644 index 00000000000..3364b1b3872 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -0,0 +1,231 @@ +import pytest + +import litellm +from litellm.interactions.utils import get_provider_interactions_api_config +from litellm.llms.gemini.interactions.transformation import ( + GoogleAIStudioInteractionsConfig, +) +from litellm.llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig, +) +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +GLOBAL_BASE = "https://aiplatform.googleapis.com/v1beta1/projects/test-proj/locations/global/interactions" + + +class MinterRecorder: + def __init__(self, resolved_project: str = "creds-proj") -> None: + self.calls: list[tuple[VERTEX_CREDENTIALS_TYPES | None, str | None]] = [] + self.resolved_project = resolved_project + + def __call__( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + self.calls.append((credentials, project_id)) + return "test-token", project_id or self.resolved_project + + +@pytest.fixture +def minter(): + return MinterRecorder() + + +@pytest.fixture +def config(minter): + return VertexAIInteractionsConfig(mint_access_token=minter) + + +@pytest.fixture +def litellm_params(): + return GenericLiteLLMParams(vertex_project="test-proj", vertex_credentials="creds.json") + + +class TestRegistration: + def test_vertex_ai_returns_vertex_config(self): + assert isinstance(get_provider_interactions_api_config("vertex_ai"), VertexAIInteractionsConfig) + + def test_vertex_ai_beta_returns_vertex_config(self): + assert isinstance(get_provider_interactions_api_config("vertex_ai_beta"), VertexAIInteractionsConfig) + + def test_gemini_still_returns_google_ai_studio_config(self): + gemini_config = get_provider_interactions_api_config("gemini") + assert isinstance(gemini_config, GoogleAIStudioInteractionsConfig) + assert not isinstance(gemini_config, VertexAIInteractionsConfig) + + def test_lazy_import_resolves(self): + assert litellm.VertexAIInteractionsConfig is VertexAIInteractionsConfig + + def test_custom_llm_provider_is_vertex_ai(self, config): + assert config.custom_llm_provider == LlmProviders.VERTEX_AI + + +class TestValidateEnvironment: + def test_sets_bearer_auth_without_gemini_headers(self, config, minter, litellm_params): + headers = config.validate_environment( + headers={}, + model="gemini-omni-flash-preview", + litellm_params=litellm_params, + ) + + assert headers["Authorization"] == "Bearer test-token" + assert headers["Content-Type"] == "application/json" + assert "x-goog-api-key" not in headers + assert "Api-Revision" not in headers + assert minter.calls == [("creds.json", "test-proj")] + + def test_caller_authorization_wins(self, config, litellm_params): + headers = config.validate_environment( + headers={"Authorization": "Bearer caller-token"}, + model="gemini-omni-flash-preview", + litellm_params=litellm_params, + ) + + assert headers["Authorization"] == "Bearer caller-token" + + +class TestGetCompleteUrl: + def test_defaults_to_global_v1beta1(self, config, litellm_params): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + ) + + assert url == GLOBAL_BASE + + def test_stream_appends_alt_sse(self, config, litellm_params): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + stream=True, + ) + + assert url == f"{GLOBAL_BASE}?alt=sse" + + def test_multi_region_location_uses_rep_host(self, config): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "us"}, + ) + + assert url == "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-proj/locations/us/interactions" + + def test_regional_location_uses_regional_host(self, config): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "us-central1"}, + ) + + assert url == ( + "https://us-central1-aiplatform.googleapis.com" + "/v1beta1/projects/test-proj/locations/us-central1/interactions" + ) + + def test_location_env_fallback_is_ignored(self, config, monkeypatch): + monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5") + + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj"}, + ) + + assert url == GLOBAL_BASE + + def test_api_base_override(self, config, litellm_params): + url = config.get_complete_url( + api_base="https://proxy.example.test", + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + ) + + assert url == "https://proxy.example.test/v1beta1/projects/test-proj/locations/global/interactions" + + def test_project_resolved_from_credentials_when_not_passed(self, config, monkeypatch): + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + monkeypatch.setattr(litellm, "vertex_project", None) + + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_credentials": "creds.json"}, + ) + + assert url == "https://aiplatform.googleapis.com/v1beta1/projects/creds-proj/locations/global/interactions" + + def test_invalid_location_rejected(self, config): + with pytest.raises(ValueError, match="Invalid vertex_location"): + config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "evil.com#"}, + ) + + def test_missing_project_rejected(self, monkeypatch): + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + monkeypatch.setattr(litellm, "vertex_project", None) + + def unresolved_minter( + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + return "test-token", "" + + with pytest.raises(ValueError, match="Vertex AI project is required"): + VertexAIInteractionsConfig(mint_access_token=unresolved_minter).get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={}, + ) + + +class TestInteractionByIdRequests: + def test_get_url(self, config, litellm_params): + url, request_body = config.transform_get_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123" + assert request_body == {} + + def test_get_url_encodes_interaction_id(self, config, litellm_params): + url, _ = config.transform_get_interaction_request( + interaction_id="id/with space", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/id%2Fwith%20space" + + def test_delete_url(self, config, litellm_params): + url, request_body = config.transform_delete_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123" + assert request_body == {} + + def test_cancel_url(self, config, litellm_params): + url, request_body = config.transform_cancel_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123:cancel" + assert request_body == {} diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index ba2f20e2337..f19e169dc9e 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -538,7 +538,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 57cd729bc90..862969abbc2 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -479,7 +479,7 @@ class TestVertexAIVideoConfig: }, } - url, data = self.config.transform_video_edit_request( + url, data, files = self.config.transform_video_edit_request( prompt="Make it brighter", video_id=operation_name, api_base=api_base, @@ -488,6 +488,7 @@ class TestVertexAIVideoConfig: prefetched_source_data=prefetched, ) + assert files is None assert url.endswith(":predictLongRunning") assert "veo-3.1-generate-001" in url instance = data["instances"][0] @@ -507,7 +508,7 @@ class TestVertexAIVideoConfig: }, } - _, data = self.config.transform_video_edit_request( + _, data, _ = self.config.transform_video_edit_request( prompt="Make it darker", video_id=operation_name, api_base=api_base, diff --git a/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py b/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py index ffc48ecfae9..bc643bdbc47 100644 --- a/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py +++ b/tests/test_litellm/llms/watsonx/test_watsonx_common_utils.py @@ -124,11 +124,7 @@ class TestGenerateIAMToken: mock_client.reset_mock() mock_cache.reset_mock() - # Configure mock to return values based on env_keys - def get_secret_side_effect(key): - return env_keys.get(key) - - mock_get_secret_str.side_effect = get_secret_side_effect + mock_get_secret_str.side_effect = env_keys.get mock_response = MagicMock() mock_response.json.return_value = { diff --git a/tests/test_litellm/ocr/test_ocr_file_input.py b/tests/test_litellm/ocr/test_ocr_file_input.py index e6216d7c580..feb98d14c03 100644 --- a/tests/test_litellm/ocr/test_ocr_file_input.py +++ b/tests/test_litellm/ocr/test_ocr_file_input.py @@ -18,6 +18,7 @@ from unittest.mock import AsyncMock, MagicMock import orjson import pytest +from starlette.datastructures import FormData from litellm.ocr.main import convert_file_document_to_url_document, get_mime_type @@ -470,7 +471,7 @@ class TestProxySecurityGuard: mock_request = MagicMock() mock_request.headers = {"content-type": "multipart/form-data; boundary=---"} - mock_request.form = AsyncMock(return_value=mock_form) + mock_request.form = AsyncMock(return_value=FormData(mock_form)) result = await self._parse_multipart(mock_request) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py index 5193ad989dc..30095796a7f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py @@ -245,13 +245,21 @@ def _fake_get_async_httpx_client_factory(captured_calls: list): return _fake_get_async_httpx_client -async def _fake_create_client(base_url, client_config=None, **kwargs): +async def _fake_create_client(agent_card, client_config=None, **kwargs): client = MagicMock() if client_config is not None: client._litellm_httpx_client = client_config.httpx_client return client +def _fake_card_resolver(httpx_client, base_url, **kwargs): + resolver = MagicMock() + card = MagicMock() + card.supported_interfaces = () + resolver.get_agent_card = AsyncMock(return_value=card) + return resolver + + @pytest.mark.asyncio async def test_create_a2a_client_leaves_the_shared_client_untouched(): """ @@ -276,6 +284,10 @@ async def test_create_a2a_client_leaves_the_shared_client_untouched(): "litellm.a2a_protocol.main.create_client", new=AsyncMock(side_effect=_fake_create_client), ), + patch( + "litellm.a2a_protocol.main.A2ACardResolver", + side_effect=_fake_card_resolver, + ), ): await create_a2a_client( base_url="http://agent-a:9999", @@ -321,6 +333,10 @@ async def test_create_a2a_client_default_timeout_matches_constant(): "litellm.a2a_protocol.main.create_client", new=AsyncMock(side_effect=_fake_create_client), ), + patch( + "litellm.a2a_protocol.main.A2ACardResolver", + side_effect=_fake_card_resolver, + ), ): await create_a2a_client(base_url="http://127.0.0.1:9") @@ -352,6 +368,10 @@ async def test_create_a2a_client_explicit_timeout_overrides_default(): "litellm.a2a_protocol.main.create_client", new=AsyncMock(side_effect=_fake_create_client), ), + patch( + "litellm.a2a_protocol.main.A2ACardResolver", + side_effect=_fake_card_resolver, + ), ): await create_a2a_client(base_url="http://127.0.0.1:9", timeout=42.5) diff --git a/tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py b/tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py index c48b8cfd5a5..4072f83511e 100644 --- a/tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py +++ b/tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py @@ -14,6 +14,7 @@ from fastapi import HTTPException from litellm.proxy.analytics_endpoints.analytics_endpoints import get_global_activity from litellm.proxy.analytics_endpoints.cache_activity import ( + ERROR_BREAKDOWN_SQL, GROUPS_SQL, CacheActivityGroup, compute_totals, @@ -37,6 +38,11 @@ GROUP_ROWS = [ "generated_completion_tokens": 0, }, ] +ERROR_ROWS = [ + {"call_type": "acompletion", "error_code": "429", "error_class": "RateLimitError", "count": 150}, + {"call_type": "acompletion", "error_code": "401", "error_class": "AuthenticationError", "count": 50}, + {"call_type": "Unknown", "error_code": "Unknown", "error_class": "Unknown", "count": 110}, +] KEY_ALIAS_ROWS = [{"key_alias": "Unnamed Key"}, {"key_alias": "my-key"}] MODEL_ROWS = [{"model": "gpt-5.1"}] @@ -49,6 +55,8 @@ def build_prisma(query_raw: AsyncMock) -> MagicMock: def dispatching_query_raw() -> AsyncMock: async def dispatch(sql: str, *params: object) -> list[dict[str, object]]: + if "error_code" in sql: + return ERROR_ROWS if "GROUP BY" in sql: return GROUP_ROWS if "key_alias" in sql: @@ -67,9 +75,7 @@ def mock_prisma(monkeypatch: pytest.MonkeyPatch) -> MagicMock: @pytest.mark.asyncio async def test_returns_groups_totals_and_filter_options(mock_prisma: MagicMock): - response = await get_global_activity( - start_date="2026-07-01", end_date="2026-07-27", key_aliases=[], models=[] - ) + response = await get_global_activity(start_date="2026-07-01", end_date="2026-07-27", key_aliases=[], models=[]) assert [group.call_type for group in response.groups] == ["acompletion", "Unknown"] assert response.groups[0].api_requests == 1000 @@ -81,6 +87,11 @@ async def test_returns_groups_totals_and_filter_options(mock_prisma: MagicMock): assert response.totals.cache_hit_ratio == pytest.approx((300 / 1610) * 100) assert response.filter_options.key_aliases == ["Unnamed Key", "my-key"] assert response.filter_options.models == ["gpt-5.1"] + assert [(bucket.error_code, bucket.error_class, bucket.count) for bucket in response.error_breakdown] == [ + ("429", "RateLimitError", 150), + ("401", "AuthenticationError", 50), + ("Unknown", "Unknown", 110), + ] @pytest.mark.asyncio @@ -92,11 +103,13 @@ async def test_filters_are_passed_to_sql_as_json_arrays(mock_prisma: MagicMock): models=["gpt-5.1", "claude-opus-4-8"], ) - groups_call = next( - call for call in mock_prisma.db.query_raw.call_args_list if "GROUP BY" in call.args[0] - ) - assert groups_call.args[3] == json.dumps(["my-key"]) - assert groups_call.args[4] == json.dumps(["gpt-5.1", "claude-opus-4-8"]) + filtered_calls = [ + call for call in mock_prisma.db.query_raw.call_args_list if call.args[0] in (GROUPS_SQL, ERROR_BREAKDOWN_SQL) + ] + assert len(filtered_calls) == 2 + for call in filtered_calls: + assert call.args[3] == json.dumps(["my-key"]) + assert call.args[4] == json.dumps(["gpt-5.1", "claude-opus-4-8"]) @pytest.mark.asyncio @@ -131,3 +144,10 @@ def test_totals_denominator_includes_failed_requests(): def test_groups_sql_splits_failures_and_labels_empty_call_type_unknown(): assert "SUM(CASE WHEN sl.\"status\" = 'failure' THEN 1 ELSE 0 END)" in GROUPS_SQL assert "CASE WHEN sl.\"call_type\" = '' THEN 'Unknown' ELSE sl.\"call_type\" END" in GROUPS_SQL + + +def test_error_breakdown_sql_counts_only_failures_bucketed_by_code_and_class(): + assert "sl.\"status\" = 'failure'" in ERROR_BREAKDOWN_SQL + assert "sl.\"metadata\"->'error_information'->>'error_code'" in ERROR_BREAKDOWN_SQL + assert "sl.\"metadata\"->'error_information'->>'error_class'" in ERROR_BREAKDOWN_SQL + assert "GROUP BY 1, 2, 3" in ERROR_BREAKDOWN_SQL diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py index 18e0f2cb559..d1ccd86044c 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_claude_code_marketplace.py @@ -18,6 +18,9 @@ from litellm.types.proxy.claude_code_endpoints import ( UpdatePluginRequest, ) from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( + delete_plugin, + disable_plugin, + enable_plugin, get_marketplace, register_plugin, update_plugin, @@ -72,6 +75,12 @@ _USER = UserAPIKeyAuth( user_id="test-user", ) +_NON_ADMIN_USER = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-5678", + user_id="regular-user", +) + _GIT_SUBDIR_SOURCE = { "source": "git-subdir", "url": "https://github.com/org/monorepo.git", @@ -151,6 +160,7 @@ async def test_update_plugin_replaces_existing_source(): response = await update_plugin( plugin_name=name, request=UpdatePluginRequest(source=new_source, version="2.0.0", description="updated"), + user_api_key_dict=_USER, ) assert response.status == "success" @@ -170,6 +180,7 @@ async def test_update_plugin_not_found(): await update_plugin( plugin_name="does-not-exist", request=UpdatePluginRequest(source=_GIT_SUBDIR_SOURCE), + user_api_key_dict=_USER, ) assert exc_info.value.status_code == 404 @@ -213,6 +224,7 @@ async def test_update_plugin_db_error_maps_to_structured_500(): await update_plugin( plugin_name=name, request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}), + user_api_key_dict=_USER, ) assert exc_info.value.status_code == 500 @@ -341,3 +353,62 @@ async def test_register_plugin_unknown_source_type(): assert exc_info.value.status_code == 400 assert "git-subdir" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_register_plugin_rejects_non_admin(): + """A non-admin key cannot add an entry to the marketplace catalog.""" + request = RegisterPluginRequest(name="attacker-plugin", source=_GIT_SUBDIR_SOURCE) + + with pytest.raises(HTTPException) as exc_info: + await register_plugin(request=request, user_api_key_dict=_NON_ADMIN_USER) + + assert exc_info.value.status_code == 403 + + table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + assert await table.find_unique(where={"name": "attacker-plugin"}) is None + + +@pytest.mark.asyncio +async def test_update_plugin_rejects_non_admin_overwrite(): + """A non-admin key cannot overwrite an existing plugin's source.""" + name = "trusted-plugin" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + + malicious_source = {"source": "github", "repo": "attacker/malicious-repo"} + with pytest.raises(HTTPException) as exc_info: + await update_plugin( + plugin_name=name, + request=UpdatePluginRequest(source=malicious_source), + user_api_key_dict=_NON_ADMIN_USER, + ) + + assert exc_info.value.status_code == 403 + + stored = await _read_stored_manifest(name) + assert stored["source"] == _GIT_SUBDIR_SOURCE + + +@pytest.mark.asyncio +async def test_enable_disable_delete_plugin_reject_non_admin(): + """Non-admin keys cannot enable, disable, or delete catalog entries.""" + name = "trusted-plugin-2" + await register_plugin( + request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"), + user_api_key_dict=_USER, + ) + + for coro in ( + enable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + disable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + delete_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER), + ): + with pytest.raises(HTTPException) as exc_info: + await coro + assert exc_info.value.status_code == 403 + + table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable + assert (await table.find_unique(where={"name": name})).enabled is True diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 04f38b5e2ed..7d2e3f45c04 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import ( _virtual_key_soft_budget_check, get_key_object, get_user_object, + invalidate_team_member_spend_state, vector_store_access_check, ) from litellm.caching.in_memory_cache import InMemoryCache @@ -4868,10 +4869,9 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): 3. When team_alias is None, NO alias-key operation happens (no delete of an empty-keyed entry, no spurious write). 4. DELETES the team_id-keyed entry from the internal usage cache - BEFORE the fresh write (LIT-4391). `_get_team_object_from_cache` - consults the internal usage cache first, so a leftover copy there - (backfilled from a Redis shared with `user_api_key_cache`) would - keep serving the pre-update team allowlist. + BEFORE the fresh write (LIT-4391). `_get_team_object_from_cache` no + longer reads the internal usage cache (LIT-5944), but the delete + protects mixed-version rolling deploys where older workers still do. """ from unittest.mock import AsyncMock, MagicMock @@ -4980,8 +4980,10 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): Regression test for LIT-4391: keys with models=["all-team-models"] kept getting 403 team_model_access_denied for models added via /team/update. - `_get_team_object_from_cache` consults the internal usage cache BEFORE - `user_api_key_cache`. When both share one Redis (enable_redis_auth_cache), + `_get_team_object_from_cache` used to consult the internal usage cache + BEFORE `user_api_key_cache` (removed in LIT-5944; this test now also + guards against reintroducing that read). + When both share one Redis (enable_redis_auth_cache), any team read backfills the internal cache's in-memory tier with the team object. `_cache_team_object` (the /team/update refresh) only wrote `user_api_key_cache`, so that backfilled copy kept shadowing the update @@ -5053,6 +5055,75 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): ) +class _CountingFakeRedis(_SharedFakeRedis): + """Counts per-key Redis round-trips so tests can pin the number of + network operations a code path issues.""" + + def __init__(self): + super().__init__() + self.get_calls: int = 0 + + async def async_get_cache(self, key, **kwargs): + self.get_calls += 1 + return await super().async_get_cache(key, **kwargs) + + +@pytest.mark.asyncio +async def test_warm_team_object_reads_issue_no_redis_ops_lit_5944(): + """ + Regression test for LIT-5944: project/team-scoped virtual-key requests + paid ~4 awaited Redis GETs per request just to re-read the team object. + + `_get_team_object_from_cache` used to consult + `proxy_logging_obj.internal_usage_cache.dual_cache` (in-memory TTL 1s, + Redis-backed) BEFORE `user_api_key_cache`. Nothing writes team objects + into that internal cache — `_cache_team_object` only DELETES the key + there — so when `user_api_key_cache` has no Redis tier the shared Redis + key stays absent forever and every team lookup in the auth hot path + (4 call sites per chat-completion request) became a guaranteed-miss + Redis round-trip, saturating the event loop at high TPS. + + Pins: once `_cache_team_object` has cached a team, repeated + `get_team_object` reads are served from `user_api_key_cache`'s in-memory + tier and issue ZERO Redis operations. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import _cache_team_object, get_team_object + + team_id = "team-lit-5944" + counting_redis = _CountingFakeRedis() + user_api_key_cache = UserApiKeyCache() + proxy_logging_obj = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache = DualCache( + redis_cache=counting_redis, + default_in_memory_ttl=1, + ) + prisma_client = MagicMock() + + await _cache_team_object( + team_id=team_id, + team_table=LiteLLM_TeamTableCachedObj(team_id=team_id, models=["model-a"]), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + for _ in range(4): + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert team_obj is not None and team_obj.models == ["model-a"] + + assert counting_redis.get_calls == 0, ( + "Warm team-object reads must be served from user_api_key_cache's " + "in-memory tier without any Redis round-trips. " + f"Got {counting_redis.get_calls} Redis GETs for 4 get_team_object calls." + ) + + @pytest.mark.asyncio async def test_cache_team_object_tolerates_cache_invalidation_failures(): """ @@ -6939,3 +7010,387 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is router = _router_with_a_group_priced_through_model_info() assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False + + +@pytest.mark.parametrize( + "user_route, expected", + [ + ("/internal-models/v1/chat/completions", True), + ("/internal-models/newly-registered-model/predict", True), + ("/internal-models-other/v1/chat/completions", False), + ("/anthropic/v1/messages", False), + ], +) +def test_team_allowed_routes_wildcard_prefix_matches_unregistered_passthrough_routes(user_route, expected): + """A `/prefix/*` entry in `team_allowed_routes` must cover every route under that prefix, so + passthrough endpoints registered after the proxy config was written are reachable without an + exact-route config change.""" + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route=user_route, + litellm_proxy_roles=LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/*"]), + ) + is expected + ) + + +def test_team_allowed_routes_exact_route_does_not_become_a_prefix_grant(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["/internal-models/model-a"]) + + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-a", litellm_proxy_roles=roles) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/internal-models/model-b", litellm_proxy_roles=roles) + is False + ) + + +def test_admin_allowed_routes_wildcard_prefix_is_honored(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(admin_allowed_routes=["/internal-models/*"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/internal-models/anything", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.PROXY_ADMIN, user_route="/other/anything", litellm_proxy_roles=roles + ) + is False + ) + + +def test_team_allowed_routes_named_route_group_still_resolves(): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.auth_checks import allowed_routes_check + + roles = LiteLLM_JWTAuth(team_allowed_routes=["openai_routes"]) + + assert ( + allowed_routes_check( + user_role=LitellmUserRoles.TEAM, user_route="/v1/chat/completions", litellm_proxy_roles=roles + ) + is True + ) + assert ( + allowed_routes_check(user_role=LitellmUserRoles.TEAM, user_route="/key/generate", litellm_proxy_roles=roles) + is False + ) + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_clears_both_membership_cache_keys(): + """A team-member budget reset (new_spend passed) must SET the spend counter to the reset + value, clear its DB-floor marker, AND invalidate both independently-keyed membership caches + (user_api_key_auth.py's admission check writes one key format, budget_reservation.py and + auth_checks.py's own get_team_membership() write the other) or a stale read keeps 429ing + after the reset. Asserted against real cache reads, not mock call args, so a change that + keeps the call but drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:user-1:team-1", value="stale-membership") + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0) + real_spend_counter_cache.in_memory_cache.set_cache( + key="spend_db_floor:spend:team_member:user-1:team-1", value=999.0 + ) + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=0.0, + ) + + assert await real_cache.async_get_cache(key="team-1_user-1") is None + assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0 + assert ( + real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") + == 0.0 + ), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up" + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_alone_without_new_spend(): + """team_member_update only changes the budget cap, not the tracked spend, so it calls + invalidate_team_member_spend_state with no new_spend. Deleting the live spend counter in that + case would force the next read to reseed from the DB's own spend column, which lags the live + counter via periodic batch writes, briefly UNDER-enforcing the raised cap against a spend + value lower than what was actually tracked (regression: PR #37971 Bugbot finding).""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership") + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0) + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + ) + + assert await real_cache.async_get_cache(key="team-1_user-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_deleting(): + """/key/{key}/reset_spend SETs its counter to the reset value rather than deleting it, so a + worker's next read reflects it directly instead of falling back through a DB reseed. A reset + caller passing new_spend must match that precedent, not merely delete the counter.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock() + real_spend_counter_cache.redis_cache = fake_redis_cache + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5 + fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60) + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set_fails(): # test-quality-ok: only observable effect is the fallback call on the same fake client + """Redis reads take priority over the local in-memory copy (get_current_spend reads Redis + first), so a failed Redis SET would otherwise leave the OLD pre-reset value authoritative + for every worker even though the reset reported success. On a failed SET, the stale Redis + entry must be deleted instead, so the next read clean-misses and reseeds from the DB.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down")) + fake_redis_cache.async_delete_cache = AsyncMock() + real_spend_counter_cache.redis_cache = fake_redis_cache + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + fake_redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1") + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_raises_503_when_both_redis_writes_fail(): + """If the Redis SET fails AND the fallback DELETE fails, the stale pre-reset counter is still + authoritative in Redis for every worker. Reporting success would silently keep 429ing the + member, so the reset must surface a 503 instead (regression: PR #37971 Greptile finding).""" + from fastapi import HTTPException + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down")) + fake_redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis still down")) + real_spend_counter_cache.redis_cache = fake_redis_cache + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ), + pytest.raises(HTTPException) as exc_info, + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + assert exc_info.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_to_remote_workers(): + """The test above only proves the handling worker's own spend counter is + cleared. A remote worker's spend counter is a separate DualCache instance; + if the reset never reaches it, that worker keeps enforcing the pre-reset + spend the moment its own Redis read for the counter fails and it falls + back to its own (now-stale) in-memory copy. Drives the actual message + published onto the invalidation channel through a second, independent + AuthCacheInvalidationSubscriber standing in for that remote worker, rather + than asserting on the publish call args.""" + from redis.asyncio import Redis + + from litellm.caching.dual_cache import DualCache + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + published: list[tuple[str, str]] = [] + + class _RecordingRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + published.append((channel, message)) + return 1 + + class _FakeRedisCache: + def __init__(self) -> None: + self.namespace = None + + def init_async_client(self) -> object: + return _RecordingRedisClient() + + local_spend_counter_cache = DualCache() + + remote_user_api_key_cache = UserApiKeyCache() + remote_spend_counter_in_memory_cache = InMemoryCache() + remote_spend_counter_in_memory_cache.set_cache("spend:team_member:user-1:team-1", 999.0) + remote_spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:team_member:user-1:team-1", 999.0) + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache + ), + patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(), + ), + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=UserApiKeyCache(), + new_spend=0.0, + ) + + def _published_message_for(cache_key: str) -> str: + matches = [message for _, message in published if json.loads(message)["cache_key"] == cache_key] + assert matches, f"{cache_key} never reached the cross-worker invalidation channel" + return matches[-1] + + remote_subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(), + user_api_key_cache=remote_user_api_key_cache, + additional_in_memory_caches=(remote_spend_counter_in_memory_cache,), + ) + for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"): + remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API + {"type": "message", "data": _published_message_for(cache_key)} + ) + + assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0 + assert ( + remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 + ), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor" + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_not_erase_the_reset(): + """The handling worker subscribes to the same invalidation channel it publishes on, so it + receives its own reset message. A delete-style broadcast would erase the post-reset counter + and floor marker the handler just wrote, reopening the stale-floor race the reset closed + (regression: PR #37971 Greptile finding). The broadcast carries the reset value as a SET, so + applying the self-delivered message must leave both keys at the post-reset value.""" + from redis.asyncio import Redis + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + published: list[tuple[str, str]] = [] + + class _RecordingRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + published.append((channel, message)) + return 1 + + class _FakeRedisCache: + def __init__(self) -> None: + self.namespace = None + + def init_async_client(self) -> object: + return _RecordingRedisClient() + + local_spend_counter_cache = DualCache() + local_user_api_key_cache = UserApiKeyCache() + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache + ), + patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(), + ), + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=local_user_api_key_cache, + new_spend=0.0, + ) + + own_subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(), + user_api_key_cache=local_user_api_key_cache, + additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,), + ) + for _, message in published: + own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API + {"type": "message", "data": message} + ) + + assert local_spend_counter_cache.in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0, ( + "the handler's self-delivered broadcast erased the post-reset spend counter" + ) + assert ( + local_spend_counter_cache.in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 + ), "the handler's self-delivered broadcast erased the post-reset floor marker, reopening the stale-floor race" diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index b0094b81112..90b3b29d919 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -701,3 +701,58 @@ async def test_auth_failure_ip_stamp_does_not_mutate_callers_request_data(): ) assert request_data == {"model": "gpt-4o"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_error,expect_traceback", + [ + pytest.param( + ProxyException( + message="Authentication Error", type=ProxyErrorTypes.auth_error, param=None, code=401 + ), + False, + id="expected_401_no_traceback", + ), + pytest.param(ValueError("unexpected internal error"), True, id="unexpected_error_keeps_traceback"), + ], +) +async def test_handle_authentication_error_traceback_only_for_unexpected_errors(auth_error, expect_traceback, caplog): + """Regression for LIT-6043: expected 4xx auth rejections must not format a + traceback via logger.exception; unexpected errors must keep it.""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + verbose_proxy_logger.propagate = True + try: + try: + raise auth_error + except (ProxyException, ValueError) as caught: + with caplog.at_level("ERROR", logger="LiteLLM Proxy"), pytest.raises(ProxyException): + await handler._handle_authentication_error( + caught, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + finally: + verbose_proxy_logger.propagate = False + + records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + assert len(records) == 1 + assert (records[0].exc_info is not None) is expect_traceback diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 468e8aabae8..7d5fc1a3544 100644 --- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest from redis.asyncio import Redis +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AUTH_CACHE_INVALIDATION_CHANNEL, AuthCacheInvalidationSubscriber, @@ -144,6 +145,37 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] +@pytest.mark.asyncio +async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None: + """ + The spend-counter half of the same cross-worker gap: a remote worker's own + spend counter can hold a stale value (its fallback path when that worker's + own Redis read for the counter fails), and only clearing user_api_key_cache + on message would leave that separate DualCache's in-memory copy untouched. + """ + cache = UserApiKeyCache() + spend_counter_in_memory_cache = InMemoryCache() + spend_counter_in_memory_cache.set_cache("spend:team_member:u-1:t-1", 999.0) + assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message("spend:team_member:u-1:t-1")]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + additional_in_memory_caches=(spend_counter_in_memory_cache,), + ) + subscriber.start() + try: + for _ in range(200): + if spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None + + @pytest.mark.asyncio async def test_subscriber_ignores_malformed_messages() -> None: cache = UserApiKeyCache() diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index 375c0d2640c..ef560bd1b7d 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -5,6 +5,7 @@ import orjson import pytest from fastapi import Request from fastapi.testclient import TestClient +from starlette.datastructures import FormData @@ -68,7 +69,7 @@ async def test_form_data_parsing(): test_data = {"name": "test_user", "message": "hello world"} # Mock the form method to return the test data as an awaitable - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "application/x-www-form-urlencoded"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -119,7 +120,7 @@ async def test_form_data_with_json_metadata(): } # Mock the form method to return the test data as an awaitable - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -160,7 +161,7 @@ async def test_form_data_with_invalid_json_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -183,7 +184,7 @@ async def test_form_data_without_metadata(): test_data = {"model": "whisper-1", "file": "audio.mp3", "language": "en"} # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "application/x-www-form-urlencoded"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -214,7 +215,7 @@ async def test_form_data_with_empty_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -249,7 +250,7 @@ async def test_form_data_with_dict_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -280,7 +281,7 @@ async def test_form_data_with_none_metadata(): } # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=test_data) + mock_request.form = AsyncMock(return_value=FormData(test_data)) mock_request.headers = {"content-type": "multipart/form-data"} mock_request.scope = {} mock_request.state._cached_headers = None @@ -495,33 +496,29 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch): @pytest.mark.asyncio async def test_get_form_data(): """ - Test that get_form_data correctly handles form data with array notation. - Tests audio transcription parameters as a specific example. + A repeated `foo[]` key is how the OpenAI SDKs send a list, so every value has to + survive. `FormData`, not a dict: a dict cannot even hold the duplicate key. """ - # Create a mock request with transcription form data mock_request = MagicMock() + mock_request.form = AsyncMock( + return_value=FormData( + [ + ("file", "file_object"), + ("model", "gpt-4o-transcribe"), + ("include[]", "logprobs"), + ("language", "en"), + ("prompt", "Transcribe this audio file"), + ("response_format", "json"), + ("stream", "false"), + ("temperature", "0.2"), + ("timestamp_granularities[]", "word"), + ("timestamp_granularities[]", "segment"), + ] + ) + ) - # Create mock form data with array notation for timestamp_granularities - mock_form_data = { - "file": "file_object", # In a real request this would be an UploadFile - "model": "gpt-4o-transcribe", - "include[]": "logprobs", # Array notation - "language": "en", - "prompt": "Transcribe this audio file", - "response_format": "json", - "stream": "false", - "temperature": "0.2", - "timestamp_granularities[]": "word", # First array item - "timestamp_granularities[]": "segment", # Second array item (would overwrite in dict, but handled by the function) - } - - # Mock the form method to return the test data - mock_request.form = AsyncMock(return_value=mock_form_data) - - # Call the function being tested result = await get_form_data(mock_request) - # Verify regular form fields are preserved assert result["file"] == "file_object" assert result["model"] == "gpt-4o-transcribe" assert result["language"] == "en" @@ -529,17 +526,8 @@ async def test_get_form_data(): assert result["response_format"] == "json" assert result["stream"] == "false" assert result["temperature"] == "0.2" - - # Verify array fields are correctly parsed - assert "include" in result - assert isinstance(result["include"], list) - assert "logprobs" in result["include"] - - assert "timestamp_granularities" in result - assert isinstance(result["timestamp_granularities"], list) - # Note: In a real MultiDict, both values would be present - # But in our mock dictionary the second value overwrites the first - assert "segment" in result["timestamp_granularities"] + assert result["include"] == ["logprobs"] + assert result["timestamp_granularities"] == ["word", "segment"] def test_get_tags_from_request_body_with_metadata_tags(): @@ -953,7 +941,7 @@ class TestReadRequestBodyNonCanonicalContentType: mock_request = MagicMock() mock_request.body = AsyncMock(return_value=orjson.dumps(payload)) - mock_request.form = AsyncMock(return_value={}) + mock_request.form = AsyncMock(return_value=FormData({})) mock_request.headers = {"content-type": content_type} mock_request.scope = {} @@ -964,7 +952,7 @@ class TestReadRequestBodyNonCanonicalContentType: @pytest.mark.asyncio async def test_real_form_post_still_parsed_as_form(self): mock_request = MagicMock() - mock_request.form = AsyncMock(return_value={"k": "v"}) + mock_request.form = AsyncMock(return_value=FormData({"k": "v"})) mock_request.body = AsyncMock(return_value=b"") mock_request.headers = {"content-type": "application/x-www-form-urlencoded"} mock_request.scope = {} @@ -1020,7 +1008,7 @@ class TestGetRequestBody: mock_request = MagicMock() mock_request.method = "POST" mock_request.headers = {"content-type": "multipart/form-data; boundary=x"} - mock_request.form = AsyncMock(return_value={"k": "v"}) + mock_request.form = AsyncMock(return_value=FormData({"k": "v"})) mock_request.scope = {} result = await get_request_body(mock_request) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 62919200d47..72c6f77a4f7 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -898,6 +898,62 @@ async def test_test_model_connection_uses_loaded_deployment_team_id_via_model_na assert passed_model_params.model_info.team_id == deployment_owner_team_id +@pytest.mark.asyncio +async def test_test_model_connection_authorizes_on_params_after_health_check_params_merge(): + """ + Regression guard for the ordering fix: health_check_params from the request + body are merged into the probe params BEFORE the authorization check, so a + caller cannot smuggle a field past auth via health_check_params. Auth is + stubbed to reject, which halts the endpoint right after it records the + params it was handed, so the outbound probe is never reached. If the merge + is moved back to after can_user_make_model_call, the marker is absent from + those params and this test fails. + """ + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelManagementAuthChecks, + ) + from litellm.types.router import Deployment + + marker = "sentinel-from-health-check-params" + mock_can_user_make_model_call = AsyncMock( + side_effect=HTTPException(status_code=403, detail="denied") + ) + + with ( + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), + patch( # test-quality-ok: proxy module global, no injection seam + "litellm.proxy.proxy_server.llm_router", None + ), + patch.object( # test-quality-ok: capturing the params handed to auth is the assertion + ModelManagementAuthChecks, + "can_user_make_model_call", + mock_can_user_make_model_call, + ), + pytest.raises(HTTPException), + ): + await health_test_model_connection( + request=MagicMock(), + mode="chat", + litellm_params={"model": "openai/gpt-4o"}, + model_info={"health_check_params": {"probe_marker": marker}}, + user_api_key_dict=UserAPIKeyAuth( + token="requester-token", + user_id="admin-user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) + + assert mock_can_user_make_model_call.called + passed_model_params = mock_can_user_make_model_call.call_args.kwargs["model_params"] + assert isinstance(passed_model_params, Deployment) + authorized_params = passed_model_params.litellm_params.model_dump() + assert authorized_params.get("probe_marker") == marker + + @pytest.mark.asyncio async def test_test_model_connection_authorized_team_admin_passes_real_auth(): """ diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 871f4b4bcd1..ca517474a5c 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -731,6 +731,318 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +async def test_track_cost_callback_defers_in_progress_background_interaction(): # test-quality-ok: writing no spend row and raising no alert is the whole observable contract of the deferral path + """ + A background=true interaction create returns in_progress with no usage + block, so its success event has a model but no standard_logging_object. + The callback must skip quietly (billing happens later via the background + poll task) instead of raising 'Cost tracking failed' and alerting. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + + kwargs = { + "call_type": "acreate_interaction", + "model": "gemini/gemini-3-flash-preview", + "litellm_call_id": "test-call-id", + "litellm_params": {}, + "stream": False, + } + in_progress_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status="in_progress", + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=in_progress_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +def _in_progress_interaction_kwargs(reservation: dict) -> dict: + return { + "call_type": "acreate_interaction", + "model": "gemini/gemini-3-flash-preview", + "litellm_call_id": "test-call-id", + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + "stream": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["in_progress", "queued"]) +async def test_track_cost_callback_keeps_reservation_open_for_in_progress_background_interaction(status): + """ + The pre-call budget reservation must stay open while a background + interaction is in flight, so concurrent creates cannot stack past the + budget; the poll task's completion event reconciles it to the actual cost. + + ``queued`` is in flight for the same reason ``in_progress`` is: it has not + reached a terminal status, so releasing its reservation here would drop the + estimate off the spend counters while the interaction is still going to run + and still going to cost money. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + in_progress_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status=status, + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=in_progress_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is False + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +@pytest.mark.asyncio +async def test_track_cost_callback_releases_reservation_for_in_progress_interaction_when_polling_disabled( + monkeypatch, +): + """ + With the poll task kill switch off nothing will ever reconcile the + reservation, so the callback must release it or the spend counters stay + pinned at the estimated cost forever. + """ + import litellm.proxy.hooks.proxy_track_cost_callback as callback_module + from litellm.types.interactions import InteractionsAPIResponse + + monkeypatch.setattr(callback_module, "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", False) + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + in_progress_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status="in_progress", + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=in_progress_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "status", + ["failed", "cancelled", "incomplete", "budget_exceeded"], +) +async def test_track_cost_callback_releases_reservation_for_unpollable_interaction(status): + """ + Only an in-progress create gets a poll task, so a create that comes back + terminal with no usage has nobody left to reconcile its reservation. The + callback must release it there and then, or the pre-call estimate stays + added to the key, user, team and org spend counters and starts refusing + traffic against budget that was never actually spent. + + None of these statuses produced output, so their missing usage is a normal + outcome rather than a cost-tracking failure, and the callback must not fire + ``failed_tracking_alert``: doing so would flood operators with false alerts + and mask real cost-tracking failures. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + terminal_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status=status, + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=terminal_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", ["completed", "requires_action"]) +async def test_track_cost_callback_alerts_when_an_interaction_that_produced_output_has_no_usage(status): + """ + ``completed`` and ``requires_action`` both mean the model produced output, + so a usage block is always expected with them. One arriving without it + means the charge for real work was lost, which is exactly what the + cost-tracking alert is for: silencing it here would let an operator's + interactions bill nothing with no signal that anything went wrong. + + The reservation still has to be released, since suppressing the alert was + never what freed it. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + usageless_response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status=status, + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=usageless_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_called_once() + + +@pytest.mark.asyncio +async def test_track_cost_callback_releases_reservation_for_interaction_without_an_id(): + """ + The scheduler also refuses a response with no id, since it has nothing to + poll for, so the callback must not defer to a poll task that will never + exist, and it must not fire ``failed_tracking_alert`` for what is a + legitimate no-usage response rather than a cost-tracking failure. + """ + from litellm.types.interactions import InteractionsAPIResponse + + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + idless_response = InteractionsAPIResponse( + id="", + model="gemini-3-flash-preview", + status="in_progress", + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=idless_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert reservation["finalized"] is True + mock_proxy_logging.failed_tracking_alert.assert_not_called() + + +@pytest.mark.asyncio +async def test_callback_handles_every_status_the_interactions_api_can_return(): + """ + Whatever status a usage-less create comes back with, exactly one of two + things has to happen to its budget reservation: the callback holds it open + for a poll task that will settle it, or it releases it on the spot. A + status that falls through both leaves the pre-call estimate pinned to the + key, user, team and org spend counters forever, refusing traffic against + budget nobody spent. + + Driven off the generated spec enum so a status Google adds later fails here + instead of quietly leaking reservations in production. + """ + from litellm.types.interactions import InteractionsAPIResponse + from litellm.types.interactions.generated import Status1 + + deferred = set() + released = set() + + for status in sorted(member.value for member in Status1): + logger = _ProxyDBLogger() + reservation = {"reserved_cost": 0.05, "entries": [], "finalized": False} + response = InteractionsAPIResponse( + id="interactions/bg-abc", + model="gemini-3-flash-preview", + status=status, + ) + + with patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + ) as mock_proxy_logging: + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=_in_progress_interaction_kwargs(reservation), + completion_response=response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + (deferred if reservation["finalized"] is False else released).add(status) + + assert deferred == {"in_progress", "queued"} + assert released == { + "completed", + "requires_action", + "failed", + "cancelled", + "incomplete", + "budget_exceeded", + } + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py index 9853ce7e1cf..135175dd29d 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_transformations.py @@ -2,7 +2,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, Member from litellm.proxy.management_endpoints.scim.scim_transformations import ( ScimTransformations, @@ -95,25 +94,17 @@ def mock_prisma_client(): class TestScimTransformations: @pytest.mark.asyncio - async def test_transform_litellm_user_to_scim_user( - self, mock_user, mock_prisma_client - ): + async def test_transform_litellm_user_to_scim_user(self, mock_user, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client # Mock the team lookup - team1 = LiteLLM_TeamTable( - team_id="team-1", team_alias="Team One", members_with_roles=[] - ) - team2 = LiteLLM_TeamTable( - team_id="team-2", team_alias="Team Two", members_with_roles=[] - ) + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[]) + team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two", members_with_roles=[]) mock_find_unique.side_effect = [team1, team2] with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - mock_user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user) assert scim_user.id == mock_user.user_id assert scim_user.userName == mock_user.user_email @@ -129,21 +120,15 @@ class TestScimTransformations: assert scim_user.groups[1].display == "Team Two" @pytest.mark.asyncio - async def test_transform_user_with_scim_metadata( - self, mock_user_with_scim_metadata, mock_prisma_client - ): + async def test_transform_user_with_scim_metadata(self, mock_user_with_scim_metadata, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client # Mock the team lookup - team1 = LiteLLM_TeamTable( - team_id="team-1", team_alias="Team One", members_with_roles=[] - ) + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[]) mock_find_unique.return_value = team1 with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - mock_user_with_scim_metadata - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user_with_scim_metadata) assert scim_user.name.givenName == "Test" assert scim_user.name.familyName == "User" @@ -160,15 +145,11 @@ class TestScimTransformations: teams=[], created_at=None, updated_at=None, - metadata={ - "scim_enterprise": {"costCenter": "CC-42", "department": "Platform"} - }, + metadata={"scim_enterprise": {"costCenter": "CC-42", "department": "Platform"}}, ) with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user) assert scim_user.enterprise_user is not None assert scim_user.enterprise_user.costCenter == "CC-42" @@ -176,9 +157,7 @@ class TestScimTransformations: assert SCIM_ENTERPRISE_USER_SCHEMA in scim_user.schemas @pytest.mark.asyncio - async def test_transform_user_with_entitlements_and_roles_metadata( - self, mock_prisma_client - ): + async def test_transform_user_with_entitlements_and_roles_metadata(self, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client mock_find_unique.return_value = None @@ -190,17 +169,13 @@ class TestScimTransformations: created_at=None, updated_at=None, metadata={ - "scim_entitlements": [ - {"value": "jira-software", "display": "Jira Software"} - ], + "scim_entitlements": [{"value": "jira-software", "display": "Jira Software"}], "scim_roles": [{"value": "engineering-admin", "primary": True}], }, ) with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user) assert scim_user.entitlements is not None assert scim_user.entitlements[0].value == "jira-software" @@ -210,9 +185,7 @@ class TestScimTransformations: assert scim_user.roles[0].primary is True @pytest.mark.asyncio - async def test_transform_user_with_malformed_directory_metadata_fails_soft( - self, mock_prisma_client - ): + async def test_transform_user_with_malformed_directory_metadata_fails_soft(self, mock_prisma_client): """Metadata is writable outside the SCIM surface; a corrupted value on one user must omit the attribute, not fail the whole directory response""" mock_client, mock_find_unique = mock_prisma_client @@ -233,9 +206,7 @@ class TestScimTransformations: ) with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user) assert scim_user.id == "user-corrupt" assert scim_user.entitlements is None @@ -244,22 +215,14 @@ class TestScimTransformations: assert SCIM_ENTERPRISE_USER_SCHEMA not in scim_user.schemas @pytest.mark.asyncio - async def test_transform_user_without_enterprise_metadata_omits_schema( - self, mock_user, mock_prisma_client - ): + async def test_transform_user_without_enterprise_metadata_omits_schema(self, mock_user, mock_prisma_client): mock_client, mock_find_unique = mock_prisma_client - team1 = LiteLLM_TeamTable( - team_id="team-1", team_alias="Team One", members_with_roles=[] - ) - team2 = LiteLLM_TeamTable( - team_id="team-2", team_alias="Team Two", members_with_roles=[] - ) + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One", members_with_roles=[]) + team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two", members_with_roles=[]) mock_find_unique.side_effect = [team1, team2] with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - mock_user - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(mock_user) assert scim_user.enterprise_user is None assert SCIM_ENTERPRISE_USER_SCHEMA not in scim_user.schemas @@ -309,36 +272,28 @@ class TestScimTransformations: assert dumped_attrs["roles"][0]["value"] == "engineering-admin" @pytest.mark.asyncio - async def test_transform_litellm_team_to_scim_group( - self, mock_team, mock_prisma_client - ): + async def test_transform_litellm_team_to_scim_group(self, mock_team, mock_prisma_client): mock_client, _ = mock_prisma_client with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( - mock_team - ) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(mock_team) assert scim_group.id == mock_team.team_id assert scim_group.displayName == mock_team.team_alias assert len(scim_group.members) == 2 - assert scim_group.members[0].value == "test@example.com" + assert scim_group.members[0].value == "user-123" assert scim_group.members[0].display == "test@example.com" - assert scim_group.members[1].value == "test2@example.com" + assert scim_group.members[1].value == "user-456" assert scim_group.members[1].display == "test2@example.com" @pytest.mark.asyncio - async def test_transform_team_marks_members_as_users( - self, mock_team, mock_prisma_client - ): + async def test_transform_team_marks_members_as_users(self, mock_team, mock_prisma_client): """A LiteLLM team only holds users, and stating the member type keeps the response from emitting a null ``type`` now that SCIMMember carries one.""" mock_client, _ = mock_prisma_client with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_group = await ScimTransformations.transform_litellm_team_to_scim_group( - mock_team - ) + scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(mock_team) assert [member.type for member in scim_group.members] == ["User", "User"] @@ -351,9 +306,7 @@ class TestScimTransformations: result = ScimTransformations._get_scim_user_name(mock_user_minimal) assert result == ScimTransformations.DEFAULT_SCIM_DISPLAY_NAME - def test_get_scim_family_name( - self, mock_user, mock_user_with_scim_metadata, mock_user_minimal - ): + def test_get_scim_family_name(self, mock_user, mock_user_with_scim_metadata, mock_user_minimal): # User with alias result = ScimTransformations._get_scim_family_name(mock_user) assert result == mock_user.user_alias @@ -366,9 +319,7 @@ class TestScimTransformations: result = ScimTransformations._get_scim_family_name(mock_user_minimal) assert result == ScimTransformations.DEFAULT_SCIM_FAMILY_NAME - def test_get_scim_given_name( - self, mock_user, mock_user_with_scim_metadata, mock_user_minimal - ): + def test_get_scim_given_name(self, mock_user, mock_user_with_scim_metadata, mock_user_minimal): # User with alias result = ScimTransformations._get_scim_given_name(mock_user) assert result == mock_user.user_alias @@ -382,14 +333,10 @@ class TestScimTransformations: assert result == ScimTransformations.DEFAULT_SCIM_NAME def test_get_scim_member_value(self): - # Member with email - member_with_email = Member( - user_id="user-123", user_email="test@example.com", role="admin" - ) + member_with_email = Member(user_id="user-123", user_email="test@example.com", role="admin") result = ScimTransformations._get_scim_member_value(member_with_email) - assert result == member_with_email.user_email + assert result == member_with_email.user_id - # Member without email should fall back to user_id member_without_email = Member(user_id="user-456", user_email=None, role="user") result = ScimTransformations._get_scim_member_value(member_without_email) assert result == member_without_email.user_id @@ -415,9 +362,7 @@ class TestScimTransformations: mock_find_unique.return_value = None with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user_with_uuid_email - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user_with_uuid_email) assert scim_user.id == user_with_uuid_email.user_id assert scim_user.emails is None or len(scim_user.emails) == 0 @@ -443,9 +388,7 @@ class TestScimTransformations: mock_find_unique.return_value = None with patch("litellm.proxy.proxy_server.prisma_client", mock_client): - scim_user = await ScimTransformations.transform_litellm_user_to_scim_user( - user_with_none_email - ) + scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user_with_none_email) assert scim_user.id == user_with_none_email.user_id assert scim_user.emails is None or len(scim_user.emails) == 0 diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py index 5f6c1a2375b..d51e3abde83 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py @@ -752,6 +752,63 @@ async def test_handle_existing_user_by_email_syncs_roster_and_dedups_teams(mocke assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"] +@pytest.mark.asyncio +async def test_handle_existing_user_by_email_without_teams_preserves_memberships(mocker): + """Adoption via POST /Users without ``groups`` must keep the user's existing teams. + + Regression: Entra manages membership exclusively through /Groups and never sends + ``groups`` on POST /Users, so the empty team list was treated as the desired + state and the adopted user was removed from every team roster and had ``teams`` + overwritten with []. + """ + existing_user = mocker.MagicMock() + existing_user.user_id = "adopted-id" + existing_user.user_email = "member@example.com" + existing_user.user_alias = "Member" + existing_user.teams = ["team-a", "team-b"] + existing_user.metadata = {} + + mock_prisma_client = mocker.MagicMock() + mock_prisma_client.db = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable = mocker.MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=existing_user) + mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={}) + + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user", + AsyncMock(return_value=None), + ) + mock_team_member_add = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) + mock_team_member_delete = mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_delete", + AsyncMock(), + ) + + new_user_request = NewUserRequest( + user_id="entra-object-id", + user_email="member@example.com", + user_alias="Member", + teams=[], + metadata={}, + auto_create_key=False, + ) + + await UserProvisionerHelpers.handle_existing_user_by_email( + prisma_client=mock_prisma_client, new_user_request=new_user_request + ) + + mock_team_member_add.assert_not_awaited() + mock_team_member_delete.assert_not_awaited() + + update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list + assert len(update_calls) == 1 + assert update_calls[0].kwargs["where"] == {"user_id": "adopted-id"} + assert update_calls[0].kwargs["data"]["teams"] == ["team-a", "team-b"] + + @pytest.mark.asyncio async def test_handle_existing_user_by_email_roster_add_failure_blocks_teams_write(mocker): """A genuine roster add failure must propagate and must not persist the teams array. @@ -872,11 +929,16 @@ async def test_handle_existing_user_by_email_roster_remove_failure_blocks_teams_ AsyncMock(side_effect=HTTPException(status_code=500, detail={"error": "No db connected"})), ) + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) + new_user_request = NewUserRequest( user_id="uid", user_email="member@example.com", user_alias="Member", - teams=[], + teams=["replacement-team"], metadata={}, auto_create_key=False, ) @@ -917,11 +979,16 @@ async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noo AsyncMock(return_value=None), ) + mocker.patch( # test-quality-ok: roster helpers are module-level, not injectable into the helper + "litellm.proxy.management_endpoints.scim.scim_v2.team_member_add", + AsyncMock(), + ) + new_user_request = NewUserRequest( user_id="uid", user_email="member@example.com", user_alias="Member", - teams=[], + teams=["replacement-team"], metadata={}, auto_create_key=False, ) @@ -933,7 +1000,7 @@ async def test_handle_existing_user_by_email_roster_remove_already_absent_is_noo mock_team_member_delete.assert_awaited_once() update_calls = mock_prisma_client.db.litellm_usertable.update.call_args_list assert len(update_calls) == 1 - assert update_calls[0].kwargs["data"]["teams"] == [] + assert update_calls[0].kwargs["data"]["teams"] == ["replacement-team"] @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 805168c84ac..3a0279ab0fa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -2,6 +2,7 @@ Unit tests for auto router management endpoints """ +from collections.abc import Mapping, Sequence from pathlib import Path from typing import Final @@ -22,11 +23,22 @@ from litellm.proxy.management_endpoints.auto_router_endpoints import ( from litellm.router import Router from litellm.types.utils import Choices, Message, ModelResponse from litellm.types.management_endpoints.auto_router_endpoints import ( + AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, ) ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test", user_id="admin") + +def _deployment(model_name: str, model: str, *, db_model: bool) -> dict[str, object]: + """One entry as `Router.model_list` holds it, for either origin.""" + return { + "model_name": model_name, + "litellm_params": {"model": model}, + "model_info": {"id": f"{model_name}-{int(db_model)}", "db_model": db_model}, + } + + TIERS = { "SIMPLE": ["cheap-model"], "MEDIUM": ["mid-model"], @@ -295,6 +307,34 @@ def test_classifier_plugin_is_not_settable_over_http(): class TestAutoRouterBenchmarks: from litellm.proxy.management_endpoints.auto_router_endpoints import _SessionAggRow + @pytest.fixture(autouse=True) + def _pin_the_router_global(self, monkeypatch: pytest.MonkeyPatch): + """Every test here reads proxy_server.llm_router, so no test may inherit a sibling's.""" + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "llm_router", None) + + @staticmethod + async def _benchmarks( + monkeypatch: pytest.MonkeyPatch, + rows: Sequence[Mapping[str, object]], + model_list: Sequence[object], + ) -> AutoRouterBenchmarksResponse: + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks + + class _DB: + async def query_raw(self, sql: str, *params: object): + return rows + + monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})()) + monkeypatch.setattr(proxy_server, "llm_router", type("R", (), {"model_list": model_list})()) + return await get_auto_router_benchmarks( + user_api_key_dict=ADMIN, + start_date="2026-07-01", + end_date="2026-08-01", + ) + ROW = _SessionAggRow( router_name="live-auto", router_type="complexity", @@ -471,6 +511,113 @@ class TestAutoRouterBenchmarks: ) assert response.groups[0].tier_turns == expected + @pytest.mark.asyncio + async def test_the_picker_lists_configured_routers_before_they_have_traffic(self, monkeypatch: pytest.MonkeyPatch): + """A router must be selectable the moment it exists, from either origin. + + `live-auto` is the only router the rollup knows about, so before this it was the only + thing the dropdown could offer. Both a config.yaml router and a DB-created one now + arrive zeroed, and neither moves the totals or duplicates the router that has traffic. + """ + response = await self._benchmarks( + monkeypatch, + rows=[self.ROW.model_dump()], + model_list=[ + _deployment("live-auto", "auto_router/complexity_router", db_model=False), + _deployment("idle-from-config", "auto_router/complexity_router", db_model=False), + _deployment("idle-from-db", "auto_router/complexity_router", db_model=True), + ], + ) + + by_name = {group.router_name: group for group in response.groups} + assert sorted(by_name) == ["idle-from-config", "idle-from-db", "live-auto"] + assert len(response.groups) == 3 + assert response.routers_in_scope == 3 + assert by_name["live-auto"].spend == 10.0 + assert response.totals.spend == 10.0 + assert response.totals.sessions == 4 + for name in ("idle-from-config", "idle-from-db"): + idle = by_name[name] + assert idle.router_type == "complexity" + assert (idle.sessions, idle.turns, idle.spend, idle.saved_spend, idle.baseline_spend) == ( + 0, + 0, + 0.0, + 0.0, + 0.0, + ) + assert (idle.saved_pct, idle.saved_per_session, idle.avg_turns_per_session) == (0.0, 0.0, 0.0) + assert (idle.cache.hit_rate_pct, idle.cache.coverage_pct) == (0.0, 0.0) + assert idle.cache.same_model.turns == idle.cache.return_to_tier.hits == 0 + assert idle.tier_turns == {} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "model, listed_as", + [ + ("auto_router/complexity_router", "complexity"), + ("auto_router/adaptive_router", "adaptive"), + ("auto_router/quality_router", "quality"), + ("auto_router/my-semantic-router", None), + ("openai/gpt-5", None), + ], + ) + async def test_only_kinds_whose_routing_the_rollup_records_are_listed( + self, model: str, listed_as: str | None, monkeypatch: pytest.MonkeyPatch + ): + """A semantic auto-router records no routing decision, so it can never own a session + row; listing it would show $0 forever even while it serves traffic.""" + response = await self._benchmarks( + monkeypatch, rows=[], model_list=[_deployment("candidate", model, db_model=True)] + ) + + assert [group.router_type for group in response.groups] == ([listed_as] if listed_as else []) + + @pytest.mark.asyncio + async def test_a_malformed_deployment_is_skipped_rather_than_failing_the_dashboard( + self, monkeypatch: pytest.MonkeyPatch + ): + response = await self._benchmarks( + monkeypatch, + rows=[self.ROW.model_dump()], + model_list=[ + "not-a-mapping", + {}, + {"model_name": "no-params"}, + {"model_name": "", "litellm_params": {"model": "auto_router/complexity_router"}}, + {"model_name": 7, "litellm_params": {"model": "auto_router/complexity_router"}}, + {"model_name": "no-model", "litellm_params": {}}, + {"model_name": "unreadable-model", "litellm_params": {"model": None}}, + ], + ) + + assert [group.router_name for group in response.groups] == ["live-auto"] + + @pytest.mark.asyncio + async def test_two_deployments_of_one_router_are_listed_once(self, monkeypatch: pytest.MonkeyPatch): + """Tagged variants share a model_name, and the picker selects by name and type.""" + response = await self._benchmarks( + monkeypatch, + rows=[], + model_list=[ + _deployment("tagged", "auto_router/complexity_router", db_model=True), + _deployment("tagged", "auto_router/complexity_router", db_model=True), + ], + ) + + assert [group.router_name for group in response.groups] == ["tagged"] + + def test_the_listed_kinds_match_the_router_types_traffic_can_record(self): + """The one reason semantic is excluded, pinned against both declarations: a kind the + rollup can record must be listable, and a kind it cannot must not be.""" + from typing import get_args, get_type_hints + + from litellm.router_utils.auto_router_model_naming import StrategyRouterKind + from litellm.types.utils import StandardLoggingRoutingDecision + + recorded = set(get_args(get_type_hints(StandardLoggingRoutingDecision)["router_type"])) + assert set(get_args(StrategyRouterKind)) - {"semantic"} == recorded + # --------------------------------------------------------------------------- # Shadow eval endpoints diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index f6d74a189bc..7f5d3eb0a14 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9,11 +9,13 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +from pydantic import ValidationError from litellm._uuid import uuid from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( + LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, LiteLLM_ModelTable, LiteLLM_OrganizationMembershipTable, @@ -27,7 +29,9 @@ from litellm.proxy._types import ( Member, ProxyErrorTypes, ProxyException, + ResetSpendRequest, TeamMemberAddRequest, + TeamMemberUpdateRequest, UpdateTeamRequest, ) from litellm.proxy.management_endpoints.team_endpoints import ( @@ -42,12 +46,15 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _transform_teams_to_deleted_records, _update_model_table, _validate_and_populate_member_user_info, + _validate_team_member_reset_spend_value, _verify_team_access, delete_team, list_available_teams, + reset_team_member_spend_fn, router, team_member_add_duplication_check, team_member_delete, + team_member_update, update_team, validate_team_org_change, ) @@ -12603,3 +12610,376 @@ async def test_invalidate_access_group_cache_deletes_the_cached_object(): "user_api_key_cache": cache, "proxy_logging_obj": logging_obj, } + + +def test_validate_team_member_reset_spend_value_rejects_non_numeric(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to="not-a-number", + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_rejects_negative(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=-1.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("reset_to", [float("nan"), float("inf"), float("-inf")]) +def test_validate_team_member_reset_spend_value_rejects_non_finite(reset_to): + """NaN and +/-inf are instances of float and compare False against every bound + below (`nan < 0`, `nan > current_spend` are both False), so an isinstance-and-range + check alone lets them through to persist as the member's spend and silently + disable every later budget comparison against it.""" + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=reset_to, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("reset_to", [True, False]) +def test_reset_spend_request_rejects_bool_reset_to(reset_to): + """bool is a subclass of int, so pydantic silently coerces True/False into 1.0/0.0 for a + ``float`` field: {"reset_to": true} would otherwise reach _validate_team_member_reset_spend_value + as an indistinguishable 1.0 and reset the member's spend instead of failing the request.""" + with pytest.raises(ValidationError): + ResetSpendRequest(reset_to=reset_to) + + +def test_validate_team_member_reset_spend_value_rejects_above_current_spend(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=20.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_rejects_above_max_budget(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=10.0, + membership=LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + spend=10.0, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=5.0), + ), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_accepts_valid_reset(): + result = _validate_team_member_reset_spend_value( + reset_to=0.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert result == 0.0 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_success(monkeypatch): + """A proxy admin resetting a stuck team member's spend must write the DB + row to reset_to AND invalidate the cached spend/membership state, or the + 429 the endpoint exists to clear keeps firing off the stale cache. + Asserted against real cache reads, not mock call args, so a change that + keeps the call but drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + mock_proxy_logging_obj = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + spend=10.0, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0), + ) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + response = await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert response["spend"] == 0.0 + assert response["previous_spend"] == 10.0 + assert response["max_budget"] == 50.0 + mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={"spend": 0.0}, + ) + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch): + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="ghost-user", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_team_not_found(monkeypatch): + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="ghost-team", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_forbidden_for_non_admin(monkeypatch): + """A caller who is neither proxy admin, org admin, nor this team's admin must be refused, + matching every other team-mutating endpoint's authorization.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user" + ), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_team_admin_cannot_reset_own_spend(monkeypatch): + """_verify_team_access authorizes a team admin over their own team with no check that the + target differs from the caller. Unchecked, that admin could target their own membership row + and repeatedly zero it right before it crosses their per-member cap, consuming the shared + team budget without the configured limit ever binding (Veria finding on PR #37971).""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-admin", user_id="team-admin-1") + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock( + return_value=LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[Member(user_id="team-admin-1", role="admin")], + ) + ), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="team-admin-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=team_admin, + ) + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_teammembership.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkeypatch): + """The self-reset guard is scoped to non-proxy-admin roles: a proxy admin resetting their + own membership spend is the platform-wide trust boundary, not a team-scoped one.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + response = await reset_team_member_spend_fn( + team_id="team-1", + user_id="admin-user", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert response["spend"] == 0.0 + + +@pytest.mark.asyncio +async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch): + """Raising a stuck member's max_budget_in_team via the documented /team/member_update + endpoint must invalidate the cached membership state, or the raised cap never reaches the + admission check and the member stays 429ing. The live spend counter itself must be left + untouched: only the cap changed, and deleting the counter would force a reseed from the + DB's own spend column, which lags the live counter via periodic batch writes, briefly + UNDER-enforcing the raised cap against a spend value lower than what was actually tracked. + Asserted against real cache reads, not mock call args, so a change that keeps the call but + drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[]) + team_info_response = { + "team_info": team_row, + "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_response), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=999999.0), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(monkeypatch): + """A role-only update carries an empty budget_patch and touches no budget state, + so the member's cached spend/membership state must be left untouched.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="still-fresh-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=1.5) + + team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[]) + team_info_response = { + "team_info": team_row, + "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_response), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership" + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index ac140abe31f..a1725005dbf 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -12,6 +12,7 @@ import httpx import pytest from fastapi import HTTPException, Request, Response from fastapi.testclient import TestClient +from starlette.datastructures import FormData import litellm @@ -35,7 +36,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( vertex_proxy_route, vllm_proxy_route, ) -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, SpecialHeaders, UserAPIKeyAuth from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials @@ -553,10 +554,9 @@ class TestVertexAIPassThroughHandler: @pytest.mark.asyncio async def test_vertex_passthrough_with_no_default_credentials(self, monkeypatch): """ - Test that when no default credentials are set, the request fails - """ - """ - Test that when passthrough credentials are set, they are correctly used in the request + With no Vertex credential matching the request, the only Authorization present + is the caller's own virtual key. It must not be forwarded to Google; the + request fails with a clean 401 instead (LIT-5997). """ from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import ( PassthroughEndpointRouter, @@ -619,31 +619,25 @@ class TestVertexAIPassThroughHandler: mock_get_token.return_value = (test_token, "") mock_auth.return_value = MagicMock() - # Call the route - try: + with pytest.raises(HTTPException) as exc_info: await vertex_proxy_route( endpoint=endpoint, request=mock_request, fastapi_response=mock_response, ) - except Exception as e: - traceback.print_exc() - print(f"Error: {e}") - # Verify create_pass_through_route was called with correct arguments - mock_create_route.assert_called_once_with( - endpoint=endpoint, - target=f"https://{test_location}-aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", - custom_headers={"authorization": f"Bearer {test_token}"}, - is_streaming_request=False, - ) + assert exc_info.value.status_code == 401 + mock_create_route.assert_not_called() @pytest.mark.asyncio async def test_async_vertex_proxy_route_api_key_auth(self): """ Critical - This is how Vertex AI JS SDK will Auth to Litellm Proxy + This is how Vertex AI JS SDK will Auth to Litellm Proxy: the virtual key + arrives in x-litellm-api-key and must reach user_api_key_auth. With no Vertex + credential configured, that virtual key must not be forwarded to Google, so + the request fails with a clean 401 (LIT-5997). """ # Mock dependencies mock_request = Mock() @@ -663,14 +657,15 @@ class TestVertexAIPassThroughHandler: return_value={"status": "success"} ) - # Call the function - result = await vertex_proxy_route( - endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", - request=mock_request, - fastapi_response=mock_response, - ) + with pytest.raises(HTTPException) as exc_info: + await vertex_proxy_route( + endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent", + request=mock_request, + fastapi_response=mock_response, + ) - # Verify user_api_key_auth was called with the correct Bearer token + assert exc_info.value.status_code == 401 + mock_pass_through.assert_not_called() mock_auth.assert_called_once() call_args = mock_auth.call_args[1] assert call_args["api_key"] == "Bearer test-key-123" @@ -1338,7 +1333,9 @@ class TestVertexAIDiscoveryPassThroughHandler: @pytest.mark.asyncio async def test_vertex_discovery_proxy_route_api_key_auth(self): """ - Test that the route correctly handles API key authentication + The virtual key arrives in x-litellm-api-key and must reach user_api_key_auth. + With no Vertex credential configured, that virtual key must not be forwarded to + Google, so the request fails with a clean 401 (LIT-5997). """ # Mock dependencies mock_request = Mock() @@ -1358,14 +1355,15 @@ class TestVertexAIDiscoveryPassThroughHandler: return_value={"status": "success"} ) - # Call the function - result = await vertex_discovery_proxy_route( - endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search", - request=mock_request, - fastapi_response=mock_response, - ) + with pytest.raises(HTTPException) as exc_info: + await vertex_discovery_proxy_route( + endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search", + request=mock_request, + fastapi_response=mock_response, + ) - # Verify user_api_key_auth was called with the correct Bearer token + assert exc_info.value.status_code == 401 + mock_pass_through.assert_not_called() mock_auth.assert_called_once() call_args = mock_auth.call_args[1] assert call_args["api_key"] == "Bearer test-key-123" @@ -1380,7 +1378,7 @@ async def test_is_streaming_request_fn(): mock_request = Mock() mock_request.method = "POST" mock_request.headers = {"content-type": "multipart/form-data"} - mock_request.form = AsyncMock(return_value={"stream": "true"}) + mock_request.form = AsyncMock(return_value=FormData({"stream": "true"})) assert await is_streaming_request_fn(mock_request) is True @@ -3312,7 +3310,11 @@ class TestVertexRawPredictStreamingClassification: "type": "http", "method": "POST", "path": f"/vertex_ai/{endpoint}", - "headers": [(b"content-type", b"application/json")], + "headers": [ + (b"content-type", b"application/json"), + (b"x-litellm-api-key", b"test-key"), + (b"authorization", b"Bearer ya29.byo-google-oauth"), + ], "query_string": b"", }, receive=receive, @@ -3445,6 +3447,327 @@ def test_is_passthrough_request_streaming_tolerates_non_object_bodies(request_bo assert is_passthrough_request_streaming(request_body) is expected +class TestVertexCredentiallessPassthroughVirtualKeyLeak: + """Regression coverage for LIT-5997. + + With no Vertex credential configured, the passthrough took the + bring-your-own-credentials branch and forwarded the whole incoming header set + to Google, including whichever header carried the caller's LiteLLM virtual key. + LiteLLM accepts that key from several headers (``Authorization``, + ``x-litellm-api-key``, ``x-goog-api-key``, ``api-key``, ``x-api-key``), and + ``x-goog-api-key`` doubles as a genuine Google credential, so any of them could + leak the proxy's own secret to an upstream provider. + + A credential-less request that carries no upstream Google credential must now + fail with a clean 401 and never reach ``create_pass_through_route``. The + proxy-only auth headers Google never consumes (``x-litellm-api-key``, + ``api-key``, ``x-api-key``) are dropped by name, and the virtual key is dropped + by value from ``Authorization`` / ``x-goog-api-key``, which may instead carry a + genuine bring-your-own Google credential that must still pass through. The + by-value strip also covers a virtual key sent in the operator-configured + ``general_settings.litellm_key_header_name``, whatever that header is named. + """ + + VKEY = "sk-litellm-victim-key" + ENDPOINT = ( + "v1/projects/my-proj/locations/us-central1/publishers/google/models/" + "gemini-2.5-flash:generateContent" + ) + + async def _run( + self, monkeypatch, headers: list[tuple[bytes, bytes]] + ) -> tuple[HTTPException | None, dict | None]: + from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import ( + PassthroughEndpointRouter, + ) + + async def receive(): + return {"type": "http.request", "body": b"{}", "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": f"/vertex_ai/{self.ENDPOINT}", + "headers": headers, + "query_string": b"", + }, + receive=receive, + ) + + captured: dict = {} + + def fake_create_pass_through_route(**kwargs): + captured.update(kwargs) + return AsyncMock(return_value={"status": "success"}) + + mock_handler = Mock() + mock_handler.get_default_base_target_url.return_value = "https://us-central1-aiplatform.googleapis.com/" + + module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints" + monkeypatch.setattr(f"{module}.passthrough_endpoint_router", PassthroughEndpointRouter()) + raised: HTTPException | None = None + with ( + mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route), + mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=UserAPIKeyAuth(token="hashed"))), + mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler), + ): + try: + await vertex_proxy_route( + endpoint=self.ENDPOINT, + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(token="hashed"), + ) + except HTTPException as exc: + raised = exc + + return raised, (captured.get("custom_headers") if captured else None) + + @pytest.mark.asyncio + async def test_authorization_bearer_virtual_key_is_rejected_not_forwarded(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [(b"authorization", f"Bearer {self.VKEY}".encode()), (b"content-type", b"application/json")], + ) + assert forwarded is None, "credential-less request must never reach the upstream forwarder" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_x_litellm_api_key_virtual_key_is_rejected_not_forwarded(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [(b"x-litellm-api-key", self.VKEY.encode()), (b"content-type", b"application/json")], + ) + assert forwarded is None, "credential-less request must never reach the upstream forwarder" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_x_goog_api_key_carrying_virtual_key_is_rejected_not_forwarded(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-litellm-api-key", self.VKEY.encode()), + (b"x-goog-api-key", self.VKEY.encode()), + (b"content-type", b"application/json"), + ], + ) + assert forwarded is None, "the virtual key in x-goog-api-key must not satisfy the gate nor be forwarded" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_virtual_key_authenticated_solely_via_x_goog_api_key_is_rejected(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-goog-api-key", self.VKEY.encode()), + (b"content-type", b"application/json"), + ], + ) + assert forwarded is None, "a virtual key that authenticated via x-goog-api-key must be stripped, not forwarded" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_byo_google_oauth_token_still_forwards_without_virtual_key(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-litellm-api-key", self.VKEY.encode()), + (b"authorization", b"Bearer ya29.google-oauth-token"), + (b"content-type", b"application/json"), + ], + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("authorization") == "Bearer ya29.google-oauth-token" + assert "x-litellm-api-key" not in forwarded + assert self.VKEY not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) + + @pytest.mark.asyncio + @pytest.mark.parametrize("scheme", ["Bearer", "bearer", "Basic"]) + async def test_virtual_key_echoed_in_authorization_with_any_scheme_is_stripped(self, monkeypatch, scheme): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-litellm-api-key", self.VKEY.encode()), + (b"authorization", f"{scheme} {self.VKEY}".encode()), + (b"content-type", b"application/json"), + ], + ) + assert forwarded is None, f"a virtual key echoed as '{scheme} ' in Authorization must be stripped, not forwarded" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_byo_x_goog_api_key_still_forwards_without_virtual_key(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-litellm-api-key", self.VKEY.encode()), + (b"x-goog-api-key", b"AIza-google-api-key"), + (b"content-type", b"application/json"), + ], + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("x-goog-api-key") == "AIza-google-api-key" + assert "x-litellm-api-key" not in forwarded + assert self.VKEY not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) + + @pytest.mark.asyncio + async def test_alternate_proxy_auth_headers_are_never_forwarded_to_google(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-litellm-api-key", self.VKEY.encode()), + (b"authorization", b"Bearer ya29.google-oauth-token"), + (b"api-key", b"azure-style-caller-secret"), + (b"x-api-key", b"anthropic-style-caller-secret"), + (b"content-type", b"application/json"), + ], + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("authorization") == "Bearer ya29.google-oauth-token" + assert "api-key" not in forwarded + assert "x-api-key" not in forwarded + assert "x-litellm-api-key" not in forwarded + forwarded_blob = " ".join(f"{name}:{value}" for name, value in forwarded.items()) + assert self.VKEY not in forwarded_blob + assert "azure-style-caller-secret" not in forwarded_blob + assert "anthropic-style-caller-secret" not in forwarded_blob + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "credential_header", + sorted( + SpecialHeaders.litellm_credential_header_names() + - {"authorization", "x-goog-api-key", "x-litellm-api-key"} + ), + ) + async def test_every_non_google_credential_header_is_dropped_by_name(self, monkeypatch, credential_header): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-litellm-api-key", self.VKEY.encode()), + (b"x-goog-api-key", b"AIza-real-google-api-key"), + (credential_header.encode(), b"some-distinct-caller-secret-value"), + (b"content-type", b"application/json"), + ], + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("x-goog-api-key") == "AIza-real-google-api-key" + assert credential_header not in forwarded + assert "x-litellm-api-key" not in forwarded + forwarded_blob = " ".join(f"{name}:{value}" for name, value in forwarded.items()) + assert self.VKEY not in forwarded_blob + assert "some-distinct-caller-secret-value" not in forwarded_blob + + @pytest.mark.asyncio + async def test_virtual_key_in_operator_configured_header_is_stripped(self, monkeypatch): + with mock.patch.dict( # test-quality-ok: general_settings is the real proxy config surface for litellm_key_header_name; no injection seam exists on this route + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + ): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-company-key", f"Bearer {self.VKEY}".encode()), + (b"x-goog-api-key", b"AIza-real-google-api-key"), + (b"content-type", b"application/json"), + ], + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("x-goog-api-key") == "AIza-real-google-api-key" + assert "x-company-key" not in forwarded + assert self.VKEY not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) + + @pytest.mark.asyncio + async def test_virtual_key_in_operator_configured_header_alone_is_rejected(self, monkeypatch): + with mock.patch.dict( # test-quality-ok: general_settings is the real proxy config surface for litellm_key_header_name; no injection seam exists on this route + "litellm.proxy.proxy_server.general_settings", + {"litellm_key_header_name": "x-company-key"}, + ): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-company-key", f"Bearer {self.VKEY}".encode()), + (b"content-type", b"application/json"), + ], + ) + assert forwarded is None, "a virtual key in the custom auth header must not satisfy the gate nor be forwarded" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_virtual_key_in_pass_through_configured_header_is_dropped_and_rejected(self, monkeypatch): + with mock.patch.dict( # test-quality-ok: general_settings is the real proxy config surface for pass_through_endpoints; no injection seam exists on this route + "litellm.proxy.proxy_server.general_settings", + {"pass_through_endpoints": [{"headers": {"litellm_user_api_key": "x-company-key"}}]}, + ): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"x-company-key", f"Bearer {self.VKEY}".encode()), + (b"content-type", b"application/json"), + ], + ) + assert forwarded is None, "a virtual key in the pass-through key header must be dropped, not forwarded" + assert raised is not None and raised.status_code == 401 + + @pytest.mark.asyncio + async def test_authenticated_authorization_is_stripped_over_a_lower_precedence_pass_through_header(self, monkeypatch): + with mock.patch.dict( # test-quality-ok: general_settings is the real proxy config surface for pass_through_endpoints; no injection seam exists on this route + "litellm.proxy.proxy_server.general_settings", + {"pass_through_endpoints": [{"headers": {"litellm_user_api_key": "x-company-key"}}]}, + ): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"authorization", f"Bearer {self.VKEY}".encode()), + (b"x-company-key", b"sk-decoy-lower-precedence-value"), + (b"x-goog-api-key", b"AIza-real-google-api-key"), + (b"content-type", b"application/json"), + ], + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("x-goog-api-key") == "AIza-real-google-api-key" + assert "authorization" not in forwarded, "Authorization authenticated (higher precedence) so its key must be stripped" + assert "x-company-key" not in forwarded + assert self.VKEY not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) + + @pytest.mark.asyncio + async def test_virtual_key_in_mapped_route_litellm_user_api_key_header_is_stripped(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"litellm_user_api_key", self.VKEY.encode()), + (b"authorization", b"Bearer ya29.byo-google-oauth"), + (b"x-goog-api-key", b"AIza-real-google-api-key"), + (b"content-type", b"application/json"), + ], + ) + assert raised is None + assert forwarded is not None + assert forwarded.get("x-goog-api-key") == "AIza-real-google-api-key" + assert forwarded.get("authorization") == "Bearer ya29.byo-google-oauth" + assert "litellm_user_api_key" not in forwarded + assert self.VKEY not in " ".join(f"{name}:{value}" for name, value in forwarded.items()) + + @pytest.mark.asyncio + async def test_virtual_key_in_mapped_route_litellm_user_api_key_header_alone_is_rejected(self, monkeypatch): + raised, forwarded = await self._run( + monkeypatch, + [ + (b"litellm_user_api_key", self.VKEY.encode()), + (b"content-type", b"application/json"), + ], + ) + assert forwarded is None, "a virtual key in the mapped-route litellm_user_api_key header must be dropped, not forwarded" + assert raised is not None and raised.status_code == 401 + + class TestGetAzureAISearchIndexFromEndpoint: """The operable index is only the segment right after ``indexes``. @@ -4055,3 +4378,95 @@ class TestVertexAILiveWebsocketPassthrough: assert "use_in_pass_through" in close_kwargs["reason"] assert "default_vertex_config" in close_kwargs["reason"] assert len(close_kwargs["reason"].encode("utf-8")) <= 123 + + +class TestPassthroughRouterModelBudgetReservation: + """ + Router-model passthrough on /vllm and /azure must thread the calling key's + metadata into ``allm_passthrough_route``. Without ``user_api_key`` the spend + is attributed to nobody, and without ``user_api_key_budget_reservation`` the + pre-call reservation is never released, so the shared spend counter drifts up + until the key falsely trips a 429 BudgetExceededError (LIT-5470). + """ + + def _key_with_reservation(self) -> UserAPIKeyAuth: + reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:hashed-token", "reserved_cost": 0.5}], + } + return UserAPIKeyAuth( + api_key="hashed-token", + user_id="u1", + team_id="t1", + budget_reservation=reservation, + agent_id="agent-xyz", + end_user_max_budget=42.0, + ) + + def _request(self) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "POST" + request.headers = {"content-type": "application/json"} + request.query_params = {} + return request + + def _install_recording_router(self, monkeypatch, body: dict) -> list[dict]: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + captured: list[dict] = [] + + class RecordingRouter: + async def allm_passthrough_route(self, **kwargs): + captured.append(kwargs) + return httpx.Response(200, json={"ok": True}) + + async def fake_get_request_body(_request): + return body + + monkeypatch.setattr(proxy_server, "llm_router", RecordingRouter()) + monkeypatch.setattr(ep, "get_request_body", fake_get_request_body) + monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True) + return captured + + def _assert_metadata_carries_attribution(self, captured: list[dict], user_api_key_dict: UserAPIKeyAuth) -> None: + assert len(captured) == 1, "the router-model branch must dispatch exactly once" + assert captured[0].get("metadata") is None, ( + "attribution must ride the litellm_metadata bucket the router canonicalizes on; " + "the plain metadata bucket is dropped for every non-user_api_key field" + ) + litellm_metadata = captured[0]["litellm_metadata"] + assert litellm_metadata["user_api_key"] == user_api_key_dict.api_key + assert litellm_metadata["user_api_key_budget_reservation"] is user_api_key_dict.budget_reservation + assert litellm_metadata["user_api_key_user_id"] == user_api_key_dict.user_id + assert litellm_metadata["user_api_key_team_id"] == user_api_key_dict.team_id + assert litellm_metadata["agent_id"] == user_api_key_dict.agent_id + assert litellm_metadata["user_api_end_user_max_budget"] == user_api_key_dict.end_user_max_budget + + @pytest.mark.asyncio + async def test_vllm_router_model_threads_key_metadata(self, monkeypatch): + user_api_key_dict = self._key_with_reservation() + captured = self._install_recording_router(monkeypatch, {"model": "router-model", "stream": False}) + + await vllm_proxy_route( + endpoint="/chat/completions", + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=user_api_key_dict, + ) + + self._assert_metadata_carries_attribution(captured, user_api_key_dict) + + @pytest.mark.asyncio + async def test_azure_router_model_threads_key_metadata(self, monkeypatch): + user_api_key_dict = self._key_with_reservation() + captured = self._install_recording_router(monkeypatch, {"model": "gpt-5", "stream": False}) + + await azure_proxy_route( + endpoint="openai/deployments/gpt-5/chat/completions", + request=self._request(), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=user_api_key_dict, + ) + + self._assert_metadata_carries_attribution(captured, user_api_key_dict) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 25d176e48bb..99a84d43c9b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -186,6 +186,43 @@ async def test_make_multipart_http_request_forwards_repeated_fields(): assert call_args["data"] == {"other_parameter": ["xxx", "yyy"]} +@pytest.mark.asyncio +async def test_make_multipart_http_request_fileless_form_stays_multipart(): + """ + Regression for #36493: a multipart form with no file parts was forwarded + through httpx's ``data=`` alone, which downgrades the request to + application/x-www-form-urlencoded. Every field must go through ``files`` + as a ``(field_name, (None, value))`` tuple so httpx keeps the + multipart/form-data encoding the client sent. + """ + request = MagicMock(spec=Request) + request.method = "POST" + form_data = FormData([("prompt", "a cat surfing"), ("model", "sora-2"), ("seconds", "4")]) + request.form = AsyncMock(return_value=form_data) + + mock_response = MagicMock() + mock_response.status_code = 200 + async_client = MagicMock() + async_client.request = AsyncMock(return_value=mock_response) + + await HttpPassThroughEndpointHelpers.make_multipart_http_request( + request=request, + async_client=async_client, + url=httpx.URL("http://test.com"), + headers={}, + requested_query_params=None, + ) + + call_args = async_client.request.call_args[1] + + assert call_args["files"] == ( + ("prompt", (None, "a cat surfing")), + ("model", (None, "sora-2")), + ("seconds", (None, "4")), + ) + assert call_args["data"] is None + + @pytest.mark.asyncio async def test_make_multipart_http_request_removes_content_type_header(): """ diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py index b08de04e801..abbf6892a98 100644 --- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py @@ -322,3 +322,92 @@ def test_rag_query_stream_returns_event_stream(client_internal_user): assert response.headers.get("content-type", "").startswith("text/event-stream") assert '"object":"chat.completion.chunk"' in response.text assert "data: [DONE]" in response.text + + +EICAR = r"X5O!P%@AP[4\PZX54(P^)7CC)7}$EICAR-STANDARD-ANTIVIRUS-TEST-FILE!$H+H*" +INGEST_REQUEST = '{"ingest_options":{"vector_store":{"custom_llm_provider":"openai"}}}' + + +def _multipart_ingest_request(*, filename: str, content: bytes, content_type: str): + from starlette.requests import Request + + boundary = "litellmuploadtestboundary" + head = ( + f"--{boundary}\r\n" + f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n' + f"Content-Type: {content_type}\r\n\r\n" + ).encode() + tail = ( + f"\r\n--{boundary}\r\n" + f'Content-Disposition: form-data; name="request"\r\n\r\n' + f"{INGEST_REQUEST}\r\n" + f"--{boundary}--\r\n" + ).encode() + body = head + content + tail + scope = { + "type": "http", + "method": "POST", + "path": "/v1/rag/ingest", + "headers": [ + (b"content-type", f"multipart/form-data; boundary={boundary}".encode()), + (b"content-length", str(len(body)).encode()), + ], + "state": {}, + } + + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + return Request(scope, receive) + + +class TestVectorStoreUploadControls: + """End-to-end enforcement of pentest M4 upload controls on /v1/rag/ingest.""" + + def test_eicar_upload_blocked_by_malware_scanner(self, client_internal_user): + response = client_internal_user.post( + "/v1/rag/ingest", + files={"file": ("clean_name.txt", io.BytesIO(EICAR.encode()), "text/plain")}, + data={"request": INGEST_REQUEST}, + ) + assert response.status_code == 400, response.text + assert response.json()["detail"]["reason"] == "malware_detected" + + def test_executable_upload_rejected(self, client_internal_user): + elf = b"\x7fELF\x02\x01\x01\x00" + b"\x00" * 40 + response = client_internal_user.post( + "/v1/rag/ingest", + files={"file": ("doc.txt", io.BytesIO(elf), "text/plain")}, + data={"request": INGEST_REQUEST}, + ) + assert response.status_code == 400, response.text + assert response.json()["detail"]["reason"] == "executable_not_allowed" + + def test_zip_archive_upload_rejected(self, client_internal_user): + response = client_internal_user.post( + "/v1/rag/ingest", + files={"file": ("doc.pdf", io.BytesIO(b"PK\x03\x04\x14\x00\x00\x00payload"), "application/pdf")}, + data={"request": INGEST_REQUEST}, + ) + assert response.status_code == 400, response.text + assert response.json()["detail"]["reason"] == "archive_not_allowed" + + async def test_clean_text_upload_gets_server_generated_filename(self): + from litellm.proxy.rag_endpoints.endpoints import parse_rag_ingest_request + from litellm.proxy.rag_endpoints.upload_security import EicarTestMalwareScanner + + request = _multipart_ingest_request( + filename="../../etc/passwd", + content=b"benign document text\n", + content_type="text/plain", + ) + _options, file_data, _url, _file_id = await parse_rag_ingest_request( + request, scanner=EicarTestMalwareScanner() + ) + assert file_data is not None + server_filename, content_bytes, secured_content_type = file_data + assert server_filename != "../../etc/passwd" + assert "/" not in server_filename and "\\" not in server_filename + assert server_filename.endswith(".txt") + assert secured_content_type == "text/plain" + assert content_bytes == b"benign document text\n" diff --git a/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py b/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py new file mode 100644 index 00000000000..84375080904 --- /dev/null +++ b/tests/test_litellm/proxy/rag_endpoints/test_upload_security.py @@ -0,0 +1,176 @@ +"""Unit tests for vector-store upload security controls. + +These pin the pentest M4 remediation: an allowlist enforced by real content +inspection (not extension/mime trust), a size cap, archive and executable +rejection, server-generated filenames, safe download headers, and a +dependency-injected malware scanner validated with the EICAR test file. +""" + +from dataclasses import dataclass + +import pytest + +from litellm.proxy.rag_endpoints.upload_security import ( + EICAR_TEST_SIGNATURE, + DetectedFormat, + EicarTestMalwareScanner, + RejectedUpload, + RejectionReason, + ScanResult, + ScanVerdict, + SecuredUpload, + generate_safe_filename, + inspect_content, + safe_download_headers, + validate_upload, +) + + +@dataclass(frozen=True) +class _StubScanner: + result: ScanResult + + def scan(self, content: bytes) -> ScanResult: + return self.result + + +_CLEAN_SCANNER = _StubScanner(ScanResult(ScanVerdict.CLEAN)) +_INFECTED_SCANNER = _StubScanner(ScanResult(ScanVerdict.INFECTED, signature="Test.Sig")) +_ERROR_SCANNER = _StubScanner(ScanResult(ScanVerdict.ERROR)) + +_PDF_BYTES = b"%PDF-1.7\n1 0 obj<<>>endobj\n" +_TEXT_BYTES = "the quick brown fox\n".encode("utf-8") +_ELF_BYTES = b"\x7fELF\x02\x01\x01\x00" + b"\x00" * 32 +_PE_BYTES = b"MZ\x90\x00\x03\x00\x00\x00\x04\x00\x00\x00" +_ZIP_BYTES = b"PK\x03\x04\x14\x00\x00\x00" +_GZIP_BYTES = b"\x1f\x8b\x08\x00\x00\x00\x00\x00" +_SHEBANG_BYTES = b"#!/bin/bash\nrm -rf /\n" + + +def _tar_bytes() -> bytes: + header = bytearray(512) + header[257:262] = b"ustar" + return bytes(header) + + +def _expect_rejected(content: bytes, reason: RejectionReason, *, max_size_bytes: int = 512 * 1024 * 1024) -> None: + result = validate_upload(content=content, scanner=_CLEAN_SCANNER, max_size_bytes=max_size_bytes) + assert isinstance(result, RejectedUpload), f"expected rejection, got {result!r}" + assert result.reason is reason, f"expected {reason}, got {result.reason}" + + +def test_empty_file_rejected(): + _expect_rejected(b"", RejectionReason.EMPTY_FILE) + + +def test_oversized_file_rejected(): + _expect_rejected(b"%PDF-" + b"a" * 100, RejectionReason.FILE_TOO_LARGE, max_size_bytes=10) + + +def test_zip_archive_rejected(): + _expect_rejected(_ZIP_BYTES, RejectionReason.ARCHIVE_NOT_ALLOWED) + + +def test_gzip_archive_rejected(): + _expect_rejected(_GZIP_BYTES, RejectionReason.ARCHIVE_NOT_ALLOWED) + + +def test_tar_archive_rejected(): + _expect_rejected(_tar_bytes(), RejectionReason.ARCHIVE_NOT_ALLOWED) + + +def test_elf_executable_rejected(): + _expect_rejected(_ELF_BYTES, RejectionReason.EXECUTABLE_NOT_ALLOWED) + + +def test_windows_pe_executable_rejected(): + _expect_rejected(_PE_BYTES, RejectionReason.EXECUTABLE_NOT_ALLOWED) + + +def test_shebang_script_rejected(): + _expect_rejected(_SHEBANG_BYTES, RejectionReason.EXECUTABLE_NOT_ALLOWED) + + +def test_unknown_binary_rejected(): + _expect_rejected(b"\x89\x01\x02\x00\xff\xfe garbage", RejectionReason.UNSUPPORTED_FORMAT) + + +def test_pdf_accepted_with_server_filename_and_content_type(): + result = validate_upload(content=_PDF_BYTES, scanner=_CLEAN_SCANNER) + assert isinstance(result, SecuredUpload) + assert result.detected_format is DetectedFormat.PDF + assert result.content_type == "application/pdf" + assert result.safe_filename.endswith(".pdf") + assert result.size_bytes == len(_PDF_BYTES) + + +def test_utf8_text_accepted(): + result = validate_upload(content=_TEXT_BYTES, scanner=_CLEAN_SCANNER) + assert isinstance(result, SecuredUpload) + assert result.detected_format is DetectedFormat.TEXT + assert result.content_type == "text/plain" + assert result.safe_filename.endswith(".txt") + + +def test_inspect_content_classifies_directly(): + from litellm.proxy.rag_endpoints.upload_security import AllowedContent, DisallowedContent, DisallowedKind + + assert inspect_content(_PDF_BYTES) == AllowedContent(DetectedFormat.PDF) + assert inspect_content(_TEXT_BYTES) == AllowedContent(DetectedFormat.TEXT) + assert inspect_content(_ZIP_BYTES) == DisallowedContent(DisallowedKind.ARCHIVE) + assert inspect_content(_ELF_BYTES) == DisallowedContent(DisallowedKind.EXECUTABLE) + + +def test_server_generated_filenames_are_unique_and_ignore_client_name(): + first = generate_safe_filename(DetectedFormat.PDF) + second = generate_safe_filename(DetectedFormat.PDF) + assert first != second + assert first.endswith(".pdf") + assert "/" not in first and "\\" not in first + + +def test_malware_hook_blocks_infected_clean_format(): + result = validate_upload(content=_TEXT_BYTES, scanner=_INFECTED_SCANNER) + assert isinstance(result, RejectedUpload) + assert result.reason is RejectionReason.MALWARE_DETECTED + assert "Test.Sig" in result.message + + +def test_malware_scan_error_fails_closed(): + result = validate_upload(content=_TEXT_BYTES, scanner=_ERROR_SCANNER) + assert isinstance(result, RejectedUpload) + assert result.reason is RejectionReason.MALWARE_SCAN_ERROR + + +def test_injected_clean_scanner_allows_valid_file(): + result = validate_upload(content=_TEXT_BYTES, scanner=_CLEAN_SCANNER) + assert isinstance(result, SecuredUpload) + + +def test_eicar_default_scanner_flags_only_eicar(): + scanner = EicarTestMalwareScanner() + assert scanner.scan(EICAR_TEST_SIGNATURE).verdict is ScanVerdict.INFECTED + assert scanner.scan(b"totally benign text").verdict is ScanVerdict.CLEAN + + +def test_eicar_upload_passes_format_but_blocked_by_scanner(): + """EICAR is valid ASCII text, so only the malware hook can stop it.""" + format_only = validate_upload(content=EICAR_TEST_SIGNATURE, scanner=_CLEAN_SCANNER) + assert isinstance(format_only, SecuredUpload) + + scanned = validate_upload(content=EICAR_TEST_SIGNATURE, scanner=EicarTestMalwareScanner()) + assert isinstance(scanned, RejectedUpload) + assert scanned.reason is RejectionReason.MALWARE_DETECTED + + +def test_safe_download_headers_force_attachment_and_nosniff(): + headers = safe_download_headers("file_abc123") + assert headers["Content-Disposition"] == 'attachment; filename="file_abc123"' + assert headers["X-Content-Type-Options"] == "nosniff" + + +@pytest.mark.parametrize("hostile", ['a"; drop', "a\r\nSet-Cookie: x=1", "../../etc/passwd", ""]) +def test_safe_download_headers_sanitize_injection(hostile): + disposition = safe_download_headers(hostile)["Content-Disposition"] + assert "\r" not in disposition and "\n" not in disposition + assert disposition.count('"') == 2 diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 9177944df2d..791d64c6428 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1353,6 +1353,8 @@ class TestParseCursorModelVariant: ("claude-opus-5-fast", "claude-opus-5", None), ("gpt-5.6-sol", "gpt-5.6-sol", None), ("foo-thinking-ultra-fast", "foo-thinking-ultra", None), + ("gpt-5.6-thinking-max", "gpt-5.6", "max"), + ("foo-thinking-mega-fast", "foo-thinking-mega", None), ("-thinking-high", "-thinking-high", None), ], ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 6c8e641642b..843e1d1296f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3,11 +3,10 @@ import datetime import json from datetime import timezone from typing import Any, Final, cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest - - -from unittest.mock import AsyncMock, MagicMock, patch +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -3500,6 +3499,50 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["custom_llm_provider"] == "" +class _ModelRouterSpendLogKwargs(TypedDict): + model: ReadOnly[str] + litellm_params: ReadOnly[dict[str, dict[str, str]]] + standard_logging_object: ReadOnly[StandardLoggingPayload] + + +def _model_router_spend_log_kwargs(slp_model: str | None) -> _ModelRouterSpendLogKwargs: + standard_logging_payload: Final = cast( + StandardLoggingPayload, + { + "model": slp_model, + "metadata": {}, + "model_map_information": StandardLoggingModelInformation( + model_map_key="azure_ai/model_router", model_map_value=None + ), + }, + ) + return { + "model": "azure_ai/model_router/model-router", + "litellm_params": {"metadata": {"user_api_key": "sk-test-key"}}, + "standard_logging_object": standard_logging_payload, + } + + +def test_get_logging_payload_uses_standard_logging_payload_model(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model="azure_ai/gpt-5-mini"), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/gpt-5-mini" + + +def test_get_logging_payload_falls_back_to_kwargs_model_when_slp_model_missing(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model=None), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/model_router/model-router" + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal(): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58714a5e319..b595c44d2ce 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2294,6 +2294,54 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == actual_model_used assert response_obj.model != requested_model + def test_override_model_preserves_model_router_model_for_alias_without_router_in_name( + self, + ): + """ + The client sends a model group alias, which carries no model_router/ prefix, so the + name check alone only fires when the operator happened to put "model-router" in the + alias. With the stamp on the response the actual model survives whatever it is named. + """ + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + + requested_model = "smart-pick" + actual_model_used = "azure_ai/grok-4-1-fast-reasoning" + + response_obj = MagicMock() + response_obj.model = actual_model_used + response_obj._hidden_params = { + "additional_headers": {}, + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: actual_model_used, + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + assert response_obj.model == actual_model_used + + def test_override_model_still_restamps_non_router_alias_without_stamp(self): + """ + Control for the test above: absent the stamp, an ordinary deployment keeps being + restamped to the requested model, so the stamp is doing the work rather than the + preserve branch having gone unconditional. + """ + requested_model = "smart-pick" + + response_obj = MagicMock() + response_obj.model = "azure_ai/grok-4-1-fast-reasoning" + response_obj._hidden_params = {"additional_headers": {}} + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + assert response_obj.model == requested_model + def test_override_model_uses_winning_model_for_fastest_response(self): """ Test that when fastest_response batch completion is used with a @@ -7295,3 +7343,31 @@ class TestRouterModelNameOnNonStreamingResponse: ) assert response[ROUTER_MODEL_NAME_RESPONSE_FIELD] == "deep-model" + + +@pytest.mark.parametrize( + "exc,expect_traceback", + [ + pytest.param(HTTPException(status_code=400, detail="Invalid model name passed in"), False, id="expected_400"), + pytest.param(ValueError("unexpected internal error"), True, id="unexpected_error"), + ], +) +def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_traceback, caplog): + """Regression for LIT-6043: expected 4xx errors log without formatting a + traceback; unexpected errors keep logger.exception behavior.""" + from litellm._logging import verbose_proxy_logger + from litellm.proxy.common_request_processing import _log_llm_api_exception + + verbose_proxy_logger.propagate = True + try: + with caplog.at_level("ERROR", logger="LiteLLM Proxy"): + try: + raise exc + except Exception as raised: + _log_llm_api_exception(raised) + finally: + verbose_proxy_logger.propagate = False + + records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()] + assert len(records) == 1 + assert (records[0].exc_info is not None) is expect_traceback diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index 5a606d5f74e..642e4bb8e11 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -1,11 +1,15 @@ +import json +import logging from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx +import litellm from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers from litellm.proxy import health_check as hc_module from litellm.proxy.health_check import ( - _is_semantic_auto_router_deployment, + _is_strategy_router_deployment, _resolve_health_check_max_tokens, _resolve_health_check_mode, _update_litellm_params_for_health_check, @@ -495,33 +499,22 @@ def test_autodetected_embedding_skips_reasoning_effort(): assert "max_tokens" not in updated -# --------------------------------------------------------------------------- -# auto_router (semantic router) deployments must be skipped by health checks. -# -# These are meta-routers that select among real LLM deployments at request -# time. They have no LLM endpoint to probe. Before this fix, the health check -# passed model="auto_router/router_1" to get_llm_provider(), which raised -# BadRequestError: "Unmapped LLM provider for this endpoint" because -# auto_router is not a real LLM provider. -# --------------------------------------------------------------------------- - - @pytest.mark.parametrize( "model, expected", [ ("auto_router/router_1", True), ("auto_router/my_router", True), - ("auto_router/complexity_router", False), - ("auto_router/adaptive_router", False), - ("auto_router/quality_router", False), - ("auto_router/adaptive_router/subpath", False), + ("auto_router/complexity_router", True), + ("auto_router/adaptive_router", True), + ("auto_router/quality_router", True), + ("auto_router/adaptive_router/subpath", True), ("gpt-4", False), ("openai/gpt-4", False), ("bedrock/claude", False), ], ) -def test_is_semantic_auto_router_deployment(model, expected): - assert _is_semantic_auto_router_deployment({"model": model}) == expected +def test_is_strategy_router_deployment(model, expected): + assert _is_strategy_router_deployment({"model": model}) == expected @pytest.mark.asyncio @@ -543,3 +536,151 @@ async def test_run_model_health_check_skips_auto_router_deployment(): fake_ahealth_check.assert_not_called() assert result == {} + + +def test_health_check_params_merge_into_probe_params(): + """health_check_params reach the probe request for the deployment that declares them.""" + media_source = {"s3Location": {"uri": "s3://my-bucket/clip.mp4"}} + + updated = _update_litellm_params_for_health_check( + {"mode": "chat", "health_check_params": {"mediaSource": media_source}}, + {"model": "bedrock/us.twelvelabs.pegasus-1-2-v1:0"}, + ) + + assert updated["mediaSource"] == media_source + assert updated["model"] == "us.twelvelabs.pegasus-1-2-v1:0" + assert updated["custom_llm_provider"] == "bedrock" + + +def test_health_check_params_lose_to_dedicated_health_check_knobs(): + """The dedicated knobs are applied after the merge, so they win on conflict.""" + model_info = { + "mode": "chat", + "health_check_params": { + "max_tokens": 4096, + "model": "openai/expensive-model", + "messages": [{"role": "user", "content": "from health_check_params"}], + "reasoning_effort": "high", + }, + "health_check_max_tokens": 5, + "health_check_model": "openai/cheap-model", + "health_check_reasoning_effort": "none", + } + + updated = _update_litellm_params_for_health_check(model_info, {"model": "openai/dummy"}) + + assert updated["max_tokens"] == 5 + assert updated["model"] == "openai/cheap-model" + assert updated["reasoning_effort"] == "none" + assert updated["messages"] != model_info["health_check_params"]["messages"] + + +def test_health_check_params_lose_to_the_audio_speech_voice_knob(): + """health_check_voice still wins for audio_speech deployments.""" + updated = _update_litellm_params_for_health_check( + { + "mode": "audio_speech", + "health_check_params": {"voice": "sage", "response_format": "wav"}, + "health_check_voice": "shimmer", + }, + {"model": "openai/tts-1"}, + ) + + assert updated["voice"] == "shimmer" + assert updated["response_format"] == "wav" + + +@pytest.mark.parametrize( + "bad_value", + ["mediaSource", ["mediaSource"], 5, True], +) +def test_health_check_params_ignored_when_not_a_dict(bad_value, caplog): + """A misconfigured health_check_params is skipped with a warning instead of breaking the probe.""" + with caplog.at_level(logging.WARNING, logger="litellm.proxy.health_check"): + updated = _update_litellm_params_for_health_check( + {"mode": "chat", "health_check_params": bad_value}, + {"model": "openai/dummy"}, + ) + + assert updated["model"] == "openai/dummy" + assert updated["max_tokens"] == 16 + assert "health_check_params" in caplog.text + + +def test_health_check_params_apply_to_non_chat_modes(): + """Non-chat probes get health_check_params too, and still no max_tokens.""" + updated = _update_litellm_params_for_health_check( + {"mode": "embedding", "health_check_params": {"dimensions": 8}}, + {"model": "bedrock/amazon.titan-embed-text-v2:0"}, + ) + + assert updated["dimensions"] == 8 + assert "max_tokens" not in updated + + +async def _pegasus_health_check_request_body( + model_info: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> dict[str, object]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + litellm_params = _update_litellm_params_for_health_check( + model_info, + { + "model": "bedrock/us.twelvelabs.pegasus-1-2-v1:0", + "aws_access_key_id": "fake-access-key", + "aws_secret_access_key": "fake-secret-key", + "aws_region_name": "us-east-1", + }, + ) + + with respx.mock(assert_all_called=True) as respx_mock: + invoke_route = respx_mock.post( + host="bedrock-runtime.us-east-1.amazonaws.com", + path__regex=r"/model/.+/invoke", + ).respond(json={"message": "a person walks a dog", "finishReason": "stop"}) + result = await litellm.ahealth_check(litellm_params, mode="chat") + + assert "error" not in result, result + return json.loads(invoke_route.calls.last.request.content) + + +@pytest.mark.asyncio +async def test_health_check_params_reach_the_bedrock_invoke_body(monkeypatch): + """The probe Bedrock actually receives carries mediaSource, which is what unblocks Pegasus.""" + media_source = {"s3Location": {"uri": "s3://my-bucket/clip.mp4"}} + + body = await _pegasus_health_check_request_body( + {"mode": "chat", "health_check_params": {"mediaSource": media_source}}, monkeypatch + ) + + assert body["mediaSource"] == media_source + assert body["maxOutputTokens"] == 16 + assert body["inputPrompt"] + + +@pytest.mark.asyncio +async def test_bedrock_invoke_body_has_no_media_source_without_health_check_params(monkeypatch): + """Negative control: the field only appears because the deployment asked for it.""" + body = await _pegasus_health_check_request_body({"mode": "chat"}, monkeypatch) + + assert "mediaSource" not in body + + +@pytest.mark.asyncio +async def test_run_model_health_check_skips_complexity_router_deployment(): + fake_ahealth_check = AsyncMock(return_value={}) + model = { + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"simple": "gpt-4o-mini"}}, + "complexity_router_default_model": "gpt-4o-mini", + }, + "model_info": {}, + } + + with patch.object(hc_module.litellm, "ahealth_check", fake_ahealth_check): + result = await hc_module._run_model_health_check(model) + + fake_ahealth_check.assert_not_called() + assert result == {} diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 31d2a6cef98..b9a31acca96 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -78,6 +78,16 @@ def client_no_auth(): return TestClient(app) +def test_cors_exposes_cache_key_header_to_browser_js(): + from fastapi.middleware.cors import CORSMiddleware + + from litellm.constants import LITELLM_UI_ALLOW_HEADERS + + cors_middleware = next(m for m in app.user_middleware if m.cls is CORSMiddleware) + assert cors_middleware.kwargs["expose_headers"] is LITELLM_UI_ALLOW_HEADERS + assert "x-litellm-cache-key" in cors_middleware.kwargs["expose_headers"] + + def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): mock_login_result = {"user_id": "test-user"} mock_prisma_client = MagicMock() @@ -11347,3 +11357,38 @@ class TestRouterModelNameOnStreamingChunks: assert len(frames) >= 3 assert '"router_model_name":"deep-model"' in frames[0] assert all('"router_model_name":"backup-tier"' in frame for frame in frames[1:]) + + +@pytest.mark.asyncio +async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the_db_read(): + """A team-member spend reset writes the post-reset floor to the spend_db_floor marker + (auth_checks.invalidate_team_member_spend_state). A floor read already in flight when the + reset commits would otherwise cache its stale pre-reset DB value over the fresh marker, + letting a budget check raise the counter right back above the just-reset spend + (regression: PR #37971 Greptile finding).""" + from litellm.proxy.proxy_server import _authoritative_floor_spend + + real_spend_counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + marker_key = f"spend_db_floor:{counter_key}" + + async def db_read_racing_with_a_reset(prisma_client, counter_key): + real_spend_counter_cache.in_memory_cache.set_cache(key=marker_key, value=0.0) + return 999.0 + + with ( + patch.object( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + proxy_server_module, "spend_counter_cache", real_spend_counter_cache + ), + patch.object( # test-quality-ok: the DB read must race the reset; no injectable seam for module-global prisma reads + proxy_server_module.SpendCounterReseed, + "from_db", + AsyncMock(side_effect=db_read_racing_with_a_reset), + ), + ): + result = await _authoritative_floor_spend(counter_key=counter_key) + + assert result == 0.0 + assert real_spend_counter_cache.in_memory_cache.get_cache(key=marker_key) == 0.0, ( + "the in-flight DB read clobbered the post-reset floor marker with the stale pre-reset value" + ) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index fb01216982f..6920cc0dae3 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -11,7 +11,6 @@ from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks - from unittest.mock import MagicMock, patch from litellm.proxy.utils import get_custom_url, join_paths @@ -82,9 +81,7 @@ async def test_proxy_only_error_log_marks_no_upstream_llm_call(): captured = {} def fake_pre_call(self, *args, **kwargs): - captured["flag"] = self.model_call_details.get( - LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL - ) + captured["flag"] = self.model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) from litellm.litellm_core_utils.litellm_logging import Logging @@ -102,9 +99,7 @@ async def test_proxy_only_error_log_marks_no_upstream_llm_call(): "model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-bad", request_route="/v1/chat/completions" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route="/v1/chat/completions"), route="/v1/chat/completions", original_exception=Exception("bad key"), ) @@ -148,13 +143,9 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): request_data={ "model": "gpt-4o", "input": "blocked prompt", - "litellm_metadata": { - "standard_logging_guardrail_information": guardrail_info - }, + "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, }, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-1234", request_route="/v1/responses" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/responses"), route="/v1/responses", original_exception=HTTPException(status_code=400, detail="blocked"), ) @@ -163,12 +154,7 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): Logging.pre_call = orig_pre_call Logging.async_failure_handler = orig_async_failure - assert ( - captured["litellm_params"]["litellm_metadata"][ - "standard_logging_guardrail_information" - ] - == guardrail_info - ) + assert captured["litellm_params"]["litellm_metadata"]["standard_logging_guardrail_information"] == guardrail_info assert "litellm_metadata" not in captured["optional_params"] @@ -206,9 +192,7 @@ def test_get_model_group_info_order(): def test_join_paths_no_duplication(): """Test that join_paths doesn't duplicate route when base_path already ends with it""" - result = join_paths( - base_path="http://0.0.0.0:4000/my-custom-path/", route="/my-custom-path" - ) + result = join_paths(base_path="http://0.0.0.0:4000/my-custom-path/", route="/my-custom-path") assert result == "http://0.0.0.0:4000/my-custom-path" @@ -814,9 +798,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=system_prompt - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=system_prompt) assert estimated.prompt_tokens == expected @pytest.mark.asyncio @@ -834,7 +818,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter( model="gpt-3.5-turbo", text="part one of the system prompt. part two of the system prompt." ) assert estimated.prompt_tokens == expected @@ -870,9 +856,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=system_prompt - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=system_prompt) assert estimated.prompt_tokens == expected @pytest.mark.asyncio @@ -890,9 +876,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=dispatched_system - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=dispatched_system) assert estimated.prompt_tokens == expected @@ -916,9 +902,7 @@ def test_create_model_info_response_includes_max_tokens_from_lookup(): model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert response["id"] == "some-model" @@ -935,9 +919,7 @@ def test_create_model_info_response_does_not_call_router_group_info(): model_id="some-model", provider="openai", llm_router=router, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) router.get_model_group_info.assert_not_called() @@ -968,9 +950,7 @@ def test_create_model_info_response_deployment_limits_override_cost_map(): model_id="gpt-4o", provider="openai", llm_router=router, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert response["max_input_tokens"] == 200000 @@ -1008,9 +988,7 @@ def test_create_model_info_response_survives_malformed_cost_map_limits(bad_value model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=bad_value, max_output_tokens=bad_value - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=bad_value, max_output_tokens=bad_value), ) assert response["id"] == "some-model" @@ -1023,9 +1001,7 @@ def test_create_model_info_response_keeps_valid_cost_map_limit_beside_malformed_ model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens="128,000", max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens="128,000", max_output_tokens=16384), ) assert "max_input_tokens" not in response @@ -1068,9 +1044,7 @@ def test_create_model_info_response_emits_integer_token_counts(): model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert isinstance(response["max_input_tokens"], int) @@ -1119,9 +1093,7 @@ def test_create_model_info_response_no_router_keeps_base_fields(): def test_create_model_info_response_reads_real_cost_map(): - response = create_model_info_response( - model_id="gpt-4o", provider="openai", llm_router=None - ) + response = create_model_info_response(model_id="gpt-4o", provider="openai", llm_router=None) assert isinstance(response["max_input_tokens"], int) assert response["max_input_tokens"] > 0 @@ -1205,10 +1177,7 @@ class TestPostCallFailureHookLLMExceptionAlerting: @pytest.mark.asyncio async def test_http_exception_does_not_alert(self): - assert ( - await self._alerted(HTTPException(status_code=400, detail="blocked")) - is False - ) + assert await self._alerted(HTTPException(status_code=400, detail="blocked")) is False @pytest.mark.asyncio async def test_genuine_llm_api_error_still_alerts(self): @@ -1241,9 +1210,7 @@ class TestPostCallFailureHookProxyExceptionLogging: await proxy_logging_obj.post_call_failure_hook( request_data={}, original_exception=exc, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-test", request_route=request_route - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=request_route), ) return handle_mock.await_count > 0 @@ -1260,20 +1227,12 @@ class TestPostCallFailureHookProxyExceptionLogging: @pytest.mark.asyncio async def test_proxy_exception_on_llm_route_is_logged(self): - assert ( - await self._logged(self._block(), request_route="/v1/chat/completions") - is True - ) + assert await self._logged(self._block(), request_route="/v1/chat/completions") is True @pytest.mark.asyncio async def test_generic_exception_on_llm_route_is_not_logged(self): # A raw provider/unknown exception is logged by the LLM call path, not here. - assert ( - await self._logged( - Exception("upstream 503"), request_route="/v1/chat/completions" - ) - is False - ) + assert await self._logged(Exception("upstream 503"), request_route="/v1/chat/completions") is False class TestShouldUseSmtpSsl: @@ -1307,9 +1266,7 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection( - smtp_host="mail.example.com", smtp_port=465 - ) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465) mock_smtp.assert_not_called() assert result is mock_smtp_ssl.return_value @@ -1329,9 +1286,7 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection( - smtp_host="mail.example.com", smtp_port=587 - ) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587) mock_smtp_ssl.assert_not_called() assert result is mock_smtp.return_value @@ -1352,9 +1307,7 @@ class TestSendEmailStartTls: monkeypatch.delenv("SMTP_USE_SSL", raising=False) mock_server = MagicMock(spec=smtplib.SMTP) - with patch( - "litellm.proxy.utils._create_smtp_connection" - ) as mock_create_connection: + with patch("litellm.proxy.utils._create_smtp_connection") as mock_create_connection: mock_create_connection.return_value.__enter__.return_value = mock_server await send_email( receiver_email="receiver@example.com", @@ -1690,9 +1643,7 @@ def test_a_failed_dispatch_is_estimated_as_input_only(): usage = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage is not None - assert usage.prompt_tokens == _count_request_input_tokens( - FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None - ) + assert usage.prompt_tokens == _count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage.completion_tokens == 0 assert usage.total_tokens == usage.prompt_tokens @@ -1759,9 +1710,7 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost(): def test_a_failure_that_cost_the_provider_nothing_lifts_nothing(model_call_details, dispatched): from litellm.proxy.utils import _failure_usage_to_lift - assert _failure_usage_to_lift( - model_call_details=model_call_details, request_body={}, dispatched=dispatched - ) is None + assert _failure_usage_to_lift(model_call_details=model_call_details, request_body={}, dispatched=dispatched) is None def test_the_no_upstream_call_key_the_module_uses_is_the_one_asserted_above(): @@ -1829,3 +1778,102 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): assert lifted["response_cost"] == 0.0 assert lifted["combined_usage_object"].prompt_tokens > 0 assert lifted["standard_logging_object"] == {"id": "log-1"} + + +@pytest.mark.asyncio +async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch): + """Regression for LIT-6043: an expected 4xx must not format a traceback for + either the async or the threaded sync failure handler.""" + import asyncio + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr(litellm, "failure_callback", []) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + captured = {} + sync_ran = asyncio.Event() + loop = asyncio.get_running_loop() + + async def fake_async_failure(self, exception, traceback_exception, *args, **kwargs): + captured["async_traceback"] = traceback_exception + + def fake_sync_failure(self, exception, traceback_exception, *args, **kwargs): + captured["sync_traceback"] = traceback_exception + loop.call_soon_threadsafe(sync_ran.set) + + orig_async_failure = Logging.async_failure_handler + orig_sync_failure = Logging.failure_handler + Logging.async_failure_handler = fake_async_failure + Logging.failure_handler = fake_sync_failure + try: + try: + raise HTTPException(status_code=400, detail="Invalid model name passed in") + except HTTPException as exc: + await proxy_logging_obj._handle_logging_proxy_only_error( + request_data={ + "model": "does-not-exist", + "messages": [{"role": "user", "content": "hi"}], + }, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=exc, + ) + await asyncio.wait_for(sync_ran.wait(), timeout=5) + finally: + Logging.async_failure_handler = orig_async_failure + Logging.failure_handler = orig_sync_failure + + assert captured["async_traceback"] == "" + assert captured["sync_traceback"] == "" + + +@pytest.mark.asyncio +async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monkeypatch): + """Unexpected (5xx) errors keep the full traceback, and a configured + sync-only failure callback still gets its threaded handler.""" + import asyncio + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + + def _custom_sync_callback(kwargs, completion_response, start_time, end_time): + pass + + monkeypatch.setattr(litellm, "failure_callback", [_custom_sync_callback]) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + captured = {} + sync_ran = asyncio.Event() + loop = asyncio.get_running_loop() + + async def fake_async_failure(self, exception, traceback_exception, *args, **kwargs): + captured["async_traceback"] = traceback_exception + + def fake_sync_failure(self, *args, **kwargs): + loop.call_soon_threadsafe(sync_ran.set) + + orig_async_failure = Logging.async_failure_handler + orig_sync_failure = Logging.failure_handler + Logging.async_failure_handler = fake_async_failure + Logging.failure_handler = fake_sync_failure + try: + try: + raise HTTPException(status_code=500, detail="internal error") + except HTTPException as exc: + await proxy_logging_obj._handle_logging_proxy_only_error( + request_data={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + }, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=exc, + ) + await asyncio.wait_for(sync_ran.wait(), timeout=5) + finally: + Logging.async_failure_handler = orig_async_failure + Logging.failure_handler = orig_sync_failure + + assert "test_proxy_utils" in captured["async_traceback"] diff --git a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py b/tests/test_litellm/proxy/video_endpoints/test_endpoints.py index a959326817c..c5996f95f54 100644 --- a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/video_endpoints/test_endpoints.py @@ -47,6 +47,7 @@ from litellm.types.videos.utils import ( ) from fastapi import Response +from starlette.datastructures import UploadFile as StarletteUploadFile # --------------------------------------------------------------------------- # # A real model-encoded video id: decodes (for real) to provider "azure", @@ -372,6 +373,7 @@ async def test_content__model_encoded_id(harness): async def call_edit( harness: Harness, *, body: Dict[str, Any], headers=None, query=None ): + harness.read_body.return_value = dict(body) return await endpoints.video_edit( request=FakeRequest(headers=headers, query=query, raw_body=orjson.dumps(body)), fastapi_response=Response(), @@ -428,6 +430,52 @@ async def test_edit__missing_video_object_defaults_to_openai(harness): assert "video" not in data +@pytest.mark.asyncio +async def test_edit__bare_string_video_id_from_form_field(harness): + await call_edit(harness, body={"prompt": "brighter", "video": "video_plain"}) + + assert harness.processor_data() == { + "prompt": "brighter", + "video_id": "video_plain", + "custom_llm_provider": "openai", + } + + +@pytest.mark.asyncio +async def test_edit__json_string_video_reference_from_form_field(harness): + await call_edit( + harness, + body={"prompt": "brighter", "video": orjson.dumps({"id": "video_plain"}).decode()}, + ) + + assert harness.processor_data()["video_id"] == "video_plain" + + +@pytest.mark.asyncio +async def test_edit__uploaded_video_file_is_forwarded_not_dropped(harness): + """A multipart-uploaded source video must be converted to bytes and attached + under ``video`` so the provider receives the file. Before the fix the upload + was popped, coerced to an empty ``video_id``, and silently dropped.""" + import io + + upload = StarletteUploadFile(file=io.BytesIO(b"rawmp4"), filename="clip.mp4") + harness.read_body.return_value = {"prompt": "make it nighttime", "video": upload} + + await endpoints.video_edit( + request=FakeRequest(raw_body=b"multipart"), + fastapi_response=Response(), + user_api_key_dict=_user(), + ) + + harness.batch_to_bytesio.assert_called_once_with((upload,)) + assert harness.processor_data() == { + "prompt": "make it nighttime", + "video": b"filebytes", + "video_id": "", + "custom_llm_provider": "openai", + } + + # =========================================================================== # # GET /v1/videos - video_list # # =========================================================================== # @@ -471,6 +519,7 @@ async def test_list__provider_from_header(harness): async def call_remix( harness: Harness, video_id: str, *, body, headers=None, query=None ): + harness.read_body.return_value = dict(body) return await endpoints.video_remix( video_id=video_id, request=FakeRequest(headers=headers, query=query, raw_body=orjson.dumps(body)), @@ -629,6 +678,7 @@ async def test_get_character__plain_id_defaults_openai_no_encode(harness): async def call_extension(harness: Harness, *, body, headers=None, query=None): + harness.read_body.return_value = dict(body) return await endpoints.video_extension( request=FakeRequest(headers=headers, query=query, raw_body=orjson.dumps(body)), fastapi_response=Response(), diff --git a/tests/test_litellm/proxy/video_endpoints/test_utils.py b/tests/test_litellm/proxy/video_endpoints/test_utils.py index efbaaff5f4b..9a2c208c075 100644 --- a/tests/test_litellm/proxy/video_endpoints/test_utils.py +++ b/tests/test_litellm/proxy/video_endpoints/test_utils.py @@ -1,8 +1,9 @@ """ Pure-logic contract tests for litellm/proxy/video_endpoints/utils.py -Three helpers the video proxy endpoints lean on: +Four helpers the video proxy endpoints lean on: - extract_model_from_target_model_names: first model from a comma string / list + - video_reference_to_id: normalize a video reference (dict / bare id / JSON string) to an id - get_custom_provider_from_data: provider precedence (top-level > extra_body) - encode_character_id_in_response: re-encode a response id in place @@ -20,6 +21,7 @@ from litellm.proxy.video_endpoints.utils import ( encode_character_id_in_response, extract_model_from_target_model_names, get_custom_provider_from_data, + video_reference_to_id, ) from litellm.types.videos.utils import ( decode_character_id_with_provider, @@ -53,6 +55,31 @@ def test_extract_model__non_str_non_list_is_none(value): assert extract_model_from_target_model_names(value) is None +# =========================================================================== # +# video_reference_to_id +# =========================================================================== # + + +@pytest.mark.parametrize( + "video_ref,expected", + [ + ({"id": "video_123"}, "video_123"), # dict reference -> its id + ({"id": ""}, ""), # dict with empty id + ({}, ""), # dict missing id -> default empty + ({"other": "x"}, ""), # dict without id key + ("video_123", "video_123"), # bare id string (not valid JSON) -> itself + ('{"id": "video_9"}', "video_9"), # JSON-encoded dict -> its id + ('{"other": 1}', ""), # JSON-encoded dict without id -> empty + ("[1, 2]", "[1, 2]"), # JSON parses to non-dict -> original string + (None, ""), # non-str, non-dict + (123, ""), # non-str, non-dict + (["video_123"], ""), # list is neither dict nor str + ], +) +def test_video_reference_to_id(video_ref, expected): + assert video_reference_to_id(video_ref) == expected + + # =========================================================================== # # get_custom_provider_from_data # =========================================================================== # diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 85777afe81c..587be59c550 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -1,6 +1,10 @@ import logging from unittest.mock import MagicMock, patch +import httpx +import pytest +import respx + import litellm @@ -62,3 +66,66 @@ def test_rerank_does_not_log_request_content_at_info(caplog): assert all( r.levelno == logging.DEBUG for r in optional_params_logs ), "optional_rerank_params must be logged at DEBUG, not INFO" + + +TOGETHER_RERANK_BODY = { + "id": "rerank-mock-id", + "results": [{"index": 0, "relevance_score": 0.95}], + "usage": {"prompt_tokens": 10, "total_tokens": 10}, +} + + +def test_together_rerank_defaults_to_together_ai_host(respx_mock: respx.MockRouter, monkeypatch): + """Regression for the Together host migration: rerank used to hardcode + https://api.together.xyz/v1/rerank. The default must now be api.together.ai.""" + monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False) + + mock_route = respx_mock.post("https://api.together.ai/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + response = litellm.rerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + ) + + assert mock_route.called + assert response.results[0]["relevance_score"] == 0.95 + + +def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter): + """Regression: a custom api_base was silently ignored by the Together rerank handler.""" + mock_route = respx_mock.post("https://custom-together.example/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + litellm.rerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + api_base="https://custom-together.example/v1", + ) + + assert mock_route.called + assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key" + + +@pytest.mark.asyncio +async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch): + """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank.""" + monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://env-together.example/v1") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + mock_route = respx_mock.post("https://env-together.example/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + response = await litellm.arerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + ) + + assert mock_route.called + assert response.results[0]["relevance_score"] == 0.95 diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py index 3a1c77d1dab..3dec1d571c4 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_tool_output_order_preserved_for_gemini.py @@ -114,3 +114,46 @@ def test_assistant_message_after_tool_call_is_folded_into_it(): tool_call_idx = next(i for i, m in enumerate(msgs) if isinstance(m, dict) and m.get("tool_calls")) assert msgs[tool_call_idx].get("role") == "assistant" assert msgs[tool_call_idx + 1].get("role") == "tool" + + +def test_assistant_message_before_function_call_keeps_one_assistant_turn(): + """The chat->responses bridge emits an assistant message ahead of its function_call. + + Round-tripping that order back to chat must fold both into a single assistant + turn, so the tool result still follows the message that made the call. + """ + msgs = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( + input=[ + { + "role": "user", + "type": "message", + "content": [{"type": "input_text", "text": "What is the weather?"}], + }, + { + "role": "assistant", + "type": "message", + "content": [{"type": "output_text", "text": "Let me check."}], + }, + { + "type": "function_call", + "name": "get_weather", + "call_id": "call_1", + "arguments": "{}", + }, + { + "type": "function_call_output", + "call_id": "call_1", + "output": "sunny", + }, + ] + ) + + assistant_msgs = [m for m in msgs if isinstance(m, dict) and m.get("role") == "assistant"] + assert len(assistant_msgs) == 1 + assistant = assistant_msgs[0] + assert assistant["content"] == [{"type": "text", "text": "Let me check."}] + assert [tc["function"]["name"] for tc in assistant["tool_calls"]] == ["get_weather"] + + assistant_idx = msgs.index(assistant) + assert msgs[assistant_idx + 1].get("role") == "tool" + assert msgs[assistant_idx + 1].get("tool_call_id") == "call_1" diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index cbaac4d89a7..a29b4d03bc5 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -5988,6 +5988,78 @@ class TestContextAwareClassifier: assert _strip_reminder_blocks(text, pairs) == "<<>> why is my tag stripped?" + @pytest.mark.parametrize( + "text,limit,expected", + [ + pytest.param("short", 10, "short", id="under-the-limit-is-untouched"), + pytest.param("exact", 5, "exact", id="exactly-the-limit-is-untouched"), + pytest.param( + "Second request with more details and longer text", + 30, + "Second re...tails and longer text", + id="over-the-limit-keeps-both-ends", + ), + pytest.param("abcdefghij", 4, "a...hij", id="tiny-limit-still-splits"), + pytest.param("abcdefghij", 1, "...j", id="limit-too-small-for-a-head-keeps-the-tail"), + pytest.param("abcdefghij", 0, "...", id="zero-limit-quotes-nothing"), + pytest.param("日本語のテキストと最後の質問", 6, "日...最後の質問", id="cjk-slices-by-character"), + ], + ) + def test_truncate_keeps_the_end_of_an_over_long_turn(self, text, limit, expected): + """A cut turn keeps its tail, because that is where a chat turn puts its ask. + + Head-only truncation was the shipped behavior and it discarded exactly the part that carries + the difficulty. The degenerate limits are here because the budget hands this function whatever + space is left rather than a configured constant, so it must stay total: a limit too small to + hold a head degrades to tail-only rather than raising or slicing with a negative index. + """ + from litellm.router_strategy.complexity_router.complexity_router import _truncate + + assert _truncate(text, limit) == expected + + def test_truncate_holds_its_length_budget(self): + """Cutting to N spends N characters plus the marker, at every N including the degenerate ones. + + The marker is the cost of having cut at all, so it is charged uniformly rather than only once + the limit is large enough to hold a head; a caller sizing a cut against a remaining budget can + therefore price it as limit plus marker without special-casing the small end. + """ + from litellm.router_strategy.complexity_router.complexity_router import _TRUNCATION_MARKER, _truncate + + text = "x" * 500 + + assert all( + len(_truncate(text, limit)) == limit + len(_TRUNCATION_MARKER) for limit in (0, 1, 2, 4, 30, 200, 499) + ) + + def test_clipped_prior_turn_still_carries_the_ask_it_closes_on(self): + """The reported defect, at the level the classifier sees it. + + A prior turn that opens with an incident report and closes with the request routed to the + cheapest tier, because the 200-character cut kept the report and dropped the request. The + quoted turn must carry both ends. + """ + from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns + + turn = ( + "We run a multi-region gateway and last night the eu-west pod returned 502s on the " + "streaming path only, for thirty minutes, while non-streaming stayed healthy the whole " + "window and the cooldown map was mid-failover. " + "Filler sentence to push past the cap. " * 4 + + "Now rewrite the streaming retry path and prove it cannot livelock." + ) + + quoted = _extract_prior_turns( + [{"role": "user", "content": turn}, {"role": "user", "content": "go ahead"}], + "go ahead", + 3, + budget_chars=10_000, + per_turn_chars=200, + include_assistant=False, + ) + + assert "multi-region gateway" in quoted[0][1] + assert "prove it cannot livelock" in quoted[0][1] + @pytest.mark.parametrize( "messages,current_ask,window,per_turn_chars,include_assistant,expected", [ @@ -6002,7 +6074,7 @@ class TestContextAwareClassifier: 2, 30, False, - (("user", "First request"), ("user", "Second request with more detai...")), + (("user", "First request"), ("user", "Second re...tails and longer text")), id="current-ask-excluded-and-long-turn-marked-as-clipped", ), pytest.param( @@ -6111,7 +6183,7 @@ class TestContextAwareClassifier: 1, 20, True, - (("assistant", "a very long plan tha..."),), + (("assistant", "a very...l past the cap"),), id="assistant-reply-is-clipped-at-per-turn-chars", ), pytest.param( @@ -6134,7 +6206,8 @@ class TestContextAwareClassifier: The current ask is excluded by matching it rather than by position, since `aclassify` takes `prompt` and `messages` separately and a caller may classify other than the newest turn. A turn - cut at per_turn_chars is marked so a clip does not read as an abandoned thought. + over per_turn_chars keeps both ends with its middle elided, so the ask it closes on survives the + cut and the marker does not read as an abandoned thought. With assistant turns enabled the window is the last N turns of the conversation rather than the last N asks, which is what makes a plan the assistant called complex visible under a bare "yes". @@ -6144,7 +6217,167 @@ class TestContextAwareClassifier: """ from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns - assert _extract_prior_turns(messages, current_ask, window, per_turn_chars, include_assistant) == expected + assert ( + _extract_prior_turns( + messages, + current_ask, + window, + budget_chars=10_000, + per_turn_chars=per_turn_chars, + include_assistant=include_assistant, + ) + == expected + ) + + @pytest.mark.parametrize( + "turn_lengths,budget_chars,expected_lengths", + [ + pytest.param((50, 50, 50), 10_000, (50, 50, 50), id="a-block-that-fits-is-quoted-whole"), + pytest.param((100, 100, 100), 250, (100, 100), id="oldest-turn-is-dropped-whole"), + pytest.param((500, 100), 400, (300, 100), id="only-the-boundary-turn-is-cut"), + pytest.param((900,), 300, (300,), id="a-turn-larger-than-the-budget-is-still-quoted"), + pytest.param((500, 100), 180, (100,), id="a-remainder-too-small-to-carry-a-sentence-is-dropped"), + pytest.param((50,), 0, (), id="a-zero-budget-quotes-nothing"), + ], + ) + def test_budget_bounds_the_block_not_each_turn(self, turn_lengths, budget_chars, expected_lengths): + """Turns are taken newest first and quoted whole while they fit. + + The defect this replaces capped every turn independently, so a 785 character turn was cut even + though the whole block it belonged to was 353 characters. Bounding the block instead means an + ordinary conversation arrives intact, and when the budget really does run out the older turns + are dropped entire rather than each arriving mangled. At most one turn is ever cut, and a + remainder too small to carry a sentence is dropped rather than quoted as two ellipses around a + fragment. A single turn bigger than the whole budget is still quoted, cut to the budget, since + dropping it would leave the classifier with no context at all. + """ + from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns + + messages = [{"role": "user", "content": f"{i}" * length} for i, length in enumerate(turn_lengths)] + + quoted = _extract_prior_turns( + [*messages, {"role": "user", "content": "go ahead"}], + "go ahead", + len(turn_lengths), + budget_chars=budget_chars, + per_turn_chars=None, + include_assistant=False, + ) + + assert tuple(len(text) for _, text in quoted) == expected_lengths + + @pytest.mark.parametrize("budget_chars", [130, 200, 351, 400, 999, 8000]) + @pytest.mark.parametrize("turn_lengths", [(900,), (500, 100), (100, 100, 100), (50, 50, 50)]) + def test_the_quoted_block_never_exceeds_the_budget(self, turn_lengths, budget_chars): + """The budget is a ceiling on what is quoted, marker included. + + Cutting the boundary turn to the remainder and then appending the marker put the block three + characters over the number an operator configured, which is the kind of drift that makes a + documented ceiling untrue. Asserted across shapes rather than at the one boundary that happened + to be wrong, so any future off-by-marker anywhere in the fill is caught here. + """ + from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns + + messages = [{"role": "user", "content": f"{i}" * length} for i, length in enumerate(turn_lengths)] + + quoted = _extract_prior_turns( + [*messages, {"role": "user", "content": "go ahead"}], + "go ahead", + len(turn_lengths), + budget_chars=budget_chars, + per_turn_chars=None, + include_assistant=False, + ) + + assert sum(len(text) for _, text in quoted) <= budget_chars + + def test_per_turn_cap_still_clamps_when_an_operator_sets_it(self): + """An operator who set the per-turn cap keeps exactly what they configured. + + The cap stopped being the default, so it has to keep working for the deployments that named it + deliberately; it applies before the block budget rather than instead of it. + """ + from litellm.router_strategy.complexity_router.complexity_router import _extract_prior_turns + + quoted = _extract_prior_turns( + [{"role": "user", "content": "z" * 900}, {"role": "user", "content": "go ahead"}], + "go ahead", + 3, + budget_chars=10_000, + per_turn_chars=200, + include_assistant=False, + ) + + assert len(quoted[0][1]) == 203 + + @pytest.mark.asyncio + async def test_a_long_turn_reaches_the_classifier_whole_by_default( + self, mock_router_instance, llm_classifier_config + ): + """The shipped defaults quote an ordinary long turn without cutting it anywhere. + + This is the whole point of the change, asserted where a deployment actually meets it: no knob + set, one turn well past the retired 200 character cap, and no truncation marker in the payload. + """ + from litellm.router_strategy.complexity_router.complexity_router import _TRUNCATION_MARKER + + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=llm_classifier_config, + ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) + turn = "The incident ran from 02:10 to 02:40 and only streaming was affected. " * 10 + "Now rewrite it" + + await router.aclassify( + "go ahead", + messages=[{"role": "user", "content": turn}, {"role": "user", "content": "go ahead"}], + ) + + user_payload = mock_router_instance.acompletion.call_args.kwargs["messages"][1]["content"] + assert turn in user_payload + assert _TRUNCATION_MARKER not in user_payload + + @pytest.mark.asyncio + async def test_a_turn_dropped_for_budget_still_counts_as_prior_conversation( + self, mock_router_instance, llm_classifier_config + ): + """Dropping turns to fit the budget must not make a long conversation look single-turn. + + The depth line gates on whether prior conversation exists, not on whether any of it was worth + quoting, exactly so a continuation is never reported as a context-free first request. A budget + tight enough to drop every turn is the newest way to reach that mismatch. + """ + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**llm_classifier_config, "classifier_context_budget_chars": 1}, + ) + mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "SIMPLE"}')) + + await router.aclassify( + "go ahead", + messages=[ + {"role": "user", "content": "a long earlier request that cannot fit a one character budget"}, + {"role": "user", "content": "go ahead"}, + ], + ) + + user_payload = mock_router_instance.acompletion.call_args.kwargs["messages"][1]["content"] + assert "Recent conversation" not in user_payload + assert "Conversation so far" in user_payload + + def test_context_defaults_bound_the_block_and_leave_turns_uncapped(self): + """The shipped defaults: a block budget, and no per-turn cap unless one is named.""" + from litellm.router_strategy.complexity_router.config import ( + DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, + ComplexityRouterConfig, + ) + + config = ComplexityRouterConfig() + + assert config.classifier_context_budget_chars == DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS + assert config.classifier_context_per_turn_chars is None def test_prior_turn_context_strips_every_configured_pair(self): """The classifier's context window is stripped with the same pairs as the ask. @@ -6163,7 +6396,7 @@ class TestContextAwareClassifier: {"role": "user", "content": "current ask"}, ] - assert _extract_prior_turns(messages, "current ask", 5, 200, False, pairs) == ( + assert _extract_prior_turns(messages, "current ask", 5, 10_000, 200, False, pairs) == ( ("user", "what about b-trees?"), ("user", "and heaps?"), ) diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py new file mode 100644 index 00000000000..6cd5b956b05 --- /dev/null +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -0,0 +1,198 @@ +import pytest + +from litellm.router_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) + + +class TestDeploymentIsCatalogMapped: + def test_a_mode_the_catalog_supplied_marks_the_deployment_mapped(self): + assert deployment_is_catalog_mapped({"mode": "chat"}, {}) is True + + def test_a_deployment_the_catalog_never_described_is_not_mapped(self): + assert deployment_is_catalog_mapped(None, {}) is False + assert deployment_is_catalog_mapped({"max_input_tokens": 200000}, {}) is False + + def test_a_mode_the_operator_wrote_does_not_make_the_deployment_mapped(self): + # Every deployment is registered in the cost map under its own id, so an operator-written + # mode reads back identically to one the catalog supplied and would otherwise let an + # off-map deployment empty the levels its mapped siblings agree on. + assert deployment_is_catalog_mapped({"mode": "chat"}, {"mode": "chat", "id": "abc"}) is False + + +class TestProvenanceSeparatesUnknownFromNonReasoning: + def test_an_off_map_deployment_resolves_to_unknown(self): + # get_model_info answers supports_reasoning None both for a deployment the map never + # described and for a mapped non-reasoning model, so reading an unset flag as () would let + # one custom deployment empty every level its mapped siblings agree on. + assert resolve_supported_reasoning_efforts({}, deployment_is_mapped=False) is None + assert resolve_supported_reasoning_efforts({"supports_reasoning": None}, deployment_is_mapped=False) is None + + def test_a_mapped_deployment_the_map_calls_non_reasoning_supports_no_efforts(self): + assert resolve_supported_reasoning_efforts({}, deployment_is_mapped=True) == () + assert resolve_supported_reasoning_efforts({"supports_reasoning": None}, deployment_is_mapped=True) == () + + def test_an_explicit_false_supports_no_efforts_off_the_map_too(self): + # The operator's own escape hatch: saying so on an off-map deployment must still empty the + # group, since nothing else can tell the resolver that model takes no effort level. + assert resolve_supported_reasoning_efforts({"supports_reasoning": False}, deployment_is_mapped=False) == () + + +class TestResolveSupportedReasoningEfforts: + def test_a_reasoning_model_with_no_flags_at_all_resolves_to_unknown(self): + # 689 of the map's 854 reasoning entries carry no effort flag, and the o-series, xai and + # bedrock nova entries among them accept neither none nor minimal, so composing a set out of + # the opt-out defaults alone would advertise levels those providers reject. + assert resolve_supported_reasoning_efforts({"supports_reasoning": True}, deployment_is_mapped=True) is None + + def test_explicit_false_removes_an_opt_out_level(self): + # The gpt-5.5-pro shape from the model map: only medium/high/xhigh are accepted upstream. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": False, + "supports_low_reasoning_effort": False, + "supports_xhigh_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("medium", "high", "xhigh") + + def test_explicit_true_adds_the_opt_in_levels(self): + # The claude-opus shape: xhigh and max explicitly true, everything else absent. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("none", "minimal", "low", "medium", "high", "xhigh", "max") + + def test_opt_in_flag_set_false_stays_excluded(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_minimal_reasoning_effort": True, + "supports_xhigh_reasoning_effort": False, + }, + deployment_is_mapped=True, + ) + assert resolved == ("none", "minimal", "low", "medium", "high") + + +class TestBareModelNameFallback: + def test_a_prefixed_entry_inherits_the_flags_of_its_unprefixed_twin(self): + """azure/gpt-5-mini carries no effort flag while gpt-5-mini carries three, and the request + path resolves capability flags through that same twin (#20885). Reading only the prefixed + entry would answer unknown for a model the map fully describes.""" + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="gpt-5-mini", custom_llm_provider="azure")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "minimal", + "low", + "medium", + "high", + ) + + def test_the_prefixed_entry_wins_over_its_twin_per_flag(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-mini", + "supports_xhigh_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("minimal", "low", "medium", "high", "xhigh") + + +class TestNoneLevelPolarity: + def test_none_stays_opt_out_off_azure(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "openai", + "key": "openai/some-reasoner", + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + def test_none_stays_opt_out_on_an_azure_model_outside_the_gpt_5_family(self): + """AzureOpenAIGPT5Config is selected by is_model_gpt_5_model, so an azure o-series or + anthropic deployment never reaches the gate that refuses none and must keep the level.""" + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/o3", + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + def test_azure_gpt_5_without_the_flag_does_not_advertise_none(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-turbo", + "supports_minimal_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("minimal", "low", "medium", "high") + + def test_azure_gpt_5_with_the_flag_advertises_none(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-turbo", + "supports_none_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + @pytest.mark.parametrize( + "model_key", + ["azure/gpt-5", "azure/gpt-5-mini", "azure/gpt-5-nano", "azure/gpt-5.2", "azure/gpt-5.6"], + ) + def test_azure_advertisement_matches_the_azure_request_gate(self, model_key): + """AzureOpenAIGPT5Config raises UnsupportedParamsError on reasoning_effort='none' for models + it does not flag, so advertising the level there would offer routing a 400.""" + from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model=model_key.split("/", 1)[1], custom_llm_provider="azure")) + resolved = resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) + + assert resolved is not None + gate_accepts_none = AzureOpenAIGPT5Config._supports_reasoning_effort_level(model_key, "none") + assert ("none" in resolved) is gate_accepts_none + + +class TestIntersectSupportedReasoningEfforts: + def test_unknown_never_narrows(self): + assert intersect_supported_reasoning_efforts(["medium", "high"], None) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, ["medium", "high"]) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, None) is None + + def test_intersection_keeps_canonical_order(self): + assert intersect_supported_reasoning_efforts( + ["max", "high", "medium", "xhigh"], ["xhigh", "medium", "minimal"] + ) == ("medium", "xhigh") + + def test_disjoint_sets_intersect_to_empty(self): + assert intersect_supported_reasoning_efforts(["max"], ["minimal"]) == () diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py new file mode 100644 index 00000000000..ccba351deaf --- /dev/null +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -0,0 +1,1705 @@ +"""Tests for tests/code_coverage_tests/check_migrations_no_data_rewrites.py. + +The checker reads migration.sql as SQL rather than as text, so the cases that matter +are the ones a grep would get wrong: the referential actions in a foreign key, of which +the shipped migrations carry 60, an `UPDATE` inside a string literal or a comment, and +an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for conditional DDL. +""" + +import importlib.util +import sys +from pathlib import Path + +_CHECKER_PATH = Path(__file__).resolve().parents[1] / "code_coverage_tests" / "check_migrations_no_data_rewrites.py" +_SPEC = importlib.util.spec_from_file_location("check_migrations_no_data_rewrites", _CHECKER_PATH) +assert _SPEC is not None and _SPEC.loader is not None +checker = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = checker +_SPEC.loader.exec_module(checker) + + +def _scan(tmp_path: Path, sql: str) -> tuple: + directory = tmp_path / "20260101000000_fixture" + directory.mkdir(exist_ok=True) + (directory / "migration.sql").write_text(sql, encoding="utf-8") + return checker.scan_migration(directory) + + +def _keywords(tmp_path: Path, sql: str) -> tuple: + return tuple(violation.keyword for violation in _scan(tmp_path, sql)) + + +class TestRowRewritesAreFlagged: + def test_update_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_delete_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'DELETE FROM "Foo" WHERE "a" IS NULL;') == ("DELETE",) + + def test_update_without_trailing_semicolon_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1') == ("UPDATE",) + + def test_lowercase_update_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'update "Foo" set "a" = 1;') == ("UPDATE",) + + def test_merge_is_flagged(self, tmp_path): + sql = 'MERGE INTO "Foo" t USING "Bar" s ON t."id" = s."id" WHEN MATCHED THEN UPDATE SET "a" = s."a";' + assert _keywords(tmp_path, sql) == ("MERGE",) + + def test_every_offending_statement_is_reported(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nDELETE FROM "Bar";' + assert _keywords(tmp_path, sql) == ("UPDATE", "DELETE") + + def test_the_incident_migration_is_flagged(self, tmp_path): + sql = ( + 'UPDATE "LiteLLM_SpendLogs"\n' + ' SET "created_at" = "endTime",\n' + ' "updated_at" = "endTime"\n' + ' WHERE "created_at" > "endTime" + interval \'1 hour\';\n' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + +class TestSchemaStatementsPass: + def test_on_delete_cascade_is_not_a_data_rewrite(self, tmp_path): + sql = ( + 'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") ' + 'REFERENCES "B"("id") ON DELETE CASCADE ON UPDATE CASCADE;' + ) + assert _keywords(tmp_path, sql) == () + + def test_on_delete_set_null_is_not_a_data_rewrite(self, tmp_path): + sql = ( + 'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") ' + 'REFERENCES "B"("id") ON DELETE SET NULL ON UPDATE CASCADE;' + ) + assert _keywords(tmp_path, sql) == () + + def test_add_column_with_default_passes(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN IF NOT EXISTS "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP;' + assert _keywords(tmp_path, sql) == () + + def test_drop_table_passes(self, tmp_path): + assert _keywords(tmp_path, 'DROP TABLE IF EXISTS "Foo";') == () + + def test_empty_file_passes(self, tmp_path): + assert _keywords(tmp_path, "") == () + + def test_only_comments_passes(self, tmp_path): + assert _keywords(tmp_path, "-- nothing to do here\n") == () + + +class TestInsert: + def test_insert_values_is_bounded_and_passes(self, tmp_path): + assert _keywords(tmp_path, "INSERT INTO \"Foo\" (\"id\") VALUES ('a'), ('b');") == () + + def test_insert_select_scans_and_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar";') == ("INSERT ... SELECT",) + + def test_insert_values_with_a_scalar_subquery_passes(self, tmp_path): + sql = 'INSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT max("id")::text FROM "Bar"));' + assert _keywords(tmp_path, sql) == () + + def test_insert_values_with_a_scalar_subquery_per_row_passes(self, tmp_path): + sql = ( + 'INSERT INTO "Config" ("k", "v") VALUES\n' + " ('a', (SELECT \"id\" FROM \"Bar\" WHERE \"n\" = 'a')),\n" + " ('b', (SELECT \"id\" FROM \"Bar\" WHERE \"n\" = 'b'));" + ) + assert _keywords(tmp_path, sql) == () + + def test_values_inside_a_subquery_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM (VALUES (1), (2)) AS "v"("id");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_values_after_a_set_operation_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" UNION ALL VALUES (1);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_values_after_an_except_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" EXCEPT VALUES (1);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_select_term_after_a_values_list_is_still_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1), (2) UNION ALL SELECT "id" FROM "Bar";' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_without_a_column_list_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_spanning_lines_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id")\n(\n SELECT "id" FROM "Bar"\n);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_over_a_values_list_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT * FROM (VALUES (1), (2)) AS "v"("id"));' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_set_operation_over_parenthesised_selects_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT 1) UNION (SELECT 2);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_joined_to_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_excepting_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) EXCEPT (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_joined_to_a_parenthesised_table_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_set_operation_inside_a_values_list_does_not_flag_it(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES ((SELECT 1 UNION SELECT 2 LIMIT 1));' + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_a_set_operated_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") VALUES ((SELECT max("id") FROM "Bar"))' + " UNION ALL VALUES (2);" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_a_parenthesised_values_list_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")));' + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_set_operated_parenthesised_values_lists_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")))' + " UNION ALL (VALUES (2));" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_set_operation_inside_a_values_list_does_not_split_the_terms(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") VALUES ((SELECT max("id") FROM "Bar"' + ' UNION SELECT max("id") FROM "Bar")) UNION ALL VALUES (2);' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_conflict_target_after_a_parenthesised_row_source_does_not_hide_it(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar")' + ' ON CONFLICT ("id") DO NOTHING;' + ) + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_returning_list_after_a_parenthesised_row_source_does_not_hide_it(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") RETURNING ("id");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_conflict_target_beside_a_bounded_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")))' + ' ON CONFLICT ("id") DO NOTHING;' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_query_term_written_before_a_values_term_is_still_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") UNION ALL (VALUES (2));' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_table_term_beside_parenthesised_values_is_still_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1)) UNION ALL (TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_table_row_source_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "Foo" TABLE "Bar";') == ("INSERT ... TABLE",) + + def test_a_table_named_in_the_insert_target_does_not_flag_it(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "audit table" ("id") VALUES (1);') == () + + def test_a_returning_subquery_after_a_wrapped_values_list_is_not_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_a_conflict_update_after_a_wrapped_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES (1))' + ' ON CONFLICT ("id") DO UPDATE SET "id" = (SELECT max("id") FROM "Bar");' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_wrapped_values_list_of_several_rows_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1), (2)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_the_row_source_names_its_own_keyword_not_a_later_subquery(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (TABLE "Bar") RETURNING (SELECT count(*) FROM "Baz");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_select_term_wrapped_beside_a_values_term_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1) UNION ALL SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_table_term_wrapped_beside_a_values_term_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1) UNION ALL TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_wrapped_set_operation_of_values_lists_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1) UNION ALL VALUES (2));' + assert _keywords(tmp_path, sql) == () + + +class TestCommonTableExpressions: + def test_cte_led_update_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) UPDATE "Foo" SET "a" = 1 FROM batch;' + assert _keywords(tmp_path, sql) == ("WITH ... UPDATE",) + + def test_cte_led_delete_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) DELETE FROM "Foo" USING batch;' + assert _keywords(tmp_path, sql) == ("WITH ... DELETE",) + + def test_cte_led_insert_select_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") SELECT "id" FROM batch;' + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_cte_led_insert_from_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") (SELECT "id" FROM batch);' + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_cte_led_insert_into_a_values_list_passes(self, tmp_path): + sql = 'WITH batch AS (SELECT max("id") FROM "Bar") INSERT INTO "Foo" ("id") VALUES (1);' + assert _keywords(tmp_path, sql) == () + + def test_read_only_cte_passes(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo") SELECT count(*) FROM batch;' + assert _keywords(tmp_path, sql) == () + + def test_cte_led_insert_values_is_bounded_and_passes(self, tmp_path): + sql = 'WITH latest AS (SELECT max("id") AS "id" FROM "Bar")\nINSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT "id"::text FROM latest));' + assert _keywords(tmp_path, sql) == () + + def test_a_writable_cte_bounded_by_values_passes(self, tmp_path): + sql = 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id") SELECT * FROM added;' + assert _keywords(tmp_path, sql) == () + + def test_a_writable_cte_copying_a_query_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_bounded_writable_cte_does_not_hide_a_copying_one_beside_it(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id"),' + ' copied AS (INSERT INTO "Baz" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added, copied;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_writable_cte_wrapping_its_row_source_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + +class TestDollarQuotedBlocks: + def test_update_inside_do_block_is_flagged(self, tmp_path): + sql = 'DO $$\nBEGIN\n UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_conditional_ddl_do_block_passes(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'x') THEN\n" + ' ALTER TABLE "Foo" DROP CONSTRAINT "x";\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_guarded_update_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo") THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_guard_with_a_nested_call_still_flags_the_update(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo" WHERE lower("a") = \'x\' UNION SELECT 1) THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_tagged_dollar_quote_is_scanned(self, tmp_path): + sql = 'DO $body$\nBEGIN\n DELETE FROM "Foo";\nEND $body$;' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_tagged_dollar_quote_holds_an_apostrophe(self, tmp_path): + sql = 'INSERT INTO "Foo" ("t") VALUES ($body$don\'t$body$);\nUPDATE "Bar" SET "b" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_semicolons_inside_do_block_do_not_split_outer_statements(self, tmp_path): + sql = 'DO $$ BEGIN PERFORM 1; END $$;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == () + + def test_line_number_inside_a_do_block_counts_from_the_top_of_the_file(self, tmp_path): + sql = ( + "-- AlterTable\n" + 'ALTER TABLE "Foo" ADD COLUMN "b" INT;\n' + "\n" + "DO $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "b" = 1;\n' + "END $$;" + ) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_line_number_inside_a_nested_body_counts_from_the_top_of_the_file(self, tmp_path): + sql = ( + "-- CreateIndex\n" + 'CREATE INDEX "i" ON "Foo"("a");\n' + "\n" + "DO $outer$\n" + "BEGIN\n" + " EXECUTE $inner$\n" + ' UPDATE "Foo" SET "a" = 1\n' + " $inner$;\n" + "END $outer$;" + ) + assert _scan(tmp_path, sql)[0].line == 7 + + +class TestStoredRoutines: + DEFINITION = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + PROCEDURE = ( + "CREATE OR REPLACE PROCEDURE sweep() AS $$\n" + "BEGIN\n" + ' DELETE FROM "Foo";\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + + def test_a_function_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION) == () + + def test_a_procedure_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE) == () + + def test_a_function_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION + "SELECT backfill();\n") == ("UPDATE",) + + def test_a_procedure_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE + "CALL sweep();\n") == ("DELETE",) + + def test_a_call_written_above_the_definition_still_counts(self, tmp_path): + assert _keywords(tmp_path, "SELECT backfill();\n" + self.DEFINITION) == ("UPDATE",) + + def test_a_call_from_inside_a_do_block_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN PERFORM backfill(); END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_through_a_quoted_identifier_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_unrelated_quoted_identifier_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'SELECT "other"();\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_column_sharing_the_routine_name_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'ALTER TABLE "Foo" ADD COLUMN "backfill" int;\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_index_target_sharing_the_routine_name_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Foo" ("backfill");\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_call_with_space_before_its_parenthesis_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill" ();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_quoted_call_with_a_block_comment_before_its_parenthesis_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill" /* reason */ ();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_quoted_call_with_a_line_comment_before_its_parenthesis_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill" -- run it\n();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_quoted_column_followed_by_a_comment_is_still_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'ALTER TABLE "Foo" ADD COLUMN "backfill" /* note */ int;\n' + assert _keywords(tmp_path, sql) == () + + def test_a_like_named_table_with_a_column_list_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TABLE "backfill" (id int);\n' + assert _keywords(tmp_path, sql) == () + + def test_an_insert_into_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO "backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_an_insert_into_a_like_named_table_with_a_commented_column_list_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO "backfill" /* cols */ ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_foreign_key_referencing_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TABLE "Bar" (id int REFERENCES "backfill" ("id"));\n' + assert _keywords(tmp_path, sql) == () + + def test_an_index_on_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "backfill" ("id");\n' + assert _keywords(tmp_path, sql) == () + + def test_an_if_not_exists_table_named_after_the_routine_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TABLE IF NOT EXISTS "backfill" (id int);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_copy_into_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'COPY "backfill" ("id") FROM stdin;\n' + assert _keywords(tmp_path, sql) == () + + def test_a_set_returning_call_in_from_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT * FROM "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_in_a_join_condition_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT 1 FROM "Bar" b JOIN "Baz" z ON "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_in_an_index_predicate_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Foo" ("a") WHERE "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_join_condition_call_after_an_earlier_index_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Foo" ("a");\nSELECT 1 FROM "Bar" b JOIN "Baz" z ON "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_insert_into_a_schema_qualified_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO public."backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_schema_qualified_index_on_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON public."backfill" ("id");\n' + assert _keywords(tmp_path, sql) == () + + def test_a_schema_qualified_quoted_call_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT public."backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_spaced_schema_qualifier_on_an_insert_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO public . "backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_spaced_schema_qualifier_on_an_index_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON public . "backfill" ("id");\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_schema_with_spaces_round_the_dot_stays_a_relation(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO "public" . "backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_spaced_schema_qualified_call_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT public . "backfill" ();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_schema_qualified_call_in_an_index_expression_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Bar" (public."backfill"("a"));\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_spaced_schema_qualified_call_in_an_index_expression_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Bar" (public . "backfill" ("a"));\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_trigger_wiring_the_function_up_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TRIGGER t AFTER INSERT ON "Foo" EXECUTE FUNCTION backfill();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_schema_qualified_definition_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION.replace("backfill()", "public.backfill()")) == () + + def test_a_schema_qualified_function_the_migration_calls_is_flagged(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", "public.backfill()") + "SELECT public.backfill();\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_the_name_written_only_in_a_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "-- backfill() is run by hand after the deploy\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_do_body_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN\n-- backfill() is run by hand after the deploy\nPERFORM 1;\nEND; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_do_body_block_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN /* backfill() runs later */ PERFORM 1; END; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_nested_body_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN EXECUTE $q$SELECT 1 -- backfill() runs later\n$q$; END; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_in_an_executed_literal_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN EXECUTE 'SELECT backfill()'; END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_after_a_literal_holding_comment_dashes_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN RAISE NOTICE '--'; PERFORM backfill(); END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_from_inside_a_single_quoted_do_block_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN PERFORM backfill(); END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_executed_literal_inside_a_single_quoted_do_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN EXECUTE ''SELECT backfill()''; END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_variable_run_by_execute_in_a_single_quoted_do_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'DECLARE q text; BEGIN q := ''SELECT backfill()''; EXECUTE q; END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_after_a_single_quoted_literal_holding_comment_dashes_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN RAISE NOTICE ''--''; PERFORM backfill(); END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_after_a_single_quoted_literal_opening_a_block_comment_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN RAISE NOTICE ''/*''; PERFORM backfill(); END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_executed_literal_after_a_single_quoted_comment_dash_string_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN RAISE NOTICE ''--''; EXECUTE ''SELECT backfill()''; END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_long_run_of_escaped_quotes_before_an_uncalled_definition_is_not_a_call(self, tmp_path): + escaped_quotes = "'" * 84 + sql = f"DO 'BEGIN RAISE NOTICE ''{escaped_quotes}''; PERFORM 1; END';\n" + self.DEFINITION + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_single_quoted_do_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN\n-- backfill() runs later\nPERFORM 1; END';\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_non_runnable_string_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "SELECT 'backfill() runs after the deploy';\n" + assert _keywords(tmp_path, sql) == () + + def test_a_recursive_call_does_not_count_as_the_migration_calling_it(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill(n int) RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " PERFORM backfill(n - 1);\n" + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_routine_name_is_read_rather_than_trusted(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", '"back fill"()') + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_do_block_is_not_a_routine_definition(self, tmp_path): + sql = 'DO $$ BEGIN UPDATE "Foo" SET "a" = 1; END; $$;\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_definition_written_after_another_statement_is_still_recognised(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;\n' + self.DEFINITION + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_rewrite_in_a_routine_the_migration_calls(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1; -- data-migration-ok: single config row\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + "SELECT backfill();\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_called_routine_reports_the_line_inside_its_body(self, tmp_path): + assert _scan(tmp_path, self.DEFINITION + "SELECT backfill();\n")[0].line == 3 + + +class TestLoopBodies: + def test_a_rewrite_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_delete_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' DELETE FROM "Foo" WHERE "id" = r."id";\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_join_using_in_the_loop_query_does_not_hide_the_body(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT a."id" FROM "A" a JOIN "B" b USING ("id") LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_executed_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_nested_under_a_guard_inside_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' IF r."id" > 0 THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_inside_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP FOR b IN SELECT "id" FROM "B" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_supplying_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP\n' + ' FOR b IN UPDATE "Foo" SET "x" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_a_loop_running_only_ddl_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' CREATE INDEX "i" ON "Foo"("a");\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_loop_over_a_rewrite_returning_rows_is_flagged_once(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN UPDATE "Foo" SET "a" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_marker_on_a_loop_exempts_the_rewrite_it_repeats(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + " -- data-migration-ok: one row\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_select_for_update_lock_is_not_read_as_a_loop(self, tmp_path): + sql = 'DO $$\nBEGIN\n PERFORM 1 FROM "Foo" FOR UPDATE;\nEND $$;' + assert _keywords(tmp_path, sql) == () + + +class TestQuotingAndComments: + def test_update_inside_string_literal_passes(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "note" TEXT NOT NULL DEFAULT \'UPDATE nothing\';' + assert _keywords(tmp_path, sql) == () + + def test_escaped_quote_inside_string_does_not_leak(self, tmp_path): + sql = "ALTER TABLE \"Foo\" ADD COLUMN \"note\" TEXT NOT NULL DEFAULT 'it''s fine';\n" + assert _keywords(tmp_path, sql) == () + + def test_update_inside_line_comment_passes(self, tmp_path): + assert _keywords(tmp_path, '-- UPDATE "Foo" SET "a" = 1;\nDROP TABLE "Bar";') == () + + def test_update_inside_block_comment_passes(self, tmp_path): + assert _keywords(tmp_path, '/* UPDATE "Foo" SET "a" = 1; */\nDROP TABLE "Bar";') == () + + def test_nested_block_comment_passes(self, tmp_path): + sql = '/* outer /* UPDATE "Foo" SET "a" = 1; */ still comment */\nDROP TABLE "Bar";' + assert _keywords(tmp_path, sql) == () + + def test_nested_block_comment_masks_past_the_inner_close(self, tmp_path): + sql = '/* outer /* inner */ UPDATE "Foo" SET "a" = 1; */\nDROP TABLE "Bar";' + assert _keywords(tmp_path, sql) == () + + def test_update_inside_quoted_identifier_passes(self, tmp_path): + assert _keywords(tmp_path, 'ALTER TABLE "UPDATE Foo" ADD COLUMN "b" TEXT;') == () + + def test_select_in_a_quoted_identifier_does_not_make_an_insert_a_rewrite(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "SELECT Foo" ("id") VALUES (\'a\');') == () + + def test_update_in_a_quoted_identifier_does_not_make_a_cte_a_rewrite(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "UPDATE Foo") SELECT count(*) FROM batch;' + assert _keywords(tmp_path, sql) == () + + def test_positional_parameter_is_not_a_dollar_quote(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nUPDATE "Foo" SET "b" = $1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + +class TestEscapeHatch: + def test_marker_with_reason_exempts_the_statement(self, tmp_path): + sql = '-- data-migration-ok: one row per tenant, at most a few hundred\nUPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == () + + def test_marker_without_reason_does_not_exempt(self, tmp_path): + assert _keywords(tmp_path, '-- data-migration-ok:\nUPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_marker_exempts_only_its_own_statement(self, tmp_path): + sql = '-- data-migration-ok: bounded to in-flight jobs\nUPDATE "Foo" SET "a" = 1;\nUPDATE "Bar" SET "b" = 2;\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_marker_works_inside_a_do_block(self, tmp_path): + sql = 'DO $$\nBEGIN\n -- data-migration-ok: single row\n UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == () + + def test_marker_below_the_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_marker_written_below_its_statement_leaves_that_statement_flagged(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_marker_inside_a_do_block_below_the_first_line_exempts(self, tmp_path): + sql = ( + "-- AlterTable\n" + 'ALTER TABLE "Foo" ADD COLUMN "b" INT;\n' + "\n" + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: one config row\n" + ' UPDATE "Foo" SET "b" = 1;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_trailing_a_statement_exempts_that_statement(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1; -- data-migration-ok: one row\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == () + + def test_a_marker_trailing_a_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1; -- data-migration-ok: one row\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 2 + + def test_a_marker_trailing_a_multiline_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = ( + 'UPDATE "Foo"\n' + ' SET "a" = 1; -- data-migration-ok: one row\n' + 'UPDATE "Bar" SET "b" = 2;' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_marker_alone_between_two_statements_belongs_to_the_one_below_it(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: one row\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_a_marker_a_blank_line_above_a_statement_does_not_exempt_it(self, tmp_path): + sql = '-- data-migration-ok: one row\n\nUPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_trailing_marker_exempts_only_the_statement_it_follows(self, tmp_path): + sql = 'DELETE FROM "Foo" WHERE "a" = 1; UPDATE "Bar" SET "b" = 2; -- data-migration-ok: one row' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_marker_above_a_shared_line_exempts_only_the_first_statement_on_it(self, tmp_path): + sql = '-- data-migration-ok: one row\nUPDATE "Foo" SET "a" = 1; DELETE FROM "Bar" WHERE "b" = 2;' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_marker_on_the_opening_line_of_a_statement_exempts_that_statement(self, tmp_path): + sql = 'UPDATE "Foo" -- data-migration-ok: one row\n SET "a" = 1;\nDELETE FROM "Bar";' + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_marker_above_a_do_block_does_not_exempt_a_rewrite_inside_it(self, tmp_path): + sql = ( + "-- data-migration-ok: bounded, this belongs to the insert below\n" + "INSERT INTO \"Config\" (\"k\") VALUES ('x');\n" + "\n" + 'DO $$ BEGIN UPDATE "Foo" SET "b" = 1; END $$;' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_marker_directly_above_a_do_block_does_not_exempt_its_body(self, tmp_path): + sql = ( + "-- data-migration-ok: seeding two default rows\n" + "DO $$\n" + "BEGIN\n" + ' INSERT INTO "Foo" ("a") VALUES (1);\n' + ' UPDATE "Foo" SET "a" = 1;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_marker_directly_above_a_one_line_do_block_does_not_exempt_its_body(self, tmp_path): + sql = '-- data-migration-ok: seeding one default row\nDO $$ BEGIN UPDATE "Foo" SET "a" = 1; END $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 2 + + def test_marker_on_the_do_line_does_not_exempt_its_body(self, tmp_path): + sql = ( + "DO $$ -- data-migration-ok: bounded to one row\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo") THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marked_rewrite_does_not_exempt_a_later_one_in_the_same_block(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + ' UPDATE "Bar" SET "b" = 2;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + +class TestDynamicSql: + def test_execute_of_a_quoted_update_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_execute_of_a_quoted_delete_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\"';\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_execute_of_a_formatted_update_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE format('UPDATE %I SET \"a\" = 1', 'Foo');\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_execute_of_a_dollar_quoted_update_is_flagged(self, tmp_path): + sql = 'DO $outer$\nBEGIN\n EXECUTE $q$UPDATE "Foo" SET "a" = 1$q$;\nEND $outer$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_doubled_quote_inside_executed_sql_does_not_hide_the_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = date_trunc(''day'', \"t\")';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_quoted_as_data_inside_executed_sql_is_not_run(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''UPDATE \"Foo\" SET \"a\" = 1''';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_doubled_quote_does_not_split_the_literal_it_sits_in(self, tmp_path): + sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('a''UPDATE \"Bar\" SET \"a\" = 1''b');" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_following_a_doubled_quote_in_the_same_payload_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''x''; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_comment_dash_inside_a_doubled_quote_does_not_hide_a_later_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''--''; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_block_comment_open_inside_a_doubled_quote_does_not_hide_a_later_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''/*''; DELETE FROM \"Foo\"';\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_genuinely_commented_out_inside_executed_sql_is_not_run(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1 -- UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_below_escaped_quotes_in_a_multiline_payload_reports_its_own_line(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE '\n" + "SELECT ''a'', ''b'', ''c'', ''d'', ''e''\n" + "; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_in_a_later_command_before_bind_values_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1; DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_bind_value_naming_a_rewrite_is_not_run(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'INSERT INTO \"Audit\" (\"note\") VALUES ($1)'" + " USING 'DELETE FROM \"Foo\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_executed_with_bind_values_is_still_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_using_written_inside_the_command_does_not_end_it(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" USING \"Bar\"" + " WHERE \"Foo\".\"a\" = \"Bar\".\"a\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_join_using_in_a_subquery_building_the_command_does_not_end_it(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) || 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_join_using_does_not_take_the_place_of_the_real_bind_values(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) || 'UPDATE \"Foo\" SET \"a\" = $1' USING 2;\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_bind_value_naming_a_rewrite_after_a_subquery_join_is_not_run(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) USING 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_execute_of_ddl_passes(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_execute_of_a_read_only_query_passes(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_executed_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n -- data-migration-ok: one row\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_literal_that_is_not_executed_is_still_inert(self, tmp_path): + sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('UPDATE \"Bar\" SET \"a\" = 1');" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_declared_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_rewrite_assigned_in_the_body_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " stmt := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_selected_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'UPDATE \"Foo\" SET \"a\" = 1' INTO stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_selected_into_a_strict_target_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'DELETE FROM \"Foo\"' INTO STRICT stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_assigned_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " stmt = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_assigned_with_a_bare_equals_after_then_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " IF true THEN stmt = 'DELETE FROM \"Foo\"'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_declared_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_insert_target_table_is_not_read_as_an_assignment(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " audit text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " INSERT INTO audit (note) VALUES ('DELETE FROM \"Foo\" is left to the app');\n" + " EXECUTE audit;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_returned_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " INSERT INTO \"Log\" (\"sql\") VALUES ('UPDATE \"Foo\" SET \"a\" = 1')\n" + " RETURNING \"sql\" INTO stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_assigned_past_an_earlier_comparison_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + "BEGIN\n" + " IF total = 1 THEN stmt = 'DELETE FROM \"Foo\" WHERE true'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_a_rewrite_assigned_past_a_loop_comparison_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 3;\n" + "BEGIN\n" + " WHILE total >= 1 LOOP stmt = 'UPDATE \"Foo\" SET \"a\" = 1'; END LOOP;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_selected_into_a_target_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'UPDATE \"Foo\" SET \"a\" = 1'\n" + " INTO\n" + " stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_selected_into_a_strict_target_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'DELETE FROM \"Foo\"' INTO STRICT\n" + " stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_executed_by_a_name_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE\n" + " stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_compared_against_is_not_an_assignment(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " IF stmt = 'UPDATE \"Foo\" SET \"a\" = 1' THEN\n" + " RAISE NOTICE 'the application owns that one';\n" + " END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_passed_to_execute_as_a_parameter_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE 'INSERT INTO \"Log\" (\"sql\") VALUES ($1)' USING stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_compared_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total = 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " RAISE NOTICE 'purge script? %', ok;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_named_as_an_argument_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := probe_match(subject => stmt, wanted => 'DELETE FROM \"Foo\"');\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_compared_after_a_wider_comparison_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total >= 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_assigned_through_a_case_expression_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + "BEGIN\n" + " stmt := CASE WHEN total = 1 THEN 'DELETE FROM \"Foo\"' ELSE 'SELECT 1' END;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_query_reaching_into_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\"'\n" + " INTO total;\n" + " SELECT (CASE WHEN total > 0 THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\"\n" + " WHERE \"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_query_reaching_using_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\" WHERE \"a\" = $1'\n" + " USING 'k1';\n" + " SELECT (CASE WHEN true THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\" x JOIN \"Foo\" y USING (\"a\")\n" + " WHERE x.\"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_walked_by_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOR stmt IN SELECT 'DELETE FROM \"Foo\"' LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_walked_by_a_foreach_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOREACH stmt IN ARRAY ARRAY['UPDATE \"Foo\" SET \"a\" = 1'] LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_loop_over_a_query_running_nothing_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " rec record;\n" + "BEGIN\n" + " FOR rec IN SELECT \"a\" FROM \"Foo\" LOOP\n" + " RAISE NOTICE 'the DELETE FROM \"Foo\" path is the application''s: %', rec;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text;\n" + "BEGIN\n" + " SELECT 'UPDATE of legacy rows is skipped' INTO msg;\n" + " RAISE NOTICE '%', msg;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_comparing_an_executed_variable_does_not_flag_the_comparison(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " IF stmt = 'DELETE FROM \"Foo\"' THEN RAISE NOTICE 'never'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_ddl_assigned_to_a_variable_passes(self, tmp_path): + sql = "DO $$\nDECLARE\n stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nBEGIN\n EXECUTE stmt;\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_rewrite_held_in_a_variable(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " stmt text := 'UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_execute_whose_sql_starts_on_a_later_line(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " EXECUTE '\n" + " UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_assignment_whose_sql_starts_on_a_later_line(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " stmt text := '\n" + " UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_an_unmarked_execute_whose_sql_starts_on_a_later_line_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE '\n" + " UPDATE \"Foo\" SET \"a\" = 1';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_message_assigned_but_never_executed_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text := 'UPDATE of legacy rows skipped, the application backfills them';\n" + "BEGIN\n" + " RAISE NOTICE '%', msg;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_notice_naming_a_delete_it_never_runs_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " note text := 'DELETE FROM legacy rows is handled by the application';\n" + "BEGIN\n" + " RAISE NOTICE '%', note;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_do_body_in_single_quotes_is_scanned(self, tmp_path): + sql = "DO 'BEGIN UPDATE \"Foo\" SET \"a\" = 1; END';" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_a_quoted_do_body_with_a_language_clause_is_scanned(self, tmp_path): + sql = "DO LANGUAGE plpgsql 'BEGIN DELETE FROM \"Foo\"; END';" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_quoted_do_body_holding_only_ddl_passes(self, tmp_path): + sql = "DO 'BEGIN ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT; END';" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_quoted_do_body(self, tmp_path): + sql = ( + "-- data-migration-ok: one config row, keyed by its primary key\n" + "DO 'BEGIN UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''; END';" + ) + assert _keywords(tmp_path, sql) == () + + def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_fragment(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE ' || quote_ident('Foo') || ' SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_later_fragment(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'WITH x AS (SELECT 1) ' || 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_only_the_variable_that_is_executed_is_read_as_sql(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text := 'UPDATE of legacy rows skipped';\n" + " stmt text := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n" + "BEGIN\n" + " RAISE NOTICE '%', msg;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marker_on_an_execute_covers_its_single_quoted_payload(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE ' -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " ';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_on_an_execute_does_not_reach_into_a_dollar_quoted_payload(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE $x$ -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " $x$;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marker_inside_a_dollar_quoted_payload_exempts_its_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE $x$\n" + " -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " $x$;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + +class TestExplain: + def test_explain_analyze_over_an_update_is_flagged(self, tmp_path): + sql = 'EXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_explain_analyze_verbose_over_a_delete_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN ANALYZE VERBOSE DELETE FROM "Foo";') == ("DELETE",) + + def test_explain_with_a_parenthesised_analyze_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN (ANALYZE, BUFFERS) UPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_explain_analyze_over_an_insert_select_is_flagged(self, tmp_path): + sql = 'EXPLAIN ANALYZE INSERT INTO "Foo" SELECT "a" FROM "Bar";' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_explain_analyze_over_a_select_passes(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN ANALYZE SELECT * FROM "Foo";') == () + + def test_a_marker_exempts_an_explained_rewrite(self, tmp_path): + sql = '-- data-migration-ok: one config row\nEXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == () + + def test_an_analyze_of_its_own_passes(self, tmp_path): + assert _keywords(tmp_path, 'ANALYZE "Foo";') == () + + def test_a_vacuum_analyze_passes(self, tmp_path): + assert _keywords(tmp_path, 'VACUUM ANALYZE "Foo";') == () + + def test_an_explained_rewrite_inside_a_block_reports_its_line(self, tmp_path): + sql = 'DO $$\nBEGIN\n EXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + +class TestReporting: + def test_line_number_points_at_the_statement_keyword(self, tmp_path): + sql = '-- CreateIndex\nCREATE INDEX "i" ON "Foo"("a");\n\nUPDATE "Foo" SET "a" = 1;' + assert _scan(tmp_path, sql)[0].line == 4 + + def test_render_names_the_migration_and_line(self, tmp_path): + violation = _scan(tmp_path, '\n\nDELETE FROM "Foo";')[0] + rendered = violation.render() + assert "20260101000000_fixture/migration.sql:3" in rendered + assert "DELETE" in rendered + + +class TestGrandfathering: + def test_every_grandfathered_migration_still_violates(self): + for name in sorted(checker.GRANDFATHERED): + directory = checker.MIGRATIONS_DIR / name + assert directory.is_dir(), f"{name} no longer exists; drop it from GRANDFATHERED" + assert checker.scan_migration(directory), f"{name} is clean; drop it from GRANDFATHERED" + + def test_stale_entry_is_reported_when_a_migration_stops_violating(self): + found = {name: () for name in checker.GRANDFATHERED} + assert checker.stale_grandfathers(found) == tuple(sorted(checker.GRANDFATHERED)) + + def test_missing_entry_is_reported(self): + assert checker.stale_grandfathers({}) == tuple(sorted(checker.GRANDFATHERED)) + + def test_no_stale_entries_against_the_real_tree(self): + found = { + path.name: checker.scan_migration(path) + for path in checker.MIGRATIONS_DIR.iterdir() + if (path / "migration.sql").is_file() + } + assert checker.stale_grandfathers(found) == () + + +class TestShippedMigrations: + def test_the_repo_is_clean(self): + assert checker.main() == 0 + + +CLEAN = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;' +DIRTY = 'UPDATE "Foo" SET "a" = 1;' +FIXTURE = "20260101000000_fixture" + + +def _tree(monkeypatch, tmp_path: Path, sql: str, grandfathered: frozenset = frozenset()) -> None: + """Stand a migrations directory holding one fixture migration in for the repo's own. The + root moves with it, since a rendered violation names the migration relative to the root and + the two are read off the same checkout everywhere but here.""" + directory = tmp_path / "migrations" / FIXTURE + directory.mkdir(parents=True) + (directory / "migration.sql").write_text(sql, encoding="utf-8") + monkeypatch.setattr(checker, "REPO_ROOT", tmp_path) + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "migrations") + monkeypatch.setattr(checker, "GRANDFATHERED", grandfathered) + + +class TestExitCode: + def test_a_clean_tree_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN) + assert checker.main() == 0 + + def test_a_violation_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY) + assert checker.main() == 1 + + def test_a_stale_grandfather_alone_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + assert checker.main() == 1 + + def test_a_grandfathered_violation_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY, frozenset({FIXTURE})) + assert checker.main() == 0 + + def test_a_missing_migrations_directory_is_an_error(self, tmp_path, monkeypatch): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + assert checker.main() == 2 + + def test_the_failure_names_the_migration_the_line_and_the_keyword( + self, tmp_path, monkeypatch, capsys + ): + _tree(monkeypatch, tmp_path, DIRTY) + checker.main() + printed = capsys.readouterr().out + assert f"migrations/{FIXTURE}/migration.sql:1" in printed + assert "UPDATE rewrites existing rows at boot" in printed + assert checker.GUIDANCE in printed + + def test_a_stale_grandfather_is_named(self, tmp_path, monkeypatch, capsys): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + checker.main() + assert f"{FIXTURE}: listed in GRANDFATHERED" in capsys.readouterr().out + + def test_a_missing_directory_is_reported_on_stderr(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + checker.main() + assert "migrations directory not found" in capsys.readouterr().err diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 8dad4bef07b..1f1c9be973f 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -3584,6 +3584,104 @@ def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) +def test_completion_cost_bills_interactions_api_response(): + from litellm.types.interactions import InteractionsAPIResponse + + model_info = litellm.get_model_info(model="gemini-2.5-flash", custom_llm_provider="gemini") + response = InteractionsAPIResponse( + id="interactions/abc123", + model="gemini-2.5-flash", + status="completed", + steps=[], + usage={ + "total_tokens": 175, + "total_input_tokens": 100, + "input_tokens_by_modality": [{"modality": "text", "tokens": 100}], + "total_cached_tokens": 0, + "total_output_tokens": 50, + "output_tokens_by_modality": [{"modality": "text", "tokens": 50}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 25, + }, + ) + + cost = completion_cost(completion_response=response, custom_llm_provider="gemini") + + reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"] + expected = ( + 100 * model_info["input_cost_per_token"] + + 50 * model_info["output_cost_per_token"] + + 25 * reasoning_rate + ) + assert cost == pytest.approx(expected) + assert cost > 0 + + +def test_completion_cost_bills_interactions_google_search_per_query(): + from litellm.types.interactions import InteractionsAPIResponse + + model_info = litellm.get_model_info(model="gemini-3-flash-preview", custom_llm_provider="gemini") + response = InteractionsAPIResponse( + id="interactions/search123", + model="gemini-3-flash-preview", + status="completed", + steps=[], + usage={ + "total_tokens": 680, + "total_input_tokens": 103, + "input_tokens_by_modality": [{"modality": "text", "tokens": 103}], + "total_cached_tokens": 0, + "total_output_tokens": 226, + "total_tool_use_tokens": 0, + "total_thought_tokens": 351, + "grounding_tool_count": [{"type": "google_search", "count": 3}], + }, + ) + + cost = completion_cost(completion_response=response, custom_llm_provider="gemini") + + per_query_cost = model_info["search_context_cost_per_query"]["search_context_size_medium"] + reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"] + expected = ( + 103 * model_info["input_cost_per_token"] + + 226 * model_info["output_cost_per_token"] + + 351 * reasoning_rate + + 3 * per_query_cost + ) + assert model_info.get("web_search_billing_unit") == "per_query" + assert cost == pytest.approx(expected) + assert cost > 3 * per_query_cost + + +def test_completion_cost_bills_interactions_video_output_at_video_rate(): + from litellm.types.interactions import InteractionsAPIResponse + + model_info = litellm.get_model_info(model="gemini-omni-flash-preview", custom_llm_provider="gemini") + video_tokens = 5792 * 8 + response = InteractionsAPIResponse( + id="interactions/video123", + model="gemini-omni-flash-preview", + status="completed", + steps=[], + usage={ + "total_tokens": 10 + video_tokens, + "total_input_tokens": 10, + "input_tokens_by_modality": [{"modality": "text", "tokens": 10}], + "total_cached_tokens": 0, + "total_output_tokens": video_tokens, + "output_tokens_by_modality": [{"modality": "video", "tokens": video_tokens}], + "total_tool_use_tokens": 0, + "total_thought_tokens": 0, + }, + ) + + cost = completion_cost(completion_response=response, custom_llm_provider="gemini") + + expected = 10 * model_info["input_cost_per_token"] + video_tokens * model_info["output_cost_per_video_token"] + assert model_info["output_cost_per_video_token"] != model_info["output_cost_per_token"] + assert cost == pytest.approx(expected) + + @pytest.mark.parametrize( "batch_rate,expected_prompt,expected_completion", [ diff --git a/tests/test_litellm/test_e2e_egress_sentinel.py b/tests/test_litellm/test_e2e_egress_sentinel.py new file mode 100644 index 00000000000..dcbd3241686 --- /dev/null +++ b/tests/test_litellm/test_e2e_egress_sentinel.py @@ -0,0 +1,172 @@ +"""Tests for .github/scripts/e2e_egress_sentinel.py. + +The replay lane's zero-egress proof is only as good as this sentinel: it pins the +provider hosts to a local sink and counts every connection that reaches them, so +a single escaped provider call turns the run red. The contract locked in here is +that the counter counts (each accepted connection is exactly one recorded hit), +that ``assert-empty`` is the pass/fail gate around that count, and that the hosts +file it edits is always handed back exactly as it was found. +""" + +from __future__ import annotations + +import importlib.util +import socket +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Final + +_REPO_ROOT: Final = Path(__file__).resolve().parents[2] +_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "e2e_egress_sentinel.py" +_spec: Final = importlib.util.spec_from_file_location("e2e_egress_sentinel", _MODULE_PATH) +sentinel: Final = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = sentinel # @dataclass(slots=True) rebuilds via sys.modules +_spec.loader.exec_module(sentinel) + + +def _hit_lines(hits_file: Path) -> list[str]: + if not hits_file.exists(): + return [] + return [line for line in hits_file.read_text(encoding="utf-8").splitlines() if line.strip()] + + +def test_pin_block_lists_every_host_against_the_sink(): + block = sentinel._pin_block("127.0.0.1", ("api.openai.com", "api.anthropic.com")) + assert "127.0.0.1\tapi.openai.com" in block + assert "127.0.0.1\tapi.anthropic.com" in block + assert sentinel._BLOCK_BEGIN in block and sentinel._BLOCK_END in block + + +def test_install_and_restore_round_trips_an_existing_hosts_file(tmp_path): + hosts = tmp_path / "hosts" + original = "127.0.0.1\tlocalhost\n255.255.255.255\tbroadcasthost\n" + hosts.write_text(original, encoding="utf-8") + + saved = sentinel._install_pins(hosts, "127.0.0.1", ("api.openai.com",)) + assert saved == original.encode() + assert "api.openai.com" in hosts.read_text(encoding="utf-8") + + sentinel._restore_pins(hosts, saved) + assert hosts.read_text(encoding="utf-8") == original + + +def test_install_on_a_missing_hosts_file_creates_then_restore_empties(tmp_path): + hosts = tmp_path / "hosts" + saved = sentinel._install_pins(hosts, "127.0.0.1", ("api.anthropic.com",)) + assert saved == b"" + assert "api.anthropic.com" in hosts.read_text(encoding="utf-8") + + sentinel._restore_pins(hosts, saved) + assert hosts.read_text(encoding="utf-8") == "" + + +def test_assert_empty_passes_when_no_calls(tmp_path): + absent = tmp_path / "absent.jsonl" + assert sentinel.assert_empty(absent) == 0 + + empty = tmp_path / "empty.jsonl" + empty.write_text("\n \n", encoding="utf-8") + assert sentinel.assert_empty(empty) == 0 + + +def test_assert_empty_fails_when_calls_recorded(tmp_path): + hits = tmp_path / "hits.jsonl" + hits.write_text('{"port": 443, "peer": ["127.0.0.1", 5]}\n', encoding="utf-8") + assert sentinel.assert_empty(hits) == 1 + + +def test_accept_loop_records_exactly_one_hit_per_connection(tmp_path): + hits_file = tmp_path / "hits.jsonl" + hits_file.write_text("", encoding="utf-8") + listener = sentinel._bind("127.0.0.1", 0) + port = listener.getsockname()[1] + stop = threading.Event() + hits = sentinel._HitLog(path=hits_file, _lock=threading.Lock()) + worker = threading.Thread(target=sentinel._serve_socket, args=(listener, port, hits, stop), daemon=True) + worker.start() + + try: + for _ in range(3): + conn = socket.create_connection(("127.0.0.1", port), timeout=2) + conn.close() + deadline = time.time() + 3 + while time.time() < deadline and len(_hit_lines(hits_file)) < 3: + time.sleep(0.02) + finally: + stop.set() + listener.close() + worker.join(timeout=3) + + assert len(_hit_lines(hits_file)) == 3 + + +def test_serve_end_to_end_pins_counts_and_restores(tmp_path): + hosts = tmp_path / "hosts" + hosts.write_text("127.0.0.1\tlocalhost\n", encoding="utf-8") + hits = tmp_path / "hits.jsonl" + ready = tmp_path / "ready" + pidf = tmp_path / "pid" + port = _free_port() + + proc = subprocess.Popen( + [ + sys.executable, + str(_MODULE_PATH), + "serve", + "--host", + "api.openai.com", + "--host", + "api.anthropic.com", + "--sink-address", + "127.0.0.1", + "--port", + str(port), + "--hits-file", + str(hits), + "--ready-file", + str(ready), + "--pid-file", + str(pidf), + "--hosts-file", + str(hosts), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + _wait_for(lambda: ready.exists(), timeout=10) + pinned = hosts.read_text(encoding="utf-8") + assert "api.openai.com" in pinned and "api.anthropic.com" in pinned + + for _ in range(2): + conn = socket.create_connection(("127.0.0.1", port), timeout=2) + conn.close() + _wait_for(lambda: len(_hit_lines(hits)) >= 2, timeout=5) + assert sentinel.assert_empty(hits) == 1 + finally: + proc.terminate() + proc.wait(timeout=10) + + assert hosts.read_text(encoding="utf-8") == "127.0.0.1\tlocalhost\n" + assert not ready.exists() + + +def _free_port() -> int: + probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + probe.close() + return port + + +def _wait_for(predicate, *, timeout: float) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return + time.sleep(0.05) + raise AssertionError("condition never became true within the timeout") diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 99b1cc826aa..8f2b06be4b3 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2944,3 +2944,16 @@ def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): assert cost == pytest.approx( _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) ) + + +@pytest.mark.asyncio +async def test_acompletion_resolves_provider_from_api_base(): + response = await litellm.acompletion( + model="deepseek-chat", + api_base="https://api.deepseek.com/v1", + api_key="fake-key", + messages=[{"role": "user", "content": "hi"}], + mock_response="resolved", + ) + + assert response.choices[0].message.content == "resolved" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d00fbf589e3..56fd6df446f 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -12,8 +12,17 @@ import pytest import litellm +from litellm import Router from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger +from litellm.router import ( + MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, + FallbackAwareAnthropicMessagesStream, + _anthropic_stream_commits_now, + _anthropic_stream_should_decline_fallback, + _anthropic_stream_should_drop_pre_content_ping, + _is_retriable_anthropic_status, +) def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -8565,6 +8574,49 @@ async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): assert "attempted_targets" not in breadcrumb +_BREADCRUMB_CREDENTIAL_CANARY = "Bearer sk-ant-oat01-RETRY-BREADCRUMB-CANARY-doNotShip" + + +@pytest.mark.parametrize( + "container_key, request_kwargs", + [ + ( + "provider_specific_header", + { + "provider_specific_header": { + "custom_llm_provider": "openai", + "extra_headers": {"authorization": _BREADCRUMB_CREDENTIAL_CANARY}, + } + }, + ), + ( + "extra_headers", + {"extra_headers": {"authorization": _BREADCRUMB_CREDENTIAL_CANARY}}, + ), + ( + "api_key", + {"api_key": _BREADCRUMB_CREDENTIAL_CANARY}, + ), + ], +) +@pytest.mark.asyncio +async def test_retry_breadcrumbs_never_carry_a_forwarded_credential(container_key, request_kwargs): + """log_retry copies kwargs into previous_models, which reaches spend logs and logging callbacks. + Any of these kwargs can carry a client's forwarded Authorization token or a provider key, and a + breadcrumb has no diagnostic use for the raw secret. A denylist of key names is always one new + credential kwarg behind, so log_retry scrubs credential-named values by pattern instead: the + container still reaches the breadcrumb, but the raw secret never does, whatever key holds it.""" + router = _cyclic_fallback_router(num_retries=1) + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback(router, capture, **request_kwargs) + + assert router.previous_models, "no retry breadcrumbs were recorded" + dumped = json.dumps(router.previous_models, default=str) + assert container_key in dumped, "the credential-bearing kwarg never reached the breadcrumb, so this test cannot see the leak" + assert _BREADCRUMB_CREDENTIAL_CANARY not in dumped + + @pytest.mark.asyncio async def test_fallback_traceback_stays_available_at_debug_level(): """Dropping the stack from the ERROR line is only safe because the fallback path still @@ -8878,3 +8930,1398 @@ class TestAzureBaseModelFallbackLogging: deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + +def test_model_group_info_intersects_supported_reasoning_efforts(): + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/mini-like"}, + "model_info": {"id": "mini-like-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + } + return { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": True, + "supports_xhigh_reasoning_effort": False, + } + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + # opus-like offers all seven levels, mini-like lacks none/xhigh/max; only the common set survives, + # so the group never advertises an effort routing could hand to a deployment that rejects it. + assert result.supported_reasoning_efforts == ("minimal", "low", "medium", "high") + + +def test_model_group_info_reasoning_efforts_ignore_a_deployment_off_the_map(): + """The router fills every ModelInfo key, so a deployment absent from the model map arrives with + supports_reasoning None rather than with the key missing. Its synthesized entry carries no mode, + which is what separates it from a mapped non-reasoning model, and nothing being known about it is + no reason to drop the levels the rest of the group agrees on.""" + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/unmapped-model"}, + "model_info": {"id": "unmapped-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_max_reasoning_effort": True, + } + return {"key": model_name, "litellm_provider": "openai", "mode": None, "supports_reasoning": None} + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high", "max") + + +def test_model_group_info_reasoning_efforts_empty_on_a_mapped_non_reasoning_deployment(): + """A group mixing a reasoning model with one the map knows is not a reasoning model shares no + level, so it advertises none and the picker offers nothing rather than a level routing would + hand to a deployment that rejects it.""" + router = litellm.Router( + model_list=[ + { + "model_name": "mixed-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "mixed-group", + "litellm_params": {"model": "openai/plain-chat"}, + "model_info": {"id": "plain-chat-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_max_reasoning_effort": True, + } + return {"key": model_name, "litellm_provider": "openai", "mode": "chat", "supports_reasoning": None} + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="mixed-group", + user_facing_model_group_name="mixed-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == () + + +def test_model_group_info_reasoning_efforts_ignore_a_value_declared_in_model_info(): + """The group's levels are computed from its deployments, so a value an operator left in one + deployment's model_info must not seed them. Seeding let the first deployment read narrow the + whole group while the same value on any other deployment was silently ignored.""" + router = litellm.Router( + model_list=[ + { + "model_name": "declared-group", + "litellm_params": {"model": "openai/first-reasoner"}, + "model_info": {"id": "first-deployment"}, + }, + { + "model_name": "declared-group", + "litellm_params": {"model": "openai/second-reasoner"}, + "model_info": {"id": "second-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + info = { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + } + if model_id == "first-deployment": + info["supported_reasoning_efforts"] = ("high",) + return info + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="declared-group", + user_facing_model_group_name="declared-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") + + +def test_model_group_info_survives_a_junk_typed_operator_effort_value(): + """A deployment's registered model_info reads back with whatever the operator wrote under any + key, so a wrong-typed supported_reasoning_efforts must not fail the group's info. Only the + constructor's trailing override keeps the junk away from ModelGroupInfo validation.""" + router = litellm.Router( + model_list=[ + { + "model_name": "junk-declared-group", + "litellm_params": {"model": "openai/lone-reasoner"}, + "model_info": {"id": "junk-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + return { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + "supported_reasoning_efforts": "high", + } + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="junk-declared-group", + user_facing_model_group_name="junk-declared-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") + + +def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared(): + """A deployment is registered in the cost map under its own id with whatever model_info the + operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only + a mode the map supplied marks the deployment as known, or an off-map deployment carrying any + mode empties the group it sits in.""" + from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + + mapped_model = "openai/gpt-5.6-sol" + expected = resolve_supported_reasoning_efforts( + litellm.get_model_info(model=mapped_model), + deployment_is_mapped=True, + ) + assert expected + + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": mapped_model, "api_key": "sk-fake"}, + "model_info": {"id": "mapped-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/a-model-the-map-never-heard-of", "api_key": "sk-fake"}, + "model_info": {"id": "off-map-deployment", "mode": "chat"}, + }, + ] + ) + + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == expected + + +class TestAddDeploymentApiBaseProviderResolution: + def test_bare_model_with_known_api_base_initializes(self): + router = litellm.Router( + model_list=[ + { + "model_name": "groq-pinned", + "litellm_params": { + "model": "llama-3.3-70b-versatile", + "api_base": "https://api.groq.com/openai/v1", + "api_key": "fake-key", + }, + }, + { + "model_name": "deepseek-pinned", + "litellm_params": { + "model": "deepseek-chat", + "api_base": "https://api.deepseek.com/v1", + "api_key": "fake-key", + }, + }, + ] + ) + + model_list = router.get_model_list() + assert model_list is not None + assert {m["model_name"] for m in model_list} == {"groq-pinned", "deepseek-pinned"} + + def test_bare_model_with_unknown_api_base_still_raises(self): + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + litellm.Router( + model_list=[ + { + "model_name": "mystery", + "litellm_params": { + "model": "some-unknown-model", + "api_base": "https://llm.internal.example.com/v1", + "api_key": "fake-key", + }, + } + ] + ) + + def test_explicit_custom_llm_provider_beats_api_base_endpoint_match(self): + router = litellm.Router( + model_list=[ + { + "model_name": "openai-via-gateway", + "litellm_params": { + "model": "gpt-3.5-turbo", + "custom_llm_provider": "openai", + "api_base": "https://api.groq.com/openai/v1", + "api_key": "fake-key", + }, + } + ] + ) + + deployment = router.get_deployment_by_model_group_name("openai-via-gateway") + assert deployment is not None + assert deployment.litellm_params.custom_llm_provider == "openai" + +# ===================================================================== +# anthropic_messages mid-stream-fallback helpers, added for #24004 +# (mid-stream fallback not supported for anthropic_messages route type). +# +# anthropic_messages goes through _ageneric_api_call_with_fallbacks rather +# than _acompletion, so its returned iterator was never wrapped by the chat +# completions fallback handler: an SSE `event: error` frame from a native +# Anthropic/Bedrock passthrough passed through to the client silently, and a +# MidStreamFallbackError raised by the completion-bridge path's +# CustomStreamWrapper (e.g. a Vertex AI transport drop) propagated +# unhandled. +# +# Targets the helpers introduced on Router: +# - _aanthropic_messages_streaming_iterator +# - _aanthropic_messages_fallback_attempt +# - _aanthropic_messages_with_streaming_fallbacks +# - _dispatch_generic_call_type +# ===================================================================== + + +async def _anthropic_messages_empty_generator(): + return + yield # pragma: no cover - makes this an async generator + + +def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream: + """A minimal wrapper for tests that call _aanthropic_messages_fallback_attempt + directly, bypassing _aanthropic_messages_streaming_iterator.""" + return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object()) + + +def _anthropic_messages_make_router() -> Router: + return Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "api_key": "sk-test", + }, + }, + { + "model_name": "fallback", + "litellm_params": { + "model": "bedrock/anthropic.claude-sonnet-4-5", + }, + }, + ] + ) + + +class _AnthropicMessagesFakeByteStream: + """Minimal AsyncIterator[bytes], carrying _hidden_params like + AnthropicMessagesStreamingResponse does.""" + + def __init__(self, chunks: list) -> None: + self._chunks = list(chunks) + self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + async def aclose(self) -> None: + self.closed = True + + +class _AnthropicMessagesRaisingByteStream: + """Simulates the completion-bridge path: no error SSE chunk is ever + yielded, the underlying CustomStreamWrapper raises MidStreamFallbackError + directly out of the iterator instead (a Vertex AI transport drop).""" + + def __init__(self, chunks: list, error: Exception) -> None: + self._chunks = list(chunks) + self._error = error + self._hidden_params: dict = {} + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if self._chunks: + return self._chunks.pop(0) + raise self._error + + async def aclose(self) -> None: + self.closed = True + + +class _AnthropicMessagesFallbackByteStream: + def __init__(self, chunks: list, hidden_params: dict | None = None) -> None: + self._chunks = list(chunks) + self._hidden_params = hidden_params if hidden_params is not None else {} + + def __aiter__(self): + return self + + async def __anext__(self) -> bytes: + if not self._chunks: + raise StopAsyncIteration + return self._chunks.pop(0) + + +def _anthropic_messages_overloaded_error_chunk() -> bytes: + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' + ) + + +def _anthropic_messages_invalid_request_error_chunk() -> bytes: + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}\n\n' + ) + + +def _anthropic_messages_rate_limit_error_chunk() -> bytes: + return ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "rate_limit_error", "message": "Too many requests"}}\n\n' + ) + + +def _anthropic_messages_content_chunk(text: str = "hi") -> bytes: + payload = f'{{"type": "content_block_delta", "delta": {{"type": "text_delta", "text": "{text}"}}}}' + return f"event: content_block_delta\ndata: {payload}\n\n".encode() + + +def _anthropic_messages_message_start_chunk() -> bytes: + """A lifecycle/bookkeeping frame Anthropic sends before any real content - + routinely the very first event before an overload error.""" + return b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_1"}}\n\n' + + +def _anthropic_messages_ping_chunk() -> bytes: + return b'event: ping\ndata: {"type": "ping"}\n\n' + + +# -------- _aanthropic_messages_streaming_iterator (passthrough) -------- + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_iterator_passthrough(): + """Without any error chunk, the wrapper forwards every chunk unchanged + and carries the source iterator's _hidden_params through (so response + headers like Bedrock's request-id keep flowing to the client).""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] + ) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + + collected = [chunk async for chunk in wrapped] + assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] + assert wrapped._hidden_params["additional_headers"]["x-amzn-requestid"] == "req-1" + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_frames_in_order(): + """Regression: lifecycle frames held back to guard against a mid-stream + fallback must still reach the client, in order, once real content + arrives - buffering them for the fallback-safety check must not silently + drop them on the happy path.""" + router = _anthropic_messages_make_router() + message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + ) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + + collected = [chunk async for chunk in wrapped] + assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] + + +@pytest.mark.asyncio +async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_stream_end(): + """Regression: if the primary stream ends with only lifecycle frames and + no content and no error, the buffered frames must still reach the + client rather than being silently swallowed.""" + router = _anthropic_messages_make_router() + message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + + collected = [chunk async for chunk in wrapped] + assert collected == [_anthropic_messages_message_start_chunk(), message_stop] + + with pytest.raises(StopAsyncIteration): + await wrapped.__anext__() + + +@pytest.mark.asyncio +async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_chunk_skips_fallback(): + """Greptile review round: transport-level buffering can coalesce a real + content_block_delta and a following retriable error into ONE physical + read from the source iterator. Since the whole chunk (content and error + together) is forwarded to the client atomically, the client genuinely + receives the content - so no fallback must be attempted, exactly as if + the two events had arrived as separate reads.""" + router = _anthropic_messages_make_router() + coalesced_chunk = _anthropic_messages_content_chunk("partial") + _anthropic_messages_overloaded_error_chunk() + source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [coalesced_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_ping_keepalive_never_buffered_or_forwarded(): + """Bugbot regression: a `ping` keepalive carries no content and must be + dropped outright before any real content arrives, rather than buffered - + otherwise a slow-starting connection sending many pings could grow the + pre-content buffer without bound.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_ping_chunk(), _anthropic_messages_content_chunk("hi")] + ) + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + collected = [chunk async for chunk in wrapped] + + assert _anthropic_messages_ping_chunk() not in collected + assert collected == [_anthropic_messages_content_chunk("hi")] + + +@pytest.mark.asyncio +async def test_anthropic_messages_pre_content_buffer_cap_forces_commit(): + """Bugbot regression: a hostile or pathological upstream that never emits + real content or an error must not grow the pre-content lifecycle buffer + without bound - hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS commits + to the primary stream early, exactly as real content arriving would.""" + router = _anthropic_messages_make_router() + lifecycle_chunk = _anthropic_messages_message_start_chunk() + error_chunk = _anthropic_messages_overloaded_error_chunk() + chunks = [lifecycle_chunk] * (MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5) + [error_chunk] + source = _AnthropicMessagesFakeByteStream(chunks) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + collected = [chunk async for chunk in wrapped] + + mock_fallback.assert_not_awaited() + assert collected.count(lifecycle_chunk) == MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5 + assert collected[-1] == error_chunk + + +@pytest.mark.asyncio +async def test_anthropic_messages_ping_coalesced_with_content_in_one_physical_chunk_is_forwarded(): + """Greptile/Bugbot regression: transport-level buffering can coalesce a + `ping` keepalive and a real content_block_delta into ONE physical read. + The pre-content ping-drop must only discard PURE ping frames - dropping + the whole coalesced chunk would silently lose generated content.""" + router = _anthropic_messages_make_router() + coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_content_chunk("hi") + source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) + + wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) + collected = [chunk async for chunk in wrapped] + + assert collected == [coalesced_chunk] + + +@pytest.mark.asyncio +async def test_anthropic_messages_ping_coalesced_with_retriable_error_still_falls_back(): + """Greptile/Bugbot regression: a physical chunk coalescing a `ping` with a + retriable `event: error` must not be discarded as a keepalive - the error + inside it must still trigger the mid-stream fallback.""" + router = _anthropic_messages_make_router() + coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_overloaded_error_chunk() + source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, initial_kwargs={"model": "primary"} + ) + collected = [chunk async for chunk in wrapped] + + mock_fallback.assert_awaited_once() + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + + +# -------- _aanthropic_messages_fallback_attempt -------- + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_yields_fallback_stream(): + """Direct-call regression: the fallback-attempt helper re-enters the + Router's fallback chain and forwards whatever the fallback produces.""" + router = _anthropic_messages_make_router() + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + collected = [ + chunk + async for chunk in router._aanthropic_messages_fallback_attempt( + error, + {"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, + _anthropic_messages_make_wrapper(), + ) + ] + + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + mock_fallback.assert_awaited_once() + assert mock_fallback.await_args.kwargs["e"] is error + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_raises_original_exception_on_double_failure(): + """Direct-call regression: when the fallback attempt itself fails with a + MidStreamFallbackError wrapping a real provider exception, that real + exception must surface rather than the internal wrapper exception.""" + router = _anthropic_messages_make_router() + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + original_exception = litellm.APIError( + status_code=503, message="fallback also overloaded", llm_provider="bedrock", model="fallback" + ) + fallback_failure = MidStreamFallbackError( + message="fallback failed", model="fallback", llm_provider="bedrock", original_exception=original_exception + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=fallback_failure), + ): + with pytest.raises(litellm.APIError) as exc_info: + async for _ in router._aanthropic_messages_fallback_attempt( + error, {"model": "primary"}, _anthropic_messages_make_wrapper() + ): + pass + + assert exc_info.value is original_exception + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_yields_non_streaming_fallback_response(): + """Bugbot regression: a fallback that resolves to a non-streaming + response (no __aiter__, e.g. an agentic tool-use interception loop) must + be synthesized into a valid SSE byte sequence, not yielded as a raw dict + into a byte stream - the generator is typed AsyncGenerator[bytes, None] + and every item reaching the client must be a real SSE frame.""" + router = _anthropic_messages_make_router() + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + non_streaming_response = {"id": "msg_1", "type": "message", "content": [{"type": "text", "text": "hi"}]} + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=non_streaming_response), + ): + collected = [ + item + async for item in router._aanthropic_messages_fallback_attempt( + error, {"model": "primary"}, _anthropic_messages_make_wrapper() + ) + ] + + assert all(isinstance(item, bytes) for item in collected) + event_types = [item.split(b"\n")[0].removeprefix(b"event: ") for item in collected] + assert event_types == [ + b"message_start", + b"content_block_start", + b"content_block_delta", + b"content_block_stop", + b"message_delta", + b"message_stop", + ] + assert b'"text": "hi"' in collected[2] + + +@pytest.mark.asyncio +async def test_aanthropic_messages_fallback_attempt_reraises_plain_exception_on_double_failure(): + """Direct-call regression: when the fallback attempt fails with a plain + exception (not a MidStreamFallbackError), that exception itself must + propagate unchanged.""" + router = _anthropic_messages_make_router() + error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") + fallback_failure = ValueError("no healthy deployments") + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=fallback_failure), + ): + with pytest.raises(ValueError, match="no healthy deployments") as exc_info: + async for _ in router._aanthropic_messages_fallback_attempt( + error, {"model": "primary"}, _anthropic_messages_make_wrapper() + ): + pass + + assert exc_info.value is fallback_failure + + +# -------- _aanthropic_messages_with_streaming_fallbacks -------- + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_non_streaming_passthrough(): + """A non-streaming response (plain dict) is returned unchanged, never wrapped.""" + router = _anthropic_messages_make_router() + plain_response = {"id": "msg_1", "type": "message"} + + async def fake_original(**_kwargs): + return plain_response + + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value=plain_response), + ): + out = await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=False, + ) + assert out is plain_response + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_wraps_streaming_iterator(): + """A streaming response is wrapped via _aanthropic_messages_streaming_iterator.""" + router = _anthropic_messages_make_router() + streaming_iter = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk()]) + wrapped_marker = object() + + async def fake_original(**_kwargs): + return streaming_iter + + with ( + patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value=streaming_iter), + ), + patch.object( + router, + "_aanthropic_messages_streaming_iterator", + new=AsyncMock(return_value=wrapped_marker), + ) as mock_wrap, + ): + out = await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=True, + ) + assert out is wrapped_marker + mock_wrap.assert_awaited_once() + + +# -------- mid-stream error handling -------- + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_on_pre_first_chunk_error_event(): + """Regression for #24004: a retriable SSE `event: error` frame + (overloaded_error/internal_server_error) that arrives before any real + content must trigger the router's fallback chain instead of passing + through to the client silently.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + mock_fallback.assert_awaited_once() + raised = mock_fallback.await_args.kwargs["e"] + assert isinstance(raised, MidStreamFallbackError) + assert raised.status_code == 503 + assert raised.is_pre_first_chunk is True + assert source.closed is True + + +@pytest.mark.asyncio +async def test_anthropic_messages_mid_stream_error_preserves_real_status_code(): + """Bugbot regression: the MidStreamFallbackError raised for a detected SSE + `event: error` frame must carry the error's REAL parsed status code + (via original_exception), not silently default to 503 for every error + type - a rate_limit_error (429) must surface as 429, not 503.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_rate_limit_error_chunk()]) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, + ) + [chunk async for chunk in wrapped] + + raised = mock_fallback.await_args.kwargs["e"] + assert isinstance(raised, MidStreamFallbackError) + assert raised.status_code == 429 + assert raised.original_exception is not None + assert raised.original_exception.status_code == 429 + assert raised.original_exception.llm_provider == "anthropic" + + +def test_merge_fallback_hidden_params_direct_call(): + """Direct-call regression: merge_fallback_hidden_params combines the + fallback's hidden params/headers with whatever was already present, + with the fallback's values winning on key collisions.""" + wrapper = FallbackAwareAnthropicMessagesStream( + _anthropic_messages_empty_generator(), + _AnthropicMessagesFakeByteStream([]), # carries {"additional_headers": {"x-amzn-requestid": "req-1"}} + ) + wrapper.merge_fallback_hidden_params( + {"model_id": "fallback-deployment"}, + {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}, + ) + assert wrapper._hidden_params["model_id"] == "fallback-deployment" + assert wrapper._hidden_params["additional_headers"] == { + "x-amzn-requestid": "req-2", + "x-fallback-only": "yes", + } + + +def test_anthropic_stream_should_drop_pre_content_ping_direct_call(): + ping = _anthropic_messages_ping_chunk() + content = _anthropic_messages_content_chunk("hi") + assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=False) is True + assert _anthropic_stream_should_drop_pre_content_ping(ping, has_generated_content=True) is False + assert _anthropic_stream_should_drop_pre_content_ping(content, has_generated_content=False) is False + + +def test_is_retriable_anthropic_status_direct_call(): + assert _is_retriable_anthropic_status(429) is True + assert _is_retriable_anthropic_status(503) is True + assert _is_retriable_anthropic_status(500) is True + assert _is_retriable_anthropic_status(400) is False + assert _is_retriable_anthropic_status(404) is False + + +def test_anthropic_stream_should_decline_fallback_direct_call(): + pre_first_chunk_error = MidStreamFallbackError( + message="overloaded", model="primary", llm_provider="anthropic", is_pre_first_chunk=True + ) + post_first_chunk_error = MidStreamFallbackError( + message="overloaded", model="primary", llm_provider="anthropic", is_pre_first_chunk=False + ) + assert _anthropic_stream_should_decline_fallback(False, pre_first_chunk_error) is False + assert _anthropic_stream_should_decline_fallback(True, pre_first_chunk_error) is True + assert _anthropic_stream_should_decline_fallback(False, post_first_chunk_error) is True + + +def test_anthropic_stream_commits_now_direct_call(): + content = _anthropic_messages_content_chunk("hi") + lifecycle_chunk = _anthropic_messages_message_start_chunk() + assert _anthropic_stream_commits_now(content, has_generated_content=False, buffered_chunk_count=0) is True + assert _anthropic_stream_commits_now(content, has_generated_content=True, buffered_chunk_count=0) is False + assert ( + _anthropic_stream_commits_now( + lifecycle_chunk, + has_generated_content=False, + buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, + ) + is True + ) + assert ( + _anthropic_stream_commits_now( + lifecycle_chunk, + has_generated_content=False, + buffered_chunk_count=MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - 1, + ) + is False + ) + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_merges_fallback_hidden_params(): + """Bugbot regression: after a successful mid-stream fallback, the + wrapper's _hidden_params must reflect the FALLBACK deployment's own + provider headers (e.g. a different Bedrock request-id), not stay + frozen on the primary's - raw bytes can't carry per-item _hidden_params + the way a ModelResponseStream/ResponsesAPI event can, so the wrapper + itself is the only place left to expose them.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_overloaded_error_chunk()] + ) # carries x-amzn-requestid: req-1 + fallback_stream = _AnthropicMessagesFallbackByteStream( + [_anthropic_messages_content_chunk("fallback answer")], + hidden_params={"additional_headers": {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}}, + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ): + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + _ = [chunk async for chunk in wrapped] + + headers = wrapped._hidden_params["additional_headers"] + assert headers["x-amzn-requestid"] == "req-2" + assert headers["x-fallback-only"] == "yes" + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_nested_metadata(): + """Bugbot regression: a shallow .copy() of kwargs still shares the + nested litellm_metadata/metadata dict objects with the primary attempt. + _update_kwargs_with_deployment mutates that dict in place with + deployment-specific fields, which must not leak into the fallback + request's metadata.""" + router = _anthropic_messages_make_router() + primary_metadata = {"model_group": "primary"} + streaming_iter_kwargs = {} + + async def fake_original(**_kwargs): + # Simulate _update_kwargs_with_deployment mutating the primary's + # litellm_metadata in place, as the real helper does. + primary_metadata["deployment"] = "primary-deployment-object" + return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) + + with patch.object( + router, + "_aanthropic_messages_streaming_iterator", + new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), + ): + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(side_effect=fake_original), + ): + await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=True, + litellm_metadata=primary_metadata, + ) + + fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] + assert fallback_kwargs["litellm_metadata"] is not primary_metadata + assert "deployment" not in fallback_kwargs["litellm_metadata"] + + +@pytest.mark.asyncio +async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_metadata_field(): + """Same regression as above for the (separate) `metadata` kwarg some + call sites use instead of `litellm_metadata`.""" + router = _anthropic_messages_make_router() + primary_metadata = {"tag": "primary"} + streaming_iter_kwargs = {} + + async def fake_original(**_kwargs): + primary_metadata["deployment"] = "primary-deployment-object" + return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) + + with patch.object( + router, + "_aanthropic_messages_streaming_iterator", + new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), + ): + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(side_effect=fake_original), + ): + await router._aanthropic_messages_with_streaming_fallbacks( + original_function=fake_original, + model="primary", + stream=True, + metadata=primary_metadata, + ) + + fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] + assert fallback_kwargs["metadata"] is not primary_metadata + assert "deployment" not in fallback_kwargs["metadata"] + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_triggers_after_lifecycle_only_frame(): + """Regression: Anthropic routinely sends a message_start lifecycle frame + before an overload error even fires. A lifecycle-only frame (no real + content) must not disqualify the fallback attempt, and must not reach + the client either - forwarding it and then appending the fallback's own + message_start would produce two overlapping message lifecycles on one + SSE stream. The primary's buffered lifecycle frame is discarded and the + client sees only the fallback's own, single, clean lifecycle.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream( + [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] + ) + fallback_message_start = b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_2"}}\n\n' + fallback_stream = _AnthropicMessagesFallbackByteStream( + [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] + assert collected.count(_anthropic_messages_message_start_chunk()) == 0, ( + "the primary's message_start must never reach the client" + ) + assert sum(1 for c in collected if c.startswith(b"event: message_start")) == 1, ( + "exactly one message_start must reach the client" + ) + mock_fallback.assert_awaited_once() + raised = mock_fallback.await_args.kwargs["e"] + assert raised.is_pre_first_chunk is True + + +@pytest.mark.asyncio +async def test_anthropic_messages_raised_error_after_real_content_does_not_restart_stream(): + """Regression: a MidStreamFallbackError raised directly by the source + iterator (the completion-bridge path's CustomStreamWrapper, e.g. a + transport drop) must not trigger a fallback once real content already + reached the client - that would append a second, overlapping message + lifecycle onto the same SSE stream. The original exception must + propagate to the caller instead.""" + router = _anthropic_messages_make_router() + content = _anthropic_messages_content_chunk("partial answer") + original_exception = litellm.APIError( + status_code=503, + message="stream reset", + llm_provider="vertex_ai", + model="primary", + ) + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + original_exception=original_exception, + is_pre_first_chunk=False, + ) + source = _AnthropicMessagesRaisingByteStream([content], raised_error) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [] + + async def _consume(): + async for chunk in wrapped: + collected.append(chunk) + + with pytest.raises(litellm.APIError) as exc_info: + await _consume() + + assert collected == [content] + assert exc_info.value is original_exception + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_also_catches_raised_midstream_error(): + """Regression for the completion-bridge path (deployments with no native + /v1/messages endpoint): its CustomStreamWrapper raises + MidStreamFallbackError directly (e.g. on a Vertex AI transport drop) + instead of yielding an SSE error chunk - the wrapper must catch that too.""" + router = _anthropic_messages_make_router() + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + is_pre_first_chunk=True, + ) + source = _AnthropicMessagesRaisingByteStream([], raised_error) + fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=fallback_stream), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_content_chunk("fallback answer")] + mock_fallback.assert_awaited_once() + assert mock_fallback.await_args.kwargs["e"] is raised_error + + +@pytest.mark.asyncio +async def test_anthropic_messages_non_retriable_client_error_skips_fallback(): + """A 4xx (non-429) error type (e.g. invalid_request_error) is a client + error a fallback attempt cannot fix, so it must be forwarded to the + client as-is rather than burning a fallback attempt.""" + router = _anthropic_messages_make_router() + error_chunk = _anthropic_messages_invalid_request_error_chunk() + source = _AnthropicMessagesFakeByteStream([error_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [error_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_post_first_chunk_error_skips_fallback(): + """Once content has already reached the caller, retrying would start a + second, overlapping Anthropic message lifecycle on the same SSE stream - + the error must be forwarded instead of triggering an invisible retry.""" + router = _anthropic_messages_make_router() + content = _anthropic_messages_content_chunk("partial answer") + error_chunk = _anthropic_messages_overloaded_error_chunk() + source = _AnthropicMessagesFakeByteStream([content, error_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [content, error_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_non_retriable_error_flushes_buffered_lifecycle_frames(): + """A non-retriable error arriving while lifecycle frames are still + buffered (no content seen yet) must flush those buffered frames before + forwarding the error, so the client still sees the whole primary + attempt rather than losing the buffered message_start silently.""" + router = _anthropic_messages_make_router() + error_chunk = _anthropic_messages_invalid_request_error_chunk() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), error_chunk]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [chunk async for chunk in wrapped] + + assert collected == [_anthropic_messages_message_start_chunk(), error_chunk] + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_raised_error_declined_flushes_buffered_lifecycle_frames(): + """When a raised MidStreamFallbackError is declined (source says content + was not pre-first-chunk) while lifecycle frames are still buffered, they + must be flushed to the client before the exception propagates.""" + router = _anthropic_messages_make_router() + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + is_pre_first_chunk=False, + ) + source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(), + ) as mock_fallback: + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [] + + async def _consume(): + async for chunk in wrapped: + collected.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _consume() + + assert collected == [_anthropic_messages_message_start_chunk()] + assert exc_info.value is raised_error + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_anthropic_messages_raised_error_without_original_exception_reraises_itself(): + """When a declined MidStreamFallbackError carries no original_exception, + the bare exception itself must propagate rather than being swallowed.""" + router = _anthropic_messages_make_router() + content = _anthropic_messages_content_chunk("partial answer") + raised_error = MidStreamFallbackError( + message="stream reset", + model="primary", + llm_provider="vertex_ai", + is_pre_first_chunk=False, + ) + source = _AnthropicMessagesRaisingByteStream([content], raised_error) + + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + collected = [] + + async def _consume(): + async for chunk in wrapped: + collected.append(chunk) + + with pytest.raises(MidStreamFallbackError) as exc_info: + await _consume() + + assert collected == [content] + assert exc_info.value is raised_error + + +@pytest.mark.asyncio +async def test_anthropic_messages_fallback_also_failing_raises_original_exception(): + """If the fallback attempt itself fails with a MidStreamFallbackError + wrapping a real provider exception, the client must see that real + exception, not the internal MidStreamFallbackError.""" + router = _anthropic_messages_make_router() + source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) + original_exception = litellm.APIError( + status_code=503, + message="fallback also overloaded", + llm_provider="bedrock", + model="fallback", + ) + fallback_failure = MidStreamFallbackError( + message="fallback failed", + model="fallback", + llm_provider="bedrock", + original_exception=original_exception, + ) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(side_effect=fallback_failure), + ): + wrapped = await router._aanthropic_messages_streaming_iterator( + response=source, + initial_kwargs={"model": "primary"}, + ) + with pytest.raises(litellm.APIError) as exc_info: + async for _ in wrapped: + pass + + assert exc_info.value is original_exception + + +# -------- _dispatch_generic_call_type -------- + + +@pytest.mark.asyncio +async def test_dispatch_generic_call_type_routes_anthropic_messages_through_streaming_fallbacks(): + router = _anthropic_messages_make_router() + + async def fake_original(**_kwargs): + return {"id": "msg_1"} + + with patch.object( + router, + "_aanthropic_messages_with_streaming_fallbacks", + new=AsyncMock(return_value="anthropic-result"), + ) as mock_anthropic: + out = await router._dispatch_generic_call_type( + call_type="anthropic_messages", + original_function=fake_original, + model="primary", + ) + assert out == "anthropic-result" + mock_anthropic.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_dispatch_generic_call_type_other_call_types_use_generic_fallback(): + router = _anthropic_messages_make_router() + + async def fake_original(**_kwargs): + return {"id": "file_1"} + + with patch.object( + router, + "_ageneric_api_call_with_fallbacks", + new=AsyncMock(return_value="generic-result"), + ) as mock_generic: + out = await router._dispatch_generic_call_type( + call_type="afile_delete", + original_function=fake_original, + model="primary", + ) + assert out == "generic-result" + mock_generic.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_factory_function_anthropic_messages_uses_streaming_fallback_dispatch(): + """anthropic_messages must be wired through the mid-stream-fallback-aware + path rather than the bare generic dispatch every other call type without + special handling uses.""" + router = _anthropic_messages_make_router() + wrapped = router.factory_function(litellm.anthropic_messages, call_type="anthropic_messages") + assert callable(wrapped) + + with patch.object( + router, + "_aanthropic_messages_with_streaming_fallbacks", + new=AsyncMock(return_value="ok"), + ) as mock_anthropic: + result = await wrapped(model="primary") + assert result == "ok" + mock_anthropic.assert_awaited_once() diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py new file mode 100644 index 00000000000..5a0aadf4737 --- /dev/null +++ b/tests/test_litellm/test_together_ai_model_metadata.py @@ -0,0 +1,161 @@ +import json +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + +REPO_ROOT: Final = Path(__file__).parents[2] + +CostMap = dict[str, dict[str, object]] +COST_MAP_ADAPTER: Final = TypeAdapter(CostMap) + +SERVERLESS_CHAT_MODELS: Final = ( + "together_ai/moonshotai/Kimi-K3", + "together_ai/zai-org/GLM-5.2", + "together_ai/deepseek-ai/DeepSeek-V4-Pro", + "together_ai/deepseek-ai/DeepSeek-V4-Pro-0813", + "together_ai/deepseek-ai/DeepSeek-V4-Flash-0731", + "together_ai/moonshotai/Kimi-K2.7-Code", + "together_ai/MiniMaxAI/MiniMax-M3", + "together_ai/thinkingmachines/Inkling", + "together_ai/thinkingmachines/Inkling-Small", + "together_ai/Qwen/Qwen3.8-2.4T-A95B", + "together_ai/Qwen/Qwen3.7-Max", + "together_ai/Qwen/Qwen3.7-Plus", + "together_ai/Qwen/Qwen3.6-Plus", + "together_ai/Qwen/Qwen3.5-9B", + "together_ai/nvidia/nemotron-3-ultra-550b-a55b", + "together_ai/meta-models/Muse-Glimmer-30B", + "together_ai/google/gemma-4-31B-it", + "together_ai/pearl-ai/gemma-4-31b-it", + "together_ai/google/gemma-3n-E4B-it", + "together_ai/arize-ai/qwen-2-1.5b-instruct", + "together_ai/Prism-ML/Ternary-Bonsai-27B", + "together_ai/meta-llama/Llama-Guard-4-12B", + "together_ai/openai/gpt-oss-120b", + "together_ai/openai/gpt-oss-20b", + "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo", +) + +DEPRECATED_MODELS: Final = { + "together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": "2026-07-10", + "together_ai/Qwen/Qwen3.5-397B-A17B": "2026-06-29", + "together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": "2026-06-04", + "together_ai/moonshotai/Kimi-K2.5": "2026-05-21", + "together_ai/deepseek-ai/DeepSeek-R1": "2026-05-14", + "together_ai/deepseek-ai/DeepSeek-V3.1": "2026-05-14", + "together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": "2026-04-16", + "together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": "2026-04-16", + "together_ai/zai-org/GLM-4.5-Air-FP8": "2026-04-02", + "together_ai/zai-org/GLM-4.7": "2026-04-02", + "together_ai/mistralai/Mistral-Small-24B-Instruct-2501": "2026-04-02", + "together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": "2026-04-02", + "together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": "2026-03-31", + "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": "2026-03-06", + "together_ai/moonshotai/Kimi-K2-Instruct-0905": "2026-03-06", + "together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": "2026-03-06", + "together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": "2026-02-25", + "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": "2026-02-25", + "together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": "2026-02-06", + "together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": "2026-02-06", + "together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": "2026-02-06", + "together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": "2026-02-06", + "together_ai/deepseek-ai/DeepSeek-R1-0528-tput": "2026-02-03", + "together_ai/mistralai/Mistral-7B-Instruct-v0.1": "2025-11-13", + "together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": "2025-11-13", +} + + +@pytest.fixture(scope="module") +def cost_map() -> CostMap: + with open(REPO_ROOT / "model_prices_and_context_window.json") as f: + return COST_MAP_ADAPTER.validate_python(json.load(f)) + + +@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS) +def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str): + info = cost_map.get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + assert info["litellm_provider"] == "together_ai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] >= 0 + assert info["output_cost_per_token"] >= info["input_cost_per_token"] + assert "deprecation_date" not in info + + routed_model, provider, _, _ = get_llm_provider(model=model) + assert routed_model == model.removeprefix("together_ai/") + assert provider == "together_ai" + + +def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap): + info = cost_map["together_ai/moonshotai/Kimi-K3"] + assert info["input_cost_per_token"] == 3e-06 + assert info["output_cost_per_token"] == 1.5e-05 + assert info["max_input_tokens"] == 1048576 + assert info["supports_function_calling"] is True + assert info["supports_tool_choice"] is True + assert info["supports_response_schema"] is True + assert info["supports_vision"] is True + assert info["supports_reasoning"] is True + + +def test_together_glm_52_pricing(cost_map: CostMap): + info = cost_map["together_ai/zai-org/GLM-5.2"] + assert info["input_cost_per_token"] == 1.4e-06 + assert info["output_cost_per_token"] == 4.4e-06 + assert info["supports_function_calling"] is True + assert info["supports_reasoning"] is True + + +def test_together_multilingual_e5_embedding_entry(cost_map: CostMap): + info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] + assert info["mode"] == "embedding" + assert info["input_cost_per_token"] == 2e-08 + assert info["max_input_tokens"] == 514 + assert info["output_vector_size"] == 1024 + + +def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap): + info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"] + assert info["input_cost_per_token"] == 1.04e-06 + assert info["output_cost_per_token"] == 1.04e-06 + assert info["max_input_tokens"] == 131072 + + +@pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS)) +def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str): + info = cost_map.get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + assert info.get("deprecation_date") == DEPRECATED_MODELS[model] + + +def _successor(info: dict[str, object]) -> str | None: + metadata = info.get("metadata") + if not isinstance(metadata, dict): + return None + successor = metadata.get("successor") + return successor if isinstance(successor, str) else None + + +def test_together_successor_metadata_points_at_live_models(cost_map: CostMap): + successors = { + model: successor + for model, info in cost_map.items() + if model.startswith("together_ai/") and (successor := _successor(info)) is not None + } + assert len(successors) >= 10 + for model, successor in successors.items(): + target = cost_map.get(successor) + assert target is not None, f"{model} names successor {successor} that is not in the map" + assert "deprecation_date" not in target, f"{model} names deprecated successor {successor}" + + +def test_together_backup_cost_map_in_sync(cost_map: CostMap): + with open(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json") as f: + backup = COST_MAP_ADAPTER.validate_python(json.load(f)) + together_main = {k: v for k, v in cost_map.items() if k.startswith("together_ai/")} + together_backup = {k: v for k, v in backup.items() if k.startswith("together_ai/")} + assert together_backup == together_main diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d655eb96a02..77dca9322c6 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -730,7 +730,9 @@ def validate_model_cost_values(model_data, exceptions=None): "output_cost_per_pixel", "input_cost_per_second", "output_cost_per_second", + "output_cost_per_second_480p", "output_cost_per_second_1080p", + "output_cost_per_second_4k", "input_cost_per_query", "input_cost_per_request", "input_cost_per_audio_token", @@ -860,7 +862,6 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_character_above_128k_tokens": {"type": "number"}, "input_cost_per_image": {"type": "number"}, "input_cost_per_image_above_128k_tokens": {"type": "number"}, - "input_cost_per_image_token": {"type": "number"}, "input_cost_per_video_token": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, @@ -944,7 +945,9 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_video_token": {"type": "number"}, "output_cost_per_pixel": {"type": "number"}, "output_cost_per_second": {"type": "number"}, + "output_cost_per_second_480p": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, + "output_cost_per_second_4k": {"type": "number"}, "output_cost_per_token": {"type": "number"}, "output_cost_per_token_above_128k_tokens": {"type": "number"}, "output_cost_per_token_above_200k_tokens": {"type": "number"}, @@ -993,6 +996,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_xhigh_reasoning_effort": {"type": "boolean"}, "supports_max_reasoning_effort": {"type": "boolean"}, "supports_adaptive_thinking": {"type": "boolean"}, + "supports_legacy_thinking": {"type": "boolean"}, "thinking_always_on": {"type": "boolean"}, "supports_mid_conversation_system": {"type": "boolean"}, "supports_sampling_params": {"type": "boolean"}, @@ -1004,7 +1008,6 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): }, "bedrock_converse_supports_strict_tools": {"type": "boolean"}, "tpm": {"type": "number"}, - "provider_specific_entry": {"type": "object"}, "supported_endpoints": { "type": "array", "items": { @@ -1124,6 +1127,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): exceptions = [ # Add any model IDs that should be exempt from the cost validation # Example: "expensive-model-id", + "runwayml/seedance2", # 4K output is 150 credits/second = $1.50/second ] is_valid, violations = validate_model_cost_values(actual_json, exceptions) diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index fb167a8624e..b166e902d6e 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -421,6 +421,117 @@ class TestVideoGeneration: ) assert cost == 0.5 + def test_completion_cost_video_custom_pricing_under_litellm_metadata(self): + """Video routes store deployment model_info under litellm_metadata, not metadata. + + Regression for https://github.com/BerriAI/litellm/issues/36483: custom video + pricing was silently ignored because completion_cost only read metadata. + """ + from litellm.cost_calculator import completion_cost + + mock_response = MagicMock() + mock_response.usage = {"duration_seconds": 10.0} + type(mock_response)._hidden_params = {} + + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "litellm_metadata": { + "model_info": { + "output_cost_per_video_per_second": 0.18, + } + } + } + + cost = completion_cost( + completion_response=mock_response, + model="runwayml/seedance2", + call_type="create_video", + custom_llm_provider="runwayml", + custom_pricing=True, + litellm_logging_obj=mock_logging_obj, + ) + assert abs(cost - 1.8) < 0.001 + + def test_completion_cost_video_uses_provider_reported_cost_without_custom_pricing(self): + """With no custom pricing, the provider's own reported cost wins over a duration estimate.""" + from litellm.cost_calculator import completion_cost + + mock_response = MagicMock() + mock_response.usage = { + "duration_seconds": 5.0, + "video_resolution": "720p", + "provider_reported_cost_usd": 0.31, + } + type(mock_response)._hidden_params = {} + + cost = completion_cost( + completion_response=mock_response, + model="runwayml/gen4_turbo", + call_type="create_video", + custom_llm_provider="runwayml", + ) + assert cost == 0.31 + + def test_completion_cost_video_custom_pricing_beats_provider_reported_cost(self): + """Deployment-level custom pricing overrides the provider's reported cost.""" + from litellm.cost_calculator import completion_cost + + mock_response = MagicMock() + mock_response.usage = { + "duration_seconds": 10.0, + "provider_reported_cost_usd": 0.31, + } + type(mock_response)._hidden_params = {} + + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "metadata": { + "model_info": { + "output_cost_per_video_per_second": 0.18, + } + } + } + + cost = completion_cost( + completion_response=mock_response, + model="runwayml/seedance2", + call_type="create_video", + custom_llm_provider="runwayml", + custom_pricing=True, + litellm_logging_obj=mock_logging_obj, + ) + assert abs(cost - 1.8) < 0.001 + + def test_completion_cost_video_resolution_tiers_from_cost_map(self, monkeypatch): + """The 480p/1080p/4k tier keys resolve from the shipped runwayml cost map entries.""" + from litellm.cost_calculator import completion_cost + + local_map_path = os.path.join( + os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json" + ) + with open(local_map_path, "r") as f: + monkeypatch.setattr(litellm, "model_cost", json.load(f)) + + def cost_for(model: str, resolution: str | None, duration: float) -> float: + mock_response = MagicMock() + mock_response.usage = { + "duration_seconds": duration, + **({"video_resolution": resolution} if resolution else {}), + } + type(mock_response)._hidden_params = {} + return completion_cost( + completion_response=mock_response, + model=model, + call_type="create_video", + custom_llm_provider="runwayml", + ) + + assert abs(cost_for("runwayml/seedance2", "4k", 8.0) - 12.0) < 0.001 + assert abs(cost_for("runwayml/seedance2", "1080p", 8.0) - 3.2) < 0.001 + assert abs(cost_for("runwayml/seedance2", "720p", 8.0) - 2.88) < 0.001 + assert abs(cost_for("runwayml/seedance2_5", "480p", 8.0) - 1.6) < 0.001 + assert abs(cost_for("runwayml/gen4.5", None, 8.0) - 0.96) < 0.001 + def test_video_generation_with_files(self): """Test video generation with file uploads.""" config = OpenAIVideoConfig() @@ -1998,7 +2109,7 @@ class TestVideoEdit: """Verify JSON body with video.id for POST /videos/edits.""" config = OpenAIVideoConfig() - url, data = config.transform_video_edit_request( + url, data, files = config.transform_video_edit_request( prompt="make it brighter", video_id="video_abc123", api_base="https://api.openai.com/v1/videos", @@ -2009,12 +2120,13 @@ class TestVideoEdit: assert url == "https://api.openai.com/v1/videos/edits" assert data["prompt"] == "make it brighter" assert data["video"]["id"] == "video_abc123" + assert files is None def test_video_edit_transform_request_with_extra_body(self): """Extra body params are merged into request data.""" config = OpenAIVideoConfig() - url, data = config.transform_video_edit_request( + url, data, files = config.transform_video_edit_request( prompt="darken it", video_id="video_abc123", api_base="https://api.openai.com/v1/videos", @@ -2024,6 +2136,7 @@ class TestVideoEdit: ) assert data["resolution"] == "1080p" + assert files is None def test_video_edit_mock_response(self): """video_edit returns VideoObject on mock_response.""" @@ -2049,7 +2162,7 @@ class TestVideoEdit: config = OpenAIVideoConfig() encoded_id = encode_video_id_with_provider("raw_video_id", "openai", None) - url, data = config.transform_video_edit_request( + url, data, files = config.transform_video_edit_request( prompt="test", video_id=encoded_id, api_base="https://api.openai.com/v1/videos", @@ -2059,6 +2172,7 @@ class TestVideoEdit: # The video.id in the request body should be the raw ID, not the encoded one assert data["video"]["id"] == "raw_video_id" + assert files is None class TestVideoExtension: @@ -2316,6 +2430,72 @@ def test_edit_and_extension_support_custom_provider_from_extra_body( assert captured_data["custom_llm_provider"] == "vertex_ai" +@pytest.mark.parametrize( + "handler_name, path, form", + [ + ( + "video_edit", + "/v1/videos/edits", + {"model": "my-video-model", "prompt": "brighter", "video": "video_123"}, + ), + ( + "video_extension", + "/v1/videos/extensions", + {"model": "my-video-model", "prompt": "continue", "seconds": "4", "video": "video_123"}, + ), + ], +) +@pytest.mark.asyncio +async def test_edit_and_extension_read_cached_body_after_auth_consumes_stream( + handler_name, path, form +): + from urllib.parse import urlencode + + from fastapi import Response + from starlette.requests import Request + + import litellm.proxy.video_endpoints.endpoints as endpoints + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + + body = urlencode(form).encode() + stream = {"sent": False} + + async def receive(): + if stream["sent"]: + return {"type": "http.request", "body": b"", "more_body": False} + stream["sent"] = True + return {"type": "http.request", "body": body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": path, + "headers": [ + (b"content-type", b"application/x-www-form-urlencoded"), + (b"content-length", str(len(body)).encode()), + ], + "query_string": b"", + }, + receive, + ) + + await _read_request_body(request=request) + + handler = getattr(endpoints, handler_name) + with pytest.raises(ProxyException) as exc_info: + await handler( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"), + ) + + message = str(exc_info.value) + assert "Stream consumed" not in message + assert "my-video-model" in message + + @pytest.mark.parametrize("endpoint", ["/v1/videos/edits", "/v1/videos/extensions"]) def test_edit_and_extension_route_with_encoded_video_ids( video_proxy_test_client, endpoint diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..7165bb53812 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22805 + "limit": 22804 }, "LIT002": { - "limit": 26873 + "limit": 26872 }, "LIT003": { "limit": 269 @@ -30,7 +30,7 @@ "limit": 16673 }, "LIT011": { - "limit": 5588 + "limit": 5587 }, "LIT012": { "limit": 4510 diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index b7c578d8ec6..a647eedc354 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -2018,11 +2018,6 @@ "count": 1 } }, - "src/components/shared/form/field.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, "src/components/shared/numerical_input.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2198,6 +2193,11 @@ "count": 1 } }, + "src/components/ui/field.tsx": { + "local/filename-pascal-case": { + "count": 1 + } + }, "src/components/ui/hover-card.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2218,11 +2218,6 @@ "count": 1 } }, - "src/components/ui/meter.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, "src/components/ui/popover.tsx": { "local/filename-pascal-case": { "count": 1 @@ -2253,11 +2248,6 @@ "count": 1 } }, - "src/components/ui/sidebar.tsx": { - "local/filename-pascal-case": { - "count": 1 - } - }, "src/components/ui/skeleton.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 564f32e2573..d2a64b93384 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -17,7 +17,8 @@ "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@types/papaparse": "5.5.2", - "cva": "1.0.0-beta.4", + "class-variance-authority": "0.7.1", + "clsx": "^2.1.1", "date-fns": "^4.4.0", "dayjs": "1.11.19", "jwt-decode": "4.0.0", @@ -5161,6 +5162,18 @@ "node": ">= 16" } }, + "node_modules/class-variance-authority": { + "version": "0.7.1", + "resolved": "https://registry.npmjs.org/class-variance-authority/-/class-variance-authority-0.7.1.tgz", + "integrity": "sha512-Ka+9Trutv7G8M6WT6SeiRWz792K5qEqIGEGzXKhAE6xOWAY6pPH8U+9IY3oCMv6kqTmLsv7Xh/2w2RigkePMsg==", + "license": "Apache-2.0", + "dependencies": { + "clsx": "^2.1.1" + }, + "funding": { + "url": "https://polar.sh/cva" + } + }, "node_modules/client-only": { "version": "0.0.1", "resolved": "https://registry.npmjs.org/client-only/-/client-only-0.0.1.tgz", @@ -5299,26 +5312,6 @@ "integrity": "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ==", "license": "MIT" }, - "node_modules/cva": { - "version": "1.0.0-beta.4", - "resolved": "https://registry.npmjs.org/cva/-/cva-1.0.0-beta.4.tgz", - "integrity": "sha512-F/JS9hScapq4DBVQXcK85l9U91M6ePeXoBMSp7vypzShoefUBxjQTo3g3935PUHgQd+IW77DjbPRIxugy4/GCQ==", - "license": "Apache-2.0", - "dependencies": { - "clsx": "^2.1.1" - }, - "funding": { - "url": "https://polar.sh/cva" - }, - "peerDependencies": { - "typescript": ">= 4.5.5" - }, - "peerDependenciesMeta": { - "typescript": { - "optional": true - } - } - }, "node_modules/d3-array": { "version": "3.2.4", "resolved": "https://registry.npmjs.org/d3-array/-/d3-array-3.2.4.tgz", @@ -12175,7 +12168,7 @@ "version": "5.9.3", "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", - "devOptional": true, + "dev": true, "license": "Apache-2.0", "bin": { "tsc": "bin/tsc", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ff6448ad75c..ededdfb4606 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -33,7 +33,8 @@ "@tanstack/react-query": "5.100.7", "@tanstack/react-table": "8.21.3", "@types/papaparse": "5.5.2", - "cva": "1.0.0-beta.4", + "class-variance-authority": "0.7.1", + "clsx": "^2.1.1", "date-fns": "^4.4.0", "dayjs": "1.11.19", "jwt-decode": "4.0.0", diff --git a/ui/litellm-dashboard/public/assets/logos/bing.png b/ui/litellm-dashboard/public/assets/logos/bing.png new file mode 100644 index 00000000000..ab1f4359281 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/bing.png differ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx index c9ba082f7a6..f8ec3b5e1e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupBaseForm.tsx @@ -7,7 +7,7 @@ import { z } from "zod/v4"; import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { ModelSelect } from "@/components/ModelSelect/ModelSelect"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.tsx b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.tsx index 624c85cc818..a7f2ee18521 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/access-group-create/AccessGroupCreateDialog.tsx @@ -9,7 +9,7 @@ import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { ModelSelect } from "@/components/ModelSelect/ModelSelect"; import { toast } from "@/lib/toast"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx index 976fb94acea..7dc3f57bd6f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx @@ -29,7 +29,7 @@ import { } from "@/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm"; import UIAccessControlForm from "@/components/UIAccessControlForm"; import { z } from "zod/v4"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Input } from "@/components/ui/input"; import { useZodForm } from "@/lib/forms/useZodForm"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx index e0152358f8c..15b85001cdc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx @@ -26,7 +26,7 @@ import { import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; import { Input } from "@/components/ui/input"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; -import { Field, FieldDescription, FieldError, FieldGroup, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldDescription, FieldError, FieldGroup, FieldLabel } from "@/components/ui/field"; export interface AgentSkillFormValue { id?: string; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx index b5c04029d69..51445ec1bdb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx @@ -5,6 +5,7 @@ import { Logo } from "@/components/molecules/logo/Logo"; import { Bot, Check, CircleCheck, Key, LayoutGrid } from "lucide-react"; import CreatedKeyDisplay from "@/components/shared/CreatedKeyDisplay"; import { Badge } from "@/components/ui/badge"; +import { StatusBadge } from "@/components/shared/table_cells/status_badge"; import { Button } from "@/components/ui/button"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { Input } from "@/components/ui/input"; @@ -14,7 +15,7 @@ import { Separator } from "@/components/ui/separator"; import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; import { TooltipProvider } from "@/components/ui/tooltip"; -import { Field, FieldGroup, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { createAgentCall, @@ -764,9 +765,7 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok Custom / Other - - GENERIC - + For agents that don't follow a standard protocol, just needs a virtual key @@ -935,7 +934,7 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok )} - Recommended + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx index 622e5ea7c5c..816df4d805c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_form_fields.tsx @@ -6,7 +6,7 @@ import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; -import { Field, FieldGroup, FieldTitle } from "@/components/shared/form/field"; +import { Field, FieldGroup, FieldTitle } from "@/components/ui/field"; import { AGENT_FORM_CONFIG, SKILL_FIELD_CONFIG } from "./agent_config"; import CostConfigFields, { COST_FIELD_NAMES } from "./cost_config_fields"; import { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index a8d78f9f53a..eddeeec674b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -8,7 +8,7 @@ import { Separator } from "@/components/ui/separator"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { TooltipProvider } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; -import { Field, FieldGroup, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; import { toast } from "@/lib/toast"; import { ArrowLeft } from "lucide-react"; import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "@/components/networking"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx index 85d066f4233..04a8b0df9d9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/dynamic_agent_form_fields.tsx @@ -2,7 +2,7 @@ import React from "react"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Textarea } from "@/components/ui/textarea"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { AgentCreateInfo, AgentCredentialFieldMetadata } from "@/components/networking"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { AGENT_FORM_CONFIG } from "./agent_config"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx index 5761d7308cd..50f9c7cda3c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx @@ -4,7 +4,7 @@ import { z } from "zod/v4"; import { useCreateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { applyBudgetPrecision } from "./budgetPrecision"; import { toast } from "@/lib/toast"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx index 3f9de881710..f4597a72fbb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx @@ -5,7 +5,7 @@ import { useUpdateBudget } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { budgetItem } from "@/app/(dashboard)/hooks/budgets/useBudgets"; import { applyBudgetPrecision } from "./budgetPrecision"; import { toast } from "@/lib/toast"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/ErrorDrilldown.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/ErrorDrilldown.test.tsx new file mode 100644 index 00000000000..a29d23d0af5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/ErrorDrilldown.test.tsx @@ -0,0 +1,55 @@ +import React from "react"; +import { describe, it, expect } from "vitest"; +import { render, screen } from "@testing-library/react"; +import { ErrorCodeTooltip, groupErrorBuckets, type CacheActivityErrorBucket } from "./ErrorDrilldown"; + +const BUCKETS: CacheActivityErrorBucket[] = [ + { call_type: "acompletion", error_code: "401", error_class: "AuthenticationError", count: 50 }, + { call_type: "acompletion", error_code: "429", error_class: "RateLimitError", count: 120 }, + { call_type: "acompletion", error_code: "429", error_class: "InternalServerError", count: 30 }, + { call_type: "aembedding", error_code: "500", error_class: "InternalServerError", count: 999 }, +]; + +describe("groupErrorBuckets", () => { + it("keeps only the requested call_type, totals per code, and sorts codes and classes by count desc", () => { + expect(groupErrorBuckets(BUCKETS, "acompletion")).toEqual([ + { + error_code: "429", + "Failed requests": 150, + classes: [ + { error_class: "RateLimitError", count: 120 }, + { error_class: "InternalServerError", count: 30 }, + ], + }, + { + error_code: "401", + "Failed requests": 50, + classes: [{ error_class: "AuthenticationError", count: 50 }], + }, + ]); + }); + + it("returns no data for a call_type without failures", () => { + expect(groupErrorBuckets(BUCKETS, "atranscription")).toEqual([]); + }); +}); + +describe("ErrorCodeTooltip", () => { + const datum = groupErrorBuckets(BUCKETS, "acompletion")[0]; + + it("shows the code total and one row per error class on hover", () => { + render(); + + expect(screen.getByText("Error code 429: 150 failed")).toBeInTheDocument(); + expect(screen.getByText("RateLimitError")).toBeInTheDocument(); + expect(screen.getByText("120")).toBeInTheDocument(); + expect(screen.getByText("InternalServerError")).toBeInTheDocument(); + expect(screen.getByText("30")).toBeInTheDocument(); + }); + + it("renders nothing when inactive", () => { + const { container } = render(); + + expect(container).toBeEmptyDOMElement(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/ErrorDrilldown.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/ErrorDrilldown.tsx new file mode 100644 index 00000000000..0bf670abc08 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/ErrorDrilldown.tsx @@ -0,0 +1,96 @@ +"use client"; + +import React from "react"; +import { X } from "lucide-react"; +import { BarChart } from "@/components/shared/charts"; +import type { ChartTooltipProps } from "@/components/shared/charts/chart_tooltip"; +import { Button } from "@/components/ui/button"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import type { components } from "@/lib/http/schema"; + +export type CacheActivityErrorBucket = components["schemas"]["CacheActivityErrorBucket"]; + +export const FAILED_REQUESTS_SERIES = "Failed requests"; + +export type ErrorClassCount = { + error_class: string; + count: number; +}; + +export type ErrorCodeDatum = { + error_code: string; + [FAILED_REQUESTS_SERIES]: number; + classes: ErrorClassCount[]; +}; + +export const groupErrorBuckets = (buckets: readonly CacheActivityErrorBucket[], callType: string): ErrorCodeDatum[] => { + const rows = buckets.filter((bucket) => bucket.call_type === callType); + return [...new Set(rows.map((row) => row.error_code))] + .map((errorCode) => { + const codeRows = rows.filter((row) => row.error_code === errorCode); + return { + error_code: errorCode, + [FAILED_REQUESTS_SERIES]: codeRows.reduce((total, row) => total + row.count, 0), + classes: codeRows + .map((row) => ({ error_class: row.error_class, count: row.count })) + .sort((a, b) => b.count - a.count), + }; + }) + .sort((a, b) => b[FAILED_REQUESTS_SERIES] - a[FAILED_REQUESTS_SERIES]); +}; + +export const ErrorCodeTooltip = ({ active, payload, label }: ChartTooltipProps) => { + if (!active || !payload || payload.length === 0) return null; + const datum = payload[0]?.payload as ErrorCodeDatum | undefined; + if (!datum) return null; + + return ( +
+

+ Error code {String(label)}: {datum[FAILED_REQUESTS_SERIES].toLocaleString()} failed +

+
+ {datum.classes.map((errorClass) => ( +
+ {errorClass.error_class} + + {errorClass.count.toLocaleString()} + +
+ ))} +
+
+ ); +}; + +interface ErrorDrilldownCardProps { + callType: string; + buckets: readonly CacheActivityErrorBucket[]; + valueFormatter: (value: number) => string; + onClose: () => void; +} + +export const ErrorDrilldownCard = ({ callType, buckets, valueFormatter, onClose }: ErrorDrilldownCardProps) => ( + + + Failed requests by error code: {callType} + + + +

Hover a bar to see the error classes behind that code.

+ + + +); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx index fd8dd011b05..deb819fdc04 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.test.tsx @@ -1,6 +1,6 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; -import { screen, waitFor, within } from "@testing-library/react"; +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; import { renderWithProviders } from "../../../../../tests/test-utils"; import CacheDashboard from "./cache_dashboard"; @@ -47,6 +47,11 @@ const cacheActivity = { key_aliases: ["my-key", "Unnamed Key"], models: ["gpt-5.1", "text-embedding-3-large"], }, + error_breakdown: [ + { call_type: "acompletion", error_code: "429", error_class: "RateLimitError", count: 150 }, + { call_type: "acompletion", error_code: "401", error_class: "AuthenticationError", count: 50 }, + { call_type: "aembedding", error_code: "500", error_class: "InternalServerError", count: 50 }, + ], }; const renderDashboard = () => @@ -185,6 +190,52 @@ describe("CacheDashboard cache analytics charts", () => { }); }); + it("opens the error-code drilldown for a call_type when its failed segment is clicked, and closes it again", async () => { + renderDashboard(); + const { requestsCard } = await findChartCards(); + + const redBar = Array.from(requestsCard.querySelectorAll(".recharts-bar")).find((bar) => + bar.querySelector("path.recharts-rectangle")?.getAttribute("fill")?.includes("red"), + ); + expect(redBar).toBeDefined(); + expect(screen.queryByText(/Failed requests by error code/)).not.toBeInTheDocument(); + + fireEvent.click(redBar!.querySelectorAll("path.recharts-rectangle")[0]); + + const drilldownCard = cardTitled("Failed requests by error code: acompletion"); + expect(within(drilldownCard).getAllByText("429").length).toBeGreaterThan(0); + expect(within(drilldownCard).getAllByText("401").length).toBeGreaterThan(0); + expect(within(drilldownCard).queryByText("500")).not.toBeInTheDocument(); + + fireEvent.click(within(drilldownCard).getByRole("button", { name: "Close error breakdown" })); + expect(screen.queryByText(/Failed requests by error code/)).not.toBeInTheDocument(); + }); + + it("dismisses an open drilldown when refetched data no longer has failures for that call_type", async () => { + const { rerender } = renderDashboard(); + const { requestsCard } = await findChartCards(); + + const redBar = Array.from(requestsCard.querySelectorAll(".recharts-bar")).find((bar) => + bar.querySelector("path.recharts-rectangle")?.getAttribute("fill")?.includes("red"), + ); + fireEvent.click(redBar!.querySelectorAll("path.recharts-rectangle")[0]); + expect(screen.getByText("Failed requests by error code: acompletion")).toBeInTheDocument(); + + useCacheActivity.mockReturnValue({ + data: { + ...cacheActivity, + groups: cacheActivity.groups.map((group) => + group.call_type === "acompletion" ? { ...group, failed_requests: 0 } : group, + ), + error_breakdown: cacheActivity.error_breakdown.filter((bucket) => bucket.call_type !== "acompletion"), + }, + refetch: vi.fn(), + }); + rerender(); + + expect(screen.queryByText(/Failed requests by error code/)).not.toBeInTheDocument(); + }); + it("formats y-axis ticks with compact notation", async () => { renderDashboard(); const { requestsCard, tokensCard } = await findChartCards(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx index cafa7a2ce37..befc0b3c5ca 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_dashboard.tsx @@ -27,6 +27,7 @@ import { useCacheActivity, type CacheActivityGroup } from "@/app/(dashboard)/hoo import { CacheHealthTab } from "./cache_health"; import CacheSettings from "./cache_settings"; import CoordinationRedisSettings from "./coordination_redis_settings"; +import { ErrorDrilldownCard } from "./ErrorDrilldown"; const REQUEST_SERIES = { apiRequests: "LLM API requests", @@ -48,6 +49,11 @@ const formatDateWithoutTZ = (date: Date | undefined) => { return date.toISOString().split("T")[0]; }; +const resolveDrilldownCallType = (selected: string | null, groups: readonly CacheActivityGroup[]): string | null => + selected !== null && groups.some((group) => group.call_type === selected && group.failed_requests > 0) + ? selected + : null; + function valueFormatterNumbers(number: number) { const formatter = new Intl.NumberFormat("en-US", { maximumFractionDigits: 0, @@ -73,6 +79,7 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole const anchor2 = useComboboxAnchor(); const [selectedApiKeys, setSelectedApiKeys] = useState([]); const [selectedModels, setSelectedModels] = useState([]); + const [errorDrilldownCallType, setErrorDrilldownCallType] = useState(null); const [dateValue, setDateValue] = useState({ from: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), @@ -96,6 +103,7 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole const uniqueApiKeys = activity?.filter_options.key_aliases ?? []; const uniqueModels = activity?.filter_options.models ?? []; const chartData = (activity?.groups ?? []).map(toChartDatum); + const activeDrilldownCallType = resolveDrilldownCallType(errorDrilldownCallType, activity?.groups ?? []); const handleRefreshClick = () => { refetch(); @@ -277,6 +285,9 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole Cache Hits vs API Requests +

+ Click a red failed-requests segment to see which error codes caused those failures. +

= ({ accessToken, token, userRole categories={[REQUEST_SERIES.apiRequests, REQUEST_SERIES.cacheHits, REQUEST_SERIES.failed]} colors={["sky", "teal", "red"]} yAxisWidth={48} + className="mt-2" + onValueChange={(item) => { + if (item.categoryClicked === REQUEST_SERIES.failed) setErrorDrilldownCallType(item.name); + }} /> + {activeDrilldownCallType !== null && ( + setErrorDrilldownCallType(null)} + /> + )} + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx index 23700470683..eb0d1ada42e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx @@ -6,7 +6,7 @@ import { z } from "zod/v4"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { createGuardrailCall, getGuardrailsList } from "@/components/networking"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx index c307c176122..4b5c3e77252 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_margin_form.tsx @@ -2,7 +2,7 @@ import React from "react"; import { CircleHelp } from "lucide-react"; import { Providers, provider_map } from "@/components/provider_info_helpers"; import { Logo } from "@/components/molecules/logo/Logo"; -import { Field, FieldLabel, FieldTitle } from "@/components/shared/form/field"; +import { Field, FieldLabel, FieldTitle } from "@/components/ui/field"; import { Button } from "@/components/ui/button"; import { Combobox, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx index 4f4edf6ff30..05748652180 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/add_provider_form.tsx @@ -3,7 +3,7 @@ import { CircleHelp } from "lucide-react"; import { Logo } from "@/components/molecules/logo/Logo"; import { Providers, provider_map } from "@/components/provider_info_helpers"; -import { Field, FieldGroup, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; import { Button } from "@/components/ui/button"; import { Combobox, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailFormField.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailFormField.tsx index 0edce78b8fd..0afd9aab0cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailFormField.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailFormField.tsx @@ -3,7 +3,7 @@ import { CircleHelp } from "lucide-react"; import React, { useId } from "react"; import { useController, type Control, type ControllerRenderProps, type RegisterOptions } from "react-hook-form"; -import { Field, FieldDescription, FieldError, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldDescription, FieldError, FieldLabel } from "@/components/ui/field"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx index fd8ed22867b..c64b5d7cb5c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx @@ -6,13 +6,15 @@ import { toast } from "@/lib/toast"; import { Card, CardContent } from "@/components/ui/card"; import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { GuardrailMode } from "@/components/guardrails/types"; +import { formatGuardrailMode } from "./guardrail_info_helpers"; interface GuardrailItem { guardrail_id?: string; guardrail_name: string | null; litellm_params: { guardrail: string; - mode: string; + mode: GuardrailMode; default_on: boolean; }; guardrail_info: Record | null; @@ -171,7 +173,9 @@ const GuardrailTestPlayground: React.FC = ({
Mode: - {guardrail.litellm_params.mode} + + {formatGuardrailMode(guardrail.litellm_params.mode)} +
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index b4c29bd9c40..7e59abf8e3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -18,7 +18,7 @@ import GuardrailTestPlayground from "./GuardrailTestPlayground"; import { toast } from "@/lib/toast"; import { Guardrail } from "@/components/guardrails/types"; import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { CustomCodeModal } from "./custom_code"; import GuardrailGarden from "./guardrail_garden"; import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; @@ -211,7 +211,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { label: "Name", value: guardrailToDelete?.guardrail_name }, { label: "ID", value: guardrailToDelete?.guardrail_id, code: true }, { label: "Provider", value: providerDisplayName }, - { label: "Mode", value: guardrailToDelete?.litellm_params.mode }, + { label: "Mode", value: formatGuardrailMode(guardrailToDelete?.litellm_params.mode) }, { label: "Default On", value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx index c9de758f50d..fae21f8dfc4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/TeamGuardrailsTab.tsx @@ -30,7 +30,7 @@ import TeamDropdown from "@/components/common_components/team_dropdown"; import { useRegisterGuardrail } from "@/app/(dashboard)/hooks/guardrails/useRegisterGuardrail"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { isProxyAdminRole } from "@/utils/roles"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx index 47326786890..29df7c8bf3d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/add_guardrail_form.tsx @@ -25,7 +25,7 @@ import { } from "./guardrail_info_helpers"; import { Logo } from "@/components/molecules/logo/Logo"; import { MultiSelect } from "@/components/shared/MultiSelect"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { Button } from "@/components/ui/button"; import { Combobox, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx index d77dd216c2a..8445495246d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx @@ -1,7 +1,7 @@ import React, { useEffect, useId, useState } from "react"; import { getMajorAirlines } from "@/components/networking"; -import { Field, FieldDescription, FieldGroup, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldDescription, FieldGroup, FieldLabel } from "@/components/ui/field"; import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Switch } from "@/components/ui/switch"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx index ec3d05a6907..53f1b1a03d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx @@ -15,7 +15,7 @@ import { } from "@/components/ui/dropdown-menu"; import { cn } from "@/lib/cva.config"; -import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; import { Logo } from "@/components/molecules/logo/Logo"; const CONFIG_DELETE_HINT = "Config guardrails are defined in the config file and cannot be deleted from the dashboard."; @@ -117,9 +117,14 @@ export const getGuardrailTableColumns = ({ header: "Mode", size: 130, enableSorting: false, - cell: ({ row }) => ( - {row.original.litellm_params.mode} - ), + cell: ({ row }) => { + const mode = formatGuardrailMode(row.original.litellm_params.mode); + return ( + + {mode || "-"} + + ); + }, }, { id: "default_on", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 4f1ac3e7d7a..2b90a1d8cbc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -82,6 +82,36 @@ describe("Guardrail Info", () => { expect(getByText("Settings")).toBeInTheDocument(); }); + it("should render a tag-based mode object rather than crashing the detail view", async () => { + vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "database", + }); + + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue({ + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], + }); + + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + + const { findAllByText } = render( + {}} accessToken="123" isAdmin={true} />, + ); + + expect(await findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); + }); + it("should render the provider logo from the bundled guardrail logo map", async () => { vi.mocked(networking.getGuardrailInfo).mockResolvedValue({ guardrail_id: "123", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 07e75511ecb..226b7b6064d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -14,7 +14,7 @@ import React, { useCallback, useEffect, useLayoutEffect, useState } from "react" import { useForm } from "react-hook-form"; import { toast } from "@/lib/toast"; import { Logo } from "@/components/molecules/logo/Logo"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; @@ -33,6 +33,7 @@ import { import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; import { + formatGuardrailMode, getGuardrailLogoAndName, guardrail_provider_map, skipSystemMessageToChoice, @@ -559,7 +560,9 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-

{guardrailData.litellm_params?.mode || "-"}

+

+ {formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"} +

{guardrailData.litellm_params?.default_on ? "Default On" : "Default Off"} @@ -856,7 +859,7 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose,

Mode

-
{guardrailData.litellm_params?.mode || "-"}
+
{formatGuardrailMode(guardrailData.litellm_params?.mode) || "-"}

Default On

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx index ec910673b8f..c5e07fe9624 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.test.tsx @@ -14,6 +14,7 @@ import { choiceToSkipSystemForCreate, skipToolMessageToChoice, choiceToSkipToolForCreate, + formatGuardrailMode, } from "./guardrail_info_helpers"; describe("guardrail_info_helpers", () => { @@ -210,6 +211,34 @@ describe("guardrail_info_helpers", () => { }); }); + describe("formatGuardrailMode", () => { + it("renders a single mode and a list of modes", () => { + expect(formatGuardrailMode("pre_call")).toBe("pre_call"); + expect(formatGuardrailMode(["pre_call", "post_call"])).toBe("pre_call, post_call"); + }); + + it("flattens a tag-based mode object into deduped modes instead of returning it verbatim", () => { + const mode = { + tags: { "Service-Type: internal-service": "post_call", "Service-Type: batch": ["during_call", "post_call"] }, + default: ["pre_call", "post_call"], + }; + + expect(formatGuardrailMode(mode)).toBe("pre_call, post_call, during_call (tag-based)"); + }); + + it("handles a tag-based mode with no default and with no tags", () => { + expect(formatGuardrailMode({ tags: { "team: a": "post_call" } })).toBe("post_call (tag-based)"); + expect(formatGuardrailMode({ default: "pre_call" })).toBe("pre_call (tag-based)"); + }); + + it("returns an empty string for missing or unusable modes", () => { + expect(formatGuardrailMode(undefined)).toBe(""); + expect(formatGuardrailMode(null)).toBe(""); + expect(formatGuardrailMode({})).toBe(""); + expect(formatGuardrailMode({ tags: {}, default: null })).toBe(""); + }); + }); + describe("skipSystemMessageToChoice / choiceToSkipSystemForCreate", () => { it("maps API values to form choices and back for create", () => { expect(skipSystemMessageToChoice(undefined)).toBe("inherit"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx index 12aaba0d696..83038b8e0e7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info_helpers.tsx @@ -110,6 +110,17 @@ export const toModeArray = (raw: unknown): string[] => { return []; }; +export const formatGuardrailMode = (raw: unknown): string => { + const flat: string[] = toModeArray(raw); + if (flat.length > 0) return flat.join(", "); + if (raw === null || typeof raw !== "object") return ""; + + const { tags, default: fallback } = raw as { tags?: Record; default?: unknown }; + const tagged: string[] = tags && typeof tags === "object" ? Object.values(tags).flatMap(toModeArray) : []; + const modes: string[] = Array.from(new Set([...toModeArray(fallback), ...tagged])); + return modes.length > 0 ? `${modes.join(", ")} (tag-based)` : ""; +}; + // Resolves the supported modes for the selected provider, falling back to the global list export const getSupportedModesForProvider = ( settings: { supported_modes?: string[]; supported_modes_by_provider?: Record } | null, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx index 632c2a6b0df..76c91d4c176 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_provider_fields.tsx @@ -9,7 +9,7 @@ import { getGuardrailProviderSpecificParams } from "@/components/networking"; import { MultiSelect } from "@/components/shared/MultiSelect"; import { PasswordInput } from "@/components/shared/PasswordInput"; import NumericalInput from "@/components/shared/numerical_input"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Slider } from "@/components/ui/slider"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx index ee619dc7468..561a89a191a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx @@ -46,6 +46,18 @@ describe("GuardrailTable", () => { expect(screen.getByText("m")).toBeInTheDocument(); }); + it("renders a tag-based mode object instead of crashing the table", () => { + const guardrail = makeGuardrail({ + litellm_params: { + guardrail: "bedrock", + mode: { tags: { "Service-Type: internal-service": "post_call" }, default: ["pre_call", "post_call"] }, + default_on: true, + }, + }); + render(); + expect(screen.getByText("pre_call, post_call (tag-based)")).toBeInTheDocument(); + }); + it("deletes a DB guardrail through the actions menu", async () => { const user = userEvent.setup(); const onDeleteClick = vi.fn(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx index a7a247e1495..256049975d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/llm_judge/LLMJudgeFields.tsx @@ -3,7 +3,7 @@ import { Plus, X } from "lucide-react"; import React from "react"; import { useController } from "react-hook-form"; -import { Field, FieldGroup, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldGroup, FieldLabel } from "@/components/ui/field"; import { Button } from "@/components/ui/button"; import { Combobox, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx index da264063839..cb423b435ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx @@ -16,7 +16,7 @@ import { type MountedFormValues, } from "@/components/common_components/MountedFormField"; import { requiredRule } from "@/components/common_components/formRules"; -import { Field, FieldLabel } from "@/components/shared/form/field"; +import { Field, FieldLabel } from "@/components/ui/field"; import { invertedSwitchControl, switchControl, tagsControl, textControl } from "./mcpFieldRules"; import { listControl } from "./mcpFormStore"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx index c077a513757..c637655d665 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx @@ -14,7 +14,7 @@ import { getProxyBaseUrl, } from "@/components/networking"; import { MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx index 14d69c358ba..90d38642a5d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/ToolArgumentsForm.tsx @@ -7,7 +7,7 @@ import { Input } from "@/components/ui/input"; import { Textarea } from "@/components/ui/textarea"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/components/ui/tooltip"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField, type FormFieldControlProps } from "@/components/shared/form/FormField"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import type { InputSchemaProperty } from "@/components/mcp_tools/types"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx index ea2addcd5f1..5d871664314 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/UserEnvVarsModal.tsx @@ -5,11 +5,12 @@ import { z } from "zod/v4"; import { MCPServer, MCPUserEnvVarsStatus, MCPUserEnvVarSpec } from "@/components/mcp_tools/types"; import { getMCPUserEnvVars, storeMCPUserEnvVars } from "@/components/networking"; import { toast } from "@/lib/toast"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Alert, AlertTitle } from "@/components/shared/Alert"; import { PasswordInput } from "@/components/shared/PasswordInput"; import { Badge } from "@/components/ui/badge"; +import { StatusBadge } from "@/components/shared/table_cells/status_badge"; import { Button } from "@/components/ui/button"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; @@ -133,7 +134,7 @@ const UserEnvVarsModal: React.FC = ({ server, open, acces
Set your credentials - Per-user +
{displayName} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx index 438caa2f5e6..5aec78ba926 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx @@ -381,6 +381,44 @@ describe("MCPServerEdit (true passthrough warning)", () => { }); }); +describe("MCPServerEdit (OAuth authorize temp payload)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("forwards issuer/authorization_url/token_url/registration_url to the temp OAuth session payload", async () => { + // Without these fields the ephemeral server the temp OAuth session endpoint builds has no + // admin-configured OAuth endpoints on it, discovery falls back to (and fails against) the + // plain server url, and Authorize & Fetch Token 400s with "authorization url is not + // configured" even though the saved server (and the visible form) has all four fields filled in. + render( + , + ); + + await waitFor(() => { + expect(mockOauth.getTemporaryPayload).toBeTruthy(); + }); + const payload = mockOauth.getTemporaryPayload!(); + expect(payload).toBeTruthy(); + expect(payload?.issuer).toBe("https://github.com/login/oauth"); + expect(payload?.authorization_url).toBe("https://github.com/login/oauth/authorize"); + expect(payload?.token_url).toBe("https://github.com/login/oauth/access_token"); + expect(payload?.registration_url).toBe("https://github.com/login/oauth/register"); + }); +}); + describe("MCPServerEdit (auth type switch)", () => { beforeEach(() => { vi.clearAllMocks(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 5e79b20825c..8793c45371a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -282,6 +282,10 @@ const MCPServerEdit: React.FC = ({ credentials: isClientForwardedTokenMode(values.auth_type) ? preservedAdminCredentials(values.credentials) : values.credentials, + issuer: values.issuer, + authorization_url: values.authorization_url, + token_url: values.token_url, + registration_url: values.registration_url, mcp_access_groups: values.mcp_access_groups || mcpServer.mcp_access_groups, static_headers: staticHeaders, command: values.command, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryEditModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryEditModal.tsx index cd238406237..33b75286969 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryEditModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryEditModal.tsx @@ -5,7 +5,7 @@ import React, { useEffect, useState } from "react"; import { z } from "zod/v4"; import type { MemoryRow } from "@/components/networking"; -import { FieldGroup } from "@/components/shared/form/field"; +import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Input } from "@/components/ui/input"; import { Textarea } from "@/components/ui/textarea"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx index aca1db0fbae..42be1b34f07 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx @@ -17,7 +17,7 @@ import { ComboboxValue, useComboboxAnchor, } from "@/components/ui/combobox"; -import { Meter, MeterIndicator, MeterTrack } from "@/components/ui/meter"; +import { Meter, MeterIndicator, MeterTrack } from "@/components/shared/Meter"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx index a06c5c885e3..1ac33a27186 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx @@ -165,6 +165,20 @@ describe("OrganizationsTable", () => { expect(screen.getByText("RPM: Unlimited")).toBeInTheDocument(); }); + it("renders a tpm/rpm limit of 0 as 0, never as Unlimited", () => { + render( + , + ); + + expect(screen.getByText("TPM: 0")).toBeInTheDocument(); + expect(screen.getByText("RPM: 0")).toBeInTheDocument(); + expect(screen.queryByText("TPM: Unlimited")).not.toBeInTheDocument(); + expect(screen.queryByText("RPM: Unlimited")).not.toBeInTheDocument(); + }); + it("renders loading skeletons instead of rows while loading", () => { render( - TPM: {tpm_limit ? tpm_limit : "Unlimited"} - RPM: {rpm_limit ? rpm_limit : "Unlimited"} + TPM: {tpm_limit ?? "Unlimited"} + RPM: {rpm_limit ?? "Unlimited"}
); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx index 6ae7b855f1e..4925e775ac7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatComposer.tsx @@ -94,7 +94,16 @@ export function ChatComposer({ /> )} - + { + if ((event.target as HTMLElement).closest("button")) { + return; + } + event.currentTarget.parentElement?.querySelector("[data-slot=input-group-control]")?.focus(); + }} + >
{tools}
{isLoading && onCancel ? ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx index c218db5b914..a83c11d1444 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.test.tsx @@ -87,6 +87,21 @@ describe("ChatMessageBubble", () => { expect(screen.getByText("Hi there")).toBeInTheDocument(); }); + it.each([ + { role: "user" as const, bubble: ["bg-info/10", "border-info/20"], avatar: "bg-info/20" }, + { role: "assistant" as const, bubble: ["bg-card", "border-border"], avatar: "bg-muted" }, + ])("should paint the $role surface from theme tokens, not fixed colours", ({ role, bubble, avatar }) => { + render(); + + const header = screen.getByText(role).closest("div") as HTMLElement; + const surface = header.parentElement as HTMLElement; + + expect(surface).toHaveClass(...bubble); + expect(surface).not.toHaveAttribute("style"); + expect(header.firstElementChild).toHaveClass(avatar); + expect(header.firstElementChild).not.toHaveAttribute("style"); + }); + it("should show model badge for assistant messages when model is provided", () => { render(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx index eb24463cfd3..8c54d9e89fa 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatMessageBubble.tsx @@ -46,20 +46,16 @@ function ChatMessageBubble({ return (
{/* Header: role icon + name + model badge */}
{isUser ? (