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/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/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/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 6103b1bf484..b94e91b3034 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -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, @@ -85,6 +88,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 +391,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 +426,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 diff --git a/litellm/constants.py b/litellm/constants.py index aaaddd063e7..0a1ada3bab2 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) @@ -1563,6 +1564,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/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/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3ad4c187b6d..5437ce52706 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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: @@ -3558,7 +3653,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 +3678,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 +5218,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 +5246,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() 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..826a890eca9 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 @@ -5383,12 +5384,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..23abca7d5f2 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 diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index c1a2384e693..654b51ce3d0 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -440,6 +440,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/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/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/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..d09055d7671 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -1,9 +1,54 @@ +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`)" + ) class AzureFoundryModelInfo(BaseLLMModelInfo): @@ -43,7 +88,7 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): @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/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index ce89c6c23e2..35e82f1961a 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -567,6 +567,10 @@ class AWSEventStreamDecoder: 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 +581,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/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/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index cdd81b24ca3..33871a60b21 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( @@ -7048,9 +7053,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, @@ -7058,9 +7061,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, @@ -7152,20 +7160,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, ) @@ -7825,6 +7839,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, @@ -7836,6 +7851,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, @@ -7889,9 +7905,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, @@ -7910,11 +7927,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( @@ -7934,6 +7950,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, @@ -7985,9 +8002,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, @@ -8006,11 +8024,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/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/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..d3967473f99 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6828,6 +6828,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 +6839,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 +8568,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 +8583,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 +8649,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 +8663,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 +8676,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 +8687,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 +8700,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 +8713,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 +8730,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 +8744,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 +8761,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..9f953e11df1 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, @@ -39056,6 +39333,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 +40593,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 +40626,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 +40992,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 +44076,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 +44089,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 +44228,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 +48901,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, @@ -49265,6 +49686,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 +50677,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/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/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/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/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/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/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/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/router.py b/litellm/router.py index 045fd32847c..6ee474730c9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -80,6 +80,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 @@ -375,10 +376,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: @@ -7347,7 +7355,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: 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/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..9fd5cfa96ef 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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 67eae2b4f21..73f46bd2181 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 ######################################################### @@ -3331,6 +3340,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 @@ -3844,6 +3855,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..012e8785321 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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 @@ -9111,6 +9096,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 +9138,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..9f953e11df1 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, @@ -39056,6 +39333,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 +40593,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 +40626,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 +40992,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 +44076,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 +44089,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 +44228,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 +48901,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, @@ -49265,6 +49686,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 +50677,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 3585c7a7bd3..69902ecfbbb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3020 + "limit": 3018 }, "ANN002": { "limit": 71 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2920 + "limit": 2919 }, "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 }, 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/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/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..1fb74b2b7bf 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 @@ -3762,3 +3762,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_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_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_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1c706be51fa..1cac0cce4ef 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4546,6 +4546,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 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 d4168d88818..9960b282f11 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..12ab536ed45 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", [ @@ -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/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_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..d4fcbc823a6 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 @@ -300,6 +300,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 +328,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/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_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index e2892a6ccee..cd305d8ed26 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, @@ -292,6 +294,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/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/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/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/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/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/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/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/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/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/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_router.py b/tests/test_litellm/test_router.py index d00fbf589e3..910b874c2ac 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8565,6 +8565,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 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/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/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/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 07e75511ecb..f1b5a0c61ec 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"; 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/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)/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 ? (