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/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/CLAUDE.md b/CLAUDE.md index b3383b4a895..b7bd06954d0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -79,6 +79,8 @@ Do not put names of customers or customer company names in code, PR descriptions CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI +Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: ` + Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors): - Composition over inheritance diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 664e1669834..f4d4e25859a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 19955 + "limit": 19949 }, "reportArgumentType": { "limit": 2566 @@ -54,7 +54,7 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5663 + "limit": 5661 }, "reportMissingTypeArgument": { "limit": 15555 @@ -105,10 +105,10 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 39011 + "limit": 39009 }, "reportUnknownParameterType": { - "limit": 19885 + "limit": 19883 }, "reportUnknownVariableType": { "limit": 30569 diff --git a/litellm/__init__.py b/litellm/__init__.py index e95b553c5d4..f42b8280761 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -199,6 +199,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = ( None # Fields to exclude from StandardLoggingPayload before callbacks receive it ) log_raw_request_response: bool = False +log_client_error_tracebacks: bool = False request_correlation_in_logs: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False @@ -1801,6 +1802,9 @@ if TYPE_CHECKING: from .llms.gemini.interactions.transformation import ( GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig, ) + from .llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig as VertexAIInteractionsConfig, + ) from .llms.openai.chat.o_series_transformation import ( OpenAIOSeriesConfig as OpenAIOSeriesConfig, OpenAIOSeriesConfig as OpenAIO1Config, diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 89c72acc06d..c34c9eefe85 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -242,6 +242,7 @@ LLM_CONFIG_NAMES: Final = ( "OpenRouterResponsesAPIConfig", "BedrockMantleResponsesAPIConfig", "GoogleAIStudioInteractionsConfig", + "VertexAIInteractionsConfig", "OpenAIOSeriesConfig", "AnthropicSkillsConfig", "BaseSkillsAPIConfig", @@ -977,6 +978,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.gemini.interactions.transformation", "GoogleAIStudioInteractionsConfig", ), + "VertexAIInteractionsConfig": ( + ".llms.vertex_ai.interactions.transformation", + "VertexAIInteractionsConfig", + ), "OpenAIOSeriesConfig": ( ".llms.openai.chat.o_series_transformation", "OpenAIOSeriesConfig", diff --git a/litellm/caching/redis_cluster_node_isolation.py b/litellm/caching/redis_cluster_node_isolation.py index 8b0c120e80c..ae8c78709d9 100644 --- a/litellm/caching/redis_cluster_node_isolation.py +++ b/litellm/caching/redis_cluster_node_isolation.py @@ -18,6 +18,14 @@ already does when one of its pooled connections errors), leaving every other nod connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-covered, retry-exhaustion) is unchanged from upstream, since those already carry real evidence the topology changed. + +redis-py 8.x fixed this upstream with gentler machinery than this override's +``node.disconnect()`` (which also kills connections other coroutines are mid-operation +on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per +killed connection): it marks in-use connections for reconnect only after their current +operation completes, disconnects only the idle pooled ones, and defers reinitialization +to the outer retry loop. When the installed ``ClusterNode`` has that per-connection +recovery API, the factory returns the base ``RedisCluster`` unmodified. """ import asyncio @@ -72,8 +80,16 @@ class _ClusterAttrs(Protocol): _VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"}) -def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: - """Builds the ``RedisCluster`` subclass with the per-node isolation fix. +def get_litellm_async_redis_cluster_class( + cluster_node_class: type | None = None, +) -> type["_AsyncRedisClusterType"]: + """Returns the base ``RedisCluster`` when the installed redis-py already recovers a + node-level connection error per-connection (8.x+), else builds the ``RedisCluster`` + subclass with the per-node isolation fix for older versions whose upstream branch + tears down the whole cluster client. + + ``cluster_node_class`` exists for dependency injection in tests; production callers + leave it unset and the installed ``ClusterNode`` is used. Imported lazily because this module is reachable from a base ``import litellm`` while redis is not a base dependency. Cheap to call repeatedly: the underlying redis @@ -81,7 +97,10 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: """ import redis from redis.asyncio.cluster import ( - RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin + ClusterNode as _AsyncClusterNode, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin + ) + from redis.asyncio.cluster import ( + RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # same stale-stub gap as the import above ) from redis.cluster import get_node_name from redis.commands import READ_COMMANDS @@ -98,6 +117,15 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]: from redis.exceptions import ConnectionError as _RedisConnectionError from redis.exceptions import TimeoutError as _RedisTimeoutError + node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode + if hasattr(node_class, "update_active_connections_for_reconnect"): + verbose_logger.debug( + "redis-py %s recovers a node-level connection error per-connection upstream; " + "using the base RedisCluster without litellm's node-isolation override.", + redis.__version__, + ) + return _BaseAsyncRedisCluster + if redis.__version__ not in _VERIFIED_REDIS_VERSIONS: verbose_logger.warning( "redis-py %s is not in the set this cluster-teardown-storm fix was verified " diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index b94e91b3034..17815976b4a 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -5,7 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args from openai.types.responses.custom_tool_param import CustomToolParam from openai.types.responses.response_input_param import ( @@ -35,6 +35,7 @@ from litellm.responses.sse_output_recovery import ( ) from litellm.responses.utils import normalize_responses_api_stream_options from litellm.types.llms.openai import ( + REASONING_EFFORT, ChatCompletionAnnotation, ChatCompletionReasoningItem, ChatCompletionToolCallChunk, @@ -1113,22 +1114,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" ) - # If string is passed, map with optional summary based on flag/env var - if reasoning_effort == "none": - return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") - elif reasoning_effort == "high": - return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high") - elif reasoning_effort == "xhigh": - return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") - elif reasoning_effort == "medium": + if reasoning_effort in get_args(REASONING_EFFORT): return ( - Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium") - ) - elif reasoning_effort == "low": - return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low") - elif reasoning_effort == "minimal": - return ( - Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal") + Reasoning(effort=reasoning_effort, summary="detailed") + if auto_summary_enabled + else Reasoning(effort=reasoning_effort) ) return None diff --git a/litellm/constants.py b/litellm/constants.py index d6a00294004..78aba30f9c0 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) @@ -749,6 +750,7 @@ openai_compatible_endpoints: Final[list] = [ "api.groq.com/openai/v1", "https://integrate.api.nvidia.com/v1", "api.deepseek.com/v1", + "api.together.ai/v1", "api.together.xyz/v1", "app.empower.dev/api/v1", "https://api.friendli.ai/serverless/v1", 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/utils.py b/litellm/interactions/utils.py index 8a1e8836894..3895a85061d 100644 --- a/litellm/interactions/utils.py +++ b/litellm/interactions/utils.py @@ -47,6 +47,13 @@ def get_provider_interactions_api_config( return GoogleAIStudioInteractionsConfig() + if provider in (LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value): + from litellm.llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig, + ) + + return VertexAIInteractionsConfig() + return None diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index de1092bc02f..b71ef4c8e06 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -58,6 +58,26 @@ def safe_divide( return numerator / denominator +def is_expected_client_error(exception: BaseException | None) -> bool: + """ + True when the exception maps to an HTTP 4xx status. + + ProxyException stores the status on .code (as a str), HTTPException and + litellm exceptions on .status_code. + """ + if exception is None: + return False + code: Final[object] = getattr(exception, "code", None) + status_code: Final[object] = code if code is not None else getattr(exception, "status_code", None) + if status_code is None or isinstance(status_code, bool): + return False + try: + status: Final = int(str(status_code)) + except ValueError: + return False + return 400 <= status < 500 + + def coerce_token_limit(value: object) -> int | None: """ Coerce a max_input_tokens / max_output_tokens value to an int, treating a diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index e674fc37673..005e94ebe82 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -272,6 +272,14 @@ def get_llm_provider( elif endpoint == "api.deepseek.com/v1": custom_llm_provider = "deepseek" dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY") + elif endpoint == "api.together.ai/v1" or endpoint == "api.together.xyz/v1": + custom_llm_provider = "together_ai" + dynamic_api_key = api_key or ( + get_secret_str("TOGETHER_API_KEY") + or get_secret_str("TOGETHER_AI_API_KEY") + or get_secret_str("TOGETHERAI_API_KEY") + or get_secret_str("TOGETHER_AI_TOKEN") + ) elif endpoint == "ollama.com": custom_llm_provider = "ollama" dynamic_api_key = get_secret_str("OLLAMA_API_KEY") @@ -707,7 +715,7 @@ def _get_openai_compatible_provider_info( dynamic_api_key, ) = litellm.ZAIChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "together_ai": - api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.xyz/v1" + api_base = api_base or get_secret_str("TOGETHER_AI_API_BASE") or "https://api.together.ai/v1" dynamic_api_key = api_key or ( get_secret_str("TOGETHER_API_KEY") or get_secret_str("TOGETHER_AI_API_KEY") diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5437ce52706..626af7530a4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -62,7 +62,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger -from litellm.litellm_core_utils.core_helpers import reconstruct_model_name +from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( cost_breakdown_with_guardrail, @@ -3124,6 +3124,13 @@ class Logging(LiteLLMLoggingBaseClass): if not hasattr(self, "model_call_details"): self.model_call_details = {} + if ( + self.model_call_details.get("log_event_type") == "failed_api_call" + and self.model_call_details.get("exception") is exception + and self.model_call_details.get("standard_logging_object") is not None + ): + return start_time, self.model_call_details["end_time"] + self.model_call_details["log_event_type"] = "failed_api_call" self.model_call_details["exception"] = exception self.model_call_details["traceback_exception"] = ( @@ -5455,9 +5462,10 @@ class StandardLoggingPayloadSetup: error_class: Final[str] = str(original_exception.__class__.__name__) if original_exception else "" _llm_provider_in_exception: Final = getattr(original_exception, "llm_provider", "") - # Get traceback information (first 100 lines) traceback_info = traceback_str or "" - if original_exception: + if original_exception and ( + litellm.log_client_error_tracebacks or not is_expected_client_error(original_exception) + ): tb: Final[TracebackType | None] = getattr(original_exception, "__traceback__", None) if tb: tb_lines: Final = traceback.format_tb(tb) @@ -5930,11 +5938,15 @@ def get_standard_logging_object_payload( response_model_name = final_response_obj.get("model") # For Azure Model Router, preserve the actual model in the top-level standard - # logging payload only when the user has opted in. + # logging payload. + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + requested_model: Final = kwargs.get("model") - if ( - isinstance(requested_model, str) - and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower()) + stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) + if stamped_selected_model is not None: + model_name = stamped_selected_model + elif ( + AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params) and isinstance(response_model_name, str) and response_model_name ): diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index e0ba2e13c6e..72ea85dfa33 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -470,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: @@ -508,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 @@ -535,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): @@ -578,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) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b676077ab0e..86cfbf70255 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -16,6 +16,7 @@ import litellm.types import litellm.types.llms from litellm import verbose_logger from litellm._uuid import uuid +from litellm.constants import REDACTED_BY_LITELLM from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client from litellm.types.files import get_file_extension_from_mime_type @@ -642,49 +643,6 @@ def claude_2_1_pt( return prompt -### TOGETHER AI - - -def get_model_info(token, model): - try: - headers: Final = {"Authorization": f"Bearer {token}"} - client: Final = HTTPHandler(concurrent_limit=1) - response: Final = client.get("https://api.together.xyz/models/info", headers=headers) - if response.status_code == 200: - model_info: Final = response.json() - for m in model_info: - if m["name"].lower().strip() == model.strip(): - return m["config"].get("prompt_format", None), m["config"].get("chat_template", None) - return None, None - else: - return None, None - except Exception: # safely fail a prompt template request - return None, None - - -## OLD TOGETHER AI FLOW -# def format_prompt_togetherai(messages, prompt_format, chat_template): -# if prompt_format is None: -# return default_pt(messages) - -# human_prompt, assistant_prompt = prompt_format.split("{prompt}") - -# if chat_template is not None: -# prompt = hf_chat_template( -# model=None, messages=messages, chat_template=chat_template -# ) -# elif prompt_format is not None: -# prompt = custom_prompt( -# role_dict={}, -# messages=messages, -# initial_prompt_value=human_prompt, -# final_prompt_value=assistant_prompt, -# ) -# else: -# prompt = default_pt(messages) -# return prompt - - ### IBM Granite @@ -5383,12 +5341,13 @@ def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> return raw if not isinstance(raw, str): return {} + normalized_raw: Final = "{}" if raw == REDACTED_BY_LITELLM else raw from litellm.litellm_core_utils.prompt_templates.common_utils import ( parse_tool_call_arguments, ) try: - parsed: Final = parse_tool_call_arguments(raw, tool_name=tool_name, context=context) + parsed: Final = parse_tool_call_arguments(normalized_raw, tool_name=tool_name, context=context) except ValueError as e: verbose_logger.warning("Failed to parse tool call arguments: %s", e) return {} diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 0d590e1ceba..be187b091b8 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -13,6 +13,7 @@ import inspect from typing import TYPE_CHECKING, Any, Final import litellm +from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -84,29 +85,29 @@ def _redact_tool_calls(tool_calls) -> None: for tool_call in tool_calls: function = getattr(tool_call, "function", None) if function is not None and hasattr(function, "arguments"): - function.arguments = "redacted-by-litellm" + function.arguments = REDACTED_BY_LITELLM def _redact_function_call(function_call) -> None: """Redact legacy assistant function_call arguments.""" if function_call is not None and hasattr(function_call, "arguments"): - function_call.arguments = "redacted-by-litellm" + function_call.arguments = REDACTED_BY_LITELLM def _redact_choice_content(choice): """Helper to redact content in a choice (message or delta).""" if isinstance(choice, litellm.Choices): - choice.message.content = "redacted-by-litellm" + choice.message.content = REDACTED_BY_LITELLM if hasattr(choice.message, "reasoning_content"): - choice.message.reasoning_content = "redacted-by-litellm" + choice.message.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.message, "thinking_blocks"): choice.message.thinking_blocks = None _redact_tool_calls(getattr(choice.message, "tool_calls", None)) _redact_function_call(getattr(choice.message, "function_call", None)) elif isinstance(choice, litellm.utils.StreamingChoices): - choice.delta.content = "redacted-by-litellm" + choice.delta.content = REDACTED_BY_LITELLM if hasattr(choice.delta, "reasoning_content"): - choice.delta.reasoning_content = "redacted-by-litellm" + choice.delta.reasoning_content = REDACTED_BY_LITELLM if hasattr(choice.delta, "thinking_blocks"): choice.delta.thinking_blocks = None _redact_tool_calls(getattr(choice.delta, "tool_calls", None)) @@ -117,22 +118,22 @@ def _redact_responses_api_output(output_items): """Helper to redact ResponsesAPIResponse output items.""" for output_item in output_items: if hasattr(output_item, "text"): - output_item.text = "redacted-by-litellm" + output_item.text = REDACTED_BY_LITELLM if hasattr(output_item, "content") and isinstance(output_item.content, list): for content_part in output_item.content: if hasattr(content_part, "text"): - content_part.text = "redacted-by-litellm" + content_part.text = REDACTED_BY_LITELLM # Redact reasoning items in output array if hasattr(output_item, "type") and output_item.type == "reasoning": if hasattr(output_item, "summary") and isinstance(output_item.summary, list): for summary_item in output_item.summary: if hasattr(summary_item, "text"): - summary_item.text = "redacted-by-litellm" + summary_item.text = REDACTED_BY_LITELLM if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"): - output_item.arguments = "redacted-by-litellm" + output_item.arguments = REDACTED_BY_LITELLM def _redact_responses_api_output_dict(output_items, redacted_str: str): @@ -164,7 +165,7 @@ def _redact_standard_logging_object(model_call_details: dict): if standard_logging_object is None: return - redacted_str: Final = "redacted-by-litellm" + redacted_str: Final = REDACTED_BY_LITELLM if standard_logging_object.get("messages") is not None: standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}] @@ -235,7 +236,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons copy via redact_streaming_responses_for_custom_logger instead. """ # Redact model_call_details - model_call_details["messages"] = [{"role": "user", "content": "redacted-by-litellm"}] + model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}] model_call_details["prompt"] = "" model_call_details["input"] = "" _redact_standard_logging_object(model_call_details) @@ -256,7 +257,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons or hasattr(result, "__anext__") # async generator ): # async iterator # For async objects, return a simple redacted response without deepcopy - return {"text": "redacted-by-litellm"} + return {"text": REDACTED_BY_LITELLM} if not ( isinstance(result, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse)) @@ -273,11 +274,11 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: - _redact_model_response_dict_choices(_result["choices"], "redacted-by-litellm") + _redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM) redact_vertex_ai_metadata_from_logged_object(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): - _redact_responses_api_output_dict(_result["output"], "redacted-by-litellm") + _redact_responses_api_output_dict(_result["output"], REDACTED_BY_LITELLM) elif isinstance(_result, litellm.ResponsesAPIResponse): if hasattr(_result, "output"): _redact_responses_api_output(_result.output) @@ -288,7 +289,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons if hasattr(_result, "data") and _result.data is not None: _result.data = [] else: - return {"text": "redacted-by-litellm"} + return {"text": REDACTED_BY_LITELLM} return _result diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 23abca7d5f2..26380ad0af8 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -1215,8 +1215,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): if reasoning_effort is None or reasoning_effort == "none": return None if AnthropicConfig._is_adaptive_thinking_model(model, custom_llm_provider): + # without display, Anthropic defaults adaptive thinking to + # display="omitted" and returns a blank thinking block return AnthropicThinkingParam( type="adaptive", + display="summarized", ) elif reasoning_effort == "low": return AnthropicThinkingParam( @@ -2144,7 +2147,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) @staticmethod - def _thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: + def thinking_tokens_from_usage(usage_object: Mapping[str, object]) -> int | None: details: Final = usage_object.get("output_tokens_details") if not isinstance(details, Mapping): return None @@ -2176,7 +2179,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): reported_thinking_tokens: Final = ( iteration_thinking_tokens if iteration_thinking_tokens is not None - else self._thinking_tokens_from_usage(usage_object) + else self.thinking_tokens_from_usage(usage_object) ) if reported_thinking_tokens is not None: capped_reported: Final = min(max(0, reported_thinking_tokens), completion_tokens) @@ -2199,7 +2202,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): def _sum_iteration_thinking_tokens(self, iterations: Sequence[object]) -> int | None: per_iteration: Final = tuple( - self._thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None + self.thinking_tokens_from_usage(iteration) if isinstance(iteration, Mapping) else None for iteration in iterations ) reported: Final = tuple(tokens for tokens in per_iteration if tokens is not None) diff --git a/litellm/llms/azure/search/__init__.py b/litellm/llms/azure/search/__init__.py new file mode 100644 index 00000000000..2414ba2b1e8 --- /dev/null +++ b/litellm/llms/azure/search/__init__.py @@ -0,0 +1,3 @@ +from litellm.llms.azure.search.transformation import BingGroundingSearchConfig + +__all__ = ("BingGroundingSearchConfig",) diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py new file mode 100644 index 00000000000..0754c9b1fda --- /dev/null +++ b/litellm/llms/azure/search/transformation.py @@ -0,0 +1,442 @@ +""" +Calls the Microsoft Foundry Responses API with the `bing_grounding` or `web_search` +tool to search the web (Grounding with Bing Search). + +Microsoft docs: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/how-to/tools/bing-grounding + +Setup: + 1. Set BING_GROUNDING_PROJECT_ENDPOINT to the Foundry project endpoint, e.g. + https://.services.ai.azure.com/api/projects/ + 2. Set BING_GROUNDING_MODEL to a model deployment in that project (e.g. gpt-4.1); + it runs the grounded search and its tokens are billed on that deployment + 3. Optional: set BING_GROUNDING_CONNECTION_ID to a Grounding with Bing Search + project connection id to use the `bing_grounding` tool; without it the + project's built-in `web_search` tool is used + 4. Auth: pass api_key (an Azure API key, sent in the api-key header), or set + BING_GROUNDING_TOKEN to an Entra bearer token for scope + https://ai.azure.com/.default, or configure azure-identity (AZURE_CLIENT_ID / + AZURE_CLIENT_SECRET / AZURE_TENANT_ID, managed identity, or any + DefaultAzureCredential source) and the token is minted automatically + +Usage: + response = litellm.search( + query="latest AI developments", + search_provider="bing_grounding", + max_results=5, + ) +""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal + +import httpx +from pydantic import BaseModel, ConfigDict, ValidationError + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig, + SearchResponse, + SearchResult, +) +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_DOCS_URL: Final = "https://learn.microsoft.com/en-us/azure/ai-foundry/agents/how-to/tools/bing-grounding" + +PROJECT_ENDPOINT_ENV: Final = "BING_GROUNDING_PROJECT_ENDPOINT" +MODEL_ENV: Final = "BING_GROUNDING_MODEL" +CONNECTION_ID_ENV: Final = "BING_GROUNDING_CONNECTION_ID" +TOKEN_ENV: Final = "BING_GROUNDING_TOKEN" + +ENTRA_SCOPE: Final = "https://ai.azure.com/.default" + +_RESPONSES_PATH: Final = "/openai/v1/responses" +_SNIPPET_FALLBACK_LENGTH: Final = 300 +_UPSTREAM_ERROR_STATUS: Final = 502 +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" + + +class _Annotation(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + url: str | None = None + title: str | None = None + start_index: int | None = None + end_index: int | None = None + + +class _ContentPart(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + text: str = "" + annotations: tuple[_Annotation, ...] = () + + +class _OutputItem(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + type: str = "" + content: tuple[_ContentPart, ...] = () + + +class _ErrorBody(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + message: str | None = None + + +class _IncompleteDetails(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + reason: str | None = None + + +class _ResponsesEnvelope(BaseModel): + """A Foundry Responses API body. `output` is required: a body without it is not a + Responses API response and must not be reported as a successful empty search. + + A 200 body can still carry `status` `failed` or `incomplete`; those are surfaced as + errors rather than reported as a successful empty search.""" + + model_config = ConfigDict(extra="ignore", frozen=True) + + output: tuple[_OutputItem, ...] + status: str | None = None + error: _ErrorBody | None = None + incomplete_details: _IncompleteDetails | None = None + + +class _ErrorEnvelope(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + error: _ErrorBody | None = None + + +def _unwrap_error_detail(error_message: str) -> str: + """ + Surface the human-readable message inside Foundry's error envelope. + + Tool failures nest a second JSON document as a string inside `error.message` + (observed live for `bing_grounding` connection errors), so the unwrap runs twice. + Falls back to the raw body for anything else. + """ + try: + envelope: Final = _ErrorEnvelope.model_validate_json(error_message) + except ValidationError: + return error_message + message: Final = envelope.error.message if envelope.error else None + if message is None: + return error_message + try: + nested: Final = _ErrorBody.model_validate_json(message) + except ValidationError: + return message + return nested.message or message + + +def _snippet(text: str, annotation: _Annotation) -> str: + """ + The text a citation supports, not the citation marker itself. + + A url_citation's start/end indices span the inline marker ("([host](url))"), + which follows the claim it backs, so the snippet is the marker's own line up + to where the marker starts. + """ + start: Final = annotation.start_index + marker_start: Final = start if start is not None and 0 <= start <= len(text) else len(text) + claim: Final = text[:marker_start].rsplit("\n", 1)[-1].strip() + if claim: + return claim[-_SNIPPET_FALLBACK_LENGTH:] + return text[:_SNIPPET_FALLBACK_LENGTH] + + +def _citation_results(envelope: _ResponsesEnvelope) -> tuple[SearchResult, ...]: + """One result per cited URL: first occurrence wins, order preserved as answered.""" + cited: Final = tuple( + SearchResult( + title=annotation.title or "", + url=annotation.url or "", + snippet=_snippet(part.text, annotation), + date=None, + last_updated=None, + ) + for item in envelope.output + if item.type == "message" + for part in item.content + if part.type == "output_text" + for annotation in part.annotations + if annotation.type == "url_citation" and annotation.url + ) + first_by_url: Final = MappingProxyType({result.url: result for result in reversed(cited)}) + return tuple(first_by_url[url] for url in dict.fromkeys(result.url for result in cited)) + + +def _valid_max_results(max_results: object) -> int | None: + """A positive-int `max_results`, else None. Rejects bools, an `int` subclass, and + non-positive values so neither the request-side `count` nor the response-side cap + forwards a value the other would silently ignore. + """ + if isinstance(max_results, bool) or not isinstance(max_results, int): + return None + return max_results if max_results > 0 else None + + +def _requested_max_results(response_kwargs: Mapping[str, object]) -> int | None: + """The unified `max_results` cap the caller asked for, if any. + + The built-in web_search tool has no server-side result-count knob, so the cap is + enforced here after the fact; connection mode also honors it as a hard ceiling on + top of the tool's `count` hint. + """ + optional_params: Final = response_kwargs.get("optional_params") + if not isinstance(optional_params, Mapping): + return None + return _valid_max_results(optional_params.get("max_results")) + + +def _capped(results: tuple[SearchResult, ...], max_results: int | None) -> tuple[SearchResult, ...]: + return results[:max_results] if max_results is not None else results + + +class _SearchConfiguration(BaseModel): + model_config = ConfigDict(frozen=True) + + project_connection_id: str + count: int | None = None + + +class _BingGroundingParams(BaseModel): + model_config = ConfigDict(frozen=True) + + search_configurations: tuple[_SearchConfiguration, ...] + + +class _BingGroundingTool(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["bing_grounding"] = "bing_grounding" + bing_grounding: _BingGroundingParams + + +class _UserLocation(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["approximate"] = "approximate" + country: str + + +class _WebSearchTool(BaseModel): + model_config = ConfigDict(frozen=True) + + type: Literal["web_search"] = "web_search" + user_location: _UserLocation | None = None + + +class _ResponsesRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + model: str + input: str + tools: tuple[_BingGroundingTool | _WebSearchTool, ...] + + +def _search_tool(optional_params: Mapping[str, object]) -> _BingGroundingTool | _WebSearchTool: + connection_id: Final = get_secret_str(CONNECTION_ID_ENV) + max_results: Final = optional_params.get("max_results") + country: Final = optional_params.get("country") + if connection_id: + configuration: Final = _SearchConfiguration( + project_connection_id=connection_id, + count=_valid_max_results(max_results), + ) + return _BingGroundingTool(bing_grounding=_BingGroundingParams(search_configurations=(configuration,))) + location: Final = _UserLocation(country=country.upper()) if isinstance(country, str) else None + return _WebSearchTool(user_location=location) + + +def _default_entra_token_minter() -> str: + from litellm.secret_managers.get_azure_ad_token_provider import get_azure_ad_token_provider + + return get_azure_ad_token_provider(azure_scope=ENTRA_SCOPE)() + + +class BingGroundingSearchConfig(BaseSearchConfig): + def __init__(self, entra_token_minter: Callable[[], str] | None = None) -> None: + super().__init__() + self._entra_token_minter = entra_token_minter + + @staticmethod + def ui_friendly_name() -> str: + return "Grounding with Bing Search" + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseSearchConfig.validate_environment signature + api_key: str | None = None, + api_base: str | None = None, + **kwargs: object, # kwargs-ok: BaseSearchConfig.validate_environment signature + ) -> dict[str, str]: # mutable-ok: the http handler passes this straight to httpx as headers + """ + Validate environment and return headers. + + Returns a new dict rather than mutating ``headers``: the http handler calls this + a second time after ``litellm/search/main.py`` already did, so it has to be idempotent. + """ + return { # mutable-ok: httpx requires a plain dict of headers + **headers, + **self._auth_header(api_key, api_base), + "Content-Type": "application/json", + } + + def _auth_header(self, api_key: str | None, api_base: str | None) -> Mapping[str, str]: + """ + A caller-supplied ``api_key`` is an Azure API key and rides the ``api-key`` header; + an Entra bearer token (``BING_GROUNDING_TOKEN`` or one minted via azure-identity) + rides ``Authorization: Bearer``. Foundry rejects the wrong scheme for each. + """ + if api_key: + return MappingProxyType({"api-key": api_key}) + token: Final = self.resolve_server_api_key( + caller_api_key=None, + caller_api_base=api_base, + key_env_vars=(TOKEN_ENV,), + base_env_var=PROJECT_ENDPOINT_ENV, + default_api_base=None, + ) or self._mint_entra_token(api_base) + return MappingProxyType({"Authorization": f"Bearer {token}"}) + + def _mint_entra_token(self, caller_api_base: str | None) -> str: + self._assert_trusted_api_base_for_server_credential( + caller_api_base, None, PROJECT_ENDPOINT_ENV, "Azure AD token" + ) + minter: Final = self._entra_token_minter or _default_entra_token_minter + try: + return minter() + except Exception as e: + raise ValueError( + f"Grounding with Bing Search: no credential available. Pass api_key, set {TOKEN_ENV} " + f"to an Entra bearer token, or configure azure-identity (AZURE_CLIENT_ID / " + f"AZURE_CLIENT_SECRET / AZURE_TENANT_ID or any DefaultAzureCredential source) " + f"for scope {ENTRA_SCOPE}. Underlying error: {e}" + ) from e + + def get_complete_url( + self, + api_base: str | None, + optional_params: dict[str, object], # mutable-ok: BaseSearchConfig.get_complete_url signature + data: dict[str, object] | list[dict[str, object]] | None = None, # mutable-ok: base signature + **kwargs: object, # kwargs-ok: BaseSearchConfig.get_complete_url signature + ) -> str: + resolved_base: Final = api_base or get_secret_str(PROJECT_ENDPOINT_ENV) + if not resolved_base: + raise ValueError( + f"{PROJECT_ENDPOINT_ENV} is not set. Set it to your Microsoft Foundry project " + f"endpoint, e.g. https://.services.ai.azure.com/api/projects/." + ) + trimmed: Final = resolved_base.rstrip("/") + if trimmed.endswith(_RESPONSES_PATH): + return trimmed + return f"{trimmed}{_RESPONSES_PATH}" + + def transform_search_request( + self, + query: str | list[str], # mutable-ok: BaseSearchConfig.transform_search_request signature + optional_params: dict[str, object], # mutable-ok: base signature + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_request signature + ) -> dict[str, object]: # mutable-ok: the http handler passes this straight to httpx as the JSON body + """ + Transform Search request to the Foundry Responses API format. + + The unified params map as far as the API allows: + - max_results -> the bing_grounding search configuration's `count`; the built-in + web_search tool has no result-count knob, so that mode instead caps the returned + results after the fact (see transform_search_response) + - country -> web_search's approximate `user_location` (bing_grounding's `market` + wants a full locale like en-US, which a bare country code cannot fill) + - search_domain_filter, max_tokens_per_page -> no API equivalent, dropped + """ + model: Final = get_secret_str(MODEL_ENV) + if not model: + raise ValueError( + f"{MODEL_ENV} is not set. Set it to a model deployment in the Foundry project " + f"that runs the grounded search, e.g. gpt-4.1." + ) + request: Final = _ResponsesRequest( + model=model, + input=" ".join(query) if isinstance(query, list) else query, + tools=(_search_tool(optional_params),), + ) + return request.model_dump(mode="json", exclude_none=True) + + def transform_search_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + **kwargs: object, # kwargs-ok: BaseSearchConfig.transform_search_response signature + ) -> SearchResponse: + try: + parsed: Final = _ResponsesEnvelope.model_validate_json(raw_response.content) + except ValidationError as e: + raise self.get_error_class( + error_message=f"response does not match the Foundry Responses API schema: {e}", + status_code=raw_response.status_code, + headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature + ) + if parsed.status == "failed": + detail: Final = ( + parsed.error.message if parsed.error and parsed.error.message else "the grounded search failed" + ) + raise self._upstream_error(detail, raw_response) + results: Final = _capped(_citation_results(parsed), _requested_max_results(kwargs)) + if not results and parsed.status == "incomplete": + reason: Final = ( + parsed.incomplete_details.reason + if parsed.incomplete_details and parsed.incomplete_details.reason + else "unknown reason" + ) + raise self._upstream_error(f"the grounded search was incomplete: {reason}", raw_response) + return self._priced(results) + + def _upstream_error(self, detail: str, raw_response: httpx.Response) -> Exception: + return self.get_error_class( + error_message=detail, + status_code=_UPSTREAM_ERROR_STATUS, + headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature + ) + + def _priced(self, results: tuple[SearchResult, ...]) -> SearchResponse: + """web_search mode runs no paid Grounding with Bing transaction, so it must not + inherit the connection-mode ``bing_grounding/search`` price; zero its per-query + cost while leaving connection mode to the cost map.""" + response: Final = SearchResponse( + results=list(results), # mutable-ok: SearchResponse.results is list[SearchResult] + object="search", + ) + if get_secret_str(CONNECTION_ID_ENV): + return response + response._hidden_params[ + "additional_headers" + ] = { # mutable-ok: response_cost_calculator writes into _hidden_params + _RESPONSE_COST_HEADER: 0.0 + } + return response + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str], # mutable-ok: BaseSearchConfig.get_error_class signature + ) -> Exception: + detail: Final = _unwrap_error_detail(error_message).rstrip(". ") + return BaseLLMException( + status_code=status_code, + message=f"Grounding with Bing Search: {detail}. See {_DOCS_URL} for details.", + headers=headers, + ) diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index 1a924088390..61cbc213b11 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -65,15 +65,24 @@ class AzureModelRouterConfig(AzureAIStudioConfig): Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07) and returns it with the azure_ai/ prefix for proper display and cost tracking. + + Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs, + response restamping) can read it instead of guessing the route from the model string. """ - from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + AzureFoundryModelInfo, + ) + from litellm.router_utils.add_retry_fallback_headers import ( + get_hidden_params_dict, + ) # Get base model for the parent call (strips routing prefixes for API compatibility) base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model) # Call parent transform_response first - this will extract the actual model # from the raw response (e.g., "gpt-5-nano-2025-08-07") - model_response = super().transform_response( + transformed_response: Final = super().transform_response( model=base_model, raw_response=raw_response, model_response=model_response, @@ -86,7 +95,15 @@ class AzureModelRouterConfig(AzureAIStudioConfig): api_key=api_key, json_mode=json_mode, ) - return model_response + selected_model: Final = transformed_response.model + if selected_model: + # Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a + # class-level dict, so an in-place write can bleed into unrelated responses. + transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter # mutable-ok: ModelResponse requires _hidden_params to be a plain dict + **get_hidden_params_dict(transformed_response), + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model, + } + return transformed_response def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None: """ diff --git a/litellm/llms/azure_ai/common_utils.py b/litellm/llms/azure_ai/common_utils.py index d09055d7671..26a90157455 100644 --- a/litellm/llms/azure_ai/common_utils.py +++ b/litellm/llms/azure_ai/common_utils.py @@ -51,6 +51,9 @@ def get_azure_ai_auth_headers( ) +AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model" + + class AzureFoundryModelInfo(BaseLLMModelInfo): """Model info for Azure AI / Azure Foundry models.""" @@ -82,6 +85,41 @@ class AzureFoundryModelInfo(BaseLLMModelInfo): return "model_router" return "default" + @staticmethod + def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None: + """The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``. + + Reading this beats re-deriving the route from a model string: the stamp is set on the + code path that was actually taken, so it holds no matter what the caller named the model. + """ + if not hidden_params: + return None + selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY) + if isinstance(selected, str) and selected: + return selected + return None + + @staticmethod + def is_model_router_call( + model: str | None = None, + hidden_params: Mapping[str, object] | None = None, + ) -> bool: + """Whether a request went down the Azure Model Router route. + + Prefers the response stamp, then the deployment's litellm model path, and only then the + caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router + name heuristic lives in exactly one place. + """ + if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None: + return True + deployment_model: Final = ( + hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None + ) + return any( + isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router" + for candidate in (deployment_model, model) + ) + @staticmethod def get_api_base(api_base: str | None = None) -> str | None: return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE") diff --git a/litellm/llms/base_llm/videos/transformation.py b/litellm/llms/base_llm/videos/transformation.py index dcecdc646ff..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 @@ -340,14 +341,18 @@ class BaseVideoConfig(ABC): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, + video_file: FileContent | None = None, extra_body: dict[str, Any] | None = None, prefetched_source_data: dict[str, Any] | None = None, - ) -> tuple[str, dict]: + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: """ - Transform the video edit request into a URL and JSON data. + Transform the video edit request into a URL plus either JSON data or + multipart form fields and files. Returns: - Tuple[str, Dict]: (url, data) for the POST request + tuple[str, Mapping[str, object], RequestFiles | None]: (url, data, + files). When files is None the handler sends data as JSON; otherwise + data holds the form fields and files holds the uploaded source video. """ raise NotImplementedError("video edit is not supported for this provider") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b437e25d24b..767677cbcbf 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1617,6 +1617,8 @@ class AmazonConverseConfig(BaseConfig): } if additional_request_params: data["additionalModelRequestFields"] = additional_request_params + if "thinking" in additional_request_params: + data["additionalModelResponseFieldPaths"] = ("/usage/output_tokens_details",) if system_content_blocks: data["system"] = system_content_blocks @@ -1801,6 +1803,17 @@ class AmazonConverseConfig(BaseConfig): thinking_blocks_list.append(_redacted_block) return thinking_blocks_list + @staticmethod + def thinking_tokens_from_additional_fields(additional_fields: object) -> int | None: + """Converse omits thinking tokens from its usage block; they only arrive under + ``additionalModelResponseFields`` when ``/usage/output_tokens_details`` is requested.""" + if not isinstance(additional_fields, Mapping): + return None + usage: Final = additional_fields.get("usage") + if not isinstance(usage, Mapping): + return None + return AnthropicConfig.thinking_tokens_from_usage(usage) + @staticmethod def is_converse_usage_shape(usage_object: Mapping[str, object]) -> bool: """Converse-family models report camelCase token counts, not Anthropic's snake_case.""" @@ -1842,6 +1855,7 @@ class AmazonConverseConfig(BaseConfig): usage: ConverseTokenUsageBlock, reasoning_content: str | None = None, thinking_ran: bool = False, + provider_reasoning_tokens: int | None = None, ) -> Usage: input_tokens = usage["inputTokens"] output_tokens: Final = usage["outputTokens"] @@ -1862,9 +1876,14 @@ class AmazonConverseConfig(BaseConfig): cache_creation_tokens=cache_creation_input_tokens, text_tokens=raw_input_tokens, ) - reasoning_tokens: Final = ( + estimated_reasoning_tokens: Final = ( token_counter(text=reasoning_content, count_response_tokens=True) if reasoning_content else 0 ) + reasoning_tokens: Final = ( + min(max(0, provider_reasoning_tokens), output_tokens) + if provider_reasoning_tokens is not None + else estimated_reasoning_tokens + ) completion_tokens_details: Final = ( CompletionTokensDetailsWrapper( reasoning_tokens=reasoning_tokens, @@ -2272,6 +2291,9 @@ class AmazonConverseConfig(BaseConfig): completion_response["usage"], reasoning_content=chat_completion_message.get("reasoning_content"), thinking_ran=reasoningContentBlocks is not None, + provider_reasoning_tokens=self.thinking_tokens_from_additional_fields( + completion_response.get("additionalModelResponseFields") + ), ) ## HANDLE TOOL CALLS diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 35e82f1961a..fc34e403beb 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -331,6 +331,7 @@ class AWSEventStreamDecoder: self.json_mode = json_mode self._current_tool_name: str | None = None self._thinking_ran = False + self._provider_reasoning_tokens: int | None = None def check_empty_tool_call_args(self) -> bool: """ @@ -559,10 +560,14 @@ class AWSEventStreamDecoder: tool_use = self._handle_converse_stop_event(content_block_index) elif "stopReason" in chunk_data: finish_reason = map_finish_reason(chunk_data.get("stopReason", "stop")) + self._provider_reasoning_tokens = AmazonConverseConfig.thinking_tokens_from_additional_fields( + chunk_data.get("additionalModelResponseFields") + ) elif "usage" in chunk_data: usage = converse_config.transform_usage( chunk_data.get("usage", {}), thinking_ran=self._thinking_ran, + provider_reasoning_tokens=self._provider_reasoning_tokens, ) if thinking_blocks: self._thinking_ran = True diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 0ce2e6f60d3..d0a3c37ffb3 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, cast from httpx import Response @@ -93,6 +94,9 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD endpoint_url, ) + def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None: + return None + def sign_request( self, headers: dict, @@ -109,6 +113,7 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD request_data=request_data or {}, api_base=api_base, model=model, + api_key=self.get_bedrock_bearer_token(optional_params), ) def logging_non_streaming_response( diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 889361cd808..d877fbb4e09 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -13,6 +13,7 @@ global state. """ import re +from collections.abc import Mapping from typing import Final from botocore.exceptions import ( @@ -31,30 +32,39 @@ BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" MANTLE_HOST_RE: Final = re.compile(r"^https?://bedrock-mantle\.([^/.]+)\.api\.aws", re.IGNORECASE) +def resolve_mantle_bearer_token(api_key: str | None) -> str | None: + return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + + +def resolve_mantle_region(params: Mapping[str, object]) -> str: + region: Final = params.get("aws_region_name") + if isinstance(region, str) and region: + BaseAWSLLM._validate_aws_region_name(region) + return region + api_base: Final = params.get("api_base") + base: Final = (api_base if isinstance(api_base, str) else None) or get_secret_str("BEDROCK_MANTLE_API_BASE") + if base: + match: Final = MANTLE_HOST_RE.match(base.rstrip("/")) + if match: + return match.group(1) + return ( + get_secret_str("BEDROCK_MANTLE_REGION") + or get_secret_str("AWS_REGION_NAME") + or get_secret_str("AWS_REGION") + or BEDROCK_MANTLE_DEFAULT_REGION + ) + + class BedrockMantleAuthMixin: _aws_signer: BaseAWSLLM @staticmethod def _resolve_bearer_token(api_key: str | None) -> str | None: - return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") + return resolve_mantle_bearer_token(api_key) @staticmethod def _resolve_region(params: dict) -> str: - region: Final = params.get("aws_region_name") - if region: - BaseAWSLLM._validate_aws_region_name(region) - return region - base: Final = params.get("api_base") or get_secret_str("BEDROCK_MANTLE_API_BASE") - if base: - match: Final = MANTLE_HOST_RE.match(base.rstrip("/")) - if match: - return match.group(1) - return ( - get_secret_str("BEDROCK_MANTLE_REGION") - or get_secret_str("AWS_REGION_NAME") - or get_secret_str("AWS_REGION") - or BEDROCK_MANTLE_DEFAULT_REGION - ) + return resolve_mantle_region(params) def sign_request( self, diff --git a/litellm/llms/bedrock_mantle/passthrough/transformation.py b/litellm/llms/bedrock_mantle/passthrough/transformation.py new file mode 100644 index 00000000000..e6b831efa57 --- /dev/null +++ b/litellm/llms/bedrock_mantle/passthrough/transformation.py @@ -0,0 +1,71 @@ +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final, Literal, Optional + +from httpx import Response + +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.llms.bedrock_mantle.common_utils import ( + MANTLE_HOST_RE, + resolve_mantle_bearer_token, + resolve_mantle_region, +) +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.types.utils import CostResponseTypes + + +class BedrockMantlePassthroughConfig(BedrockPassthroughConfig): + """Native Bedrock runtime passthrough (InvokeModel, Converse) for deployments declared as bedrock_mantle. + + The Mantle host only serves the OpenAI-compatible surface, so a Mantle api_base lends its region and the + request itself goes to bedrock-runtime, signed with the deployment's Bearer token or SigV4 credentials. + """ + + def _get_aws_region_name( + self, + optional_params: Mapping[str, object], + model: str | None = None, + model_id: str | None = None, + ) -> str: + return resolve_mantle_region(optional_params) + + def get_runtime_endpoint( + self, + api_base: str | None, + aws_bedrock_runtime_endpoint: str | None, + aws_region_name: str, + endpoint_type: Literal["runtime", "agent", "agentcore"] | None = "runtime", + ) -> tuple[str, str]: + is_mantle_host: Final = api_base is not None and MANTLE_HOST_RE.match(api_base.rstrip("/")) is not None + return super().get_runtime_endpoint( + api_base=None if is_mantle_host else api_base, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, + aws_region_name=aws_region_name, + endpoint_type=endpoint_type, + ) + + def get_bedrock_bearer_token(self, litellm_params: Mapping[str, object]) -> str | None: + api_key: Final = litellm_params.get("api_key") + return resolve_mantle_bearer_token(api_key if isinstance(api_key, str) else None) + + def logging_non_streaming_response( + self, + model: str, + custom_llm_provider: str, + httpx_response: Response, + request_data: dict, # mutable-ok: mirrors the inherited BedrockPassthroughConfig signature + logging_obj: Logging, + endpoint: str, + ) -> Optional["CostResponseTypes"]: + is_converse: Final = "invoke" not in endpoint and "converse" in endpoint + shape_provider: Final = LlmProviders.BEDROCK.value if is_converse else custom_llm_provider + return super().logging_non_streaming_response( + model=model, + custom_llm_provider=shape_provider, + httpx_response=httpx_response, + request_data=request_data, + logging_obj=logging_obj, + endpoint=endpoint, + ) diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 92da5835b2d..3e5dd4ff87d 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -15,8 +15,12 @@ role / access key / profile / web identity), signed via the shared BaseAWSLLM._sign_request after the request body is finalized. """ +import json +from collections.abc import Mapping from typing import Any, Final +from typing_extensions import ReadOnly, TypedDict + import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -50,6 +54,33 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"}) _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools" +_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message" +_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction" +_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call" + + +class _RewrittenOutputTextBlock(TypedDict): + type: ReadOnly[str] + text: ReadOnly[str] + + +class _RewrittenAssistantMessageItem(TypedDict): + type: ReadOnly[str] + role: ReadOnly[str] + content: ReadOnly[tuple[_RewrittenOutputTextBlock, ...]] + + +class _RewrittenCompactionItem(TypedDict): + type: ReadOnly[str] + encrypted_content: ReadOnly[str] + + +class _RewrittenFunctionCallItem(TypedDict): + type: ReadOnly[str] + call_id: ReadOnly[str] + name: ReadOnly[str] + arguments: ReadOnly[str] + class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig): def __init__( @@ -155,6 +186,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI headers: dict, ) -> dict: remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input) + normalized_input: Final = self._normalize_codex_input_items(remaining_input) request_params: Final = ( { **response_api_optional_request_params, @@ -168,7 +200,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) return super().transform_responses_api_request( model=model, - input=remaining_input, + input=normalized_input, response_api_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, @@ -210,6 +242,91 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI ) return remaining_input, cls._filter_unsupported_tools(hoisted_tools) + @staticmethod + def _agent_message_text(item: "Mapping[str, Any]") -> str: + content: Final = item.get("content") + if not isinstance(content, list): + return "" + return "".join( + str(block.get("text") or block.get("encrypted_content") or "") + for block in content + if isinstance(block, dict) + ) + + @classmethod + def _normalize_agent_message_item(cls, item: "Mapping[str, Any]") -> "_RewrittenAssistantMessageItem | None": + text: Final = cls._agent_message_text(item) + if not text: + return None + rewritten: Final[_RewrittenAssistantMessageItem] = { + "type": "message", + "role": "assistant", + "content": ({"type": "output_text", "text": text},), + } + return rewritten + + @staticmethod + def _normalize_context_compaction_item(item: "Mapping[str, Any]") -> "_RewrittenCompactionItem | None": + encrypted_content: Final = item.get("encrypted_content") + if not isinstance(encrypted_content, str) or not encrypted_content: + return None + rewritten: Final[_RewrittenCompactionItem] = {"type": "compaction", "encrypted_content": encrypted_content} + return rewritten + + @staticmethod + def _normalize_local_shell_call_item(item: "Mapping[str, Any]") -> "_RewrittenFunctionCallItem | None": + call_id: Final = item.get("call_id") + if not isinstance(call_id, str) or not call_id: + return None + action: Final = item.get("action") + rewritten: Final[_RewrittenFunctionCallItem] = { + "type": "function_call", + "call_id": call_id, + "name": "local_shell", + "arguments": json.dumps(action) if isinstance(action, dict) else "{}", + } + return rewritten + + @classmethod + def _normalize_codex_input_item(cls, item: object) -> "tuple[object, str | None]": + """Returns (normalized item or None to drop it, original type when rewritten).""" + if not isinstance(item, dict): + return item, None + item_type: Final = item.get("type") + if item_type == _CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: + return cls._normalize_agent_message_item(item), item_type + if item_type == _CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: + return cls._normalize_context_compaction_item(item), item_type + if item_type == _CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: + return cls._normalize_local_shell_call_item(item), item_type + return item, None + + @classmethod + def _normalize_codex_input_items( + cls, + input: "str | ResponseInputParam", + ) -> "str | ResponseInputParam": + """Rewrite Codex history item types Mantle rejects with 400 "Invalid + 'input': value did not match any expected variant" into supported + equivalents. `agent_message` (Codex multi-agent traffic; its + encrypted_content slot carries the plaintext payload when the model + never issued encrypted args) becomes an assistant message, + `context_compaction` becomes the `compaction` spelling Mantle accepts, + and `local_shell_call` becomes the function_call its recorded + function_call_output already pairs with. + """ + if not isinstance(input, list): + return input + normalized: Final = tuple(cls._normalize_codex_input_item(item) for item in input) + rewritten_types: Final = sorted(frozenset(item_type for _, item_type in normalized if item_type is not None)) + if rewritten_types: + verbose_logger.warning( + "Bedrock Mantle Responses API: rewrote Codex input item type(s) %s that Mantle rejects.", + rewritten_types, + ) + kept: Final = [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list + return kept # pyright: ignore[reportReturnType] # Codex passthrough items sit outside the OpenAI input union + def map_openai_params( self, response_api_optional_params: ResponsesAPIOptionalRequestParams, diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 52f30e31641..777ab576de2 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -9,7 +9,7 @@ import threading import time from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy -from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional, TypeAlias, TypedDict import certifi import httpx @@ -933,11 +933,83 @@ class AsyncHTTPHandler: response.raise_for_status() return response + # Strong references to finalizer-scheduled client-close tasks. A bare + # create_task() result may be garbage-collected before it runs, leaving + # the underlying aiohttp session unclosed ("Unclosed client session"). + # Mirrors LiteLLMAiohttpTransport._background_close_tasks. + _finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes + + @classmethod + def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None: + cls._finalizer_close_tasks.discard(task) + if task.cancelled(): + return + exc: Final = task.exception() + if exc is not None: + verbose_logger.debug("Error closing client at finalization: %s", exc) + + def _aiohttp_session_bound_elsewhere(self, loop: asyncio.AbstractEventLoop) -> bool: + """True when the wrapped aiohttp session is bound to a loop other than + ``loop`` — awaiting ``aclose()`` here would touch that loop's internals.""" + from litellm.llms.custom_httpx.aiohttp_transport import ( + LiteLLMAiohttpTransport, + ) + + transport: Final = getattr(self._client, "_transport", None) + if not isinstance(transport, LiteLLMAiohttpTransport): + return False + session: Final = transport.client + if not isinstance(session, ClientSession) or session.closed: + return False + return getattr(session, "_loop", None) is not loop + + def _dispose_wrapped_aiohttp_session(self) -> None: + """Dispose the wrapped aiohttp session when ``aclose()`` cannot run here. + + Finalization either has no running loop, or a loop the session is not + bound to. Delegating to the transport's lifecycle-aware disposal picks + the safe path per session state (async close on its own loop, threadsafe + handoff to a loop running elsewhere, or the synchronous connector + teardown that flips the flags ``ClientSession.__del__`` checks), so no + "Unclosed client session" / "Unclosed connector" warnings fire at + garbage collection. + """ + from litellm.llms.custom_httpx.aiohttp_transport import ( + LiteLLMAiohttpTransport, + ) + + transport: Final = getattr(self._client, "_transport", None) + if not isinstance(transport, LiteLLMAiohttpTransport): + return + # A shared session (e.g. the proxy's) is never this handler's to close. + if not getattr(transport, "_owns_session", False): + return + session: Final = transport.client + if isinstance(session, ClientSession) and not session.closed: + transport._close_recycled_session(session) # pyright: ignore[reportPrivateUsage] # deliberate reuse of the transport's lifecycle-aware disposal; an async close can never run in this context + def __del__(self) -> None: try: if not _handler_may_close_client(sys.getrefcount(self._client), self._owns_client): return - asyncio.get_running_loop().create_task(self._client.aclose()) + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + # No running loop at finalization time (worker threads after + # their loop closed, interpreter/worker shutdown, GC in a + # sync context). An async close can never run here. + self._dispose_wrapped_aiohttp_session() + return + if self._aiohttp_session_bound_elsewhere(loop): + # GC ran on a live loop (e.g. the app's) but the session + # belongs to another, possibly dead, loop — awaiting aclose() + # here is the cross-loop path the transport refuses. + self._dispose_wrapped_aiohttp_session() + return + task: Final = loop.create_task(self._client.aclose()) + cls: Final = type(self) + cls._finalizer_close_tasks.add(task) + task.add_done_callback(cls._on_finalizer_close_done) except Exception: pass diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 76f109f0eb6..b7ddb55ae89 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, Type from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx +from httpx._types import FileContent from openai.types.file_deleted import FileDeleted import litellm @@ -1846,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( @@ -1944,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( @@ -7838,6 +7841,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, + video_file: FileContent | None = None, extra_headers: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, timeout: float | None = None, @@ -7849,6 +7853,7 @@ class BaseLLMHTTPHandler: return self.async_video_edit_handler( prompt=prompt, video_id=video_id, + video_file=video_file, video_provider_config=video_provider_config, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, @@ -7902,9 +7907,10 @@ class BaseLLMHTTPHandler: prefetched_source_data = prefetch_resp.json() try: - url, data = video_provider_config.transform_video_edit_request( + url, data, files = video_provider_config.transform_video_edit_request( prompt=prompt, video_id=video_id, + video_file=video_file, api_base=api_base, litellm_params=litellm_params, headers=headers, @@ -7923,11 +7929,10 @@ class BaseLLMHTTPHandler: }, ) - response: Final = sync_httpx_client.post( - url=url, - headers=headers, - json=data, - timeout=timeout, + response: Final = ( + sync_httpx_client.post(url=url, headers=headers, data=data, files=files, timeout=timeout) + if files + else sync_httpx_client.post(url=url, headers=headers, json=data, timeout=timeout) ) response.raise_for_status() return video_provider_config.transform_video_edit_response( @@ -7947,6 +7952,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, litellm_params, logging_obj, + video_file: FileContent | None = None, extra_headers: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, timeout: float | None = None, @@ -7998,9 +8004,10 @@ class BaseLLMHTTPHandler: prefetched_source_data = prefetch_resp.json() try: - url, data = video_provider_config.transform_video_edit_request( + url, data, files = video_provider_config.transform_video_edit_request( prompt=prompt, video_id=video_id, + video_file=video_file, api_base=api_base, litellm_params=litellm_params, headers=headers, @@ -8019,11 +8026,10 @@ class BaseLLMHTTPHandler: }, ) - response: Final = await async_httpx_client.post( - url=url, - headers=headers, - json=data, - timeout=timeout, + response: Final = await ( + async_httpx_client.post(url=url, headers=headers, data=data, files=files, timeout=timeout) + if files + else async_httpx_client.post(url=url, headers=headers, json=data, timeout=timeout) ) response.raise_for_status() return video_provider_config.transform_video_edit_response( diff --git a/litellm/llms/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/openai/chat/gpt_5_transformation.py b/litellm/llms/openai/chat/gpt_5_transformation.py index ffa3de0d5c6..3a65e4a9426 100644 --- a/litellm/llms/openai/chat/gpt_5_transformation.py +++ b/litellm/llms/openai/chat/gpt_5_transformation.py @@ -16,7 +16,7 @@ def _normalize_reasoning_effort_for_chat_completion( ) -> str | None: """Convert reasoning_effort to the string format expected by OpenAI chat completion API. - The chat completion API expects a simple string: 'none', 'low', 'medium', 'high', or 'xhigh'. + The chat completion API expects an effort string such as 'low' or 'high'. Config/deployments may pass the Responses API format: {'effort': 'high', 'summary': 'detailed'}. """ if value is None: diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index 4fd0429c182..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 @@ -502,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, @@ -570,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 2264f52e2a3..8dd77db08c0 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -672,6 +672,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): api_base, litellm_params, headers, + video_file=None, extra_body=None, prefetched_source_data=None, ): diff --git a/litellm/llms/together_ai/rerank/handler.py b/litellm/llms/together_ai/rerank/handler.py index 10246451a9d..b8079e52c97 100644 --- a/litellm/llms/together_ai/rerank/handler.py +++ b/litellm/llms/together_ai/rerank/handler.py @@ -16,11 +16,16 @@ from litellm.llms.together_ai.rerank.transformation import TogetherAIRerankConfi from litellm.types.rerank import RerankRequest, RerankResponse +def _rerank_url(api_base: str) -> str: + return f"{api_base.rstrip('/')}/rerank" + + class TogetherAIRerank(BaseLLM): def rerank( self, model: str, api_key: str, + api_base: str, query: str, documents: list[str | dict[str, Any]], top_n: int | None = None, @@ -46,10 +51,10 @@ class TogetherAIRerank(BaseLLM): raise ValueError("TogetherAI does not support max_chunks_per_doc") if _is_async: - return self.async_rerank(request_data_dict, api_key) # Call async method + return self.async_rerank(request_data_dict, api_key, api_base) response: Final = client.post( - "https://api.together.xyz/v1/rerank", + _rerank_url(api_base), headers={ "accept": "application/json", "content-type": "application/json", @@ -69,11 +74,12 @@ class TogetherAIRerank(BaseLLM): self, request_data_dict: dict[str, Any], api_key: str, + api_base: str, ) -> RerankResponse: client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.TOGETHER_AI) # Use async client response: Final = await client.post( - "https://api.together.xyz/v1/rerank", + _rerank_url(api_base), headers={ "accept": "application/json", "content-type": "application/json", diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py new file mode 100644 index 00000000000..0764a8bea62 --- /dev/null +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -0,0 +1,149 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Final + +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig +from litellm.llms.vertex_ai.common_utils import validate_vertex_location +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +VERTEX_INTERACTIONS_API_VERSION: Final = "v1beta1" +VERTEX_INTERACTIONS_DEFAULT_LOCATION: Final = "global" + + +@dataclass(frozen=True, slots=True) +class VertexInteractionsTarget: + base_url: str + project_id: str + location: str + + @property + def collection_url(self) -> str: + return ( + f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}" + f"/projects/{self.project_id}/locations/{self.location}/interactions" + ) + + def interaction_url(self, interaction_id: str) -> str: + encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") + return f"{self.collection_url}/{encoded_interaction_id}" + + +class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig): + def __init__( + self, + mint_access_token: Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]] | None = None, + ) -> None: + super().__init__() + self._mint_access_token: Final[Callable[[VERTEX_CREDENTIALS_TYPES | None, str | None], tuple[str, str]]] = ( + mint_access_token or self._mint_access_token_with_vertex_base + ) + + def _mint_access_token_with_vertex_base( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + return self._ensure_access_token( + credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai" + ) + + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.VERTEX_AI + + @property + def api_version(self) -> str: + return VERTEX_INTERACTIONS_API_VERSION + + def get_default_vertex_location(self) -> str: + return VERTEX_INTERACTIONS_DEFAULT_LOCATION + + def _mint(self, litellm_params: GenericLiteLLMParams) -> tuple[str, str]: + raw_params: Final = litellm_params.model_dump() + return self._mint_access_token( + self.safe_get_vertex_ai_credentials(raw_params), + self.safe_get_vertex_ai_project(raw_params), + ) + + def _target(self, api_base: str | None, litellm_params: GenericLiteLLMParams) -> VertexInteractionsTarget: + _, project_id = self._mint(litellm_params) + if not project_id: + raise ValueError( + "Vertex AI project is required. Set vertex_project, litellm.vertex_project, or VERTEXAI_PROJECT" + ) + location: Final = validate_vertex_location( + self.explicit_vertex_ai_location(litellm_params.model_dump()) or VERTEX_INTERACTIONS_DEFAULT_LOCATION + ) + return VertexInteractionsTarget( + base_url=self.get_api_base(api_base or None, location), + project_id=project_id, + location=location, + ) + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict: # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers + access_token, _ = self._mint(litellm_params or GenericLiteLLMParams()) + return { # mutable-ok: BaseInteractionsAPIConfig declares plain-dict headers + "Content-Type": "application/json", + "Authorization": f"Bearer {access_token}", + **headers, + } + + def get_complete_url( + self, + api_base: str | None, + model: str | None, + agent: str | None = None, + litellm_params: Mapping[str, object] | None = None, + stream: bool | None = None, + ) -> str: + params: Final = ( + GenericLiteLLMParams.model_validate(litellm_params) if litellm_params else GenericLiteLLMParams() + ) + collection_url: Final = self._target(api_base, params).collection_url + return f"{collection_url}?alt=sse" if stream else collection_url + + def _interaction_by_id_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + url_suffix: str = "", + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + target: Final = self._target(api_base or None, litellm_params) + return f"{target.interaction_url(interaction_id)}{url_suffix}", {} # mutable-ok: same base contract + + def transform_get_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params) + + def transform_delete_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params) + + def transform_cancel_interaction_request( + self, + interaction_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body + return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel") diff --git a/litellm/llms/vertex_ai/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/main.py b/litellm/main.py index d3967473f99..fcf33653cb3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -416,7 +416,7 @@ async def acompletion( logprobs: bool | None = None, top_logprobs: int | None = None, deployment_id=None, - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, safety_identifier: str | None = None, service_tier: str | None = None, @@ -602,7 +602,7 @@ async def acompletion( _, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, - api_base=base_url, + api_base=kwargs.get("api_base") or base_url, ) fallbacks = fallbacks or litellm.model_fallbacks @@ -4920,7 +4920,7 @@ def completion( logit_bias: dict | None = None, user: str | None = None, # openai v1.0+ new params - reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None, + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None, verbosity: Literal["low", "medium", "high"] | None = None, response_format: dict | type[BaseModel] | None = None, seed: int | None = None, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c36018feb9d..1aa4c7cd060 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -17155,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", @@ -49008,12 +49016,13 @@ "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49040,12 +49049,13 @@ "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49072,12 +49082,13 @@ "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 8a2ee2a3af8..4b30afb2f98 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -199,7 +199,7 @@ def llm_passthrough_route( api_key=api_key, ) - litellm_params_dict: Final = get_litellm_params(**kwargs) + litellm_params_dict: Final = get_litellm_params(api_key=api_key, api_base=api_base, **kwargs) if client is None: from litellm.llms.custom_httpx.http_handler import ( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0840d37ffa1..628e569e1b8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -815,6 +815,7 @@ class LiteLLMRoutes(enum.Enum): "/team/member_add", "/team/member_delete", "/team/member_update", + "/team/{team_id}/member/{user_id}/reset_spend", "/team/permissions_list", "/team/permissions_update", "/team/daily/activity", @@ -1287,6 +1288,16 @@ class RegenerateKeyRequest(GenerateKeyRequest): class ResetSpendRequest(LiteLLMPydanticObjectBase): reset_to: float + @field_validator("reset_to", mode="before") + @classmethod + def reject_bool_reset_to(cls, v): + # bool is a subclass of int, so pydantic silently coerces True/False into + # 1.0/0.0 for a `float` field: a caller who accidentally sends a boolean + # would otherwise get an unintended spend reset instead of a 422. + if isinstance(v, bool): + raise ValueError("reset_to must be a number, not a boolean") # noqa: TRY004 # pydantic needs ValueError + return v + class KeyRequest(LiteLLMPydanticObjectBase): keys: list[str] | None = None diff --git a/litellm/proxy/analytics_endpoints/cache_activity.py b/litellm/proxy/analytics_endpoints/cache_activity.py index 751d6c15dfa..b87b8eac3ef 100644 --- a/litellm/proxy/analytics_endpoints/cache_activity.py +++ b/litellm/proxy/analytics_endpoints/cache_activity.py @@ -34,10 +34,18 @@ class CacheActivityFilterOptions(BaseModel): models: list[str] +class CacheActivityErrorBucket(BaseModel): + call_type: str + error_code: str + error_class: str + count: int + + class CacheActivityResponse(BaseModel): groups: list[CacheActivityGroup] totals: CacheActivityTotals filter_options: CacheActivityFilterOptions + error_breakdown: tuple[CacheActivityErrorBucket, ...] GROUPS_SQL: Final = """ @@ -65,6 +73,26 @@ GROUPS_SQL: Final = """ ORDER BY (COUNT(*)) DESC """ +ERROR_BREAKDOWN_SQL: Final = """ + SELECT + CASE WHEN sl."call_type" = '' THEN 'Unknown' ELSE sl."call_type" END AS call_type, + COALESCE(NULLIF(sl."metadata"->'error_information'->>'error_code', ''), 'Unknown') AS error_code, + COALESCE(NULLIF(sl."metadata"->'error_information'->>'error_class', ''), 'Unknown') AS error_class, + COUNT(*)::int AS count + FROM "LiteLLM_SpendLogs" sl + LEFT JOIN "LiteLLM_VerificationToken" vt ON sl."api_key" = vt."token" + WHERE + sl."status" = 'failure' + AND sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') + AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC') + AND ($3::jsonb = '[]'::jsonb + OR COALESCE(vt."key_alias", 'Unnamed Key') IN (SELECT jsonb_array_elements_text($3::jsonb))) + AND ($4::jsonb = '[]'::jsonb + OR sl."model" IN (SELECT jsonb_array_elements_text($4::jsonb))) + GROUP BY 1, 2, 3 + ORDER BY (COUNT(*)) DESC +""" + KEY_ALIAS_OPTIONS_SQL: Final = """ SELECT DISTINCT COALESCE(vt."key_alias", 'Unnamed Key') AS key_alias FROM "LiteLLM_SpendLogs" sl @@ -95,6 +123,7 @@ class _ModelRow(BaseModel): _groups_adapter: Final = TypeAdapter(list[CacheActivityGroup]) +_error_buckets_adapter: Final = TypeAdapter(tuple[CacheActivityErrorBucket, ...]) _key_alias_rows_adapter: Final = TypeAdapter(list[_KeyAliasRow]) _model_rows_adapter: Final = TypeAdapter(list[_ModelRow]) @@ -120,10 +149,11 @@ async def get_cache_activity( key_aliases: Sequence[str], models: Sequence[str], ) -> CacheActivityResponse: - group_rows, key_alias_rows, model_rows = await asyncio.gather( - prisma_client.db.query_raw( - GROUPS_SQL, start_date, end_date, json.dumps(list(key_aliases)), json.dumps(list(models)) - ), + key_aliases_json: Final = json.dumps(list(key_aliases)) + models_json: Final = json.dumps(list(models)) + group_rows, error_rows, key_alias_rows, model_rows = await asyncio.gather( + prisma_client.db.query_raw(GROUPS_SQL, start_date, end_date, key_aliases_json, models_json), + prisma_client.db.query_raw(ERROR_BREAKDOWN_SQL, start_date, end_date, key_aliases_json, models_json), prisma_client.db.query_raw(KEY_ALIAS_OPTIONS_SQL, start_date, end_date), prisma_client.db.query_raw(MODEL_OPTIONS_SQL, start_date, end_date), ) @@ -135,4 +165,5 @@ async def get_cache_activity( key_aliases=[row.key_alias for row in _key_alias_rows_adapter.validate_python(key_alias_rows or [])], models=[row.model for row in _model_rows_adapter.validate_python(model_rows or [])], ), + error_breakdown=_error_buckets_adapter.validate_python(error_rows or []), ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e7b98b3cc7f..5d52af3b364 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -71,7 +71,6 @@ from litellm.proxy.auth.budget_throttle import ( ) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation -from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec from litellm.proxy.common_utils.http_parsing_utils import ( _safe_get_request_headers, _safe_get_request_query_params, @@ -87,6 +86,8 @@ from litellm.proxy.common_utils.user_api_key_cache import ( object_permission_cache_key, tag_cache_key, tag_registry_cache_key, + team_membership_auth_cache_key, + team_membership_reservation_cache_key, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -1967,7 +1968,7 @@ async def get_team_membership( if user_id is None or team_id is None: return None - _key: Final = f"team_membership:{user_id}:{team_id}" + _key: Final = team_membership_reservation_cache_key(user_id=user_id, team_id=team_id) # check if in cache cached_membership_obj: Final = await user_api_key_cache.async_get_cache( @@ -2402,6 +2403,116 @@ async def _cache_team_object( ) +async def invalidate_team_member_spend_state( + user_id: str, + team_id: str, + user_api_key_cache: UserApiKeyCache, + new_spend: float | None = None, +) -> None: + """ + Clear every cached read path for one team member's budget so a spend + reset or a raised cap takes effect on the next request instead of + waiting on the membership cache's TTL. + + Two independently-keyed cache entries hold the same LiteLLM_TeamMembership + row: user_api_key_auth.py's admission check writes ``{team_id}_{user_id}``, + while budget_reservation.py's pre-call reservation and auth_checks.py's own + get_team_membership() (used by _check_team_member_budget) both write + ``team_membership:{user_id}:{team_id}``. Both formats must be invalidated + explicitly; writing one does not refresh the other. All keys are also + broadcast (LIT-3803): each worker's own in-memory copy (membership object, + spend counter, or the counter's own short-TTL DB-floor marker) survives + eviction elsewhere until its TTL, so the handling worker alone clearing its + copy leaves every other worker still enforcing the pre-reset budget. + + ``new_spend`` is only passed by reset_team_member_spend_fn, which knows the + exact post-reset value: it is SET everywhere (matching /key/{key}/reset_spend's + own precedent) rather than deleted, so a worker's next read reflects it + directly instead of re-deriving it through a DB reseed. team_member_update + only changes the budget cap, not the tracked spend, so it passes no + new_spend; the live spend counter is untouched in that case (deleting it + would force a reseed from the DB's own spend column, which lags the live + counter via periodic batch writes, briefly under-enforcing the raised cap + against a spend value lower than what was actually tracked) and only the + membership caches carrying the new cap are invalidated. + + The floor marker (``spend_db_floor:``, proxy_server.py's + _authoritative_floor_spend) caches the pre-reset DB spend for + SPEND_DB_FLOOR_CACHE_TTL_SECONDS; left stale after a real reset, a request + landing on the pod that cached it can read that higher floor and raise the + counter right back above the just-reset spend. It is overwritten here with + the post-reset floor (not merely deleted) and _authoritative_floor_spend + re-checks the marker after its DB read, so a floor read already in flight + on this pod when the reset commits cannot clobber it with the pre-reset + value. Both keys are broadcast as SETs carrying new_spend, not deletes: + every subscriber (remote pods AND this pod's own, which receives its own + message) writes the post-reset value, so the self-delivered message cannot + erase the guard just written here. + + Raises HTTPException(503) if Redis still holds the stale pre-reset counter + after both the SET and the fallback DELETE fail: budget checks read Redis + first, so returning success would leave the old value authoritative for + every worker despite the DB write having committed. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + evict_and_broadcast, + publish_auth_cache_invalidation, + ) + + if new_spend is not None: + from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache + + spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}" + spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}" + + spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60) + except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up + verbose_proxy_logger.warning( + "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next " + "read reseeds from the DB rather than keeping the stale pre-reset value authoritative", + spend_counter_key, + e, + ) + try: + await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key) + except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success + verbose_proxy_logger.warning( + "Failed to delete stale spend counter %s in Redis after a failed reset write", + spend_counter_key, + exc_info=True, + ) + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={ # mutable-ok: HTTPException.detail takes a dict + "error": "Spend was reset in the database, but Redis is unreachable and still " + "holds the pre-reset counter. Retry once Redis is reachable." + }, + ) from e + + spend_counter_cache.in_memory_cache.set_cache( + key=spend_db_floor_key, + value=new_spend, + ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ) + await publish_auth_cache_invalidation(cache_key=spend_counter_key, new_value=new_spend, ttl=60) + await publish_auth_cache_invalidation( + cache_key=spend_db_floor_key, + new_value=new_spend, + ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ) + + await evict_and_broadcast( + cache_keys=( + team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + ), + user_api_key_cache=user_api_key_cache, + ) + + async def delete_cache_team_object( team_id: str, team_alias: str | None, @@ -2629,20 +2740,9 @@ async def _get_team_object_from_user_api_key_cache( async def _get_team_object_from_cache( key: str, - proxy_logging_obj: ProxyLogging | None, user_api_key_cache: UserApiKeyCache, parent_otel_span: Span | None, ) -> LiteLLM_TeamTableCachedObj | None: - ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ## - if proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache: - cached_raw: Final = await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache( - key=key, parent_otel_span=parent_otel_span - ) - if cached_raw is not None: - from_internal: Final = CacheCodec.deserialize(cached_raw, LiteLLM_TeamTableCachedObj) - if from_internal is not None: - return from_internal - decoded: Final = await user_api_key_cache.async_get_cache( key=key, parent_otel_span=parent_otel_span, @@ -2678,7 +2778,6 @@ async def get_team_object( if not check_db_only: cached_team_obj: Final = await _get_team_object_from_cache( key=key, - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) @@ -2841,7 +2940,6 @@ async def get_team_object_by_alias( cached_team_obj: Final = await _get_team_object_from_cache( key=cache_key, - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=parent_otel_span, ) diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 233679126f8..a42187b3a44 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import EMPTY_MAPPING from litellm.integrations.otel.runtime import seed_request_identity +from litellm.litellm_core_utils.core_helpers import is_expected_client_error from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, @@ -109,7 +110,12 @@ class UserAPIKeyAuthExceptionHandler: request=request, use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True, ) - verbose_proxy_logger.exception( + log_fn: Final = ( + verbose_proxy_logger.error + if is_expected_client_error(e) and not litellm.log_client_error_tracebacks + else verbose_proxy_logger.exception + ) + log_fn( "litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s", e, requester_ip, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 658d176f6a7..28d76e6799c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -87,7 +87,10 @@ from litellm.proxy.common_utils.http_parsing_utils import ( populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_auth_cache_key, +) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import ( @@ -1970,8 +1973,10 @@ async def _user_api_key_auth_builder( # Check 3. Check if user is in their team budget if not skip_budget_checks and valid_token.team_member_spend is not None: - if prisma_client is not None: - _cache_key: Final = f"{valid_token.team_id}_{valid_token.user_id}" + _user_id: Final = valid_token.user_id + _team_id: Final = valid_token.team_id + if prisma_client is not None and _user_id is not None and _team_id is not None: + _cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id) team_member_info = await user_api_key_cache.async_get_cache( key=_cache_key, @@ -1979,25 +1984,21 @@ async def _user_api_key_auth_builder( ) if team_member_info is None: # read from DB - _user_id: Final = valid_token.user_id - _team_id: Final = valid_token.team_id - - if _user_id is not None and _team_id is not None: - _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, - include={"litellm_budget_table": True}, + _db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first( + where={ + "user_id": _user_id, + "team_id": _team_id, + }, + include={"litellm_budget_table": True}, + ) + if _db_member is not None: + team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=team_member_info, + model_type=LiteLLM_TeamMembership, + ttl=5, ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) if team_member_info is not None and team_member_info.litellm_budget_table is not None: team_member_budget: Final = team_member_info.litellm_budget_table.max_budget @@ -2013,11 +2014,16 @@ async def _user_api_key_auth_builder( max_budget=team_member_budget, ) if team_member_spend > team_member_budget: + _entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}" raise litellm.BudgetExceededError( current_cost=team_member_spend, max_budget=team_member_budget, + message=( + f"Budget has been exceeded! TeamMember={_entity_id} " + f"Current cost: {team_member_spend}, Max budget: {team_member_budget}" + ), entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{valid_token.team_id}", + entity_id=_entity_id, ) # Check 3. If token is expired diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index dbbf9cb673e..6fd743deb44 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -33,7 +33,7 @@ from litellm.constants import ( UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -1138,24 +1138,25 @@ async def open_sse_before_first_byte( ) -def _is_azure_model_router_request(model: str) -> bool: +def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool: """ - Check if the requested model is an Azure Model Router. + Check if a request went down the Azure Model Router route. - Azure Model Router models follow the pattern: - - azure_ai/model_router/ - - azure_ai/model-router - - model_router/ - - model-router + ``model`` here is what the *client* sent, a model group alias with no ``model_router/`` + prefix, so matching on it alone only works when the operator happened to put "model-router" + in the alias. Where the response is in hand its stamp answers this outright, so callers + should pass ``hidden_params``. Args: model: The requested model name + hidden_params: ``_hidden_params`` from the response, when the caller has it Returns: bool: True if this is an Azure Model Router request """ - model_lower: Final = model.lower() - return "model-router" in model_lower or "model_router" in model_lower + from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo + + return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params) def _override_openai_response_model( @@ -1223,7 +1224,7 @@ def _override_openai_response_model( return # Check if this is an Azure Model Router request - if so, preserve the actual model used - if _is_azure_model_router_request(requested_model): + if _is_azure_model_router_request(requested_model, hidden_params): verbose_proxy_logger.debug( "%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.", log_context, @@ -1379,7 +1380,12 @@ def _log_llm_api_exception(e: Exception) -> None: "litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled" ) return - verbose_proxy_logger.exception("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e) + log_fn: Final = ( + verbose_proxy_logger.error + if is_expected_client_error(e) and not litellm.log_client_error_tracebacks + else verbose_proxy_logger.exception + ) + log_fn("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e) async def _cancel_llm_call_on_client_disconnect( diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index acdc9728390..fb2ca6372c0 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -12,6 +12,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( ) if TYPE_CHECKING: + from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -30,15 +31,24 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @dataclass(frozen=True, slots=True) class _CacheInvalidationMessage: cache_key: str + new_value: float | None = None + ttl: float | None = None -def _cache_invalidation_message_json(cache_key: str) -> str: - return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key))) +def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: + message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) + return json.dumps({field: value for field, value in message.items() if value is not None}) -def _cache_key_from_message_data(data: object) -> str | None: +def _finite_number_or_none(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _message_from_data(data: object) -> _CacheInvalidationMessage | None: if isinstance(data, bytes): - data = data.decode("utf-8", errors="replace") + data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str if not isinstance(data, str): return None try: @@ -48,14 +58,28 @@ def _cache_key_from_message_data(data: object) -> str | None: if not isinstance(parsed, dict): return None cache_key: Final = parsed.get("cache_key") - return cache_key if isinstance(cache_key, str) else None + if not isinstance(cache_key, str): + return None + return _CacheInvalidationMessage( + cache_key=cache_key, + new_value=_finite_number_or_none(parsed.get("new_value")), + ttl=_finite_number_or_none(parsed.get("ttl")), + ) -async def publish_auth_cache_invalidation(cache_key: str) -> None: +async def publish_auth_cache_invalidation( + cache_key: str, new_value: float | None = None, ttl: float | None = None +) -> None: """ Best-effort broadcast so every worker drops its local in-memory copy of a mutated management object; without this, only the handling worker and Redis are evicted and other workers keep serving the stale object until its TTL. + + Passing ``new_value`` broadcasts a SET instead of a delete: every subscriber + (including the publishing worker's own, which receives its own message) + writes the value into its additional in-memory caches rather than deleting + the key. A spend reset uses this so the handler's self-delivered message + cannot erase the freshly-written post-reset counter or floor marker. """ redis_cache: Final = coordination_redis_cache() if redis_cache is None: @@ -68,7 +92,10 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None: cache_key, ) return - await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key)) + await client.publish( + auth_cache_invalidation_channel(redis_cache), + _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl), + ) except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) @@ -95,15 +122,17 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") + __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache") def __init__( self, redis_cache: "RedisCache", user_api_key_cache: "UserApiKeyCache", + additional_in_memory_caches: Sequence["InMemoryCache"] = (), ) -> None: self._redis_cache = redis_cache self._user_api_key_cache = user_api_key_cache + self._additional_in_memory_caches = tuple(additional_in_memory_caches) self._task: asyncio.Task[None] | None = None def start(self) -> None: @@ -160,12 +189,18 @@ class AuthCacheInvalidationSubscriber: def _apply_message(self, message: object) -> None: data: Final = message.get("data") if isinstance(message, dict) else None - cache_key: Final = _cache_key_from_message_data(data) - if cache_key is None: + parsed: Final = _message_from_data(data) + if parsed is None: + return + if parsed.new_value is not None: + for additional_cache in self._additional_in_memory_caches: + additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl) return in_memory_cache: Final = self._user_api_key_cache.in_memory_cache if in_memory_cache is not None: - in_memory_cache.delete_cache(cache_key) + in_memory_cache.delete_cache(parsed.cache_key) + for additional_cache in self._additional_in_memory_caches: + additional_cache.delete_cache(parsed.cache_key) @staticmethod async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 93d51bdd461..b8df0105b7b 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -200,6 +200,21 @@ def end_user_restricted_registry_cache_key() -> str: return "end_user_restricted_registry" +def team_membership_auth_cache_key(team_id: str, user_id: str) -> str: + """Cache key one team member's ``LiteLLM_TeamMembership`` row is stored under for the admission check.""" + return f"{team_id}_{user_id}" + + +def team_membership_reservation_cache_key(user_id: str, team_id: str) -> str: + """Cache key the pre-call budget reservation stores the same ``LiteLLM_TeamMembership`` row under. + + Deliberately not unified with ``team_membership_auth_cache_key``: the two readers wrote independent + keys before this file existed, so a fix that invalidates one must invalidate both explicitly rather + than assume a single write is visible to both. + """ + return f"team_membership:{user_id}:{team_id}" + + def get_management_object_ttl(cache: DualCache) -> float: """ In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...). diff --git a/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml b/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml new file mode 100644 index 00000000000..ac8a3db2d21 --- /dev/null +++ b/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml @@ -0,0 +1,30 @@ +# Web search via Microsoft Foundry (Grounding with Bing Search / the built-in +# web_search tool), called through the Foundry Responses API. +# +# Configure the provider with env vars (setup and pricing are in the LiteLLM docs; +# the code lives in litellm/llms/azure/search/transformation.py): +# BING_GROUNDING_PROJECT_ENDPOINT (required) the Foundry project endpoint +# BING_GROUNDING_MODEL (required) a model deployment in that project +# BING_GROUNDING_CONNECTION_ID (optional) a Grounding with Bing connection id; +# without it the built-in web_search tool is used +# BING_GROUNDING_TOKEN (optional) an Entra bearer token; without it (and +# without api_key) azure-identity mints one + +model_list: + - model_name: claude-sonnet + litellm_params: + model: bedrock/us.anthropic.claude-sonnet-5 + aws_region_name: us-east-1 + +search_tools: + - search_tool_name: bing-grounding-search + litellm_params: + search_provider: bing_grounding + # Optional: an Azure API key instead of BING_GROUNDING_TOKEN / azure-identity + # api_key: os.environ/AZURE_AI_API_KEY + +litellm_settings: + callbacks: ["websearch_interception"] + websearch_interception_params: + enabled_providers: ["bedrock"] + search_tool_name: bing-grounding-search diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 414beabe014..83caa92ede5 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,6 +1,7 @@ import asyncio import io import traceback +from collections.abc import Sequence from typing import Final import orjson @@ -33,10 +34,10 @@ async def uploadfile_to_bytesio(upload: UploadFile) -> io.BytesIO: async def batch_to_bytesio( - uploads: list[UploadFile] | None, + uploads: Sequence[UploadFile] | None, ) -> list[io.BytesIO] | None: """ - Convert a list of UploadFiles to a list of BytesIO buffers, or None. + Convert a sequence of UploadFiles to a list of BytesIO buffers, or None. """ if not uploads: return None diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 2271501d480..c803d87e436 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -210,7 +210,6 @@ async def _patch_team_caches_add_access_group( for team_id in team_ids: cached_team = await _get_team_object_from_cache( key=f"team_id:{team_id}", - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=None, ) @@ -240,7 +239,6 @@ async def _patch_team_caches_remove_access_group( for team_id in team_ids: cached_team = await _get_team_object_from_cache( key=f"team_id:{team_id}", - proxy_logging_obj=proxy_logging_obj, user_api_key_cache=user_api_key_cache, parent_otel_span=None, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 01254d5c064..49461d7841d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -16,7 +16,7 @@ import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import Annotated, Final, NamedTuple, Protocol, TypedDict, TypeVar, cast +from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -56,6 +56,7 @@ from litellm.proxy._types import ( PatchTeamRequest, ProxyErrorTypes, ProxyException, + ResetSpendRequest, SpecialManagementEndpointEnums, SpecialModelNames, SpecialProxyStrings, @@ -84,6 +85,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_membership, get_team_object, get_user_object, + invalidate_team_member_spend_state, ) from litellm.proxy.auth.auth_utils import ( enforce_batch_enqueued_token_limit_is_admin_only, @@ -3392,7 +3394,7 @@ async def team_member_update( Update team member budgets and team member role """ - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) @@ -3491,6 +3493,12 @@ async def team_member_update( budget_patch=budget_patch, team_default_budget_id=team_default_budget_id, ) + if budget_patch: + await invalidate_team_member_spend_state( + user_id=received_user_id, + team_id=data.team_id, + user_api_key_cache=user_api_key_cache, + ) ### update team member role if data.role is not None: @@ -3527,6 +3535,125 @@ async def team_member_update( ) +def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAuth) -> None: + """ + _verify_team_access authorizes a team admin (or org admin) over their own + team, with no check that the target user_id differs from the caller. Left + unchecked, that admin could target their own LiteLLM_TeamMembership row and + repeatedly reset it to 0 right before it crosses their per-member cap, + consuming the shared team budget without the configured limit ever binding. + Only a proxy admin may reset an admin's own spend. + """ + if user_id == user_api_key_dict.user_id and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + _raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.") + + +def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn: + detail: Final = {"error": message} # mutable-ok: HTTPException.detail takes a dict + raise HTTPException(status_code=status_code, detail=detail) + + +def _validate_team_member_reset_spend_value( + reset_to: object, + membership: LiteLLM_TeamMembership, +) -> float: + if not isinstance(reset_to, (int, float)): + _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a float") + + reset_to_float: Final = float(reset_to) + if not math.isfinite(reset_to_float) or reset_to_float < 0: + _raise_reset_spend_error(status.HTTP_400_BAD_REQUEST, "reset_to must be a finite number >= 0") + + current_spend: Final = membership.spend or 0.0 + if reset_to_float > current_spend: + _raise_reset_spend_error( + status.HTTP_400_BAD_REQUEST, + f"reset_to ({reset_to_float}) must be <= current spend ({current_spend})", + ) + + max_budget: Final = membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None + if max_budget is not None and reset_to_float > max_budget: + _raise_reset_spend_error( + status.HTTP_400_BAD_REQUEST, + f"reset_to ({reset_to_float}) must be <= budget ({max_budget})", + ) + + return reset_to_float + + +@router.post( + "/team/{team_id}/member/{user_id}/reset_spend", + tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence + dependencies=(Depends(user_api_key_auth),), +) +@management_endpoint_wrapper +async def reset_team_member_spend_fn( + team_id: str, + user_id: str, + data: ResetSpendRequest, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +): + """ + Reset a team member's tracked spend against their per-member budget. + + A member's spend is tracked separately from both their own personal + budget and the team's own budget (LiteLLM_TeamMembership.spend), so + neither /user/update nor /team/update can clear it: this is the only + endpoint that does. The cross-pod spend counter and cached membership + reads are invalidated so the reset takes effect on the member's next + request rather than waiting on the membership cache's TTL. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + _raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None") + + team_obj: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict) + _check_not_resetting_own_spend(user_id=user_id, user_api_key_dict=user_api_key_dict) + + membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument + "user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument + } + _membership_row: Final = await _team_membership_db(prisma_client).find_unique( + where=membership_where, + include={"litellm_budget_table": True}, # mutable-ok: prisma client requires a plain dict include= argument + ) + if _membership_row is None: + _raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.") + membership: Final = LiteLLM_TeamMembership.model_validate(_membership_row.model_dump()) + + current_spend: Final = membership.spend or 0.0 + reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership) + + await _team_membership_db(prisma_client).update( + where=membership_where, + data={"spend": reset_to}, # mutable-ok: prisma client requires a plain dict data= argument + ) + + await invalidate_team_member_spend_state( + user_id=user_id, + team_id=team_id, + user_api_key_cache=user_api_key_cache, + new_spend=reset_to, + ) + + return { # mutable-ok: matches this router's established untyped-response-dict convention + "team_id": team_id, + "user_id": user_id, + "spend": reset_to, + "previous_spend": current_spend, + "max_budget": membership.litellm_budget_table.max_budget if membership.litellm_budget_table else None, + } + + def _create_results_from_response( members: list[Member], response: TeamAddMemberResponse, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dced4e26b6..0abcdeaf3f6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2555,6 +2555,12 @@ async def _authoritative_floor_spend( if db_spend is None: return None + # a spend reset that committed during the DB read above wrote the post-reset + # floor to the marker; keep it over this read's now-stale pre-commit value + rechecked: Final = spend_counter_cache.in_memory_cache.get_cache(key=marker_key) + if rechecked is not None: + return float(rechecked) + spend_counter_cache.in_memory_cache.set_cache( key=marker_key, value=db_spend, @@ -6798,6 +6804,7 @@ class ProxyConfig: subscriber: Final = AuthCacheInvalidationSubscriber( redis_cache=redis_cache, user_api_key_cache=user_api_key_cache, + additional_in_memory_caches=(spend_counter_cache.in_memory_cache,), ) self.auth_cache_invalidation_subscriber = subscriber subscriber.start() diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index ce6c9330620..149f9b960a1 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -25,7 +25,11 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_utils import get_model_from_request from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.user_api_key_cache import end_user_cache_key, tag_cache_key +from litellm.proxy.common_utils.user_api_key_cache import ( + end_user_cache_key, + tag_cache_key, + team_membership_reservation_cache_key, +) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -546,7 +550,9 @@ async def _get_team_member_budget_counter( if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None - membership_cache_key: Final = f"team_membership:{valid_token.user_id}:{team_object.team_id}" + membership_cache_key: Final = team_membership_reservation_cache_key( + user_id=valid_token.user_id, team_id=team_object.team_id + ) cached_team_membership: Final = await user_api_key_cache.async_get_cache(key=membership_cache_key) team_membership: LiteLLM_TeamMembership | None = None if isinstance(cached_team_membership, LiteLLM_TeamMembership): diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b6f695db512..d931d712a92 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -444,7 +444,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs or None ) raw_model: Final = cast(str, kwargs.get("model") or "") - model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) + model_name: Final = ( + standard_logging_payload.get("model") if standard_logging_payload is not None else None + ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: Final[SpendLogsPayload] = SpendLogsPayload( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 86d954c0913..2bbdad9c487 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -91,7 +91,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert -from litellm.litellm_core_utils.core_helpers import coerce_token_limit +from litellm.litellm_core_utils.core_helpers import coerce_token_limit, is_expected_client_error from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -2575,20 +2575,36 @@ class ProxyLogging: api_key="", ) - # log the custom exception - await litellm_logging_obj.async_failure_handler( - exception=original_exception, - traceback_exception=traceback.format_exc(), + await self._dispatch_proxy_only_failure_handlers( + litellm_logging_obj=litellm_logging_obj, + original_exception=original_exception, ) - threading.Thread( - target=litellm_logging_obj.failure_handler, - args=( - original_exception, - traceback.format_exc(), - ), - daemon=True, - ).start() + @staticmethod + async def _dispatch_proxy_only_failure_handlers( + litellm_logging_obj: Logging, + original_exception: Exception | None, + ) -> None: + """Runs the async failure handler plus the threaded sync handler. Expected + client (4xx) errors skip traceback formatting unless + litellm.log_client_error_tracebacks is set.""" + include_traceback: Final = litellm.log_client_error_tracebacks or not is_expected_client_error( + original_exception + ) + traceback_str: Final = traceback.format_exc() if include_traceback else "" + await litellm_logging_obj.async_failure_handler( + exception=original_exception, + traceback_exception=traceback_str, + ) + + threading.Thread( + target=litellm_logging_obj.failure_handler, + args=( + original_exception, + traceback_str, + ), + daemon=True, + ).start() async def post_call_success_hook( self, diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index 74688c58eaf..d985a546fa7 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -4,6 +4,7 @@ from typing import Any, Final 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 @@ -759,7 +760,14 @@ async def video_edit( ) data: Final = await _read_request_body(request=request) - data["video_id"] = video_reference_to_id(data.pop("video", None)) + 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) decoded: Final = decode_video_id_with_provider(data["video_id"]) provider_from_id: Final = decoded.get("custom_llm_provider") diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 15a6f18a6bb..c8f7842aebf 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -277,6 +277,8 @@ def rerank( if api_key is None: raise ValueError("TogetherAI API key is required, please set 'TOGETHERAI_API_KEY' in your environment") + api_base = dynamic_api_base or optional_params.api_base or litellm.api_base or "https://api.together.ai/v1" + response = together_rerank.rerank( model=model, query=query, @@ -286,6 +288,7 @@ def rerank( return_documents=return_documents, max_chunks_per_doc=max_chunks_per_doc, api_key=api_key, + api_base=api_base, _is_async=_is_async, ) elif _custom_llm_provider == litellm.LlmProviders.JINA_AI: diff --git a/litellm/router.py b/litellm/router.py index 6ee474730c9..c658bc441ba 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -168,6 +168,11 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, ) +from litellm.router_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -8341,6 +8346,7 @@ class Router: ) = litellm.get_llm_provider( model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), + api_base=deployment.litellm_params.api_base, ) # done reading model["litellm_params"] # Check if provider is supported: either in enum or JSON-configured @@ -9448,6 +9454,8 @@ class Router: except Exception: model_info = None + deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict) + # get llm provider litellm_model, llm_provider = "", "" try: @@ -9490,6 +9498,7 @@ class Router: "model_group": user_facing_model_group_name, "providers": [llm_provider], **model_info, + "supported_reasoning_efforts": None, } ) else: @@ -9567,6 +9576,11 @@ class Router: if model_info.get("rpm", None) is not None and _deployment_rpm is None: _deployment_rpm = model_info.get("rpm") + model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( + model_group_info.supported_reasoning_efforts, + resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=deployment_is_mapped), + ) + if _deployment_tpm is not None: if total_tpm is None: total_tpm = 0 diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 4c07679c5ca..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 @@ -275,6 +275,7 @@ _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]") @@ -593,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], ...]: @@ -612,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: @@ -1363,6 +1403,7 @@ class ComplexityRouter(CustomLogger): messages, current_ask=prompt, window_size=self.config.classifier_context_window_size, + budget_chars=self.config.classifier_context_budget_chars, per_turn_chars=self.config.classifier_context_per_turn_chars, include_assistant=include_assistant, marker_pairs=self._reminder_markers, diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index d3c4bd7938b..9907407d84d 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -49,7 +49,7 @@ TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = ( DEFAULT_TIER_DISTANCE_PENALTY: Final[float] = 0.5 DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE: Final[int] = 3 -DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS: Final[int] = 200 +DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS: Final[int] = 8000 class KeywordTierRule(BaseModel): @@ -645,12 +645,30 @@ class ComplexityRouterConfig(BaseModel): "classifier_type is 'llm'." ), ) - classifier_context_per_turn_chars: int = Field( - default=DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS, + classifier_context_budget_chars: int = Field( + default=DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, + ge=0, + description=( + "Maximum characters of prior-turn text quoted to the LLM classifier, across the whole " + "context window, per classification call. Turns are taken newest first and quoted whole " + "while they fit, so a conversation small enough to quote entirely is never cut; once the " + "budget runs out the older turns are dropped whole and only the turn straddling the " + "boundary is truncated, into whatever space is left. The current ask and the caller's " + "system prompt sit outside this budget and are always sent in full, as does the numbering " + "each quoted turn carries. A budget under 120 leaves no room to quote a turn and " + "suppresses the block; set classifier_context_window_size to 0 to turn context off " + "deliberately. Only applies when classifier_type is 'llm'." + ), + ) + classifier_context_per_turn_chars: int | None = Field( + default=None, gt=0, description=( - "Maximum character length for each prior turn's text in the classifier context window. " - "Turns exceeding this are truncated. Only applies when classifier_type is 'llm'." + "Optional cap on each individual prior turn's text, applied before " + "classifier_context_budget_chars bounds the block. Unset by default, so one long turn may " + "spend the whole budget, which is usually what a follow-up needs; set it when no single " + "turn should dominate the context the classifier sees. A capped turn keeps its opening " + "and its ending with the middle elided. Only applies when classifier_type is 'llm'." ), ) classifier_context_include_assistant_turns: bool = Field( @@ -662,9 +680,9 @@ class ComplexityRouterConfig(BaseModel): "word 'yes'. When enabled, classifier_context_window_size counts the last N turns of the " "conversation across both roles rather than the last N user turns, and assistant text is " "sent to the classifier model, which may be a different deployment or provider than the " - "routed completion model. Assistant replies share classifier_context_per_turn_chars with " - "user turns, so raise it if replies are truncated before the part that carries the " - "difficulty. Off by default because enabling it shifts tier decisions, and therefore " + "routed completion model. Assistant replies spend classifier_context_budget_chars " + "alongside user turns, so raise it if the oldest turns stop being quoted once replies " + "join the window. Off by default because enabling it shifts tier decisions, and therefore " "spend, for an already-deployed router. Only applies when classifier_type is 'llm'." ), ) diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py new file mode 100644 index 00000000000..3e4478e5e21 --- /dev/null +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -0,0 +1,146 @@ +"""Resolve which reasoning_effort values a deployment, and by intersection a model group, accepts. + +The model map's supports_*_reasoning_effort flags are the only signal, and each level's polarity +mirrors how a request path reads that same flag. medium and high are unconditional for a reasoning +model. minimal and low are opt-out: openai/chat/gpt_5_transformation.py refuses them only when the +map says false. xhigh and max are opt-in. none is opt-out everywhere except the azure gpt-5 family, +whose config raises UnsupportedParamsError without an explicit true. + +xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at +all: every entry carrying supports_max_reasoning_effort is Claude-family, and +anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort +path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an +explicit flag is the only signal that the tier is a real one rather than litellm rounding the level +to a budget, and a missing flag costs advisory metadata rather than a rejected request. + +A deployment the map describes with no effort flags at all resolves to None rather than to the +opt-out defaults. 689 of the map's 854 reasoning entries carry no flag, and the o-series, xai and +bedrock nova entries among them take neither none nor minimal, so composing a set out of the +defaults alone would advertise levels those providers reject. + +The advertisement order is the REASONING_EFFORT declaration order, which is presentation only. It +is not a strength scale and does not reconcile with bedrock's output_config ceiling order in +llms/bedrock/common_utils.py, which ranks max below xhigh while the thinking-budget constants rank +it above. +""" + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, get_args + +import litellm +from litellm.types.llms.openai import REASONING_EFFORT + +REASONING_EFFORT_ADVERTISEMENT_ORDER: Final = get_args(REASONING_EFFORT) +_EMPTY_ENTRY: Final[Mapping[str, object]] = MappingProxyType({}) + +_EFFORT_FLAGS: Final = ( + ("none", "supports_none_reasoning_effort"), + ("minimal", "supports_minimal_reasoning_effort"), + ("low", "supports_low_reasoning_effort"), + ("xhigh", "supports_xhigh_reasoning_effort"), + ("max", "supports_max_reasoning_effort"), +) +_OPT_OUT_EFFORTS: Final = ("minimal", "low") +_OPT_IN_EFFORTS: Final = ("xhigh", "max") +_UNCONDITIONAL_EFFORTS: Final = frozenset(("medium", "high")) + + +def _bare_model_entry(model_info: Mapping[str, object]) -> Mapping[str, object]: + """The unprefixed twin of a provider-prefixed map entry, which is where the flags often live: + azure/gpt-5-mini carries none of them while gpt-5-mini carries all three. The request-path + gates resolve through the same twin (_supports_factory, #20885), so reading it here is what + keeps the advertisement and the gate on the same answer.""" + key: Final = model_info.get("key") + provider: Final = model_info.get("litellm_provider") + if not isinstance(key, str) or not isinstance(provider, str) or not key.startswith(f"{provider}/"): + return _EMPTY_ENTRY + entry: Final[Mapping[str, object] | None] = litellm.model_cost.get(key.removeprefix(f"{provider}/")) + return entry if entry is not None else _EMPTY_ENTRY + + +def _declared_effort_flags(model_info: Mapping[str, object]) -> Mapping[str, object]: + bare: Final = _bare_model_entry(model_info) + return MappingProxyType( + { + effort: model_info.get(flag) if model_info.get(flag) is not None else bare.get(flag) + for effort, flag in _EFFORT_FLAGS + } + ) + + +def _supports_none_reasoning_effort(model_info: Mapping[str, object], flag: object) -> bool: + """Opt-in only where a request path refuses the level. AzureOpenAIGPT5Config raises + UnsupportedParamsError on reasoning_effort='none' without an explicit true, and it is selected + only for the gpt-5 family, so every other azure deployment keeps the opt-out default.""" + if model_info.get("litellm_provider") != "azure": + return flag is not False + + from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + + key: Final = model_info.get("key") + if not isinstance(key, str) or not AzureOpenAIGPT5Config.is_model_gpt_5_model(key): + return flag is not False + return flag is True + + +def deployment_is_catalog_mapped( + resolved_model_info: Mapping[str, object] | None, + operator_model_info: Mapping[str, object], +) -> bool: + """Whether the model map described this deployment, as opposed to the operator describing it. + + Every deployment is registered in the cost map under its own id, so a mode the operator wrote + on an off-map deployment reads back here exactly like one the catalog supplied. Excluding it is + what stops such a deployment from claiming to be a known non-reasoning model and emptying the + levels its mapped siblings agree on. + """ + if resolved_model_info is None or resolved_model_info.get("mode") is None: + return False + return operator_model_info.get("mode") is None + + +def resolve_supported_reasoning_efforts( + model_info: Mapping[str, object], + *, + deployment_is_mapped: bool, +) -> tuple[str, ...] | None: + """None = nothing is known about this deployment, so it must not narrow its group; () = a known + model that accepts no effort level, which correctly empties the group. + + Telling those apart needs provenance the flattened ModelInfo does not carry. A deployment the + map does not describe arrives with supports_reasoning None, exactly like a mapped non-reasoning + model: 2273 of the map's 3165 entries omit the key rather than setting it false, so reading an + unset flag as () would let one custom deployment empty every level its mapped siblings agree + on. deployment_is_mapped is that provenance, and an operator who wants either answer for an + off-map deployment gets it by setting supports_reasoning explicitly. + """ + supports_reasoning: Final = model_info.get("supports_reasoning") + if supports_reasoning is not True: + return () if supports_reasoning is False or deployment_is_mapped else None + + flags: Final = _declared_effort_flags(model_info) + if all(value is None for value in flags.values()): + return None + + opt_out: Final = frozenset(effort for effort in _OPT_OUT_EFFORTS if flags[effort] is not False) + opt_in: Final = frozenset(effort for effort in _OPT_IN_EFFORTS if flags[effort] is True) + none_level: Final = ( + frozenset(("none",)) if _supports_none_reasoning_effort(model_info, flags["none"]) else frozenset() + ) + allowed: Final = opt_out | _UNCONDITIONAL_EFFORTS | opt_in | none_level + return tuple(effort for effort in REASONING_EFFORT_ADVERTISEMENT_ORDER if effort in allowed) + + +def intersect_supported_reasoning_efforts( + current: Sequence[str] | None, + resolved: Sequence[str] | None, +) -> tuple[str, ...] | None: + """Deployments without metadata (None) never narrow the group; an effort survives only when + every deployment with metadata accepts it, so the group offers nothing routing could reject.""" + if resolved is None: + return tuple(current) if current is not None else None + if current is None: + return tuple(resolved) + keep: Final = frozenset(current) & frozenset(resolved) + return tuple(effort for effort in REASONING_EFFORT_ADVERTISEMENT_ORDER if effort in keep) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index cc6eccbf3e0..d3b0f334163 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -685,6 +685,7 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai class AnthropicThinkingParam(TypedDict, total=False): type: ReadOnly[Literal["enabled", "adaptive", "disabled"]] budget_tokens: int + display: ReadOnly[Literal["summarized", "omitted"]] class ANTHROPIC_HOSTED_TOOLS(str, Enum): diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5665aa3277a..6ae2e31fe60 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from collections.abc import Sequence from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal @@ -396,7 +397,7 @@ class OutputConfigBlock(TypedDict, total=False): class CommonRequestObject(TypedDict, total=False): # common request object across sync + async flows additionalModelRequestFields: dict - additionalModelResponseFieldPaths: list[str] + additionalModelResponseFieldPaths: Sequence[str] inferenceConfig: InferenceConfig system: list[SystemContentBlock] toolConfig: ToolConfigBlock diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index e7a3f825455..4a6c4a5bbb5 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1840,7 +1840,7 @@ ResponsesAPIStreamingResponse = Annotated[ ] -REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh"] +REASONING_EFFORT = Literal["none", "minimal", "low", "medium", "high", "xhigh", "max"] class OpenAIRealtimeStreamSession(TypedDict, total=False): diff --git a/litellm/types/router.py b/litellm/types/router.py index 9fd5cfa96ef..d4c735387a5 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -637,6 +637,7 @@ class ModelGroupInfo(BaseModel): supports_url_context: bool = Field(default=False) supports_reasoning: bool = Field(default=False) supports_function_calling: bool = Field(default=False) + supported_reasoning_efforts: tuple[str, ...] | None = Field(default=None) supported_openai_params: list[str] | None = Field(default=[]) configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4d59650f410..73f46bd2181 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3855,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 7fed0bab935..5cbd0519032 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8610,6 +8610,12 @@ class ProviderConfigManager: ) return BedrockPassthroughConfig() + elif LlmProviders.BEDROCK_MANTLE == provider: + from litellm.llms.bedrock_mantle.passthrough.transformation import ( + BedrockMantlePassthroughConfig, + ) + + return BedrockMantlePassthroughConfig() elif LlmProviders.VLLM == provider or LlmProviders.HOSTED_VLLM == provider: from litellm.llms.vllm.passthrough.transformation import ( VLLMPassthroughConfig, @@ -9096,6 +9102,7 @@ class ProviderConfigManager: from litellm.llms.apiserpent.search.transformation import ( APISerpentSearchConfig, ) + from litellm.llms.azure.search.transformation import BingGroundingSearchConfig from litellm.llms.bedrock.search.transformation import AgentCoreSearchConfig from litellm.llms.brave.search.transformation import BraveSearchConfig from litellm.llms.dataforseo.search.transformation import DataForSEOSearchConfig @@ -9137,6 +9144,7 @@ class ProviderConfigManager: SearchProviders.TINYFISH: TinyfishSearchConfig, SearchProviders.AGENTCORE: AgentCoreSearchConfig, SearchProviders.NIMBLE: NimbleSearchConfig, + SearchProviders.BING_GROUNDING: BingGroundingSearchConfig, } config_class: Final = PROVIDER_TO_CONFIG_MAP.get(provider, None) if config_class is None: diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 978849ac006..445435a30fa 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -5,6 +5,8 @@ from collections.abc import Coroutine from functools import partial from typing import Final, Literal, overload +from httpx._types import FileContent + import litellm from litellm.constants import DEFAULT_VIDEO_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT @@ -1344,6 +1346,8 @@ async def avideo_edit( extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, + *, + video: FileContent | None = None, **kwargs, ) -> VideoObject: """ @@ -1359,6 +1363,7 @@ async def avideo_edit( video_edit, video_id=video_id, prompt=prompt, + video=video, timeout=timeout, custom_llm_provider=custom_llm_provider, extra_headers=extra_headers, @@ -1396,6 +1401,8 @@ def video_edit( extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, + *, + video: FileContent | None = None, **kwargs, ) -> VideoObject | Coroutine[object, object, VideoObject]: """ @@ -1444,6 +1451,7 @@ def video_edit( return base_llm_http_handler.video_edit_handler( prompt=prompt, video_id=video_id, + video_file=video, video_provider_config=provider_config, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index c36018feb9d..1aa4c7cd060 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -17155,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", @@ -49008,12 +49016,13 @@ "output_cost_per_token": 3.3e-05, "output_cost_per_token_above_272k_tokens": 4.95e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49040,12 +49049,13 @@ "output_cost_per_token": 1.32e-05, "output_cost_per_token_above_272k_tokens": 1.98e-05, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ @@ -49072,12 +49082,13 @@ "output_cost_per_token": 1.32e-06, "output_cost_per_token_above_272k_tokens": 1.98e-06, "litellm_provider": "bedrock_mantle", - "max_input_tokens": 1000000, + "max_input_tokens": 1050000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", "use_openai_responses_path": true, "supported_endpoints": [ + "/v1/chat/completions", "/v1/responses" ], "supported_modalities": [ diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 03318718fb5..bbfbbbc1cfb 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3018 + "limit": 3016 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2016 + "limit": 2015 }, "ANN202": { "limit": 852 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2919 + "limit": 2918 }, "C401": { "limit": 8 @@ -168,7 +168,7 @@ "limit": 3 }, "RET504": { - "limit": 176 + "limit": 175 }, "RUF012": { "limit": 240 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/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index e0ee40a65c5..cf912ddaa25 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -94,9 +94,9 @@ E2E_FIXTURE_MODE=record E2E_FIXTURE_DIR=/tmp/e2e-fixtures E2E_RESET_SPEND_LOGS=1 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: 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 75261b301bd..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 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, CI wiring) +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/conftest.py b/tests/e2e/conftest.py index 1bc3ed98eb6..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", 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 d43446e4b5a..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): @@ -171,7 +171,7 @@ 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) diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index 8ef9afbfa1c..c3f8865f218 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -5,6 +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/mcp_tests/conftest.py b/tests/mcp_tests/conftest.py index d1dc3ec7216..5823893afc0 100644 --- a/tests/mcp_tests/conftest.py +++ b/tests/mcp_tests/conftest.py @@ -7,6 +7,7 @@ import pytest import litellm import asyncio +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @pytest.fixture(scope="session") @@ -38,6 +39,8 @@ def setup_and_teardown(): yield # Teardown code (executes after the yield point) + # LoggingWorker carries still-queued coroutines onto the next test's loop, where they'd log into that test's callbacks + asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue()) loop.close() # Close the loop created earlier asyncio.set_event_loop(None) # Remove the reference to the loop diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 7ee745b311e..fc9f675f837 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -1,6 +1,9 @@ import os import pytest import asyncio +import subprocess +import sys +from pathlib import Path from typing import Optional from unittest.mock import AsyncMock, patch @@ -24,12 +27,20 @@ from mcp.types import Tool as MCPTool, CallToolResult, TextContent class TestMCPLogger(CustomLogger): def __init__(self): self.standard_logging_payload = None + self.mcp_tool_call_payloads = [] super().__init__() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): print("success event") - self.standard_logging_payload = kwargs.get("standard_logging_object", None) - print(f"Captured standard_logging_payload: {self.standard_logging_payload}") + payload = kwargs.get("standard_logging_object", None) + self.standard_logging_payload = payload + # Async success events from other calls (e.g. a mocked acompletion whose + # log task is delivered late) race with the MCP event for the single + # last-writer slot; keep MCP tool calls in their own list so assertions + # are order-independent. + if payload is not None and payload.get("call_type") == "call_mcp_tool": + self.mcp_tool_call_payloads.append(payload) + print(f"Captured standard_logging_payload: {payload}") def _set_authorized_user(server_ids): @@ -138,7 +149,11 @@ async def test_mcp_cost_tracking(): # wait 1-2 seconds for logging to be processed await asyncio.sleep(2) - logged_standard_logging_payload = test_logger.standard_logging_payload + logged_standard_logging_payload = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print("logged_standard_logging_payload", logged_standard_logging_payload) # Add assertions @@ -277,7 +292,11 @@ async def test_mcp_cost_tracking_per_tool(): # wait for logging to be processed await asyncio.sleep(2) - logged_standard_logging_payload_1 = test_logger.standard_logging_payload + logged_standard_logging_payload_1 = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print( "logged_standard_logging_payload_1", logged_standard_logging_payload_1 ) @@ -290,6 +309,7 @@ async def test_mcp_cost_tracking_per_tool(): # Reset logger for second test test_logger.standard_logging_payload = None + test_logger.mcp_tool_call_payloads.clear() # Test 2: Call cheap_tool - should cost 0.1 response2 = await mcp_server_tool_call( @@ -300,7 +320,11 @@ async def test_mcp_cost_tracking_per_tool(): # wait for logging to be processed await asyncio.sleep(2) - logged_standard_logging_payload_2 = test_logger.standard_logging_payload + logged_standard_logging_payload_2 = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print( "logged_standard_logging_payload_2", logged_standard_logging_payload_2 ) @@ -329,16 +353,7 @@ async def test_mcp_cost_tracking_per_tool(): assert mock_client.call_tool.call_count == 2 -class MCPLoggerHook(CustomLogger): - def __init__(self): - self.standard_logging_payload = None - super().__init__() - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - print("success event") - self.standard_logging_payload = kwargs.get("standard_logging_object", None) - print(f"Captured standard_logging_payload: {self.standard_logging_payload}") - +class MCPLoggerHook(TestMCPLogger): async def async_post_mcp_tool_call_hook( self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time ) -> Optional[MCPPostCallResponseObject]: @@ -436,9 +451,55 @@ async def test_mcp_tool_call_hook(): await asyncio.sleep(2) # check logged standard logging payload - logged_standard_logging_payload = test_logger.standard_logging_payload + logged_standard_logging_payload = ( + test_logger.mcp_tool_call_payloads[-1] + if test_logger.mcp_tool_call_payloads + else None + ) print("logged_standard_logging_payload", logged_standard_logging_payload) assert ( logged_standard_logging_payload is not None ), "Standard logging payload should not be None" assert logged_standard_logging_payload["response_cost"] == 1.42 + + +_QUEUED_LOGGING_OUTLIVES_TEST = ''' +import time + +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + +ran_at = [] + + +async def _record_run(): + ran_at.append(time.monotonic()) + + +async def test_1_leaves_logging_queued_behind_a_stopped_worker(): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run()) + await GLOBAL_LOGGING_WORKER.stop() + assert ran_at == [] + + +async def test_2_starts_after_the_previous_tests_logging_ran(): + started_at = time.monotonic() + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(_record_run()) + await GLOBAL_LOGGING_WORKER.flush() + assert [t < started_at for t in ran_at] == [True, False] +''' + + +def test_logging_queued_by_one_test_is_drained_before_the_next(tmp_path: Path): + """Regression: a logging coroutine queued by one test must not run inside a later test (it would log into that + test's callbacks, which is how test_mcp_tool_call_hook captured a gpt-4o-mini payload under xdist).""" + (tmp_path / "conftest.py").write_text((Path(__file__).parent / "conftest.py").read_text()) + (tmp_path / "pyproject.toml").write_text('[tool.pytest.ini_options]\nasyncio_mode = "auto"\n') + (tmp_path / "test_queued_logging.py").write_text(_QUEUED_LOGGING_OUTLIVES_TEST) + result = subprocess.run( + [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", "test_queued_logging.py"], + cwd=tmp_path, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/proxy_behavior/management/test_team_member_reset_spend.py b/tests/proxy_behavior/management/test_team_member_reset_spend.py new file mode 100644 index 00000000000..ec2c78139fe --- /dev/null +++ b/tests/proxy_behavior/management/test_team_member_reset_spend.py @@ -0,0 +1,152 @@ +import uuid + +import pytest + +from .actors import Actor +from .conftest import create_scratch_team + +pytestmark = pytest.mark.asyncio(loop_scope="session") + +_SEED_SPEND = 5.0 +_RESET_TO = 2.0 + + +# POST /team/{team_id}/member/{user_id}/reset_spend. The handler gate is +# _verify_team_access (proxy admin / team admin of this team / org admin of +# the team's org) — the same gate /team/member_update uses, so this mirrors +# that file's matrix exactly. +_MATRIX = [ + ("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200), + ("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200), + ("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200), + ("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403), + ("alpha/owner", Actor.OWNER, "alpha", 403), + ("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403), + ("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403), + ("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403), + ("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403), + ("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200), + ("beta/org_admin", Actor.ORG_ADMIN, "beta", 403), + ("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403), + ("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200), +] + + +async def _seed_target(prisma, world, shape: str, team_id: str, member_id: str) -> None: + if shape == "alpha": + await create_scratch_team( + prisma, + team_id, + organization_id=world.org_a_id, + admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id], + ) + elif shape == "beta": + await create_scratch_team(prisma, team_id, organization_id=world.org_b_id) + else: # pragma: no cover - guard + pytest.fail(f"unknown shape={shape}") + await prisma.db.litellm_teammembership.create( + data={"user_id": member_id, "team_id": team_id, "spend": _SEED_SPEND} + ) + + +@pytest.mark.parametrize( + "actor,shape,expected_status", + [(a, sh, s) for (_id, a, sh, s) in _MATRIX], + ids=[s[0] for s in _MATRIX], +) +async def test_team_member_reset_spend_authz_matrix( + actor: Actor, + shape: str, + expected_status: int, + proxy_client, + prisma, + scratch, + world, +): + member_id = scratch.tag("member") + await _seed_target(prisma, world, shape, scratch.prefix, member_id) + caller = world.keys[actor] + + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_spend", + headers={"Authorization": f"Bearer {caller.cleartext}"}, + json={"reset_to": _RESET_TO}, + ) + assert ( + resp.status_code == expected_status + ), f"{actor.value} {shape}: {resp.status_code} {resp.text}" + + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": member_id, "team_id": scratch.prefix}} + ) + assert row is not None + if expected_status == 200: + assert row.spend == _RESET_TO + else: + assert row.spend == _SEED_SPEND, "denied but spend reset" + + +async def test_team_member_reset_spend_missing_team_is_404(proxy_client, world): + resp = await proxy_client.post( + f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_spend_missing_membership_is_404( + proxy_client, prisma, scratch, world +): + """A well-formed team but a user_id with no LiteLLM_TeamMembership row is 404.""" + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 404, resp.text + + +async def test_team_member_reset_spend_above_current_spend_is_400( + proxy_client, prisma, scratch, world +): + member_id = scratch.tag("member") + await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id) + await prisma.db.litellm_teammembership.create( + data={"user_id": member_id, "team_id": scratch.prefix, "spend": 1.0} + ) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{member_id}/reset_spend", + headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"}, + json={"reset_to": 5.0}, + ) + assert resp.status_code == 400, resp.text + + +async def test_team_member_reset_spend_team_admin_cannot_reset_own_spend( + proxy_client, prisma, scratch, world +): + """A team admin targeting their own LiteLLM_TeamMembership row is 403: unchecked, an + admin could repeatedly zero their own spend right before it crosses their per-member + cap, consuming the shared team budget without the configured limit ever binding.""" + team_admin = world.keys[Actor.TEAM_ADMIN] + await create_scratch_team( + prisma, + scratch.prefix, + organization_id=world.org_a_id, + admin_user_ids=[team_admin.user_id], + ) + await prisma.db.litellm_teammembership.create( + data={"user_id": team_admin.user_id, "team_id": scratch.prefix, "spend": _SEED_SPEND} + ) + resp = await proxy_client.post( + f"/team/{scratch.prefix}/member/{team_admin.user_id}/reset_spend", + headers={"Authorization": f"Bearer {team_admin.cleartext}"}, + json={"reset_to": 0.0}, + ) + assert resp.status_code == 403, resp.text + row = await prisma.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": team_admin.user_id, "team_id": scratch.prefix}} + ) + assert row is not None and row.spend == _SEED_SPEND, "denied but spend reset" diff --git a/tests/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/caching/test_redis_cluster_node_isolation.py b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py index f4cd3ab20ef..c16ceec8c31 100644 --- a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py +++ b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py @@ -28,6 +28,14 @@ if TYPE_CHECKING: from redis.asyncio.cluster import RedisCluster as _AsyncRedisClusterType +class _NodeClassWithPerConnectionRecovery: + def update_active_connections_for_reconnect(self) -> None: ... + + +class _NodeClassWithoutPerConnectionRecovery: + pass + + class _FakeClusterNode: def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None: self.name = name @@ -47,7 +55,9 @@ class _FakeNodesManager: def _build_cluster_instance() -> "_AsyncRedisClusterType": - cluster_cls = get_litellm_async_redis_cluster_class() + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithoutPerConnectionRecovery + ) instance = cluster_cls.__new__(cluster_cls) instance.RedisClusterRequestTTL = 1 instance.reinitialize_counter = 0 @@ -58,6 +68,33 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType": return instance +def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None: + """Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level + connection error per-connection, the factory must NOT install the copied override, + whose node.disconnect() also kills connections other coroutines are mid-operation on.""" + from redis.asyncio.cluster import RedisCluster + + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithPerConnectionRecovery + ) + + assert cluster_cls is RedisCluster + + +def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None: + """Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(), + so those versions must keep litellm's per-node isolation override.""" + from redis.asyncio.cluster import RedisCluster + + cluster_cls = get_litellm_async_redis_cluster_class( + cluster_node_class=_NodeClassWithoutPerConnectionRecovery + ) + + assert cluster_cls is not RedisCluster + assert issubclass(cluster_cls, RedisCluster) + assert "_execute_command" in cluster_cls.__dict__ + + @pytest.mark.asyncio @pytest.mark.parametrize("error_cls", [RedisConnectionError, RedisTimeoutError]) async def test_node_level_error_resets_only_that_node_not_the_whole_client(error_cls: type[Exception]) -> None: diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 1fb74b2b7bf..6ca48ce63b8 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -2,7 +2,7 @@ import datetime import json import os import unittest -from typing import TYPE_CHECKING, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple from unittest.mock import ANY, MagicMock, Mock, patch import httpx @@ -1585,10 +1585,16 @@ def test_map_reasoning_effort_adds_summary_detailed(monkeypatch): assert result_dict["summary"] == "custom_summary" print("✓ Dict input is passed through without modification") - # Test 5: None/unknown values return None - result_unknown = handler._map_reasoning_effort("unknown_value") - assert result_unknown is None - print("✓ Unknown reasoning_effort values return None") + # Test 5: every REASONING_EFFORT level reaches the provider, and anything else (a typo, an + # unshipped level, "default") is dropped so the request still succeeds at the provider default + from litellm.types.llms.openai import Reasoning + + for effort in ("max", "xhigh", "none"): + result_passthrough = handler._map_reasoning_effort(effort) + assert result_passthrough == Reasoning(effort=effort) + for dropped in ("ultra", "hgih", "unknown_value", "", "default"): + assert handler._map_reasoning_effort(dropped) is None + print("✓ Enumerated levels pass through and unknown ones are dropped") print( "✓ All reasoning_effort behaviors work correctly with flag/env var control" @@ -2438,6 +2444,32 @@ def test_map_optional_params_preserves_reasoning_summary(): assert responses_api_request["reasoning"]["summary"] == "detailed" +@pytest.mark.parametrize("reasoning_effort", ["max", "high"]) +def test_transform_request_bedrock_mantle_tools_keeps_reasoning_effort(monkeypatch, reasoning_effort): + """Regression for reasoning_effort=max being dropped on the chat -> Responses bridge (issue #38084).""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + monkeypatch.setattr(litellm, "reasoning_auto_summary", False) + monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False) + handler: Final = LiteLLMResponsesTransformationHandler() + + result: Final = handler.transform_request( + model="openai.gpt-5.6-sol", + messages=[{"role": "user", "content": "Say pong"}], + optional_params={ + "reasoning_effort": reasoning_effort, + "tools": [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + }, + litellm_params={"custom_llm_provider": "bedrock_mantle"}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert result["reasoning"] == {"effort": reasoning_effort} + + def test_map_optional_params_tool_choice_chat_nested_to_responses_api(): """Chat tool_choice must become Responses ToolChoiceFunction (top-level name).""" from litellm.completion_extras.litellm_responses_transformation.transformation import ( 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/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index c8c36032793..6f513ce1bd4 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -478,10 +478,10 @@ def test_generic_cost_per_token_minimax_m3_above_512k_tokens(_local_model_cost_m ], ) def test_generic_cost_per_token_bedrock_mantle_gpt56_long_context(_local_model_cost_map, model): - """Bedrock GPT-5.6 supports a 1M context window, billed at the long-context rates above 272K.""" + """Bedrock GPT-5.6 enforces a 1,050,000-token context window, billed at the long-context rates above 272K.""" model_cost_map = litellm.model_cost[model] - assert model_cost_map["max_input_tokens"] == 1000000 + assert model_cost_map["max_input_tokens"] == 1050000 cached_tokens = 100000 completion_tokens = 1000 diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 44fa8fc8ae0..44f91c98d81 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -19,9 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( def test_get_format_from_file_id(): - unified_file_id = ( - "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf" - ) + unified_file_id = "litellm_proxy:application/pdf;unified_id,cbbe3534-8bf8-4386-af00-f5f6b7e370bf" format = get_format_from_file_id(unified_file_id) @@ -48,9 +46,7 @@ def test_update_messages_with_model_file_ids(): model_file_id_mapping = {file_id: {"my_model_id": "provider_file_id"}} - updated_messages = update_messages_with_model_file_ids( - messages, model_id, model_file_id_mapping - ) + updated_messages = update_messages_with_model_file_ids(messages, model_id, model_file_id_mapping) assert updated_messages == [ { @@ -143,9 +139,7 @@ def test_add_system_prompt_to_messages_merge_with_first_system(): {"role": "system", "content": "Existing system prompt."}, {"role": "user", "content": "Hello"}, ] - result = add_system_prompt_to_messages( - messages, "You are helpful.", merge_with_first_system=True - ) + result = add_system_prompt_to_messages(messages, "You are helpful.", merge_with_first_system=True) assert result == [ {"role": "system", "content": "You are helpful.\n\nExisting system prompt."}, {"role": "user", "content": "Hello"}, @@ -155,9 +149,7 @@ def test_add_system_prompt_to_messages_merge_with_first_system(): def test_add_system_prompt_to_messages_merge_with_first_system_adds_new_when_no_system(): """When merge_with_first_system=True but no system message, adds new one at start.""" messages = [{"role": "user", "content": "Hello"}] - result = add_system_prompt_to_messages( - messages, "You are helpful.", merge_with_first_system=True - ) + result = add_system_prompt_to_messages(messages, "You are helpful.", merge_with_first_system=True) assert result == [ {"role": "system", "content": "You are helpful."}, {"role": "user", "content": "Hello"}, @@ -492,14 +484,8 @@ def test_update_messages_with_model_file_ids_tolerates_non_dict_content_items(): messages_token_ids_batch = [{"role": "user", "content": [[15496, 995], [9906, 0]]}] # Both should pass through unchanged without raising. - assert ( - update_messages_with_model_file_ids(messages_token_ids, "model-A", {}) - == messages_token_ids - ) - assert ( - update_messages_with_model_file_ids(messages_token_ids_batch, "model-A", {}) - == messages_token_ids_batch - ) + assert update_messages_with_model_file_ids(messages_token_ids, "model-A", {}) == messages_token_ids + assert update_messages_with_model_file_ids(messages_token_ids_batch, "model-A", {}) == messages_token_ids_batch class TestExtractFileDataBareStr: @@ -645,9 +631,7 @@ class TestUnpackLegacyDefs: definitions = { f"L{i}": { "type": "object", - "properties": { - f"x{j}": {"$ref": f"#/definitions/L{i + 1}"} for j in range(fanout) - }, + "properties": {f"x{j}": {"$ref": f"#/definitions/L{i + 1}"} for j in range(fanout)}, } for i in range(depth) } @@ -712,9 +696,7 @@ class TestUnpackLegacyDefs: schema = { "type": "object", - "properties": { - f"r{i}": {"$ref": f"#/components/schemas/T{i}"} for i in range(50) - }, + "properties": {f"r{i}": {"$ref": f"#/components/schemas/T{i}"} for i in range(50)}, "components": { "schemas": { f"T{i}": { @@ -739,9 +721,7 @@ class TestTextCompletionPromptToMessages: text_completion_prompt_to_messages, ) - assert text_completion_prompt_to_messages("summarize this") == ( - {"role": "user", "content": "summarize this"}, - ) + assert text_completion_prompt_to_messages("summarize this") == ({"role": "user", "content": "summarize this"},) def test_list_of_strings_becomes_one_message_each(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -970,3 +950,80 @@ class TestCustomToolFormatShapeConversion: for weird in ({}, {"type": "grammar"}, {"type": "future_format", "x": 1}): assert convert_custom_tool_format_to_chat_shape(dict(weird)) in (weird, {"type": "grammar", "grammar": {}}) assert convert_custom_tool_format_to_responses_shape(dict(weird)) == weird + + +# --- x-litellm-model upload-path decoding (litellm #29830) ------------------- + + +def _xlitellm_encoded(raw_id: str, model: str) -> str: + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + return encode_file_id_with_model(raw_id, model) + + +def test_update_messages_with_model_file_ids_decodes_xlitellm_encoded_id(): + """x-litellm-model upload returns `file-;model,)>`. + Without decoding, the encoded id leaks to upstream OpenAI and errors as + 'Files [...] were not found'. Decode it back to raw provider id.""" + raw_id = "file-ExTuCawUqxEMjVFK6xwR9B" + encoded_id = _xlitellm_encoded(raw_id, "gpt-5.1") + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this."}, + {"type": "file", "file": {"file_id": encoded_id}}, + ], + } + ] + + updated = update_messages_with_model_file_ids(messages, "model-A", {}) + + assert updated[0]["content"][1]["file"]["file_id"] == raw_id + + +def test_update_responses_input_with_model_file_ids_decodes_xlitellm_encoded_id(): + """Same bug on /v1/responses path. Without decoding the encoded id (>64 + chars), OpenAI rejects with 'string too long. Expected ... maximum length + 64'.""" + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + update_responses_input_with_model_file_ids, + ) + + raw_id = "file-ExTuCawUqxEMjVFK6xwR9B" + encoded_id = _xlitellm_encoded(raw_id, "gpt-5.1") + input_items = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize."}, + {"type": "input_file", "file_id": encoded_id}, + ], + } + ] + + updated = update_responses_input_with_model_file_ids(input_items) + + assert updated[0]["content"][1]["file_id"] == raw_id + + +def test_update_messages_xlitellm_decode_does_not_override_mapping(): + """If the call-site already resolved a provider id via the mapping, that + wins. The new decode fallback runs only when no mapping match.""" + raw_id = "file-ExTuCawUqxEMjVFK6xwR9B" + encoded_id = _xlitellm_encoded(raw_id, "gpt-5.1") + mapping = {encoded_id: {"model-A": "provider-explicit-id"}} + messages = [ + { + "role": "user", + "content": [ + {"type": "file", "file": {"file_id": encoded_id}}, + ], + } + ] + + updated = update_messages_with_model_file_ids(messages, "model-A", mapping) + + assert updated[0]["content"][0]["file"]["file_id"] == "provider-explicit-id" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 3a7e06d085a..6265779b90d 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1,6 +1,8 @@ import base64 import json +import logging import os +from typing import Final from unittest.mock import MagicMock, patch import pytest @@ -3309,6 +3311,66 @@ def test_get_tool_calls_from_response_include_all_choices_reads_every_choice(): assert names == ["tool_alpha", "tool_beta"] +def test_get_tool_calls_from_response_silences_redacted_arguments(caplog): + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) + + response: Final = { + "choices": [ + { + "message": { + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "Read", + "arguments": "redacted-by-litellm", + }, + } + ] + } + } + ] + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tool_calls: Final = get_tool_calls_from_response(response) + + assert tool_calls == [{"id": "call_1", "name": "Read", "arguments": {}}] + assert "Failed to parse tool call arguments" not in caplog.text + + +def test_get_tool_calls_from_response_warns_for_malformed_arguments(caplog): + from litellm.litellm_core_utils.prompt_templates.factory import ( + get_tool_calls_from_response, + ) + + response: Final = { + "choices": [ + { + "message": { + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "Read", + "arguments": "not-json", + }, + } + ] + } + } + ] + } + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + tool_calls: Final = get_tool_calls_from_response(response) + + assert tool_calls == [{"id": "call_1", "name": "Read", "arguments": {}}] + assert "Failed to parse tool call arguments" in caplog.text + + def test_group_tool_exchanges_pairs_assistant_with_its_tool_rows(): from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index aa7f2990c13..33a33711698 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -255,3 +255,26 @@ class TestRedactNestedMatchAndRegexKeys: def test_passes_through_none_and_str(self): assert redact_nested_match_and_regex_keys(None) is None assert redact_nested_match_and_regex_keys("plain") == "plain" + + +class TestIsExpectedClientError: + def test_status_ranges(self): + from litellm.litellm_core_utils.core_helpers import is_expected_client_error + + class WithStatusCode(Exception): + def __init__(self, status_code): + self.status_code = status_code + + class WithCode(Exception): + def __init__(self, code): + self.code = code + + assert is_expected_client_error(WithStatusCode(400)) is True + assert is_expected_client_error(WithStatusCode(429)) is True + assert is_expected_client_error(WithStatusCode(499)) is True + assert is_expected_client_error(WithStatusCode(500)) is False + assert is_expected_client_error(WithStatusCode(399)) is False + assert is_expected_client_error(WithCode("403")) is True + assert is_expected_client_error(WithCode("invalid_request_error")) is False + assert is_expected_client_error(Exception("no status")) is False + assert is_expected_client_error(None) is False diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py index bda7ab4afc6..6cacd119030 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py +++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_endpoint_match.py @@ -133,3 +133,54 @@ class TestGetLlmProviderRejectsAttackerSmuggledApiBase: assert provider == "groq" assert dynamic_api_key == "server-real-groq-key" + + +class TestTogetherApiBaseResolvesProvider: + """ + Regression for the Together host migration: both the current + ``api.together.ai`` host and the legacy ``api.together.xyz`` host must + resolve to ``together_ai`` when passed as ``api_base``. Before the fix + the endpoint list carried the legacy host but the provider-mapping + chain had no branch for it, so the match fell through with a None + provider and the deployment failed with "LLM Provider NOT provided". + """ + + @pytest.mark.parametrize( + "api_base", + [ + "https://api.together.ai/v1", + "https://api.together.xyz/v1", + ], + ) + def test_together_api_base_resolves_to_together_ai(self, api_base, monkeypatch): + monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env") + + model, provider, dynamic_api_key, returned_api_base = get_llm_provider( + model="some-model", + api_base=api_base, + ) + + assert provider == "together_ai" + assert dynamic_api_key == "together-key-from-env" + assert returned_api_base == api_base + assert model == "some-model" + + def test_explicit_api_key_beats_together_env_key(self, monkeypatch): + monkeypatch.setenv("TOGETHER_API_KEY", "together-key-from-env") + + _, provider, dynamic_api_key, _ = get_llm_provider( + model="some-model", + api_base="https://api.together.ai/v1", + api_key="explicit-caller-key", + ) + + assert provider == "together_ai" + assert dynamic_api_key == "explicit-caller-key" + + def test_together_default_api_base_is_together_ai(self, monkeypatch): + monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False) + + _, provider, _, api_base = get_llm_provider(model="together_ai/some-model") + + assert provider == "together_ai" + assert api_base == "https://api.together.ai/v1" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1cac0cce4ef..f93daa61570 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - import time import httpx @@ -3867,6 +3866,90 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob assert payload["litellm_call_id"] == call_id +# ── Azure Model Router selected-model attribution ──────────────────────────── + + +def _model_router_response(selected_model: str, stamp: bool): + """A ModelResponse as AzureModelRouterConfig hands it back, with or without the stamp.""" + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + from litellm.types.utils import ModelResponse + + response = ModelResponse(model=selected_model) + response._hidden_params = ( + {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {} + ) + return response + + +def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj): + """ + The selected model must win off the stamp, not off "model-router" appearing in the + requested model. An operator whose model group is named anything else was invisible + to the name check, so their logs and spend rows named the router instead. + """ + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "azure_ai/smart-pick", + "custom_llm_provider": "azure_ai", + "messages": [], + "litellm_params": {"metadata": {}}, + }, + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=True + ), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning" + + +def test_standard_logging_payload_keeps_requested_model_without_router_stamp( + logging_obj, +): + """ + Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp + is what redirects attribution rather than the response model winning unconditionally. + """ + import datetime + + from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, + ) + + now = datetime.datetime.now() + payload = get_standard_logging_object_payload( + kwargs={ + "model": "azure_ai/smart-pick", + "custom_llm_provider": "azure_ai", + "messages": [], + "litellm_params": {"metadata": {}}, + }, + init_response_obj=_model_router_response( + "azure_ai/grok-4-1-fast-reasoning", stamp=False + ), + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["model"] == "azure_ai/smart-pick" + + def _make_dict_logging_obj(): """Build a Logging instance configured for a non-streaming dict result.""" obj = LitellmLogging( @@ -5595,3 +5678,62 @@ def test_get_custom_logger_compatible_class_finds_v2_newrelic(monkeypatch): logging_module._in_memory_loggers.clear() monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) is_otel_v2_enabled.cache_clear() + + +class _ClientError(Exception): + def __init__(self, status_code, message): + self.status_code = status_code + self.message = message + super().__init__(message) + + +def _raise_and_catch(exc): + try: + raise exc + except Exception as caught: + return caught + + +def test_get_error_information_skips_traceback_for_expected_4xx(monkeypatch): + """Regression for LIT-6043: expected client (4xx) errors must not pay for + traceback.format_tb on every rejected request unless + litellm.log_client_error_tracebacks is enabled.""" + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + client_exc = _raise_and_catch(_ClientError(status_code=403, message="team does not allow model")) + assert client_exc.__traceback__ is not None + result = StandardLoggingPayloadSetup.get_error_information(client_exc) + assert result["traceback"] == "" + + server_exc = _raise_and_catch(_ClientError(status_code=500, message="boom")) + result = StandardLoggingPayloadSetup.get_error_information(server_exc) + assert "test_litellm_logging" in result["traceback"] + + monkeypatch.setattr(litellm, "log_client_error_tracebacks", True) + result = StandardLoggingPayloadSetup.get_error_information(client_exc) + assert "test_litellm_logging" in result["traceback"] + + +def test_failure_handler_helper_fn_builds_payload_once_per_exception(): + """Regression for LIT-6043: async and sync failure handlers both call + _failure_handler_helper_fn for the same failed request; the standardized + payload must be built once, not once per handler.""" + obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="lit-6043-1", + function_id="f", + ) + exc = _raise_and_catch(_ClientError(status_code=400, message="invalid model")) + obj._failure_handler_helper_fn(exception=exc, traceback_exception="") + first_payload = obj.model_call_details["standard_logging_object"] + assert first_payload is not None + obj._failure_handler_helper_fn(exception=exc, traceback_exception="") + assert obj.model_call_details["standard_logging_object"] is first_payload + + other_exc = _raise_and_catch(_ClientError(status_code=429, message="rate limited")) + obj._failure_handler_helper_fn(exception=other_exc, traceback_exception="") + assert obj.model_call_details["standard_logging_object"] is not first_payload diff --git a/tests/test_litellm/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 12ab536ed45..c9170efd18a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_reasoning_effort_translation.py @@ -42,7 +42,7 @@ def test_reasoning_effort_maps_to_output_config_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py index ef74249ca8e..288817dff07 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_reasoning_effort.py @@ -5,6 +5,8 @@ Verifies that reasoning_effort=None returns None for all models, including Claude Opus 4.6. """ +import pytest + from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -35,6 +37,16 @@ class TestMapReasoningEffort: ) assert result["type"] == "adaptive" + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_adaptive_mapping_requests_summarized_display(self, effort): + """Regression LIT-5714: adaptive thinking without ``display`` makes Anthropic + return a blank thinking block, so reasoning_effort callers always got + ``reasoning_content: ""``.""" + result = AnthropicConfig._map_reasoning_effort( + reasoning_effort=effort, model="claude-opus-4-6", custom_llm_provider="anthropic" + ) + assert result["display"] == "summarized" + def test_other_model_low_returns_enabled_with_budget(self): result = AnthropicConfig._map_reasoning_effort( reasoning_effort="low", model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" diff --git a/tests/test_litellm/llms/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_ai/chat/test_azure_ai_transformation.py b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py index d4fcbc823a6..11a727c9635 100644 --- a/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -201,6 +201,90 @@ def test_azure_model_router_response_shows_actual_model(): ) +def test_azure_model_router_stamps_selected_model_on_hidden_params(): + """ + The selected model must be stamped on _hidden_params, not left for downstream code to + re-derive by looking for "model-router" in the model string. Deployments whose alias + does not contain that text are invisible to the string check. + """ + from httpx import Response + + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + AzureFoundryModelInfo, + ) + from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj + from litellm.types.utils import ModelResponse + + raw_response_json = { + "id": "chatcmpl-test456", + "object": "chat.completion", + "created": 1234567890, + "model": "grok-4-1-fast-reasoning", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "pong"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + mock_response = MagicMock(spec=Response) + mock_response.json.return_value = raw_response_json + mock_response.text = json.dumps(raw_response_json) + mock_response.headers = {} + + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + logging_obj.post_call = MagicMock() + logging_obj.model_call_details = {} + + result = AzureModelRouterConfig().transform_response( + model="smart-pick", + raw_response=mock_response, + model_response=ModelResponse(), + logging_obj=logging_obj, + request_data={}, + messages=[{"role": "user", "content": "Reply with just pong"}], + optional_params={}, + litellm_params={"model": "azure_ai/model_router/smart-pick"}, + encoding=None, + api_key="test-key", + json_mode=False, + ) + + assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model + assert ( + result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] + == "azure_ai/grok-4-1-fast-reasoning" + ) + assert AzureFoundryModelInfo.get_model_router_selected_model( + result._hidden_params + ) == ("azure_ai/grok-4-1-fast-reasoning") + assert ( + AzureFoundryModelInfo.is_model_router_call( + model="smart-pick", hidden_params=result._hidden_params + ) + is True + ) + + +def test_azure_model_router_stamp_does_not_leak_across_responses(): + """ + ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written + as a fresh dict. Mutating in place would bleed the selected model into unrelated responses. + """ + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + from litellm.types.utils import ModelResponse + + untouched = ModelResponse() + + assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {}) + + def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name(): """ Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name. diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 53a432427d3..326edde743d 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -341,7 +341,7 @@ def test_messages_thinking_shape_follows_exact_azure_entry_flag(local_model_cost ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( diff --git a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py index e4402bbec49..e6aad7688d1 100644 --- a/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py +++ b/tests/test_litellm/llms/base_llm/search/test_base_search_transformation.py @@ -16,6 +16,7 @@ import pytest import litellm from litellm.llms.apiserpent.search.transformation import APISerpentSearchConfig +from litellm.llms.azure.search.transformation import BingGroundingSearchConfig from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, _is_trusted_search_api_base, @@ -59,6 +60,7 @@ _BASE_ENV_VARS = ( "TINYFISH_API_BASE", "CRW_API_BASE", "NIMBLE_API_BASE", + "BING_GROUNDING_PROJECT_ENDPOINT", ) @@ -99,6 +101,7 @@ PROVIDERS: Tuple[ProviderSpec, ...] = ( (TinyfishSearchConfig, {"TINYFISH_API_KEY": "srv"}, "caller-key", {}), (FastCRWSearchConfig, {"CRW_API_KEY": "srv"}, "caller-key", {}), (NimbleSearchConfig, {"NIMBLE_API_KEY": "srv"}, "caller-key", {}), + (BingGroundingSearchConfig, {"BING_GROUNDING_TOKEN": "srv"}, "caller-key", {}), ) _IDS = tuple(spec[0].__name__ for spec in PROVIDERS) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 604f3414775..4d2c077b548 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -366,6 +366,96 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +def test_reasoning_effort_requests_summarized_display_converse(): + """Regression LIT-5714: adaptive thinking synthesized from reasoning_effort must + request the summarized display, otherwise the provider returns a blank thinking + block and reasoning_content is always empty.""" + config = AmazonConverseConfig() + + optional_params = config.map_openai_params( + non_default_params={"reasoning_effort": "high"}, + optional_params={}, + model="bedrock/converse/us.anthropic.claude-opus-4-7", + drop_params=False, + ) + + assert optional_params["thinking"]["type"] == "adaptive" + assert optional_params["thinking"]["display"] == "summarized" + + +def test_thinking_request_adds_output_tokens_details_response_path(): + """Regression LIT-5714: the Converse usage block has no thinking-token field, so + thinking requests must ask for ``/usage/output_tokens_details`` via + ``additionalModelResponseFieldPaths``.""" + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "high"}, + }, + litellm_params={}, + headers={}, + ) + + assert result["additionalModelResponseFieldPaths"] == ("/usage/output_tokens_details",) + + +def test_request_without_thinking_omits_response_field_paths(): + config = AmazonConverseConfig() + + result = config._transform_request( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + messages=[{"role": "user", "content": "hi"}], + optional_params={"maxTokens": 256}, + litellm_params={}, + headers={}, + ) + + assert "additionalModelResponseFieldPaths" not in result + + +def test_transform_usage_prefers_provider_reasoning_tokens(): + """Regression LIT-5714: provider-reported thinking tokens must win over the + token_counter estimate derived from visible reasoning text.""" + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + provider_reasoning_tokens=1033, + ) + + assert usage.completion_tokens_details.reasoning_tokens == 1033 + assert usage.completion_tokens_details.text_tokens == 3002 - 1033 + + +def test_transform_usage_falls_back_to_estimate_without_provider_tokens(): + config = AmazonConverseConfig() + + usage = config.transform_usage( + {"inputTokens": 40, "outputTokens": 300, "totalTokens": 340}, + reasoning_content="a short reasoning summary", + thinking_ran=True, + ) + + assert usage.completion_tokens_details.reasoning_tokens > 0 + assert usage.completion_tokens_details.reasoning_tokens < 300 + + +def test_thinking_tokens_parsed_from_additional_model_response_fields(): + parsed = AmazonConverseConfig.thinking_tokens_from_additional_fields( + {"usage": {"output_tokens_details": {"thinking_tokens": 92}}} + ) + assert parsed == 92 + assert AmazonConverseConfig.thinking_tokens_from_additional_fields(None) is None + assert AmazonConverseConfig.thinking_tokens_from_additional_fields({"usage": {}}) is None + + @pytest.mark.parametrize( "model,effort,expected_effort", [ diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index cd305d8ed26..4bef59842f1 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -208,6 +208,29 @@ def test_bedrock_converse_streaming_consistent_id(): ), "All chunk IDs must match the one captured from the messageStart event" +def test_converse_streaming_usage_uses_provider_thinking_tokens(): + """Regression LIT-5714: the messageStop event carries provider thinking tokens + under ``additionalModelResponseFields``; the usage chunk must report them instead + of a token_counter estimate.""" + chunks = [ + { + "contentBlockIndex": 0, + "delta": {"reasoningContent": {"text": "thinking about it"}}, + }, + { + "stopReason": "end_turn", + "additionalModelResponseFields": {"usage": {"output_tokens_details": {"thinking_tokens": 1033}}}, + }, + {"usage": {"inputTokens": 40, "outputTokens": 3002, "totalTokens": 3042}}, + ] + + decoder = AWSEventStreamDecoder(model="bedrock/anthropic.claude-opus-4-7") + parsed = [decoder.converse_chunk_parser(chunk) for chunk in chunks] + + usage = parsed[-1].usage + assert usage.completion_tokens_details.reasoning_tokens == 1033 + + @pytest.mark.asyncio async def test_make_call_does_not_rechunk_stream_by_default(): """Re-chunking the event stream into fixed 1024-byte blocks holds small diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index d3c28302bf9..1e09afd6919 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1387,7 +1387,7 @@ def test_bedrock_messages_maps_reasoning_effort_for_adaptive_model( ) assert "reasoning_effort" not in result - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": expected_effort} @@ -2935,7 +2935,7 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem(litellm.model_cost[model], "supports_adaptive_thinking", False) diff --git a/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py new file mode 100644 index 00000000000..8c6eda605ca --- /dev/null +++ b/tests/test_litellm/llms/bedrock_mantle/passthrough/test_bedrock_mantle_passthrough_transformation.py @@ -0,0 +1,197 @@ +import json +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from botocore.credentials import Credentials + +from litellm.llms.bedrock.passthrough.transformation import BedrockPassthroughConfig +from litellm.llms.bedrock_mantle.passthrough.transformation import BedrockMantlePassthroughConfig +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.passthrough.main import llm_passthrough_route +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +MANTLE_API_BASE = "https://bedrock-mantle.us-east-2.api.aws" +INVOKE_ENDPOINT = "model/us.openai.gpt-5.6-sol/invoke" +CONVERSE_ENDPOINT = "model/us.openai.gpt-5.6-sol/converse" +REQUEST_BODY = {"messages": [{"role": "user", "content": "say pong"}], "max_completion_tokens": 64} + + +@pytest.fixture +def no_ambient_aws(monkeypatch): + for name in ( + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_KEY", + "BEDROCK_MANTLE_API_BASE", + "BEDROCK_MANTLE_REGION", + "AWS_BEDROCK_RUNTIME_ENDPOINT", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_DEFAULT_REGION", + ): + monkeypatch.delenv(name, raising=False) + + +def test_bedrock_mantle_registers_its_own_bedrock_passthrough_config(): + config = ProviderConfigManager.get_provider_passthrough_config( + model="us.openai.gpt-5.6-sol", provider=LlmProviders.BEDROCK_MANTLE + ) + assert isinstance(config, BedrockMantlePassthroughConfig) + assert isinstance(config, BedrockPassthroughConfig) + + +def test_mantle_api_base_only_lends_its_region_to_the_runtime_url(no_ambient_aws): + url, base_url = BedrockMantlePassthroughConfig().get_complete_url( + api_base=MANTLE_API_BASE, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={"api_base": MANTLE_API_BASE}, + ) + assert str(url) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}" + assert base_url == "https://bedrock-runtime.us-east-2.amazonaws.com" + + +def test_explicit_region_and_non_mantle_api_base_are_kept(no_ambient_aws): + vpc_endpoint = "https://vpce-0123.bedrock-runtime.us-east-1.vpce.amazonaws.com" + url, base_url = BedrockMantlePassthroughConfig().get_complete_url( + api_base=vpc_endpoint, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={"api_base": vpc_endpoint, "aws_region_name": "us-east-1"}, + ) + assert str(url) == f"{vpc_endpoint}/{INVOKE_ENDPOINT}" + assert base_url == vpc_endpoint + + +def test_region_falls_back_to_the_mantle_default_without_any_hint(no_ambient_aws): + url, _ = BedrockMantlePassthroughConfig().get_complete_url( + api_base=None, + api_key=None, + model="us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + request_query_params=None, + litellm_params={}, + ) + assert str(url) == f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}" + + +@pytest.mark.parametrize( + ("litellm_params", "env", "expected_bearer"), + [ + ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"), + ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"), + ({}, {"AWS_BEARER_TOKEN_BEDROCK": "aws-env-key"}, "aws-env-key"), + ], +) +def test_sign_request_uses_the_deployment_bearer_token(no_ambient_aws, monkeypatch, litellm_params, env, expected_bearer): + for name, value in env.items(): + monkeypatch.setenv(name, value) + headers, body = BedrockMantlePassthroughConfig().sign_request( + headers={}, + litellm_params=litellm_params, + request_data=REQUEST_BODY, + api_base=f"https://bedrock-runtime.us-east-1.amazonaws.com/{INVOKE_ENDPOINT}", + model="us.openai.gpt-5.6-sol", + ) + assert headers["Authorization"] == f"Bearer {expected_bearer}" + assert body is not None + assert json.loads(body) == REQUEST_BODY + + +def test_sign_request_falls_back_to_sigv4_scoped_to_the_mantle_region(no_ambient_aws): + config = BedrockMantlePassthroughConfig() + with patch.object(config, "get_credentials", return_value=Credentials("AKIA", "secret")): + headers, body = config.sign_request( + headers={}, + litellm_params={"api_base": MANTLE_API_BASE}, + request_data=REQUEST_BODY, + api_base=f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}", + model="us.openai.gpt-5.6-sol", + ) + assert headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=AKIA/") + assert "/us-east-2/bedrock/aws4_request" in headers["Authorization"] + assert body is not None + assert json.loads(body) == REQUEST_BODY + + +@pytest.mark.parametrize( + ("route_kwargs", "env", "expected_bearer"), + [ + ({"api_key": "deployment-bedrock-api-key"}, {}, "deployment-bedrock-api-key"), + ({}, {"BEDROCK_MANTLE_API_KEY": "mantle-env-key"}, "mantle-env-key"), + ], +) +def test_invoke_passthrough_route_reaches_bedrock_runtime_for_a_mantle_deployment( + no_ambient_aws, monkeypatch, route_kwargs, env, expected_bearer +): + for name, value in env.items(): + monkeypatch.setenv(name, value) + client = HTTPHandler() + with ( + patch.object(client.client, "send", return_value=MagicMock(status_code=200)), + patch.object(client.client, "build_request", wraps=client.client.build_request) as build_request, + ): + response = llm_passthrough_route( + model="bedrock_mantle/us.openai.gpt-5.6-sol", + endpoint=INVOKE_ENDPOINT, + method="POST", + api_base=MANTLE_API_BASE, + json=dict(REQUEST_BODY), + client=client, + litellm_logging_obj=MagicMock(), + **route_kwargs, + ) + assert response.status_code == 200 + sent = build_request.call_args.kwargs + assert str(sent["url"]) == f"https://bedrock-runtime.us-east-2.amazonaws.com/{INVOKE_ENDPOINT}" + assert sent["headers"]["Authorization"] == f"Bearer {expected_bearer}" + assert json.loads(sent["content"]) == REQUEST_BODY + + +def _logged_model_response(endpoint, body): + request = httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com/{endpoint}") + return BedrockMantlePassthroughConfig().logging_non_streaming_response( + model="us.openai.gpt-5.6-sol", + custom_llm_provider="bedrock_mantle", + httpx_response=httpx.Response(200, json=body, request=request), + request_data={"messages": [{"role": "user", "content": [{"text": "say pong"}]}]}, + logging_obj=MagicMock(), + endpoint=endpoint, + ) + + +def test_converse_logging_parses_the_converse_response_shape(): + result = _logged_model_response( + CONVERSE_ENDPOINT, + { + "metrics": {"latencyMs": 800.0}, + "output": {"message": {"content": [{"text": "pong"}], "role": "assistant"}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 8, "outputTokens": 5, "totalTokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5 + + +def test_invoke_logging_parses_the_openai_chat_response_shape(): + result = _logged_model_response( + INVOKE_ENDPOINT, + { + "choices": [{"finish_reason": "stop", "index": 0, "message": {"content": "pong", "role": "assistant"}}], + "created": 1787677792, + "id": "chatcmpl-regression", + "model": "us.openai.gpt-5.6-sol", + "object": "chat.completion", + "usage": {"completion_tokens": 5, "prompt_tokens": 8, "total_tokens": 13}, + }, + ) + assert result.choices[0].message.content == "pong" + assert result.usage.prompt_tokens == 8 + assert result.usage.completion_tokens == 5 diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 9e05d48a18f..00a319f99f0 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -8,7 +8,9 @@ gate, the URL construction for both paths, and the shared Bearer auth. """ import copy - +import json +import logging +from pathlib import Path import pytest from botocore.exceptions import ( @@ -623,6 +625,181 @@ class TestBedrockMantleCodexAdditionalTools: assert "additional_tools" in str(mock_debug.call_args) +class TestBedrockMantleCodexInputItemNormalization: + """Mantle 400s ("Invalid 'input': value did not match any expected variant") + on the Codex history item types agent_message, context_compaction, and + local_shell_call (verified against bedrock-mantle.us-east-1.api.aws with + openai.gpt-5.6-sol), so the config must rewrite them into supported + equivalents. agent_message is what every Codex multi-agent v2 session sends, + and its encrypted_content slot carries the verbatim plaintext payload when + the upstream model never issued encrypted args, so that slot must be + preserved, not dropped. Mantle also rejects assistant messages with + input_text content, so the rewrite must use output_text.""" + + _USER_MESSAGE = { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Continue."}], + } + + def _transform(self, input): + cfg = BedrockMantleResponsesAPIConfig() + return cfg.transform_responses_api_request( + model="openai.gpt-5.6-sol", + input=input, + response_api_optional_request_params={}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + def test_plaintext_agent_message_becomes_assistant_output_text_message(self): + body = self._transform( + input=[ + self._USER_MESSAGE, + { + "type": "agent_message", + "id": "amsg_1", + "author": "/root/arithmetic", + "recipient": "/root", + "content": [{"type": "input_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."}], + }, + ] + ) + assert body["input"] == [ + self._USER_MESSAGE, + { + "type": "message", + "role": "assistant", + "content": ({"type": "output_text", "text": "Message Type: FINAL_ANSWER\nPayload:\n2+2 is 4."},), + }, + ] + + def test_agent_message_encrypted_content_payload_is_preserved(self): + body = self._transform( + input=[ + { + "type": "agent_message", + "author": "/root", + "recipient": "/root/arithmetic", + "content": [ + {"type": "input_text", "text": "Message Type: NEW_TASK\nPayload:\n"}, + {"type": "encrypted_content", "encrypted_content": "Answer the question 'what is 2+2'."}, + ], + }, + self._USER_MESSAGE, + ] + ) + assert body["input"][0] == { + "type": "message", + "role": "assistant", + "content": ( + { + "type": "output_text", + "text": "Message Type: NEW_TASK\nPayload:\nAnswer the question 'what is 2+2'.", + }, + ), + } + + def test_agent_message_without_any_text_is_dropped(self): + body = self._transform( + input=[ + {"type": "agent_message", "author": "/root", "recipient": "/root/a", "content": []}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_context_compaction_becomes_compaction_with_same_ciphertext(self): + body = self._transform( + input=[ + {"type": "context_compaction", "id": "cc_1", "encrypted_content": "smry_abc123"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [ + {"type": "compaction", "encrypted_content": "smry_abc123"}, + self._USER_MESSAGE, + ] + + def test_context_compaction_without_ciphertext_is_dropped(self): + body = self._transform( + input=[ + {"type": "context_compaction", "id": "cc_1"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_local_shell_call_becomes_function_call_keeping_call_id_pairing(self): + body = self._transform( + input=[ + { + "type": "local_shell_call", + "id": "lsh_1", + "call_id": "call_1", + "status": "completed", + "action": {"type": "exec", "command": ["echo", "hi"]}, + }, + {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [ + { + "type": "function_call", + "call_id": "call_1", + "name": "local_shell", + "arguments": '{"type": "exec", "command": ["echo", "hi"]}', + }, + {"type": "function_call_output", "call_id": "call_1", "output": "hi\n"}, + self._USER_MESSAGE, + ] + + def test_local_shell_call_without_call_id_is_dropped(self): + body = self._transform( + input=[ + {"type": "local_shell_call", "status": "completed", "action": {"type": "exec", "command": ["ls"]}}, + self._USER_MESSAGE, + ] + ) + assert body["input"] == [self._USER_MESSAGE] + + def test_mantle_supported_item_types_pass_through_untouched(self): + supported_items = [ + self._USER_MESSAGE, + {"type": "compaction", "encrypted_content": "smry_abc123"}, + {"type": "function_call", "name": "shell", "arguments": "{}", "call_id": "call_2"}, + {"type": "function_call_output", "call_id": "call_2", "output": "ok"}, + {"type": "tool_search_call", "call_id": "call_3", "execution": "server", "arguments": {"query": "x"}}, + {"type": "tool_search_output", "call_id": "call_3", "status": "completed", "execution": "server", "tools": []}, + {"type": "compaction_trigger"}, + ] + body = self._transform(input=copy.deepcopy(supported_items)) + assert body["input"] == supported_items + + def test_string_input_passes_through(self): + body = self._transform(input="Say hi.") + assert body["input"] == "Say hi." + + def test_rewrite_is_logged_as_warning_naming_the_types(self, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + body = self._transform( + input=[ + {"type": "agent_message", "author": "a", "recipient": "b", "content": [{"type": "input_text", "text": "hi"}]}, + self._USER_MESSAGE, + ] + ) + assert body["input"][0]["role"] == "assistant" + rewrite_warnings = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.WARNING and "rewrote Codex input item type" in record.getMessage() + ] + assert rewrite_warnings == [ + "Bedrock Mantle Responses API: rewrote Codex input item type(s) ['agent_message'] that Mantle rejects." + ] + + class TestBedrockMantleResponsesRegistry: def test_registry_returns_config_for_gpt_5_5(self, local_cost_map): # gpt-5.x advertises /v1/responses in supported_endpoints (capability) @@ -1523,7 +1700,7 @@ class TestBedrockMantleResponsesPricing: assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost) assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost) assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == 1000000 + assert info["max_input_tokens"] == 1050000 assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2) assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2) @@ -1565,3 +1742,42 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models + + +def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]: + repo_root = Path(__file__).resolve().parents[4] + paths = { + "root": repo_root / "model_prices_and_context_window.json", + "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json", + } + return json.loads(paths[map_name].read_text()) + + +class TestGpt56MantleRegistryEntries: + """Locks the gpt-5.6 frontier entries to Bedrock Mantle's live behavior. + + Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna + (oversize requests 400 with "prompt tokens (N) exceed model maximum + (1050000)", and a 1,030,590-token request completes), matching the OpenAI + Bedrock guide. mode must stay "responses": Mantle's native + /v1/chat/completions rejects function tools unless reasoning_effort is + "none", so chat traffic has to keep bridging to the Responses API + (see the responses_api_bridge tests above). + """ + + @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) + @pytest.mark.parametrize( + "key", + ( + "bedrock_mantle/openai.gpt-5.6-sol", + "bedrock_mantle/openai.gpt-5.6-terra", + "bedrock_mantle/openai.gpt-5.6-luna", + ), + ) + def test_entry_matches_mantle_enforced_limits(self, map_name, key): + entry = _repo_cost_map(map_name)[key] + assert entry["max_input_tokens"] == 1050000 + assert entry["max_output_tokens"] == 128000 + assert entry["mode"] == "responses" + assert entry["use_openai_responses_path"] is True + assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index f7f89cd1d8d..16d57437043 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -56,9 +56,7 @@ async def test_async_post_streaming_status_error_should_not_wait_forever_for_bod litellm_handler = AsyncHTTPHandler() await litellm_handler.client.aclose() - litellm_handler.client = httpx.AsyncClient( - transport=httpx.MockTransport(mock_handler) - ) + litellm_handler.client = httpx.AsyncClient(transport=httpx.MockTransport(mock_handler)) try: with pytest.raises(MaskedHTTPStatusError) as exc_info: await asyncio.wait_for( @@ -202,9 +200,7 @@ async def test_ssl_verification_with_aiohttp_transport(monkeypatch: pytest.Monke transport_connector = transport._get_valid_client_session().connector assert isinstance(transport_connector, TCPConnector) - aiohttp_session = aiohttp.ClientSession( - connector=aiohttp.TCPConnector(ssl=False) - ) + aiohttp_session = aiohttp.ClientSession(connector=aiohttp.TCPConnector(ssl=False)) try: aiohttp_connector = aiohttp_session.connector assert isinstance(aiohttp_connector, aiohttp.TCPConnector) @@ -378,7 +374,8 @@ async def test_get_async_httpx_client_with_shared_session(): # Test with shared session client = get_async_httpx_client( - llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session # type: ignore + llm_provider=LlmProviders.ANTHROPIC, + shared_session=mock_session, # type: ignore ) # Verify the client was created successfully @@ -397,9 +394,7 @@ async def test_get_async_httpx_client_without_shared_session(): from litellm.types.utils import LlmProviders # Test without shared session - client = get_async_httpx_client( - llm_provider=LlmProviders.ANTHROPIC, shared_session=None - ) + client = get_async_httpx_client(llm_provider=LlmProviders.ANTHROPIC, shared_session=None) # Verify the client was created successfully assert client is not None @@ -476,11 +471,13 @@ async def test_session_reuse_integration(): # Create two clients with the same session client1 = get_async_httpx_client( - llm_provider=LlmProviders.ANTHROPIC, shared_session=mock_session # type: ignore + llm_provider=LlmProviders.ANTHROPIC, + shared_session=mock_session, # type: ignore ) client2 = get_async_httpx_client( - llm_provider=LlmProviders.OPENAI, shared_session=mock_session # type: ignore + llm_provider=LlmProviders.OPENAI, + shared_session=mock_session, # type: ignore ) # Both clients should be created successfully @@ -512,9 +509,7 @@ async def test_session_reuse_integration(): (None, None, None, False), # None value - skip configuration ], ) -def test_ssl_ecdh_curve( - env_curve, litellm_curve, expected_curve, should_call, monkeypatch -): +def test_ssl_ecdh_curve(env_curve, litellm_curve, expected_curve, should_call, monkeypatch): """Test SSL ECDH curve configuration with valid curves and precedence""" from litellm.llms.custom_httpx.http_handler import _ssl_context_cache @@ -717,9 +712,7 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout: _default_cached_client_timeout, ) - monkeypatch.setattr( - litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS - ) + monkeypatch.setattr(litellm, "request_timeout", litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) assert _default_cached_client_timeout() is _DEFAULT_TIMEOUT @@ -734,9 +727,7 @@ class TestDefaultCachedClientTimeoutHonorsRequestTimeout: assert resolved.read == 300.0 assert resolved.connect == 5.0 - def test_cached_async_client_built_with_explicit_request_timeout( - self, monkeypatch: pytest.MonkeyPatch - ): + def test_cached_async_client_built_with_explicit_request_timeout(self, monkeypatch: pytest.MonkeyPatch): from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.utils import LlmProviders @@ -1195,3 +1186,131 @@ async def test_aiohttp_session_never_replays_one_upstreams_cookie_to_another(): assert len(jar) == 0 assert dict(jar.filter_cookies(URL("https://upstream-a.example.com"))) == {} await session.close() + + +def _mint_session_on_dead_loop(handler: AsyncHTTPHandler) -> ClientSession: + """Create the transport's real ClientSession on a loop that then closes. + + This is the lifecycle of every client minted for a short-lived event loop + (the loop-id-keyed LLM client cache creates one handler per loop): the + session outlives its loop and can only ever be disposed loop-lessly. + """ + transport = handler.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport) + loop = asyncio.new_event_loop() + + async def _create() -> ClientSession: + return transport._get_valid_client_session() + + session = loop.run_until_complete(_create()) + loop.close() + return session + + +def test_finalizer_without_running_loop_closes_dead_loop_session(): + """A handler finalized with no running event loop must still dispose its + aiohttp session. + + The async close can never run in that context; without the synchronous + fallback the session and its connector are abandoned to GC and emit + "Unclosed client session" / "Unclosed connector" warnings.""" + handler = AsyncHTTPHandler(timeout=61.0) + session = _mint_session_on_dead_loop(handler) + assert not session.closed + + del handler + gc.collect() + + assert session.closed + + +@pytest.mark.asyncio +async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref(): + """With a running loop, finalization schedules an async close and must keep + a strong reference to the task until it completes — a bare create_task() + result may be collected before it runs, leaving the session unclosed.""" + handler = AsyncHTTPHandler(timeout=61.0) + transport = handler.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport) + session = transport._get_valid_client_session() + assert not session.closed + del transport + + baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) + del handler + gc.collect() + + scheduled = AsyncHTTPHandler._finalizer_close_tasks - baseline_tasks + assert len(scheduled) == 1 + + await asyncio.gather(*scheduled) + assert session.closed + assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled) + + +@pytest.mark.asyncio +async def test_sync_close_helper_respects_session_ownership(): + """The loop-less fallback closes only sessions the transport owns; a + shared session (e.g. the proxy's) must never be closed by a handler.""" + owned_handler = AsyncHTTPHandler(timeout=61.0) + owned_transport = owned_handler.client._transport + assert isinstance(owned_transport, LiteLLMAiohttpTransport) + owned_session = owned_transport._get_valid_client_session() + + baseline = set(LiteLLMAiohttpTransport._background_close_tasks) + owned_handler._dispose_wrapped_aiohttp_session() + scheduled = LiteLLMAiohttpTransport._background_close_tasks - baseline + await asyncio.gather(*scheduled) + assert owned_session.closed + + shared_session = ClientSession() + shared_handler = AsyncHTTPHandler(timeout=61.0, shared_session=shared_session) + shared_transport = shared_handler.client._transport + assert isinstance(shared_transport, LiteLLMAiohttpTransport) + assert shared_transport._owns_session is False + + shared_handler._dispose_wrapped_aiohttp_session() + assert not shared_session.closed + + await shared_session.close() + await shared_handler.close() + await owned_handler.close() + + +@pytest.mark.asyncio +async def test_finalizer_close_done_consumes_exception(): + """A failing finalizer close must have its exception retrieved by the done + callback, or asyncio emits "Task exception was never retrieved" at GC — + the same log noise the finalizer path exists to eliminate.""" + + async def failing_close() -> None: + raise RuntimeError("close failed") + + task = asyncio.get_running_loop().create_task(failing_close()) + AsyncHTTPHandler._finalizer_close_tasks.add(task) + await asyncio.sleep(0) + + AsyncHTTPHandler._on_finalizer_close_done(task) + assert task not in AsyncHTTPHandler._finalizer_close_tasks + + cancelled = asyncio.get_running_loop().create_task(asyncio.sleep(30)) + cancelled.cancel() + await asyncio.sleep(0) + AsyncHTTPHandler._on_finalizer_close_done(cancelled) + + +@pytest.mark.asyncio +async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_scheduling(): + """GC on a live loop (e.g. the app's) of a handler whose session belongs to + another, dead loop must not schedule aclose() here — that is the cross-loop + path the transport refuses — and must still dispose the session.""" + handler = AsyncHTTPHandler(timeout=61.0) + session = await asyncio.to_thread(_mint_session_on_dead_loop, handler) + assert not session.closed + + baseline_tasks = set(AsyncHTTPHandler._finalizer_close_tasks) + del handler + gc.collect() + + assert AsyncHTTPHandler._finalizer_close_tasks == baseline_tasks + assert session.closed diff --git a/tests/test_litellm/llms/openai/test_gpt5_transformation.py b/tests/test_litellm/llms/openai/test_gpt5_transformation.py index d279b119efe..a35b75a6106 100644 --- a/tests/test_litellm/llms/openai/test_gpt5_transformation.py +++ b/tests/test_litellm/llms/openai/test_gpt5_transformation.py @@ -1309,3 +1309,79 @@ def test_responses_gpt54_allow_temperature_effort_none( drop_params=False, ) assert params["temperature"] == 0.7 + + +@pytest.mark.parametrize("model", ["gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"]) +def test_gpt5_6_forwards_reasoning_effort_max_for_the_responses_bridge(config: OpenAIConfig, model: str): + """A chat request carrying tools or a reasoning summary is converted to /v1/responses further + down main.py, and that surface accepts max. This runs before litellm has decided to bridge, so + refusing max here would break the cursor thinking-max shape that works today. Plain chat + completions still answer max with a provider 400, and the capability list below is what keeps + the level out of the picker.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "max"}, + optional_params={}, + model=model, + drop_params=False, + ) + assert params["reasoning_effort"] == "max" + + +@pytest.mark.parametrize("model", ["gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"]) +def test_gpt5_6_never_advertises_reasoning_effort_max(model: str): + """/v1/chat/completions answers max with "Unsupported value: 'reasoning_effort' does not support + 'max' with this model. Supported values are: 'none', 'low', 'medium', 'high', and 'xhigh'", so no + gpt-5.6 entry asserts supports_max_reasoning_effort and the advertised set stops at xhigh.""" + from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + + resolved = resolve_supported_reasoning_efforts(litellm.get_model_info(model), deployment_is_mapped=True) + assert resolved is not None + assert "max" not in resolved + assert "xhigh" in resolved + + +def test_gpt5_6_keeps_reasoning_effort_max_on_the_responses_api( + responses_config: OpenAIResponsesAPIConfig, +): + """/v1/responses accepts max for gpt-5.6, and that is the surface the cursor thinking-max + variant resolves onto, so the responses path keeps carrying the level chat refuses.""" + params = responses_config.map_openai_params( + response_api_optional_params=ResponsesAPIOptionalRequestParams( + reasoning={"effort": "max"}, + ), + model="gpt-5.6", + drop_params=False, + ) + assert params["reasoning"] == {"effort": "max"} + + +def test_gpt5_forwards_levels_the_chat_gate_does_not_own(config: OpenAIConfig): + """Only xhigh is gated on this surface. max reaches the provider (or the responses bridge) and + is answered there, which is what happened before per-group capabilities existed.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "max"}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + assert params["reasoning_effort"] == "max" + + +def test_gpt5_rejects_xhigh_for_models_without_the_flag(config: OpenAIConfig): + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model="gpt-5.1", + drop_params=False, + ) + + +def test_gpt5_drops_xhigh_when_requested(config: OpenAIConfig): + params = config.map_openai_params( + non_default_params={"reasoning_effort": "xhigh"}, + optional_params={}, + model="gpt-5.1", + drop_params=True, + ) + assert "reasoning_effort" not in params diff --git a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py index c15554a46a5..89f52261cba 100644 --- a/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py +++ b/tests/test_litellm/llms/openai/videos/test_openai_video_transformation.py @@ -1,3 +1,5 @@ +import io + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.utils import encode_character_id_with_provider @@ -68,3 +70,48 @@ def test_wrapped_character_id_is_decoded_then_encoded_as_path_segment(): == "https://api.openai.com/v1/videos/characters/..%2F..%2Fcharacters%3Fx%3D1%23frag" ) assert params == {} + + +def test_video_edit_request_forwards_uploaded_file_as_multipart(): + """An uploaded source video must leave as a multipart ``video`` file part, + not be dropped in favor of a JSON id reference.""" + config = OpenAIVideoConfig() + source = io.BytesIO(b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isomBODY") + source.name = "clip.mp4" + + url, data, files = config.transform_video_edit_request( + prompt="make it nighttime", + video_id="", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + video_file=source, + ) + + assert url == "https://api.openai.com/v1/videos/edits" + assert data == {"prompt": "make it nighttime"} + assert files is not None + field_names = [field for field, _ in files] + assert field_names == ["video"] + _, (filename, content, content_type) = files[0] + assert filename == "clip.mp4" + assert content is source + assert content_type == "video/mp4" + + +def test_video_edit_request_without_file_sends_json_id_reference(): + """The id-reference path must stay JSON (files is None) so existing + remix/edit-by-id callers keep working.""" + config = OpenAIVideoConfig() + + url, data, files = config.transform_video_edit_request( + prompt="brighter", + video_id="video_abc123", + api_base="https://api.openai.com/v1/videos", + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "https://api.openai.com/v1/videos/edits" + assert data == {"prompt": "brighter", "video": {"id": "video_abc123"}} + assert files is None diff --git a/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py new file mode 100644 index 00000000000..3364b1b3872 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -0,0 +1,231 @@ +import pytest + +import litellm +from litellm.interactions.utils import get_provider_interactions_api_config +from litellm.llms.gemini.interactions.transformation import ( + GoogleAIStudioInteractionsConfig, +) +from litellm.llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig, +) +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +GLOBAL_BASE = "https://aiplatform.googleapis.com/v1beta1/projects/test-proj/locations/global/interactions" + + +class MinterRecorder: + def __init__(self, resolved_project: str = "creds-proj") -> None: + self.calls: list[tuple[VERTEX_CREDENTIALS_TYPES | None, str | None]] = [] + self.resolved_project = resolved_project + + def __call__( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + self.calls.append((credentials, project_id)) + return "test-token", project_id or self.resolved_project + + +@pytest.fixture +def minter(): + return MinterRecorder() + + +@pytest.fixture +def config(minter): + return VertexAIInteractionsConfig(mint_access_token=minter) + + +@pytest.fixture +def litellm_params(): + return GenericLiteLLMParams(vertex_project="test-proj", vertex_credentials="creds.json") + + +class TestRegistration: + def test_vertex_ai_returns_vertex_config(self): + assert isinstance(get_provider_interactions_api_config("vertex_ai"), VertexAIInteractionsConfig) + + def test_vertex_ai_beta_returns_vertex_config(self): + assert isinstance(get_provider_interactions_api_config("vertex_ai_beta"), VertexAIInteractionsConfig) + + def test_gemini_still_returns_google_ai_studio_config(self): + gemini_config = get_provider_interactions_api_config("gemini") + assert isinstance(gemini_config, GoogleAIStudioInteractionsConfig) + assert not isinstance(gemini_config, VertexAIInteractionsConfig) + + def test_lazy_import_resolves(self): + assert litellm.VertexAIInteractionsConfig is VertexAIInteractionsConfig + + def test_custom_llm_provider_is_vertex_ai(self, config): + assert config.custom_llm_provider == LlmProviders.VERTEX_AI + + +class TestValidateEnvironment: + def test_sets_bearer_auth_without_gemini_headers(self, config, minter, litellm_params): + headers = config.validate_environment( + headers={}, + model="gemini-omni-flash-preview", + litellm_params=litellm_params, + ) + + assert headers["Authorization"] == "Bearer test-token" + assert headers["Content-Type"] == "application/json" + assert "x-goog-api-key" not in headers + assert "Api-Revision" not in headers + assert minter.calls == [("creds.json", "test-proj")] + + def test_caller_authorization_wins(self, config, litellm_params): + headers = config.validate_environment( + headers={"Authorization": "Bearer caller-token"}, + model="gemini-omni-flash-preview", + litellm_params=litellm_params, + ) + + assert headers["Authorization"] == "Bearer caller-token" + + +class TestGetCompleteUrl: + def test_defaults_to_global_v1beta1(self, config, litellm_params): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + ) + + assert url == GLOBAL_BASE + + def test_stream_appends_alt_sse(self, config, litellm_params): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + stream=True, + ) + + assert url == f"{GLOBAL_BASE}?alt=sse" + + def test_multi_region_location_uses_rep_host(self, config): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "us"}, + ) + + assert url == "https://aiplatform.us.rep.googleapis.com/v1beta1/projects/test-proj/locations/us/interactions" + + def test_regional_location_uses_regional_host(self, config): + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "us-central1"}, + ) + + assert url == ( + "https://us-central1-aiplatform.googleapis.com" + "/v1beta1/projects/test-proj/locations/us-central1/interactions" + ) + + def test_location_env_fallback_is_ignored(self, config, monkeypatch): + monkeypatch.setenv("VERTEXAI_LOCATION", "us-east5") + + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj"}, + ) + + assert url == GLOBAL_BASE + + def test_api_base_override(self, config, litellm_params): + url = config.get_complete_url( + api_base="https://proxy.example.test", + model="gemini-omni-flash-preview", + litellm_params=dict(litellm_params), + ) + + assert url == "https://proxy.example.test/v1beta1/projects/test-proj/locations/global/interactions" + + def test_project_resolved_from_credentials_when_not_passed(self, config, monkeypatch): + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + monkeypatch.setattr(litellm, "vertex_project", None) + + url = config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_credentials": "creds.json"}, + ) + + assert url == "https://aiplatform.googleapis.com/v1beta1/projects/creds-proj/locations/global/interactions" + + def test_invalid_location_rejected(self, config): + with pytest.raises(ValueError, match="Invalid vertex_location"): + config.get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={"vertex_project": "test-proj", "vertex_location": "evil.com#"}, + ) + + def test_missing_project_rejected(self, monkeypatch): + monkeypatch.delenv("VERTEXAI_PROJECT", raising=False) + monkeypatch.setattr(litellm, "vertex_project", None) + + def unresolved_minter( + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + return "test-token", "" + + with pytest.raises(ValueError, match="Vertex AI project is required"): + VertexAIInteractionsConfig(mint_access_token=unresolved_minter).get_complete_url( + api_base=None, + model="gemini-omni-flash-preview", + litellm_params={}, + ) + + +class TestInteractionByIdRequests: + def test_get_url(self, config, litellm_params): + url, request_body = config.transform_get_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123" + assert request_body == {} + + def test_get_url_encodes_interaction_id(self, config, litellm_params): + url, _ = config.transform_get_interaction_request( + interaction_id="id/with space", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/id%2Fwith%20space" + + def test_delete_url(self, config, litellm_params): + url, request_body = config.transform_delete_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123" + assert request_body == {} + + def test_cancel_url(self, config, litellm_params): + url, request_body = config.transform_cancel_interaction_request( + interaction_id="abc123", + api_base="", + litellm_params=litellm_params, + headers={}, + ) + + assert url == f"{GLOBAL_BASE}/abc123:cancel" + assert request_body == {} diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index ba2f20e2337..f19e169dc9e 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -538,7 +538,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos ) result = transform() - assert result.get("thinking") == {"type": "adaptive"} + assert result.get("thinking") == {"type": "adaptive", "display": "summarized"} assert result.get("output_config") == {"effort": "medium"} monkeypatch.setitem( diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 57cd729bc90..862969abbc2 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -479,7 +479,7 @@ class TestVertexAIVideoConfig: }, } - url, data = self.config.transform_video_edit_request( + url, data, files = self.config.transform_video_edit_request( prompt="Make it brighter", video_id=operation_name, api_base=api_base, @@ -488,6 +488,7 @@ class TestVertexAIVideoConfig: prefetched_source_data=prefetched, ) + assert files is None assert url.endswith(":predictLongRunning") assert "veo-3.1-generate-001" in url instance = data["instances"][0] @@ -507,7 +508,7 @@ class TestVertexAIVideoConfig: }, } - _, data = self.config.transform_video_edit_request( + _, data, _ = self.config.transform_video_edit_request( prompt="Make it darker", video_id=operation_name, api_base=api_base, diff --git a/tests/test_litellm/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/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 04f38b5e2ed..06fd3730f7d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -47,6 +47,7 @@ from litellm.proxy.auth.auth_checks import ( _virtual_key_soft_budget_check, get_key_object, get_user_object, + invalidate_team_member_spend_state, vector_store_access_check, ) from litellm.caching.in_memory_cache import InMemoryCache @@ -4868,10 +4869,9 @@ async def test_cache_team_object_writes_team_id_and_invalidates_team_alias(): 3. When team_alias is None, NO alias-key operation happens (no delete of an empty-keyed entry, no spurious write). 4. DELETES the team_id-keyed entry from the internal usage cache - BEFORE the fresh write (LIT-4391). `_get_team_object_from_cache` - consults the internal usage cache first, so a leftover copy there - (backfilled from a Redis shared with `user_api_key_cache`) would - keep serving the pre-update team allowlist. + BEFORE the fresh write (LIT-4391). `_get_team_object_from_cache` no + longer reads the internal usage cache (LIT-5944), but the delete + protects mixed-version rolling deploys where older workers still do. """ from unittest.mock import AsyncMock, MagicMock @@ -4980,8 +4980,10 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): Regression test for LIT-4391: keys with models=["all-team-models"] kept getting 403 team_model_access_denied for models added via /team/update. - `_get_team_object_from_cache` consults the internal usage cache BEFORE - `user_api_key_cache`. When both share one Redis (enable_redis_auth_cache), + `_get_team_object_from_cache` used to consult the internal usage cache + BEFORE `user_api_key_cache` (removed in LIT-5944; this test now also + guards against reintroducing that read). + When both share one Redis (enable_redis_auth_cache), any team read backfills the internal cache's in-memory tier with the team object. `_cache_team_object` (the /team/update refresh) only wrote `user_api_key_cache`, so that backfilled copy kept shadowing the update @@ -5053,6 +5055,75 @@ async def test_team_update_not_shadowed_by_internal_usage_cache_lit_4391(): ) +class _CountingFakeRedis(_SharedFakeRedis): + """Counts per-key Redis round-trips so tests can pin the number of + network operations a code path issues.""" + + def __init__(self): + super().__init__() + self.get_calls: int = 0 + + async def async_get_cache(self, key, **kwargs): + self.get_calls += 1 + return await super().async_get_cache(key, **kwargs) + + +@pytest.mark.asyncio +async def test_warm_team_object_reads_issue_no_redis_ops_lit_5944(): + """ + Regression test for LIT-5944: project/team-scoped virtual-key requests + paid ~4 awaited Redis GETs per request just to re-read the team object. + + `_get_team_object_from_cache` used to consult + `proxy_logging_obj.internal_usage_cache.dual_cache` (in-memory TTL 1s, + Redis-backed) BEFORE `user_api_key_cache`. Nothing writes team objects + into that internal cache — `_cache_team_object` only DELETES the key + there — so when `user_api_key_cache` has no Redis tier the shared Redis + key stays absent forever and every team lookup in the auth hot path + (4 call sites per chat-completion request) became a guaranteed-miss + Redis round-trip, saturating the event loop at high TPS. + + Pins: once `_cache_team_object` has cached a team, repeated + `get_team_object` reads are served from `user_api_key_cache`'s in-memory + tier and issue ZERO Redis operations. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamTableCachedObj + from litellm.proxy.auth.auth_checks import _cache_team_object, get_team_object + + team_id = "team-lit-5944" + counting_redis = _CountingFakeRedis() + user_api_key_cache = UserApiKeyCache() + proxy_logging_obj = MagicMock() + proxy_logging_obj.internal_usage_cache.dual_cache = DualCache( + redis_cache=counting_redis, + default_in_memory_ttl=1, + ) + prisma_client = MagicMock() + + await _cache_team_object( + team_id=team_id, + team_table=LiteLLM_TeamTableCachedObj(team_id=team_id, models=["model-a"]), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + for _ in range(4): + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + assert team_obj is not None and team_obj.models == ["model-a"] + + assert counting_redis.get_calls == 0, ( + "Warm team-object reads must be served from user_api_key_cache's " + "in-memory tier without any Redis round-trips. " + f"Got {counting_redis.get_calls} Redis GETs for 4 get_team_object calls." + ) + + @pytest.mark.asyncio async def test_cache_team_object_tolerates_cache_invalidation_failures(): """ @@ -6939,3 +7010,307 @@ def test_model_has_no_cost_mapping_alias_to_a_group_priced_through_model_info_is router = _router_with_a_group_priced_through_model_info() assert model_has_no_cost_mapping(model="model-info-priced-alias", llm_router=router) is False + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_sets_the_spend_counter_and_clears_both_membership_cache_keys(): + """A team-member budget reset (new_spend passed) must SET the spend counter to the reset + value, clear its DB-floor marker, AND invalidate both independently-keyed membership caches + (user_api_key_auth.py's admission check writes one key format, budget_reservation.py and + auth_checks.py's own get_team_membership() write the other) or a stale read keeps 429ing + after the reset. Asserted against real cache reads, not mock call args, so a change that + keeps the call but drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:user-1:team-1", value="stale-membership") + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0) + real_spend_counter_cache.in_memory_cache.set_cache( + key="spend_db_floor:spend:team_member:user-1:team-1", value=999.0 + ) + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=0.0, + ) + + assert await real_cache.async_get_cache(key="team-1_user-1") is None + assert await real_cache.async_get_cache(key="team_membership:user-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 0.0 + assert ( + real_spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:team_member:user-1:team-1") + == 0.0 + ), "the DB-floor marker kept the pre-reset value; a stale-floor read can raise the counter right back up" + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_leaves_the_live_spend_counter_alone_without_new_spend(): + """team_member_update only changes the budget cap, not the tracked spend, so it calls + invalidate_team_member_spend_state with no new_spend. Deleting the live spend counter in that + case would force the next read to reseed from the DB's own spend column, which lags the live + counter via periodic batch writes, briefly UNDER-enforcing the raised cap against a spend + value lower than what was actually tracked (regression: PR #37971 Bugbot finding).""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_user-1", value="stale-membership") + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:user-1:team-1", value=999.0) + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + ) + + assert await real_cache.async_get_cache(key="team-1_user-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_deleting(): + """/key/{key}/reset_spend SETs its counter to the reset value rather than deleting it, so a + worker's next read reflects it directly instead of falling back through a DB reseed. A reset + caller passing new_spend must match that precedent, not merely delete the counter.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock() + real_spend_counter_cache.redis_cache = fake_redis_cache + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5 + fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60) + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_deletes_redis_counter_when_set_fails(): # test-quality-ok: only observable effect is the fallback call on the same fake client + """Redis reads take priority over the local in-memory copy (get_current_spend reads Redis + first), so a failed Redis SET would otherwise leave the OLD pre-reset value authoritative + for every worker even though the reset reported success. On a failed SET, the stale Redis + entry must be deleted instead, so the next read clean-misses and reseeds from the DB.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down")) + fake_redis_cache.async_delete_cache = AsyncMock() + real_spend_counter_cache.redis_cache = fake_redis_cache + + with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + fake_redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1") + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_raises_503_when_both_redis_writes_fail(): + """If the Redis SET fails AND the fallback DELETE fails, the stale pre-reset counter is still + authoritative in Redis for every worker. Reporting success would silently keep 429ing the + member, so the reset must surface a 503 instead (regression: PR #37971 Greptile finding).""" + from fastapi import HTTPException + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + real_cache = UserApiKeyCache() + real_spend_counter_cache = DualCache() + fake_redis_cache = MagicMock() + fake_redis_cache.async_set_cache = AsyncMock(side_effect=ConnectionError("redis down")) + fake_redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis still down")) + real_spend_counter_cache.redis_cache = fake_redis_cache + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache + ), + pytest.raises(HTTPException) as exc_info, + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=real_cache, + new_spend=2.5, + ) + + assert exc_info.value.status_code == 503 + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_to_remote_workers(): + """The test above only proves the handling worker's own spend counter is + cleared. A remote worker's spend counter is a separate DualCache instance; + if the reset never reaches it, that worker keeps enforcing the pre-reset + spend the moment its own Redis read for the counter fails and it falls + back to its own (now-stale) in-memory copy. Drives the actual message + published onto the invalidation channel through a second, independent + AuthCacheInvalidationSubscriber standing in for that remote worker, rather + than asserting on the publish call args.""" + from redis.asyncio import Redis + + from litellm.caching.dual_cache import DualCache + from litellm.caching.in_memory_cache import InMemoryCache + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + published: list[tuple[str, str]] = [] + + class _RecordingRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + published.append((channel, message)) + return 1 + + class _FakeRedisCache: + def __init__(self) -> None: + self.namespace = None + + def init_async_client(self) -> object: + return _RecordingRedisClient() + + local_spend_counter_cache = DualCache() + + remote_user_api_key_cache = UserApiKeyCache() + remote_spend_counter_in_memory_cache = InMemoryCache() + remote_spend_counter_in_memory_cache.set_cache("spend:team_member:user-1:team-1", 999.0) + remote_spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:team_member:user-1:team-1", 999.0) + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache + ), + patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(), + ), + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=UserApiKeyCache(), + new_spend=0.0, + ) + + def _published_message_for(cache_key: str) -> str: + matches = [message for _, message in published if json.loads(message)["cache_key"] == cache_key] + assert matches, f"{cache_key} never reached the cross-worker invalidation channel" + return matches[-1] + + remote_subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(), + user_api_key_cache=remote_user_api_key_cache, + additional_in_memory_caches=(remote_spend_counter_in_memory_cache,), + ) + for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"): + remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API + {"type": "message", "data": _published_message_for(cache_key)} + ) + + assert remote_spend_counter_in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0 + assert ( + remote_spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 + ), "the DB-floor marker was not broadcast; a remote worker can re-raise the counter off its stale floor" + + +@pytest.mark.asyncio +async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_not_erase_the_reset(): + """The handling worker subscribes to the same invalidation channel it publishes on, so it + receives its own reset message. A delete-style broadcast would erase the post-reset counter + and floor marker the handler just wrote, reopening the stale-floor race the reset closed + (regression: PR #37971 Greptile finding). The broadcast carries the reset value as a SET, so + applying the self-delivered message must leave both keys at the post-reset value.""" + from redis.asyncio import Redis + + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import AuthCacheInvalidationSubscriber + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + published: list[tuple[str, str]] = [] + + class _RecordingRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + published.append((channel, message)) + return 1 + + class _FakeRedisCache: + def __init__(self) -> None: + self.namespace = None + + def init_async_client(self) -> object: + return _RecordingRedisClient() + + local_spend_counter_cache = DualCache() + local_user_api_key_cache = UserApiKeyCache() + + with ( + patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + "litellm.proxy.proxy_server.spend_counter_cache", local_spend_counter_cache + ), + patch( # test-quality-ok: injects a fake pub/sub-capable redis cache; no live redis in this unit test + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(), + ), + ): + await invalidate_team_member_spend_state( + user_id="user-1", + team_id="team-1", + user_api_key_cache=local_user_api_key_cache, + new_spend=0.0, + ) + + own_subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(), + user_api_key_cache=local_user_api_key_cache, + additional_in_memory_caches=(local_spend_counter_cache.in_memory_cache,), + ) + for _, message in published: + own_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API + {"type": "message", "data": message} + ) + + assert local_spend_counter_cache.in_memory_cache.get_cache("spend:team_member:user-1:team-1") == 0.0, ( + "the handler's self-delivered broadcast erased the post-reset spend counter" + ) + assert ( + local_spend_counter_cache.in_memory_cache.get_cache("spend_db_floor:spend:team_member:user-1:team-1") == 0.0 + ), "the handler's self-delivered broadcast erased the post-reset floor marker, reopening the stale-floor race" diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index b0094b81112..90b3b29d919 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -701,3 +701,58 @@ async def test_auth_failure_ip_stamp_does_not_mutate_callers_request_data(): ) assert request_data == {"model": "gpt-4o"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "auth_error,expect_traceback", + [ + pytest.param( + ProxyException( + message="Authentication Error", type=ProxyErrorTypes.auth_error, param=None, code=401 + ), + False, + id="expected_401_no_traceback", + ), + pytest.param(ValueError("unexpected internal error"), True, id="unexpected_error_keeps_traceback"), + ], +) +async def test_handle_authentication_error_traceback_only_for_unexpected_errors(auth_error, expect_traceback, caplog): + """Regression for LIT-6043: expected 4xx auth rejections must not format a + traceback via logger.exception; unexpected errors must keep it.""" + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.auth.auth_exception_handler.seed_request_identity", + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + verbose_proxy_logger.propagate = True + try: + try: + raise auth_error + except (ProxyException, ValueError) as caught: + with caplog.at_level("ERROR", logger="LiteLLM Proxy"), pytest.raises(ProxyException): + await handler._handle_authentication_error( + caught, + MagicMock(), + {}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + finally: + verbose_proxy_logger.propagate = False + + records = [r for r in caplog.records if "user_api_key_auth(): Exception occured" in r.getMessage()] + assert len(records) == 1 + assert (records[0].exc_info is not None) is expect_traceback diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 468e8aabae8..7d5fc1a3544 100644 --- a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest from redis.asyncio import Redis +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AUTH_CACHE_INVALIDATION_CHANNEL, AuthCacheInvalidationSubscriber, @@ -144,6 +145,37 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None: assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] +@pytest.mark.asyncio +async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() -> None: + """ + The spend-counter half of the same cross-worker gap: a remote worker's own + spend counter can hold a stale value (its fallback path when that worker's + own Redis read for the counter fails), and only clearing user_api_key_cache + on message would leave that separate DualCache's in-memory copy untouched. + """ + cache = UserApiKeyCache() + spend_counter_in_memory_cache = InMemoryCache() + spend_counter_in_memory_cache.set_cache("spend:team_member:u-1:t-1", 999.0) + assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message("spend:team_member:u-1:t-1")]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + additional_in_memory_caches=(spend_counter_in_memory_cache,), + ) + subscriber.start() + try: + for _ in range(200): + if spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None + + @pytest.mark.asyncio async def test_subscriber_ignores_malformed_messages() -> None: cache = UserApiKeyCache() diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index f6d74a189bc..7f5d3eb0a14 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9,11 +9,13 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +from pydantic import ValidationError from litellm._uuid import uuid from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( + LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, LiteLLM_ModelTable, LiteLLM_OrganizationMembershipTable, @@ -27,7 +29,9 @@ from litellm.proxy._types import ( Member, ProxyErrorTypes, ProxyException, + ResetSpendRequest, TeamMemberAddRequest, + TeamMemberUpdateRequest, UpdateTeamRequest, ) from litellm.proxy.management_endpoints.team_endpoints import ( @@ -42,12 +46,15 @@ from litellm.proxy.management_endpoints.team_endpoints import ( _transform_teams_to_deleted_records, _update_model_table, _validate_and_populate_member_user_info, + _validate_team_member_reset_spend_value, _verify_team_access, delete_team, list_available_teams, + reset_team_member_spend_fn, router, team_member_add_duplication_check, team_member_delete, + team_member_update, update_team, validate_team_org_change, ) @@ -12603,3 +12610,376 @@ async def test_invalidate_access_group_cache_deletes_the_cached_object(): "user_api_key_cache": cache, "proxy_logging_obj": logging_obj, } + + +def test_validate_team_member_reset_spend_value_rejects_non_numeric(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to="not-a-number", + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_rejects_negative(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=-1.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("reset_to", [float("nan"), float("inf"), float("-inf")]) +def test_validate_team_member_reset_spend_value_rejects_non_finite(reset_to): + """NaN and +/-inf are instances of float and compare False against every bound + below (`nan < 0`, `nan > current_spend` are both False), so an isinstance-and-range + check alone lets them through to persist as the member's spend and silently + disable every later budget comparison against it.""" + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=reset_to, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("reset_to", [True, False]) +def test_reset_spend_request_rejects_bool_reset_to(reset_to): + """bool is a subclass of int, so pydantic silently coerces True/False into 1.0/0.0 for a + ``float`` field: {"reset_to": true} would otherwise reach _validate_team_member_reset_spend_value + as an indistinguishable 1.0 and reset the member's spend instead of failing the request.""" + with pytest.raises(ValidationError): + ResetSpendRequest(reset_to=reset_to) + + +def test_validate_team_member_reset_spend_value_rejects_above_current_spend(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=20.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_rejects_above_max_budget(): + with pytest.raises(HTTPException) as exc: + _validate_team_member_reset_spend_value( + reset_to=10.0, + membership=LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + spend=10.0, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=5.0), + ), + ) + assert exc.value.status_code == 400 + + +def test_validate_team_member_reset_spend_value_accepts_valid_reset(): + result = _validate_team_member_reset_spend_value( + reset_to=0.0, + membership=LiteLLM_TeamMembership(user_id="u1", team_id="t1", spend=10.0), + ) + assert result == 0.0 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_success(monkeypatch): + """A proxy admin resetting a stuck team member's spend must write the DB + row to reset_to AND invalidate the cached spend/membership state, or the + 429 the endpoint exists to clear keeps firing off the stale cache. + Asserted against real cache reads, not mock call args, so a change that + keeps the call but drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + mock_proxy_logging_obj = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + membership_row = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + spend=10.0, + litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0), + ) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + response = await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert response["spend"] == 0.0 + assert response["previous_spend"] == 10.0 + assert response["max_budget"] == 50.0 + mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with( + where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}}, + data={"spend": 0.0}, + ) + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch): + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="ghost-user", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_team_not_found(monkeypatch): + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist in db."})), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="ghost-team", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_forbidden_for_non_admin(monkeypatch): + """A caller who is neither proxy admin, org admin, nor this team's admin must be refused, + matching every other team-mutating endpoint's authorization.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="member-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user" + ), + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_team_admin_cannot_reset_own_spend(monkeypatch): + """_verify_team_access authorizes a team admin over their own team with no check that the + target differs from the caller. Unchecked, that admin could target their own membership row + and repeatedly zero it right before it crosses their per-member cap, consuming the shared + team budget without the configured limit ever binding (Veria finding on PR #37971).""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-admin", user_id="team-admin-1") + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock( + return_value=LiteLLM_TeamTable( + team_id="team-1", + members_with_roles=[Member(user_id="team-admin-1", role="admin")], + ) + ), + ): + with pytest.raises(HTTPException) as exc: + await reset_team_member_spend_fn( + team_id="team-1", + user_id="team-admin-1", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=team_admin, + ) + assert exc.value.status_code == 403 + mock_prisma_client.db.litellm_teammembership.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkeypatch): + """The self-reset guard is scoped to non-proxy-admin roles: a proxy admin resetting their + own membership spend is the platform-wide trust boundary, not a team-scoped one.""" + mock_prisma_client = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + + membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0) + mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row) + + with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")), + ): + response = await reset_team_member_spend_fn( + team_id="team-1", + user_id="admin-user", + data=ResetSpendRequest(reset_to=0.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + assert response["spend"] == 0.0 + + +@pytest.mark.asyncio +async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch): + """Raising a stuck member's max_budget_in_team via the documented /team/member_update + endpoint must invalidate the cached membership state, or the raised cap never reaches the + admission check and the member stays 429ing. The live spend counter itself must be left + untouched: only the cap changed, and deleting the counter would force a reseed from the + DB's own spend column, which lags the live counter via periodic batch writes, briefly + UNDER-enforcing the raised cap against a spend value lower than what was actually tracked. + Asserted against real cache reads, not mock call args, so a change that keeps the call but + drops its effect still fails.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership") + await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0) + + team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[]) + team_info_response = { + "team_info": team_row, + "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_response), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1", max_budget_in_team=999999.0), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert await real_cache.async_get_cache(key="team-1_member-1") is None + assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0 + + +@pytest.mark.asyncio +async def test_team_member_update_skips_invalidation_when_no_budget_fields_sent(monkeypatch): + """A role-only update carries an empty budget_patch and touches no budget state, + so the member's cached spend/membership state must be left untouched.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + mock_prisma_client = MagicMock() + real_cache = UserApiKeyCache() + await real_cache.async_set_cache(key="team-1_member-1", value="still-fresh-membership") + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=1.5) + + team_row = LiteLLM_TeamTable(team_id="team-1", metadata={}, members_with_roles=[]) + team_info_response = { + "team_info": team_row, + "team_memberships": [LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id=None)], + } + + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache) + + mock_tx = AsyncMock() + mock_prisma_client.tx.return_value.__aenter__ = AsyncMock(return_value=mock_tx) + mock_prisma_client.tx.return_value.__aexit__ = AsyncMock(return_value=None) + + with ( + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints.team_info", + AsyncMock(return_value=team_info_response), + ), + patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests + "litellm.proxy.management_endpoints.team_endpoints._upsert_budget_and_membership", + AsyncMock(), + ), + ): + await team_member_update( + data=TeamMemberUpdateRequest(team_id="team-1", user_id="member-1"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + ) + + assert await real_cache.async_get_cache(key="team-1_member-1") == "still-fresh-membership" + assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 1.5 diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 9177944df2d..791d64c6428 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1353,6 +1353,8 @@ class TestParseCursorModelVariant: ("claude-opus-5-fast", "claude-opus-5", None), ("gpt-5.6-sol", "gpt-5.6-sol", None), ("foo-thinking-ultra-fast", "foo-thinking-ultra", None), + ("gpt-5.6-thinking-max", "gpt-5.6", "max"), + ("foo-thinking-mega-fast", "foo-thinking-mega", None), ("-thinking-high", "-thinking-high", None), ], ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 6c8e641642b..843e1d1296f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -3,11 +3,10 @@ import datetime import json from datetime import timezone from typing import Any, Final, cast +from unittest.mock import AsyncMock, MagicMock, patch import pytest - - -from unittest.mock import AsyncMock, MagicMock, patch +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import ( @@ -3500,6 +3499,50 @@ def test_get_logging_payload_failed_request_without_standard_logging_payload_lea assert payload["custom_llm_provider"] == "" +class _ModelRouterSpendLogKwargs(TypedDict): + model: ReadOnly[str] + litellm_params: ReadOnly[dict[str, dict[str, str]]] + standard_logging_object: ReadOnly[StandardLoggingPayload] + + +def _model_router_spend_log_kwargs(slp_model: str | None) -> _ModelRouterSpendLogKwargs: + standard_logging_payload: Final = cast( + StandardLoggingPayload, + { + "model": slp_model, + "metadata": {}, + "model_map_information": StandardLoggingModelInformation( + model_map_key="azure_ai/model_router", model_map_value=None + ), + }, + ) + return { + "model": "azure_ai/model_router/model-router", + "litellm_params": {"metadata": {"user_api_key": "sk-test-key"}}, + "standard_logging_object": standard_logging_payload, + } + + +def test_get_logging_payload_uses_standard_logging_payload_model(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model="azure_ai/gpt-5-mini"), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/gpt-5-mini" + + +def test_get_logging_payload_falls_back_to_kwargs_model_when_slp_model_missing(): + payload = get_logging_payload( + kwargs=_model_router_spend_log_kwargs(slp_model=None), + response_obj={}, + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + assert payload["model"] == "azure_ai/model_router/model-router" + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_empty_key_slp_none_is_empty_string_not_none_literal(): diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 58714a5e319..b595c44d2ce 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2294,6 +2294,54 @@ class TestOverrideOpenAIResponseModel: assert response_obj.model == actual_model_used assert response_obj.model != requested_model + def test_override_model_preserves_model_router_model_for_alias_without_router_in_name( + self, + ): + """ + The client sends a model group alias, which carries no model_router/ prefix, so the + name check alone only fires when the operator happened to put "model-router" in the + alias. With the stamp on the response the actual model survives whatever it is named. + """ + from litellm.llms.azure_ai.common_utils import ( + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY, + ) + + requested_model = "smart-pick" + actual_model_used = "azure_ai/grok-4-1-fast-reasoning" + + response_obj = MagicMock() + response_obj.model = actual_model_used + response_obj._hidden_params = { + "additional_headers": {}, + AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: actual_model_used, + } + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + assert response_obj.model == actual_model_used + + def test_override_model_still_restamps_non_router_alias_without_stamp(self): + """ + Control for the test above: absent the stamp, an ordinary deployment keeps being + restamped to the requested model, so the stamp is doing the work rather than the + preserve branch having gone unconditional. + """ + requested_model = "smart-pick" + + response_obj = MagicMock() + response_obj.model = "azure_ai/grok-4-1-fast-reasoning" + response_obj._hidden_params = {"additional_headers": {}} + + _override_openai_response_model( + response_obj=response_obj, + requested_model=requested_model, + log_context="test_context", + ) + assert response_obj.model == requested_model + def test_override_model_uses_winning_model_for_fastest_response(self): """ Test that when fastest_response batch completion is used with a @@ -7295,3 +7343,31 @@ class TestRouterModelNameOnNonStreamingResponse: ) assert response[ROUTER_MODEL_NAME_RESPONSE_FIELD] == "deep-model" + + +@pytest.mark.parametrize( + "exc,expect_traceback", + [ + pytest.param(HTTPException(status_code=400, detail="Invalid model name passed in"), False, id="expected_400"), + pytest.param(ValueError("unexpected internal error"), True, id="unexpected_error"), + ], +) +def test_log_llm_api_exception_traceback_only_for_unexpected_errors(exc, expect_traceback, caplog): + """Regression for LIT-6043: expected 4xx errors log without formatting a + traceback; unexpected errors keep logger.exception behavior.""" + from litellm._logging import verbose_proxy_logger + from litellm.proxy.common_request_processing import _log_llm_api_exception + + verbose_proxy_logger.propagate = True + try: + with caplog.at_level("ERROR", logger="LiteLLM Proxy"): + try: + raise exc + except Exception as raised: + _log_llm_api_exception(raised) + finally: + verbose_proxy_logger.propagate = False + + records = [r for r in caplog.records if "_handle_llm_api_exception(): Exception occured" in r.getMessage()] + assert len(records) == 1 + assert (records[0].exc_info is not None) is expect_traceback diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 31d2a6cef98..3383527e932 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11347,3 +11347,38 @@ class TestRouterModelNameOnStreamingChunks: assert len(frames) >= 3 assert '"router_model_name":"deep-model"' in frames[0] assert all('"router_model_name":"backup-tier"' in frame for frame in frames[1:]) + + +@pytest.mark.asyncio +async def test_authoritative_floor_spend_keeps_a_reset_marker_written_during_the_db_read(): + """A team-member spend reset writes the post-reset floor to the spend_db_floor marker + (auth_checks.invalidate_team_member_spend_state). A floor read already in flight when the + reset commits would otherwise cache its stale pre-reset DB value over the fresh marker, + letting a budget check raise the counter right back above the just-reset spend + (regression: PR #37971 Greptile finding).""" + from litellm.proxy.proxy_server import _authoritative_floor_spend + + real_spend_counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + marker_key = f"spend_db_floor:{counter_key}" + + async def db_read_racing_with_a_reset(prisma_client, counter_key): + real_spend_counter_cache.in_memory_cache.set_cache(key=marker_key, value=0.0) + return 999.0 + + with ( + patch.object( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock + proxy_server_module, "spend_counter_cache", real_spend_counter_cache + ), + patch.object( # test-quality-ok: the DB read must race the reset; no injectable seam for module-global prisma reads + proxy_server_module.SpendCounterReseed, + "from_db", + AsyncMock(side_effect=db_read_racing_with_a_reset), + ), + ): + result = await _authoritative_floor_spend(counter_key=counter_key) + + assert result == 0.0 + assert real_spend_counter_cache.in_memory_cache.get_cache(key=marker_key) == 0.0, ( + "the in-flight DB read clobbered the post-reset floor marker with the stale pre-reset value" + ) diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index fb01216982f..6920cc0dae3 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -11,7 +11,6 @@ from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks - from unittest.mock import MagicMock, patch from litellm.proxy.utils import get_custom_url, join_paths @@ -82,9 +81,7 @@ async def test_proxy_only_error_log_marks_no_upstream_llm_call(): captured = {} def fake_pre_call(self, *args, **kwargs): - captured["flag"] = self.model_call_details.get( - LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL - ) + captured["flag"] = self.model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) from litellm.litellm_core_utils.litellm_logging import Logging @@ -102,9 +99,7 @@ async def test_proxy_only_error_log_marks_no_upstream_llm_call(): "model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-bad", request_route="/v1/chat/completions" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-bad", request_route="/v1/chat/completions"), route="/v1/chat/completions", original_exception=Exception("bad key"), ) @@ -148,13 +143,9 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): request_data={ "model": "gpt-4o", "input": "blocked prompt", - "litellm_metadata": { - "standard_logging_guardrail_information": guardrail_info - }, + "litellm_metadata": {"standard_logging_guardrail_information": guardrail_info}, }, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-1234", request_route="/v1/responses" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/responses"), route="/v1/responses", original_exception=HTTPException(status_code=400, detail="blocked"), ) @@ -163,12 +154,7 @@ async def test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params(): Logging.pre_call = orig_pre_call Logging.async_failure_handler = orig_async_failure - assert ( - captured["litellm_params"]["litellm_metadata"][ - "standard_logging_guardrail_information" - ] - == guardrail_info - ) + assert captured["litellm_params"]["litellm_metadata"]["standard_logging_guardrail_information"] == guardrail_info assert "litellm_metadata" not in captured["optional_params"] @@ -206,9 +192,7 @@ def test_get_model_group_info_order(): def test_join_paths_no_duplication(): """Test that join_paths doesn't duplicate route when base_path already ends with it""" - result = join_paths( - base_path="http://0.0.0.0:4000/my-custom-path/", route="/my-custom-path" - ) + result = join_paths(base_path="http://0.0.0.0:4000/my-custom-path/", route="/my-custom-path") assert result == "http://0.0.0.0:4000/my-custom-path" @@ -814,9 +798,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=system_prompt - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=system_prompt) assert estimated.prompt_tokens == expected @pytest.mark.asyncio @@ -834,7 +818,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter( model="gpt-3.5-turbo", text="part one of the system prompt. part two of the system prompt." ) assert estimated.prompt_tokens == expected @@ -870,9 +856,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=system_prompt - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=system_prompt) assert estimated.prompt_tokens == expected @pytest.mark.asyncio @@ -890,9 +876,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: estimated = request_data["combined_usage_object"] assert isinstance(estimated, Usage) - expected = litellm_module.token_counter(model="gpt-3.5-turbo", messages=messages) + litellm_module.token_counter( - model="gpt-3.5-turbo", text=dispatched_system - ) + expected = litellm_module.token_counter( + model="gpt-3.5-turbo", messages=messages + ) + litellm_module.token_counter(model="gpt-3.5-turbo", text=dispatched_system) assert estimated.prompt_tokens == expected @@ -916,9 +902,7 @@ def test_create_model_info_response_includes_max_tokens_from_lookup(): model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert response["id"] == "some-model" @@ -935,9 +919,7 @@ def test_create_model_info_response_does_not_call_router_group_info(): model_id="some-model", provider="openai", llm_router=router, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) router.get_model_group_info.assert_not_called() @@ -968,9 +950,7 @@ def test_create_model_info_response_deployment_limits_override_cost_map(): model_id="gpt-4o", provider="openai", llm_router=router, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert response["max_input_tokens"] == 200000 @@ -1008,9 +988,7 @@ def test_create_model_info_response_survives_malformed_cost_map_limits(bad_value model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=bad_value, max_output_tokens=bad_value - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=bad_value, max_output_tokens=bad_value), ) assert response["id"] == "some-model" @@ -1023,9 +1001,7 @@ def test_create_model_info_response_keeps_valid_cost_map_limit_beside_malformed_ model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens="128,000", max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens="128,000", max_output_tokens=16384), ) assert "max_input_tokens" not in response @@ -1068,9 +1044,7 @@ def test_create_model_info_response_emits_integer_token_counts(): model_id="some-model", provider="openai", llm_router=None, - get_model_info=lambda _model: _fake_model_info( - max_input_tokens=128000, max_output_tokens=16384 - ), + get_model_info=lambda _model: _fake_model_info(max_input_tokens=128000, max_output_tokens=16384), ) assert isinstance(response["max_input_tokens"], int) @@ -1119,9 +1093,7 @@ def test_create_model_info_response_no_router_keeps_base_fields(): def test_create_model_info_response_reads_real_cost_map(): - response = create_model_info_response( - model_id="gpt-4o", provider="openai", llm_router=None - ) + response = create_model_info_response(model_id="gpt-4o", provider="openai", llm_router=None) assert isinstance(response["max_input_tokens"], int) assert response["max_input_tokens"] > 0 @@ -1205,10 +1177,7 @@ class TestPostCallFailureHookLLMExceptionAlerting: @pytest.mark.asyncio async def test_http_exception_does_not_alert(self): - assert ( - await self._alerted(HTTPException(status_code=400, detail="blocked")) - is False - ) + assert await self._alerted(HTTPException(status_code=400, detail="blocked")) is False @pytest.mark.asyncio async def test_genuine_llm_api_error_still_alerts(self): @@ -1241,9 +1210,7 @@ class TestPostCallFailureHookProxyExceptionLogging: await proxy_logging_obj.post_call_failure_hook( request_data={}, original_exception=exc, - user_api_key_dict=UserAPIKeyAuth( - api_key="sk-test", request_route=request_route - ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", request_route=request_route), ) return handle_mock.await_count > 0 @@ -1260,20 +1227,12 @@ class TestPostCallFailureHookProxyExceptionLogging: @pytest.mark.asyncio async def test_proxy_exception_on_llm_route_is_logged(self): - assert ( - await self._logged(self._block(), request_route="/v1/chat/completions") - is True - ) + assert await self._logged(self._block(), request_route="/v1/chat/completions") is True @pytest.mark.asyncio async def test_generic_exception_on_llm_route_is_not_logged(self): # A raw provider/unknown exception is logged by the LLM call path, not here. - assert ( - await self._logged( - Exception("upstream 503"), request_route="/v1/chat/completions" - ) - is False - ) + assert await self._logged(Exception("upstream 503"), request_route="/v1/chat/completions") is False class TestShouldUseSmtpSsl: @@ -1307,9 +1266,7 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection( - smtp_host="mail.example.com", smtp_port=465 - ) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=465) mock_smtp.assert_not_called() assert result is mock_smtp_ssl.return_value @@ -1329,9 +1286,7 @@ class TestCreateSmtpConnection: patch("smtplib.SMTP_SSL") as mock_smtp_ssl, patch("smtplib.SMTP") as mock_smtp, ): - result = _create_smtp_connection( - smtp_host="mail.example.com", smtp_port=587 - ) + result = _create_smtp_connection(smtp_host="mail.example.com", smtp_port=587) mock_smtp_ssl.assert_not_called() assert result is mock_smtp.return_value @@ -1352,9 +1307,7 @@ class TestSendEmailStartTls: monkeypatch.delenv("SMTP_USE_SSL", raising=False) mock_server = MagicMock(spec=smtplib.SMTP) - with patch( - "litellm.proxy.utils._create_smtp_connection" - ) as mock_create_connection: + with patch("litellm.proxy.utils._create_smtp_connection") as mock_create_connection: mock_create_connection.return_value.__enter__.return_value = mock_server await send_email( receiver_email="receiver@example.com", @@ -1690,9 +1643,7 @@ def test_a_failed_dispatch_is_estimated_as_input_only(): usage = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage is not None - assert usage.prompt_tokens == _count_request_input_tokens( - FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None - ) + assert usage.prompt_tokens == _count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage.completion_tokens == 0 assert usage.total_tokens == usage.prompt_tokens @@ -1759,9 +1710,7 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost(): def test_a_failure_that_cost_the_provider_nothing_lifts_nothing(model_call_details, dispatched): from litellm.proxy.utils import _failure_usage_to_lift - assert _failure_usage_to_lift( - model_call_details=model_call_details, request_body={}, dispatched=dispatched - ) is None + assert _failure_usage_to_lift(model_call_details=model_call_details, request_body={}, dispatched=dispatched) is None def test_the_no_upstream_call_key_the_module_uses_is_the_one_asserted_above(): @@ -1829,3 +1778,102 @@ def test_a_dispatched_failure_lifts_the_four_fields_the_spend_log_needs(): assert lifted["response_cost"] == 0.0 assert lifted["combined_usage_object"].prompt_tokens > 0 assert lifted["standard_logging_object"] == {"id": "log-1"} + + +@pytest.mark.asyncio +async def test_proxy_only_error_expected_4xx_skips_traceback_for_both_handlers(monkeypatch): + """Regression for LIT-6043: an expected 4xx must not format a traceback for + either the async or the threaded sync failure handler.""" + import asyncio + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr(litellm, "failure_callback", []) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + captured = {} + sync_ran = asyncio.Event() + loop = asyncio.get_running_loop() + + async def fake_async_failure(self, exception, traceback_exception, *args, **kwargs): + captured["async_traceback"] = traceback_exception + + def fake_sync_failure(self, exception, traceback_exception, *args, **kwargs): + captured["sync_traceback"] = traceback_exception + loop.call_soon_threadsafe(sync_ran.set) + + orig_async_failure = Logging.async_failure_handler + orig_sync_failure = Logging.failure_handler + Logging.async_failure_handler = fake_async_failure + Logging.failure_handler = fake_sync_failure + try: + try: + raise HTTPException(status_code=400, detail="Invalid model name passed in") + except HTTPException as exc: + await proxy_logging_obj._handle_logging_proxy_only_error( + request_data={ + "model": "does-not-exist", + "messages": [{"role": "user", "content": "hi"}], + }, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=exc, + ) + await asyncio.wait_for(sync_ran.wait(), timeout=5) + finally: + Logging.async_failure_handler = orig_async_failure + Logging.failure_handler = orig_sync_failure + + assert captured["async_traceback"] == "" + assert captured["sync_traceback"] == "" + + +@pytest.mark.asyncio +async def test_proxy_only_error_5xx_keeps_traceback_and_runs_sync_callbacks(monkeypatch): + """Unexpected (5xx) errors keep the full traceback, and a configured + sync-only failure callback still gets its threaded handler.""" + import asyncio + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy._types import UserAPIKeyAuth + + def _custom_sync_callback(kwargs, completion_response, start_time, end_time): + pass + + monkeypatch.setattr(litellm, "failure_callback", [_custom_sync_callback]) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + captured = {} + sync_ran = asyncio.Event() + loop = asyncio.get_running_loop() + + async def fake_async_failure(self, exception, traceback_exception, *args, **kwargs): + captured["async_traceback"] = traceback_exception + + def fake_sync_failure(self, *args, **kwargs): + loop.call_soon_threadsafe(sync_ran.set) + + orig_async_failure = Logging.async_failure_handler + orig_sync_failure = Logging.failure_handler + Logging.async_failure_handler = fake_async_failure + Logging.failure_handler = fake_sync_failure + try: + try: + raise HTTPException(status_code=500, detail="internal error") + except HTTPException as exc: + await proxy_logging_obj._handle_logging_proxy_only_error( + request_data={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + }, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/chat/completions"), + route="/v1/chat/completions", + original_exception=exc, + ) + await asyncio.wait_for(sync_ran.wait(), timeout=5) + finally: + Logging.async_failure_handler = orig_async_failure + Logging.failure_handler = orig_sync_failure + + assert "test_proxy_utils" in captured["async_traceback"] diff --git a/tests/test_litellm/proxy/video_endpoints/test_endpoints.py b/tests/test_litellm/proxy/video_endpoints/test_endpoints.py index 0dd163b7b94..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", @@ -450,6 +451,31 @@ async def test_edit__json_string_video_reference_from_form_field(harness): 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 # # =========================================================================== # diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 85777afe81c..587be59c550 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -1,6 +1,10 @@ import logging from unittest.mock import MagicMock, patch +import httpx +import pytest +import respx + import litellm @@ -62,3 +66,66 @@ def test_rerank_does_not_log_request_content_at_info(caplog): assert all( r.levelno == logging.DEBUG for r in optional_params_logs ), "optional_rerank_params must be logged at DEBUG, not INFO" + + +TOGETHER_RERANK_BODY = { + "id": "rerank-mock-id", + "results": [{"index": 0, "relevance_score": 0.95}], + "usage": {"prompt_tokens": 10, "total_tokens": 10}, +} + + +def test_together_rerank_defaults_to_together_ai_host(respx_mock: respx.MockRouter, monkeypatch): + """Regression for the Together host migration: rerank used to hardcode + https://api.together.xyz/v1/rerank. The default must now be api.together.ai.""" + monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False) + + mock_route = respx_mock.post("https://api.together.ai/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + response = litellm.rerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + ) + + assert mock_route.called + assert response.results[0]["relevance_score"] == 0.95 + + +def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter): + """Regression: a custom api_base was silently ignored by the Together rerank handler.""" + mock_route = respx_mock.post("https://custom-together.example/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + litellm.rerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + api_base="https://custom-together.example/v1", + ) + + assert mock_route.called + assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key" + + +@pytest.mark.asyncio +async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch): + """Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank.""" + monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://env-together.example/v1") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + + mock_route = respx_mock.post("https://env-together.example/v1/rerank") + mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY) + + response = await litellm.arerank( + model="together_ai/mixedbread-ai/mxbai-rerank-large-v2", + query=MARKER_QUERY, + documents=[MARKER_DOC], + api_key="fake-together-key", + ) + + assert mock_route.called + assert response.results[0]["relevance_score"] == 0.95 diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index d4333fde163..a29b4d03bc5 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -6052,8 +6052,9 @@ class TestContextAwareClassifier: [{"role": "user", "content": turn}, {"role": "user", "content": "go ahead"}], "go ahead", 3, - 200, - False, + budget_chars=10_000, + per_turn_chars=200, + include_assistant=False, ) assert "multi-region gateway" in quoted[0][1] @@ -6216,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. @@ -6235,7 +6396,7 @@ class TestContextAwareClassifier: {"role": "user", "content": "current ask"}, ] - assert _extract_prior_turns(messages, "current ask", 5, 200, False, pairs) == ( + assert _extract_prior_turns(messages, "current ask", 5, 10_000, 200, False, pairs) == ( ("user", "what about b-trees?"), ("user", "and heaps?"), ) diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py new file mode 100644 index 00000000000..6cd5b956b05 --- /dev/null +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -0,0 +1,198 @@ +import pytest + +from litellm.router_utils.reasoning_effort_capability import ( + deployment_is_catalog_mapped, + intersect_supported_reasoning_efforts, + resolve_supported_reasoning_efforts, +) + + +class TestDeploymentIsCatalogMapped: + def test_a_mode_the_catalog_supplied_marks_the_deployment_mapped(self): + assert deployment_is_catalog_mapped({"mode": "chat"}, {}) is True + + def test_a_deployment_the_catalog_never_described_is_not_mapped(self): + assert deployment_is_catalog_mapped(None, {}) is False + assert deployment_is_catalog_mapped({"max_input_tokens": 200000}, {}) is False + + def test_a_mode_the_operator_wrote_does_not_make_the_deployment_mapped(self): + # Every deployment is registered in the cost map under its own id, so an operator-written + # mode reads back identically to one the catalog supplied and would otherwise let an + # off-map deployment empty the levels its mapped siblings agree on. + assert deployment_is_catalog_mapped({"mode": "chat"}, {"mode": "chat", "id": "abc"}) is False + + +class TestProvenanceSeparatesUnknownFromNonReasoning: + def test_an_off_map_deployment_resolves_to_unknown(self): + # get_model_info answers supports_reasoning None both for a deployment the map never + # described and for a mapped non-reasoning model, so reading an unset flag as () would let + # one custom deployment empty every level its mapped siblings agree on. + assert resolve_supported_reasoning_efforts({}, deployment_is_mapped=False) is None + assert resolve_supported_reasoning_efforts({"supports_reasoning": None}, deployment_is_mapped=False) is None + + def test_a_mapped_deployment_the_map_calls_non_reasoning_supports_no_efforts(self): + assert resolve_supported_reasoning_efforts({}, deployment_is_mapped=True) == () + assert resolve_supported_reasoning_efforts({"supports_reasoning": None}, deployment_is_mapped=True) == () + + def test_an_explicit_false_supports_no_efforts_off_the_map_too(self): + # The operator's own escape hatch: saying so on an off-map deployment must still empty the + # group, since nothing else can tell the resolver that model takes no effort level. + assert resolve_supported_reasoning_efforts({"supports_reasoning": False}, deployment_is_mapped=False) == () + + +class TestResolveSupportedReasoningEfforts: + def test_a_reasoning_model_with_no_flags_at_all_resolves_to_unknown(self): + # 689 of the map's 854 reasoning entries carry no effort flag, and the o-series, xai and + # bedrock nova entries among them accept neither none nor minimal, so composing a set out of + # the opt-out defaults alone would advertise levels those providers reject. + assert resolve_supported_reasoning_efforts({"supports_reasoning": True}, deployment_is_mapped=True) is None + + def test_explicit_false_removes_an_opt_out_level(self): + # The gpt-5.5-pro shape from the model map: only medium/high/xhigh are accepted upstream. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": False, + "supports_low_reasoning_effort": False, + "supports_xhigh_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("medium", "high", "xhigh") + + def test_explicit_true_adds_the_opt_in_levels(self): + # The claude-opus shape: xhigh and max explicitly true, everything else absent. + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("none", "minimal", "low", "medium", "high", "xhigh", "max") + + def test_opt_in_flag_set_false_stays_excluded(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "supports_minimal_reasoning_effort": True, + "supports_xhigh_reasoning_effort": False, + }, + deployment_is_mapped=True, + ) + assert resolved == ("none", "minimal", "low", "medium", "high") + + +class TestBareModelNameFallback: + def test_a_prefixed_entry_inherits_the_flags_of_its_unprefixed_twin(self): + """azure/gpt-5-mini carries no effort flag while gpt-5-mini carries three, and the request + path resolves capability flags through that same twin (#20885). Reading only the prefixed + entry would answer unknown for a model the map fully describes.""" + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="gpt-5-mini", custom_llm_provider="azure")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "minimal", + "low", + "medium", + "high", + ) + + def test_the_prefixed_entry_wins_over_its_twin_per_flag(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-mini", + "supports_xhigh_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("minimal", "low", "medium", "high", "xhigh") + + +class TestNoneLevelPolarity: + def test_none_stays_opt_out_off_azure(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "openai", + "key": "openai/some-reasoner", + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + def test_none_stays_opt_out_on_an_azure_model_outside_the_gpt_5_family(self): + """AzureOpenAIGPT5Config is selected by is_model_gpt_5_model, so an azure o-series or + anthropic deployment never reaches the gate that refuses none and must keep the level.""" + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/o3", + "supports_max_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + def test_azure_gpt_5_without_the_flag_does_not_advertise_none(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-turbo", + "supports_minimal_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved == ("minimal", "low", "medium", "high") + + def test_azure_gpt_5_with_the_flag_advertises_none(self): + resolved = resolve_supported_reasoning_efforts( + { + "supports_reasoning": True, + "litellm_provider": "azure", + "key": "azure/gpt-5-turbo", + "supports_none_reasoning_effort": True, + }, + deployment_is_mapped=True, + ) + assert resolved is not None and "none" in resolved + + @pytest.mark.parametrize( + "model_key", + ["azure/gpt-5", "azure/gpt-5-mini", "azure/gpt-5-nano", "azure/gpt-5.2", "azure/gpt-5.6"], + ) + def test_azure_advertisement_matches_the_azure_request_gate(self, model_key): + """AzureOpenAIGPT5Config raises UnsupportedParamsError on reasoning_effort='none' for models + it does not flag, so advertising the level there would offer routing a 400.""" + from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model=model_key.split("/", 1)[1], custom_llm_provider="azure")) + resolved = resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) + + assert resolved is not None + gate_accepts_none = AzureOpenAIGPT5Config._supports_reasoning_effort_level(model_key, "none") + assert ("none" in resolved) is gate_accepts_none + + +class TestIntersectSupportedReasoningEfforts: + def test_unknown_never_narrows(self): + assert intersect_supported_reasoning_efforts(["medium", "high"], None) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, ["medium", "high"]) == ("medium", "high") + assert intersect_supported_reasoning_efforts(None, None) is None + + def test_intersection_keeps_canonical_order(self): + assert intersect_supported_reasoning_efforts( + ["max", "high", "medium", "xhigh"], ["xhigh", "medium", "minimal"] + ) == ("medium", "xhigh") + + def test_disjoint_sets_intersect_to_empty(self): + assert intersect_supported_reasoning_efforts(["max"], ["minimal"]) == () diff --git a/tests/test_litellm/test_check_migrations_no_data_rewrites.py b/tests/test_litellm/test_check_migrations_no_data_rewrites.py new file mode 100644 index 00000000000..ccba351deaf --- /dev/null +++ b/tests/test_litellm/test_check_migrations_no_data_rewrites.py @@ -0,0 +1,1705 @@ +"""Tests for tests/code_coverage_tests/check_migrations_no_data_rewrites.py. + +The checker reads migration.sql as SQL rather than as text, so the cases that matter +are the ones a grep would get wrong: the referential actions in a foreign key, of which +the shipped migrations carry 60, an `UPDATE` inside a string literal or a comment, and +an `UPDATE` hidden in the `DO $$ ... $$` block this repo uses for conditional DDL. +""" + +import importlib.util +import sys +from pathlib import Path + +_CHECKER_PATH = Path(__file__).resolve().parents[1] / "code_coverage_tests" / "check_migrations_no_data_rewrites.py" +_SPEC = importlib.util.spec_from_file_location("check_migrations_no_data_rewrites", _CHECKER_PATH) +assert _SPEC is not None and _SPEC.loader is not None +checker = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = checker +_SPEC.loader.exec_module(checker) + + +def _scan(tmp_path: Path, sql: str) -> tuple: + directory = tmp_path / "20260101000000_fixture" + directory.mkdir(exist_ok=True) + (directory / "migration.sql").write_text(sql, encoding="utf-8") + return checker.scan_migration(directory) + + +def _keywords(tmp_path: Path, sql: str) -> tuple: + return tuple(violation.keyword for violation in _scan(tmp_path, sql)) + + +class TestRowRewritesAreFlagged: + def test_update_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_delete_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'DELETE FROM "Foo" WHERE "a" IS NULL;') == ("DELETE",) + + def test_update_without_trailing_semicolon_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'UPDATE "Foo" SET "a" = 1') == ("UPDATE",) + + def test_lowercase_update_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'update "Foo" set "a" = 1;') == ("UPDATE",) + + def test_merge_is_flagged(self, tmp_path): + sql = 'MERGE INTO "Foo" t USING "Bar" s ON t."id" = s."id" WHEN MATCHED THEN UPDATE SET "a" = s."a";' + assert _keywords(tmp_path, sql) == ("MERGE",) + + def test_every_offending_statement_is_reported(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nDELETE FROM "Bar";' + assert _keywords(tmp_path, sql) == ("UPDATE", "DELETE") + + def test_the_incident_migration_is_flagged(self, tmp_path): + sql = ( + 'UPDATE "LiteLLM_SpendLogs"\n' + ' SET "created_at" = "endTime",\n' + ' "updated_at" = "endTime"\n' + ' WHERE "created_at" > "endTime" + interval \'1 hour\';\n' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + +class TestSchemaStatementsPass: + def test_on_delete_cascade_is_not_a_data_rewrite(self, tmp_path): + sql = ( + 'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") ' + 'REFERENCES "B"("id") ON DELETE CASCADE ON UPDATE CASCADE;' + ) + assert _keywords(tmp_path, sql) == () + + def test_on_delete_set_null_is_not_a_data_rewrite(self, tmp_path): + sql = ( + 'ALTER TABLE "A" ADD CONSTRAINT "A_b_fkey" FOREIGN KEY ("b") ' + 'REFERENCES "B"("id") ON DELETE SET NULL ON UPDATE CASCADE;' + ) + assert _keywords(tmp_path, sql) == () + + def test_add_column_with_default_passes(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN IF NOT EXISTS "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP;' + assert _keywords(tmp_path, sql) == () + + def test_drop_table_passes(self, tmp_path): + assert _keywords(tmp_path, 'DROP TABLE IF EXISTS "Foo";') == () + + def test_empty_file_passes(self, tmp_path): + assert _keywords(tmp_path, "") == () + + def test_only_comments_passes(self, tmp_path): + assert _keywords(tmp_path, "-- nothing to do here\n") == () + + +class TestInsert: + def test_insert_values_is_bounded_and_passes(self, tmp_path): + assert _keywords(tmp_path, "INSERT INTO \"Foo\" (\"id\") VALUES ('a'), ('b');") == () + + def test_insert_select_scans_and_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar";') == ("INSERT ... SELECT",) + + def test_insert_values_with_a_scalar_subquery_passes(self, tmp_path): + sql = 'INSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT max("id")::text FROM "Bar"));' + assert _keywords(tmp_path, sql) == () + + def test_insert_values_with_a_scalar_subquery_per_row_passes(self, tmp_path): + sql = ( + 'INSERT INTO "Config" ("k", "v") VALUES\n' + " ('a', (SELECT \"id\" FROM \"Bar\" WHERE \"n\" = 'a')),\n" + " ('b', (SELECT \"id\" FROM \"Bar\" WHERE \"n\" = 'b'));" + ) + assert _keywords(tmp_path, sql) == () + + def test_values_inside_a_subquery_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM (VALUES (1), (2)) AS "v"("id");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_values_after_a_set_operation_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" UNION ALL VALUES (1);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_values_after_an_except_does_not_bound_an_insert_select(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" EXCEPT VALUES (1);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_select_term_after_a_values_list_is_still_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1), (2) UNION ALL SELECT "id" FROM "Bar";' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_without_a_column_list_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_row_source_spanning_lines_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id")\n(\n SELECT "id" FROM "Bar"\n);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_parenthesised_select_over_a_values_list_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT * FROM (VALUES (1), (2)) AS "v"("id"));' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_set_operation_over_parenthesised_selects_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT 1) UNION (SELECT 2);' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_joined_to_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_excepting_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) EXCEPT (SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_values_list_joined_to_a_parenthesised_table_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES (1) UNION ALL (TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_set_operation_inside_a_values_list_does_not_flag_it(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") VALUES ((SELECT 1 UNION SELECT 2 LIMIT 1));' + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_a_set_operated_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") VALUES ((SELECT max("id") FROM "Bar"))' + " UNION ALL VALUES (2);" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_a_parenthesised_values_list_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")));' + assert _keywords(tmp_path, sql) == () + + def test_a_scalar_subquery_in_set_operated_parenthesised_values_lists_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")))' + " UNION ALL (VALUES (2));" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_set_operation_inside_a_values_list_does_not_split_the_terms(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") VALUES ((SELECT max("id") FROM "Bar"' + ' UNION SELECT max("id") FROM "Bar")) UNION ALL VALUES (2);' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_conflict_target_after_a_parenthesised_row_source_does_not_hide_it(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar")' + ' ON CONFLICT ("id") DO NOTHING;' + ) + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_returning_list_after_a_parenthesised_row_source_does_not_hide_it(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") RETURNING ("id");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_conflict_target_beside_a_bounded_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES ((SELECT max("id") FROM "Bar")))' + ' ON CONFLICT ("id") DO NOTHING;' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_query_term_written_before_a_values_term_is_still_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") UNION ALL (VALUES (2));' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_table_term_beside_parenthesised_values_is_still_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1)) UNION ALL (TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_table_row_source_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "Foo" TABLE "Bar";') == ("INSERT ... TABLE",) + + def test_a_table_named_in_the_insert_target_does_not_flag_it(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "audit table" ("id") VALUES (1);') == () + + def test_a_returning_subquery_after_a_wrapped_values_list_is_not_the_row_source(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_a_conflict_update_after_a_wrapped_values_list_stays_bounded(self, tmp_path): + sql = ( + 'INSERT INTO "Foo" ("id") (VALUES (1))' + ' ON CONFLICT ("id") DO UPDATE SET "id" = (SELECT max("id") FROM "Bar");' + ) + assert _keywords(tmp_path, sql) == () + + def test_a_wrapped_values_list_of_several_rows_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1), (2)) RETURNING (SELECT count(*) FROM "Bar");' + assert _keywords(tmp_path, sql) == () + + def test_the_row_source_names_its_own_keyword_not_a_later_subquery(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (TABLE "Bar") RETURNING (SELECT count(*) FROM "Baz");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_select_term_wrapped_beside_a_values_term_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1) UNION ALL SELECT "id" FROM "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_a_table_term_wrapped_beside_a_values_term_is_flagged(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1) UNION ALL TABLE "Bar");' + assert _keywords(tmp_path, sql) == ("INSERT ... TABLE",) + + def test_a_wrapped_set_operation_of_values_lists_stays_bounded(self, tmp_path): + sql = 'INSERT INTO "Foo" ("id") (VALUES (1) UNION ALL VALUES (2));' + assert _keywords(tmp_path, sql) == () + + +class TestCommonTableExpressions: + def test_cte_led_update_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) UPDATE "Foo" SET "a" = 1 FROM batch;' + assert _keywords(tmp_path, sql) == ("WITH ... UPDATE",) + + def test_cte_led_delete_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo" LIMIT 100) DELETE FROM "Foo" USING batch;' + assert _keywords(tmp_path, sql) == ("WITH ... DELETE",) + + def test_cte_led_insert_select_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") SELECT "id" FROM batch;' + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_cte_led_insert_from_a_parenthesised_select_is_flagged(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Bar") INSERT INTO "Foo" ("id") (SELECT "id" FROM batch);' + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_cte_led_insert_into_a_values_list_passes(self, tmp_path): + sql = 'WITH batch AS (SELECT max("id") FROM "Bar") INSERT INTO "Foo" ("id") VALUES (1);' + assert _keywords(tmp_path, sql) == () + + def test_read_only_cte_passes(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "Foo") SELECT count(*) FROM batch;' + assert _keywords(tmp_path, sql) == () + + def test_cte_led_insert_values_is_bounded_and_passes(self, tmp_path): + sql = 'WITH latest AS (SELECT max("id") AS "id" FROM "Bar")\nINSERT INTO "Config" ("k", "v") VALUES (\'rev\', (SELECT "id"::text FROM latest));' + assert _keywords(tmp_path, sql) == () + + def test_a_writable_cte_bounded_by_values_passes(self, tmp_path): + sql = 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id") SELECT * FROM added;' + assert _keywords(tmp_path, sql) == () + + def test_a_writable_cte_copying_a_query_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_bounded_writable_cte_does_not_hide_a_copying_one_beside_it(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") VALUES (1) RETURNING "id"),' + ' copied AS (INSERT INTO "Baz" ("id") SELECT "id" FROM "Bar" RETURNING "id")' + " SELECT * FROM added, copied;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + def test_a_writable_cte_wrapping_its_row_source_is_flagged(self, tmp_path): + sql = ( + 'WITH added AS (INSERT INTO "Foo" ("id") (SELECT "id" FROM "Bar") RETURNING "id")' + " SELECT * FROM added;" + ) + assert _keywords(tmp_path, sql) == ("WITH ... INSERT ... SELECT",) + + +class TestDollarQuotedBlocks: + def test_update_inside_do_block_is_flagged(self, tmp_path): + sql = 'DO $$\nBEGIN\n UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_conditional_ddl_do_block_passes(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'x') THEN\n" + ' ALTER TABLE "Foo" DROP CONSTRAINT "x";\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_guarded_update_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo") THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_guard_with_a_nested_call_still_flags_the_update(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo" WHERE lower("a") = \'x\' UNION SELECT 1) THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_tagged_dollar_quote_is_scanned(self, tmp_path): + sql = 'DO $body$\nBEGIN\n DELETE FROM "Foo";\nEND $body$;' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_tagged_dollar_quote_holds_an_apostrophe(self, tmp_path): + sql = 'INSERT INTO "Foo" ("t") VALUES ($body$don\'t$body$);\nUPDATE "Bar" SET "b" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_semicolons_inside_do_block_do_not_split_outer_statements(self, tmp_path): + sql = 'DO $$ BEGIN PERFORM 1; END $$;\nALTER TABLE "Foo" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == () + + def test_line_number_inside_a_do_block_counts_from_the_top_of_the_file(self, tmp_path): + sql = ( + "-- AlterTable\n" + 'ALTER TABLE "Foo" ADD COLUMN "b" INT;\n' + "\n" + "DO $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "b" = 1;\n' + "END $$;" + ) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_line_number_inside_a_nested_body_counts_from_the_top_of_the_file(self, tmp_path): + sql = ( + "-- CreateIndex\n" + 'CREATE INDEX "i" ON "Foo"("a");\n' + "\n" + "DO $outer$\n" + "BEGIN\n" + " EXECUTE $inner$\n" + ' UPDATE "Foo" SET "a" = 1\n' + " $inner$;\n" + "END $outer$;" + ) + assert _scan(tmp_path, sql)[0].line == 7 + + +class TestStoredRoutines: + DEFINITION = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + PROCEDURE = ( + "CREATE OR REPLACE PROCEDURE sweep() AS $$\n" + "BEGIN\n" + ' DELETE FROM "Foo";\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + + def test_a_function_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION) == () + + def test_a_procedure_body_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE) == () + + def test_a_function_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION + "SELECT backfill();\n") == ("UPDATE",) + + def test_a_procedure_the_migration_calls_is_flagged(self, tmp_path): + assert _keywords(tmp_path, self.PROCEDURE + "CALL sweep();\n") == ("DELETE",) + + def test_a_call_written_above_the_definition_still_counts(self, tmp_path): + assert _keywords(tmp_path, "SELECT backfill();\n" + self.DEFINITION) == ("UPDATE",) + + def test_a_call_from_inside_a_do_block_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN PERFORM backfill(); END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_through_a_quoted_identifier_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_unrelated_quoted_identifier_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'SELECT "other"();\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_column_sharing_the_routine_name_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'ALTER TABLE "Foo" ADD COLUMN "backfill" int;\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_index_target_sharing_the_routine_name_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Foo" ("backfill");\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_call_with_space_before_its_parenthesis_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill" ();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_quoted_call_with_a_block_comment_before_its_parenthesis_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill" /* reason */ ();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_quoted_call_with_a_line_comment_before_its_parenthesis_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT "backfill" -- run it\n();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_quoted_column_followed_by_a_comment_is_still_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'ALTER TABLE "Foo" ADD COLUMN "backfill" /* note */ int;\n' + assert _keywords(tmp_path, sql) == () + + def test_a_like_named_table_with_a_column_list_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TABLE "backfill" (id int);\n' + assert _keywords(tmp_path, sql) == () + + def test_an_insert_into_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO "backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_an_insert_into_a_like_named_table_with_a_commented_column_list_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO "backfill" /* cols */ ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_foreign_key_referencing_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TABLE "Bar" (id int REFERENCES "backfill" ("id"));\n' + assert _keywords(tmp_path, sql) == () + + def test_an_index_on_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "backfill" ("id");\n' + assert _keywords(tmp_path, sql) == () + + def test_an_if_not_exists_table_named_after_the_routine_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TABLE IF NOT EXISTS "backfill" (id int);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_copy_into_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'COPY "backfill" ("id") FROM stdin;\n' + assert _keywords(tmp_path, sql) == () + + def test_a_set_returning_call_in_from_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT * FROM "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_in_a_join_condition_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT 1 FROM "Bar" b JOIN "Baz" z ON "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_in_an_index_predicate_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Foo" ("a") WHERE "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_join_condition_call_after_an_earlier_index_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Foo" ("a");\nSELECT 1 FROM "Bar" b JOIN "Baz" z ON "backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_insert_into_a_schema_qualified_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO public."backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_schema_qualified_index_on_a_like_named_table_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON public."backfill" ("id");\n' + assert _keywords(tmp_path, sql) == () + + def test_a_schema_qualified_quoted_call_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT public."backfill"();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_spaced_schema_qualifier_on_an_insert_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO public . "backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_spaced_schema_qualifier_on_an_index_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON public . "backfill" ("id");\n' + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_schema_with_spaces_round_the_dot_stays_a_relation(self, tmp_path): + sql = self.DEFINITION + 'INSERT INTO "public" . "backfill" ("id") VALUES (1);\n' + assert _keywords(tmp_path, sql) == () + + def test_a_spaced_schema_qualified_call_still_counts(self, tmp_path): + sql = self.DEFINITION + 'SELECT public . "backfill" ();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_schema_qualified_call_in_an_index_expression_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Bar" (public."backfill"("a"));\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_spaced_schema_qualified_call_in_an_index_expression_still_counts(self, tmp_path): + sql = self.DEFINITION + 'CREATE INDEX ON "Bar" (public . "backfill" ("a"));\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_trigger_wiring_the_function_up_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + 'CREATE TRIGGER t AFTER INSERT ON "Foo" EXECUTE FUNCTION backfill();\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_schema_qualified_definition_nothing_calls_passes(self, tmp_path): + assert _keywords(tmp_path, self.DEFINITION.replace("backfill()", "public.backfill()")) == () + + def test_a_schema_qualified_function_the_migration_calls_is_flagged(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", "public.backfill()") + "SELECT public.backfill();\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_the_name_written_only_in_a_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "-- backfill() is run by hand after the deploy\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_do_body_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN\n-- backfill() is run by hand after the deploy\nPERFORM 1;\nEND; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_do_body_block_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN /* backfill() runs later */ PERFORM 1; END; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_nested_body_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN EXECUTE $q$SELECT 1 -- backfill() runs later\n$q$; END; $$;\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_in_an_executed_literal_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN EXECUTE 'SELECT backfill()'; END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_after_a_literal_holding_comment_dashes_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO $$ BEGIN RAISE NOTICE '--'; PERFORM backfill(); END; $$;\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_from_inside_a_single_quoted_do_block_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN PERFORM backfill(); END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_executed_literal_inside_a_single_quoted_do_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN EXECUTE ''SELECT backfill()''; END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_variable_run_by_execute_in_a_single_quoted_do_counts_as_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'DECLARE q text; BEGIN q := ''SELECT backfill()''; EXECUTE q; END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_after_a_single_quoted_literal_holding_comment_dashes_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN RAISE NOTICE ''--''; PERFORM backfill(); END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_call_after_a_single_quoted_literal_opening_a_block_comment_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN RAISE NOTICE ''/*''; PERFORM backfill(); END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_executed_literal_after_a_single_quoted_comment_dash_string_still_counts(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN RAISE NOTICE ''--''; EXECUTE ''SELECT backfill()''; END';\n" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_long_run_of_escaped_quotes_before_an_uncalled_definition_is_not_a_call(self, tmp_path): + escaped_quotes = "'" * 84 + sql = f"DO 'BEGIN RAISE NOTICE ''{escaped_quotes}''; PERFORM 1; END';\n" + self.DEFINITION + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_single_quoted_do_comment_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "DO 'BEGIN\n-- backfill() runs later\nPERFORM 1; END';\n" + assert _keywords(tmp_path, sql) == () + + def test_the_name_written_only_in_a_non_runnable_string_is_not_a_call(self, tmp_path): + sql = self.DEFINITION + "SELECT 'backfill() runs after the deploy';\n" + assert _keywords(tmp_path, sql) == () + + def test_a_recursive_call_does_not_count_as_the_migration_calling_it(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill(n int) RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " PERFORM backfill(n - 1);\n" + "END;\n" + "$$ LANGUAGE plpgsql;\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_quoted_routine_name_is_read_rather_than_trusted(self, tmp_path): + sql = self.DEFINITION.replace("backfill()", '"back fill"()') + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_do_block_is_not_a_routine_definition(self, tmp_path): + sql = 'DO $$ BEGIN UPDATE "Foo" SET "a" = 1; END; $$;\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_definition_written_after_another_statement_is_still_recognised(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;\n' + self.DEFINITION + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_rewrite_in_a_routine_the_migration_calls(self, tmp_path): + sql = ( + "CREATE FUNCTION backfill() RETURNS void AS $$\n" + "BEGIN\n" + ' UPDATE "Foo" SET "a" = 1; -- data-migration-ok: single config row\n' + "END;\n" + "$$ LANGUAGE plpgsql;\n" + "SELECT backfill();\n" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_called_routine_reports_the_line_inside_its_body(self, tmp_path): + assert _scan(tmp_path, self.DEFINITION + "SELECT backfill();\n")[0].line == 3 + + +class TestLoopBodies: + def test_a_rewrite_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_delete_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' DELETE FROM "Foo" WHERE "id" = r."id";\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_join_using_in_the_loop_query_does_not_hide_the_body(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT a."id" FROM "A" a JOIN "B" b USING ("id") LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_executed_in_a_query_driven_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + " EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_nested_under_a_guard_inside_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' IF r."id" > 0 THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_inside_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP FOR b IN SELECT "id" FROM "B" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_supplying_a_nested_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE a record;\n" + "DECLARE b record;\n" + "BEGIN\n" + ' FOR a IN SELECT "id" FROM "A" LOOP\n' + ' FOR b IN UPDATE "Foo" SET "x" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP; END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_a_loop_running_only_ddl_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' CREATE INDEX "i" ON "Foo"("a");\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_loop_over_a_rewrite_returning_rows_is_flagged_once(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + ' FOR r IN UPDATE "Foo" SET "a" = 1 RETURNING "id" LOOP\n' + " NULL;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_marker_on_a_loop_exempts_the_rewrite_it_repeats(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE r record;\n" + "BEGIN\n" + " -- data-migration-ok: one row\n" + ' FOR r IN SELECT "id" FROM "Bar" LOOP\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_select_for_update_lock_is_not_read_as_a_loop(self, tmp_path): + sql = 'DO $$\nBEGIN\n PERFORM 1 FROM "Foo" FOR UPDATE;\nEND $$;' + assert _keywords(tmp_path, sql) == () + + +class TestQuotingAndComments: + def test_update_inside_string_literal_passes(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "note" TEXT NOT NULL DEFAULT \'UPDATE nothing\';' + assert _keywords(tmp_path, sql) == () + + def test_escaped_quote_inside_string_does_not_leak(self, tmp_path): + sql = "ALTER TABLE \"Foo\" ADD COLUMN \"note\" TEXT NOT NULL DEFAULT 'it''s fine';\n" + assert _keywords(tmp_path, sql) == () + + def test_update_inside_line_comment_passes(self, tmp_path): + assert _keywords(tmp_path, '-- UPDATE "Foo" SET "a" = 1;\nDROP TABLE "Bar";') == () + + def test_update_inside_block_comment_passes(self, tmp_path): + assert _keywords(tmp_path, '/* UPDATE "Foo" SET "a" = 1; */\nDROP TABLE "Bar";') == () + + def test_nested_block_comment_passes(self, tmp_path): + sql = '/* outer /* UPDATE "Foo" SET "a" = 1; */ still comment */\nDROP TABLE "Bar";' + assert _keywords(tmp_path, sql) == () + + def test_nested_block_comment_masks_past_the_inner_close(self, tmp_path): + sql = '/* outer /* inner */ UPDATE "Foo" SET "a" = 1; */\nDROP TABLE "Bar";' + assert _keywords(tmp_path, sql) == () + + def test_update_inside_quoted_identifier_passes(self, tmp_path): + assert _keywords(tmp_path, 'ALTER TABLE "UPDATE Foo" ADD COLUMN "b" TEXT;') == () + + def test_select_in_a_quoted_identifier_does_not_make_an_insert_a_rewrite(self, tmp_path): + assert _keywords(tmp_path, 'INSERT INTO "SELECT Foo" ("id") VALUES (\'a\');') == () + + def test_update_in_a_quoted_identifier_does_not_make_a_cte_a_rewrite(self, tmp_path): + sql = 'WITH batch AS (SELECT "id" FROM "UPDATE Foo") SELECT count(*) FROM batch;' + assert _keywords(tmp_path, sql) == () + + def test_positional_parameter_is_not_a_dollar_quote(self, tmp_path): + sql = 'ALTER TABLE "Foo" ADD COLUMN "b" TEXT;\nUPDATE "Foo" SET "b" = $1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + +class TestEscapeHatch: + def test_marker_with_reason_exempts_the_statement(self, tmp_path): + sql = '-- data-migration-ok: one row per tenant, at most a few hundred\nUPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == () + + def test_marker_without_reason_does_not_exempt(self, tmp_path): + assert _keywords(tmp_path, '-- data-migration-ok:\nUPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_marker_exempts_only_its_own_statement(self, tmp_path): + sql = '-- data-migration-ok: bounded to in-flight jobs\nUPDATE "Foo" SET "a" = 1;\nUPDATE "Bar" SET "b" = 2;\n' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_marker_works_inside_a_do_block(self, tmp_path): + sql = 'DO $$\nBEGIN\n -- data-migration-ok: single row\n UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == () + + def test_marker_below_the_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_marker_written_below_its_statement_leaves_that_statement_flagged(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: bounded\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_marker_inside_a_do_block_below_the_first_line_exempts(self, tmp_path): + sql = ( + "-- AlterTable\n" + 'ALTER TABLE "Foo" ADD COLUMN "b" INT;\n' + "\n" + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: one config row\n" + ' UPDATE "Foo" SET "b" = 1;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_trailing_a_statement_exempts_that_statement(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1; -- data-migration-ok: one row\nALTER TABLE "Bar" ADD COLUMN "b" TEXT;' + assert _keywords(tmp_path, sql) == () + + def test_a_marker_trailing_a_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1; -- data-migration-ok: one row\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 2 + + def test_a_marker_trailing_a_multiline_statement_does_not_exempt_the_next_one(self, tmp_path): + sql = ( + 'UPDATE "Foo"\n' + ' SET "a" = 1; -- data-migration-ok: one row\n' + 'UPDATE "Bar" SET "b" = 2;' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_marker_alone_between_two_statements_belongs_to_the_one_below_it(self, tmp_path): + sql = 'UPDATE "Foo" SET "a" = 1;\n-- data-migration-ok: one row\nUPDATE "Bar" SET "b" = 2;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_a_marker_a_blank_line_above_a_statement_does_not_exempt_it(self, tmp_path): + sql = '-- data-migration-ok: one row\n\nUPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_trailing_marker_exempts_only_the_statement_it_follows(self, tmp_path): + sql = 'DELETE FROM "Foo" WHERE "a" = 1; UPDATE "Bar" SET "b" = 2; -- data-migration-ok: one row' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_marker_above_a_shared_line_exempts_only_the_first_statement_on_it(self, tmp_path): + sql = '-- data-migration-ok: one row\nUPDATE "Foo" SET "a" = 1; DELETE FROM "Bar" WHERE "b" = 2;' + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_marker_on_the_opening_line_of_a_statement_exempts_that_statement(self, tmp_path): + sql = 'UPDATE "Foo" -- data-migration-ok: one row\n SET "a" = 1;\nDELETE FROM "Bar";' + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_marker_above_a_do_block_does_not_exempt_a_rewrite_inside_it(self, tmp_path): + sql = ( + "-- data-migration-ok: bounded, this belongs to the insert below\n" + "INSERT INTO \"Config\" (\"k\") VALUES ('x');\n" + "\n" + 'DO $$ BEGIN UPDATE "Foo" SET "b" = 1; END $$;' + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_marker_directly_above_a_do_block_does_not_exempt_its_body(self, tmp_path): + sql = ( + "-- data-migration-ok: seeding two default rows\n" + "DO $$\n" + "BEGIN\n" + ' INSERT INTO "Foo" ("a") VALUES (1);\n' + ' UPDATE "Foo" SET "a" = 1;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_marker_directly_above_a_one_line_do_block_does_not_exempt_its_body(self, tmp_path): + sql = '-- data-migration-ok: seeding one default row\nDO $$ BEGIN UPDATE "Foo" SET "a" = 1; END $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 2 + + def test_marker_on_the_do_line_does_not_exempt_its_body(self, tmp_path): + sql = ( + "DO $$ -- data-migration-ok: bounded to one row\n" + "BEGIN\n" + ' IF EXISTS (SELECT 1 FROM "Foo") THEN\n' + ' UPDATE "Foo" SET "a" = 1;\n' + " END IF;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marked_rewrite_does_not_exempt_a_later_one_in_the_same_block(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + ' UPDATE "Bar" SET "b" = 2;\n' + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + +class TestDynamicSql: + def test_execute_of_a_quoted_update_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_execute_of_a_quoted_delete_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\"';\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_execute_of_a_formatted_update_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE format('UPDATE %I SET \"a\" = 1', 'Foo');\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_execute_of_a_dollar_quoted_update_is_flagged(self, tmp_path): + sql = 'DO $outer$\nBEGIN\n EXECUTE $q$UPDATE "Foo" SET "a" = 1$q$;\nEND $outer$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_doubled_quote_inside_executed_sql_does_not_hide_the_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = date_trunc(''day'', \"t\")';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_quoted_as_data_inside_executed_sql_is_not_run(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''UPDATE \"Foo\" SET \"a\" = 1''';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_doubled_quote_does_not_split_the_literal_it_sits_in(self, tmp_path): + sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('a''UPDATE \"Bar\" SET \"a\" = 1''b');" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_following_a_doubled_quote_in_the_same_payload_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''x''; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_comment_dash_inside_a_doubled_quote_does_not_hide_a_later_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''--''; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_block_comment_open_inside_a_doubled_quote_does_not_hide_a_later_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT ''/*''; DELETE FROM \"Foo\"';\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_genuinely_commented_out_inside_executed_sql_is_not_run(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1 -- UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_below_escaped_quotes_in_a_multiline_payload_reports_its_own_line(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE '\n" + "SELECT ''a'', ''b'', ''c'', ''d'', ''e''\n" + "; UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_in_a_later_command_before_bind_values_is_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1; DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_bind_value_naming_a_rewrite_is_not_run(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'INSERT INTO \"Audit\" (\"note\") VALUES ($1)'" + " USING 'DELETE FROM \"Foo\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_executed_with_bind_values_is_still_flagged(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" WHERE \"a\" = $1' USING 1;\nEND $$;" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_using_written_inside_the_command_does_not_end_it(self, tmp_path): + sql = ( + "DO $$\nBEGIN\n EXECUTE 'DELETE FROM \"Foo\" USING \"Bar\"" + " WHERE \"Foo\".\"a\" = \"Bar\".\"a\"';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_join_using_in_a_subquery_building_the_command_does_not_end_it(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) || 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_join_using_does_not_take_the_place_of_the_real_bind_values(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) || 'UPDATE \"Foo\" SET \"a\" = $1' USING 2;\nEND $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_bind_value_naming_a_rewrite_after_a_subquery_join_is_not_run(self, tmp_path): + sql = ( + "DO $$\nDECLARE v text;\nBEGIN\n EXECUTE (SELECT v FROM \"A\" x JOIN \"A\" y" + " USING (\"id\")) USING 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_execute_of_ddl_passes(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_execute_of_a_read_only_query_passes(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'SELECT 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_executed_rewrite(self, tmp_path): + sql = "DO $$\nBEGIN\n -- data-migration-ok: one row\n EXECUTE 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_literal_that_is_not_executed_is_still_inert(self, tmp_path): + sql = "INSERT INTO \"Foo\" (\"note\") VALUES ('UPDATE \"Bar\" SET \"a\" = 1');" + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_declared_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_a_rewrite_assigned_in_the_body_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " stmt := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_selected_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'UPDATE \"Foo\" SET \"a\" = 1' INTO stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_selected_into_a_strict_target_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'DELETE FROM \"Foo\"' INTO STRICT stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_assigned_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " stmt = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_assigned_with_a_bare_equals_after_then_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " IF true THEN stmt = 'DELETE FROM \"Foo\"'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_declared_with_a_bare_equals_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text = 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_an_insert_target_table_is_not_read_as_an_assignment(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " audit text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " INSERT INTO audit (note) VALUES ('DELETE FROM \"Foo\" is left to the app');\n" + " EXECUTE audit;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_returned_into_a_variable_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " INSERT INTO \"Log\" (\"sql\") VALUES ('UPDATE \"Foo\" SET \"a\" = 1')\n" + " RETURNING \"sql\" INTO stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_assigned_past_an_earlier_comparison_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + "BEGIN\n" + " IF total = 1 THEN stmt = 'DELETE FROM \"Foo\" WHERE true'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 6 + + def test_a_rewrite_assigned_past_a_loop_comparison_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 3;\n" + "BEGIN\n" + " WHILE total >= 1 LOOP stmt = 'UPDATE \"Foo\" SET \"a\" = 1'; END LOOP;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_selected_into_a_target_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'UPDATE \"Foo\" SET \"a\" = 1'\n" + " INTO\n" + " stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_selected_into_a_strict_target_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " SELECT 'DELETE FROM \"Foo\"' INTO STRICT\n" + " stmt;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_rewrite_executed_by_a_name_a_line_down_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE\n" + " stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_rewrite_compared_against_is_not_an_assignment(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " IF stmt = 'UPDATE \"Foo\" SET \"a\" = 1' THEN\n" + " RAISE NOTICE 'the application owns that one';\n" + " END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_passed_to_execute_as_a_parameter_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'UPDATE \"Foo\" SET \"a\" = 1';\n" + "BEGIN\n" + " EXECUTE 'INSERT INTO \"Log\" (\"sql\") VALUES ($1)' USING stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_compared_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total = 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " RAISE NOTICE 'purge script? %', ok;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_named_as_an_argument_beside_an_assignment_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := probe_match(subject => stmt, wanted => 'DELETE FROM \"Foo\"');\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_compared_after_a_wider_comparison_is_not_run(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + " ok boolean;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " ok := total >= 1 AND stmt = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_assigned_through_a_case_expression_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int := 1;\n" + "BEGIN\n" + " stmt := CASE WHEN total = 1 THEN 'DELETE FROM \"Foo\"' ELSE 'SELECT 1' END;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_query_reaching_into_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\"'\n" + " INTO total;\n" + " SELECT (CASE WHEN total > 0 THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\"\n" + " WHERE \"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_query_reaching_using_past_an_execute_is_not_a_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + " total int;\n" + "BEGIN\n" + " stmt := 'CREATE INDEX IF NOT EXISTS \"ix_a\" ON \"Foo\" (\"a\")';\n" + " EXECUTE 'SELECT count(*) FROM \"Foo\" WHERE \"a\" = $1'\n" + " USING 'k1';\n" + " SELECT (CASE WHEN true THEN 1 ELSE 2 END) INTO total\n" + " FROM \"Foo\" x JOIN \"Foo\" y USING (\"a\")\n" + " WHERE x.\"a\" = 'DELETE FROM \"Foo\"';\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_rewrite_walked_by_a_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOR stmt IN SELECT 'DELETE FROM \"Foo\"' LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 5 + + def test_a_rewrite_walked_by_a_foreach_loop_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text;\n" + "BEGIN\n" + " FOREACH stmt IN ARRAY ARRAY['UPDATE \"Foo\" SET \"a\" = 1'] LOOP\n" + " EXECUTE stmt;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_a_loop_over_a_query_running_nothing_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " rec record;\n" + "BEGIN\n" + " FOR rec IN SELECT \"a\" FROM \"Foo\" LOOP\n" + " RAISE NOTICE 'the DELETE FROM \"Foo\" path is the application''s: %', rec;\n" + " END LOOP;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_literal_selected_into_a_variable_nothing_runs_is_inert(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text;\n" + "BEGIN\n" + " SELECT 'UPDATE of legacy rows is skipped' INTO msg;\n" + " RAISE NOTICE '%', msg;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_comparing_an_executed_variable_does_not_flag_the_comparison(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\n" + "BEGIN\n" + " IF stmt = 'DELETE FROM \"Foo\"' THEN RAISE NOTICE 'never'; END IF;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_ddl_assigned_to_a_variable_passes(self, tmp_path): + sql = "DO $$\nDECLARE\n stmt text := 'ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT';\nBEGIN\n EXECUTE stmt;\nEND $$;" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_rewrite_held_in_a_variable(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " stmt text := 'UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_execute_whose_sql_starts_on_a_later_line(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " EXECUTE '\n" + " UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_an_assignment_whose_sql_starts_on_a_later_line(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " -- data-migration-ok: one config row, keyed by its primary key\n" + " stmt text := '\n" + " UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''';\n" + "BEGIN\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_an_unmarked_execute_whose_sql_starts_on_a_later_line_is_flagged(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE '\n" + " UPDATE \"Foo\" SET \"a\" = 1';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_message_assigned_but_never_executed_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text := 'UPDATE of legacy rows skipped, the application backfills them';\n" + "BEGIN\n" + " RAISE NOTICE '%', msg;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_notice_naming_a_delete_it_never_runs_passes(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " note text := 'DELETE FROM legacy rows is handled by the application';\n" + "BEGIN\n" + " RAISE NOTICE '%', note;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_do_body_in_single_quotes_is_scanned(self, tmp_path): + sql = "DO 'BEGIN UPDATE \"Foo\" SET \"a\" = 1; END';" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_a_quoted_do_body_with_a_language_clause_is_scanned(self, tmp_path): + sql = "DO LANGUAGE plpgsql 'BEGIN DELETE FROM \"Foo\"; END';" + assert _keywords(tmp_path, sql) == ("DELETE",) + + def test_a_quoted_do_body_holding_only_ddl_passes(self, tmp_path): + sql = "DO 'BEGIN ALTER TABLE \"Foo\" ADD COLUMN \"b\" TEXT; END';" + assert _keywords(tmp_path, sql) == () + + def test_a_marker_exempts_a_quoted_do_body(self, tmp_path): + sql = ( + "-- data-migration-ok: one config row, keyed by its primary key\n" + "DO 'BEGIN UPDATE \"Config\" SET \"v\" = 1 WHERE \"k\" = ''rev''; END';" + ) + assert _keywords(tmp_path, sql) == () + + def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_fragment(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'UPDATE ' || quote_ident('Foo') || ' SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + def test_concatenated_sql_is_flagged_when_the_keyword_leads_a_later_fragment(self, tmp_path): + sql = "DO $$\nBEGIN\n EXECUTE 'WITH x AS (SELECT 1) ' || 'UPDATE \"Foo\" SET \"a\" = 1';\nEND $$;" + assert _keywords(tmp_path, sql) == ("UPDATE",) + + def test_only_the_variable_that_is_executed_is_read_as_sql(self, tmp_path): + sql = ( + "DO $$\n" + "DECLARE\n" + " msg text := 'UPDATE of legacy rows skipped';\n" + " stmt text := 'DELETE FROM \"Foo\" WHERE \"a\" = 1';\n" + "BEGIN\n" + " RAISE NOTICE '%', msg;\n" + " EXECUTE stmt;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("DELETE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marker_on_an_execute_covers_its_single_quoted_payload(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE ' -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " ';\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + def test_a_marker_on_an_execute_does_not_reach_into_a_dollar_quoted_payload(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE $x$ -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " $x$;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 4 + + def test_a_marker_inside_a_dollar_quoted_payload_exempts_its_rewrite(self, tmp_path): + sql = ( + "DO $$\n" + "BEGIN\n" + " EXECUTE $x$\n" + " -- data-migration-ok: bounded to one row\n" + ' UPDATE "Foo" SET "a" = 1;\n' + " $x$;\n" + "END $$;" + ) + assert _keywords(tmp_path, sql) == () + + +class TestExplain: + def test_explain_analyze_over_an_update_is_flagged(self, tmp_path): + sql = 'EXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 1 + + def test_explain_analyze_verbose_over_a_delete_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN ANALYZE VERBOSE DELETE FROM "Foo";') == ("DELETE",) + + def test_explain_with_a_parenthesised_analyze_is_flagged(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN (ANALYZE, BUFFERS) UPDATE "Foo" SET "a" = 1;') == ("UPDATE",) + + def test_explain_analyze_over_an_insert_select_is_flagged(self, tmp_path): + sql = 'EXPLAIN ANALYZE INSERT INTO "Foo" SELECT "a" FROM "Bar";' + assert _keywords(tmp_path, sql) == ("INSERT ... SELECT",) + + def test_explain_analyze_over_a_select_passes(self, tmp_path): + assert _keywords(tmp_path, 'EXPLAIN ANALYZE SELECT * FROM "Foo";') == () + + def test_a_marker_exempts_an_explained_rewrite(self, tmp_path): + sql = '-- data-migration-ok: one config row\nEXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;' + assert _keywords(tmp_path, sql) == () + + def test_an_analyze_of_its_own_passes(self, tmp_path): + assert _keywords(tmp_path, 'ANALYZE "Foo";') == () + + def test_a_vacuum_analyze_passes(self, tmp_path): + assert _keywords(tmp_path, 'VACUUM ANALYZE "Foo";') == () + + def test_an_explained_rewrite_inside_a_block_reports_its_line(self, tmp_path): + sql = 'DO $$\nBEGIN\n EXPLAIN ANALYZE UPDATE "Foo" SET "a" = 1;\nEND $$;' + assert _keywords(tmp_path, sql) == ("UPDATE",) + assert _scan(tmp_path, sql)[0].line == 3 + + +class TestReporting: + def test_line_number_points_at_the_statement_keyword(self, tmp_path): + sql = '-- CreateIndex\nCREATE INDEX "i" ON "Foo"("a");\n\nUPDATE "Foo" SET "a" = 1;' + assert _scan(tmp_path, sql)[0].line == 4 + + def test_render_names_the_migration_and_line(self, tmp_path): + violation = _scan(tmp_path, '\n\nDELETE FROM "Foo";')[0] + rendered = violation.render() + assert "20260101000000_fixture/migration.sql:3" in rendered + assert "DELETE" in rendered + + +class TestGrandfathering: + def test_every_grandfathered_migration_still_violates(self): + for name in sorted(checker.GRANDFATHERED): + directory = checker.MIGRATIONS_DIR / name + assert directory.is_dir(), f"{name} no longer exists; drop it from GRANDFATHERED" + assert checker.scan_migration(directory), f"{name} is clean; drop it from GRANDFATHERED" + + def test_stale_entry_is_reported_when_a_migration_stops_violating(self): + found = {name: () for name in checker.GRANDFATHERED} + assert checker.stale_grandfathers(found) == tuple(sorted(checker.GRANDFATHERED)) + + def test_missing_entry_is_reported(self): + assert checker.stale_grandfathers({}) == tuple(sorted(checker.GRANDFATHERED)) + + def test_no_stale_entries_against_the_real_tree(self): + found = { + path.name: checker.scan_migration(path) + for path in checker.MIGRATIONS_DIR.iterdir() + if (path / "migration.sql").is_file() + } + assert checker.stale_grandfathers(found) == () + + +class TestShippedMigrations: + def test_the_repo_is_clean(self): + assert checker.main() == 0 + + +CLEAN = 'ALTER TABLE "Foo" ADD COLUMN "a" INT;' +DIRTY = 'UPDATE "Foo" SET "a" = 1;' +FIXTURE = "20260101000000_fixture" + + +def _tree(monkeypatch, tmp_path: Path, sql: str, grandfathered: frozenset = frozenset()) -> None: + """Stand a migrations directory holding one fixture migration in for the repo's own. The + root moves with it, since a rendered violation names the migration relative to the root and + the two are read off the same checkout everywhere but here.""" + directory = tmp_path / "migrations" / FIXTURE + directory.mkdir(parents=True) + (directory / "migration.sql").write_text(sql, encoding="utf-8") + monkeypatch.setattr(checker, "REPO_ROOT", tmp_path) + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "migrations") + monkeypatch.setattr(checker, "GRANDFATHERED", grandfathered) + + +class TestExitCode: + def test_a_clean_tree_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN) + assert checker.main() == 0 + + def test_a_violation_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY) + assert checker.main() == 1 + + def test_a_stale_grandfather_alone_fails_the_check(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + assert checker.main() == 1 + + def test_a_grandfathered_violation_passes(self, tmp_path, monkeypatch): + _tree(monkeypatch, tmp_path, DIRTY, frozenset({FIXTURE})) + assert checker.main() == 0 + + def test_a_missing_migrations_directory_is_an_error(self, tmp_path, monkeypatch): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + assert checker.main() == 2 + + def test_the_failure_names_the_migration_the_line_and_the_keyword( + self, tmp_path, monkeypatch, capsys + ): + _tree(monkeypatch, tmp_path, DIRTY) + checker.main() + printed = capsys.readouterr().out + assert f"migrations/{FIXTURE}/migration.sql:1" in printed + assert "UPDATE rewrites existing rows at boot" in printed + assert checker.GUIDANCE in printed + + def test_a_stale_grandfather_is_named(self, tmp_path, monkeypatch, capsys): + _tree(monkeypatch, tmp_path, CLEAN, frozenset({FIXTURE})) + checker.main() + assert f"{FIXTURE}: listed in GRANDFATHERED" in capsys.readouterr().out + + def test_a_missing_directory_is_reported_on_stderr(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr(checker, "MIGRATIONS_DIR", tmp_path / "absent") + checker.main() + assert "migrations directory not found" in capsys.readouterr().err diff --git a/tests/test_litellm/test_e2e_egress_sentinel.py b/tests/test_litellm/test_e2e_egress_sentinel.py new file mode 100644 index 00000000000..dcbd3241686 --- /dev/null +++ b/tests/test_litellm/test_e2e_egress_sentinel.py @@ -0,0 +1,172 @@ +"""Tests for .github/scripts/e2e_egress_sentinel.py. + +The replay lane's zero-egress proof is only as good as this sentinel: it pins the +provider hosts to a local sink and counts every connection that reaches them, so +a single escaped provider call turns the run red. The contract locked in here is +that the counter counts (each accepted connection is exactly one recorded hit), +that ``assert-empty`` is the pass/fail gate around that count, and that the hosts +file it edits is always handed back exactly as it was found. +""" + +from __future__ import annotations + +import importlib.util +import socket +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Final + +_REPO_ROOT: Final = Path(__file__).resolve().parents[2] +_MODULE_PATH: Final = _REPO_ROOT / ".github" / "scripts" / "e2e_egress_sentinel.py" +_spec: Final = importlib.util.spec_from_file_location("e2e_egress_sentinel", _MODULE_PATH) +sentinel: Final = importlib.util.module_from_spec(_spec) +sys.modules[_spec.name] = sentinel # @dataclass(slots=True) rebuilds via sys.modules +_spec.loader.exec_module(sentinel) + + +def _hit_lines(hits_file: Path) -> list[str]: + if not hits_file.exists(): + return [] + return [line for line in hits_file.read_text(encoding="utf-8").splitlines() if line.strip()] + + +def test_pin_block_lists_every_host_against_the_sink(): + block = sentinel._pin_block("127.0.0.1", ("api.openai.com", "api.anthropic.com")) + assert "127.0.0.1\tapi.openai.com" in block + assert "127.0.0.1\tapi.anthropic.com" in block + assert sentinel._BLOCK_BEGIN in block and sentinel._BLOCK_END in block + + +def test_install_and_restore_round_trips_an_existing_hosts_file(tmp_path): + hosts = tmp_path / "hosts" + original = "127.0.0.1\tlocalhost\n255.255.255.255\tbroadcasthost\n" + hosts.write_text(original, encoding="utf-8") + + saved = sentinel._install_pins(hosts, "127.0.0.1", ("api.openai.com",)) + assert saved == original.encode() + assert "api.openai.com" in hosts.read_text(encoding="utf-8") + + sentinel._restore_pins(hosts, saved) + assert hosts.read_text(encoding="utf-8") == original + + +def test_install_on_a_missing_hosts_file_creates_then_restore_empties(tmp_path): + hosts = tmp_path / "hosts" + saved = sentinel._install_pins(hosts, "127.0.0.1", ("api.anthropic.com",)) + assert saved == b"" + assert "api.anthropic.com" in hosts.read_text(encoding="utf-8") + + sentinel._restore_pins(hosts, saved) + assert hosts.read_text(encoding="utf-8") == "" + + +def test_assert_empty_passes_when_no_calls(tmp_path): + absent = tmp_path / "absent.jsonl" + assert sentinel.assert_empty(absent) == 0 + + empty = tmp_path / "empty.jsonl" + empty.write_text("\n \n", encoding="utf-8") + assert sentinel.assert_empty(empty) == 0 + + +def test_assert_empty_fails_when_calls_recorded(tmp_path): + hits = tmp_path / "hits.jsonl" + hits.write_text('{"port": 443, "peer": ["127.0.0.1", 5]}\n', encoding="utf-8") + assert sentinel.assert_empty(hits) == 1 + + +def test_accept_loop_records_exactly_one_hit_per_connection(tmp_path): + hits_file = tmp_path / "hits.jsonl" + hits_file.write_text("", encoding="utf-8") + listener = sentinel._bind("127.0.0.1", 0) + port = listener.getsockname()[1] + stop = threading.Event() + hits = sentinel._HitLog(path=hits_file, _lock=threading.Lock()) + worker = threading.Thread(target=sentinel._serve_socket, args=(listener, port, hits, stop), daemon=True) + worker.start() + + try: + for _ in range(3): + conn = socket.create_connection(("127.0.0.1", port), timeout=2) + conn.close() + deadline = time.time() + 3 + while time.time() < deadline and len(_hit_lines(hits_file)) < 3: + time.sleep(0.02) + finally: + stop.set() + listener.close() + worker.join(timeout=3) + + assert len(_hit_lines(hits_file)) == 3 + + +def test_serve_end_to_end_pins_counts_and_restores(tmp_path): + hosts = tmp_path / "hosts" + hosts.write_text("127.0.0.1\tlocalhost\n", encoding="utf-8") + hits = tmp_path / "hits.jsonl" + ready = tmp_path / "ready" + pidf = tmp_path / "pid" + port = _free_port() + + proc = subprocess.Popen( + [ + sys.executable, + str(_MODULE_PATH), + "serve", + "--host", + "api.openai.com", + "--host", + "api.anthropic.com", + "--sink-address", + "127.0.0.1", + "--port", + str(port), + "--hits-file", + str(hits), + "--ready-file", + str(ready), + "--pid-file", + str(pidf), + "--hosts-file", + str(hosts), + ], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + try: + _wait_for(lambda: ready.exists(), timeout=10) + pinned = hosts.read_text(encoding="utf-8") + assert "api.openai.com" in pinned and "api.anthropic.com" in pinned + + for _ in range(2): + conn = socket.create_connection(("127.0.0.1", port), timeout=2) + conn.close() + _wait_for(lambda: len(_hit_lines(hits)) >= 2, timeout=5) + assert sentinel.assert_empty(hits) == 1 + finally: + proc.terminate() + proc.wait(timeout=10) + + assert hosts.read_text(encoding="utf-8") == "127.0.0.1\tlocalhost\n" + assert not ready.exists() + + +def _free_port() -> int: + probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + probe.bind(("127.0.0.1", 0)) + port = probe.getsockname()[1] + probe.close() + return port + + +def _wait_for(predicate, *, timeout: float) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + if predicate(): + return + time.sleep(0.05) + raise AssertionError("condition never became true within the timeout") diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 99b1cc826aa..8f2b06be4b3 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -2944,3 +2944,16 @@ def test_a_stream_that_reported_no_usage_is_still_billed(local_cost_map): assert cost == pytest.approx( _priced_at(rebuilt.usage.prompt_tokens, rebuilt.usage.completion_tokens) ) + + +@pytest.mark.asyncio +async def test_acompletion_resolves_provider_from_api_base(): + response = await litellm.acompletion( + model="deepseek-chat", + api_base="https://api.deepseek.com/v1", + api_key="fake-key", + messages=[{"role": "user", "content": "hi"}], + mock_response="resolved", + ) + + assert response.choices[0].message.content == "resolved" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 910b874c2ac..0b82ab971bf 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8921,3 +8921,306 @@ class TestAzureBaseModelFallbackLogging: deployment=None, received_model_name="my-group", id="azure-base-model-test-id" ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] + +def test_model_group_info_intersects_supported_reasoning_efforts(): + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/mini-like"}, + "model_info": {"id": "mini-like-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_xhigh_reasoning_effort": True, + "supports_max_reasoning_effort": True, + } + return { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": False, + "supports_minimal_reasoning_effort": True, + "supports_xhigh_reasoning_effort": False, + } + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + # opus-like offers all seven levels, mini-like lacks none/xhigh/max; only the common set survives, + # so the group never advertises an effort routing could hand to a deployment that rejects it. + assert result.supported_reasoning_efforts == ("minimal", "low", "medium", "high") + + +def test_model_group_info_reasoning_efforts_ignore_a_deployment_off_the_map(): + """The router fills every ModelInfo key, so a deployment absent from the model map arrives with + supports_reasoning None rather than with the key missing. Its synthesized entry carries no mode, + which is what separates it from a mapped non-reasoning model, and nothing being known about it is + no reason to drop the levels the rest of the group agrees on.""" + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/unmapped-model"}, + "model_info": {"id": "unmapped-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_max_reasoning_effort": True, + } + return {"key": model_name, "litellm_provider": "openai", "mode": None, "supports_reasoning": None} + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high", "max") + + +def test_model_group_info_reasoning_efforts_empty_on_a_mapped_non_reasoning_deployment(): + """A group mixing a reasoning model with one the map knows is not a reasoning model shares no + level, so it advertises none and the picker offers nothing rather than a level routing would + hand to a deployment that rejects it.""" + router = litellm.Router( + model_list=[ + { + "model_name": "mixed-group", + "litellm_params": {"model": "anthropic/opus-like"}, + "model_info": {"id": "opus-like-deployment"}, + }, + { + "model_name": "mixed-group", + "litellm_params": {"model": "openai/plain-chat"}, + "model_info": {"id": "plain-chat-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + if model_id == "opus-like-deployment": + return { + "key": model_name, + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": True, + "supports_max_reasoning_effort": True, + } + return {"key": model_name, "litellm_provider": "openai", "mode": "chat", "supports_reasoning": None} + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="mixed-group", + user_facing_model_group_name="mixed-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == () + + +def test_model_group_info_reasoning_efforts_ignore_a_value_declared_in_model_info(): + """The group's levels are computed from its deployments, so a value an operator left in one + deployment's model_info must not seed them. Seeding let the first deployment read narrow the + whole group while the same value on any other deployment was silently ignored.""" + router = litellm.Router( + model_list=[ + { + "model_name": "declared-group", + "litellm_params": {"model": "openai/first-reasoner"}, + "model_info": {"id": "first-deployment"}, + }, + { + "model_name": "declared-group", + "litellm_params": {"model": "openai/second-reasoner"}, + "model_info": {"id": "second-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + info = { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + } + if model_id == "first-deployment": + info["supported_reasoning_efforts"] = ("high",) + return info + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="declared-group", + user_facing_model_group_name="declared-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") + + +def test_model_group_info_survives_a_junk_typed_operator_effort_value(): + """A deployment's registered model_info reads back with whatever the operator wrote under any + key, so a wrong-typed supported_reasoning_efforts must not fail the group's info. Only the + constructor's trailing override keeps the junk away from ModelGroupInfo validation.""" + router = litellm.Router( + model_list=[ + { + "model_name": "junk-declared-group", + "litellm_params": {"model": "openai/lone-reasoner"}, + "model_info": {"id": "junk-deployment"}, + }, + ] + ) + + def _model_info(model_id: str, model_name: str): + return { + "key": model_name, + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_none_reasoning_effort": True, + "supported_reasoning_efforts": "high", + } + + with patch.object(router, "get_deployment_model_info", side_effect=_model_info): + result = router._set_model_group_info( + model_group="junk-declared-group", + user_facing_model_group_name="junk-declared-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") + + +def test_model_group_info_reasoning_efforts_ignore_a_mode_the_operator_declared(): + """A deployment is registered in the cost map under its own id with whatever model_info the + operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only + a mode the map supplied marks the deployment as known, or an off-map deployment carrying any + mode empties the group it sits in.""" + from litellm.router_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts + + mapped_model = "openai/gpt-5.6-sol" + expected = resolve_supported_reasoning_efforts( + litellm.get_model_info(model=mapped_model), + deployment_is_mapped=True, + ) + assert expected + + router = litellm.Router( + model_list=[ + { + "model_name": "smart-group", + "litellm_params": {"model": mapped_model, "api_key": "sk-fake"}, + "model_info": {"id": "mapped-deployment"}, + }, + { + "model_name": "smart-group", + "litellm_params": {"model": "openai/a-model-the-map-never-heard-of", "api_key": "sk-fake"}, + "model_info": {"id": "off-map-deployment", "mode": "chat"}, + }, + ] + ) + + result = router._set_model_group_info( + model_group="smart-group", + user_facing_model_group_name="smart-group", + ) + + assert result is not None + assert result.supported_reasoning_efforts == expected + + +class TestAddDeploymentApiBaseProviderResolution: + def test_bare_model_with_known_api_base_initializes(self): + router = litellm.Router( + model_list=[ + { + "model_name": "groq-pinned", + "litellm_params": { + "model": "llama-3.3-70b-versatile", + "api_base": "https://api.groq.com/openai/v1", + "api_key": "fake-key", + }, + }, + { + "model_name": "deepseek-pinned", + "litellm_params": { + "model": "deepseek-chat", + "api_base": "https://api.deepseek.com/v1", + "api_key": "fake-key", + }, + }, + ] + ) + + model_list = router.get_model_list() + assert model_list is not None + assert {m["model_name"] for m in model_list} == {"groq-pinned", "deepseek-pinned"} + + def test_bare_model_with_unknown_api_base_still_raises(self): + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + litellm.Router( + model_list=[ + { + "model_name": "mystery", + "litellm_params": { + "model": "some-unknown-model", + "api_base": "https://llm.internal.example.com/v1", + "api_key": "fake-key", + }, + } + ] + ) + + def test_explicit_custom_llm_provider_beats_api_base_endpoint_match(self): + router = litellm.Router( + model_list=[ + { + "model_name": "openai-via-gateway", + "litellm_params": { + "model": "gpt-3.5-turbo", + "custom_llm_provider": "openai", + "api_base": "https://api.groq.com/openai/v1", + "api_key": "fake-key", + }, + } + ] + ) + + deployment = router.get_deployment_by_model_group_name("openai-via-gateway") + assert deployment is not None + assert deployment.litellm_params.custom_llm_provider == "openai" diff --git a/tests/test_litellm/test_video_generation.py b/tests/test_litellm/test_video_generation.py index bc23b586279..b166e902d6e 100644 --- a/tests/test_litellm/test_video_generation.py +++ b/tests/test_litellm/test_video_generation.py @@ -2109,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", @@ -2120,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", @@ -2135,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.""" @@ -2160,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", @@ -2170,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: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..7165bb53812 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22805 + "limit": 22804 }, "LIT002": { - "limit": 26873 + "limit": 26872 }, "LIT003": { "limit": 269 @@ -30,7 +30,7 @@ "limit": 16673 }, "LIT011": { - "limit": 5588 + "limit": 5587 }, "LIT012": { "limit": 4510 diff --git a/ui/litellm-dashboard/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)/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)/mcp-servers/_components/mcp_server_edit.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx index 438caa2f5e6..5aec78ba926 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.test.tsx @@ -381,6 +381,44 @@ describe("MCPServerEdit (true passthrough warning)", () => { }); }); +describe("MCPServerEdit (OAuth authorize temp payload)", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("forwards issuer/authorization_url/token_url/registration_url to the temp OAuth session payload", async () => { + // Without these fields the ephemeral server the temp OAuth session endpoint builds has no + // admin-configured OAuth endpoints on it, discovery falls back to (and fails against) the + // plain server url, and Authorize & Fetch Token 400s with "authorization url is not + // configured" even though the saved server (and the visible form) has all four fields filled in. + render( + , + ); + + await waitFor(() => { + expect(mockOauth.getTemporaryPayload).toBeTruthy(); + }); + const payload = mockOauth.getTemporaryPayload!(); + expect(payload).toBeTruthy(); + expect(payload?.issuer).toBe("https://github.com/login/oauth"); + expect(payload?.authorization_url).toBe("https://github.com/login/oauth/authorize"); + expect(payload?.token_url).toBe("https://github.com/login/oauth/access_token"); + expect(payload?.registration_url).toBe("https://github.com/login/oauth/register"); + }); +}); + describe("MCPServerEdit (auth type switch)", () => { beforeEach(() => { vi.clearAllMocks(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx index 5e79b20825c..8793c45371a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_server_edit.tsx @@ -282,6 +282,10 @@ const MCPServerEdit: React.FC = ({ credentials: isClientForwardedTokenMode(values.auth_type) ? preservedAdminCredentials(values.credentials) : values.credentials, + issuer: values.issuer, + authorization_url: values.authorization_url, + token_url: values.token_url, + registration_url: values.registration_url, mcp_access_groups: values.mcp_access_groups || mcpServer.mcp_access_groups, static_headers: staticHeaders, command: values.command, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx index c094cd903fc..d58edd84f91 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/CreateSearchTools.tsx @@ -27,6 +27,7 @@ import { useZodForm } from "@/lib/forms/useZodForm"; import SearchConnectionTest from "./SearchConnectionTest"; import { buildSearchToolPayload } from "./searchToolPayload"; import { AvailableSearchProvider, SearchTool } from "./types"; +import bingLogo from "../../../../../public/assets/logos/bing.png"; import dataforseoLogo from "../../../../../public/assets/logos/dataforseo.png"; import exaAiLogo from "../../../../../public/assets/logos/exa_ai.png"; import googlePseLogo from "../../../../../public/assets/logos/google_pse.png"; @@ -44,6 +45,7 @@ const searchProviderLogoMap: Record = { google_pse: googlePseLogo.src, dataforseo: dataforseoLogo.src, nimble: nimbleLogo.src, + bing_grounding: bingLogo.src, }; interface SearchProviderLabelProps { diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index ef6e521de42..c3fc23034f1 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -16,7 +16,8 @@ import { ClassifierFallback, ClassifierType, ComplexityRouterConfigValue, - DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS, + DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, + MIN_QUOTED_CONTEXT_TURN_CHARS, DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, DEFAULT_CLASSIFIER_FALLBACK, DEFAULT_CLASSIFIER_TIMEOUT_MS, @@ -149,6 +150,8 @@ const ClassificationMethodConfig: React.FC = ({ const classifierModelMissing = showValidationErrors && value.classifier_type === "llm" && !value.classifier_llm_config?.model; const usesCustomPrompt = Boolean(value.classifier_llm_config?.system_prompt?.trim()); + const contextBudget = value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS; + const contextBudgetQuotesNothing = contextBudget > 0 && contextBudget < MIN_QUOTED_CONTEXT_TURN_CHARS; const classificationRubric = value.classifier_llm_config?.classification_rubric ?? DEFAULT_CLASSIFICATION_RUBRIC; const handleClassifierTypeChange = (classifierType: ClassifierType) => { @@ -167,9 +170,9 @@ const ClassificationMethodConfig: React.FC = ({ classifierType === "llm" ? value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE : undefined, - classifier_context_per_turn_chars: + classifier_context_budget_chars: classifierType === "llm" - ? value.classifier_context_per_turn_chars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS + ? value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS : undefined, classifier_context_include_assistant_turns: classifierType === "llm" ? value.classifier_context_include_assistant_turns : undefined, @@ -235,10 +238,10 @@ const ClassificationMethodConfig: React.FC = ({ }); }; - const handleClassifierContextPerTurnCharsChange = (perTurnChars: number | null) => { + const handleClassifierContextBudgetCharsChange = (budgetChars: number | null) => { onChange({ ...value, - classifier_context_per_turn_chars: perTurnChars ?? DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS, + classifier_context_budget_chars: budgetChars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, }); }; @@ -411,17 +414,27 @@ const ClassificationMethodConfig: React.FC = ({
- Context Per-Turn Character Limit + Context Character Budget - handleClassifierContextPerTurnCharsChange(event.target.value === "" ? null : event.target.valueAsNumber) + handleClassifierContextBudgetCharsChange(event.target.value === "" ? null : event.target.valueAsNumber) } - min={1} + min={0} className="w-full" /> - Prior turns longer than this are truncated. + + Total characters of prior conversation sent to the classifier. Turns are taken newest first and quoted + whole while they fit, so a short conversation is never cut. + + {contextBudgetQuotesNothing && ( + + Under {MIN_QUOTED_CONTEXT_TURN_CHARS} characters there is no room to quote a turn that does not already + fit, so a long conversation reaches the classifier with no context at all. Set Context Window Size to 0 + to turn context off deliberately. + + )}
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 0a848ed7deb..45a967c0537 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -8,7 +8,12 @@ vi.mock( ); const mockModelInfo = [ - { model_group: "gpt-4", mode: "chat", supports_reasoning: true }, + { + model_group: "gpt-4", + mode: "chat", + supports_reasoning: true, + supported_reasoning_efforts: ["medium", "high", "xhigh"], + }, { model_group: "gpt-3.5-turbo", mode: "chat" }, { model_group: "claude-3-opus", mode: "chat", supports_reasoning: true }, { model_group: "text-embedding-3-small", mode: "embedding" }, @@ -108,7 +113,7 @@ describe("ComplexityRouterConfig", () => { classifier_type: "llm", classifier_llm_config: { model: "", timeout_ms: 3000, classification_rubric: "agentic" }, classifier_context_window_size: 3, - classifier_context_per_turn_chars: 200, + classifier_context_budget_chars: 8000, }; expect(onChange).toHaveBeenCalledWith(expectedValue); }); @@ -130,11 +135,10 @@ describe("ComplexityRouterConfig", () => { expect(screen.getByDisplayValue("750")).toBeInTheDocument(); expect(screen.getByText("Context Window Size")).toBeInTheDocument(); expect(screen.getByDisplayValue("5")).toBeInTheDocument(); - expect(screen.getByText("Context Per-Turn Character Limit")).toBeInTheDocument(); - expect(screen.getByDisplayValue("400")).toBeInTheDocument(); + expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument(); }); - it("should default classifier context fields to 3 and 200 when llm is selected without explicit values", () => { + it("should default the context window and budget when llm is selected", () => { const llmValue: ComplexityRouterConfigValue = { ...defaultValue, classifier_type: "llm", @@ -147,8 +151,42 @@ describe("ComplexityRouterConfig", () => { const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement; expect(within(windowSizeSection).getByDisplayValue("3")).toBeInTheDocument(); - const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement; - expect(within(perTurnCharsSection).getByDisplayValue("200")).toBeInTheDocument(); + const budgetSection = screen.getByText("Context Character Budget").closest("div") as HTMLElement; + expect(within(budgetSection).getByDisplayValue("8000")).toBeInTheDocument(); + }); + + it("should warn when the budget is too small to quote any turn that does not already fit", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + classifier_context_budget_chars: 50, + }; + renderWithProviders(); + + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + expect(screen.getByText(/no room to quote a turn/i)).toBeInTheDocument(); + }); + + it("should not warn on a budget large enough to quote a turn, nor on a deliberate zero", () => { + for (const budget of [120, 8000, 0]) { + const { unmount } = renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + expect(screen.queryByText(/no room to quote a turn/i)).not.toBeInTheDocument(); + unmount(); + } }); it("should show the assistant-turns switch with its configured value when classifier_type is llm", () => { @@ -229,26 +267,6 @@ describe("ComplexityRouterConfig", () => { }); }); - it("should call onChange with the updated classifier_context_per_turn_chars when edited", () => { - const onChange = vi.fn(); - const llmValue: ComplexityRouterConfigValue = { - ...defaultValue, - classifier_type: "llm", - classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, - }; - renderWithProviders(); - fireEvent.click(screen.getByText("Advanced: Classification Method")); - - const perTurnCharsSection = screen.getByText("Context Per-Turn Character Limit").closest("div") as HTMLElement; - const input = within(perTurnCharsSection).getByRole("spinbutton"); - fireEvent.change(input, { target: { value: "500" } }); - - expect(onChange).toHaveBeenCalledWith({ - ...llmValue, - classifier_context_per_turn_chars: 500, - }); - }); - it("should render the custom technical keywords field", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Classification Method")); @@ -943,3 +961,54 @@ describe("ComplexityRouterConfig reasoning effort gating", () => { ).toHaveTextContent("low"); }); }); + +describe("ComplexityRouterConfig per-model effort filtering", () => { + it("offers only the efforts the model group supports", async () => { + renderWithProviders(); + const user = userEvent.setup(); + await user.click(screen.getByRole("combobox", { name: "Reasoning effort for gpt-4 in the Complex tier" })); + const options = (await screen.findAllByRole("option")).map((option) => option.textContent); + expect(options).toEqual(["Default", "medium", "high", "xhigh"]); + }); + + it("falls back to every effort when the group only reports supports_reasoning", async () => { + renderWithProviders(); + const user = userEvent.setup(); + await user.click( + screen.getByRole("combobox", { name: "Reasoning effort for claude-3-opus in the Reasoning tier" }), + ); + const options = (await screen.findAllByRole("option")).map((option) => option.textContent); + expect(options).toEqual(["Default", "none", "minimal", "low", "medium", "high", "xhigh"]); + }); + + // An empty list is the group's own answer that its deployments share no level, which is different + // from the field being absent, so the control is dropped rather than falling back to every level. + it("offers no effort at all when the group intersects to nothing", () => { + renderWithProviders( + model.model_group !== "claude-3-opus"), + { model_group: "claude-3-opus", mode: "chat", supports_reasoning: true, supported_reasoning_efforts: [] }, + ]} + />, + ); + expect( + screen.queryByRole("combobox", { name: "Reasoning effort for claude-3-opus in the Reasoning tier" }), + ).not.toBeInTheDocument(); + }); + + // Hand-authored configs can carry a level outside the supported set (e.g. max); it must render + // and stay clearable rather than being masked as Default. + it("keeps showing a stored effort outside the supported set", () => { + renderWithProviders( + , + ); + expect(screen.getByRole("combobox", { name: "Reasoning effort for gpt-4 in the Complex tier" })).toHaveTextContent( + "max", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index fc731e2c77f..78e5b572492 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -13,6 +13,7 @@ import { ModelGroup } from "@/components/llm_calls/fetch_models"; import AdaptiveRoutingConfig from "./AdaptiveRoutingConfig"; import ClassificationMethodConfig from "./ClassificationMethodConfig"; import { + REASONING_EFFORT_OPTIONS, ReasoningEffort, TierModelParamsByTier, pruneTierModelParams, @@ -31,7 +32,8 @@ export type { DimensionWeights, TierBoundaries, TokenThresholds }; export const DEFAULT_CLASSIFIER_TIMEOUT_MS = 3000; export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5; export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3; -export const DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS = 200; +export const DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS = 8000; +export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120; export const DEFAULT_SESSION_AFFINITY = false; export const DEFAULT_DEPLOYMENT_AFFINITY = true; @@ -137,6 +139,7 @@ export interface ComplexityRouterConfigValue { classifier_type: ClassifierType; classifier_llm_config?: ClassifierLLMConfig; classifier_context_window_size?: number; + classifier_context_budget_chars?: number; classifier_context_per_turn_chars?: number; classifier_context_include_assistant_turns?: boolean; classifier_fallback?: ClassifierFallback; @@ -250,11 +253,16 @@ const ComplexityRouterConfig: React.FC = ({ const derivedDefaultModel = resolveComplexityDefaultModel(value.tiers); const defaultModel = resolveComplexityDefaultModel(value.tiers, value.default_model); - // Embedding models can't serve a chat-completion role, so they're excluded here. - const reasoningModels = new Set( - modelInfo.filter((model) => model.supports_reasoning).map((model) => model.model_group), + // An absent list means the proxy does not send the field yet, so every level is offered as before. + // An empty list is the group's own answer that its deployments share no level, and is left empty. + const effortOptionsByModel: Record = Object.fromEntries( + modelInfo.map((model) => [ + model.model_group, + model.supported_reasoning_efforts ?? (model.supports_reasoning ? [...REASONING_EFFORT_OPTIONS] : []), + ]), ); + // Embedding models can't serve a chat-completion role, so they're excluded here. const modelOptions = modelInfo .filter((model) => model.mode !== "embedding") .map((model) => ({ @@ -365,7 +373,7 @@ const ComplexityRouterConfig: React.FC = ({ handleTierModelEffortChange(tier, model, effort)} /> diff --git a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.test.tsx b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.test.tsx new file mode 100644 index 00000000000..218a18b805e --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.test.tsx @@ -0,0 +1,78 @@ +import { describe, expect, it } from "vitest"; +import { tierEffortRows } from "./TierModelEffortRows"; + +describe("tierEffortRows", () => { + it("offers the levels the proxy reports for the model", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low", "medium", "high"] }, + paramsByModel: undefined, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: undefined, options: ["low", "medium", "high"] }]); + }); + + it("keeps a row whose model reports no level but already has an effort stored, so it can be cleared", () => { + const rows = tierEffortRows({ + models: ["off-map-model"], + effortOptionsByModel: {}, + paramsByModel: { "off-map-model": { reasoning_effort: "high" } }, + }); + + expect(rows).toEqual([{ model: "off-map-model", effort: "high", options: ["high"] }]); + }); + + it("drops a row whose model reports no level and has nothing stored", () => { + const rows = tierEffortRows({ + models: ["plain-chat-model"], + effortOptionsByModel: { "plain-chat-model": [] }, + paramsByModel: { "plain-chat-model": { temperature: 0.5 } }, + }); + + expect(rows).toEqual([]); + }); + + it("lists a stored level the model no longer reports without duplicating the ones it does", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low", "high"] }, + paramsByModel: { "gpt-5-mini": { reasoning_effort: "xhigh" } }, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: "xhigh", options: ["low", "high", "xhigh"] }]); + }); + + it("does not repeat a stored level the model already reports", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low", "high"] }, + paramsByModel: { "gpt-5-mini": { reasoning_effort: "high" } }, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: "high", options: ["low", "high"] }]); + }); + + it.each([ + ["an unset key", {}], + ["an explicit null", { reasoning_effort: null }], + ["an empty string", { reasoning_effort: "" }], + ])("reads %s as no stored effort", (_label, params) => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low"] }, + paramsByModel: { "gpt-5-mini": params }, + }); + + expect(rows[0].effort).toBeUndefined(); + }); + + it("renders a non-string stored value as a string so the select can show and clear it", () => { + const rows = tierEffortRows({ + models: ["gpt-5-mini"], + effortOptionsByModel: { "gpt-5-mini": ["low"] }, + paramsByModel: { "gpt-5-mini": { reasoning_effort: 3 } }, + }); + + expect(rows).toEqual([{ model: "gpt-5-mini", effort: "3", options: ["low", "3"] }]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx index 67f583bd894..ec9705b9451 100644 --- a/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx +++ b/ui/litellm-dashboard/src/components/add_model/TierModelEffortRows.tsx @@ -2,35 +2,58 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@ import { SimpleTooltip } from "@/components/ui/tooltip"; import { Info } from "lucide-react"; import React from "react"; -import { REASONING_EFFORT_OPTIONS, ReasoningEffort, TierModelParams } from "./complexity_router_tiers"; +import { ReasoningEffort, TierModelParams } from "./complexity_router_tiers"; const PROVIDER_DEFAULT = "__provider_default__"; -const asEffort = (params: TierModelParams | undefined): ReasoningEffort | undefined => { +const storedEffort = (params: TierModelParams | undefined): ReasoningEffort | undefined => { const stored = params?.reasoning_effort; - if (typeof stored !== "string") return undefined; - return REASONING_EFFORT_OPTIONS.find((option) => option === stored); + if (stored === undefined || stored === null || stored === "") return undefined; + return typeof stored === "string" ? stored : String(stored); }; interface TierModelEffortRowsProps { tierLabel: string; models: string[]; - reasoningModels: ReadonlySet; + effortOptionsByModel: Record; paramsByModel: Record | undefined; onEffortChange: (model: string, effort: ReasoningEffort | undefined) => void; } +export interface TierEffortRow { + model: string; + effort: ReasoningEffort | undefined; + options: string[]; +} + +/** + * A stored effort outside the model's supported set (hand-authored, or capabilities changed since + * it was saved) is listed anyway, so the row renders with its value selected and can be cleared. + * Only a model with no supported level and nothing stored drops out. + */ +export const tierEffortRows = ({ + models, + effortOptionsByModel, + paramsByModel, +}: Pick): TierEffortRow[] => + models + .map((model) => { + const effort = storedEffort(paramsByModel?.[model]); + const supported = effortOptionsByModel[model] ?? []; + const listed = effort !== undefined && !supported.includes(effort) ? [...supported, effort] : supported; + return { model, effort, options: Array.from(new Set(listed)) }; + }) + .filter(({ options }) => options.length > 0); + const TierModelEffortRows: React.FC = ({ tierLabel, models, - reasoningModels, + effortOptionsByModel, paramsByModel, onEffortChange, }) => { - const shown = models.filter( - (model) => reasoningModels.has(model) || Object.keys(paramsByModel?.[model] ?? {}).length > 0, - ); - if (shown.length === 0) return null; + const rows = tierEffortRows({ models, effortOptionsByModel, paramsByModel }); + if (rows.length === 0) return null; return (
@@ -41,18 +64,17 @@ const TierModelEffortRows: React.FC = ({
- {shown.map((model) => ( + {rows.map(({ model, effort, options }) => (
{model}