From 4a646dd9a0d7acd2ea7c3fbd57de1e17ead7cec8 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 14:53:40 -0700 Subject: [PATCH 01/63] ci(e2e): run a PR's changed e2e tests three times behind a human-approved environment Adds a required-check candidate that selects the tests/e2e test files a PR added or modified, boots a stage-mirror stack on the runner (migrations, backend, two gateway processes behind nginx, Postgres, Jaeger, TLS cluster Valkey), and runs those files three times with retries off. The run job sits behind the e2e-changed GitHub environment, so a reviewer approves each run before the OIDC token that reads the provider keys from AWS Secrets Manager exists. Supersedes #34981 --- .github/e2e-stack/down.sh | 17 ++ .github/e2e-stack/secrets_to_env.py | 28 +++ .github/e2e-stack/up.sh | 198 +++++++++++++++++++ .github/workflows/test-e2e-changed.yml | 171 ++++++++++++++++ tests/e2e/CONTRIBUTING.md | 4 + tests/e2e/gateway/stage_mirror_ci_config.yml | 63 ++++++ 6 files changed, 481 insertions(+) create mode 100755 .github/e2e-stack/down.sh create mode 100644 .github/e2e-stack/secrets_to_env.py create mode 100755 .github/e2e-stack/up.sh create mode 100644 .github/workflows/test-e2e-changed.yml create mode 100644 tests/e2e/gateway/stage_mirror_ci_config.yml diff --git a/.github/e2e-stack/down.sh b/.github/e2e-stack/down.sh new file mode 100755 index 00000000000..9f72f2d6e64 --- /dev/null +++ b/.github/e2e-stack/down.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +set -uo pipefail + +STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}" + +for pid_file in "${STACK_DIR}"/pids/*.pid; do + [[ -f "${pid_file}" ]] || continue + pkill -TERM -P "$(cat "${pid_file}")" 2>/dev/null + kill -TERM "$(cat "${pid_file}")" 2>/dev/null + rm -f "${pid_file}" +done + +for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do + docker rm -f "${container}" >/dev/null 2>&1 +done + +exit 0 diff --git a/.github/e2e-stack/secrets_to_env.py b/.github/e2e-stack/secrets_to_env.py new file mode 100644 index 00000000000..25b22555c29 --- /dev/null +++ b/.github/e2e-stack/secrets_to_env.py @@ -0,0 +1,28 @@ +import sys +from pathlib import Path + +from pydantic import TypeAdapter + +secrets_adapter: TypeAdapter[dict[str, str]] = TypeAdapter(dict[str, str]) + + +def main() -> int: + env_path = Path(sys.argv[1]) + secrets = secrets_adapter.validate_json(sys.stdin.read()) + unwritable = tuple( + key for key, value in secrets.items() if "'" in value or "\n" in value or "\r" in value + ) + if unwritable: + _ = sys.stderr.write(f"values contain characters unsafe for both bash and dotenv: {', '.join(unwritable)}\n") + return 1 + lines = tuple(f"{key}='{value}'" for key, value in secrets.items() if value) + with env_path.open("a") as handle: + _ = handle.write("\n".join(lines) + "\n") + for value in secrets.values(): + if value: + _ = sys.stdout.write(f"::add-mask::{value}\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/e2e-stack/up.sh b/.github/e2e-stack/up.sh new file mode 100755 index 00000000000..06fce35a896 --- /dev/null +++ b/.github/e2e-stack/up.sh @@ -0,0 +1,198 @@ +#!/usr/bin/env bash +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}" +CERTS_DIR="${STACK_DIR}/certs" +LOGS_DIR="${STACK_DIR}/logs" +PIDS_DIR="${STACK_DIR}/pids" + +POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}" +VALKEY_IMAGE="${E2E_VALKEY_IMAGE:-valkey/valkey:8.1}" +JAEGER_IMAGE="${E2E_JAEGER_IMAGE:-jaegertracing/jaeger:2.10.0}" +NGINX_IMAGE="${E2E_NGINX_IMAGE:-nginx:1.29-alpine}" + +LB_PORT="${E2E_LB_PORT:-4000}" +GATEWAY_PORT_1="${E2E_GATEWAY_PORT_1:-4010}" +GATEWAY_PORT_2="${E2E_GATEWAY_PORT_2:-4011}" +BACKEND_PORT="${E2E_BACKEND_PORT:-4001}" +REDIS_PORT="${E2E_REDIS_PORT:-6379}" +DATABASE_HOST="${E2E_DATABASE_HOST:-127.0.0.1}" +DATABASE_PORT="${E2E_DATABASE_PORT:-5432}" +DATABASE_USER="${E2E_DATABASE_USER:-litellm}" +DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}" +DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}" +JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}" +JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}" + +MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}" + +mkdir -p "${CERTS_DIR}" "${LOGS_DIR}" "${PIDS_DIR}" + +log() { printf 'e2e-stack: %s\n' "$*"; } + +port_open() { (exec 3<>"/dev/tcp/127.0.0.1/$1") 2>/dev/null; } + +wait_for() { + local label="$1" check="$2" deadline=$((SECONDS + ${3:-120})) + until eval "${check}"; do + if ((SECONDS >= deadline)); then + log "timed out waiting for ${label}" + tail -n 60 "${LOGS_DIR}"/*.log 2>/dev/null || true + exit 1 + fi + sleep 2 + done + log "${label} is up" +} + +if [[ -f "${REPO_ROOT}/tests/e2e/.env" ]]; then + set -a + source "${REPO_ROOT}/tests/e2e/.env" + set +a +fi + +if ! port_open "${DATABASE_PORT}"; then + docker run -d --name e2e-postgres -p "${DATABASE_PORT}:5432" \ + -e "POSTGRES_USER=${DATABASE_USER}" -e "POSTGRES_PASSWORD=${DATABASE_PASSWORD}" -e "POSTGRES_DB=${DATABASE_NAME}" \ + "${POSTGRES_IMAGE}" >/dev/null +fi +wait_for "postgres" "port_open ${DATABASE_PORT}" + +if ! port_open "${JAEGER_QUERY_PORT}"; then + docker run -d --name e2e-jaeger -p "${JAEGER_OTLP_PORT}:4318" -p "${JAEGER_QUERY_PORT}:16686" \ + "${JAEGER_IMAGE}" >/dev/null +fi +wait_for "jaeger" "curl -fs http://127.0.0.1:${JAEGER_QUERY_PORT}/api/services >/dev/null" + +openssl genrsa -out "${CERTS_DIR}/ca.key" 2048 2>/dev/null +openssl req -x509 -new -nodes -key "${CERTS_DIR}/ca.key" -sha256 -days 7 \ + -subj "/CN=litellm-e2e-ca" \ + -addext "basicConstraints=critical,CA:TRUE" -addext "keyUsage=critical,keyCertSign,cRLSign" \ + -out "${CERTS_DIR}/ca.crt" 2>/dev/null +openssl genrsa -out "${CERTS_DIR}/server.key" 2048 2>/dev/null +openssl req -new -key "${CERTS_DIR}/server.key" -subj "/CN=localhost" -out "${CERTS_DIR}/server.csr" 2>/dev/null +openssl x509 -req -in "${CERTS_DIR}/server.csr" -CA "${CERTS_DIR}/ca.crt" -CAkey "${CERTS_DIR}/ca.key" \ + -CAcreateserial -days 7 -sha256 \ + -extfile <(printf 'basicConstraints=CA:FALSE\nkeyUsage=critical,digitalSignature,keyEncipherment\nextendedKeyUsage=serverAuth\nsubjectAltName=DNS:localhost,IP:127.0.0.1\n') \ + -out "${CERTS_DIR}/server.crt" 2>/dev/null +chmod 644 "${CERTS_DIR}"/*.key "${CERTS_DIR}"/*.crt + +CERTIFI_BUNDLE="$(cd "${REPO_ROOT}" && uv run --no-sync python -c 'import certifi; print(certifi.where())')" +cat "${CERTIFI_BUNDLE}" "${CERTS_DIR}/ca.crt" > "${CERTS_DIR}/ca-bundle.pem" + +docker rm -f e2e-valkey >/dev/null 2>&1 || true +docker run -d --name e2e-valkey -p "${REDIS_PORT}:${REDIS_PORT}" -v "${CERTS_DIR}:/certs:ro" \ + "${VALKEY_IMAGE}" valkey-server \ + --cluster-enabled yes --port 0 --tls-port "${REDIS_PORT}" \ + --tls-cert-file /certs/server.crt --tls-key-file /certs/server.key --tls-ca-cert-file /certs/ca.crt \ + --tls-auth-clients no --cluster-announce-ip 127.0.0.1 >/dev/null +VALKEY_CLI="docker exec e2e-valkey valkey-cli --tls --cacert /certs/ca.crt -h 127.0.0.1 -p ${REDIS_PORT}" +wait_for "valkey" "${VALKEY_CLI} ping 2>/dev/null | grep -q PONG" +${VALKEY_CLI} cluster addslotsrange 0 16383 >/dev/null +wait_for "valkey cluster" "${VALKEY_CLI} cluster info 2>/dev/null | grep -q cluster_state:ok" + +CONFIG_SOURCE="${REPO_ROOT}/tests/e2e/gateway/stage_mirror_ci_config.yml" +CONFIG_PATH="${CONFIG_SOURCE}" +if [[ "${REDIS_PORT}" != "6379" ]]; then + CONFIG_PATH="${STACK_DIR}/litellm-config.yml" + sed "s/port: 6379/port: ${REDIS_PORT}/" "${CONFIG_SOURCE}" > "${CONFIG_PATH}" +fi + +SERVER_ENV=( + "LITELLM_MASTER_KEY=${MASTER_KEY}" + "DATABASE_HOST=${DATABASE_HOST}" + "DATABASE_PORT=${DATABASE_PORT}" + "DATABASE_USER=${DATABASE_USER}" + "DATABASE_PASSWORD=${DATABASE_PASSWORD}" + "DATABASE_NAME=${DATABASE_NAME}" + "DISABLE_SCHEMA_UPDATE=true" + "REDIS_HOST=127.0.0.1" + "REDIS_PORT=${REDIS_PORT}" + "REDIS_CLUSTER_NODES=[{\"host\":\"127.0.0.1\",\"port\":${REDIS_PORT}}]" + "CONFIG_FILE_PATH=${CONFIG_PATH}" + "STORE_MODEL_IN_DB=True" + "OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf" + "OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}" + "SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem" + "PYTHONPATH=${REPO_ROOT}" +) +if [[ -n "${VERTEXAI_CREDENTIALS:-}" ]]; then + printf '%s' "${VERTEXAI_CREDENTIALS}" > "${STACK_DIR}/vertex-adc.json" + SERVER_ENV+=("GOOGLE_APPLICATION_CREDENTIALS=${STACK_DIR}/vertex-adc.json") +fi + +cd "${REPO_ROOT}" + +log "running migrations" +env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/migrations.log" 2>&1 + +start_server() { + local name="$1"; shift + env "${SERVER_ENV[@]}" "$@" >"${LOGS_DIR}/${name}.log" 2>&1 & + echo $! > "${PIDS_DIR}/${name}.pid" +} + +start_server backend uv run --no-sync uvicorn backend.main:app --host 0.0.0.0 --port "${BACKEND_PORT}" +start_server gateway-1 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_1}" +start_server gateway-2 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_2}" + +if [[ "$(uname)" == "Linux" ]]; then + NGINX_UPSTREAM_HOST=127.0.0.1 + NGINX_DOCKER_ARGS=(--network host) +else + NGINX_UPSTREAM_HOST=host.docker.internal + NGINX_DOCKER_ARGS=(-p "${LB_PORT}:${LB_PORT}") +fi + +cat > "${STACK_DIR}/nginx.conf" </dev/null 2>&1 || true +docker run -d --name e2e-nginx "${NGINX_DOCKER_ARGS[@]}" \ + -v "${STACK_DIR}/nginx.conf:/etc/nginx/nginx.conf:ro" "${NGINX_IMAGE}" >/dev/null + +wait_for "backend" "curl -fs http://127.0.0.1:${BACKEND_PORT}/health/liveliness >/dev/null" 300 +wait_for "gateway-1" "curl -fs http://127.0.0.1:${GATEWAY_PORT_1}/health/liveliness >/dev/null" 300 +wait_for "gateway-2" "curl -fs http://127.0.0.1:${GATEWAY_PORT_2}/health/liveliness >/dev/null" 300 +wait_for "load balancer" "curl -fs http://127.0.0.1:${LB_PORT}/health/liveliness >/dev/null" 60 + +cat > "${STACK_DIR}/stack.env" <> "${GITHUB_OUTPUT}" + if [ -n "${tests}" ]; then + echo "any=true" >> "${GITHUB_OUTPUT}" + echo "selected e2e tests: ${tests}" + else + echo "any=false" >> "${GITHUB_OUTPUT}" + echo "no e2e changes; nothing to run" + fi + + run: + name: Run changed e2e tests against the stage-mirror stack + needs: detect + if: needs.detect.outputs.any == 'true' && github.event.pull_request.head.repo.full_name == github.repository + runs-on: ubuntu-latest + timeout-minutes: 90 + environment: e2e-changed + permissions: + contents: read + id-token: write + services: + postgres: + image: postgres:16.6 + env: + POSTGRES_USER: litellm + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U litellm" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + jaeger: + image: jaegertracing/jaeger:2.10.0 + ports: + - 4318:4318 + - 16686:16686 + steps: + - name: Validate configuration + env: + ROLE: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }} + run: test -n "${ROLE}" || { echo "::error::Set repo variable E2E_AWS_ROLE_TO_ASSUME to an OIDC role with read access to the e2e secrets"; exit 1; } + + - name: Checkout + 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.13" + + - 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 \ + --extra proxy --extra proxy-runtime --extra extra_proxy \ + --extra semantic-router --extra bedrock-realtime \ + --group ci --group proxy-dev --group e2e-dev + uv pip install "pipecat-ai[openai]==1.4.0" + + - 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: Install Playwright chromium + run: uv run --no-sync playwright install --with-deps chromium + + - name: Configure AWS credentials + uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0 + with: + role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }} + aws-region: us-east-1 + role-session-name: litellm-e2e-changed-${{ github.run_id }} + + - name: Fetch provider credentials from AWS Secrets Manager + run: | + aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-provider-keys \ + --query SecretString --output text \ + | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env + aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license \ + --query SecretString --output text \ + | jq -R -s '{"LITELLM_LICENSE": rtrimstr("\n")}' \ + | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env + + - name: Boot the stage-mirror stack + run: bash .github/e2e-stack/up.sh + + - name: Export stack environment + run: | + master_key="$(grep '^LITELLM_MASTER_KEY=' "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" | cut -d= -f2-)" + echo "::add-mask::${master_key}" + cat "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" >> "${GITHUB_ENV}" + + - name: Run the selected tests three times with retries off + env: + TESTS: ${{ needs.detect.outputs.tests }} + run: | + read -r -a test_files <<< "${TESTS}" + for pass in 1 2 3; do + echo "::group::pass ${pass} of 3" + set +e + uv run --no-sync pytest "${test_files[@]}" --reruns 0 -v -rA --tb=short -p no:cacheprovider + status=$? + set -e + echo "::endgroup::" + if [ "${status}" = "5" ]; then + echo "selected files collected no runnable tests" + exit 0 + fi + if [ "${status}" != "0" ]; then + echo "::error::pass ${pass} of 3 failed with exit code ${status}" + exit "${status}" + fi + done + + - name: Show stack logs on failure + if: failure() + run: tail -n 200 "${RUNNER_TEMP}/litellm-e2e-stack/logs"/*.log diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 871d6b3904c..e04781ab205 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -52,6 +52,10 @@ The suites run against a live proxy, so bring one up first by running the litell Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy +### The pull request check + +Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, and `load/`, which have their own lanes) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down + ### Record and replay Record/replay scopes to the proxy's provider-bound traffic only. In `E2E_FIXTURE_MODE=record` the harness boots a local provider-edge server, edge-wired tests register their deployments with an `api_base` pointing at it, and every provider call the proxy makes is forwarded verbatim and written to a fixture bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`). `E2E_FIXTURE_MODE=replay` runs the same tests against the same live proxy and database, but the edge answers the proxy's provider calls from the bundle instead of the provider, so the run makes zero provider calls and spends nothing while key auth, routing, cost calculation, and spend-log writes all still execute for real. Unset (or `live`) behaves exactly as before the knob existed. Both record and replay need the proxy up; only the provider is taken out of the loop diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml new file mode 100644 index 00000000000..57a92fd47fb --- /dev/null +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -0,0 +1,63 @@ +general_settings: + store_prompts_in_spend_logs: true + database_connection_pool_limit: 10 + forward_client_headers_to_llm_api: false + maximum_spend_logs_retention_period: "60d" + maximum_spend_logs_cleanup_cron: "0 1 * * *" + proxy_budget_rescheduler_min_time: 15 + proxy_budget_rescheduler_max_time: 20 + +litellm_settings: + drop_params: true + default_redis_ttl: 20 + request_timeout: 600 + num_retries: 3 + json_logs: true + store_audit_logs: true + cache: true + cache_params: + type: redis + host: 127.0.0.1 + port: 6379 + redis_startup_nodes: + - host: 127.0.0.1 + port: 6379 + ssl: true + callbacks: ["arize_phoenix", "datadog", "smtp_email", "prometheus", "otel"] + require_auth_for_metrics_endpoint: false + +router_settings: + routing_strategy: simple-shuffle + num_retries: 3 + allowed_fails: 5 + cooldown_time: 30 + +model_list: + - model_name: gpt-5.5 + litellm_params: + model: openai/gpt-5.5 + api_key: os.environ/OPENAI_API_KEY + - model_name: claude-haiku-4-5 + litellm_params: + model: anthropic/claude-haiku-4-5 + api_key: os.environ/ANTHROPIC_API_KEY + - model_name: gemini-2.5-flash-vertex + litellm_params: + model: vertex_ai/gemini-2.5-flash + vertex_project: os.environ/VERTEXAI_PROJECT + vertex_location: us-central1 + vertex_credentials: os.environ/VERTEXAI_CREDENTIALS + - model_name: gemini-2.5-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + - model_name: gemini-2.5-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + +mcp_servers: + devin: + url: "https://mcp.devin.ai/mcp" + auth_type: api_key + auth_value: os.environ/DEVIN_API_KEY From b348ed7f09709647d3f9bce38be3ba49741be307 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 14:56:52 -0700 Subject: [PATCH 02/63] ci(e2e): only tail stack logs when the stack actually booted --- .github/workflows/test-e2e-changed.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index 6ad252e7246..a4466bba3fd 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -136,6 +136,7 @@ jobs: | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env - name: Boot the stage-mirror stack + id: boot run: bash .github/e2e-stack/up.sh - name: Export stack environment @@ -167,5 +168,5 @@ jobs: done - name: Show stack logs on failure - if: failure() + if: failure() && steps.boot.conclusion != 'skipped' run: tail -n 200 "${RUNNER_TEMP}/litellm-e2e-stack/logs"/*.log From 743f94f82abf9d4963c087809fab782d47761b09 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 14:59:15 -0700 Subject: [PATCH 03/63] ci(e2e): fail when nothing collects and ignore lanes with their own checks in the smoke trigger --- .github/workflows/test-e2e-changed.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index a4466bba3fd..7a638801161 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -34,7 +34,7 @@ jobs: | grep -E '^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$' \ | grep -vE '^tests/e2e/(ui|claude_code|load)/' \ | sort -u | tr '\n' ' ' | sed 's/ $//')" || true - if [ -z "${tests}" ] && printf '%s\n' "${files}" | grep -v '^tests/e2e/ui/' \ + if [ -z "${tests}" ] && printf '%s\n' "${files}" | grep -vE '^tests/e2e/(ui|claude_code|load)/' \ | grep -qE '^(tests/e2e/|\.github/e2e-stack/|\.github/workflows/test-e2e-changed\.yml$)'; then tests="${SMOKE_TESTS}" echo "harness or stack changed without a test file; running the smoke suite" @@ -158,8 +158,8 @@ jobs: set -e echo "::endgroup::" if [ "${status}" = "5" ]; then - echo "selected files collected no runnable tests" - exit 0 + echo "::error::the selected files collected no runnable tests, so nothing was verified" + exit 1 fi if [ "${status}" != "0" ]; then echo "::error::pass ${pass} of 3 failed with exit code ${status}" From c62643f0cf72c2c02f08c2073a621e81d906b056 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 15:39:21 -0700 Subject: [PATCH 04/63] ci(e2e): drop trailing newlines from fetched secret values before writing the env --- .github/e2e-stack/secrets_to_env.py | 2 +- .github/workflows/test-e2e-changed.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/e2e-stack/secrets_to_env.py b/.github/e2e-stack/secrets_to_env.py index 25b22555c29..65022931d81 100644 --- a/.github/e2e-stack/secrets_to_env.py +++ b/.github/e2e-stack/secrets_to_env.py @@ -8,7 +8,7 @@ secrets_adapter: TypeAdapter[dict[str, str]] = TypeAdapter(dict[str, str]) def main() -> int: env_path = Path(sys.argv[1]) - secrets = secrets_adapter.validate_json(sys.stdin.read()) + secrets = {key: value.rstrip("\r\n") for key, value in secrets_adapter.validate_json(sys.stdin.read()).items()} unwritable = tuple( key for key, value in secrets.items() if "'" in value or "\n" in value or "\r" in value ) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index 7a638801161..b20c84af02d 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -132,7 +132,7 @@ jobs: | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license \ --query SecretString --output text \ - | jq -R -s '{"LITELLM_LICENSE": rtrimstr("\n")}' \ + | jq -R -s '{"LITELLM_LICENSE": .}' \ | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env - name: Boot the stage-mirror stack From 2bce27cfa6a0277c7237649c1c322883887e781a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 15:52:42 -0700 Subject: [PATCH 05/63] docs(e2e): describe the dedicated, least-privilege, capped credentials behind the pull request check --- tests/e2e/CONTRIBUTING.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index e04781ab205..b984791c738 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -56,6 +56,8 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, and `load/`, which have their own lanes) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down +Credentials: the lane never sees the stage keys. Everything the runner writes into `tests/e2e/.env` comes from two AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`, and the first holds only the names that the tests and `tests/e2e/gateway/stage_mirror_ci_config.yml` read. Each credential in it is dedicated to this lane and can do only what the tests do: the AWS pair belongs to the `litellm-e2e-changed` IAM user, whose inline policy allows Bedrock inference, the one guardrail the tests use, the batch job APIs, the batch S3 bucket, and assuming the batch test's role; the Vertex key belongs to a service account holding only `roles/aiplatform.user` on the Vertex project; the OpenAI, Anthropic, Gemini, Cohere, and Devin keys are per-lane keys with hard monthly spend caps; and the Datadog application key is scoped to log search. A leaked key can therefore spend at most one month of a small capped budget or call a handful of Bedrock APIs, and rotating one is a single `aws secretsmanager put-secret-value` on the secret, with no workflow change + ### Record and replay Record/replay scopes to the proxy's provider-bound traffic only. In `E2E_FIXTURE_MODE=record` the harness boots a local provider-edge server, edge-wired tests register their deployments with an `api_base` pointing at it, and every provider call the proxy makes is forwarded verbatim and written to a fixture bundle (default `tests/e2e/.fixtures`, override with `E2E_FIXTURE_DIR`). `E2E_FIXTURE_MODE=replay` runs the same tests against the same live proxy and database, but the edge answers the proxy's provider calls from the bundle instead of the provider, so the run makes zero provider calls and spends nothing while key auth, routing, cost calculation, and spend-log writes all still execute for real. Unset (or `live`) behaves exactly as before the knob existed. Both record and replay need the proxy up; only the provider is taken out of the loop From 1c4e46f17e813d5eb8e5fa0c2577f7bef3a32b9a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 17:15:36 -0700 Subject: [PATCH 06/63] ci(e2e): reload the stack config every 7s so the harness propagation budget holds, and only mask credential-named values --- .github/e2e-stack/secrets_to_env.py | 7 +++++-- tests/e2e/gateway/stage_mirror_ci_config.yml | 1 + 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/.github/e2e-stack/secrets_to_env.py b/.github/e2e-stack/secrets_to_env.py index 65022931d81..d2d6675d690 100644 --- a/.github/e2e-stack/secrets_to_env.py +++ b/.github/e2e-stack/secrets_to_env.py @@ -1,9 +1,12 @@ +import re import sys from pathlib import Path +from typing import Final from pydantic import TypeAdapter secrets_adapter: TypeAdapter[dict[str, str]] = TypeAdapter(dict[str, str]) +SECRET_NAME: Final = re.compile(r"KEY|SECRET|TOKEN|PASS|CREDENTIAL|LICENSE|AUTH") def main() -> int: @@ -18,8 +21,8 @@ def main() -> int: lines = tuple(f"{key}='{value}'" for key, value in secrets.items() if value) with env_path.open("a") as handle: _ = handle.write("\n".join(lines) + "\n") - for value in secrets.values(): - if value: + for key, value in secrets.items(): + if value and SECRET_NAME.search(key): _ = sys.stdout.write(f"::add-mask::{value}\n") return 0 diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 57a92fd47fb..21664c2d0a1 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -1,4 +1,5 @@ general_settings: + proxy_config_reload_interval_seconds: 7 store_prompts_in_spend_logs: true database_connection_pool_limit: 10 forward_client_headers_to_llm_api: false From 3ead9d16884c652ccbabe618f7526d74b3f4743e Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Fri, 19 Jun 2026 16:43:02 -0500 Subject: [PATCH 07/63] feat(vertex): add Lyria model support --- .../llms/vertex_ai/interactions/__init__.py | 3 + ...odel_prices_and_context_window_backup.json | 81 +++++++++++++++++++ .../vertex_passthrough_logging_handler.py | 69 ++++++++++++++++ model_prices_and_context_window.json | 81 +++++++++++++++++++ ...test_vertex_passthrough_logging_handler.py | 68 ++++++++++++++++ tests/test_litellm/test_utils.py | 37 +++++++++ 6 files changed, 339 insertions(+) create mode 100644 tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py index e69de29bb2d..3e2309d21ef 100644 --- a/litellm/llms/vertex_ai/interactions/__init__.py +++ b/litellm/llms/vertex_ai/interactions/__init__.py @@ -0,0 +1,3 @@ +from .transformation import VertexAIInteractionsConfig + +__all__ = ["VertexAIInteractionsConfig"] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d8a8f84b032..e854e8fc9bb 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45511,6 +45511,87 @@ "output_cost_per_token": 4e-07, "supports_tool_choice": true }, + "vertex_ai/lyria-002": { + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 0.009111111111111111, + "max_audio_per_prompt": 4, + "mode": "chat", + "output_cost_per_second": 0.002, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, + "vertex_ai/lyria-3-clip-preview": { + "input_cost_per_token": 0, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_image": 0.04, + "output_cost_per_token": 0, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "audio" + ], + "supported_regions": [ + "global" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_image_input": true, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": true, + "supports_web_search": false + }, + "vertex_ai/lyria-3-pro-preview": { + "input_cost_per_token": 0, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_image": 0.08, + "output_cost_per_token": 0, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "audio" + ], + "supported_regions": [ + "global" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_image_input": true, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": true, + "supports_web_search": false + }, "vertex_ai/meta/llama-3.1-405b-instruct-maas": { "input_cost_per_token": 5e-06, "litellm_provider": "vertex_ai-llama_models", diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 49ec18013b5..a9f68f8380b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -46,6 +46,7 @@ else: # Define EndpointType locally to avoid import issues EndpointType = Any +_LYRIA_SECONDS_PER_AUDIO_PREDICTION = 30 class VertexPassthroughLoggingHandler: @@ -270,6 +271,16 @@ class VertexPassthroughLoggingHandler: _json_response: Final[dict[str, object]] = httpx_response.json() litellm_prediction_response: ModelResponse | EmbeddingResponse | ImageResponse = ModelResponse() + if VertexPassthroughLoggingHandler._is_lyria_predict_response( + model=model, + json_response=_json_response, + ): + return VertexPassthroughLoggingHandler._handle_lyria_predict_response( + json_response=_json_response, + logging_obj=logging_obj, + model=model, + kwargs=kwargs, + ) if vertex_image_generation_class.is_image_generation_response(_json_response): litellm_prediction_response = vertex_image_generation_class.process_image_generation_response( _json_response, @@ -323,6 +334,64 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } + @staticmethod + def _handle_lyria_predict_response( + json_response: dict, + logging_obj: LiteLLMLoggingObj, + model: str, + kwargs: dict, + ) -> PassThroughEndpointLoggingTypedDict: + prediction_count: Final = ( + VertexPassthroughLoggingHandler._get_lyria_audio_prediction_count( + json_response=json_response + ) + ) + model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}", {}) + response_cost: Final = ( + model_info.get("output_cost_per_second", 0.0) + * _LYRIA_SECONDS_PER_AUDIO_PREDICTION + * prediction_count + ) + + logging_obj.model = model + logging_obj.model_call_details["model"] = model + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" + logging_obj.custom_llm_provider = "vertex_ai" + logging_obj.model_call_details["response_cost"] = response_cost + + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = "vertex_ai" + + standard_pass_through_response_object: Final[StandardPassThroughResponseObject] = { + "response": json_response, + } + return { + "result": standard_pass_through_response_object, + "kwargs": kwargs, + } + + @staticmethod + def _is_lyria_predict_response(model: str, json_response: dict) -> bool: + return ( + model == "lyria-002" + and VertexPassthroughLoggingHandler._get_lyria_audio_prediction_count( + json_response=json_response + ) + > 0 + ) + + @staticmethod + def _get_lyria_audio_prediction_count(json_response: dict) -> int: + predictions: Final = json_response.get("predictions") + if not isinstance(predictions, list): + return 0 + return sum( + 1 + for prediction in predictions + if isinstance(prediction, dict) and prediction.get("audioContent") + ) + @staticmethod def _extract_embed_content_input(request_body: dict | None, batch: bool) -> str: """Extract raw input text from an :embedContent or :batchEmbedContents request body for token counting.""" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d8a8f84b032..e854e8fc9bb 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45511,6 +45511,87 @@ "output_cost_per_token": 4e-07, "supports_tool_choice": true }, + "vertex_ai/lyria-002": { + "litellm_provider": "vertex_ai", + "max_audio_length_hours": 0.009111111111111111, + "max_audio_per_prompt": 4, + "mode": "chat", + "output_cost_per_second": 0.002, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, + "vertex_ai/lyria-3-clip-preview": { + "input_cost_per_token": 0, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_image": 0.04, + "output_cost_per_token": 0, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "audio" + ], + "supported_regions": [ + "global" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_image_input": true, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": true, + "supports_web_search": false + }, + "vertex_ai/lyria-3-pro-preview": { + "input_cost_per_token": 0, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_image": 0.08, + "output_cost_per_token": 0, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "audio" + ], + "supported_regions": [ + "global" + ], + "supports_audio_input": false, + "supports_audio_output": true, + "supports_function_calling": false, + "supports_image_input": true, + "supports_prompt_caching": false, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": true, + "supports_web_search": false + }, "vertex_ai/meta/llama-3.1-405b-instruct-maas": { "input_cost_per_token": 5e-06, "litellm_provider": "vertex_ai-llama_models", diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py new file mode 100644 index 00000000000..7af67cc0795 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -0,0 +1,68 @@ +from datetime import datetime +from unittest.mock import MagicMock + +import httpx +import litellm +import pytest + +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( + VertexPassthroughLoggingHandler, +) + + +def test_lyria_predict_response_preserves_audio_response_and_logs_cost( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/lyria-002", + {"output_cost_per_second": 0.002}, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + response = httpx.Response( + status_code=200, + json={ + "predictions": [ + { + "audioContent": "clip-1", + "mimeType": "audio/wav", + }, + { + "audioContent": "clip-2", + "mimeType": "audio/wav", + }, + ] + }, + ) + + result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=response, + logging_obj=logging_obj, + url_route="/v1/projects/test/locations/us-central1/publishers/google/models/lyria-002:predict", + result=response.text, + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"instances": [{"prompt": "ambient piano"}]}, + ) + + assert result["result"] == { + "response": { + "predictions": [ + { + "audioContent": "clip-1", + "mimeType": "audio/wav", + }, + { + "audioContent": "clip-2", + "mimeType": "audio/wav", + }, + ] + } + } + assert result["kwargs"]["model"] == "lyria-002" + assert result["kwargs"]["custom_llm_provider"] == "vertex_ai" + assert result["kwargs"]["response_cost"] == pytest.approx(0.12) + assert logging_obj.model == "lyria-002" + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.12) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 200cfd02197..d8255e9ff1b 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1067,6 +1067,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "/v1/images/variations", "/v1/images/edits", "/v1/batch", + "/v1beta/interactions", "/v1/audio/transcriptions", "/v1/audio/speech", "/v1/ocr", @@ -2833,6 +2834,42 @@ def test_gemini_lyria_3_preview_models_in_cost_map(): assert clip["output_cost_per_image"] == 0.04 +def test_vertex_ai_lyria_models_in_cost_map(): + import json + from pathlib import Path + + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(json_path) as f: + model_cost = json.load(f) + + lyria_2 = model_cost.get("vertex_ai/lyria-002") + clip = model_cost.get("vertex_ai/lyria-3-clip-preview") + pro = model_cost.get("vertex_ai/lyria-3-pro-preview") + + assert lyria_2 is not None + assert clip is not None + assert pro is not None + assert lyria_2["litellm_provider"] == "vertex_ai" + assert clip["litellm_provider"] == "vertex_ai" + assert pro["litellm_provider"] == "vertex_ai" + assert lyria_2["output_cost_per_second"] == 0.002 + assert lyria_2["supported_modalities"] == ["text"] + assert lyria_2["supported_output_modalities"] == ["audio"] + assert lyria_2["supports_audio_output"] is True + assert clip["output_cost_per_image"] == 0.04 + assert pro["output_cost_per_image"] == 0.08 + assert clip["supported_endpoints"] == ["/v1beta/interactions"] + assert pro["supported_endpoints"] == ["/v1beta/interactions"] + assert clip["supported_modalities"] == ["text", "image"] + assert pro["supported_modalities"] == ["text", "image"] + assert clip["supported_regions"] == ["global"] + assert pro["supported_regions"] == ["global"] + assert clip["supports_audio_output"] is True + assert pro["supports_audio_output"] is True + assert clip["supports_image_input"] is True + assert pro["supports_image_input"] is True + + def test_model_info_for_fireworks_short_form_models(): """ Test that fireworks_ai short-form model entries (fireworks_ai/) From 514e9a1ee62e52480343d084b8f87b54c67f5a41 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Fri, 19 Jun 2026 17:44:54 -0500 Subject: [PATCH 08/63] fix(vertex): address lyria review feedback --- ...odel_prices_and_context_window_backup.json | 1 + .../vertex_passthrough_logging_handler.py | 44 ++++++++++-------- model_prices_and_context_window.json | 1 + ...test_vertex_passthrough_logging_handler.py | 46 ++++++++++++++++++- tests/test_litellm/test_utils.py | 2 + 5 files changed, 75 insertions(+), 19 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index e854e8fc9bb..1a1d92805e2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45512,6 +45512,7 @@ "supports_tool_choice": true }, "vertex_ai/lyria-002": { + "audio_seconds_per_prediction": 30, "litellm_provider": "vertex_ai", "max_audio_length_hours": 0.009111111111111111, "max_audio_per_prompt": 4, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index a9f68f8380b..6b2d7763fa0 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -44,9 +44,7 @@ else: PassThroughEndpointLogging = Any LiteLLMBatch = Any -# Define EndpointType locally to avoid import issues EndpointType = Any -_LYRIA_SECONDS_PER_AUDIO_PREDICTION = 30 class VertexPassthroughLoggingHandler: @@ -271,11 +269,11 @@ class VertexPassthroughLoggingHandler: _json_response: Final[dict[str, object]] = httpx_response.json() litellm_prediction_response: ModelResponse | EmbeddingResponse | ImageResponse = ModelResponse() - if VertexPassthroughLoggingHandler._is_lyria_predict_response( + if VertexPassthroughLoggingHandler._is_audio_predict_response( model=model, json_response=_json_response, ): - return VertexPassthroughLoggingHandler._handle_lyria_predict_response( + return VertexPassthroughLoggingHandler._handle_audio_predict_response( json_response=_json_response, logging_obj=logging_obj, model=model, @@ -335,23 +333,19 @@ class VertexPassthroughLoggingHandler: } @staticmethod - def _handle_lyria_predict_response( + def _handle_audio_predict_response( json_response: dict, logging_obj: LiteLLMLoggingObj, model: str, kwargs: dict, ) -> PassThroughEndpointLoggingTypedDict: - prediction_count: Final = ( - VertexPassthroughLoggingHandler._get_lyria_audio_prediction_count( - json_response=json_response - ) + prediction_count: Final = VertexPassthroughLoggingHandler._get_audio_prediction_count( + json_response=json_response ) - model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}", {}) response_cost: Final = ( - model_info.get("output_cost_per_second", 0.0) - * _LYRIA_SECONDS_PER_AUDIO_PREDICTION - * prediction_count - ) + VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) + or 0.0 + ) * prediction_count logging_obj.model = model logging_obj.model_call_details["model"] = model @@ -372,17 +366,31 @@ class VertexPassthroughLoggingHandler: } @staticmethod - def _is_lyria_predict_response(model: str, json_response: dict) -> bool: + def _is_audio_predict_response(model: str, json_response: dict) -> bool: return ( - model == "lyria-002" - and VertexPassthroughLoggingHandler._get_lyria_audio_prediction_count( + VertexPassthroughLoggingHandler._get_audio_prediction_count( json_response=json_response ) > 0 + and VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost( + model=model + ) + is not None ) @staticmethod - def _get_lyria_audio_prediction_count(json_response: dict) -> int: + def _get_audio_prediction_unit_cost(model: str) -> float | None: + model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}", {}) + output_cost_per_second: Final = model_info.get("output_cost_per_second") + audio_seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction") + if not isinstance(output_cost_per_second, (int, float)) or not isinstance( + audio_seconds_per_prediction, (int, float) + ): + return None + return float(output_cost_per_second * audio_seconds_per_prediction) + + @staticmethod + def _get_audio_prediction_count(json_response: dict) -> int: predictions: Final = json_response.get("predictions") if not isinstance(predictions, list): return 0 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index e854e8fc9bb..1a1d92805e2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45512,6 +45512,7 @@ "supports_tool_choice": true }, "vertex_ai/lyria-002": { + "audio_seconds_per_prediction": 30, "litellm_provider": "vertex_ai", "max_audio_length_hours": 0.009111111111111111, "max_audio_per_prompt": 4, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index 7af67cc0795..9f65fa09b1b 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -16,7 +16,10 @@ def test_lyria_predict_response_preserves_audio_response_and_logs_cost( monkeypatch.setitem( litellm.model_cost, "vertex_ai/lyria-002", - {"output_cost_per_second": 0.002}, + { + "audio_seconds_per_prediction": 30, + "output_cost_per_second": 0.002, + }, ) logging_obj = MagicMock() logging_obj.model_call_details = {} @@ -66,3 +69,44 @@ def test_lyria_predict_response_preserves_audio_response_and_logs_cost( assert result["kwargs"]["response_cost"] == pytest.approx(0.12) assert logging_obj.model == "lyria-002" assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.12) + + +def test_audio_predict_response_uses_model_map_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/music-audio-preview", + { + "audio_seconds_per_prediction": 12, + "output_cost_per_second": 0.5, + }, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + response = httpx.Response( + status_code=200, + json={ + "predictions": [ + { + "audioContent": "clip", + "mimeType": "audio/wav", + } + ] + }, + ) + + result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=response, + logging_obj=logging_obj, + url_route="/v1/projects/test/locations/us-central1/publishers/google/models/music-audio-preview:predict", + result=response.text, + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"instances": [{"prompt": "ambient piano"}]}, + ) + + assert result["kwargs"]["model"] == "music-audio-preview" + assert result["kwargs"]["response_cost"] == pytest.approx(6.0) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(6.0) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d8255e9ff1b..6ae7c27b758 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -949,6 +949,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "max_tokens": {"type": "number"}, "metadata": {"type": "object"}, "provider_specific_entry": {"type": "object"}, + "audio_seconds_per_prediction": {"type": "number"}, "mode": { "type": "string", "enum": [ @@ -2852,6 +2853,7 @@ def test_vertex_ai_lyria_models_in_cost_map(): assert lyria_2["litellm_provider"] == "vertex_ai" assert clip["litellm_provider"] == "vertex_ai" assert pro["litellm_provider"] == "vertex_ai" + assert lyria_2["audio_seconds_per_prediction"] == 30 assert lyria_2["output_cost_per_second"] == 0.002 assert lyria_2["supported_modalities"] == ["text"] assert lyria_2["supported_output_modalities"] == ["audio"] From e00fe023a973860345f56278233398a3f8a5b4ff Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 15 Jul 2026 19:25:19 -0500 Subject: [PATCH 09/63] test(models): validate Lyria audio metadata --- tests/test_litellm/test_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 6ae7c27b758..7bec649099a 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -944,6 +944,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "code_interpreter_cost_per_session": {"type": "number"}, "inference_geo": {"type": "string"}, "litellm_provider": {"type": "string"}, + "max_audio_length_hours": {"type": "number"}, + "max_audio_per_prompt": {"type": "number"}, "max_input_tokens": {"type": "number"}, "max_output_tokens": {"type": "number"}, "max_tokens": {"type": "number"}, From b96844dd0c23a084108191df7ff37423a40f862f Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 15 Jul 2026 19:56:53 -0500 Subject: [PATCH 10/63] feat(vertex): expose Lyria through audio speech --- litellm/cost_calculator.py | 12 +- .../text_to_speech/transformation.py | 160 ++++++++++ litellm/main.py | 6 +- ...odel_prices_and_context_window_backup.json | 15 +- .../vertex_passthrough_logging_handler.py | 12 +- litellm/proxy/proxy_server.py | 11 +- litellm/types/utils.py | 4 + litellm/utils.py | 6 + model_prices_and_context_window.json | 15 +- ...test_vertex_passthrough_logging_handler.py | 33 ++ .../text_to_speech/test_transformation.py | 298 +++++++++++++++++- tests/test_litellm/test_cost_calculator.py | 18 ++ tests/test_litellm/test_utils.py | 15 +- 13 files changed, 565 insertions(+), 40 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index b83e9b395a8..2bac5e234bc 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -496,9 +496,19 @@ def cost_per_token( # see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models if call_type == "speech" or call_type == "aspeech": speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) - cost_metric: Final = select_cost_metric_for_model(speech_model_info) prompt_cost: float = 0.0 completion_cost: float = 0.0 + if not speech_model_info.get("input_cost_per_character") and not speech_model_info.get( + "input_cost_per_token" + ): + output_cost_per_generation: Final = speech_model_info.get("output_cost_per_image") + output_cost_per_second: Final = speech_model_info.get("output_cost_per_second") + audio_seconds_per_prediction: Final = speech_model_info.get("audio_seconds_per_prediction") + if output_cost_per_generation is not None: + return prompt_cost, float(output_cost_per_generation) + if output_cost_per_second is not None and audio_seconds_per_prediction is not None: + return prompt_cost, float(output_cost_per_second) * float(audio_seconds_per_prediction) + cost_metric: Final = select_cost_metric_for_model(speech_model_info) if cost_metric == "cost_per_character": if prompt_characters is None: raise ValueError( diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 332f892ae6b..ad39aa35f8c 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -12,6 +12,8 @@ from typing import TYPE_CHECKING, Any, Final, Union import httpx +import litellm +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.audio_utils.utils import ( speech_media_type_from_audio_bytes, ) @@ -471,3 +473,161 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): # Initialize the HttpxBinaryResponseContent instance return HttpxBinaryResponseContent(response) + + +class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): + LYRIA_MODELS = { + "lyria-002", + "lyria-3-clip-preview", + "lyria-3-pro-preview", + } + + @classmethod + def is_lyria_model(cls, model: str) -> bool: + return model.removeprefix("vertex_ai/") in cls.LYRIA_MODELS + + def get_supported_openai_params(self, model: str) -> list: + return ["response_format"] + + def map_openai_params( + self, + model: str, + optional_params: dict, + voice: str | dict | None = None, + drop_params: bool = False, + kwargs: dict = {}, + ) -> tuple[str | None, dict]: + mapped_params = dict(optional_params) + base_model = model.removeprefix("vertex_ai/") + unsupported_params = [param for param in ("speed", "instructions") if mapped_params.get(param) is not None] + if unsupported_params: + if drop_params or litellm.drop_params: + for param in unsupported_params: + mapped_params.pop(param, None) + else: + raise UnsupportedParamsError( + status_code=400, + message=( + f"Vertex AI {base_model} does not support the OpenAI parameters: " + f"{', '.join(unsupported_params)}. To drop unsupported openai params " + "from the call, set `litellm.drop_params = True`" + ), + ) + response_format = mapped_params.get("response_format") + supported_formats = ( + {"wav"} if base_model == "lyria-002" else {"mp3", "wav"} if base_model == "lyria-3-pro-preview" else {"mp3"} + ) + if response_format is not None and response_format not in supported_formats: + if drop_params or litellm.drop_params: + mapped_params.pop("response_format", None) + else: + raise UnsupportedParamsError( + status_code=400, + message=( + f"Vertex AI {base_model} does not support response_format={response_format!r}. " + f"Supported values: {', '.join(sorted(supported_formats))}. " + "To drop unsupported openai params from the call, set `litellm.drop_params = True`" + ), + ) + return voice if isinstance(voice, str) else None, mapped_params + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, + ) -> str: + base_model = model.removeprefix("vertex_ai/") + project = self.safe_get_vertex_ai_project(litellm_params) + if project is None: + _, project = self._ensure_access_token( + credentials=self.safe_get_vertex_ai_credentials(litellm_params), + project_id=None, + custom_llm_provider="vertex_ai", + ) + if base_model.startswith("lyria-3-"): + from litellm.llms.vertex_ai.interactions.transformation import ( + VertexAIInteractionsConfig, + ) + + return VertexAIInteractionsConfig().get_complete_url( + api_base=api_base, + model=base_model, + litellm_params={**litellm_params, "vertex_project": project}, + ) + location = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location() + base_url = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/") + return f"{base_url}/v1/projects/{project}/locations/{location}/publishers/google/models/{base_model}:predict" + + def transform_text_to_speech_request( + self, + model: str, + input: str, + voice: str | None, + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> TextToSpeechRequestData: + access_token, project = self._ensure_access_token( + credentials=self.safe_get_vertex_ai_credentials(litellm_params), + project_id=self.safe_get_vertex_ai_project(litellm_params), + custom_llm_provider="vertex_ai", + ) + headers.update( + { + "Authorization": f"Bearer {access_token}", + "x-goog-user-project": project, + "Content-Type": "application/json", + } + ) + base_model = model.removeprefix("vertex_ai/") + if base_model == "lyria-002": + request_body = { + "instances": [{"prompt": input}], + "parameters": {"sample_count": 1}, + } + else: + request_body = {"model": base_model, "input": input} + if optional_params.get("response_format") == "wav": + request_body["response_format"] = { + "type": "audio", + "mime_type": "audio/wav", + } + return TextToSpeechRequestData(dict_body=request_body, headers=headers) + + def transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + from litellm.types.llms.openai import HttpxBinaryResponseContent + + response_json = raw_response.json() + base_model = model.removeprefix("vertex_ai/") + audio_data: str | None = None + mime_type: str | None = None + if base_model == "lyria-002": + predictions = response_json.get("predictions") or [] + if predictions: + audio_data = predictions[0].get("audioContent") or predictions[0].get("bytesBase64Encoded") + mime_type = predictions[0].get("mimeType") + else: + for step in response_json.get("steps") or response_json.get("outputs") or []: + content_items = step.get("content") or [] if step.get("type") == "model_output" else [step] + for content in content_items: + if content.get("type") == "audio" and content.get("data"): + audio_data = content["data"] + mime_type = content.get("mime_type") + if audio_data is None: + raise ValueError(f"No generated audio found in Vertex AI {base_model} response") + mime_type = mime_type or ("audio/wav" if base_model == "lyria-002" else "audio/mpeg") + response = HttpxBinaryResponseContent( + httpx.Response( + status_code=raw_response.status_code, + content=base64.b64decode(audio_data), + headers={"content-type": mime_type}, + ) + ) + response._hidden_params = {"audio_mime_type": mime_type} + return response diff --git a/litellm/main.py b/litellm/main.py index 0128e4defe5..55db92d44b3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8235,6 +8235,7 @@ def speech( ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": from litellm.llms.vertex_ai.text_to_speech.transformation import ( + VertexAILyriaTextToSpeechConfig, VertexAITextToSpeechConfig, ) @@ -8259,7 +8260,10 @@ def speech( # Vertex AI Text-to-Speech (Google Cloud TTS) if text_to_speech_provider_config is None: - text_to_speech_provider_config = VertexAITextToSpeechConfig() + if VertexAILyriaTextToSpeechConfig.is_lyria_model(model): + text_to_speech_provider_config = VertexAILyriaTextToSpeechConfig() + else: + text_to_speech_provider_config = VertexAITextToSpeechConfig() # Cast to specific Vertex AI config type to access dispatch method vertex_config: Final = cast(VertexAITextToSpeechConfig, text_to_speech_provider_config) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1a1d92805e2..2955a0ffa4e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45516,9 +45516,12 @@ "litellm_provider": "vertex_ai", "max_audio_length_hours": 0.009111111111111111, "max_audio_per_prompt": 4, - "mode": "chat", + "mode": "audio_speech", "output_cost_per_second": 0.002, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_endpoints": [ + "/v1/audio/speech" + ], "supported_modalities": [ "text" ], @@ -45533,12 +45536,13 @@ "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, - "mode": "chat", + "mode": "audio_speech", "output_cost_per_image": 0.04, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_endpoints": [ - "/v1beta/interactions" + "/v1beta/interactions", + "/v1/audio/speech" ], "supported_modalities": [ "text", @@ -45566,12 +45570,13 @@ "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, - "mode": "chat", + "mode": "audio_speech", "output_cost_per_image": 0.08, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_endpoints": [ - "/v1beta/interactions" + "/v1beta/interactions", + "/v1/audio/speech" ], "supported_modalities": [ "text", diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 6b2d7763fa0..ab3b24f470f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -368,14 +368,8 @@ class VertexPassthroughLoggingHandler: @staticmethod def _is_audio_predict_response(model: str, json_response: dict) -> bool: return ( - VertexPassthroughLoggingHandler._get_audio_prediction_count( - json_response=json_response - ) - > 0 - and VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost( - model=model - ) - is not None + VertexPassthroughLoggingHandler._get_audio_prediction_count(json_response=json_response) > 0 + and VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) is not None ) @staticmethod @@ -397,7 +391,7 @@ class VertexPassthroughLoggingHandler: return sum( 1 for prediction in predictions - if isinstance(prediction, dict) and prediction.get("audioContent") + if isinstance(prediction, dict) and (prediction.get("audioContent") or prediction.get("bytesBase64Encoded")) ) @staticmethod diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27132c90e05..f11cbc92224 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11176,9 +11176,14 @@ async def audio_speech( upstream_content_type: Final = ( response.response.headers.get("content-type") if isinstance(response, HttpxBinaryResponseContent) else None ) - media_type: Final = resolve_speech_media_type( - upstream_content_type=upstream_content_type, - response_format=requested_format if isinstance(requested_format, str) else None, + hidden_audio_mime_type: Final = hidden_params.get("audio_mime_type") + media_type: Final = ( + hidden_audio_mime_type + if isinstance(hidden_audio_mime_type, str) + else resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=requested_format if isinstance(requested_format, str) else None, + ) ) return StreamingResponse( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ee6f09e05dc..d2a09639362 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -310,6 +310,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_video_per_second: float | None # only for vertex ai models output_cost_per_audio_per_second: float | None # only for vertex ai models output_cost_per_second: float | None # for OpenAI Speech models + audio_seconds_per_prediction: float | None + max_audio_length_hours: float | None + max_audio_per_prompt: int | None output_cost_per_second_1080p: ( float | None ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) @@ -333,6 +336,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): "image_generation", "chat", "audio_transcription", + "audio_speech", "responses", "ocr", "realtime", diff --git a/litellm/utils.py b/litellm/utils.py index ba456fc353b..7c8974906ed 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5880,6 +5880,9 @@ def _get_model_info_helper( "output_cost_per_token_above_512k_tokens", None ), output_cost_per_second=_model_info.get("output_cost_per_second", None), + audio_seconds_per_prediction=_model_info.get("audio_seconds_per_prediction", None), + max_audio_length_hours=_model_info.get("max_audio_length_hours", None), + max_audio_per_prompt=_model_info.get("max_audio_per_prompt", None), output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), @@ -9415,9 +9418,12 @@ class ProviderConfigManager: # mapping would drop response_format before the bridge sees it (LIT-6501) return None from litellm.llms.vertex_ai.text_to_speech.transformation import ( + VertexAILyriaTextToSpeechConfig, VertexAITextToSpeechConfig, ) + if VertexAILyriaTextToSpeechConfig.is_lyria_model(model): + return VertexAILyriaTextToSpeechConfig() return VertexAITextToSpeechConfig() elif litellm.LlmProviders.MINIMAX == provider: from litellm.llms.minimax.text_to_speech.transformation import ( diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1a1d92805e2..2955a0ffa4e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45516,9 +45516,12 @@ "litellm_provider": "vertex_ai", "max_audio_length_hours": 0.009111111111111111, "max_audio_per_prompt": 4, - "mode": "chat", + "mode": "audio_speech", "output_cost_per_second": 0.002, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_endpoints": [ + "/v1/audio/speech" + ], "supported_modalities": [ "text" ], @@ -45533,12 +45536,13 @@ "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, - "mode": "chat", + "mode": "audio_speech", "output_cost_per_image": 0.04, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_endpoints": [ - "/v1beta/interactions" + "/v1beta/interactions", + "/v1/audio/speech" ], "supported_modalities": [ "text", @@ -45566,12 +45570,13 @@ "max_input_tokens": 131072, "max_output_tokens": 8192, "max_tokens": 8192, - "mode": "chat", + "mode": "audio_speech", "output_cost_per_image": 0.08, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_endpoints": [ - "/v1beta/interactions" + "/v1beta/interactions", + "/v1/audio/speech" ], "supported_modalities": [ "text", diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index 9f65fa09b1b..ccce19f1634 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -110,3 +110,36 @@ def test_audio_predict_response_uses_model_map_metadata( assert result["kwargs"]["model"] == "music-audio-preview" assert result["kwargs"]["response_cost"] == pytest.approx(6.0) assert logging_obj.model_call_details["response_cost"] == pytest.approx(6.0) + + +def test_audio_predict_response_supports_bytes_base64_encoded( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/lyria-002", + { + "audio_seconds_per_prediction": 30, + "output_cost_per_second": 0.002, + }, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + response = httpx.Response( + status_code=200, + json={"predictions": [{"bytesBase64Encoded": "clip"}]}, + ) + + result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=response, + logging_obj=logging_obj, + url_route="/v1/projects/test/locations/us-central1/publishers/google/models/lyria-002:predict", + result=response.text, + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"instances": [{"prompt": "ambient piano"}]}, + ) + + assert result["kwargs"]["response_cost"] == pytest.approx(0.06) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06) diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index fba337b5f2c..910bc977a79 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -4,11 +4,13 @@ from unittest.mock import MagicMock, Mock, patch import httpx import pytest - import litellm from litellm.llms.vertex_ai.text_to_speech.transformation import ( + VertexAILyriaTextToSpeechConfig, VertexAITextToSpeechConfig, ) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager class TestVertexAITextToSpeechConfig: @@ -41,9 +43,7 @@ class TestVertexAITextToSpeechConfig: @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") - def test_transform_text_to_speech_request_body( - self, mock_get_token, mock_ensure_token - ): + def test_transform_text_to_speech_request_body(self, mock_get_token, mock_ensure_token): """Test that transform_text_to_speech_request generates correct request body""" # Mock authentication mock_ensure_token.return_value = ("mock-token", "test-project") @@ -104,9 +104,7 @@ class TestVertexAITextToSpeechConfig: config = VertexAITextToSpeechConfig() # Test with a Chirp3 HD voice - voice_str, voice_dict = config._map_voice_to_vertex_format( - "en-US-Chirp3-HD-Charon" - ) + voice_str, voice_dict = config._map_voice_to_vertex_format("en-US-Chirp3-HD-Charon") assert voice_str == "en-US-Chirp3-HD-Charon" assert voice_dict is not None @@ -169,6 +167,284 @@ def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled(): assert result.response.content == raw_pcm +class TestVertexAILyriaTextToSpeechConfig: + @pytest.mark.parametrize( + "model", + ["lyria-002", "vertex_ai/lyria-3-clip-preview", "lyria-3-pro-preview"], + ) + def test_provider_config_manager_selects_lyria_config(self, model): + config = ProviderConfigManager.get_provider_text_to_speech_config( + model=model, + provider=LlmProviders.VERTEX_AI, + ) + + assert isinstance(config, VertexAILyriaTextToSpeechConfig) + + def test_get_complete_url_for_lyria_2(self): + config = VertexAILyriaTextToSpeechConfig() + + url = config.get_complete_url( + model="lyria-002", + api_base=None, + litellm_params={ + "vertex_project": "music-project", + "vertex_location": "europe-west4", + }, + ) + + assert url == ( + "https://europe-west4-aiplatform.googleapis.com/v1/projects/music-project/" + "locations/europe-west4/publishers/google/models/lyria-002:predict" + ) + + def test_get_complete_url_for_lyria_3(self): + config = VertexAILyriaTextToSpeechConfig() + + url = config.get_complete_url( + model="lyria-3-pro-preview", + api_base=None, + litellm_params={"vertex_project": "music-project"}, + ) + + assert url == ("https://aiplatform.googleapis.com/v1beta1/projects/music-project/locations/global/interactions") + + @pytest.mark.parametrize( + ("model", "response_format", "expected_body"), + [ + ( + "lyria-002", + "wav", + { + "instances": [{"prompt": "A bright synth track"}], + "parameters": {"sample_count": 1}, + }, + ), + ( + "lyria-3-clip-preview", + "mp3", + { + "model": "lyria-3-clip-preview", + "input": "A bright synth track", + }, + ), + ( + "lyria-3-pro-preview", + "wav", + { + "model": "lyria-3-pro-preview", + "input": "A bright synth track", + "response_format": { + "type": "audio", + "mime_type": "audio/wav", + }, + }, + ), + ], + ) + @patch.object(VertexAILyriaTextToSpeechConfig, "_ensure_access_token") + def test_transform_request( + self, + mock_ensure_token, + model, + response_format, + expected_body, + ): + mock_ensure_token.return_value = ("mock-token", "music-project") + config = VertexAILyriaTextToSpeechConfig() + + request = config.transform_text_to_speech_request( + model=model, + input="A bright synth track", + voice="alloy", + optional_params={"response_format": response_format}, + litellm_params={"vertex_project": "music-project"}, + headers={}, + ) + + assert request["dict_body"] == expected_body + assert request["headers"]["Authorization"] == "Bearer mock-token" + assert request["headers"]["x-goog-user-project"] == "music-project" + + @pytest.mark.parametrize( + ("model", "response_json", "expected_audio", "expected_mime_type"), + [ + ( + "lyria-002", + { + "predictions": [ + { + "bytesBase64Encoded": "bHlyaWEtMi1hdWRpbw==", + } + ] + }, + b"lyria-2-audio", + "audio/wav", + ), + ( + "lyria-3-pro-preview", + { + "steps": [ + { + "type": "model_output", + "content": [ + {"type": "text", "text": "Generated lyrics"}, + { + "type": "audio", + "data": "bHlyaWEtMy1hdWRpbw==", + "mime_type": "audio/mpeg", + }, + ], + } + ] + }, + b"lyria-3-audio", + "audio/mpeg", + ), + ( + "lyria-3-clip-preview", + { + "outputs": [ + {"type": "text", "text": "Generated lyrics"}, + { + "type": "audio", + "data": "bHlyaWEtMy1hdWRpbw==", + "mime_type": "audio/mpeg", + }, + ] + }, + b"lyria-3-audio", + "audio/mpeg", + ), + ], + ) + def test_transform_response( + self, + model, + response_json, + expected_audio, + expected_mime_type, + ): + config = VertexAILyriaTextToSpeechConfig() + raw_response = httpx.Response(200, json=response_json) + + response = config.transform_text_to_speech_response( + model=model, + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert response.content == expected_audio + assert response._hidden_params["audio_mime_type"] == expected_mime_type + + @pytest.mark.parametrize( + ("model", "response_format"), + [ + ("lyria-002", "mp3"), + ("lyria-3-clip-preview", "wav"), + ("lyria-3-pro-preview", "opus"), + ], + ) + def test_rejects_unsupported_response_format(self, model, response_format): + config = VertexAILyriaTextToSpeechConfig() + + with pytest.raises(litellm.UnsupportedParamsError): + config.map_openai_params( + model=model, + optional_params={"response_format": response_format}, + ) + + @pytest.mark.parametrize("param", ["speed", "instructions"]) + def test_rejects_unsupported_openai_params(self, param): + config = VertexAILyriaTextToSpeechConfig() + + with pytest.raises(litellm.UnsupportedParamsError): + config.map_openai_params( + model="lyria-3-pro-preview", + optional_params={param: "unsupported"}, + ) + + @pytest.mark.parametrize( + ("model", "response_format", "response_json", "expected_url", "expected_body"), + [ + ( + "lyria-002", + "wav", + { + "predictions": [ + { + "audioContent": "bHlyaWEtMi1hdWRpbw==", + "mimeType": "audio/wav", + } + ] + }, + "https://us-central1-aiplatform.googleapis.com/v1/projects/music-project/locations/us-central1/publishers/google/models/lyria-002:predict", + { + "instances": [{"prompt": "A bright synth track"}], + "parameters": {"sample_count": 1}, + }, + ), + ( + "lyria-3-pro-preview", + "mp3", + { + "steps": [ + { + "type": "model_output", + "content": [ + { + "type": "audio", + "data": "bHlyaWEtMy1hdWRpbw==", + "mime_type": "audio/mpeg", + } + ], + } + ] + }, + "https://aiplatform.googleapis.com/v1beta1/projects/music-project/locations/global/interactions", + { + "model": "lyria-3-pro-preview", + "input": "A bright synth track", + }, + ), + ], + ) + def test_litellm_speech_dispatches_to_lyria_api( + self, + model, + response_format, + response_json, + expected_url, + expected_body, + ): + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = response_json + with ( + patch.object( + VertexAILyriaTextToSpeechConfig, + "_ensure_access_token", + return_value=("mock-token", "music-project"), + ), + patch( + "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post", + return_value=mock_response, + ) as mock_post, + ): + response = litellm.speech( + model=f"vertex_ai/{model}", + input="A bright synth track", + voice="alloy", + response_format=response_format, + vertex_project="music-project", + vertex_location="us-central1", + ) + + assert response.content in {b"lyria-2-audio", b"lyria-3-audio"} + mock_post.assert_called_once() + assert mock_post.call_args.kwargs["url"] == expected_url + assert mock_post.call_args.kwargs["json"] == expected_body + + @patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") @@ -182,9 +458,7 @@ def test_litellm_speech_vertex_ai_chirp(mock_get_token, mock_ensure_token, mock_ # Mock HTTP response mock_response = Mock(spec=httpx.Response) - mock_response.content = ( - b'{"audioContent": "SGVsbG8gV29ybGQ="}' # base64 encoded "Hello World" - ) + mock_response.content = b'{"audioContent": "SGVsbG8gV29ybGQ="}' # base64 encoded "Hello World" mock_response.status_code = 200 mock_response.headers = {"content-type": "application/json"} mock_response.json.return_value = {"audioContent": "SGVsbG8gV29ybGQ="} @@ -203,9 +477,7 @@ def test_litellm_speech_vertex_ai_chirp(mock_get_token, mock_ensure_token, mock_ call_kwargs = mock_post.call_args.kwargs # Verify the URL is the Google Cloud TTS API - assert ( - call_kwargs["url"] == "https://texttospeech.googleapis.com/v1/text:synthesize" - ) + assert call_kwargs["url"] == "https://texttospeech.googleapis.com/v1/text:synthesize" # Verify request body structure assert "json" in call_kwargs diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 7c2174018e8..ffc71c76d03 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -146,6 +146,24 @@ def test_cost_calculator_with_response_cost_in_additional_headers(): assert result == 1000 +@pytest.mark.parametrize( + ("model", "expected_cost"), + [ + ("vertex_ai/lyria-002", 0.06), + ("vertex_ai/lyria-3-clip-preview", 0.04), + ("vertex_ai/lyria-3-pro-preview", 0.08), + ], +) +def test_vertex_lyria_speech_cost(model, expected_cost, _local_model_cost_map): + cost = completion_cost( + model=model, + prompt="A bright synth track", + call_type="speech", + ) + + assert cost == pytest.approx(expected_cost) + + def test_baseten_model_api_pricing_entries(_local_model_cost_map): expected_pricing = { diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 7bec649099a..1b0b27032ce 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2855,15 +2855,25 @@ def test_vertex_ai_lyria_models_in_cost_map(): assert lyria_2["litellm_provider"] == "vertex_ai" assert clip["litellm_provider"] == "vertex_ai" assert pro["litellm_provider"] == "vertex_ai" + assert lyria_2["mode"] == "audio_speech" + assert clip["mode"] == "audio_speech" + assert pro["mode"] == "audio_speech" assert lyria_2["audio_seconds_per_prediction"] == 30 assert lyria_2["output_cost_per_second"] == 0.002 assert lyria_2["supported_modalities"] == ["text"] assert lyria_2["supported_output_modalities"] == ["audio"] assert lyria_2["supports_audio_output"] is True + assert lyria_2["supported_endpoints"] == ["/v1/audio/speech"] assert clip["output_cost_per_image"] == 0.04 assert pro["output_cost_per_image"] == 0.08 - assert clip["supported_endpoints"] == ["/v1beta/interactions"] - assert pro["supported_endpoints"] == ["/v1beta/interactions"] + assert clip["supported_endpoints"] == [ + "/v1beta/interactions", + "/v1/audio/speech", + ] + assert pro["supported_endpoints"] == [ + "/v1beta/interactions", + "/v1/audio/speech", + ] assert clip["supported_modalities"] == ["text", "image"] assert pro["supported_modalities"] == ["text", "image"] assert clip["supported_regions"] == ["global"] @@ -2873,7 +2883,6 @@ def test_vertex_ai_lyria_models_in_cost_map(): assert clip["supports_image_input"] is True assert pro["supports_image_input"] is True - def test_model_info_for_fireworks_short_form_models(): """ Test that fireworks_ai short-form model entries (fireworks_ai/) From f18cb0cdb48ce0338a30af940d557842eee4fb19 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 15 Jul 2026 20:28:00 -0500 Subject: [PATCH 11/63] fix(vertex): make Lyria routing and billing data-driven --- litellm/llms/vertex_ai/common_utils.py | 35 +++++++++++ .../text_to_speech/transformation.py | 36 ++++++----- ...odel_prices_and_context_window_backup.json | 19 +++++- litellm/types/utils.py | 2 + litellm/utils.py | 2 + model_prices_and_context_window.json | 19 +++++- .../text_to_speech/test_transformation.py | 62 +++++++++++++++++++ tests/test_litellm/test_utils.py | 17 +++++ 8 files changed, 172 insertions(+), 20 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index a36c920dda0..8a4c1e68623 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,9 +1,12 @@ import re from copy import deepcopy from enum import Enum +from functools import lru_cache from typing import Any, Final, Literal, cast, get_type_hints import httpx +from pydantic import TypeAdapter, ValidationError +from typing_extensions import NotRequired, TypedDict import litellm from litellm._logging import verbose_logger @@ -21,6 +24,38 @@ from litellm.types.utils import TokenCountResponse from litellm.utils import supports_response_schema, supports_system_messages +class VertexAILyriaModelInfo(TypedDict): + vertex_ai_audio_api: Literal["lyria_predict", "lyria_interactions"] + supported_audio_formats: tuple[Literal["mp3", "wav"], ...] + output_cost_per_image: NotRequired[float] + + +_VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER = TypeAdapter(VertexAILyriaModelInfo) + + +def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyriaModelInfo | None: + if raw_model_info is None: + return None + try: + return _VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER.validate_python(raw_model_info) + except ValidationError: + return None + + +@lru_cache(maxsize=32) +def _get_bundled_vertex_ai_lyria_model_info(model_key: str) -> VertexAILyriaModelInfo | None: + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + + bundled_model_info = GetModelCostMap.load_local_model_cost_map().get(model_key) + return _validate_vertex_ai_lyria_model_info(bundled_model_info) + + +def get_vertex_ai_lyria_model_info(model: str) -> VertexAILyriaModelInfo | None: + model_key = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" + runtime_model_info = _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key)) + return runtime_model_info or _get_bundled_vertex_ai_lyria_model_info(model_key) + + class VertexAIError(BaseLLMException): def __init__( self, diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index ad39aa35f8c..d56f151a8c2 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -21,6 +21,10 @@ from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, ) +from litellm.llms.vertex_ai.common_utils import ( + VertexAILyriaModelInfo, + get_vertex_ai_lyria_model_info, +) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.llms.vertex_ai_text_to_speech import ( @@ -476,15 +480,16 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): - LYRIA_MODELS = { - "lyria-002", - "lyria-3-clip-preview", - "lyria-3-pro-preview", - } - @classmethod def is_lyria_model(cls, model: str) -> bool: - return model.removeprefix("vertex_ai/") in cls.LYRIA_MODELS + return get_vertex_ai_lyria_model_info(model=model) is not None + + @staticmethod + def _get_model_info(model: str) -> VertexAILyriaModelInfo: + model_info = get_vertex_ai_lyria_model_info(model=model) + if model_info is None: + raise ValueError(f"Vertex AI model {model!r} does not declare a Lyria audio API") + return model_info def get_supported_openai_params(self, model: str) -> list: return ["response_format"] @@ -499,6 +504,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): ) -> tuple[str | None, dict]: mapped_params = dict(optional_params) base_model = model.removeprefix("vertex_ai/") + model_info = self._get_model_info(model=model) unsupported_params = [param for param in ("speed", "instructions") if mapped_params.get(param) is not None] if unsupported_params: if drop_params or litellm.drop_params: @@ -514,9 +520,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): ), ) response_format = mapped_params.get("response_format") - supported_formats = ( - {"wav"} if base_model == "lyria-002" else {"mp3", "wav"} if base_model == "lyria-3-pro-preview" else {"mp3"} - ) + supported_formats = frozenset(model_info["supported_audio_formats"]) if response_format is not None and response_format not in supported_formats: if drop_params or litellm.drop_params: mapped_params.pop("response_format", None) @@ -538,6 +542,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): litellm_params: dict, ) -> str: base_model = model.removeprefix("vertex_ai/") + model_info = self._get_model_info(model=model) project = self.safe_get_vertex_ai_project(litellm_params) if project is None: _, project = self._ensure_access_token( @@ -545,7 +550,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): project_id=None, custom_llm_provider="vertex_ai", ) - if base_model.startswith("lyria-3-"): + if model_info["vertex_ai_audio_api"] == "lyria_interactions": from litellm.llms.vertex_ai.interactions.transformation import ( VertexAIInteractionsConfig, ) @@ -581,7 +586,8 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): } ) base_model = model.removeprefix("vertex_ai/") - if base_model == "lyria-002": + model_info = self._get_model_info(model=model) + if model_info["vertex_ai_audio_api"] == "lyria_predict": request_body = { "instances": [{"prompt": input}], "parameters": {"sample_count": 1}, @@ -605,9 +611,10 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): response_json = raw_response.json() base_model = model.removeprefix("vertex_ai/") + model_info = self._get_model_info(model=model) audio_data: str | None = None mime_type: str | None = None - if base_model == "lyria-002": + if model_info["vertex_ai_audio_api"] == "lyria_predict": predictions = response_json.get("predictions") or [] if predictions: audio_data = predictions[0].get("audioContent") or predictions[0].get("bytesBase64Encoded") @@ -621,7 +628,8 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): mime_type = content.get("mime_type") if audio_data is None: raise ValueError(f"No generated audio found in Vertex AI {base_model} response") - mime_type = mime_type or ("audio/wav" if base_model == "lyria-002" else "audio/mpeg") + default_format = model_info["supported_audio_formats"][0] + mime_type = mime_type or {"mp3": "audio/mpeg", "wav": "audio/wav"}[default_format] response = HttpxBinaryResponseContent( httpx.Response( status_code=raw_response.status_code, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 2955a0ffa4e..9986d5056d1 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45519,6 +45519,9 @@ "mode": "audio_speech", "output_cost_per_second": 0.002, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_audio_formats": [ + "wav" + ], "supported_endpoints": [ "/v1/audio/speech" ], @@ -45528,7 +45531,8 @@ "supported_output_modalities": [ "audio" ], - "supports_audio_output": true + "supports_audio_output": true, + "vertex_ai_audio_api": "lyria_predict" }, "vertex_ai/lyria-3-clip-preview": { "input_cost_per_token": 0, @@ -45540,6 +45544,9 @@ "output_cost_per_image": 0.04, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_audio_formats": [ + "mp3" + ], "supported_endpoints": [ "/v1beta/interactions", "/v1/audio/speech" @@ -45562,7 +45569,8 @@ "supports_response_schema": false, "supports_system_messages": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "vertex_ai_audio_api": "lyria_interactions" }, "vertex_ai/lyria-3-pro-preview": { "input_cost_per_token": 0, @@ -45574,6 +45582,10 @@ "output_cost_per_image": 0.08, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_audio_formats": [ + "mp3", + "wav" + ], "supported_endpoints": [ "/v1beta/interactions", "/v1/audio/speech" @@ -45596,7 +45608,8 @@ "supports_response_schema": false, "supports_system_messages": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "vertex_ai_audio_api": "lyria_interactions" }, "vertex_ai/meta/llama-3.1-405b-instruct-maas": { "input_cost_per_token": 5e-06, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d2a09639362..bb0485be7c1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -168,6 +168,8 @@ class ProviderSpecificModelInfo(TypedDict, total=False): default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None] supports_output_config: bool | None supports_image_size: bool | None + supported_audio_formats: list[Literal["mp3", "wav"]] | None + vertex_ai_audio_api: Literal["lyria_predict", "lyria_interactions"] | None bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None bedrock_converse_supports_strict_tools: bool | None diff --git a/litellm/utils.py b/litellm/utils.py index 7c8974906ed..d6a68b3ce6e 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5940,6 +5940,8 @@ def _get_model_info_helper( provider_specific_entry=_model_info.get("provider_specific_entry", None), uses_embed_content=_model_info.get("uses_embed_content", None), supports_image_size=_model_info.get("supports_image_size", None), + supported_audio_formats=_model_info.get("supported_audio_formats", None), + vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None), ) for cost_key, cost_value in _model_info.items(): if cost_key not in returned_model_info and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 2955a0ffa4e..9986d5056d1 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45519,6 +45519,9 @@ "mode": "audio_speech", "output_cost_per_second": 0.002, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_audio_formats": [ + "wav" + ], "supported_endpoints": [ "/v1/audio/speech" ], @@ -45528,7 +45531,8 @@ "supported_output_modalities": [ "audio" ], - "supports_audio_output": true + "supports_audio_output": true, + "vertex_ai_audio_api": "lyria_predict" }, "vertex_ai/lyria-3-clip-preview": { "input_cost_per_token": 0, @@ -45540,6 +45544,9 @@ "output_cost_per_image": 0.04, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_audio_formats": [ + "mp3" + ], "supported_endpoints": [ "/v1beta/interactions", "/v1/audio/speech" @@ -45562,7 +45569,8 @@ "supports_response_schema": false, "supports_system_messages": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "vertex_ai_audio_api": "lyria_interactions" }, "vertex_ai/lyria-3-pro-preview": { "input_cost_per_token": 0, @@ -45574,6 +45582,10 @@ "output_cost_per_image": 0.08, "output_cost_per_token": 0, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", + "supported_audio_formats": [ + "mp3", + "wav" + ], "supported_endpoints": [ "/v1beta/interactions", "/v1/audio/speech" @@ -45596,7 +45608,8 @@ "supports_response_schema": false, "supports_system_messages": false, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "vertex_ai_audio_api": "lyria_interactions" }, "vertex_ai/meta/llama-3.1-405b-instruct-maas": { "input_cost_per_token": 5e-06, diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index 910bc977a79..a1bb203e67e 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -180,6 +180,68 @@ class TestVertexAILyriaTextToSpeechConfig: assert isinstance(config, VertexAILyriaTextToSpeechConfig) + @pytest.mark.parametrize( + ("model", "vertex_ai_audio_api", "supported_audio_formats", "expected_url"), + [ + ( + "future-lyria-predict", + "lyria_predict", + ["wav"], + "https://us-central1-aiplatform.googleapis.com/v1/projects/music-project/locations/" + "us-central1/publishers/google/models/future-lyria-predict:predict", + ), + ( + "future-music-interactions", + "lyria_interactions", + ["mp3", "wav"], + "https://aiplatform.googleapis.com/v1beta1/projects/music-project/locations/global/interactions", + ), + ], + ) + def test_dispatches_from_model_metadata( + self, + monkeypatch, + model, + vertex_ai_audio_api, + supported_audio_formats, + expected_url, + ): + monkeypatch.setitem( + litellm.model_cost, + f"vertex_ai/{model}", + { + "vertex_ai_audio_api": vertex_ai_audio_api, + "supported_audio_formats": supported_audio_formats, + }, + ) + + config = ProviderConfigManager.get_provider_text_to_speech_config( + model=model, + provider=LlmProviders.VERTEX_AI, + ) + + assert isinstance(config, VertexAILyriaTextToSpeechConfig) + assert ( + config.get_complete_url( + model=model, + api_base=None, + litellm_params={ + "vertex_project": "music-project", + "vertex_location": "us-central1", + }, + ) + == expected_url + ) + + def test_vertex_chirp_does_not_select_lyria_config(self): + config = ProviderConfigManager.get_provider_text_to_speech_config( + model="chirp", + provider=LlmProviders.VERTEX_AI, + ) + + assert isinstance(config, VertexAITextToSpeechConfig) + assert not isinstance(config, VertexAILyriaTextToSpeechConfig) + def test_get_complete_url_for_lyria_2(self): config = VertexAILyriaTextToSpeechConfig() diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1b0b27032ce..a179a82fcb2 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1048,6 +1048,17 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_sampling_params": {"type": "boolean"}, "supports_output_config": {"type": "boolean"}, "supports_speed": {"type": "boolean"}, + "supported_audio_formats": { + "type": "array", + "items": { + "type": "string", + "enum": ["mp3", "wav"], + }, + }, + "vertex_ai_audio_api": { + "type": "string", + "enum": ["lyria_predict", "lyria_interactions"], + }, "bedrock_output_config_effort_ceiling": { "type": "string", "enum": ["low", "medium", "high", "max", "xhigh"], @@ -2863,9 +2874,15 @@ def test_vertex_ai_lyria_models_in_cost_map(): assert lyria_2["supported_modalities"] == ["text"] assert lyria_2["supported_output_modalities"] == ["audio"] assert lyria_2["supports_audio_output"] is True + assert lyria_2["supported_audio_formats"] == ["wav"] + assert lyria_2["vertex_ai_audio_api"] == "lyria_predict" assert lyria_2["supported_endpoints"] == ["/v1/audio/speech"] assert clip["output_cost_per_image"] == 0.04 assert pro["output_cost_per_image"] == 0.08 + assert clip["supported_audio_formats"] == ["mp3"] + assert pro["supported_audio_formats"] == ["mp3", "wav"] + assert clip["vertex_ai_audio_api"] == "lyria_interactions" + assert pro["vertex_ai_audio_api"] == "lyria_interactions" assert clip["supported_endpoints"] == [ "/v1beta/interactions", "/v1/audio/speech", From e9e2fbb4385e412273ce03c6e0038d013316ce8a Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 12 Aug 2026 12:31:44 -0500 Subject: [PATCH 12/63] style(vertex-ai): modernize Lyria tests --- .../llms/vertex_ai/test_vertex_passthrough_logging_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index ccce19f1634..1e8d569d3d1 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -2,9 +2,9 @@ from datetime import datetime from unittest.mock import MagicMock import httpx -import litellm import pytest +import litellm from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) From bd5123564c66d1eb68c253ae6d2405c1e383104c Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 12 Aug 2026 12:56:06 -0500 Subject: [PATCH 13/63] fix(vertex-ai): classify Lyria model metadata --- ci_cd/generate_model_prices_schema.py | 21 ++++++++++++ .../text_to_speech/transformation.py | 2 +- .../vertex_passthrough_logging_handler.py | 3 +- model_prices_and_context_window.schema.json | 33 +++++++++++++++++++ 4 files changed, 56 insertions(+), 3 deletions(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 57cc742d5c4..ee0dad25c81 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -58,6 +58,11 @@ OBJECT_KEYS: dict[str, JsonSchema] = { } ARRAY_KEYS: dict[str, JsonSchema] = { + "supported_audio_formats": { + "type": "array", + "description": "Audio container formats the model can return.", + "items": {"type": "string", "enum": ["mp3", "wav"]}, + }, "supported_endpoints": { "type": "array", "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.", @@ -116,6 +121,10 @@ ARRAY_KEYS: dict[str, JsonSchema] = { } INTEGER_KEYS: dict[str, JsonSchema] = { + "max_audio_per_prompt": { + **NONNEG_INTEGER, + "description": "Maximum number of audio outputs accepted or generated per prompt.", + }, "max_tokens": { **NONNEG_INTEGER, "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.", @@ -141,6 +150,14 @@ INTEGER_KEYS: dict[str, JsonSchema] = { } NUMBER_KEYS: dict[str, JsonSchema] = { + "audio_seconds_per_prediction": { + **NONNEG_NUMBER, + "description": "Audio duration, in seconds, produced by one prediction.", + }, + "max_audio_length_hours": { + **NONNEG_NUMBER, + "description": "Maximum generated audio duration, expressed in hours.", + }, "regional_processing_uplift_multiplier_eu": { "type": "number", "minimum": 1, @@ -231,6 +248,10 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: }, "comment": STRING, "audio_transcription_config": STRING, + "vertex_ai_audio_api": { + "type": "string", + "enum": ["lyria_predict", "lyria_interactions"], + }, } diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index d56f151a8c2..373f6c28f9e 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -500,7 +500,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): optional_params: dict, voice: str | dict | None = None, drop_params: bool = False, - kwargs: dict = {}, + kwargs: dict | None = None, ) -> tuple[str | None, dict]: mapped_params = dict(optional_params) base_model = model.removeprefix("vertex_ai/") diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index ab3b24f470f..4d348b51055 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -343,8 +343,7 @@ class VertexPassthroughLoggingHandler: json_response=json_response ) response_cost: Final = ( - VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) - or 0.0 + VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) or 0.0 ) * prediction_count logging_obj.model = model diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 9e370e5406a..58ca91f3977 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -53,6 +53,11 @@ "type": "number", "minimum": 0 }, + "audio_seconds_per_prediction": { + "type": "number", + "minimum": 0, + "description": "Audio duration, in seconds, produced by one prediction." + }, "audio_transcription_config": { "type": "string" }, @@ -363,6 +368,16 @@ "type": "string", "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers." }, + "max_audio_length_hours": { + "type": "number", + "minimum": 0, + "description": "Maximum generated audio duration, expressed in hours." + }, + "max_audio_per_prompt": { + "type": "integer", + "minimum": 0, + "description": "Maximum number of audio outputs accepted or generated per prompt." + }, "max_input_tokens": { "type": "integer", "minimum": 0, @@ -603,6 +618,17 @@ "type": "string", "description": "URL of the provider pricing/model page this entry was taken from." }, + "supported_audio_formats": { + "type": "array", + "description": "Audio container formats the model can return.", + "items": { + "type": "string", + "enum": [ + "mp3", + "wav" + ] + } + }, "supported_endpoints": { "type": "array", "description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.", @@ -826,6 +852,13 @@ "uses_embed_content": { "type": "boolean" }, + "vertex_ai_audio_api": { + "type": "string", + "enum": [ + "lyria_predict", + "lyria_interactions" + ] + }, "web_search_billing_unit": { "type": "string", "description": "Whether web search is billed per query or per prompt.", From 6ec53f284617cc52adfb78f5c8e0b03b5a3c125e Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 12 Aug 2026 13:33:33 -0500 Subject: [PATCH 14/63] style(vertex-ai): satisfy Lyria quality gates --- litellm/cost_calculator.py | 10 +- litellm/llms/vertex_ai/common_utils.py | 16 +- .../llms/vertex_ai/interactions/__init__.py | 2 +- .../text_to_speech/transformation.py | 147 +++++++++++------- litellm/main.py | 8 +- .../vertex_passthrough_logging_handler.py | 49 ++++-- litellm/types/utils.py | 10 +- 7 files changed, 153 insertions(+), 89 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 2bac5e234bc..5dcbb1d5f37 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -498,16 +498,14 @@ def cost_per_token( speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) prompt_cost: float = 0.0 completion_cost: float = 0.0 - if not speech_model_info.get("input_cost_per_character") and not speech_model_info.get( - "input_cost_per_token" - ): + if not speech_model_info.get("input_cost_per_character") and not speech_model_info.get("input_cost_per_token"): output_cost_per_generation: Final = speech_model_info.get("output_cost_per_image") - output_cost_per_second: Final = speech_model_info.get("output_cost_per_second") + speech_output_cost_per_second: Final = speech_model_info.get("output_cost_per_second") audio_seconds_per_prediction: Final = speech_model_info.get("audio_seconds_per_prediction") if output_cost_per_generation is not None: return prompt_cost, float(output_cost_per_generation) - if output_cost_per_second is not None and audio_seconds_per_prediction is not None: - return prompt_cost, float(output_cost_per_second) * float(audio_seconds_per_prediction) + if speech_output_cost_per_second is not None and audio_seconds_per_prediction is not None: + return prompt_cost, float(speech_output_cost_per_second) * float(audio_seconds_per_prediction) cost_metric: Final = select_cost_metric_for_model(speech_model_info) if cost_metric == "cost_per_character": if prompt_characters is None: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 8a4c1e68623..8885d19c1c0 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -6,7 +6,7 @@ from typing import Any, Final, Literal, cast, get_type_hints import httpx from pydantic import TypeAdapter, ValidationError -from typing_extensions import NotRequired, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -25,12 +25,12 @@ from litellm.utils import supports_response_schema, supports_system_messages class VertexAILyriaModelInfo(TypedDict): - vertex_ai_audio_api: Literal["lyria_predict", "lyria_interactions"] - supported_audio_formats: tuple[Literal["mp3", "wav"], ...] - output_cost_per_image: NotRequired[float] + vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"]] + supported_audio_formats: ReadOnly[tuple[Literal["mp3", "wav"], ...]] + output_cost_per_image: NotRequired[ReadOnly[float]] -_VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER = TypeAdapter(VertexAILyriaModelInfo) +_VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER: Final = TypeAdapter(VertexAILyriaModelInfo) def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyriaModelInfo | None: @@ -46,13 +46,13 @@ def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyri def _get_bundled_vertex_ai_lyria_model_info(model_key: str) -> VertexAILyriaModelInfo | None: from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap - bundled_model_info = GetModelCostMap.load_local_model_cost_map().get(model_key) + bundled_model_info: Final = GetModelCostMap.load_local_model_cost_map().get(model_key) return _validate_vertex_ai_lyria_model_info(bundled_model_info) def get_vertex_ai_lyria_model_info(model: str) -> VertexAILyriaModelInfo | None: - model_key = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" - runtime_model_info = _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key)) + model_key: Final = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" + runtime_model_info: Final = _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key)) return runtime_model_info or _get_bundled_vertex_ai_lyria_model_info(model_key) diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py index 3e2309d21ef..f6f86f65e87 100644 --- a/litellm/llms/vertex_ai/interactions/__init__.py +++ b/litellm/llms/vertex_ai/interactions/__init__.py @@ -1,3 +1,3 @@ from .transformation import VertexAIInteractionsConfig -__all__ = ["VertexAIInteractionsConfig"] +__all__ = ("VertexAIInteractionsConfig",) diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 373f6c28f9e..81f1cbd5157 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -8,7 +8,7 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s import base64 from collections.abc import Coroutine from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union +from typing import TYPE_CHECKING, Any, Final, TypeAlias, Union import httpx @@ -40,6 +40,10 @@ else: LiteLLMLoggingObj = Any HttpxBinaryResponseContent = Any +_LyriaVoice: TypeAlias = ( + str | dict | None +) # mutable-ok: inherited interface supports structured provider voice dictionaries + class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): """ @@ -486,26 +490,34 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): @staticmethod def _get_model_info(model: str) -> VertexAILyriaModelInfo: - model_info = get_vertex_ai_lyria_model_info(model=model) + model_info: Final = get_vertex_ai_lyria_model_info(model=model) if model_info is None: raise ValueError(f"Vertex AI model {model!r} does not declare a Lyria audio API") return model_info - def get_supported_openai_params(self, model: str) -> list: - return ["response_format"] + def get_supported_openai_params( + self, model: str + ) -> list: # mutable-ok: inherited provider interface returns a concrete parameter list + return [ # mutable-ok: inherited provider interface requires a concrete parameter list + "response_format" + ] def map_openai_params( self, model: str, - optional_params: dict, - voice: str | dict | None = None, + optional_params: dict, # mutable-ok: inherited provider interface accepts a concrete parameter dictionary + voice: _LyriaVoice = None, drop_params: bool = False, - kwargs: dict | None = None, - ) -> tuple[str | None, dict]: - mapped_params = dict(optional_params) - base_model = model.removeprefix("vertex_ai/") - model_info = self._get_model_info(model=model) - unsupported_params = [param for param in ("speed", "instructions") if mapped_params.get(param) is not None] + kwargs: dict | None = None, # mutable-ok: inherited provider interface accepts a concrete keyword dictionary + ) -> tuple[str | None, dict]: # mutable-ok: inherited provider interface returns concrete mapped parameters + mapped_params: Final = dict( # mutable-ok: mapping drops unsupported parameters before provider dispatch + optional_params + ) + base_model: Final = model.removeprefix("vertex_ai/") + model_info: Final = self._get_model_info(model=model) + unsupported_params: Final = tuple( + param for param in ("speed", "instructions") if mapped_params.get(param) is not None + ) if unsupported_params: if drop_params or litellm.drop_params: for param in unsupported_params: @@ -519,8 +531,8 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): "from the call, set `litellm.drop_params = True`" ), ) - response_format = mapped_params.get("response_format") - supported_formats = frozenset(model_info["supported_audio_formats"]) + response_format: Final = mapped_params.get("response_format") + supported_formats: Final = frozenset(model_info["supported_audio_formats"]) if response_format is not None and response_format not in supported_formats: if drop_params or litellm.drop_params: mapped_params.pop("response_format", None) @@ -539,17 +551,20 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): self, model: str, api_base: str | None, - litellm_params: dict, + litellm_params: dict, # mutable-ok: inherited provider interface accepts concrete LiteLLM parameters ) -> str: - base_model = model.removeprefix("vertex_ai/") - model_info = self._get_model_info(model=model) - project = self.safe_get_vertex_ai_project(litellm_params) - if project is None: - _, project = self._ensure_access_token( + base_model: Final = model.removeprefix("vertex_ai/") + model_info: Final = self._get_model_info(model=model) + configured_project: Final = self.safe_get_vertex_ai_project(litellm_params) + project: Final = ( + self._ensure_access_token( credentials=self.safe_get_vertex_ai_credentials(litellm_params), project_id=None, custom_llm_provider="vertex_ai", - ) + )[1] + if configured_project is None + else configured_project + ) if model_info["vertex_ai_audio_api"] == "lyria_interactions": from litellm.llms.vertex_ai.interactions.transformation import ( VertexAIInteractionsConfig, @@ -558,10 +573,13 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): return VertexAIInteractionsConfig().get_complete_url( api_base=api_base, model=base_model, - litellm_params={**litellm_params, "vertex_project": project}, + litellm_params={ # mutable-ok: interactions dispatch expects a concrete parameter dictionary + **litellm_params, + "vertex_project": project, + }, ) - location = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location() - base_url = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/") + location: Final = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location() + base_url: Final = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/") return f"{base_url}/v1/projects/{project}/locations/{location}/publishers/google/models/{base_model}:predict" def transform_text_to_speech_request( @@ -569,9 +587,9 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): model: str, input: str, voice: str | None, - optional_params: dict, - litellm_params: dict, - headers: dict, + optional_params: dict, # mutable-ok: inherited provider interface accepts concrete mapped parameters + litellm_params: dict, # mutable-ok: inherited provider interface accepts concrete LiteLLM parameters + headers: dict, # mutable-ok: inherited provider interface accepts and updates concrete HTTP headers ) -> TextToSpeechRequestData: access_token, project = self._ensure_access_token( credentials=self.safe_get_vertex_ai_credentials(litellm_params), @@ -579,23 +597,32 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): custom_llm_provider="vertex_ai", ) headers.update( - { + { # mutable-ok: HTTP dispatch requires a concrete header dictionary "Authorization": f"Bearer {access_token}", "x-goog-user-project": project, "Content-Type": "application/json", } ) - base_model = model.removeprefix("vertex_ai/") - model_info = self._get_model_info(model=model) + base_model: Final = model.removeprefix("vertex_ai/") + model_info: Final = self._get_model_info(model=model) if model_info["vertex_ai_audio_api"] == "lyria_predict": - request_body = { - "instances": [{"prompt": input}], - "parameters": {"sample_count": 1}, + request_body = { # mutable-ok: predict dispatch requires a concrete provider request dictionary; rebind-ok: exactly one provider API shape initializes the request + "instances": [ # mutable-ok: predict dispatch requires a concrete instances list + {"prompt": input} # mutable-ok: predict dispatch requires a concrete instance dictionary + ], + "parameters": { # mutable-ok: predict dispatch requires a concrete parameters dictionary + "sample_count": 1 + }, } else: - request_body = {"model": base_model, "input": input} + request_body = { # mutable-ok: interactions dispatch requires a concrete provider request dictionary; rebind-ok: exactly one provider API shape initializes the request + "model": base_model, + "input": input, + } if optional_params.get("response_format") == "wav": - request_body["response_format"] = { + request_body[ + "response_format" + ] = { # mutable-ok: interactions dispatch requires a nested response-format dictionary "type": "audio", "mime_type": "audio/wav", } @@ -609,33 +636,49 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): ) -> "HttpxBinaryResponseContent": from litellm.types.llms.openai import HttpxBinaryResponseContent - response_json = raw_response.json() - base_model = model.removeprefix("vertex_ai/") - model_info = self._get_model_info(model=model) - audio_data: str | None = None - mime_type: str | None = None + response_json: Final = raw_response.json() + base_model: Final = model.removeprefix("vertex_ai/") + model_info: Final = self._get_model_info(model=model) + audio_data: str | None = None # rebind-ok: response parsing discovers audio data in provider-specific shapes + mime_type: str | None = None # rebind-ok: response parsing discovers the MIME type beside the audio payload if model_info["vertex_ai_audio_api"] == "lyria_predict": - predictions = response_json.get("predictions") or [] + predictions: Final = response_json.get("predictions") or () if predictions: - audio_data = predictions[0].get("audioContent") or predictions[0].get("bytesBase64Encoded") - mime_type = predictions[0].get("mimeType") + audio_data = predictions[0].get("audioContent") or predictions[0].get( + "bytesBase64Encoded" + ) # rebind-ok: predict response supplies the generated audio value + mime_type = predictions[0].get("mimeType") # rebind-ok: predict response supplies its audio MIME type else: - for step in response_json.get("steps") or response_json.get("outputs") or []: - content_items = step.get("content") or [] if step.get("type") == "model_output" else [step] + for step in response_json.get("steps") or response_json.get("outputs") or (): + content_items = step.get("content") or () if step.get("type") == "model_output" else (step,) for content in content_items: if content.get("type") == "audio" and content.get("data"): - audio_data = content["data"] - mime_type = content.get("mime_type") + audio_data = content[ + "data" + ] # rebind-ok: interactions response supplies the generated audio value + mime_type = content.get( + "mime_type" + ) # rebind-ok: interactions response supplies its audio MIME type if audio_data is None: raise ValueError(f"No generated audio found in Vertex AI {base_model} response") - default_format = model_info["supported_audio_formats"][0] - mime_type = mime_type or {"mp3": "audio/mpeg", "wav": "audio/wav"}[default_format] - response = HttpxBinaryResponseContent( + default_format: Final = model_info["supported_audio_formats"][0] + mime_type = ( + mime_type + or { # mutable-ok: short-lived lookup selects the default response MIME type; rebind-ok: absent provider MIME type falls back to model metadata + "mp3": "audio/mpeg", + "wav": "audio/wav", + }[default_format] + ) + response: Final = HttpxBinaryResponseContent( httpx.Response( status_code=raw_response.status_code, content=base64.b64decode(audio_data), - headers={"content-type": mime_type}, + headers={ # mutable-ok: httpx requires a concrete response header dictionary + "content-type": mime_type + }, ) ) - response._hidden_params = {"audio_mime_type": mime_type} + response._hidden_params = { # mutable-ok: response metadata is a concrete dictionary + "audio_mime_type": mime_type + } return response diff --git a/litellm/main.py b/litellm/main.py index 55db92d44b3..040af897256 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8261,9 +8261,13 @@ def speech( # Vertex AI Text-to-Speech (Google Cloud TTS) if text_to_speech_provider_config is None: if VertexAILyriaTextToSpeechConfig.is_lyria_model(model): - text_to_speech_provider_config = VertexAILyriaTextToSpeechConfig() + text_to_speech_provider_config = ( + VertexAILyriaTextToSpeechConfig() + ) # rebind-ok: model metadata selects the Lyria provider implementation else: - text_to_speech_provider_config = VertexAITextToSpeechConfig() + text_to_speech_provider_config = ( + VertexAITextToSpeechConfig() + ) # rebind-ok: non-Lyria Vertex models use the standard TTS implementation # Cast to specific Vertex AI config type to access dispatch method vertex_config: Final = cast(VertexAITextToSpeechConfig, text_to_speech_provider_config) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 4d348b51055..53df964e541 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -334,10 +334,10 @@ class VertexPassthroughLoggingHandler: @staticmethod def _handle_audio_predict_response( - json_response: dict, + json_response: dict, # mutable-ok: passthrough logging receives the decoded provider response dictionary logging_obj: LiteLLMLoggingObj, model: str, - kwargs: dict, + kwargs: dict, # mutable-ok: passthrough logging enriches the shared callback metadata dictionary ) -> PassThroughEndpointLoggingTypedDict: prediction_count: Final = VertexPassthroughLoggingHandler._get_audio_prediction_count( json_response=json_response @@ -346,26 +346,41 @@ class VertexPassthroughLoggingHandler: VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) or 0.0 ) * prediction_count - logging_obj.model = model - logging_obj.model_call_details["model"] = model - logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" - logging_obj.custom_llm_provider = "vertex_ai" - logging_obj.model_call_details["response_cost"] = response_cost + logging_obj.model = model # rebind-ok: passthrough attribution records the resolved Vertex model + logging_obj.model_call_details[ # rebind-ok: passthrough attribution enriches callback metadata + "model" + ] = model + logging_obj.model_call_details[ # rebind-ok: passthrough attribution enriches callback metadata + "custom_llm_provider" + ] = "vertex_ai" + logging_obj.custom_llm_provider = ( # rebind-ok: attribution records the resolved provider + "vertex_ai" + ) + logging_obj.model_call_details[ # rebind-ok: passthrough attribution enriches callback metadata + "response_cost" + ] = response_cost - kwargs["response_cost"] = response_cost - kwargs["model"] = model - kwargs["custom_llm_provider"] = "vertex_ai" + kwargs[ # rebind-ok: callback metadata is enriched for downstream hooks + "response_cost" + ] = response_cost + kwargs["model"] = model # rebind-ok: callback metadata records the resolved model + kwargs["custom_llm_provider"] = "vertex_ai" # rebind-ok: callback metadata records the resolved provider - standard_pass_through_response_object: Final[StandardPassThroughResponseObject] = { + standard_pass_through_response_object: Final[ + StandardPassThroughResponseObject + ] = { # mutable-ok: callback contract requires a concrete response dictionary "response": json_response, } - return { + return { # mutable-ok: passthrough logging contract requires a concrete result dictionary "result": standard_pass_through_response_object, "kwargs": kwargs, } @staticmethod - def _is_audio_predict_response(model: str, json_response: dict) -> bool: + def _is_audio_predict_response( + model: str, + json_response: dict, # mutable-ok: predicate inspects the decoded provider response dictionary without mutation + ) -> bool: return ( VertexPassthroughLoggingHandler._get_audio_prediction_count(json_response=json_response) > 0 and VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) is not None @@ -373,7 +388,9 @@ class VertexPassthroughLoggingHandler: @staticmethod def _get_audio_prediction_unit_cost(model: str) -> float | None: - model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}", {}) + model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}") + if model_info is None: + return None output_cost_per_second: Final = model_info.get("output_cost_per_second") audio_seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction") if not isinstance(output_cost_per_second, (int, float)) or not isinstance( @@ -383,7 +400,9 @@ class VertexPassthroughLoggingHandler: return float(output_cost_per_second * audio_seconds_per_prediction) @staticmethod - def _get_audio_prediction_count(json_response: dict) -> int: + def _get_audio_prediction_count( + json_response: dict, # mutable-ok: counter inspects the decoded provider response dictionary without mutation + ) -> int: predictions: Final = json_response.get("predictions") if not isinstance(predictions, list): return 0 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index bb0485be7c1..6c645e70bee 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -168,8 +168,8 @@ class ProviderSpecificModelInfo(TypedDict, total=False): default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None] supports_output_config: bool | None supports_image_size: bool | None - supported_audio_formats: list[Literal["mp3", "wav"]] | None - vertex_ai_audio_api: Literal["lyria_predict", "lyria_interactions"] | None + supported_audio_formats: ReadOnly[Sequence[Literal["mp3", "wav"]] | None] + vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"] | None] bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None bedrock_converse_supports_strict_tools: bool | None @@ -312,9 +312,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_video_per_second: float | None # only for vertex ai models output_cost_per_audio_per_second: float | None # only for vertex ai models output_cost_per_second: float | None # for OpenAI Speech models - audio_seconds_per_prediction: float | None - max_audio_length_hours: float | None - max_audio_per_prompt: int | None + audio_seconds_per_prediction: ReadOnly[float | None] + max_audio_length_hours: ReadOnly[float | None] + max_audio_per_prompt: ReadOnly[int | None] output_cost_per_second_1080p: ( float | None ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) From abb655daa8851f15ce534a378761c268c811d942 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sun, 30 Aug 2026 09:58:28 -0500 Subject: [PATCH 15/63] fix(vertex-ai): encode Lyria predict URL path segments Percent-encode project, location, and model as single path segments. Lyria 3 speech builds the interactions URL through staging's minter with the already resolved project --- .../text_to_speech/transformation.py | 20 ++++++++++-- .../text_to_speech/test_transformation.py | 32 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 81f1cbd5157..17c4f0ed0c0 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -17,6 +17,7 @@ from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.audio_utils.utils import ( speech_media_type_from_audio_bytes, ) +from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -570,17 +571,32 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): VertexAIInteractionsConfig, ) - return VertexAIInteractionsConfig().get_complete_url( + resolved_project: Final = project + + def mint_access_token( + _credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + ) -> tuple[str, str]: + return "", project_id or resolved_project + + return VertexAIInteractionsConfig(mint_access_token=mint_access_token).get_complete_url( api_base=api_base, model=base_model, litellm_params={ # mutable-ok: interactions dispatch expects a concrete parameter dictionary **litellm_params, "vertex_project": project, + "vertex_location": "global", }, ) location: Final = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location() base_url: Final = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/") - return f"{base_url}/v1/projects/{project}/locations/{location}/publishers/google/models/{base_model}:predict" + encoded_project: Final = encode_url_path_segment(project, field_name="project") + encoded_location: Final = encode_url_path_segment(location, field_name="location") + encoded_model: Final = encode_url_path_segment(base_model, field_name="model") + return ( + f"{base_url}/v1/projects/{encoded_project}/locations/{encoded_location}" + f"/publishers/google/models/{encoded_model}:predict" + ) def transform_text_to_speech_request( self, diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index a1bb203e67e..3fc5278d3dd 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -1,4 +1,5 @@ import base64 +from typing import Final from unittest.mock import MagicMock, Mock, patch import httpx @@ -259,6 +260,37 @@ class TestVertexAILyriaTextToSpeechConfig: "locations/europe-west4/publishers/google/models/lyria-002:predict" ) + def test_get_complete_url_encodes_injected_predict_path_segments(self, monkeypatch: pytest.MonkeyPatch) -> None: + injected: Final = ( + "victim-project/locations/us-central1/publishers/google/models/other-model:predict?ignored=" + ) + encoded: Final = ( + "victim-project%2Flocations%2Fus-central1%2Fpublishers%2Fgoogle" + "%2Fmodels%2Fother-model%3Apredict%3Fignored%3D" + ) + monkeypatch.setitem( + litellm.model_cost, + f"vertex_ai/{injected}", + { + "vertex_ai_audio_api": "lyria_predict", + "supported_audio_formats": ["wav"], + }, + ) + + url: Final = VertexAILyriaTextToSpeechConfig().get_complete_url( + model=injected, + api_base="https://us-central1-aiplatform.googleapis.com", + litellm_params={ + "vertex_project": injected, + "vertex_location": injected, + }, + ) + + assert url == ( + "https://us-central1-aiplatform.googleapis.com" + f"/v1/projects/{encoded}/locations/{encoded}/publishers/google/models/{encoded}:predict" + ) + def test_get_complete_url_for_lyria_3(self): config = VertexAILyriaTextToSpeechConfig() From e18df8f09e7aef60ddb76979ec7a4ddeae5f5aeb Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sun, 30 Aug 2026 09:58:28 -0500 Subject: [PATCH 16/63] fix(vertex-ai): fall back to bundled Lyria costs Prefer runtime model_cost when both numeric fields are present, then bundled Lyria metadata so stale maps still bill 0.002 * 30 --- litellm/llms/vertex_ai/common_utils.py | 2 + .../vertex_passthrough_logging_handler.py | 20 ++++++++-- ...test_vertex_passthrough_logging_handler.py | 39 +++++++++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 8885d19c1c0..b649d8dac0e 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -28,6 +28,8 @@ class VertexAILyriaModelInfo(TypedDict): vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"]] supported_audio_formats: ReadOnly[tuple[Literal["mp3", "wav"], ...]] output_cost_per_image: NotRequired[ReadOnly[float]] + output_cost_per_second: NotRequired[ReadOnly[float]] + audio_seconds_per_prediction: NotRequired[ReadOnly[float]] _VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER: Final = TypeAdapter(VertexAILyriaModelInfo) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 53df964e541..f7625d71168 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -1,5 +1,6 @@ import asyncio import re +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse @@ -10,7 +11,10 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.vertex_ai.common_utils import get_vertex_location_from_url +from litellm.llms.vertex_ai.common_utils import ( + get_vertex_ai_lyria_model_info, + get_vertex_location_from_url, +) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as VertexModelResponseIterator, ) @@ -388,8 +392,18 @@ class VertexPassthroughLoggingHandler: @staticmethod def _get_audio_prediction_unit_cost(model: str) -> float | None: - model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}") - if model_info is None: + runtime_unit_cost: Final = VertexPassthroughLoggingHandler._audio_prediction_unit_cost_from_model_info( + model_info=litellm.model_cost.get(f"vertex_ai/{model}") + ) + if runtime_unit_cost is not None: + return runtime_unit_cost + return VertexPassthroughLoggingHandler._audio_prediction_unit_cost_from_model_info( + model_info=get_vertex_ai_lyria_model_info(model=model) + ) + + @staticmethod + def _audio_prediction_unit_cost_from_model_info(model_info: object) -> float | None: + if not isinstance(model_info, Mapping): return None output_cost_per_second: Final = model_info.get("output_cost_per_second") audio_seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index 1e8d569d3d1..f3d28d58c22 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -1,4 +1,5 @@ from datetime import datetime +from typing import Final from unittest.mock import MagicMock import httpx @@ -143,3 +144,41 @@ def test_audio_predict_response_supports_bytes_base64_encoded( assert result["kwargs"]["response_cost"] == pytest.approx(0.06) assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06) + + +def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_map_omits_model( + monkeypatch: pytest.MonkeyPatch, +) -> None: + stale_runtime_model_cost: Final = { + key: value for key, value in litellm.model_cost.items() if key != "vertex_ai/lyria-002" + } + monkeypatch.setattr(litellm, "model_cost", stale_runtime_model_cost) + logging_obj = MagicMock() + logging_obj.model_call_details = {} + response = httpx.Response( + status_code=200, + json={ + "predictions": [ + { + "audioContent": "clip", + "mimeType": "audio/wav", + } + ] + }, + ) + + result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=response, + logging_obj=logging_obj, + url_route="/v1/projects/test/locations/us-central1/publishers/google/models/lyria-002:predict", + result=response.text, + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"instances": [{"prompt": "ambient piano"}]}, + ) + + assert "vertex_ai/lyria-002" not in litellm.model_cost + assert result["kwargs"]["model"] == "lyria-002" + assert result["kwargs"]["response_cost"] == pytest.approx(0.06) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06) From c784e604acc7791e882391a512b4836366fef0e1 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sun, 30 Aug 2026 10:15:50 -0500 Subject: [PATCH 17/63] test(vertex-ai): drop Lyria auth patches from transform tests Subclass the Lyria transformer to stub token minting, and mark the remaining litellm.speech patches so TQ008 stays within budget. --- .../vertex_ai/text_to_speech/test_transformation.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index 3fc5278d3dd..457c3f76dbb 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -335,16 +335,17 @@ class TestVertexAILyriaTextToSpeechConfig: ), ], ) - @patch.object(VertexAILyriaTextToSpeechConfig, "_ensure_access_token") def test_transform_request( self, - mock_ensure_token, model, response_format, expected_body, ): - mock_ensure_token.return_value = ("mock-token", "music-project") - config = VertexAILyriaTextToSpeechConfig() + class _LyriaConfig(VertexAILyriaTextToSpeechConfig): + def _ensure_access_token(self, *args: object, **kwargs: object) -> tuple[str, str]: + return "mock-token", "music-project" + + config = _LyriaConfig() request = config.transform_text_to_speech_request( model=model, @@ -514,12 +515,12 @@ class TestVertexAILyriaTextToSpeechConfig: mock_response.status_code = 200 mock_response.json.return_value = response_json with ( - patch.object( + patch.object( # test-quality-ok: litellm.speech has no seam for Vertex token minting VertexAILyriaTextToSpeechConfig, "_ensure_access_token", return_value=("mock-token", "music-project"), ), - patch( + patch( # test-quality-ok: litellm.speech has no seam for the HTTP handler "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post", return_value=mock_response, ) as mock_post, From cab3801f23aff8c5ac73f0f38d212ab74c9392b6 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sun, 30 Aug 2026 12:00:42 -0500 Subject: [PATCH 18/63] fix(vertex-ai): simplify Lyria provider typing --- .../llms/vertex_ai/interactions/__init__.py | 3 -- .../text_to_speech/transformation.py | 30 ++++++++++--------- litellm/types/llms/openai.py | 3 ++ 3 files changed, 19 insertions(+), 17 deletions(-) diff --git a/litellm/llms/vertex_ai/interactions/__init__.py b/litellm/llms/vertex_ai/interactions/__init__.py index f6f86f65e87..e69de29bb2d 100644 --- a/litellm/llms/vertex_ai/interactions/__init__.py +++ b/litellm/llms/vertex_ai/interactions/__init__.py @@ -1,3 +0,0 @@ -from .transformation import VertexAIInteractionsConfig - -__all__ = ("VertexAIInteractionsConfig",) diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 17c4f0ed0c0..2047be2ea37 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -621,8 +621,8 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): ) base_model: Final = model.removeprefix("vertex_ai/") model_info: Final = self._get_model_info(model=model) - if model_info["vertex_ai_audio_api"] == "lyria_predict": - request_body = { # mutable-ok: predict dispatch requires a concrete provider request dictionary; rebind-ok: exactly one provider API shape initializes the request + request_body: Final[dict[str, object]] = ( # mutable-ok: HTTP dispatch requires a concrete provider payload + { # mutable-ok: predict dispatch requires a concrete provider request dictionary "instances": [ # mutable-ok: predict dispatch requires a concrete instances list {"prompt": input} # mutable-ok: predict dispatch requires a concrete instance dictionary ], @@ -630,18 +630,22 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): "sample_count": 1 }, } - else: - request_body = { # mutable-ok: interactions dispatch requires a concrete provider request dictionary; rebind-ok: exactly one provider API shape initializes the request + if model_info["vertex_ai_audio_api"] == "lyria_predict" + else { # mutable-ok: interactions dispatch requires a concrete provider request dictionary "model": base_model, "input": input, + **( + { # mutable-ok: interactions dispatch requires a nested response-format dictionary + "response_format": { # mutable-ok: interactions response format is a concrete provider payload + "type": "audio", + "mime_type": "audio/wav", + } + } + if optional_params.get("response_format") == "wav" + else {} # mutable-ok: no response override is merged for non-WAV output + ), } - if optional_params.get("response_format") == "wav": - request_body[ - "response_format" - ] = { # mutable-ok: interactions dispatch requires a nested response-format dictionary - "type": "audio", - "mime_type": "audio/wav", - } + ) return TextToSpeechRequestData(dict_body=request_body, headers=headers) def transform_text_to_speech_response( @@ -694,7 +698,5 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): }, ) ) - response._hidden_params = { # mutable-ok: response metadata is a concrete dictionary - "audio_mime_type": mime_type - } + response.set_audio_mime_type(mime_type) return response diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 32d88da0085..0b13191d977 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -119,6 +119,9 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent): return self._hidden_params["response_cost"] = response_cost + def set_audio_mime_type(self, audio_mime_type: str) -> None: + self._hidden_params["audio_mime_type"] = audio_mime_type + class NotGiven: """ From 1af9c229fa452c1a226bcc22f34f74cfffd54c54 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 17:29:56 -0700 Subject: [PATCH 19/63] ci(e2e): leave the managed-files opt-in file to its own lane instead of failing on an empty collection --- .github/workflows/test-e2e-changed.yml | 5 +++-- tests/e2e/CONTRIBUTING.md | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index b20c84af02d..c1a74f007ce 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -27,14 +27,15 @@ jobs: REPO: ${{ github.repository }} PR_NUMBER: ${{ github.event.pull_request.number }} SMOKE_TESTS: tests/e2e/access_control + OWN_LANE: '^tests/e2e/(ui|claude_code|load)/|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$' run: | files="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}/files" --paginate \ --jq '.[] | select(.status != "removed") | .filename')" tests="$(printf '%s\n' "${files}" \ | grep -E '^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$' \ - | grep -vE '^tests/e2e/(ui|claude_code|load)/' \ + | grep -vE "${OWN_LANE}" \ | sort -u | tr '\n' ' ' | sed 's/ $//')" || true - if [ -z "${tests}" ] && printf '%s\n' "${files}" | grep -vE '^tests/e2e/(ui|claude_code|load)/' \ + if [ -z "${tests}" ] && printf '%s\n' "${files}" | grep -vE "${OWN_LANE}" \ | grep -qE '^(tests/e2e/|\.github/e2e-stack/|\.github/workflows/test-e2e-changed\.yml$)'; then tests="${SMOKE_TESTS}" echo "harness or stack changed without a test file; running the smoke suite" diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index b984791c738..81fcdf84e74 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -54,7 +54,7 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT ### The pull request check -Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, and `load/`, which have their own lanes) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down +Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, `load/`, and `batches/test_managed_files_enforcement_e2e.py`, which have their own lanes or need a differently configured proxy) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down Credentials: the lane never sees the stage keys. Everything the runner writes into `tests/e2e/.env` comes from two AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`, and the first holds only the names that the tests and `tests/e2e/gateway/stage_mirror_ci_config.yml` read. Each credential in it is dedicated to this lane and can do only what the tests do: the AWS pair belongs to the `litellm-e2e-changed` IAM user, whose inline policy allows Bedrock inference, the one guardrail the tests use, the batch job APIs, the batch S3 bucket, and assuming the batch test's role; the Vertex key belongs to a service account holding only `roles/aiplatform.user` on the Vertex project; the OpenAI, Anthropic, Gemini, Cohere, and Devin keys are per-lane keys with hard monthly spend caps; and the Datadog application key is scoped to log search. A leaked key can therefore spend at most one month of a small capped budget or call a handful of Bedrock APIs, and rotating one is a single `aws secretsmanager put-secret-value` on the secret, with no workflow change From a3ee6b25667d7640214465b58364d04554c02f20 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 17:50:35 -0700 Subject: [PATCH 20/63] ci(e2e): fail a pass whose every collected test was skipped --- .github/e2e-stack/assert_tests_ran.py | 21 +++++++++++++++++++++ .github/workflows/test-e2e-changed.yml | 4 +++- tests/e2e/CONTRIBUTING.md | 2 +- 3 files changed, 25 insertions(+), 2 deletions(-) create mode 100644 .github/e2e-stack/assert_tests_ran.py diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py new file mode 100644 index 00000000000..092a9db7256 --- /dev/null +++ b/.github/e2e-stack/assert_tests_ran.py @@ -0,0 +1,21 @@ +import sys +import xml.etree.ElementTree as ET +from pathlib import Path +from typing import Final + + +def main() -> int: + report: Final = ET.parse(Path(sys.argv[1])).getroot() + suites: Final = tuple(report.iter("testsuite")) + collected: Final = sum(int(suite.get("tests", "0")) for suite in suites) + skipped: Final = sum(int(suite.get("skipped", "0")) for suite in suites) + executed: Final = collected - skipped + _ = sys.stdout.write(f"executed {executed} of {collected} collected tests ({skipped} skipped)\n") + if executed > 0: + return 0 + _ = sys.stdout.write("::error::every selected test was skipped, so nothing was verified\n") + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index c1a74f007ce..e08660d0412 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -152,9 +152,10 @@ jobs: run: | read -r -a test_files <<< "${TESTS}" for pass in 1 2 3; do + report="${RUNNER_TEMP}/e2e-pass-${pass}.xml" echo "::group::pass ${pass} of 3" set +e - uv run --no-sync pytest "${test_files[@]}" --reruns 0 -v -rA --tb=short -p no:cacheprovider + uv run --no-sync pytest "${test_files[@]}" --reruns 0 -v -rA --tb=short -p no:cacheprovider --junitxml="${report}" status=$? set -e echo "::endgroup::" @@ -166,6 +167,7 @@ jobs: echo "::error::pass ${pass} of 3 failed with exit code ${status}" exit "${status}" fi + uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" done - name: Show stack logs on failure diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 81fcdf84e74..ef6fbb5405d 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -54,7 +54,7 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT ### The pull request check -Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, `load/`, and `batches/test_managed_files_enforcement_e2e.py`, which have their own lanes or need a differently configured proxy) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down +Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, `load/`, and `batches/test_managed_files_enforcement_e2e.py`, which have their own lanes or need a differently configured proxy) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. A pass that executes nothing is also red: each pass writes a JUnit report and fails when every collected test was skipped, so a skipped-out file cannot pass on paper. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down Credentials: the lane never sees the stage keys. Everything the runner writes into `tests/e2e/.env` comes from two AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`, and the first holds only the names that the tests and `tests/e2e/gateway/stage_mirror_ci_config.yml` read. Each credential in it is dedicated to this lane and can do only what the tests do: the AWS pair belongs to the `litellm-e2e-changed` IAM user, whose inline policy allows Bedrock inference, the one guardrail the tests use, the batch job APIs, the batch S3 bucket, and assuming the batch test's role; the Vertex key belongs to a service account holding only `roles/aiplatform.user` on the Vertex project; the OpenAI, Anthropic, Gemini, Cohere, and Devin keys are per-lane keys with hard monthly spend caps; and the Datadog application key is scoped to log search. A leaked key can therefore spend at most one month of a small capped budget or call a handful of Bedrock APIs, and rotating one is a single `aws secretsmanager put-secret-value` on the secret, with no workflow change From a9bef8d370ebd553071729b32a9f27f729f1d17e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 2 Sep 2026 18:01:19 -0700 Subject: [PATCH 21/63] docs(e2e): drop the spend-cap claim from the credentials paragraph --- tests/e2e/CONTRIBUTING.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index ef6fbb5405d..1e88f6505cf 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -56,7 +56,7 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, `load/`, and `batches/test_managed_files_enforcement_e2e.py`, which have their own lanes or need a differently configured proxy) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. A pass that executes nothing is also red: each pass writes a JUnit report and fails when every collected test was skipped, so a skipped-out file cannot pass on paper. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down -Credentials: the lane never sees the stage keys. Everything the runner writes into `tests/e2e/.env` comes from two AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`, and the first holds only the names that the tests and `tests/e2e/gateway/stage_mirror_ci_config.yml` read. Each credential in it is dedicated to this lane and can do only what the tests do: the AWS pair belongs to the `litellm-e2e-changed` IAM user, whose inline policy allows Bedrock inference, the one guardrail the tests use, the batch job APIs, the batch S3 bucket, and assuming the batch test's role; the Vertex key belongs to a service account holding only `roles/aiplatform.user` on the Vertex project; the OpenAI, Anthropic, Gemini, Cohere, and Devin keys are per-lane keys with hard monthly spend caps; and the Datadog application key is scoped to log search. A leaked key can therefore spend at most one month of a small capped budget or call a handful of Bedrock APIs, and rotating one is a single `aws secretsmanager put-secret-value` on the secret, with no workflow change +Credentials: everything the runner writes into `tests/e2e/.env` comes from two AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`, and the first holds only the names that the tests and `tests/e2e/gateway/stage_mirror_ci_config.yml` read. The cloud credentials in it are dedicated to this lane and can do only what the tests do: the AWS pair belongs to the `litellm-e2e-changed` IAM user, whose inline policy allows Bedrock inference, the one guardrail the tests use, the batch job APIs, the batch S3 bucket, and assuming the batch test's role, and the Vertex key belongs to a service account holding only `roles/aiplatform.user` on the Vertex project. The provider API keys carry no lane-specific spend cap, since a cap that trips mid-month would fail every run until it resets, so the reviewer's approval of the `e2e-changed` environment is the control on how a PR's tests use them. Rotating any value is a single `aws secretsmanager put-secret-value` on the secret, with no workflow change ### Record and replay From 200e2901d661d9bcea83d6e69c3839ea71b07e2c Mon Sep 17 00:00:00 2001 From: Atharva-Kanherkar <142440039+Atharva-Kanherkar@users.noreply.github.com> Date: Fri, 4 Sep 2026 14:33:10 +0530 Subject: [PATCH 22/63] fix(anthropic_responses): preserve Responses refusal blocks in Anthropic translation When OpenAI Responses returns a refusal content block, Anthropic /v1/messages erased the refusal text into an empty content array and emitted stop_reason 'end_turn'. Translate refusal blocks to Anthropic text blocks, map stop_reason to 'refusal', and add 'refusal' to AnthropicFinishReason. Fixes #39721 --- .../responses_adapters/streaming_iterator.py | 22 +++++- .../responses_adapters/transformation.py | 35 +++++++-- litellm/types/llms/anthropic.py | 2 +- ...t_responses_adapters_streaming_iterator.py | 34 +++++++++ .../test_responses_adapters_transformation.py | 76 +++++++++++++------ 5 files changed, 138 insertions(+), 31 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index a97ce18d179..5f6c5bad190 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -111,7 +111,7 @@ class AnthropicResponsesStreamWrapper: item_type: Final = getattr(item, "type", None) or (item.get("type") if isinstance(item, dict) else None) item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) - if item_type == "message": + if item_type in ("message", "refusal"): self._open_block(item_id, {"type": "text", "text": ""}) elif item_type == "function_call": call_id: Final = ( @@ -132,7 +132,7 @@ class AnthropicResponsesStreamWrapper: return # ---- text delta ---- - if event_type == "response.output_text.delta": + if event_type in ("response.output_text.delta", "response.refusal.delta"): item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index @@ -238,6 +238,24 @@ class AnthropicResponsesStreamWrapper: if out_type == "function_call": stop_reason = "tool_use" break + elif out_type == "refusal": + stop_reason = "refusal" + break + elif out_type == "message": + content_parts = getattr(out_item, "content", ()) or ( + out_item.get("content") or () if isinstance(out_item, dict) else () + ) + for part in content_parts: + part_type = getattr(part, "type", None) or ( + part.get("type") if isinstance(part, dict) else None + ) + if ( + part_type == "refusal" + or hasattr(part, "refusal") + or (isinstance(part, dict) and "refusal" in part) + ): + stop_reason = "refusal" + break self._chunk_queue.append( { diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 0eb0e38a46e..6cabd35117d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -631,10 +631,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter: elif isinstance(item, ResponseOutputMessage): for part in item.content: - if getattr(part, "type", None) == "output_text": + part_type = getattr(part, "type", None) + if part_type == "output_text": content.append( AnthropicResponseContentBlockText(type="text", text=getattr(part, "text", "")).model_dump() ) + elif part_type == "refusal" or hasattr(part, "refusal"): + content.append( + AnthropicResponseContentBlockText( + type="text", text=getattr(part, "refusal", "") or "" + ).model_dump() + ) + stop_reason = "refusal" elif isinstance(item, ResponseFunctionToolCall): try: @@ -654,11 +662,22 @@ class LiteLLMAnthropicToResponsesAPIAdapter: elif isinstance(item, dict): item_type = item.get("type") if item_type == "message": - for part in item.get("content", []): - if isinstance(part, dict) and part.get("type") == "output_text": - content.append( - AnthropicResponseContentBlockText(type="text", text=part.get("text", "")).model_dump() - ) + for part in item.get("content", ()): + if isinstance(part, dict): + part_type = part.get("type") + if part_type == "output_text": + content.append( + AnthropicResponseContentBlockText( + type="text", text=part.get("text", "") + ).model_dump() + ) + elif part_type == "refusal" or "refusal" in part: + content.append( + AnthropicResponseContentBlockText( + type="text", text=part.get("refusal", "") or "" + ).model_dump() + ) + stop_reason = "refusal" elif item_type == "reasoning": content.extend( self._thinking_blocks_from_reasoning_item( @@ -679,6 +698,10 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ).model_dump() ) stop_reason = "tool_use" + elif item_type == "refusal": + refusal_text = item.get("refusal") or item.get("text", "") or "" + content.append(AnthropicResponseContentBlockText(type="text", text=refusal_text).model_dump()) + stop_reason = "refusal" # status -> stop_reason override if response.status == "incomplete": diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index b3462203c4b..f1107c87c73 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -658,7 +658,7 @@ class AnthropicOutputTokensDetails(BaseModel): thinking_tokens: int | None = None -AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] +AnthropicFinishReason = Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] class AnthropicResponse(BaseModel): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index aebbed88c70..77ecfb73ff6 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -308,3 +308,37 @@ class TestResponseCompletedUsage: "cache_creation_input_tokens": 10, "cache_read_input_tokens": 4004, } + + +class TestRefusalStreamEvents: + def test_refusal_delta_emits_text_delta(self): + chunks = _process_all( + [ + {"type": "response.created"}, + {"type": "response.refusal.delta", "item_id": "ref_1", "delta": "I cannot fulfill this."}, + ] + ) + assert any( + c.get("type") == "content_block_delta" and c.get("delta", {}).get("text") == "I cannot fulfill this." + for c in chunks + ) + + def test_response_completed_with_refusal_sets_stop_reason_refusal(self): + response = SimpleNamespace( + status="completed", + output=[{"type": "message", "content": [{"type": "refusal", "refusal": "Policy violation"}]}], + usage=None, + ) + chunks = _process_all([{"type": "response.completed", "response": response}]) + message_delta = next(c for c in chunks if c["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + + def test_response_completed_with_standalone_refusal_item_sets_stop_reason_refusal(self): + response = SimpleNamespace( + status="completed", + output=[{"type": "refusal", "refusal": "Standalone refusal"}], + usage=None, + ) + chunks = _process_all([{"type": "response.completed", "response": response}]) + message_delta = next(c for c in chunks if c["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index 5ecf604f096..f48c31ea5ed 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -147,9 +147,7 @@ class TestOutputConfigStructuredOutput: def test_output_config_format_explicit_strict_true_is_preserved(self): """Nested output_config.format with explicit strict=True is preserved.""" - req = _make_request( - output_config={"format": {"type": "json_schema", "schema": self._SCHEMA, "strict": True}} - ) + req = _make_request(output_config={"format": {"type": "json_schema", "schema": self._SCHEMA, "strict": True}}) kwargs = _ADAPTER.translate_request(req) assert kwargs["text"]["format"]["strict"] is True @@ -1207,6 +1205,18 @@ def _make_output_message(texts: List[str]) -> MagicMock: return msg +def _make_refusal_message(refusal_text: str) -> MagicMock: + from openai.types.responses import ResponseOutputMessage + + part = MagicMock() + part.type = "refusal" + part.refusal = refusal_text + + msg = MagicMock(spec=ResponseOutputMessage) + msg.content = [part] + return msg + + def _make_function_call_item(call_id: str, name: str, arguments: str) -> MagicMock: """Build a mock ResponseFunctionToolCall.""" from openai.types.responses import ResponseFunctionToolCall # type: ignore[import] @@ -1279,6 +1289,38 @@ class TestTranslateResponse: result: Any = _ADAPTER.translate_response(response) assert result["stop_reason"] == "end_turn" + def test_refusal_part_becomes_text_block_and_sets_stop_reason_refusal(self): + response = _make_mock_response(output=[_make_refusal_message("I cannot fulfill this request.")]) + result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "I cannot fulfill this request." + assert result["stop_reason"] == "refusal" + + def test_dict_refusal_part_in_message_becomes_text_block(self): + output_item = { + "type": "message", + "content": [{"type": "refusal", "refusal": "Refused by policy"}], + } + response = _make_mock_response(output=[output_item]) + result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Refused by policy" + assert result["stop_reason"] == "refusal" + + def test_dict_refusal_item_becomes_text_block(self): + output_item = { + "type": "refusal", + "refusal": "Standalone refusal", + } + response = _make_mock_response(output=[output_item]) + result: Any = _ADAPTER.translate_response(response) + assert len(result["content"]) == 1 + assert result["content"][0]["type"] == "text" + assert result["content"][0]["text"] == "Standalone refusal" + assert result["stop_reason"] == "refusal" + def test_incomplete_status_sets_max_tokens(self): """status='incomplete' overrides stop_reason to 'max_tokens'.""" response = _make_mock_response( @@ -1337,9 +1379,7 @@ class TestTranslateResponse: ] ) result: Any = _ADAPTER.translate_response(response) - assert result["content"] == [ - {"type": "thinking", "thinking": "Weighing the options.", "signature": None} - ] + assert result["content"] == [{"type": "thinking", "thinking": "Weighing the options.", "signature": None}] def test_thinking_blocks_are_dropped_when_replayed_to_anthropic(self): """Replaying this turn to an Anthropic model must not send a signature it cannot verify.""" @@ -1481,9 +1521,7 @@ class TestToolResultImages: }, { "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content} - ], + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}], }, ] @@ -1630,9 +1668,7 @@ class TestToolResultDocuments: }, { "role": "user", - "content": [ - {"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content} - ], + "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": tool_result_content}], }, ] @@ -1665,9 +1701,7 @@ class TestToolResultDocuments: def test_document_title_becomes_filename(self): output = self._tool_output(self._translate([self._base64_document(title="quarterly-report.pdf")])) - assert output == [ - {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI} - ] + assert output == [{"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI}] def test_url_document_becomes_file_url_part(self): output = self._tool_output( @@ -1776,9 +1810,7 @@ class TestUserContentDocuments: def test_document_title_becomes_filename(self): content = self._user_content(self._translate([self._base64_document(title="quarterly-report.pdf")])) - assert content == [ - {"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI} - ] + assert content == [{"type": "input_file", "filename": "quarterly-report.pdf", "file_data": self.PDF_DATA_URI}] def test_url_document_becomes_file_url_part(self): content = self._user_content( @@ -1808,9 +1840,7 @@ class TestUserContentDocuments: assert content == [{"type": "input_text", "text": "still here"}] def test_document_breakpoint_rides_on_the_file_part(self): - content = self._user_content( - self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)]) - ) + content = self._user_content(self._translate([self._base64_document(prompt_cache_breakpoint=self.EXPLICIT)])) assert content == [ { "type": "input_file", @@ -1857,7 +1887,9 @@ class TestPromptCacheBreakpointToResponses: ] def test_system_without_breakpoint_still_becomes_instructions(self): - request = _make_request(system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}]) + request = _make_request( + system=[{"type": "text", "text": "Be concise."}, {"type": "text", "text": "Be helpful."}] + ) kwargs = _ADAPTER.translate_request(request) assert kwargs["instructions"] == "Be concise.\nBe helpful." assert kwargs["input"] == [ From 0b34abe8fe3904453ff4796cfdb8b8981dfaf7e8 Mon Sep 17 00:00:00 2001 From: Atharva-Kanherkar <142440039+Atharva-Kanherkar@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:02:07 +0530 Subject: [PATCH 23/63] fix(anthropic): harden refusal translation --- .../adapters/streaming_iterator.py | 39 +++++- .../adapters/transformation.py | 37 +++++- .../responses_adapters/streaming_iterator.py | 114 +++++++++++------- .../responses_adapters/transformation.py | 58 +++++++-- litellm/types/llms/anthropic.py | 7 ++ .../anthropic_messages/anthropic_response.py | 11 +- ...al_pass_through_adapters_transformation.py | 45 +++++++ .../test_streaming_iterator_first_delta.py | 87 +++++++++++++ ...t_responses_adapters_streaming_iterator.py | 60 +++++++-- .../test_responses_adapters_transformation.py | 41 ++++--- 10 files changed, 397 insertions(+), 102 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 78ff83cafbf..41db4f12143 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -4,13 +4,14 @@ import copy import json import traceback from collections import deque -from collections.abc import AsyncIterator, Iterator, Sequence +from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from typing import ( TYPE_CHECKING, Any, Final, Literal, Protocol, + cast, get_args, ) @@ -305,6 +306,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Synthesized compaction block from compact_20260112 polyfill (streaming). self.compaction_block = compaction_block self.iterations_usage = iterations_usage + self._refusal_text_parts: list[str] = [] self.sent_compaction_block: bool = False # Per-phase flags so the compaction block's start/delta/stop events # are emitted (and the public state machine is advanced) in @@ -572,6 +574,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): current_content_block_index=self.current_content_block_index, applied_edits=(self.applied_edits if is_final_chunk and not will_merge_into_held else None), ) + processed_chunk = self._with_refusal_stop_details(processed_chunk) # Check if this is a usage chunk and we have a held stop_reason chunk if will_merge_into_held: @@ -806,6 +809,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): current_content_block_index=self.current_content_block_index, applied_edits=(self.applied_edits if is_final_chunk and not will_merge_into_held else None), ) + processed_chunk = self._with_refusal_stop_details(processed_chunk) # Check if this is a usage chunk and we have a held stop_reason chunk if will_merge_into_held: @@ -993,6 +997,31 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): def _increment_content_block_index(self): self.current_content_block_index += 1 + def _with_refusal_stop_details( + self, + processed_chunk: ContentBlockDelta | MessageBlockDelta, + ) -> ContentBlockDelta | MessageBlockDelta: + if processed_chunk.get("type") != "message_delta" or not self._refusal_text_parts: + return processed_chunk + delta: Final = cast(Mapping[str, object], processed_chunk["delta"]) + if delta.get("stop_reason") == "max_tokens": + return processed_chunk + return cast( + ContentBlockDelta | MessageBlockDelta, + { + **processed_chunk, + "delta": { + **delta, + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": None, + "explanation": "".join(self._refusal_text_parts), + }, + }, + }, + ) + @staticmethod def _delta_has_content(processed_chunk: dict[str, Any]) -> bool: """Return True if a translated chunk carries a non-empty @@ -1044,6 +1073,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): return False if getattr(delta, "content", None): return False + if getattr(delta, "refusal", None): + return False if getattr(delta, "reasoning_content", None): return False # thinking_blocks whose entries are all empty AND unsigned must not @@ -1069,8 +1100,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): """ from .transformation import LiteLLMAnthropicMessagesAdapter - # Example logic - customize based on your needs: - # If chunk indicates a tool call + refusal_text: Final = LiteLLMAnthropicMessagesAdapter._refusal_text(chunk.choices[0].delta) + if refusal_text is not None: + self._refusal_text_parts.append(refusal_text) + if chunk.choices[0].finish_reason is not None: return False diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 573a461e89e..5655f59a4cc 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -307,6 +307,17 @@ class LiteLLMAnthropicMessagesAdapter: def __init__(self): pass + @staticmethod + def _refusal_text(message_or_delta: object) -> str | None: + refusal: Final = getattr(message_or_delta, "refusal", None) + if isinstance(refusal, str): + return refusal + provider_specific_fields: Final = getattr(message_or_delta, "provider_specific_fields", None) + if isinstance(provider_specific_fields, Mapping): + provider_refusal: Final = provider_specific_fields.get("refusal") + return provider_refusal if isinstance(provider_refusal, str) else None + return None + ### FOR [BETA] `/v1/messages` endpoint support def _extract_signature_from_tool_call(self, tool_call: object) -> str | None: @@ -1324,6 +1335,8 @@ class LiteLLMAnthropicMessagesAdapter: new_content.append( AnthropicResponseContentBlockText(type="text", text=choice.message.content).model_dump() ) + if (refusal_text := self._refusal_text(choice.message)) is not None: + new_content.append(AnthropicResponseContentBlockText(type="text", text=refusal_text).model_dump()) # Handle tool calls (in parallel to text content) if choice.message.tool_calls is not None and len(choice.message.tool_calls) > 0: for tool_call in choice.message.tool_calls: @@ -1482,14 +1495,23 @@ class LiteLLMAnthropicMessagesAdapter: choices=response.choices, tool_name_mapping=tool_name_mapping, ) + refusal_text: Final = next( + (text for choice in response.choices if (text := self._refusal_text(choice.message)) is not None), + None, + ) if polyfill_result is not None and polyfill_result.compaction_block is not None: anthropic_content.insert(0, polyfill_result.compaction_block) ## extract finish reason - anthropic_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic( + translated_finish_reason: Final = self._translate_openai_finish_reason_to_anthropic( openai_finish_reason=response.choices[0].finish_reason ) + anthropic_finish_reason: Final = ( + "refusal" + if refusal_text is not None and translated_finish_reason != "max_tokens" + else translated_finish_reason + ) # extract usage usage: Final[Usage] = getattr(response, "usage") anthropic_usage: Final = self._translate_openai_usage_to_anthropic_usage(usage) @@ -1511,6 +1533,15 @@ class LiteLLMAnthropicMessagesAdapter: usage=anthropic_usage, content=anthropic_content, stop_reason=anthropic_finish_reason, + stop_details=( + { + "type": "refusal", + "category": None, + "explanation": refusal_text, + } + if anthropic_finish_reason == "refusal" + else None + ), ) applied_edits: Final = polyfill_result.applied_edits_for_response() if polyfill_result else None @@ -1551,7 +1582,9 @@ class LiteLLMAnthropicMessagesAdapter: "signature": thought_sig, } return "tool_use", cast("ContentBlockContentBlockDict", tool_block) - elif choice.delta.content is not None and len(choice.delta.content) > 0: + elif (choice.delta.content is not None and len(choice.delta.content) > 0) or self._refusal_text( + choice.delta + ) is not None: return "text", TextBlock(type="text", text="") elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"): thinking_blocks = choice.delta.thinking_blocks or [] diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 5f6c5bad190..1b0ab5923d8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -1,5 +1,6 @@ # What is this? ## Translates OpenAI call to Anthropic `/v1/messages` format +import asyncio import json import traceback from collections import deque @@ -49,6 +50,8 @@ class AnthropicResponsesStreamWrapper: self._sent_message_start = False self._sent_message_stop = False self._chunk_queue: deque = deque() + self._refusal_text_parts: list[str] = [] + self._sync_responses_iterator: Any = None def _make_message_start(self) -> dict[str, Any]: return { @@ -111,7 +114,7 @@ class AnthropicResponsesStreamWrapper: item_type: Final = getattr(item, "type", None) or (item.get("type") if isinstance(item, dict) else None) item_id = getattr(item, "id", None) or (item.get("id") if isinstance(item, dict) else None) - if item_type in ("message", "refusal"): + if item_type == "message": self._open_block(item_id, {"type": "text", "text": ""}) elif item_type == "function_call": call_id: Final = ( @@ -131,8 +134,14 @@ class AnthropicResponsesStreamWrapper: ) return + if event_type == "response.refusal.delta": + delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") + if isinstance(delta, str): + self._refusal_text_parts.append(delta) + return + # ---- text delta ---- - if event_type in ("response.output_text.delta", "response.refusal.delta"): + if event_type == "response.output_text.delta": item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index @@ -215,52 +224,53 @@ class AnthropicResponsesStreamWrapper: response_obj: Final = getattr(event, "response", None) or ( event.get("response") if isinstance(event, dict) else None ) - stop_reason = "end_turn" - anthropic_usage: AnthropicUsage = AnthropicUsage(input_tokens=0, output_tokens=0) - - if response_obj is not None: - status: Final = getattr(response_obj, "status", None) - if status == "incomplete": - stop_reason = "max_tokens" - anthropic_usage = ( - LiteLLMAnthropicToResponsesAPIAdapter.translate_responses_api_usage_to_anthropic_usage( - getattr(response_obj, "usage", None) - ) + output: Final = (getattr(response_obj, "output", None) or ()) if response_obj is not None else () + refusal_text: Final = LiteLLMAnthropicToResponsesAPIAdapter._refusal_text_from_output(output) or ( + "".join(self._refusal_text_parts) or None + ) + status: Final = getattr(response_obj, "status", None) if response_obj is not None else None + has_tool_call: Final = any( + getattr(item, "type", None) == "function_call" + or (isinstance(item, dict) and item.get("type") == "function_call") + for item in output + ) + stop_reason: Final = ( + "max_tokens" + if status == "incomplete" + else "refusal" + if refusal_text is not None + else "tool_use" + if has_tool_call + else "end_turn" + ) + anthropic_usage: Final[AnthropicUsage] = ( + LiteLLMAnthropicToResponsesAPIAdapter.translate_responses_api_usage_to_anthropic_usage( + getattr(response_obj, "usage", None) ) + if response_obj is not None + else AnthropicUsage(input_tokens=0, output_tokens=0) + ) - # Check if tool_use was in the output to override stop_reason - if response_obj is not None: - output: Final = getattr(response_obj, "output", []) or [] - for out_item in output: - out_type = getattr(out_item, "type", None) or ( - out_item.get("type") if isinstance(out_item, dict) else None - ) - if out_type == "function_call": - stop_reason = "tool_use" - break - elif out_type == "refusal": - stop_reason = "refusal" - break - elif out_type == "message": - content_parts = getattr(out_item, "content", ()) or ( - out_item.get("content") or () if isinstance(out_item, dict) else () - ) - for part in content_parts: - part_type = getattr(part, "type", None) or ( - part.get("type") if isinstance(part, dict) else None - ) - if ( - part_type == "refusal" - or hasattr(part, "refusal") - or (isinstance(part, dict) and "refusal" in part) - ): - stop_reason = "refusal" - break + message_delta_payload: Final = { + "stop_reason": stop_reason, + "stop_sequence": None, + **( + { + "stop_details": { + "type": "refusal", + "category": None, + "explanation": refusal_text, + } + } + if stop_reason == "refusal" + else {} + ), + } self._chunk_queue.append( { "type": "message_delta", - "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "delta": message_delta_payload, "usage": dict(anthropic_usage), } ) @@ -284,10 +294,22 @@ class AnthropicResponsesStreamWrapper: # Consume the upstream stream try: - async for event in self.responses_stream: - self._process_event(event) - if self._chunk_queue: - return self._chunk_queue.popleft() + if hasattr(self.responses_stream, "__aiter__"): + async for event in self.responses_stream: + self._process_event(event) + if self._chunk_queue: + return self._chunk_queue.popleft() + else: + if self._sync_responses_iterator is None: + self._sync_responses_iterator = iter(self.responses_stream) + missing: Final = object() + while True: + event = await asyncio.to_thread(next, self._sync_responses_iterator, missing) + if event is missing: + break + self._process_event(event) + if self._chunk_queue: + return self._chunk_queue.popleft() except StopAsyncIteration: pass except Exception as e: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 6cabd35117d..92f3e08f24d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -6,7 +6,7 @@ path used for OpenAI and Azure models. """ import json -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from itertools import groupby from typing import Any, Final, cast @@ -69,6 +69,38 @@ class LiteLLMAnthropicToResponsesAPIAdapter: chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage) return LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage(chat_usage) + @staticmethod + def _refusal_text_from_output(output: Iterable[object]) -> str | None: + from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal + + def refusal_text_from_item(item: object) -> str | None: + if isinstance(item, ResponseOutputMessage): + return next( + (part.refusal for part in item.content if isinstance(part, ResponseOutputRefusal)), + None, + ) + if not isinstance(item, Mapping): + return None + item_mapping: Final = cast(Mapping[str, object], item) + raw_parts: Final = item_mapping.get("content") + if item_mapping.get("type") != "message" or not isinstance(raw_parts, Sequence): + return None + return next( + ( + refusal + for part in cast(Sequence[object], raw_parts) + if isinstance(part, Mapping) + and cast(Mapping[str, object], part).get("type") == "refusal" + and isinstance((refusal := cast(Mapping[str, object], part).get("refusal")), str) + ), + None, + ) + + return next( + (text for item in output if (text := refusal_text_from_item(item)) is not None), + None, + ) + # ------------------------------------------------------------------ # # Request translation: Anthropic -> Responses API # # ------------------------------------------------------------------ # @@ -624,6 +656,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content: Final[list[dict[str, object]]] = [] stop_reason: AnthropicFinishReason = "end_turn" + refusal_text: Final = self._refusal_text_from_output(cast(Iterable[object], response.output)) for item in response.output: if isinstance(item, ResponseReasoningItem): @@ -636,13 +669,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content.append( AnthropicResponseContentBlockText(type="text", text=getattr(part, "text", "")).model_dump() ) - elif part_type == "refusal" or hasattr(part, "refusal"): + elif part_type == "refusal": content.append( AnthropicResponseContentBlockText( type="text", text=getattr(part, "refusal", "") or "" ).model_dump() ) - stop_reason = "refusal" elif isinstance(item, ResponseFunctionToolCall): try: @@ -671,13 +703,12 @@ class LiteLLMAnthropicToResponsesAPIAdapter: type="text", text=part.get("text", "") ).model_dump() ) - elif part_type == "refusal" or "refusal" in part: + elif part_type == "refusal": content.append( AnthropicResponseContentBlockText( type="text", text=part.get("refusal", "") or "" ).model_dump() ) - stop_reason = "refusal" elif item_type == "reasoning": content.extend( self._thinking_blocks_from_reasoning_item( @@ -698,14 +729,10 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ).model_dump() ) stop_reason = "tool_use" - elif item_type == "refusal": - refusal_text = item.get("refusal") or item.get("text", "") or "" - content.append(AnthropicResponseContentBlockText(type="text", text=refusal_text).model_dump()) - stop_reason = "refusal" - - # status -> stop_reason override if response.status == "incomplete": stop_reason = "max_tokens" + elif refusal_text is not None: + stop_reason = "refusal" anthropic_usage: Final = self.translate_responses_api_usage_to_anthropic_usage(response.usage) @@ -718,4 +745,13 @@ class LiteLLMAnthropicToResponsesAPIAdapter: usage=anthropic_usage, content=content, stop_reason=stop_reason, + stop_details=( + { + "type": "refusal", + "category": None, + "explanation": refusal_text, + } + if stop_reason == "refusal" + else None + ), ) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index f1107c87c73..cdfc5227e06 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -520,8 +520,15 @@ ContentBlockContentBlockDict = ToolUseBlock | TextBlock | ChatCompletionThinking ContentBlockStart = ContentBlockStartToolUse | ContentBlockStartText +class AnthropicStopDetails(TypedDict, total=False): + type: ReadOnly[Literal["refusal"]] + category: ReadOnly[str | None] + explanation: ReadOnly[str | None] + + class MessageDelta(TypedDict, total=False): stop_reason: str | None + stop_details: AnthropicStopDetails class ServerToolUsage(TypedDict, total=False): diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 4fe1dafc73b..038a23a3ca2 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -5,6 +5,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.types.llms.anthropic import ( AnthropicResponseContentBlockText, AnthropicResponseContentBlockToolUse, + AnthropicStopDetails, ContextManagementResponse, ServerToolUsage, ) @@ -78,16 +79,6 @@ class AnthropicUsage(TypedDict, total=False): server_tool_use: NotRequired[ReadOnly[ServerToolUsage]] -class AnthropicStopDetails(TypedDict, total=False): - """ - Safeguard verdict accompanying a `stop_reason: "refusal"` response: - https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback - """ - - category: ReadOnly[str | None] - explanation: ReadOnly[str | None] - - class AnthropicMessagesResponse(TypedDict, total=False): """ Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 2d74c00071b..fba5532d1e5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -40,6 +40,51 @@ from litellm.types.utils import ( ) +def test_translate_chat_refusal_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-refusal", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="stop", + message=Message(content=None, role="assistant", refusal="I cannot fulfill this request."), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [{"type": "text", "text": "I cannot fulfill this request."}] + assert result["stop_reason"] == "refusal" + assert result.get("stop_details") == { + "type": "refusal", + "category": None, + "explanation": "I cannot fulfill this request.", + } + + +def test_translate_chat_length_takes_precedence_over_refusal(): + response = ModelResponse( + id="chatcmpl-partial-refusal", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="length", + message=Message(content=None, role="assistant", refusal="Partial refusal"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["stop_reason"] == "max_tokens" + assert result.get("stop_details") is None + + def test_translate_streaming_openai_chunk_to_anthropic_content_block(): choices = [ StreamingChoices( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 17d42f55ae0..e7381986ec2 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -108,6 +108,93 @@ def _text_deltas(events: List[dict]) -> List[str]: ] +def test_streaming_chat_refusal_emits_only_refusal_stop_details(): + chunks = [ + _make_chunk(Delta(content=None, refusal="I cannot fulfill this request.")), + _make_chunk(Delta(content=None), finish_reason="stop"), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="openai-model") + + events = _drain_sync(wrapper) + + assert _text_deltas(events) == [] + message_delta = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"] == { + "stop_reason": "refusal", + "stop_details": { + "type": "refusal", + "category": None, + "explanation": "I cannot fulfill this request.", + }, + } + + +@pytest.mark.asyncio +async def test_streaming_chat_refusal_emits_only_refusal_stop_details_async(): + chunks = [ + _make_chunk(Delta(content=None, refusal="I cannot fulfill this request.")), + _make_chunk(Delta(content=None), finish_reason="stop"), + ] + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="openai-model") + + events = await _drain_async(wrapper) + + assert _text_deltas(events) == [] + message_delta = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." + + +def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved(): + chunks = [ + _make_chunk( + Delta(content=None, refusal="I cannot fulfill this request."), + finish_reason="stop", + ) + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="openai-model") + + events = _drain_sync(wrapper) + + message_delta = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." + + +@pytest.mark.asyncio +async def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved_async(): + chunks = [ + _make_chunk( + Delta(content=None, refusal="I cannot fulfill this request."), + finish_reason="stop", + ) + ] + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="openai-model") + + events = await _drain_async(wrapper) + + message_delta = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." + + +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.asyncio +async def test_streaming_chat_length_takes_precedence_over_refusal(async_mode: bool): + chunks = [ + _make_chunk(Delta(content=None, refusal="Partial refusal")), + _make_chunk(Delta(content=None), finish_reason="length"), + ] + stream = _AsyncStream(chunks) if async_mode else iter(chunks) + wrapper = AnthropicStreamWrapper(completion_stream=stream, model="openai-model") + + events = await _drain_async(wrapper) if async_mode else _drain_sync(wrapper) + + message_delta = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "max_tokens" + assert "stop_details" not in message_delta["delta"] + + def _input_json_deltas(events: List[dict]) -> List[str]: return [ e["delta"]["partial_json"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index 77ecfb73ff6..d8f16dde3e7 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -34,6 +34,14 @@ def _drain_async(events: list) -> list: return asyncio.run(_run()) +def _drain_sync_upstream(events: list) -> list: + async def _run() -> list: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=iter(events), model="m") + return [chunk async for chunk in wrapper] + + return asyncio.run(_run()) + + class TestMessageStartEmittedExactlyOnce: """The ``__anext__`` fallback emits ``message_start`` before consuming the stream, so ``_process_event`` must not emit a second one when @@ -55,6 +63,15 @@ class TestMessageStartEmittedExactlyOnce: chunks = _drain_async([{"type": "response.created"}]) assert chunks[0]["type"] == "message_start" + def test_sync_upstream_iterator_is_consumed(self): + chunks = _drain_sync_upstream( + [ + {"type": "response.created"}, + {"type": "response.output_text.delta", "item_id": "m1", "delta": "hi"}, + ] + ) + assert any(chunk.get("delta", {}).get("text") == "hi" for chunk in chunks) + class TestProcessEventResponseCreatedGuard: """``_process_event`` must emit ``message_start`` exactly once even if @@ -311,17 +328,37 @@ class TestResponseCompletedUsage: class TestRefusalStreamEvents: - def test_refusal_delta_emits_text_delta(self): + def test_refusal_event_sequence_emits_only_stop_details(self): + response = SimpleNamespace( + status="completed", + output=[{"type": "message", "content": [{"type": "refusal", "refusal": "I cannot fulfill this."}]}], + usage=None, + ) chunks = _process_all( [ {"type": "response.created"}, - {"type": "response.refusal.delta", "item_id": "ref_1", "delta": "I cannot fulfill this."}, + {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}}, + {"type": "response.refusal.delta", "item_id": "msg_1", "delta": "I cannot fulfill this."}, + {"type": "response.output_item.done", "item": {"type": "message", "id": "msg_1"}}, + {"type": "response.completed", "response": response}, ] ) - assert any( - c.get("type") == "content_block_delta" and c.get("delta", {}).get("text") == "I cannot fulfill this." - for c in chunks - ) + assert [chunk["type"] for chunk in chunks] == [ + "message_start", + "content_block_start", + "content_block_stop", + "message_delta", + "message_stop", + ] + assert chunks[3]["delta"] == { + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": { + "type": "refusal", + "category": None, + "explanation": "I cannot fulfill this.", + }, + } def test_response_completed_with_refusal_sets_stop_reason_refusal(self): response = SimpleNamespace( @@ -333,12 +370,13 @@ class TestRefusalStreamEvents: message_delta = next(c for c in chunks if c["type"] == "message_delta") assert message_delta["delta"]["stop_reason"] == "refusal" - def test_response_completed_with_standalone_refusal_item_sets_stop_reason_refusal(self): + def test_incomplete_status_takes_precedence_over_refusal(self): response = SimpleNamespace( - status="completed", - output=[{"type": "refusal", "refusal": "Standalone refusal"}], + status="incomplete", + output=[{"type": "message", "content": [{"type": "refusal", "refusal": "Partial refusal"}]}], usage=None, ) - chunks = _process_all([{"type": "response.completed", "response": response}]) + chunks = _process_all([{"type": "response.incomplete", "response": response}]) message_delta = next(c for c in chunks if c["type"] == "message_delta") - assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_reason"] == "max_tokens" + assert "stop_details" not in message_delta["delta"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py index f48c31ea5ed..8205e993d26 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py @@ -1205,16 +1205,16 @@ def _make_output_message(texts: List[str]) -> MagicMock: return msg -def _make_refusal_message(refusal_text: str) -> MagicMock: - from openai.types.responses import ResponseOutputMessage +def _make_refusal_message(refusal_text: str): + from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal - part = MagicMock() - part.type = "refusal" - part.refusal = refusal_text - - msg = MagicMock(spec=ResponseOutputMessage) - msg.content = [part] - return msg + return ResponseOutputMessage( + id="msg_refusal", + content=[ResponseOutputRefusal(type="refusal", refusal=refusal_text)], + role="assistant", + status="completed", + type="message", + ) def _make_function_call_item(call_id: str, name: str, arguments: str) -> MagicMock: @@ -1296,6 +1296,11 @@ class TestTranslateResponse: assert result["content"][0]["type"] == "text" assert result["content"][0]["text"] == "I cannot fulfill this request." assert result["stop_reason"] == "refusal" + assert result.get("stop_details") == { + "type": "refusal", + "category": None, + "explanation": "I cannot fulfill this request.", + } def test_dict_refusal_part_in_message_becomes_text_block(self): output_item = { @@ -1308,18 +1313,16 @@ class TestTranslateResponse: assert result["content"][0]["type"] == "text" assert result["content"][0]["text"] == "Refused by policy" assert result["stop_reason"] == "refusal" + assert result.get("stop_details", {}).get("explanation") == "Refused by policy" - def test_dict_refusal_item_becomes_text_block(self): - output_item = { - "type": "refusal", - "refusal": "Standalone refusal", - } - response = _make_mock_response(output=[output_item]) + def test_incomplete_status_takes_precedence_over_refusal(self): + response = _make_mock_response( + output=[_make_refusal_message("Partial refusal")], + status="incomplete", + ) result: Any = _ADAPTER.translate_response(response) - assert len(result["content"]) == 1 - assert result["content"][0]["type"] == "text" - assert result["content"][0]["text"] == "Standalone refusal" - assert result["stop_reason"] == "refusal" + assert result["stop_reason"] == "max_tokens" + assert result.get("stop_details") is None def test_incomplete_status_sets_max_tokens(self): """status='incomplete' overrides stop_reason to 'max_tokens'.""" From 7399b3844dfc45ab0c9ddbb116b26235ce4f7df0 Mon Sep 17 00:00:00 2001 From: Atharva-Kanherkar <142440039+Atharva-Kanherkar@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:33:06 +0530 Subject: [PATCH 24/63] fix(anthropic): satisfy lint budget gates for refusal translation --- .../adapters/streaming_iterator.py | 14 +++++----- .../adapters/transformation.py | 2 +- .../responses_adapters/streaming_iterator.py | 10 +++---- .../responses_adapters/transformation.py | 28 ++++++++++--------- litellm/types/llms/anthropic.py | 2 +- 5 files changed, 29 insertions(+), 27 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 41db4f12143..66fcad474d5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -11,7 +11,7 @@ from typing import ( Final, Literal, Protocol, - cast, + cast, # noqa: TID251 # rebuilt message_delta dict spans the ContentBlockDelta/MessageBlockDelta union get_args, ) @@ -306,7 +306,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Synthesized compaction block from compact_20260112 polyfill (streaming). self.compaction_block = compaction_block self.iterations_usage = iterations_usage - self._refusal_text_parts: list[str] = [] + self._refusal_text_parts: list[str] = [] # mutable-ok: accumulates streamed refusal delta text across chunks self.sent_compaction_block: bool = False # Per-phase flags so the compaction block's start/delta/stop events # are emitted (and the public state machine is advanced) in @@ -1003,17 +1003,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): ) -> ContentBlockDelta | MessageBlockDelta: if processed_chunk.get("type") != "message_delta" or not self._refusal_text_parts: return processed_chunk - delta: Final = cast(Mapping[str, object], processed_chunk["delta"]) + delta: Final = cast(Mapping[str, object], processed_chunk["delta"]) # cast-ok: keys checked before use if delta.get("stop_reason") == "max_tokens": return processed_chunk - return cast( + return cast( # cast-ok: rebuilt dict matches the message_delta TypedDict shape for this branch ContentBlockDelta | MessageBlockDelta, - { + { # mutable-ok: fresh translation payload; never mutated after construction **processed_chunk, - "delta": { + "delta": { # mutable-ok: fresh message_delta payload; never mutated after construction **delta, "stop_reason": "refusal", - "stop_details": { + "stop_details": { # mutable-ok: fresh stop_details payload; never mutated after construction "type": "refusal", "category": None, "explanation": "".join(self._refusal_text_parts), diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 5655f59a4cc..2d950f007a5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -1534,7 +1534,7 @@ class LiteLLMAnthropicMessagesAdapter: content=anthropic_content, stop_reason=anthropic_finish_reason, stop_details=( - { + { # mutable-ok: fresh refusal stop_details payload built per response "type": "refusal", "category": None, "explanation": refusal_text, diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 1b0ab5923d8..8e5ce77e48f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -50,7 +50,7 @@ class AnthropicResponsesStreamWrapper: self._sent_message_start = False self._sent_message_stop = False self._chunk_queue: deque = deque() - self._refusal_text_parts: list[str] = [] + self._refusal_text_parts: list[str] = [] # mutable-ok: accumulates streamed refusal delta text across chunks self._sync_responses_iterator: Any = None def _make_message_start(self) -> dict[str, Any]: @@ -251,19 +251,19 @@ class AnthropicResponsesStreamWrapper: else AnthropicUsage(input_tokens=0, output_tokens=0) ) - message_delta_payload: Final = { + message_delta_payload: Final = { # mutable-ok: fresh message_delta payload built per chunk "stop_reason": stop_reason, "stop_sequence": None, **( - { - "stop_details": { + { # mutable-ok: fresh refusal stop_details payload built per chunk + "stop_details": { # mutable-ok: fresh refusal stop_details payload built per chunk "type": "refusal", "category": None, "explanation": refusal_text, } } if stop_reason == "refusal" - else {} + else {} # mutable-ok: empty spread placeholder for non-refusal stop ), } diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 92f3e08f24d..318065148b8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -81,20 +81,20 @@ class LiteLLMAnthropicToResponsesAPIAdapter: ) if not isinstance(item, Mapping): return None - item_mapping: Final = cast(Mapping[str, object], item) + item_mapping: Final = cast(Mapping[str, object], item) # cast-ok: keys re-checked before use raw_parts: Final = item_mapping.get("content") if item_mapping.get("type") != "message" or not isinstance(raw_parts, Sequence): return None - return next( - ( - refusal - for part in cast(Sequence[object], raw_parts) - if isinstance(part, Mapping) - and cast(Mapping[str, object], part).get("type") == "refusal" - and isinstance((refusal := cast(Mapping[str, object], part).get("refusal")), str) - ), - None, - ) + for part in cast(Sequence[object], raw_parts): # cast-ok: members re-validated below + if not isinstance(part, Mapping): + continue + part_mapping = cast(Mapping[str, object], part) # cast-ok: keys re-checked before use + if part_mapping.get("type") != "refusal": + continue + refusal = part_mapping.get("refusal") + if isinstance(refusal, str): + return refusal + return None return next( (text for item in output if (text := refusal_text_from_item(item)) is not None), @@ -656,7 +656,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content: Final[list[dict[str, object]]] = [] stop_reason: AnthropicFinishReason = "end_turn" - refusal_text: Final = self._refusal_text_from_output(cast(Iterable[object], response.output)) + refusal_text: Final = self._refusal_text_from_output( + cast(Iterable[object], response.output) # cast-ok: output items re-validated per item + ) for item in response.output: if isinstance(item, ResponseReasoningItem): @@ -746,7 +748,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content=content, stop_reason=stop_reason, stop_details=( - { + { # mutable-ok: fresh refusal stop_details payload built per response "type": "refusal", "category": None, "explanation": refusal_text, diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index cdfc5227e06..365d59a179b 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -528,7 +528,7 @@ class AnthropicStopDetails(TypedDict, total=False): class MessageDelta(TypedDict, total=False): stop_reason: str | None - stop_details: AnthropicStopDetails + stop_details: ReadOnly[AnthropicStopDetails] class ServerToolUsage(TypedDict, total=False): From 0f59b6fb7a996debcb350f6bf10a18e5ba276a62 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 12:03:42 -0700 Subject: [PATCH 25/63] ci(e2e): refine changed-test selection and runner lifecycle --- .github/e2e-stack/assert_tests_ran.py | 29 +++++-- .github/e2e-stack/secrets_to_env.py | 44 ++++++---- .github/e2e-stack/up.sh | 8 +- .github/workflows/test-e2e-changed.yml | 80 ++++++++++++++++--- .../test_e2e_changed_gate.py | 55 +++++++++++++ tests/e2e/CONTRIBUTING.md | 12 ++- tests/e2e/gateway/stage_mirror_ci_config.yml | 4 - 7 files changed, 186 insertions(+), 46 deletions(-) create mode 100644 tests/code_coverage_tests/test_e2e_changed_gate.py diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py index 092a9db7256..7fd3f7c7c33 100644 --- a/.github/e2e-stack/assert_tests_ran.py +++ b/.github/e2e-stack/assert_tests_ran.py @@ -5,15 +5,28 @@ from typing import Final def main() -> int: - report: Final = ET.parse(Path(sys.argv[1])).getroot() - suites: Final = tuple(report.iter("testsuite")) - collected: Final = sum(int(suite.get("tests", "0")) for suite in suites) - skipped: Final = sum(int(suite.get("skipped", "0")) for suite in suites) - executed: Final = collected - skipped - _ = sys.stdout.write(f"executed {executed} of {collected} collected tests ({skipped} skipped)\n") - if executed > 0: + selected: Final = tuple(sys.argv[2:]) + try: + report: Final = ET.parse(Path(sys.argv[1])).getroot() + except (ET.ParseError, OSError): + _ = sys.stdout.write("::error::could not read the test execution report\n") + return 1 + cases: Final = tuple(report.iter("testcase")) + passed: Final = frozenset( + case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error")) + ) + missing: Final = tuple(path for path in selected if path not in passed) + for path in selected: + collected: Final = sum(case.get("file") == path for case in cases) + skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases) + _ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n") + if ( + selected + and not missing + and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error")) + ): return 0 - _ = sys.stdout.write("::error::every selected test was skipped, so nothing was verified\n") + _ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n") return 1 diff --git a/.github/e2e-stack/secrets_to_env.py b/.github/e2e-stack/secrets_to_env.py index d2d6675d690..99e717b271d 100644 --- a/.github/e2e-stack/secrets_to_env.py +++ b/.github/e2e-stack/secrets_to_env.py @@ -1,29 +1,41 @@ +import os import re import sys from pathlib import Path from typing import Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError -secrets_adapter: TypeAdapter[dict[str, str]] = TypeAdapter(dict[str, str]) -SECRET_NAME: Final = re.compile(r"KEY|SECRET|TOKEN|PASS|CREDENTIAL|LICENSE|AUTH") +secrets_adapter: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) +ENV_NAME: Final = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") def main() -> int: - env_path = Path(sys.argv[1]) - secrets = {key: value.rstrip("\r\n") for key, value in secrets_adapter.validate_json(sys.stdin.read()).items()} - unwritable = tuple( - key for key, value in secrets.items() if "'" in value or "\n" in value or "\r" in value - ) - if unwritable: - _ = sys.stderr.write(f"values contain characters unsafe for both bash and dotenv: {', '.join(unwritable)}\n") + env_path: Final = Path(sys.argv[1]) + try: + secrets: Final = { + key: value.rstrip("\r\n") for key, value in secrets_adapter.validate_json(sys.stdin.read()).items() + } + except (ValidationError, UnicodeError): + _ = sys.stderr.write("expected a JSON object containing string environment values\n") + return 1 + if any( + ENV_NAME.fullmatch(key) is None or any(char in value for char in "'\n\r\0") for key, value in secrets.items() + ): + _ = sys.stderr.write("environment names or values cannot be represented in both bash and dotenv\n") + return 1 + for value in secrets.values(): + if value: + _ = sys.stdout.write(f"::add-mask::{value.replace('%', '%25')}\n") + sys.stdout.flush() + lines: Final = tuple(f"{key}='{value}'" for key, value in secrets.items() if value) + try: + with os.fdopen(os.open(env_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT | os.O_NOFOLLOW, 0o600), "w") as handle: + os.fchmod(handle.fileno(), 0o600) + _ = handle.write("\n".join(lines) + "\n") + except OSError: + _ = sys.stderr.write("could not write the environment file\n") return 1 - lines = tuple(f"{key}='{value}'" for key, value in secrets.items() if value) - with env_path.open("a") as handle: - _ = handle.write("\n".join(lines) + "\n") - for key, value in secrets.items(): - if value and SECRET_NAME.search(key): - _ = sys.stdout.write(f"::add-mask::{value}\n") return 0 diff --git a/.github/e2e-stack/up.sh b/.github/e2e-stack/up.sh index 06fce35a896..e891add341d 100755 --- a/.github/e2e-stack/up.sh +++ b/.github/e2e-stack/up.sh @@ -1,5 +1,6 @@ #!/usr/bin/env bash set -euo pipefail +umask 077 REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}" @@ -8,9 +9,9 @@ LOGS_DIR="${STACK_DIR}/logs" PIDS_DIR="${STACK_DIR}/pids" POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}" -VALKEY_IMAGE="${E2E_VALKEY_IMAGE:-valkey/valkey:8.1}" +VALKEY_IMAGE="${E2E_VALKEY_IMAGE:-valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa}" JAEGER_IMAGE="${E2E_JAEGER_IMAGE:-jaegertracing/jaeger:2.10.0}" -NGINX_IMAGE="${E2E_NGINX_IMAGE:-nginx:1.29-alpine}" +NGINX_IMAGE="${E2E_NGINX_IMAGE:-nginx:1.29.1-alpine@sha256:42a516af16b852e33b7682d5ef8acbd5d13fe08fecadc7ed98605ba5e3b26ab8}" LB_PORT="${E2E_LB_PORT:-4000}" GATEWAY_PORT_1="${E2E_GATEWAY_PORT_1:-4010}" @@ -28,6 +29,8 @@ JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}" MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}" mkdir -p "${CERTS_DIR}" "${LOGS_DIR}" "${PIDS_DIR}" +chmod 700 "${STACK_DIR}" "${LOGS_DIR}" "${PIDS_DIR}" +chmod 755 "${CERTS_DIR}" log() { printf 'e2e-stack: %s\n' "$*"; } @@ -38,7 +41,6 @@ wait_for() { until eval "${check}"; do if ((SECONDS >= deadline)); then log "timed out waiting for ${label}" - tail -n 60 "${LOGS_DIR}"/*.log 2>/dev/null || true exit 1 fi sleep 2 diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index e08660d0412..b802e9173f5 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -26,27 +26,26 @@ jobs: GH_TOKEN: ${{ github.token }} REPO: ${{ github.repository }} PR_NUMBER: ${{ github.event.pull_request.number }} - SMOKE_TESTS: tests/e2e/access_control + HEAD_SHA: ${{ github.event.pull_request.head.sha }} OWN_LANE: '^tests/e2e/(ui|claude_code|load)/|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$' run: | + gh api "repos/${REPO}/pulls/${PR_NUMBER}" \ + --jq 'select(.head.sha == env.HEAD_SHA and .changed_files < 3000) | .head.sha' \ + | grep -Fxq "${HEAD_SHA}" files="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}/files" --paginate \ --jq '.[] | select(.status != "removed") | .filename')" + gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}" tests="$(printf '%s\n' "${files}" \ | grep -E '^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$' \ | grep -vE "${OWN_LANE}" \ | sort -u | tr '\n' ' ' | sed 's/ $//')" || true - if [ -z "${tests}" ] && printf '%s\n' "${files}" | grep -vE "${OWN_LANE}" \ - | grep -qE '^(tests/e2e/|\.github/e2e-stack/|\.github/workflows/test-e2e-changed\.yml$)'; then - tests="${SMOKE_TESTS}" - echo "harness or stack changed without a test file; running the smoke suite" - fi echo "tests=${tests}" >> "${GITHUB_OUTPUT}" if [ -n "${tests}" ]; then echo "any=true" >> "${GITHUB_OUTPUT}" echo "selected e2e tests: ${tests}" else echo "any=false" >> "${GITHUB_OUTPUT}" - echo "no e2e changes; nothing to run" + echo "no changed e2e test files supported by this stack; nothing to run" fi run: @@ -88,6 +87,7 @@ jobs: uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false + ref: ${{ github.sha }} - name: Set up Python uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 @@ -120,14 +120,24 @@ jobs: run: uv run --no-sync playwright install --with-deps chromium - name: Configure AWS credentials + id: aws uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0 with: role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }} aws-region: us-east-1 role-session-name: litellm-e2e-changed-${{ github.run_id }} + role-duration-seconds: 900 + output-env-credentials: false + output-credentials: true - name: Fetch provider credentials from AWS Secrets Manager + env: + AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }} + AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }} + AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }} + AWS_DEFAULT_REGION: us-east-1 run: | + umask 077 aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-provider-keys \ --query SecretString --output text \ | uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env @@ -138,7 +148,12 @@ jobs: - name: Boot the stage-mirror stack id: boot - run: bash .github/e2e-stack/up.sh + run: | + umask 077 + if ! bash .github/e2e-stack/up.sh > "${RUNNER_TEMP}/e2e-boot.log" 2>&1; then + echo "::error::stage-mirror stack failed to boot; raw logs are not published" + exit 1 + fi - name: Export stack environment run: | @@ -149,13 +164,16 @@ jobs: - name: Run the selected tests three times with retries off env: TESTS: ${{ needs.detect.outputs.tests }} + E2E_FIXTURE_MODE: live run: | + umask 077 read -r -a test_files <<< "${TESTS}" for pass in 1 2 3; do report="${RUNNER_TEMP}/e2e-pass-${pass}.xml" echo "::group::pass ${pass} of 3" set +e - uv run --no-sync pytest "${test_files[@]}" --reruns 0 -v -rA --tb=short -p no:cacheprovider --junitxml="${report}" + uv run --no-sync pytest "${test_files[@]}" --rootdir=. --reruns 0 -v -p no:cacheprovider \ + -o junit_family=xunit1 --junitxml="${report}" > "${RUNNER_TEMP}/e2e-pass-${pass}.log" 2>&1 status=$? set -e echo "::endgroup::" @@ -163,13 +181,49 @@ jobs: echo "::error::the selected files collected no runnable tests, so nothing was verified" exit 1 fi + if ! uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"; then + echo "::error::pass ${pass} of 3 did not verify every selected file" + exit 1 + fi if [ "${status}" != "0" ]; then echo "::error::pass ${pass} of 3 failed with exit code ${status}" exit "${status}" fi - uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" + echo "pass ${pass} of 3 passed" done - - name: Show stack logs on failure - if: failure() && steps.boot.conclusion != 'skipped' - run: tail -n 200 "${RUNNER_TEMP}/litellm-e2e-stack/logs"/*.log + - name: Stop the stack + if: always() && steps.boot.outcome != 'skipped' + run: bash .github/e2e-stack/down.sh + + - name: Remove credentials and raw output + if: always() + run: | + rm -f tests/e2e/.env "${RUNNER_TEMP}/e2e-boot.log" "${RUNNER_TEMP}"/e2e-pass-*.log "${RUNNER_TEMP}"/e2e-pass-*.xml + rm -rf "${RUNNER_TEMP}/litellm-e2e-stack" + + gate: + name: e2e-changed-tests + needs: [detect, run] + if: always() + runs-on: ubuntu-latest + timeout-minutes: 5 + steps: + - name: Require three successful passes when tests changed + env: + DETECT_RESULT: ${{ needs.detect.result }} + ANY_TESTS: ${{ needs.detect.outputs.any }} + RUN_RESULT: ${{ needs.run.result }} + run: | + if [ "${DETECT_RESULT}" != "success" ]; then + echo "::error::changed-test detection did not succeed" + exit 1 + fi + if [ "${ANY_TESTS}" = "false" ]; then + echo "no changed e2e test files supported by this stack; nothing to run" + exit 0 + fi + if [ "${ANY_TESTS}" != "true" ] || [ "${RUN_RESULT}" != "success" ]; then + echo "::error::selected e2e tests require an approved, successful run; fork PRs must run from a reviewed same-repository branch" + exit 1 + fi diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py new file mode 100644 index 00000000000..7a82a8ebb60 --- /dev/null +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -0,0 +1,55 @@ +import subprocess +import sys +import xml.etree.ElementTree as ET +from pathlib import Path +from typing import Final + +import pytest + +GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py" +SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py") + + +@pytest.mark.parametrize( + ("second_outcome", "expected_status"), + (("passed", 0), ("skipped", 1), ("failure", 1), ("error", 1), ("deselected", 1)), +) +def test_each_changed_file_must_run(tmp_path: Path, second_outcome: str, expected_status: int) -> None: + suite: Final = ET.Element("testsuite") + _ = ET.SubElement(suite, "testcase", file=SELECTED[0]) + if second_outcome != "deselected": + second: Final = ET.SubElement(suite, "testcase", file=SELECTED[1]) + if second_outcome != "passed": + _ = ET.SubElement(second, second_outcome) + report: Final = tmp_path / "report.xml" + ET.ElementTree(suite).write(report) + + result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) + + assert result.returncode == expected_status, result.stdout + + +@pytest.mark.parametrize("outcome", ("failure", "error")) +def test_passing_case_does_not_hide_a_failure_in_the_same_file(tmp_path: Path, outcome: str) -> None: + suite: Final = ET.Element("testsuite") + _ = ET.SubElement(suite, "testcase", file=SELECTED[0]) + failed: Final = ET.SubElement(suite, "testcase", file=SELECTED[0]) + _ = ET.SubElement(failed, outcome) + report: Final = tmp_path / "report.xml" + ET.ElementTree(suite).write(report) + + result: Final = subprocess.run( + [sys.executable, str(GATE), str(report), SELECTED[0]], capture_output=True, text=True + ) + + assert result.returncode == 1 + + +@pytest.mark.parametrize("contents", ("", "')) +def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None: + report: Final = tmp_path / "report.xml" + _ = report.write_text(contents) + + result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) + + assert result.returncode == 1 diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 1e88f6505cf..d51e729fd1c 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -54,9 +54,17 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT ### The pull request check -Every PR that adds or modifies a `tests/e2e/**/test_*.py` file (outside `ui/`, `claude_code/`, `load/`, and `batches/test_managed_files_enforcement_e2e.py`, which have their own lanes or need a differently configured proxy) runs exactly those files three times, with retries off, against a stage-mirror stack booted on a GitHub Actions runner: migrations, a control-plane backend, two gateway processes behind an nginx load balancer, Postgres, Jaeger, and a TLS cluster-mode Valkey, wired the way stage is deployed. A PR that only touches the harness or the stack itself runs the `access_control/` suite as a smoke instead. Three green passes are the bar because the check exists to catch a flaky test before it reaches the release gate, so a red pass is a failure to fix, not a retry candidate. A pass that executes nothing is also red: each pass writes a JUnit report and fails when every collected test was skipped, so a skipped-out file cannot pass on paper. The same stack boots on a laptop with `bash .github/e2e-stack/up.sh`: it reads provider keys from `tests/e2e/.env`, writes the pytest environment to `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}/stack.env`, every port is overridable through `E2E_*_PORT` variables, and `down.sh` tears it all down +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times with pytest retries off. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. Documentation, harness, configuration, deleted-file, and workflow-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack -Credentials: everything the runner writes into `tests/e2e/.env` comes from two AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`, and the first holds only the names that the tests and `tests/e2e/gateway/stage_mirror_ci_config.yml` read. The cloud credentials in it are dedicated to this lane and can do only what the tests do: the AWS pair belongs to the `litellm-e2e-changed` IAM user, whose inline policy allows Bedrock inference, the one guardrail the tests use, the batch job APIs, the batch S3 bucket, and assuming the batch test's role, and the Vertex key belongs to a service account holding only `roles/aiplatform.user` on the Vertex project. The provider API keys carry no lane-specific spend cap, since a cap that trips mid-month would fail every run until it resets, so the reviewer's approval of the `e2e-changed` environment is the control on how a PR's tests use them. Rotating any value is a single `aws secretsmanager put-secret-value` on the secret, with no workflow change +Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A failed pass stops the run without retrying. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch + +Repository admins must require the `e2e-changed-tests` status check for merging and configure the `e2e-changed` environment with required reviewers, self-review disabled, and admin bypass disabled. Each push cancels the previous run; a new run that selects tests needs a fresh approval. Reviewers must inspect the entire executable PR diff, including application code, dependencies, tests, and workflow helpers, before approving the exact revision. Approved code executes with provider credentials, so environment approval is a trust decision about that code + +Credentials come from the existing AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`. The OIDC role must trust only `repo:BerriAI/litellm:environment:e2e-changed` with audience `sts.amazonaws.com` and have read access only to these secrets. The short-lived reader credentials are scoped to the fetch step. Provider credentials must cover the selected suites, including Datadog credentials when logging or MCP tests need them; missing credentials fail the run. Keep provider credentials dedicated to this lane with only the permissions those tests need + +Fetched values are masked before use, and credential files and raw output are private to the runner. Public logs contain selected file names, counts, and pass status; raw pytest output, reports, and stack logs are not uploaded or printed. The workflow removes them and the credential files during cleanup. To diagnose a failed pass, reproduce the selected files locally with the appropriate credentials and inspect the local logs + +To reproduce the CI topology on a dedicated machine, `bash .github/e2e-stack/up.sh` reads `tests/e2e/.env`, writes `stack.env` under `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}`, and `bash .github/e2e-stack/down.sh` stops it. Keep this directory private and remove its credential files and logs after use ### Record and replay diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 21664c2d0a1..6b3756b5f9f 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -52,10 +52,6 @@ model_list: litellm_params: model: gemini/gemini-2.5-flash api_key: os.environ/GEMINI_API_KEY - - model_name: gemini-2.5-flash - litellm_params: - model: gemini/gemini-2.5-flash - api_key: os.environ/GEMINI_API_KEY mcp_servers: devin: From 5051e6d44acb627d912cba11b4afdc91ae5ff8c6 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 12:08:11 -0700 Subject: [PATCH 26/63] ci(e2e): include execution gate checks in code quality --- .github/workflows/test-code-quality.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index c112bf2bb22..a0e0ce6665a 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -74,6 +74,9 @@ jobs: - name: check_workflow_startup_safety run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py + - name: test_e2e_changed_gate + run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py + - name: router_code_coverage run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py From c55248d113800f32283532cb811a2509b27ba94a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:07:29 -0700 Subject: [PATCH 27/63] fix(router): treat a routing entry with no latency samples as zero latency Latency-based routing averaged a deployment's cached samples with total / len(samples) and raised ZeroDivisionError once an entry held none, which the proxy answered as a 500 for every later request on that model group. Cost-based routing writes the same {model_group}_map entry with minute counters only, so a group used by both strategies hit this on every latency-routed request. A deployment with no samples now counts as 0 latency, the same as one the router has never seen Resolves LIT-7053 --- litellm/router_strategy/lowest_latency.py | 23 +++++++------- .../router_strategy/test_lowest_latency.py | 30 +++++++++++++++++++ 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index a1b67eaeaf9..4a1fb4f325a 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -1,6 +1,7 @@ #### What this does #### # picks based on response time (for streaming, this is time to first token) import random +from collections.abc import Sequence from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final @@ -25,6 +26,12 @@ class RoutingArgs(LiteLLMPydanticObjectBase): max_latency_list_size: int = 10 +def _average_latency(samples: Sequence[float | int]) -> float: + if not samples: + return 0.0 + return sum(sample for sample in samples if isinstance(sample, float)) / len(samples) + + class LowestLatencyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -431,23 +438,13 @@ class LowestLatencyLoggingHandler(CustomLogger): item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) # get average latency or average ttft (depending on streaming/non-streaming) - total: float = 0.0 use_ttft = ( request_kwargs is not None and request_kwargs.get("stream", None) is not None and request_kwargs["stream"] is True and len(item_ttft_latency) > 0 ) - if use_ttft: - for _call_latency in item_ttft_latency: - if isinstance(_call_latency, float): - total += _call_latency - item_latency = total / len(item_ttft_latency) - else: - for _call_latency in item_latency: - if isinstance(_call_latency, float): - total += _call_latency - item_latency = total / len(item_latency) + average_latency = _average_latency(item_ttft_latency if use_ttft else item_latency) # -------------- # # Debugging Logic @@ -456,7 +453,7 @@ class LowestLatencyLoggingHandler(CustomLogger): # this helps a user to debug why the router picked a specfic deployment # _deployment_api_base = _deployment.get("litellm_params", {}).get("api_base", "") if _deployment_api_base is not None: - _latency_per_deployment[_deployment_api_base] = item_latency + _latency_per_deployment[_deployment_api_base] = average_latency # -------------- # # End of Debugging Logic # -------------- # @@ -466,7 +463,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ): # if user passed in tpm / rpm in the model_list continue else: - potential_deployments.append((_deployment, item_latency)) + potential_deployments.append((_deployment, average_latency)) if len(potential_deployments) == 0: return None diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index 6701f4a7aa2..251b760a8d8 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -163,3 +163,33 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): assert latencies and latencies[-1] == pytest.approx(2.0) assert not isinstance(latencies[-1], timedelta) json.dumps({"latency": latencies}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cached_entry", + [{"latency": []}, {"2026-09-05-15-39": {"tpm": 28, "rpm": 1}}], + ids=["empty_latency_list", "minute_bucket_only_as_cost_based_routing_writes"], +) +async def test_async_get_available_deployments_treats_missing_samples_as_zero_latency(cached_entry): + """A cached entry with no latency samples (cost-based routing shares the group's map key and writes + minute buckets only) must count as 0 latency, like an unseen deployment, instead of dividing by zero.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + cache.set_cache( + key="gemini-embedding-001_map", + value={DEPLOYMENT_ID: cached_entry, "slower": {"latency": [0.5]}}, + ) + healthy_deployments = [ + {"model_info": {"id": DEPLOYMENT_ID}, "litellm_params": {}}, + {"model_info": {"id": "slower"}, "litellm_params": {}}, + ] + + picked = await handler.async_get_available_deployments( + model_group="gemini-embedding-001", + healthy_deployments=healthy_deployments, + request_kwargs={"stream": False, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == DEPLOYMENT_ID From 65d8bbb8ac4610c9722ecae03ffa16b352c47256 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:10:40 -0700 Subject: [PATCH 28/63] fix(e2e): wait for every gateway before using a new model and keep the network rerun The changed-tests workflow overrode the suite's `--reruns 1` with `--reruns 0`, so a transport blip failed a pass that pytest.ini already scopes to network errors and 5xx responses. Pass 2 of run 33692484803 also went red 15s after a model write with "no healthy deployments": the barrier only polled /v1/models through nginx, which proves one gateway converged, and the next request rolled the other. The stack now exports LITELLM_PROXY_REPLICA_URLS, the barrier polls every replica with the full budget before settling, and up.sh refuses to boot without DD_API_KEY, since the gateway config enables the datadog callback on every run --- .github/e2e-stack/up.sh | 7 ++ .github/workflows/test-e2e-changed.yml | 4 +- tests/e2e/CONTRIBUTING.md | 6 +- tests/e2e/claude_code/_env.py | 1 + tests/e2e/claude_code/conftest.py | 1 + tests/e2e/e2e_config.py | 18 +++- tests/e2e/proxy_client.py | 117 +++++++++++++++++++------ tests/e2e/test_proxy_client.py | 87 ++++++++++++++++++ 8 files changed, 206 insertions(+), 35 deletions(-) create mode 100644 tests/e2e/test_proxy_client.py diff --git a/.github/e2e-stack/up.sh b/.github/e2e-stack/up.sh index e891add341d..2f2e6c6f9a8 100755 --- a/.github/e2e-stack/up.sh +++ b/.github/e2e-stack/up.sh @@ -54,6 +54,12 @@ if [[ -f "${REPO_ROOT}/tests/e2e/.env" ]]; then set +a fi +if [[ -z "${DD_API_KEY:-}" ]]; then + log "DD_API_KEY is empty; the gateway config enables the datadog callback, so put a Datadog API key in tests/e2e/.env" + exit 1 +fi +export DD_SITE="${DD_SITE:-datadoghq.com}" + if ! port_open "${DATABASE_PORT}"; then docker run -d --name e2e-postgres -p "${DATABASE_PORT}:5432" \ -e "POSTGRES_USER=${DATABASE_USER}" -e "POSTGRES_PASSWORD=${DATABASE_PASSWORD}" -e "POSTGRES_DB=${DATABASE_NAME}" \ @@ -189,6 +195,7 @@ wait_for "load balancer" "curl -fs http://127.0.0.1:${LB_PORT}/health/liveliness cat > "${STACK_DIR}/stack.env" <> "${GITHUB_ENV}" - - name: Run the selected tests three times with retries off + - name: Run the selected tests three times env: TESTS: ${{ needs.detect.outputs.tests }} E2E_FIXTURE_MODE: live @@ -172,7 +172,7 @@ jobs: report="${RUNNER_TEMP}/e2e-pass-${pass}.xml" echo "::group::pass ${pass} of 3" set +e - uv run --no-sync pytest "${test_files[@]}" --rootdir=. --reruns 0 -v -p no:cacheprovider \ + uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \ -o junit_family=xunit1 --junitxml="${report}" > "${RUNNER_TEMP}/e2e-pass-${pass}.log" 2>&1 status=$? set -e diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index d51e729fd1c..7bfd02023dd 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -54,13 +54,13 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times with pytest retries off. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. Documentation, harness, configuration, deleted-file, and workflow-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, harness, configuration, deleted-file, and workflow-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack -Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A failed pass stops the run without retrying. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch +Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A failed pass stops the run. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch Repository admins must require the `e2e-changed-tests` status check for merging and configure the `e2e-changed` environment with required reviewers, self-review disabled, and admin bypass disabled. Each push cancels the previous run; a new run that selects tests needs a fresh approval. Reviewers must inspect the entire executable PR diff, including application code, dependencies, tests, and workflow helpers, before approving the exact revision. Approved code executes with provider credentials, so environment approval is a trust decision about that code -Credentials come from the existing AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`. The OIDC role must trust only `repo:BerriAI/litellm:environment:e2e-changed` with audience `sts.amazonaws.com` and have read access only to these secrets. The short-lived reader credentials are scoped to the fetch step. Provider credentials must cover the selected suites, including Datadog credentials when logging or MCP tests need them; missing credentials fail the run. Keep provider credentials dedicated to this lane with only the permissions those tests need +Credentials come from the existing AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`. The OIDC role must trust only `repo:BerriAI/litellm:environment:e2e-changed` with audience `sts.amazonaws.com` and have read access only to these secrets. The short-lived reader credentials are scoped to the fetch step. Provider credentials must cover the selected suites, including Datadog credentials when logging or MCP tests need them; missing credentials fail the run. `up.sh` refuses to start without `DD_API_KEY`, because the stack's gateway config enables the Datadog callback for every run and a gateway booted without the key fails readiness. Keep provider credentials dedicated to this lane with only the permissions those tests need Fetched values are masked before use, and credential files and raw output are private to the runner. Public logs contain selected file names, counts, and pass status; raw pytest output, reports, and stack logs are not uploaded or printed. The workflow removes them and the credential files during cleanup. To diagnose a failed pass, reproduce the selected files locally with the appropriate credentials and inspect the local logs diff --git a/tests/e2e/claude_code/_env.py b/tests/e2e/claude_code/_env.py index ed2da1fe1e2..889d8f848dd 100644 --- a/tests/e2e/claude_code/_env.py +++ b/tests/e2e/claude_code/_env.py @@ -99,5 +99,6 @@ def require_proxy_client( base_url=cfg.base_url, master_key=cfg.api_key, control_plane_base_url=cfg.base_url, + replica_urls=(cfg.base_url,), ) return ProxyClientConfig(client=client, api_key=cfg.api_key) diff --git a/tests/e2e/claude_code/conftest.py b/tests/e2e/claude_code/conftest.py index aaad7667936..6e3dce0377e 100644 --- a/tests/e2e/claude_code/conftest.py +++ b/tests/e2e/claude_code/conftest.py @@ -594,6 +594,7 @@ def _build_control_plane_client(proxy_config: ProxyConfig): base_url=proxy_config.base_url, master_key=proxy_config.api_key, control_plane_base_url=proxy_config.base_url, + replica_urls=(proxy_config.base_url,), ) diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 21c5a338dc3..70d1180cb38 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -10,6 +10,7 @@ import os import time import uuid from pathlib import Path +from typing import Final from dotenv import load_dotenv @@ -32,6 +33,14 @@ CONTROL_PLANE_BASE_URL = os.environ.get( "LITELLM_CONTROL_PLANE_URL", PROXY_BASE_URL ).rstrip("/") + +def parse_replica_urls(raw: str, fallback: str) -> tuple[str, ...]: + urls: Final = tuple(url.strip().rstrip("/") for url in raw.split(",") if url.strip()) + return urls or (fallback,) + + +PROXY_REPLICA_URLS: Final = parse_replica_urls(os.environ.get("LITELLM_PROXY_REPLICA_URLS", ""), PROXY_BASE_URL) + UI_USERNAME = os.environ.get("E2E_UI_USERNAME", "admin") UI_PASSWORD = os.environ.get("E2E_UI_PASSWORD", MASTER_KEY) @@ -86,10 +95,11 @@ SLOW_PROVIDER_TIMEOUT_SECONDS = float(os.environ.get("E2E_SLOW_PROVIDER_TIMEOUT" # (`proxy_config_reload_interval_seconds`, 30s by default and 7s on the e2e stack) # plus margin. # -# The barriers below wait this out instead of returning on first sight, because a -# single successful read only proves ONE replica converged: every request opens a -# fresh connection, so a load-balanced Service routes each one independently and -# the next call re-rolls. See ProxyClient._await_model_servable. +# The barriers below wait this out on top of polling /v1/models on every replica in +# PROXY_REPLICA_URLS: that poll proves each addressed gateway converged, but not the +# workers behind it, and behind a load balancer (PROXY_REPLICA_URLS unset) a +# successful read only proves ONE replica converged, because every request opens a +# fresh connection and the next call re-rolls. See ProxyClient._await_model_servable. PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15")) EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes") diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 2d382a610e1..c5b34a534b7 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -10,12 +10,15 @@ from __future__ import annotations import time import warnings -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime +from types import MappingProxyType +from typing import Final from e2e_http import ( AnthropicHeaders, + AuthHeaders, NoBody, ProbeResult, Result, @@ -71,6 +74,7 @@ from e2e_config import ( POLL_INTERVAL, POLL_TIMEOUT, PROXY_BASE_URL, + PROXY_REPLICA_URLS, REQUEST_TIMEOUT, SLOW_PROVIDER_TIMEOUT_SECONDS, settle_propagation, @@ -79,10 +83,11 @@ from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] -# After /model/new, poll data-plane /v1/models until the model is listed (or fail). -# Bound by MODEL_SERVABLE_TIMEOUT so a stuck reload does not burn the spend -# poll_timeout (120s). Return on first listing; settle_propagation owns the separate -# wait that lets every worker and replica reload before the caller uses the model. +# After /model/new, poll /v1/models on every replica in PROXY_REPLICA_URLS until each +# lists the model (or fail). Bound by MODEL_SERVABLE_TIMEOUT per replica so a stuck +# reload does not burn the spend poll_timeout (120s). Return on first listing; +# settle_propagation owns the separate wait that lets the workers behind each replica +# reload before the caller uses the model. MODEL_SERVABLE_TIMEOUT = 40.0 MODEL_SERVABLE_DB_SYNC_SECONDS = 0.0 MODEL_SERVABLE_INTERVAL = 2.0 @@ -108,6 +113,16 @@ class NotServable: ServableOutcome = Servable | NotServable +type ModelsPoller = Callable[[float], Result[ModelsListResponse]] + + +@dataclass(frozen=True, slots=True) +class NotServableOn: + """`NotServable` labeled with the replica whose /v1/models never listed the model.""" + + replica: str + last_result: Result[ModelsListResponse] | None + def await_servable( list_models: Callable[[float], Result[ModelsListResponse]], @@ -171,9 +186,41 @@ def await_servable( sleep(wait) +def await_servable_everywhere( + pollers: Mapping[str, ModelsPoller], + *, + model_name: str, + timeout: float, + interval: float, + request_timeout: float, + db_sync_seconds: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> Servable | NotServableOn: + """`await_servable` against every replica in turn, each with the full budget, so + the model is only servable once every replica has listed it.""" + for replica, list_models in pollers.items(): + match await_servable( + list_models, + model_name=model_name, + timeout=timeout, + interval=interval, + request_timeout=request_timeout, + db_sync_seconds=db_sync_seconds, + now=now, + sleep=sleep, + ): + case NotServable(last_result=last_result): + return NotServableOn(replica=replica, last_result=last_result) + case Servable(): + continue + return Servable() + + def servable_timeout_message( *, model_name: str, + replica: str, timeout: float, db_sync_seconds: float, last_result: Result[ModelsListResponse] | None, @@ -184,8 +231,8 @@ def servable_timeout_message( else "" ) return ( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " + f"model {model_name!r} was created but never became servable on {replica} " + f"within {timeout}s of first listing (plus {db_sync_seconds}s continuous " f"DB sync) after /model/new (control/data-plane propagation or " f"STORE_MODEL_IN_DB reload issue){last_error}" ) @@ -194,6 +241,7 @@ def servable_timeout_message( @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport + replicas: Mapping[str, Transport] poll_timeout: float = 120.0 poll_interval: float = 5.0 model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT @@ -310,12 +358,13 @@ class ProxyClient: model name passed". We poll the data-plane /v1/models until the model appears, then settle for the remainder of the propagation budget. - Both steps are needed, and the second is the one that matters at >1 replica. - The poll proves *a* replica is serving the model; it cannot prove they all - are, because every request opens a fresh connection and a load-balanced - Service routes each one independently -- so the caller's next request - re-rolls and can land on a replica that has not reloaded yet. Waiting out - PROPAGATION_TIMEOUT is what makes the model safe to use anywhere.""" + Both steps are needed. The poll asks every replica in PROXY_REPLICA_URLS + directly, so behind the stack's load balancer it proves each gateway serves + the model rather than whichever one the balancer routed the poll to. It still + cannot see the workers behind a gateway, nor any replica when only the + balancer address is configured (every request opens a fresh connection, so + the caller's next request re-rolls), so waiting out PROPAGATION_TIMEOUT is + what makes the model safe to use anywhere.""" model_id = unwrap( self.transport.post( "/model/new", @@ -330,16 +379,10 @@ class ProxyClient: return model_id def _await_model_servable(self, model_name: str, listed_for: str | None = None) -> None: - """Block until the data plane lists `model_name`, or fail at model_servable_timeout.""" - headers = self.transport.master if listed_for is None else self.transport.bearer(listed_for) - outcome = await_servable( - lambda poll_timeout: self.transport.get( - "/v1/models", - headers=headers, - params=NoBody(), - response_type=ModelsListResponse, - timeout=poll_timeout, - ), + """Block until every replica lists `model_name`, or fail at model_servable_timeout.""" + headers: Final = self.transport.master if listed_for is None else self.transport.bearer(listed_for) + outcome: Final = await_servable_everywhere( + {url: self._models_poller(transport, headers) for url, transport in self.replicas.items()}, model_name=model_name, timeout=self.model_servable_timeout, interval=self.model_servable_interval, @@ -351,16 +394,27 @@ class ProxyClient: match outcome: case Servable(): return - case NotServable(last_result=last_result): + case NotServableOn(replica=replica, last_result=last_result): raise AssertionError( servable_timeout_message( model_name=model_name, + replica=replica, timeout=self.model_servable_timeout, db_sync_seconds=self.model_servable_db_sync_seconds, last_result=last_result, ) ) + @staticmethod + def _models_poller(transport: Transport, headers: AuthHeaders) -> ModelsPoller: + return lambda poll_timeout: transport.get( + "/v1/models", + headers=headers, + params=NoBody(), + response_type=ModelsListResponse, + timeout=poll_timeout, + ) + def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via POST /model/update. The proxy overlays only the non-null fields and clears @@ -547,16 +601,20 @@ def build_proxy_client( base_url: str = PROXY_BASE_URL, master_key: str = MASTER_KEY, control_plane_base_url: str = CONTROL_PLANE_BASE_URL, + replica_urls: tuple[str, ...] = PROXY_REPLICA_URLS, ) -> ProxyClient: """The ProxyClient every suite's client is built from: a SplitTransport that routes LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two base URLs are the same for a monolithic proxy, so routing is then a no-op. + ``replica_urls`` (PROXY_REPLICA_URLS) names every data-plane replica the model + barrier polls directly; it is the data-plane URL itself unless the stack + exports each gateway's own address. The endpoints are injectable for callers that resolve the proxy some other way than ``e2e_config``'s env names (see ``claude_code/_env.py``); they must - pass all three together, since a caller that overrides only the data plane - would leave management calls pointed at the env default. + pass all four together, since a caller that overrides only the data plane + would leave management calls and the replica poll pointed at the env defaults. Test-to-proxy traffic always goes over the wire, in every E2E_FIXTURE_MODE: record and replay scope to the proxy's provider-bound calls via the @@ -573,8 +631,15 @@ def build_proxy_client( request_timeout=REQUEST_TIMEOUT, ), ) + replicas: Final = MappingProxyType( + { + url: HttpTransport(base_url=url, master_key=master_key, request_timeout=REQUEST_TIMEOUT) + for url in replica_urls + } + ) return ProxyClient( transport=split, + replicas=replicas, poll_timeout=POLL_TIMEOUT, poll_interval=POLL_INTERVAL, ) diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py new file mode 100644 index 00000000000..a508c97b9fb --- /dev/null +++ b/tests/e2e/test_proxy_client.py @@ -0,0 +1,87 @@ +"""Harness coverage for the model barrier that gates on every replica. + +No proxy needed and no ``e2e`` marker: this pins that a model registered through +the control plane only counts as servable once every configured replica lists it +on /v1/models, which is what keeps a two-gateway stack from handing a test a +model that one gateway has not reloaded yet. The fakes are plain pollers and an +injected clock, so nothing here monkeypatches anything. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from itertools import chain, repeat +from typing import Final + +import pytest + +from e2e_config import parse_replica_urls +from e2e_http import Success +from models import ModelListEntry, ModelsListResponse +from proxy_client import ModelsPoller, NotServableOn, Servable, await_servable_everywhere + +MODEL: Final = "gpt-under-test" +TIMEOUT: Final = 10.0 +INTERVAL: Final = 2.0 + + +@dataclass +class FakeClock: + elapsed: float = 0.0 + + def now(self) -> float: + return self.elapsed + + def sleep(self, seconds: float) -> None: + self.elapsed += seconds + + +def _listing(*model_ids: str) -> Success[ModelsListResponse]: + entries: Final = tuple(ModelListEntry(id=model_id) for model_id in model_ids) + return Success(status_code=200, data=ModelsListResponse(data=entries)) + + +def _poller(results: Iterable[Success[ModelsListResponse]]) -> ModelsPoller: + it: Final = iter(results) + return lambda _timeout: next(it) + + +def _await(pollers: Mapping[str, ModelsPoller]) -> Servable | NotServableOn: + clock: Final = FakeClock() + return await_servable_everywhere( + pollers, + model_name=MODEL, + timeout=TIMEOUT, + interval=INTERVAL, + request_timeout=5.0, + db_sync_seconds=0.0, + now=clock.now, + sleep=clock.sleep, + ) + + +class TestAwaitServableEverywhere: + @pytest.mark.parametrize("missing", ["gateway-1", "gateway-2"]) + def test_fails_on_the_replica_that_never_lists_the_model(self, missing: str) -> None: + pollers: Final = { + "gateway-1": _poller(repeat(_listing(MODEL))), + "gateway-2": _poller(repeat(_listing(MODEL))), + } | {missing: _poller(repeat(_listing()))} + assert _await(pollers) == NotServableOn(replica=missing, last_result=_listing()) + + def test_passes_once_every_replica_lists_the_model(self) -> None: + pollers: Final = { + "gateway-1": _poller(repeat(_listing(MODEL))), + "gateway-2": _poller(chain(repeat(_listing(), 2), repeat(_listing(MODEL)))), + } + assert _await(pollers) == Servable() + + +class TestParseReplicaUrls: + def test_splits_and_trims_the_gateway_addresses(self) -> None: + raw: Final = " http://127.0.0.1:4010/, http://127.0.0.1:4011 " + assert parse_replica_urls(raw, "http://lb") == ("http://127.0.0.1:4010", "http://127.0.0.1:4011") + + def test_falls_back_to_the_data_plane_address_when_unset(self) -> None: + assert parse_replica_urls("", "http://lb") == ("http://lb",) From be9a2ea7c97f945dc961c357b80f8b44f30c7746 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:20:20 -0700 Subject: [PATCH 29/63] fix(router): average every latency sample and drop the cost handler's dead division _average_latency skipped integer samples in the sum while counting them in the denominator, which contradicted its own Sequence[float | int] signature; it now averages every sample. Both success loggers in cost-based routing computed response_ms / completion_tokens and threw the result away, so a chat response with zero completion tokens raised ZeroDivisionError inside the handler. The proxy swallows and logs it, but the handler then skips that request's tpm and rpm update, so cost-based routing undercounts the deployment's usage. The QA run for the latency fix hit it on real gpt-5.5 traffic through /v1/chat/completions and /v1/messages --- litellm/router_strategy/lowest_cost.py | 11 +--- litellm/router_strategy/lowest_latency.py | 4 +- .../router_strategy/test_lowest_cost.py | 59 +++++++++++++++++++ .../router_strategy/test_lowest_latency.py | 2 - 4 files changed, 62 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_lowest_cost.py diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index b927df0c438..6eb4d86d280 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -1,6 +1,6 @@ #### What this does #### # picks based on response time (for streaming, this is time to first token) -from datetime import datetime, timedelta +from datetime import datetime from typing import Final import litellm @@ -52,16 +52,12 @@ class LowestCostLoggingHandler(CustomLogger): precise_minute: Final = f"{current_date}-{current_hour}-{current_minute}" cost_key: Final = f"{model_group}_map" - response_ms: Final[timedelta] = end_time - start_time - total_tokens = 0 if isinstance(response_obj, ModelResponse): _usage: Final = getattr(response_obj, "usage", None) if _usage is not None and isinstance(_usage, litellm.Usage): - completion_tokens: Final = _usage.completion_tokens total_tokens = _usage.total_tokens - float(response_ms.total_seconds() / completion_tokens) # ------------ # Update usage @@ -131,18 +127,13 @@ class LowestCostLoggingHandler(CustomLogger): current_minute: Final = datetime.now().strftime("%M") precise_minute: Final = f"{current_date}-{current_hour}-{current_minute}" - response_ms: Final[timedelta] = end_time - start_time - total_tokens = 0 if isinstance(response_obj, ModelResponse): _usage: Final = getattr(response_obj, "usage", None) if _usage is not None and isinstance(_usage, litellm.Usage): - completion_tokens: Final = _usage.completion_tokens total_tokens = _usage.total_tokens - float(response_ms.total_seconds() / completion_tokens) - # ------------ # Update usage # ------------ diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 4a1fb4f325a..bb6877a032b 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -26,10 +26,10 @@ class RoutingArgs(LiteLLMPydanticObjectBase): max_latency_list_size: int = 10 -def _average_latency(samples: Sequence[float | int]) -> float: +def _average_latency(samples: Sequence[float]) -> float: if not samples: return 0.0 - return sum(sample for sample in samples if isinstance(sample, float)) / len(samples) + return sum(samples) / len(samples) class LowestLatencyLoggingHandler(CustomLogger): diff --git a/tests/test_litellm/router_strategy/test_lowest_cost.py b/tests/test_litellm/router_strategy/test_lowest_cost.py new file mode 100644 index 00000000000..108053dddd9 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_lowest_cost.py @@ -0,0 +1,59 @@ +from datetime import datetime + +import pytest + +import litellm +from litellm.caching.caching import DualCache +from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler + +DEPLOYMENT_ID = "9876" +KWARGS = { + "litellm_params": { + "metadata": {"model_group": "gpt-5.5-pool"}, + "model_info": {"id": DEPLOYMENT_ID}, + } +} + + +def _chat_response_with_no_completion_tokens() -> litellm.ModelResponse: + return litellm.ModelResponse( + model="gpt-5.5", + choices=[{"index": 0, "message": {"role": "assistant", "content": ""}, "finish_reason": "length"}], + usage=litellm.Usage(prompt_tokens=12, completion_tokens=0, total_tokens=12), + ) + + +def _recorded_minute_counters(cache: DualCache) -> dict[str, int]: + cached = cache.get_cache(key="gpt-5.5-pool_map") or {} + minute_buckets = cached.get(DEPLOYMENT_ID, {}) + assert len(minute_buckets) == 1, f"expected one minute bucket, got {minute_buckets}" + return next(iter(minute_buckets.values())) + + +def test_log_success_event_counts_a_response_with_no_completion_tokens(): + cache = DualCache() + handler = LowestCostLoggingHandler(router_cache=cache) + + handler.log_success_event( + kwargs=KWARGS, + response_obj=_chat_response_with_no_completion_tokens(), + start_time=datetime(2026, 1, 1, 12, 0, 0), + end_time=datetime(2026, 1, 1, 12, 0, 2), + ) + + assert _recorded_minute_counters(cache) == {"tpm": 12, "rpm": 1} + + +@pytest.mark.asyncio +async def test_async_log_success_event_counts_a_response_with_no_completion_tokens(): + cache = DualCache() + handler = LowestCostLoggingHandler(router_cache=cache) + + await handler.async_log_success_event( + kwargs=KWARGS, + response_obj=_chat_response_with_no_completion_tokens(), + start_time=datetime(2026, 1, 1, 12, 0, 0), + end_time=datetime(2026, 1, 1, 12, 0, 2), + ) + + assert _recorded_minute_counters(cache) == {"tpm": 12, "rpm": 1} diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index 251b760a8d8..eb02459be68 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -172,8 +172,6 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): ids=["empty_latency_list", "minute_bucket_only_as_cost_based_routing_writes"], ) async def test_async_get_available_deployments_treats_missing_samples_as_zero_latency(cached_entry): - """A cached entry with no latency samples (cost-based routing shares the group's map key and writes - minute buckets only) must count as 0 latency, like an unseen deployment, instead of dividing by zero.""" cache = DualCache() handler = LowestLatencyLoggingHandler(router_cache=cache) cache.set_cache( From 5a06845db18a2e1709bcd5e696cafdd4c190bd7e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:33:15 -0700 Subject: [PATCH 30/63] fix(ci): mask only credential-length values in the e2e-changed log A one-character value in the provider secret bundle was masked too, which turned every 1 in the run log into ***, including the pass numbers and the gateway addresses, so the only public diagnostics were unreadable --- .github/e2e-stack/secrets_to_env.py | 3 ++- .../code_coverage_tests/test_e2e_changed_gate.py | 16 ++++++++++++++++ tests/e2e/CONTRIBUTING.md | 2 +- 3 files changed, 19 insertions(+), 2 deletions(-) diff --git a/.github/e2e-stack/secrets_to_env.py b/.github/e2e-stack/secrets_to_env.py index 99e717b271d..e7870a32e42 100644 --- a/.github/e2e-stack/secrets_to_env.py +++ b/.github/e2e-stack/secrets_to_env.py @@ -8,6 +8,7 @@ from pydantic import TypeAdapter, ValidationError secrets_adapter: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) ENV_NAME: Final = re.compile(r"[A-Za-z_][A-Za-z0-9_]*") +MIN_MASKED_LENGTH: Final = 8 def main() -> int: @@ -25,7 +26,7 @@ def main() -> int: _ = sys.stderr.write("environment names or values cannot be represented in both bash and dotenv\n") return 1 for value in secrets.values(): - if value: + if len(value) >= MIN_MASKED_LENGTH: _ = sys.stdout.write(f"::add-mask::{value.replace('%', '%25')}\n") sys.stdout.flush() lines: Final = tuple(f"{key}='{value}'" for key, value in secrets.items() if value) diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 7a82a8ebb60..2e95c33cf20 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -7,6 +7,7 @@ from typing import Final import pytest GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py" +SECRETS_TO_ENV: Final = GATE.with_name("secrets_to_env.py") SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py") @@ -53,3 +54,18 @@ def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) assert result.returncode == 1 + + +def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_path: Path) -> None: + env_path: Final = tmp_path / ".env" + + result: Final = subprocess.run( + [sys.executable, str(SECRETS_TO_ENV), str(env_path)], + input='{"FLAG": "1", "API_KEY": "sk-0123456789abcdef"}', + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stderr + assert result.stdout == "::add-mask::sk-0123456789abcdef\n" + assert env_path.read_text() == "FLAG='1'\nAPI_KEY='sk-0123456789abcdef'\n" diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 7b90b9de5c8..0b8859d9260 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -64,7 +64,7 @@ Repository admins must require the `e2e-changed-tests` status check for merging Credentials come from the existing AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`. The OIDC role must trust only `repo:BerriAI/litellm:environment:e2e-changed` with audience `sts.amazonaws.com` and have read access only to these secrets. The short-lived reader credentials are scoped to the fetch step. Provider credentials must cover the selected suites, including Datadog credentials when logging or MCP tests need them; missing credentials fail the run. `up.sh` refuses to start without `DD_API_KEY`, because the stack's gateway config enables the Datadog callback for every run and a gateway booted without the key fails readiness. Keep provider credentials dedicated to this lane with only the permissions those tests need -Fetched values are masked before use, and credential files and raw output are private to the runner. Public logs contain selected file names, counts, and pass status; raw pytest output, reports, and stack logs are not uploaded or printed. The workflow removes them and the credential files during cleanup. To diagnose a failed pass, reproduce the selected files locally with the appropriate credentials and inspect the local logs +Fetched values of eight characters or more are masked before use, while shorter values such as flags stay unmasked because masking a one-character value would blank every matching digit in the log, and credential files and raw output are private to the runner. Public logs contain selected file names, counts, and pass status; raw pytest output, reports, and stack logs are not uploaded or printed. The workflow removes them and the credential files during cleanup. To diagnose a failed pass, reproduce the selected files locally with the appropriate credentials and inspect the local logs To reproduce the CI topology on a dedicated machine, `bash .github/e2e-stack/up.sh` reads `tests/e2e/.env`, writes `stack.env` under `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}`, and `bash .github/e2e-stack/down.sh` stops it. Keep this directory private and remove its credential files and logs after use From 4c00a6e189d95f34a72036802937b38769f561b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 16:35:32 -0700 Subject: [PATCH 31/63] fix(batches): account a batch's cost once, from the first retrieve that sees it final Every retrieve of a batch through the proxy shares one spend row, the batch id plus the batch cost suffix, and spend log inserts skip duplicates. A poll that landed while the batch was still validating or in progress wrote that row at $0 and no later retrieve could overwrite it, and every completed retrieve after the first added the cost to the key, team, and user counters again with no new row to show for it. The cost callback now writes nothing for a batch retrieve until the batch is final, releasing the poll's budget reservation instead, and once it is final it charges only when no spend row for that batch is queued for flush or already stored. Batch cost rows are flushed to the database right away so a second instance sees them, and the logger prices a batch only once it is final, which also covers a failed batch that never produced an output file. --- litellm/batches/batch_utils.py | 20 ++ litellm/litellm_core_utils/litellm_logging.py | 13 +- litellm/proxy/db/db_spend_update_writer.py | 3 +- .../proxy/hooks/proxy_track_cost_callback.py | 56 ++++- .../openai_files_endpoints/common_utils.py | 8 +- litellm/proxy/utils.py | 10 +- .../test_litellm/batches/test_batch_utils.py | 55 +++- .../test_litellm_logging.py | 82 ++++++ .../proxy/db/test_db_spend_update_writer.py | 10 +- .../hooks/test_proxy_track_cost_callback.py | 238 +++++++++++++----- .../prisma_and_spend/test_spend_functions.py | 15 ++ 11 files changed, 428 insertions(+), 82 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 97be5f77d79..eaac3bf0e9f 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -25,6 +25,26 @@ class BatchCostUsageResult: failed_requests: int +_TERMINAL_BATCH_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) + + +def batch_cost_is_final(batch: Batch) -> bool: + """Whether this retrieve of the batch is the one to account its cost from. + + A batch still in flight has nothing to price, and a "completed" batch can report + no output_file_id for a moment before the output populates; pricing either records + $0 under the batch's single spend row and pins it there. Final means a completed + batch whose output file has arrived or whose counts prove no line succeeded, or + any other terminal status (failed, cancelled, expired). + """ + if batch.status not in _TERMINAL_BATCH_STATUSES: + return False + if batch.status != "completed" or batch.output_file_id is not None: + return True + request_counts: Final = batch.request_counts + return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 + + async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c31c4323157..09ddd1b9720 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -36,7 +36,7 @@ from litellm._logging import ( verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import _handle_completed_batch +from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final from litellm.caching.caching import DualCache, InMemoryCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.constants import ( @@ -2899,13 +2899,6 @@ class Logging(LiteLLMLoggingBaseClass): ): # polling job will query these frequently, don't spam db logs return - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) - - # check if file id is a unified file id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(result.id) - batch_cost: Final = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) @@ -2913,9 +2906,7 @@ class Logging(LiteLLMLoggingBaseClass): batch_failed_requests: Final = kwargs.get("batch_failed_requests", None) has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models)) - should_compute_batch_data: Final = ( - not is_base64_unified_file_id or not has_explicit_batch_data and result.status == "completed" - ) + should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result) if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e6880d521f1..ff48b00dc70 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -82,6 +82,7 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) +IMMEDIATE_FLUSH_CALL_TYPES: Final = RESPONSES_SESSION_CALL_TYPES | frozenset({CallTypes.aretrieve_batch.value}) class _SpendBatch(Protocol): @@ -939,7 +940,7 @@ class DBSpendUpdateWriter: from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush await enqueue_spend_logs(prisma_client, (payload,)) - if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: + if payload.get("call_type") in IMMEDIATE_FLUSH_CALL_TYPES: request_spend_log_flush() else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 7254b05db2e..95e61fdd98c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger +from litellm.batches.batch_utils import batch_cost_is_final from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( @@ -33,17 +34,21 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, get_request_model_access_groups, + get_spend_logs_id, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( CallTypes, + LiteLLMBatch, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: - from litellm.proxy.utils import ProxyLogging + from prisma.types import LiteLLM_SpendLogsWhereUniqueInput + + from litellm.proxy.utils import PrismaClient, ProxyLogging _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { @@ -224,6 +229,7 @@ class _ProxyDBLogger(CustomLogger): ): from litellm.proxy.proxy_server import ( increment_spend_counters, + prisma_client, proxy_logging_obj, update_cache, ) @@ -248,6 +254,18 @@ class _ProxyDBLogger(CustomLogger): ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) + if ( + isinstance(completion_response, LiteLLMBatch) + and kwargs.get("call_type") == CallTypes.aretrieve_batch.value + ): + batch_spend_log_id: Final = get_spend_logs_id( + CallTypes.aretrieve_batch.value, completion_response.model_dump(), kwargs + ) + if not await _batch_cost_is_trackable_now( + batch=completion_response, spend_log_id=batch_spend_log_id, prisma_client=prisma_client + ): + await _release_budget_reservation(budget_reservation=budget_reservation) + return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) @@ -491,6 +509,42 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value +async def _batch_cost_is_trackable_now( + batch: LiteLLMBatch, spend_log_id: str | None, prisma_client: "PrismaClient | None" +) -> bool: + """A batch is billed exactly once, from the first retrieve that sees it final. + + Every retrieve of one batch shares a single spend row (its id plus the batch cost + suffix), so a poll that lands before the output exists would write that row at $0 + and pin it there, and every retrieve after the first would add the cost to the + key, team, and user counters again. + """ + if not batch_cost_is_final(batch): + verbose_proxy_logger.debug("Cost tracking deferred for batch %s still in status %s", batch.id, batch.status) + return False + if prisma_client is None or spend_log_id is None: + return True + if not await _spend_log_already_recorded(prisma_client=prisma_client, request_id=spend_log_id): + return True + verbose_proxy_logger.debug( + "Cost tracking skipped for batch %s: spend row %s already recorded", batch.id, spend_log_id + ) + return False + + +async def _spend_log_already_recorded(prisma_client: "PrismaClient", request_id: str) -> bool: + from litellm.proxy.utils import spend_log_is_queued + + if await spend_log_is_queued(prisma_client, request_id): + return True + spend_log_row: Final[LiteLLM_SpendLogsWhereUniqueInput] = {"request_id": request_id} + try: + return await prisma_client.db.litellm_spendlogs.find_unique(where=spend_log_row) is not None + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreadable DB must not drop the batch's only spend row + verbose_proxy_logger.warning("Could not check for an existing spend row %s, tracking anyway: %s", request_id, e) + return False + + def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 15eeddbc489..b1f282a0978 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -15,6 +15,7 @@ from typing import ( runtime_checkable, ) +from litellm.batches.batch_utils import batch_cost_is_final from litellm.proxy._types import ProxyException from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -1357,12 +1358,7 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: enumerated the batch and none succeeded. A zero or unknown total means counts are unreported, so stay eligible and let the next poller pass revisit it. (#37713) """ - if response.output_file_id is not None: - return True - request_counts = response.request_counts - if request_counts is None: - return False - return request_counts.total > 0 and request_counts.completed == 0 + return batch_cost_is_final(response) async def update_batch_in_database( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index accf7b720fb..f64b51bc6c6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6251,7 +6251,9 @@ def request_spend_log_flush() -> None: The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. - Repeated requests coalesce into the monitor's next pass, so the batching holds. + A batch's cost row is what every other worker checks before charging the same batch + again, so it cannot wait either. Repeated requests coalesce into the monitor's next + pass, so the batching holds. """ PrismaClient.spend_log_flush_requested.set() @@ -6266,6 +6268,12 @@ async def _wait_for_spend_log_flush_request(interval: float) -> bool: return True +async def spend_log_is_queued(prisma_client: PrismaClient, request_id: str) -> bool: + """Whether a spend log with ``request_id`` is still waiting for the next flush.""" + async with prisma_client._spend_log_transactions_lock: + return any(row.get("request_id") == request_id for row in prisma_client.spend_log_transactions) + + async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c86c7c4df03..8d4f68164b4 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -21,11 +21,12 @@ from types import MappingProxyType import httpx import pytest import respx +from openai.types.batch import BatchRequestCounts import litellm import litellm.batches.batch_utils as bu -from litellm.types.utils import Usage +from litellm.types.utils import LiteLLMBatch, Usage # --------------------------------------------------------------------------- # # Builders for batch OUTPUT file rows. @@ -1718,3 +1719,55 @@ def test_unparsable_bedrock_batch_usage_warns(caplog): assert usage.total_tokens == 0 assert "does not understand" in caplog.text assert "inputTextTokenCount" in caplog.text + + +# --------------------------------------------------------------------------- # +# batch_cost_is_final +# --------------------------------------------------------------------------- # + +def _retrieved_batch( + status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None +) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + request_counts=counts, + ) + + +class TestBatchCostIsFinal: + """Every retrieve of one batch writes the same spend row, so the first retrieve + that prices it decides the row for good. A poll before the output exists must + therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048).""" + + @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) + def test_in_flight_batch_is_not_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status)) is False + + def test_completed_with_output_is_final(self): + assert bu.batch_cost_is_final(_retrieved_batch("completed", output_file_id="file-out")) is True + + def test_completed_without_output_and_unknown_counts_is_not_final(self): + assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False + + def test_completed_without_output_and_zero_counts_is_not_final(self): + counts = BatchRequestCounts(total=0, completed=0, failed=0) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + def test_completed_without_output_but_successful_lines_is_not_final(self): + counts = BatchRequestCounts(total=2, completed=2, failed=0) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False + + def test_completed_without_output_and_every_line_failed_is_final(self): + counts = BatchRequestCounts(total=2, completed=0, failed=2) + assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is True + + @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) + def test_other_terminal_statuses_are_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status)) is True 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 16a99713a06..4583429bd11 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -632,6 +632,88 @@ class TestRetrieveBatchCostPassesModelIdentity: assert captured["model_info"]["input_cost_per_token"] == 0.0 +class TestRetrieveBatchPricesOnlyFinalBatches: + """Regression (LIT-7048): retrieving a provider-id batch priced it on every poll. + + Every retrieve of one batch logs under the same spend row, so pricing a poll + that landed before the output existed wrote that row at $0 and pinned it there. + Only a final batch gets priced; an in-flight poll carries no cost at all. + """ + + @staticmethod + def _logging_obj() -> LitellmLogging: + obj = LitellmLogging( + model="gpt-5.6-luna", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="batch-call-2", + function_id="f", + ) + obj.custom_llm_provider = "openai" + return obj + + @staticmethod + def _batch(status: str, output_file_id: str | None): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch_6a9c99e185588190877d391f8b9d7f8a", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("status", "output_file_id"), + [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None)], + ) + async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None: + from litellm.litellm_core_utils import litellm_logging as logging_module + + handle_completed_batch = AsyncMock() + monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + batch = self._batch(status, output_file_id) + + with contextlib.suppress(Exception): + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + + handle_completed_batch.assert_not_awaited() + assert "response_cost" not in batch._hidden_params + + @pytest.mark.asyncio + async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None: + from litellm.batches.batch_utils import BatchCostUsageResult + from litellm.litellm_core_utils import litellm_logging as logging_module + from litellm.types.utils import Usage + + handle_completed_batch = AsyncMock( + return_value=BatchCostUsageResult( + cost=8e-06, + usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35), + models=["gpt-5.6-luna"], + successful_requests=2, + failed_requests=0, + ) + ) + monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + batch = self._batch("completed", "file-out") + + with contextlib.suppress(Exception): + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + + handle_completed_batch.assert_awaited_once() + assert batch._hidden_params["response_cost"] == 8e-06 + assert batch.usage is not None + assert batch.usage.total_tokens == 35 + + class TestAnthropicPassthroughCustomPricing: """Verify the Anthropic pass-through handler forwards custom pricing.""" diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 11ef911de3e..e1b2d151c6d 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2934,12 +2934,16 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey @pytest.mark.asyncio @pytest.mark.parametrize( "call_type, expects_flush", - [("aresponses", True), ("responses", True), ("acompletion", False)], + [("aresponses", True), ("responses", True), ("aretrieve_batch", True), ("acompletion", False)], ) -async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(call_type: str, expects_flush: bool): +async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back( + call_type: str, expects_flush: bool +): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a - Responses row cannot sit in this worker's queue until the monitor's next poll. + Responses row cannot sit in this worker's queue until the monitor's next poll. A + batch's cost row is what another worker checks before charging the same batch again + (LIT-7048), so it cannot wait either. """ from litellm.proxy.utils import PrismaClient diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 8043a1aca3f..b7037e8d621 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,4 +1,3 @@ - from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -7,6 +6,7 @@ import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( + _batch_cost_is_trackable_now, _get_budget_reservation_from_metadata, _ProxyDBLogger, _should_track_cost_callback, @@ -70,9 +70,7 @@ async def test_async_post_call_failure_hook(): # Check that metadata was properly updated assert "litellm_params" in call_args["kwargs"] - assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == { - "request_id": "test_request_id" - } + assert call_args["kwargs"]["litellm_params"]["proxy_server_request"] == {"request_id": "test_request_id"} metadata = call_args["kwargs"]["litellm_params"]["metadata"] assert metadata["user_api_key"] == "test_api_key" assert metadata["status"] == "failure" @@ -336,9 +334,7 @@ async def test_should_continue_failure_tracking_when_budget_release_fails(): ) assert mock_invalidate_budget_reservation_counters.await_count == 1 assert ( - mock_invalidate_budget_reservation_counters.await_args.kwargs[ - "budget_reservation" - ] + mock_invalidate_budget_reservation_counters.await_args.kwargs["budget_reservation"] is user_api_key_dict.budget_reservation ) assert user_api_key_dict.budget_reservation["finalized"] is True @@ -433,36 +429,21 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): "entries": [{"counter_key": "spend:key:test_api_key"}], } + assert _get_budget_reservation_from_metadata(metadata={"user_api_key_auth": dict(UserAPIKeyAuth())}) is None assert ( _get_budget_reservation_from_metadata( - metadata={"user_api_key_auth": dict(UserAPIKeyAuth())} - ) - is None - ) - assert ( - _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": UserAPIKeyAuth( - budget_reservation=budget_reservation - ) - } + metadata={"user_api_key_auth": UserAPIKeyAuth(budget_reservation=budget_reservation)} ) == budget_reservation ) assert ( _get_budget_reservation_from_metadata( - metadata={ - "user_api_key_auth": dict( - UserAPIKeyAuth(budget_reservation=budget_reservation) - ) - } + metadata={"user_api_key_auth": dict(UserAPIKeyAuth(budget_reservation=budget_reservation))} ) == budget_reservation ) assert ( - _get_budget_reservation_from_metadata( - metadata={"user_api_key_budget_reservation": budget_reservation} - ) + _get_budget_reservation_from_metadata(metadata={"user_api_key_budget_reservation": budget_reservation}) is budget_reservation ) @@ -470,9 +451,7 @@ def test_get_budget_reservation_from_metadata_handles_dict_auth_object(): @pytest.mark.asyncio async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails(): proxy_logging_obj = MagicMock() - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=Exception("db unavailable") - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=Exception("db unavailable")) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -508,9 +487,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u async def test_update_database_and_spend_counters_preserves_db_exception_when_release_fails(): proxy_logging_obj = MagicMock() db_exception = RuntimeError("db unavailable") - proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock( - side_effect=db_exception - ) + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) increment_spend_counters = AsyncMock() budget_reservation = {"reserved_cost": 0.5, "entries": []} @@ -554,12 +531,8 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re budget_reservation=budget_reservation, ) assert mock_log_exception.call_count == 2 - mock_log_exception.assert_any_call( - "Failed to release budget reservation after database update failed" - ) - mock_log_exception.assert_any_call( - "Failed to invalidate budget reservation counters after release failed" - ) + mock_log_exception.assert_any_call("Failed to release budget reservation after database update failed") + mock_log_exception.assert_any_call("Failed to invalidate budget reservation counters after release failed") increment_spend_counters.assert_not_awaited() @@ -778,6 +751,169 @@ async def test_track_cost_callback_defers_in_progress_background_interaction(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +def _batch_retrieve_kwargs(call_type: str, reservation: dict | None = None) -> dict: + metadata = { + "user_api_key": "hashed_key", + "user_api_key_user_id": "user-1", + "user_api_key_team_id": "team-1", + **({"user_api_key_budget_reservation": reservation} if reservation is not None else {}), + } + return { + "call_type": call_type, + "model": "gpt-5.6-luna", + "litellm_call_id": "test-call-id", + "litellm_params": {"metadata": metadata}, + "standard_logging_object": {"response_cost": 0.0, "request_tags": None}, + "stream": False, + } + + +def _retrieved_batch(status: str, output_file_id: str | None): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch_abc", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status=status, + output_file_id=output_file_id, + ) + + +def _prisma_client_with(queued_request_ids: tuple[str, ...], stored_row: object) -> MagicMock: + import asyncio + + prisma_client = MagicMock() + prisma_client._spend_log_transactions_lock = asyncio.Lock() + prisma_client.spend_log_transactions = [{"request_id": request_id} for request_id in queued_request_ids] + prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(return_value=stored_row) + return prisma_client + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("status", "output_file_id", "spend_log_id", "prisma_client", "trackable"), + [ + ("in_progress", None, "batch_abc_batch_cost", None, False), + ("in_progress", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), + ("completed", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), + ("completed", "file-out", "batch_abc_batch_cost", None, True), + ("completed", "file-out", None, _prisma_client_with((), None), True), + ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with(("batch_abc_batch_cost",), None), False), + ( + "completed", + "file-out", + "batch_abc_batch_cost", + _prisma_client_with((), {"request_id": "batch_abc_batch_cost"}), + False, + ), + ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with((), None), True), + ("failed", None, "batch_abc_batch_cost", _prisma_client_with((), None), True), + ], + ids=[ + "in_progress_without_db", + "in_progress_never_consults_db", + "completed_without_output_yet", + "final_without_db", + "final_without_spend_log_id", + "final_row_queued_for_flush", + "final_row_already_stored", + "final_first_sighting", + "failed_first_sighting", + ], +) +async def test_batch_cost_is_trackable_now(status, output_file_id, spend_log_id, prisma_client, trackable): + """ + A batch is billed from the first retrieve that sees it final and never again: + a poll before that wrote the shared spend row at $0 and pinned it there, and + every completed retrieve after the first charged the key again (LIT-7048). + """ + assert ( + await _batch_cost_is_trackable_now( + batch=_retrieved_batch(status, output_file_id), spend_log_id=spend_log_id, prisma_client=prisma_client + ) + is trackable + ) + + +@pytest.mark.asyncio +async def test_batch_cost_is_trackable_now_when_the_spend_row_lookup_fails(): + """An unreadable spend log table must not drop the batch's only spend row.""" + prisma_client = _prisma_client_with((), None) + prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert ( + await _batch_cost_is_trackable_now( + batch=_retrieved_batch("completed", "file-out"), + spend_log_id="batch_abc_batch_cost", + prisma_client=prisma_client, + ) + is True + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "status", "output_file_id", "stored_row", "charged"), + [ + ("aretrieve_batch", "in_progress", None, None, False), + ("aretrieve_batch", "completed", "file-out", {"request_id": "batch_abc_batch_cost"}, False), + ("aretrieve_batch", "completed", "file-out", None, True), + ("acreate_batch", "validating", None, None, True), + ], + ids=["retrieve_before_final", "retrieve_already_recorded", "retrieve_first_final", "create_before_final"], +) +async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs and whether the poll's reservation is handed back is the whole observable contract of the gate + call_type, status, output_file_id, stored_row, charged +): + """ + Only retrieves are gated, since creating a batch is its own billable request. + A retrieve that writes nothing hands its budget reservation back instead. + """ + logger = _ProxyDBLogger() + budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} + kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) + + with ( + patch( # test-quality-ok: prisma_client is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.prisma_client", _prisma_client_with((), stored_row) + ), + patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock + ), + patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam + "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock + ), + patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam + "litellm.proxy.proxy_server.proxy_logging_obj" + ) as mock_proxy_logging, + patch( # test-quality-ok: the release is imported inside the callback's helper, no seam + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", new_callable=AsyncMock + ) as mock_release_budget_reservation, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=_retrieved_batch(status, output_file_id), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + mock_proxy_logging.failed_tracking_alert.assert_not_called() + if charged: + mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() + mock_release_budget_reservation.assert_not_awaited() + else: + mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() + mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation) + + def _in_progress_interaction_kwargs(reservation: dict) -> dict: return { "call_type": "acreate_interaction", @@ -1101,10 +1237,7 @@ async def test_async_post_call_failure_hook_propagates_trace_id_from_logging_obj # standard_logging_object should have been propagated from logging obj assert call_kwargs.get("standard_logging_object") is not None - assert ( - call_kwargs["standard_logging_object"]["trace_id"] - == "trace-id-from-logging-obj" - ) + assert call_kwargs["standard_logging_object"]["trace_id"] == "trace-id-from-logging-obj" # litellm_trace_id should also be propagated as a fallback assert call_kwargs.get("litellm_trace_id") == "trace-id-from-logging-obj" @@ -1691,9 +1824,7 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): "metadata": {}, "proxy_server_request": {"request_id": "rid"}, "response_cost": 3.5e-05, - "combined_usage_object": Usage( - prompt_tokens=30, completion_tokens=1, total_tokens=31 - ), + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), } with patch( @@ -1772,15 +1903,10 @@ async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): assert mock_increment.call_args.kwargs["team_id"] == "team-123" assert mock_increment.call_args.kwargs["org_id"] == "org-456" - update_kwargs = ( - mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs - ) + update_kwargs = mock_proxy_logging.db_spend_update_writer.update_database.await_args.kwargs assert update_kwargs["user_id"] == "mcp-user@example.com" assert update_kwargs["team_id"] == "team-123" - assert ( - kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] - == "mcp-user@example.com" - ) + assert kwargs["litellm_params"]["metadata"]["user_api_key_user_id"] == "mcp-user@example.com" @pytest.mark.parametrize( @@ -1828,9 +1954,7 @@ def test_should_track_cost_callback_pass_through_without_owner(call_type, expect ], ) @pytest.mark.asyncio -async def test_track_cost_callback_logs_unauthenticated_pass_through_request( - call_type, expect_spend_log -): +async def test_track_cost_callback_logs_unauthenticated_pass_through_request(call_type, expect_spend_log): """Regression for LIT-3782: a pass-through request with auth=false reaches the cost callback with no key/user/team/end-user. Before the fix the spend-log write was skipped and the request never appeared in request/usage logs. It @@ -1876,9 +2000,7 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request( end_time=datetime.now(), ) - assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == ( - 1 if expect_spend_log else 0 - ) + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if expect_spend_log else 0) class _FakeDeploymentLookup: diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a1eb88a7834..fc97d760226 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -6,6 +6,7 @@ Symbols pinned here: - ``update_spend_logs_job`` - ``_monitor_spend_logs_queue`` - ``_raise_failed_update_spend_exception`` + - ``spend_log_is_queued`` """ from __future__ import annotations @@ -22,6 +23,7 @@ from litellm.proxy.utils import ( _monitor_spend_logs_queue, _raise_failed_update_spend_exception, drain_spend_logs_queue, + spend_log_is_queued, update_daily_tag_spend, update_spend, update_spend_logs_job, @@ -629,3 +631,16 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None: with pytest.raises(ValueError, match="specific"): asyncio.run(_runner()) + + +@pytest.mark.asyncio +async def test_spend_log_is_queued_matches_only_rows_awaiting_flush( + mock_prisma_client: Any, make_spend_log_row: Any +) -> None: + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="batch_abc_batch_cost")] + + assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is True + assert await spend_log_is_queued(mock_prisma_client, "batch_abc") is False + + mock_prisma_client.spend_log_transactions = [] + assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is False From 635bb3a2096fbb4f4c8899574807c049bb8e4825 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 17:08:42 -0700 Subject: [PATCH 32/63] feat(cost-map): add azure_ai/gpt-6-astra Foundry pricing A gpt-6-astra deployment on a Foundry project reached through the azure_ai route had no cost map entry of its own, so it resolved to the OpenAI gpt-6-astra card: missing from the azure_ai/* wildcard listing, flex and priority prices and /v1/batch it does not sell, and no none reasoning effort. Add azure_ai/gpt-6-astra mirroring the azure/gpt-6-astra Standard Global sheet the way azure_ai/gpt-5.5 mirrors azure/gpt-5.5, and extend the cost, reasoning-effort, and wildcard listing tests to the Foundry route. --- ...odel_prices_and_context_window_backup.json | 43 +++++++++++++++++++ model_prices_and_context_window.json | 43 +++++++++++++++++++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 15 +++++-- .../proxy/auth/test_model_checks.py | 19 ++++++++ .../test_reasoning_effort_capability.py | 16 +++++-- 5 files changed, 129 insertions(+), 7 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4273ec54472..ac7407c2608 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3485,6 +3485,49 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "source": "https://ai.azure.com/catalog/models/gpt-6-astra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4273ec54472..ac7407c2608 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3485,6 +3485,49 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "azure_ai/gpt-6-astra": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_above_272k_tokens": 2e-06, + "input_cost_per_token": 1e-05, + "input_cost_per_token_above_272k_tokens": 2e-05, + "litellm_provider": "azure_ai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_above_272k_tokens": 7.5e-05, + "source": "https://ai.azure.com/catalog/models/gpt-6-astra", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, "azure_ai/gpt-5.5": { "deprecation_date": "2027-10-26", "cache_read_input_token_cost": 5e-07, 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 df680b7cb0e..73f1a19d85c 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 @@ -2008,7 +2008,14 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, assert round(completion_cost, 10) == round(output_cost * completion_tokens, 10) -@pytest.mark.parametrize("model,zone_multiplier", [("azure/gpt-6-astra", 1.0), ("azure/us/gpt-6-astra", 1.1)]) +@pytest.mark.parametrize( + "model,custom_llm_provider,zone_multiplier", + [ + ("azure/gpt-6-astra", "azure", 1.0), + ("azure/us/gpt-6-astra", "azure", 1.1), + ("azure_ai/gpt-6-astra", "azure_ai", 1.0), + ], +) @pytest.mark.parametrize( "prompt_tokens,input_side_multiplier,output_multiplier", [(100000, 1.0, 1.0), (300000, 2.0, 1.5)], @@ -2016,6 +2023,7 @@ def test_generic_cost_per_token_azure_gpt56(_local_model_cost_map, def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( _local_model_cost_map, model, + custom_llm_provider, zone_multiplier, prompt_tokens, input_side_multiplier, @@ -2023,7 +2031,8 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( ): """Microsoft Foundry sells gpt-6-astra at the OpenAI rates: $10 input, $1 cache read, $12.50 cache write, $50 output per 1M tokens on Standard Global, with the input side doubling and output 1.5x above 272K - prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate. + prompt tokens. Standard US Data Zone carries the usual 10% uplift on every rate. A Foundry + deployment reached through the azure_ai route bills the same Standard Global sheet. """ cached_tokens = 50000 cache_write_tokens = 40000 @@ -2041,7 +2050,7 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( prompt_cost, completion_cost = generic_cost_per_token( model=model, usage=usage, - custom_llm_provider="azure", + custom_llm_provider=custom_llm_provider, ) input_side = zone_multiplier * input_side_multiplier diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index d58683fd1e5..eb48f70d5da 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -857,6 +857,25 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): litellm.add_known_models(model_cost_map={}) assert fake_model not in litellm.models_by_provider["vertex_ai"] + +def test_azure_ai_wildcard_lists_the_foundry_gpt_6_astra_entry(monkeypatch): + """A Foundry (azure_ai) deployment of gpt-6-astra only shows up under an azure_ai/* wildcard + when the cost map carries its own azure_ai/ entry; the azure/ entry from the OpenAI-on-Azure + price sheet never reaches the Foundry provider list (LIT-7081).""" + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + foundry_key = "azure_ai/gpt-6-astra" + local_entry = litellm.get_model_cost_map(url="")[foundry_key] + try: + litellm.add_known_models(model_cost_map={foundry_key: local_entry}) + assert foundry_key in get_known_models_from_wildcard("azure_ai/*") + finally: + litellm.azure_ai_models.discard(foundry_key) + litellm.add_known_models(model_cost_map={}) + + def test_get_complete_model_list_drops_no_default_models_sentinel(): from litellm.proxy.auth.model_checks import get_complete_model_list diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index f181370455d..b7499d1c975 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -389,14 +389,22 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: "max", ) - @pytest.mark.parametrize("model", ["azure/gpt-6-astra", "azure/us/gpt-6-astra"]) - def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model): + @pytest.mark.parametrize( + "model,custom_llm_provider", + [ + ("azure/gpt-6-astra", "azure"), + ("azure/us/gpt-6-astra", "azure"), + ("azure_ai/gpt-6-astra", "azure_ai"), + ], + ) + def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): """Microsoft Foundry serves the same model but its API accepts reasoning_effort none (verified live: 200 with zero reasoning tokens, and it unlocks temperature), which - OpenAI's rejects, so an Azure deployment offers none on top of low through max.""" + OpenAI's rejects, so an Azure deployment offers none on top of low through max, whether + it is reached through the azure route or the azure_ai (Foundry) route.""" from litellm.utils import _get_model_info_helper - model_info = dict(_get_model_info_helper(model=model, custom_llm_provider="azure")) + model_info = dict(_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)) assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( "none", From da705c947512e436bc740f8637f34246a1ee7c47 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 17:49:29 -0700 Subject: [PATCH 33/63] refactor(ui): render the Virtual Keys page without the legacy user dashboard The API Keys route mounted the pre-App-Router UserDashboard component, whose beforeunload handler cleared sessionStorage on every refresh of the Virtual Keys page. That wiped the Playground chat history and model, the logs live-tail preference, and everything else other pages keep in session storage. The same component also re-decoded the login token, re-fetched teams, and wrote cache entries nothing read. ApiKeysDashboard now renders VirtualKeysTable and the Create Key button directly, taking identity and role from useAuthorized like every other page. Create Key is hidden for view-only roles, which the proxy already rejects on /key/generate. The legacy component, its test, the fetch_teams helper, and their grandfathered eslint suppressions are removed, and the ProxySettings type moves to useProxySettings. --- ui/litellm-dashboard/eslint-budgets.json | 4 +- ui/litellm-dashboard/eslint-suppressions.json | 24 +- .../api-keys/ApiKeysDashboard.test.tsx | 109 +++++--- .../(dashboard)/api-keys/ApiKeysDashboard.tsx | 44 ++-- .../hooks/proxySettings/useProxySettings.ts | 2 + .../old-usage/_components/usage.tsx | 2 +- .../common_components/fetch_teams.tsx | 18 -- .../src/components/networking.tsx | 9 +- .../src/components/user_dashboard.test.tsx | 133 ---------- .../src/components/user_dashboard.tsx | 239 ------------------ 10 files changed, 99 insertions(+), 485 deletions(-) delete mode 100644 ui/litellm-dashboard/src/components/common_components/fetch_teams.tsx delete mode 100644 ui/litellm-dashboard/src/components/user_dashboard.test.tsx delete mode 100644 ui/litellm-dashboard/src/components/user_dashboard.tsx diff --git a/ui/litellm-dashboard/eslint-budgets.json b/ui/litellm-dashboard/eslint-budgets.json index e98cea9261e..44294b5fa97 100644 --- a/ui/litellm-dashboard/eslint-budgets.json +++ b/ui/litellm-dashboard/eslint-budgets.json @@ -3,8 +3,8 @@ "no-console": { "max": 12, "target": 0 }, "complexity": { "max": 140, "target": 80 }, "max-depth": { "max": 70, "target": 30 }, - "local/no-large-inline-object-arg": { "max": 554, "target": 300 }, - "local/no-long-condition-chain": { "max": 265, "target": 120 }, + "local/no-large-inline-object-arg": { "max": 551, "target": 300 }, + "local/no-long-condition-chain": { "max": 196, "target": 120 }, "testing-library/no-container": { "max": 133, "target": 50 }, "testing-library/no-node-access": { "max": 707, "target": 500 }, "testing-library/prefer-screen-queries": { "max": 18, "target": 18 } diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index d7475173b90..76ac60a6453 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -1619,14 +1619,6 @@ "count": 1 } }, - "src/components/common_components/fetch_teams.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "max-params": { - "count": 1 - } - }, "src/components/common_components/simple_table.tsx": { "local/filename-pascal-case": { "count": 1 @@ -1823,7 +1815,7 @@ "count": 5 }, "no-restricted-syntax": { - "count": 152 + "count": 150 }, "prefer-const": { "count": 32 @@ -1871,9 +1863,6 @@ "src/components/per_user_usage.tsx": { "local/filename-pascal-case": { "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 } }, "src/components/permissions/MCPServerPermissions.tsx": { @@ -2303,17 +2292,6 @@ "count": 1 } }, - "src/components/user_dashboard.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "prefer-const": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/components/vector_store_management/types.tsx": { "local/filename-pascal-case": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.test.tsx index 13689afcb52..2e3b000dfec 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.test.tsx @@ -1,59 +1,90 @@ -import { render } from "@testing-library/react"; -import { describe, it, expect, vi } from "vitest"; +import { render, screen } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; -const { userDashboardSpy } = vi.hoisted(() => ({ - userDashboardSpy: vi.fn((_props: Record) => null), +const { teamListCall, authorizedSession } = vi.hoisted(() => ({ + teamListCall: vi.fn(() => new Promise(() => {})), + authorizedSession: vi.fn(), })); -vi.mock("@/components/user_dashboard", () => ({ - default: (props: Record) => userDashboardSpy(props), -})); +const session = (overrides: { userRole?: string; isViewOnly?: boolean } = {}) => ({ + isLoading: false, + isAuthorized: true, + token: "jwt", + accessToken: "sk-access", + userId: "u-123", + userEmail: "admin@example.com", + userRole: "Admin", + isViewOnly: false, + premiumUser: false, + disabledPersonalKeyCreation: false, + showSSOBanner: false, + ...overrides, +}); -// AuthContext is still hydrating: userID has not been populated yet (the regression). -vi.mock("@/contexts/AuthContext", () => ({ - useAuth: () => ({ - userID: null, - userRole: "", - userEmail: null, - accessToken: null, - premiumUser: false, - setUserRole: vi.fn(), - setUserEmail: vi.fn(), - }), -})); - -// useAuthorized decodes the cookie synchronously, so identity is already available. vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ - isLoading: false, - isAuthorized: true, - token: "jwt", - accessToken: "sk-access", - userId: "u-123", - userEmail: "admin@example.com", - userRole: "Admin", - premiumUser: false, - disabledPersonalKeyCreation: false, - showSSOBanner: false, - }), + default: () => authorizedSession(), })); vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ - teamListCall: vi.fn(() => new Promise(() => {})), + teamListCall, })); vi.mock("next/navigation", () => ({ useSearchParams: () => new URLSearchParams(""), })); +vi.mock("@/components/VirtualKeysPage/VirtualKeysTable", () => ({ + VirtualKeysTable: ({ headerActions }: { headerActions?: React.ReactNode }) => ( +
+ {headerActions} + + + ), +})); + +vi.mock("@/components/organisms/create_key_button", () => ({ + default: () => , +})); + import ApiKeysDashboard from "./ApiKeysDashboard"; -describe("ApiKeysDashboard identity source", () => { - it("passes the useAuthorized userID through even while AuthContext.userID is still null", () => { +describe("ApiKeysDashboard", () => { + beforeEach(() => { + teamListCall.mockClear(); + authorizedSession.mockReturnValue(session()); + sessionStorage.clear(); + }); + + it("scopes the team list to the signed-in user for non-admin roles", () => { + authorizedSession.mockReturnValue(session({ userRole: "Internal User" })); render(); - expect(userDashboardSpy).toHaveBeenCalled(); - const props = userDashboardSpy.mock.calls[0][0]; - expect(props.userID).toBe("u-123"); + expect(teamListCall).toHaveBeenCalledWith("sk-access", 1, 100, { userID: "u-123" }); + }); + + it("renders the keys table with a Create Key action for roles that can write", () => { + render(); + + expect(screen.getByRole("table", { name: "Virtual Keys" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Create Key" })).toBeInTheDocument(); + }); + + it("hides Create Key for view-only roles", () => { + authorizedSession.mockReturnValue(session({ isViewOnly: true })); + render(); + + expect(screen.getByRole("table", { name: "Virtual Keys" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Create Key" })).not.toBeInTheDocument(); + }); + + it("leaves other pages' session state intact when the tab reloads", () => { + sessionStorage.setItem("chatHistory", '[{"role":"user","content":"hi"}]'); + sessionStorage.setItem("selectedModel", "gpt-5.5"); + render(); + + window.dispatchEvent(new Event("beforeunload")); + + expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]'); + expect(sessionStorage.getItem("selectedModel")).toBe("gpt-5.5"); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.tsx index ae0c443910a..376fee72b88 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.tsx @@ -3,22 +3,17 @@ import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { KeyResponse, Team } from "@/components/key_team_helpers/key_list"; -import { CreateKeyPrefillData } from "@/components/organisms/create_key_button"; -import UserDashboard from "@/components/user_dashboard"; -import { useAuth } from "@/contexts/AuthContext"; +import CreateKey, { CreateKeyPrefillData } from "@/components/organisms/create_key_button"; +import { VirtualKeysTable } from "@/components/VirtualKeysPage/VirtualKeysTable"; import { useSearchParams } from "next/navigation"; import { useEffect, useMemo, useState } from "react"; export default function ApiKeysDashboard() { - // Identity comes from useAuthorized (synchronous cookie decode) so userID is set whenever the - // route is authorized; useAuth only supplies the backfill setters UserDashboard still expects. - const { userId: userID, userRole, userEmail, accessToken, premiumUser } = useAuthorized(); - const { setUserRole, setUserEmail } = useAuth(); + const { userId: userID, userRole, accessToken, isViewOnly } = useAuthorized(); const searchParams = useSearchParams()!; const [teams, setTeams] = useState(null); const [keys, setKeys] = useState([]); - const [createClicked, setCreateClicked] = useState(false); const autoOpenCreate = searchParams.get("create") === "true"; const prefillData: CreateKeyPrefillData | undefined = useMemo(() => { @@ -63,7 +58,6 @@ export default function ApiKeysDashboard() { const addKey = (data: KeyResponse) => { setKeys((prevData) => (prevData ? [...prevData, data] : [data])); - setCreateClicked((prev) => !prev); }; useEffect(() => { @@ -77,21 +71,21 @@ export default function ApiKeysDashboard() { }, [accessToken, userID, userRole]); return ( - +
+ + ) + } + /> +
); } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts index 82cefd800f4..7925af223ae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts @@ -8,6 +8,8 @@ export interface ProxySettings { PROXY_BASE_URL: string; PROXY_LOGOUT_URL: string; LITELLM_UI_API_DOC_BASE_URL?: string | null; + DISABLE_EXPENSIVE_DB_QUERIES?: boolean; + NUM_SPEND_LOGS_ROWS?: number; } const EMPTY_PROXY_SETTINGS: ProxySettings = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx index 42be1b34f07..889a17bc88d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx @@ -1,7 +1,7 @@ import React, { useState, useEffect } from "react"; import ViewUserSpend from "@/components/view_user_spend"; -import { ProxySettings } from "@/components/user_dashboard"; +import { ProxySettings } from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { Button } from "@/components/ui/button"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; diff --git a/ui/litellm-dashboard/src/components/common_components/fetch_teams.tsx b/ui/litellm-dashboard/src/components/common_components/fetch_teams.tsx deleted file mode 100644 index ca82fdfb144..00000000000 --- a/ui/litellm-dashboard/src/components/common_components/fetch_teams.tsx +++ /dev/null @@ -1,18 +0,0 @@ -import { teamListCall, Organization } from "../networking"; - -export const fetchTeams = async ( - accessToken: string, - userID: string | null, - userRole: string | null, - currentOrg: Organization | null, - setTeams: (teams: any[]) => void, -) => { - let givenTeams; - if (userRole != "Admin" && userRole != "Admin Viewer") { - givenTeams = await teamListCall(accessToken, currentOrg?.organization_id || null, userID); - } else { - givenTeams = await teamListCall(accessToken, currentOrg?.organization_id || null); - } - - setTeams(givenTeams); -}; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 144c02dfcd1..a7597c30404 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -140,7 +140,7 @@ const resolveDefaultBase = (fallback: string | null): string | null => const defaultProxyBaseUrl = resolveDefaultBase(null); const WORKER_URL_KEY = "litellm_worker_url"; // If a worker URL is in localStorage, use it as the initial proxyBaseUrl. -// This survives page navigation and the sessionStorage.clear() in user_dashboard. +// This survives page navigation. const _rawWorkerUrl = typeof window !== "undefined" ? window.localStorage.getItem(WORKER_URL_KEY) : null; // Validate stored worker URL — reject non-HTTP schemes to prevent exfiltration const _initialWorkerUrl = (() => { @@ -195,10 +195,9 @@ export const getProxyBaseUrl = (): string => { /** * Switch API calls to point at a worker (or back to the control plane). - * Persists to localStorage so it survives page navigation and the - * sessionStorage.clear() in user_dashboard. Also updates the module-level - * proxyBaseUrl so in-flight code in this JS execution sees the new value - * immediately. + * Persists to localStorage so it survives page navigation. Also updates the + * module-level proxyBaseUrl so in-flight code in this JS execution sees the + * new value immediately. */ function isValidHttpUrl(url: string): boolean { try { diff --git a/ui/litellm-dashboard/src/components/user_dashboard.test.tsx b/ui/litellm-dashboard/src/components/user_dashboard.test.tsx deleted file mode 100644 index 01764613bf5..00000000000 --- a/ui/litellm-dashboard/src/components/user_dashboard.test.tsx +++ /dev/null @@ -1,133 +0,0 @@ -import { vi, describe, it, expect, beforeEach, afterEach } from "vitest"; -import { cleanup } from "@testing-library/react"; -import React from "react"; -import { renderWithProviders } from "../../tests/test-utils"; - -// Track addEventListener/removeEventListener calls for "beforeunload" -const addEventListenerSpy = vi.spyOn(window, "addEventListener"); -const removeEventListenerSpy = vi.spyOn(window, "removeEventListener"); - -// Mock next/navigation -vi.mock("next/navigation", () => ({ - useSearchParams: () => new URLSearchParams(), -})); - -// Mock networking with importOriginal so all exports are available -vi.mock("./networking", async (importOriginal) => { - const actual = await importOriginal(); - return { - ...actual, - getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"), - getProxyUISettings: vi.fn().mockResolvedValue({}), - keyInfoCall: vi.fn().mockResolvedValue({}), - modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), - userGetInfoV2: vi.fn().mockResolvedValue({ - user_id: "user-1", - user_email: "test@example.com", - spend: 0, - max_budget: null, - models: [], - teams: [], - }), - }; -}); - -// Mock jwt-decode to return a valid token structure -vi.mock("jwt-decode", () => ({ - jwtDecode: vi.fn().mockReturnValue({ - key: "test-access-token", - user_role: "proxy_admin", - user_email: "test@example.com", - exp: Math.floor(Date.now() / 1000) + 3600, - }), -})); - -// Mock cookie utility -vi.mock("@/utils/cookieUtils", () => ({ - clearTokenCookies: vi.fn(), - getCookie: vi.fn().mockReturnValue("fake-jwt-token"), -})); - -// Mock fetchTeams -vi.mock("./common_components/fetch_teams", () => ({ - fetchTeams: vi.fn(), -})); - -// Mock heavy child components to isolate UserDashboard behavior -vi.mock("./organisms/create_key_button", () => ({ - default: () =>
, -})); - -vi.mock("./VirtualKeysPage/VirtualKeysTable", () => ({ - VirtualKeysTable: () =>
, -})); - -vi.mock("../app/onboarding/page", () => ({ - default: () =>
, -})); - -// Provide a token cookie so the component doesn't redirect to login -Object.defineProperty(document, "cookie", { - writable: true, - value: "token=fake-jwt-token", -}); - -import UserDashboard from "./user_dashboard"; - -const defaultProps = { - userID: "user-1", - userRole: "Admin", - userEmail: "test@example.com", - teams: [] as any[], - keys: [] as any[], - setUserRole: vi.fn(), - setUserEmail: vi.fn(), - setTeams: vi.fn(), - setKeys: vi.fn(), - premiumUser: false, - addKey: vi.fn(), - createClicked: false, -}; - -function renderDashboard(props = {}) { - return renderWithProviders(); -} - -describe("UserDashboard beforeunload listener", () => { - beforeEach(() => { - addEventListenerSpy.mockClear(); - removeEventListenerSpy.mockClear(); - }); - - afterEach(() => { - cleanup(); - }); - - it("registers exactly one beforeunload listener on mount", () => { - renderDashboard(); - - const beforeUnloadCalls = addEventListenerSpy.mock.calls.filter(([event]) => event === "beforeunload"); - expect(beforeUnloadCalls).toHaveLength(1); - }); - - it("does not add duplicate listeners on re-render", () => { - const { rerender } = renderWithProviders(); - - addEventListenerSpy.mockClear(); - - // Re-render with different props to trigger a render cycle - rerender(); - - const beforeUnloadCalls = addEventListenerSpy.mock.calls.filter(([event]) => event === "beforeunload"); - expect(beforeUnloadCalls).toHaveLength(0); - }); - - it("removes the beforeunload listener on unmount", () => { - const { unmount } = renderDashboard(); - - unmount(); - - const removeCalls = removeEventListenerSpy.mock.calls.filter(([event]) => event === "beforeunload"); - expect(removeCalls).toHaveLength(1); - }); -}); diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx deleted file mode 100644 index dcadc141103..00000000000 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ /dev/null @@ -1,239 +0,0 @@ -"use client"; -import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; -import { jwtDecode } from "jwt-decode"; -import React, { useEffect, useState } from "react"; -import { fetchTeams } from "./common_components/fetch_teams"; -import { KeyResponse, Team } from "./key_team_helpers/key_list"; -import { effectiveSessionRole } from "@/utils/roles"; -import { getProxyBaseUrl, keyInfoCall, modelAvailableCall, Organization, userGetInfoV2 } from "./networking"; -import CreateKey, { CreateKeyPrefillData } from "./organisms/create_key_button"; -import { VirtualKeysTable } from "./VirtualKeysPage/VirtualKeysTable"; - -export interface ProxySettings { - PROXY_BASE_URL: string | null; - PROXY_LOGOUT_URL: string | null; - LITELLM_UI_API_DOC_BASE_URL?: string | null; - DEFAULT_TEAM_DISABLED: boolean; - SSO_ENABLED: boolean; - DISABLE_EXPENSIVE_DB_QUERIES: boolean; - NUM_SPEND_LOGS_ROWS: number; -} - -export type UserInfo = { - models: string[]; - max_budget?: number | null; - spend: number; -}; - -interface UserDashboardProps { - userID: string | null; - userRole: string | null; - userEmail: string | null; - teams: Team[] | null; - keys: any[] | null; - setUserRole: React.Dispatch>; - setUserEmail: React.Dispatch>; - setTeams: React.Dispatch>; - setKeys: (keys: KeyResponse[]) => void; - premiumUser: boolean; - addKey: (data: any) => void; - createClicked: boolean; - autoOpenCreate?: boolean; - prefillData?: CreateKeyPrefillData; -} - -const UserDashboard: React.FC = ({ - userID, - userRole, - teams, - keys, - setUserRole, - userEmail, - setUserEmail, - setTeams, - setKeys, - premiumUser, - addKey, - createClicked, - autoOpenCreate, - prefillData, -}) => { - const [userSpendData, setUserSpendData] = useState(null); - const [currentOrg] = useState(null); - - const token = getCookie("token"); - - const [accessToken, setAccessToken] = useState(null); - const [selectedTeam] = useState(null); - - // Clear session storage on page unload so next load fetches fresh data. - // Note: MCP auth tokens are persistent and should not be cleared on page refresh - // They are only cleared on logout - useEffect(() => { - const handleBeforeUnload = () => { - const token = sessionStorage.getItem("token"); - sessionStorage.clear(); - if (token) { - sessionStorage.setItem("token", token); - } - }; - window.addEventListener("beforeunload", handleBeforeUnload); - return () => window.removeEventListener("beforeunload", handleBeforeUnload); - }, []); - - // console.log(`selectedTeam: ${Object.entries(selectedTeam)}`); - // Moved useEffect inside the component and used a condition to run fetch only if the params are available - useEffect(() => { - if (token) { - const decoded = jwtDecode(token) as { [key: string]: any }; - if (decoded) { - // cast decoded to dictionary - - // set accessToken - setAccessToken(decoded.key); - - // check if userRole is defined - if (decoded.user_role) { - setUserRole(effectiveSessionRole(decoded.user_role)); - } else { - } - - if (decoded.user_email) { - setUserEmail(decoded.user_email); - } else { - } - } - } - if (userID && accessToken && userRole && !userSpendData) { - const cachedUserModels = sessionStorage.getItem("userModels" + userID); - if (!cachedUserModels) { - const fetchData = async () => { - try { - const response = await userGetInfoV2(accessToken, userID); - - setUserSpendData(response); - - sessionStorage.setItem("userSpendData" + userID, JSON.stringify(response)); - - const model_available = await modelAvailableCall(accessToken, userID, userRole); - // loop through model_info["data"] and create an array of element.model_name - let available_model_names = model_available["data"].map((element: { id: string }) => element.id); - - sessionStorage.setItem("userModels" + userID, JSON.stringify(available_model_names)); - } catch (error: any) { - console.error("There was an error fetching the data", error); - if (error.message.includes("Invalid proxy server token passed")) { - gotoLogin(); - } - // Optionally, update your UI to reflect the error state here as well - } - }; - fetchData(); - fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); - } - } - }, [userID, token, accessToken, userRole]); - - useEffect(() => { - // check key health - if it's invalid, redirect to login - if (accessToken) { - const fetchKeyInfo = async () => { - try { - await keyInfoCall(accessToken, [accessToken]); - } catch (error: any) { - if (error.message.includes("Invalid proxy server token passed")) { - gotoLogin(); - } - } - }; - fetchKeyInfo(); - } - }, [accessToken]); - - useEffect(() => { - if (accessToken) { - fetchTeams(accessToken, userID, userRole, currentOrg, setTeams); - } - }, [currentOrg]); - - function gotoLogin() { - // Clear token cookies using the utility function - clearTokenCookies(); - - const baseUrl = getProxyBaseUrl(); - - const url = baseUrl ? `${baseUrl}/sso/key/generate` : `/sso/key/generate`; - - window.location.href = url; - - return null; - } - - if (token == null) { - // user is not logged in as yet - - // Clear token cookies using the utility function - gotoLogin(); - return null; - } else { - // Check if token is expired - try { - const decoded = jwtDecode(token) as { [key: string]: any }; - const expTime = decoded.exp; - const currentTime = Math.floor(Date.now() / 1000); - - if (expTime && currentTime >= expTime) { - gotoLogin(); - - return null; - } - } catch (error) { - console.error("Error decoding token:", error); - // If there's an error decoding the token, consider it invalid - clearTokenCookies(); - - gotoLogin(); - - return null; - } - - if (accessToken == null) { - return null; - } - } - - if (userID == null) { - return

User ID is not set

; - } - - if (userRole == null) { - setUserRole("App Owner"); - } - - // Admin Viewer can view keys read-only — gate "Create Key" but render the - // virtual-keys table the same as for Proxy Admin (read parity). Every - // other role keeps its existing ability to create keys. - const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer"; - - return ( -
- - ) : undefined - } - /> -
- ); -}; - -export default UserDashboard; From 8bfaaffba556b3f231479cc49ea1e9647f9829e6 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 5 Sep 2026 17:57:38 -0700 Subject: [PATCH 34/63] test(ui): stub the Virtual Keys dashboard in the expired-token page test --- .../tests/CreateKeyPage.expiredToken.test.tsx | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index 211a754dad4..6a4fa2e32c3 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -105,7 +105,7 @@ vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { // Super-light stubs for all heavy components so rendering doesn't explode vi.mock("@/components/navbar", () => ({ default: stub("navbar") })); -vi.mock("@/components/user_dashboard", () => ({ default: stub("user-dashboard") })); +vi.mock("@/app/(dashboard)/api-keys/ApiKeysDashboard", () => ({ default: stub("api-keys-dashboard") })); vi.mock("@/components/templates/model_dashboard", () => ({ default: stub("model-dashboard") })); vi.mock("@/components/teams", () => ({ default: stub("teams") })); vi.mock("@/app/(dashboard)/organizations/_components/organizations", () => ({ @@ -135,7 +135,6 @@ vi.mock("@/app/(dashboard)/tag-management/_components", () => ({ default: stub(" vi.mock("@/app/(dashboard)/vector-stores/_components", () => ({ default: stub("vector-stores") })); vi.mock("@/components/ui_theme_settings", () => ({ default: stub("ui-theme-settings") })); vi.mock("@/components/organisms/create_key_button", () => ({ fetchUserModels: vi.fn() })); -vi.mock("@/components/common_components/fetch_teams", () => ({ fetchTeams: vi.fn() })); vi.mock("@/components/ui/ui-loading-spinner", () => ({ UiLoadingSpinner: stub("spinner"), })); @@ -270,9 +269,9 @@ describe("CreateKeyPage auth behavior", () => { expect(window.location.replace).not.toHaveBeenCalled(); }); - // And the default page content appears (UserDashboard stub; chrome now lives in the layout) + // And the default page content appears (ApiKeysDashboard stub; chrome now lives in the layout) await waitFor(() => { - expect(screen.getByTestId("user-dashboard")).toBeInTheDocument(); + expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument(); }); }); From 60440ee4d3e309a780c88ebbf550603e29e71cef Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 18:04:35 -0700 Subject: [PATCH 35/63] feat(mcp): add opt-in per-server oauth relay discovery (#39936) Resolves LIT-7074 --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/mcp_server.py | 1 + .../mcp_server/auth/user_api_key_auth_mcp.py | 2 +- .../mcp_server/discoverable_endpoints.py | 20 ++++-- .../mcp_server/mcp_server_manager.py | 34 +++++++++ litellm/proxy/_lazy_openapi_snapshot.json | 35 +++++++++ litellm/proxy/_types.py | 43 +++++++++++ .../mcp_management_endpoints.py | 22 ++++++ litellm/proxy/schema.prisma | 1 + .../types/mcp_server/mcp_server_manager.py | 11 +++ schema.prisma | 1 + .../auth/test_user_api_key_auth_mcp.py | 1 + .../mcp_server/test_db_credentials.py | 72 +++++++++++++++++++ .../mcp_server/test_discoverable_endpoints.py | 12 +++- .../mcp_server/test_mcp_server_manager.py | 44 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 ++++ 17 files changed, 309 insertions(+), 8 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql new file mode 100644 index 00000000000..2fa1234bbfc --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 1c43668f227..06ac177cca4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable { delegate_auth_to_upstream Boolean @default(false) oauth_passthrough Boolean @default(false) dcr_bridge Boolean? + per_server_oauth_discovery Boolean @default(false) is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index 6cc4a765e46..6bf21a19896 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -98,6 +98,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): delegate_auth_to_upstream: bool = False oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = Field(default_factory=list) byok_api_key_help_url: str | None = None diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index b9bfb062ec7..ad9622c9cd0 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -197,7 +197,7 @@ def _gateway_dcr_challenge_target( if targets is None: return None server: Final = global_mcp_server_manager.get_mcp_server_by_name(targets[0], client_ip=client_ip) - if server is None or not server.is_gateway_managed_oauth2: + if server is None or not server.advertises_gateway_authorization_server: return None return targets[0] diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3da5950ce7f..e10bfd41ed6 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1481,11 +1481,19 @@ async def _persist_dcr_client_registration( ) updated_row: Final = await update_mcp_server( prisma_client=prisma_client, - data=UpdateMCPServerRequest( - server_id=mcp_server.server_id, - credentials=credentials, - oauth2_flow="authorization_code", - **({"token_url": mcp_server.token_url} if mcp_server.token_url else {}), + data=( + UpdateMCPServerRequest( + server_id=mcp_server.server_id, + credentials=credentials, + oauth2_flow="authorization_code", + token_url=mcp_server.token_url, + ) + if mcp_server.token_url + else UpdateMCPServerRequest( + server_id=mcp_server.server_id, + credentials=credentials, + oauth2_flow="authorization_code", + ) ), touched_by="mcp_oauth_dcr", ) @@ -2367,7 +2375,7 @@ async def _build_oauth_protected_resource_response( if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource") - if explicitly_named and mcp_server is not None and mcp_server.is_gateway_managed_oauth2: + if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], "resource": resource_url, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index dc1e8db1628..612596bc803 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -155,6 +155,7 @@ from litellm.proxy._types import ( MCPTransportType, SpecialMCPServerNames, UserAPIKeyAuth, + is_per_server_oauth_discovery_eligible, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper @@ -344,6 +345,7 @@ class MCPServerConfig(TypedDict, total=False): token_endpoint_auth_method: MCPTokenEndpointAuthMethod scopes: str | Sequence[str] dcr_bridge: object + per_server_oauth_discovery: ReadOnly[object] extra_headers: _StringList allowed_tools: _StringList disallowed_tools: _StringList @@ -414,6 +416,31 @@ def _blank_to_none(value: str | None) -> str | None: return value.strip() or None +def _config_per_server_oauth_discovery( + server_config: MCPServerConfig, + server_ref: str, + auth_type: MCPAuthType | None, + oauth2_flow: object, +) -> bool: + match server_config.get("per_server_oauth_discovery", False): + case bool() as enabled: + pass + case other: + raise ValueError( + f"Invalid config for MCP server '{server_ref}': per_server_oauth_discovery must be a boolean " + f"(got {other!r})." + ) + relay_eligible: Final = is_per_server_oauth_discovery_eligible( + auth_type, oauth2_flow, server_config.get("delegate_auth_to_upstream", False) + ) + if enabled and not relay_eligible: + raise ValueError( + f"Invalid config for MCP server '{server_ref}': per_server_oauth_discovery is only supported for " + "auth_type oauth2 with oauth2_flow authorization_code and without delegate_auth_to_upstream." + ) + return enabled + + def _pinned_config_server_id(raw_server_id: object, server_name: str) -> str | None: """Return the ``server_id`` an admin pinned for this config.yaml server, or ``None`` when absent. @@ -2307,6 +2334,9 @@ class MCPServerManager: ) config_dcr_bridge = server_config.get("dcr_bridge", None) + config_per_server_oauth_discovery = _config_per_server_oauth_discovery( + server_config, server_name or server_id, auth_type, config_oauth2_flow + ) if config_dcr_bridge is not None and not isinstance(config_dcr_bridge, bool): raise ValueError( f"Invalid config for MCP server '{server_name or server_id}': dcr_bridge " @@ -2378,6 +2408,7 @@ class MCPServerManager: delegate_auth_to_upstream=bool(server_config.get("delegate_auth_to_upstream", False)), oauth_passthrough=bool(server_config.get("oauth_passthrough", False)), dcr_bridge=config_dcr_bridge, + per_server_oauth_discovery=config_per_server_oauth_discovery, # AWS SigV4 fields aws_access_key_id=server_config.get("aws_access_key_id", None), aws_secret_access_key=server_config.get("aws_secret_access_key", None), @@ -2903,6 +2934,7 @@ class MCPServerManager: delegate_auth_to_upstream=bool(getattr(mcp_server, "delegate_auth_to_upstream", False)), oauth_passthrough=bool(getattr(mcp_server, "oauth_passthrough", False)), dcr_bridge=getattr(mcp_server, "dcr_bridge", None), + per_server_oauth_discovery=bool(getattr(mcp_server, "per_server_oauth_discovery", False)), created_at=getattr(mcp_server, "created_at", None), updated_at=getattr(mcp_server, "updated_at", None), tool_name_to_display_name=_deserialize_json_dict(getattr(mcp_server, "tool_name_to_display_name", None)), @@ -6692,6 +6724,7 @@ class MCPServerManager: registration_url=server.configured_registration_url or server.registration_url, oauth2_flow=server.oauth2_flow, dcr_bridge=server.dcr_bridge, + per_server_oauth_discovery=server.per_server_oauth_discovery, token_exchange_endpoint=server.token_exchange_endpoint, audience=server.audience, subject_token_type=server.subject_token_type, @@ -6810,6 +6843,7 @@ class MCPServerManager: delegate_auth_to_upstream=server.delegate_auth_to_upstream, oauth_passthrough=getattr(server, "oauth_passthrough", False), dcr_bridge=server.dcr_bridge, + per_server_oauth_discovery=server.per_server_oauth_discovery, is_byok=server.is_byok, byok_description=server.byok_description, byok_api_key_help_url=server.byok_api_key_help_url, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index ea757edb215..91a97ad6544 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -16355,6 +16355,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -17978,6 +17983,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -18859,6 +18869,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -20868,6 +20883,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -22342,6 +22362,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -22862,6 +22887,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -25371,6 +25401,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f83011835fd..c28ac8848ba 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1379,6 +1379,35 @@ def _dcr_bridge_auth_type_error(auth_type: object) -> ValueError: ) +def _per_server_oauth_discovery_error() -> ValueError: + return ValueError( + "per_server_oauth_discovery is only supported for auth_type oauth2 with oauth2_flow " + "authorization_code and without delegate_auth_to_upstream." + ) + + +def is_per_server_oauth_discovery_eligible( + auth_type: object, oauth2_flow: object, delegate_auth_to_upstream: object +) -> bool: + return auth_type == MCPAuth.oauth2 and oauth2_flow == "authorization_code" and not delegate_auth_to_upstream + + +def _reject_unsupported_per_server_oauth_discovery(values: object, require_auth_type: bool) -> None: + """Partial updates may omit eligibility fields; those are checked against the stored row by the + update endpoint. Every field the payload does carry must be eligible on its own.""" + if not isinstance(values, dict) or not values.get("per_server_oauth_discovery"): + return + auth_type_ok: Final = values.get("auth_type") == MCPAuth.oauth2 or ( + not require_auth_type and "auth_type" not in values + ) + oauth2_flow_ok: Final = values.get("oauth2_flow") == "authorization_code" or ( + not require_auth_type and "oauth2_flow" not in values + ) + if auth_type_ok and oauth2_flow_ok and not values.get("delegate_auth_to_upstream"): + return + raise _per_server_oauth_discovery_error() + + class NewMCPServerRequest(LiteLLMPydanticObjectBase): server_id: str | None = None server_name: str | None = None @@ -1420,6 +1449,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): delegate_auth_to_upstream: bool = False oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = Field(default_factory=list) byok_api_key_help_url: str | None = None @@ -1484,6 +1514,12 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): return values raise _dcr_bridge_auth_type_error(auth_type) + @model_validator(mode="before") + @classmethod + def validate_per_server_oauth_discovery_auth_type(cls, values: object) -> object: + _reject_unsupported_per_server_oauth_discovery(values, require_auth_type=True) + return values + class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): server_id: str @@ -1526,6 +1562,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): delegate_auth_to_upstream: bool = False oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = Field(default_factory=list) byok_api_key_help_url: str | None = None @@ -1570,6 +1607,12 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): return values raise _dcr_bridge_auth_type_error(auth_type) + @model_validator(mode="before") + @classmethod + def validate_per_server_oauth_discovery_auth_type(cls, values: object) -> object: + _reject_unsupported_per_server_oauth_discovery(values, require_auth_type=False) + return values + from litellm.models.mcp_server import ( # noqa: E402 LiteLLM_MCPServerTable as LiteLLM_MCPServerTable, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index ae266792391..d5c3427f29a 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -193,6 +193,7 @@ if MCP_AVAILABLE: UpdateMCPServerRequest, UserAPIKeyAuth, UserMCPManagementMode, + is_per_server_oauth_discovery_eligible, ) from litellm.proxy.auth.user_api_key_auth import ( _user_api_key_auth_builder, @@ -2714,6 +2715,27 @@ if MCP_AVAILABLE: old_server_record = None old_server_record_read_failed = True + if payload.per_server_oauth_discovery and (old_server_record is not None or old_server_record_read_failed): + relay_eligible: Final = old_server_record is not None and is_per_server_oauth_discovery_eligible( + payload.auth_type if "auth_type" in payload_fields_set else old_server_record.auth_type, + payload.oauth2_flow if "oauth2_flow" in payload_fields_set else old_server_record.oauth2_flow, + ( + payload.delegate_auth_to_upstream + if "delegate_auth_to_upstream" in payload_fields_set + else old_server_record.delegate_auth_to_upstream + ), + ) + if not relay_eligible: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict + "error": ( + "per_server_oauth_discovery is only supported for auth_type oauth2 with oauth2_flow " + "authorization_code and without delegate_auth_to_upstream." + ) + }, + ) + if ( payload.dcr_bridge and payload.auth_type is None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1c43668f227..06ac177cca4 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable { delegate_auth_to_upstream Boolean @default(false) oauth_passthrough Boolean @default(false) dcr_bridge Boolean? + per_server_oauth_discovery Boolean @default(false) is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 9bf3acc601c..84ffd50eea1 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -158,6 +158,7 @@ class MCPServer(BaseModel): # be set explicitly to avoid regressing servers that did not opt in. oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = [] byok_api_key_help_url: str | None = None @@ -241,6 +242,16 @@ class MCPServer(BaseModel): so they are excluded by construction.""" return self.auth_type == MCPAuth.oauth2 and not self.delegate_auth_to_upstream + @property + def uses_per_server_oauth_relay(self) -> bool: + """Whether named discovery should advertise the configured per-server OAuth relay.""" + return self.per_server_oauth_discovery and self.auth_type == MCPAuth.oauth2 and not self.has_client_credentials + + @property + def advertises_gateway_authorization_server(self) -> bool: + """Whether named discovery should advertise the aggregate gateway authorization server.""" + return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + @property def is_true_passthrough(self) -> bool: """True for the transparent-proxy mode: LiteLLM performs no admission auth and forwards the diff --git a/schema.prisma b/schema.prisma index 1c43668f227..06ac177cca4 100644 --- a/schema.prisma +++ b/schema.prisma @@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable { delegate_auth_to_upstream Boolean @default(false) oauth_passthrough Boolean @default(false) dcr_bridge Boolean? + per_server_oauth_discovery Boolean @default(false) is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index f1e299802fb..44eb7795659 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -7184,6 +7184,7 @@ class TestAggregateGatewayDcrChallenge: cases = [ (_server(MCPAuth.oauth2), "srv"), + (_server(MCPAuth.oauth2, per_server_oauth_discovery=True), None), (_server(MCPAuth.oauth2, oauth2_flow="client_credentials"), "srv"), (_server(MCPAuth.oauth2, delegate_auth_to_upstream=True), None), (_server(MCPAuth.oauth2_token_exchange), None), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 4d9142ad4c5..e20f6646310 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -1362,3 +1362,75 @@ async def test_refresh_user_oauth_token_uses_admin_entered_token_url_when_issuer assert result is not None assert captured["url"] == "https://idp.example.com/token" + + +def test_prepare_mcp_server_data_carries_per_server_oauth_discovery(): + request = NewMCPServerRequest( + server_name="relay_create", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + per_server_oauth_discovery=True, + ) + + data = _prepare_mcp_server_data(request) + + assert data["per_server_oauth_discovery"] is True + + +def test_prepare_mcp_server_data_update_carries_per_server_oauth_discovery(): + request = UpdateMCPServerRequest( + server_id="relay-update", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + per_server_oauth_discovery=True, + ) + + data = _prepare_mcp_server_data(request, exclude_unset=True) + + assert data["per_server_oauth_discovery"] is True + + +@pytest.mark.parametrize( + "request_cls, extra, overrides", + [ + (NewMCPServerRequest, {"server_name": "relay_create"}, {"auth_type": MCPAuth.oauth_delegate}), + (NewMCPServerRequest, {"server_name": "relay_create"}, {"oauth2_flow": "client_credentials"}), + (UpdateMCPServerRequest, {"server_id": "relay-update"}, {"delegate_auth_to_upstream": True}), + ], +) +def test_request_models_reject_unsupported_per_server_oauth_discovery(request_cls, extra, overrides): + payload = { + "url": "https://upstream.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + "per_server_oauth_discovery": True, + **extra, + **overrides, + } + + with pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"): + request_cls(**payload) + + +@pytest.mark.parametrize( + "partial_payload", + [ + {"oauth2_flow": "client_credentials"}, + {"delegate_auth_to_upstream": True}, + {"auth_type": MCPAuth.api_key}, + ], +) +def test_partial_update_rejects_ineligible_field_alongside_per_server_oauth_discovery(partial_payload): + with pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"): + UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True, **partial_payload) + + +def test_partial_update_defers_omitted_eligibility_fields_to_the_stored_row(): + request = UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True) + + assert request.per_server_oauth_discovery is True diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 775f6e5f3b8..a9bb24bbef9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -3342,12 +3342,13 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa mock_request.headers = {} interactive = _oauth2_server("github_mcp") + relay = _oauth2_server("relay_mcp", per_server_oauth_discovery=True) m2m = _oauth2_server("m2m_mcp", oauth2_flow="client_credentials", client_id="cid", client_secret="cs") delegated = _oauth2_server("delegated_mcp", delegate_auth_to_upstream=True) global_mcp_server_manager.registry.clear() try: - for server in (interactive, m2m, delegated): + for server in (interactive, relay, m2m, delegated): global_mcp_server_manager.registry[server.server_id] = server for name in ("github_mcp", "m2m_mcp"): @@ -3363,6 +3364,15 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa assert legacy["authorization_servers"] == ["https://litellm.example.com/mcp"], name assert legacy["resource"] == f"https://litellm.example.com/{name}/mcp" + relay_response = await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name="relay_mcp", use_standard_pattern=True + ) + assert relay_response["authorization_servers"] == ["https://litellm.example.com/relay_mcp"] + relay_legacy_response = await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name="relay_mcp", use_standard_pattern=False + ) + assert relay_legacy_response["authorization_servers"] == ["https://litellm.example.com/relay_mcp"] + delegated_response = await _build_oauth_protected_resource_response( request=mock_request, mcp_server_name="delegated_mcp", use_standard_pattern=True ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e3ed48713aa..764e2bb0e99 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1229,6 +1229,50 @@ class TestMCPServerManager: base.update(overrides) return {"bridgeserver": base} + @pytest.mark.asyncio + async def test_load_servers_from_config_accepts_per_server_oauth_discovery_for_oauth2(self): + manager = MCPServerManager() + + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)): + await manager.load_servers_from_config( + self._oauth2_config(oauth2_flow="authorization_code", per_server_oauth_discovery=True) + ) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.per_server_oauth_discovery is True + assert server.uses_per_server_oauth_relay is True + assert server.advertises_gateway_authorization_server is False + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "config", + [ + {"auth_type": MCPAuth.oauth_delegate}, + {"oauth2_flow": "client_credentials"}, + {"oauth2_flow": "authorization_code", "delegate_auth_to_upstream": True}, + ], + ) + async def test_load_servers_from_config_rejects_unsupported_per_server_oauth_discovery(self, config): + manager = MCPServerManager() + + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"), + ): + await manager.load_servers_from_config(self._oauth2_config(per_server_oauth_discovery=True, **config)) + + @pytest.mark.asyncio + async def test_load_servers_from_config_rejects_non_boolean_per_server_oauth_discovery(self): + manager = MCPServerManager() + + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + pytest.raises(ValueError, match="per_server_oauth_discovery.*must be a boolean"), + ): + await manager.load_servers_from_config( + self._oauth2_config(oauth2_flow="authorization_code", per_server_oauth_discovery="yes") + ) + @pytest.mark.asyncio async def test_load_servers_from_config_rejects_dcr_bridge_on_gateway_managed_auth_type(self): manager = MCPServerManager() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8a850598259..b1534c19670 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28870,6 +28870,11 @@ export interface components { * @default false */ oauth_passthrough: boolean; + /** + * Per Server Oauth Discovery + * @default false + */ + per_server_oauth_discovery: boolean; /** Registration Url */ registration_url?: string | null; /** Review Notes */ @@ -32010,6 +32015,11 @@ export interface components { * @default false */ oauth_passthrough: boolean; + /** + * Per Server Oauth Discovery + * @default false + */ + per_server_oauth_discovery: boolean; /** Registration Url */ registration_url?: string | null; /** Server Id */ @@ -37833,6 +37843,11 @@ export interface components { * @default false */ oauth_passthrough: boolean; + /** + * Per Server Oauth Discovery + * @default false + */ + per_server_oauth_discovery: boolean; /** Registration Url */ registration_url?: string | null; /** Server Id */ From 37e8b99a9f674ac7cfe2f146abb65f4708b1c354 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 18:03:24 -0700 Subject: [PATCH 36/63] fix(docker): ship pymongo in the proxy images for the MongoDB vector store The MongoDB Atlas vector store provider imports pymongo lazily from the opt-in `mongodb` extra, but none of the shipped images installed that extra. Any image-based deployment that configured a MongoDB vector store failed at search time with "requires the 'pymongo' package", which the user cannot fix without extending the image Adds `--extra mongodb` to every uv sync in the root Dockerfile, Dockerfile.database, Dockerfile.non_root, and the gateway component image. The backend component does not serve /vector_stores so it is left as is. The extra resolves from the existing uv.lock to pymongo 4.17.0 plus dnspython 2.8.0, no lock change needed (cherry picked from commit 16fd14f53705dab68ee8c4b0fd2c9093c8e3cdac) --- Dockerfile | 2 ++ docker/Dockerfile.database | 2 ++ docker/Dockerfile.non_root | 3 +++ gateway/Dockerfile | 2 ++ 4 files changed, 9 insertions(+) diff --git a/Dockerfile b/Dockerfile index 0a92aa9a68c..1648ec69d13 100644 --- a/Dockerfile +++ b/Dockerfile @@ -67,6 +67,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Copy full source tree @@ -89,6 +90,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index e9ad2849bb2..cc81ad6b3d3 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -65,6 +65,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Copy full source tree @@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index edf20e8bbff..358425af901 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -71,6 +71,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Copy full source tree @@ -99,6 +100,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 \ --no-sources-package litellm-proxy-extras; \ else \ @@ -109,6 +111,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13; \ fi diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 308d70a6b26..e42e488d57f 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -47,6 +47,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Stage 2 — copy source and install the project + workspace members. @@ -59,6 +60,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ From a9e918577bb9cc5c35bd4fcf90948aa4522ea059 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:46:51 -0700 Subject: [PATCH 37/63] ci(e2e): run the access_control canary on harness changes and name failed tests A harness-only change (proxy_client.py, conftest.py, pytest.ini, the gateway config, .github/e2e-stack, or the workflow) selected nothing, so the stack was never exercised by the change that touched it. select_tests.py keeps the changed-file rule and adds the access_control suite whenever a harness file changes. The run step now reports the pytest exit code before the evidence check, prints pytest's summary line per pass so the rerun count is visible, and assert_tests_ran.py names each failed or errored test as classname::name --- .github/e2e-stack/assert_tests_ran.py | 4 + .github/e2e-stack/select_tests.py | 31 +++++++ .github/workflows/test-e2e-changed.yml | 29 +++++-- .../test_e2e_changed_gate.py | 85 +++++++++++++++++++ tests/e2e/CONTRIBUTING.md | 6 +- 5 files changed, 143 insertions(+), 12 deletions(-) create mode 100644 .github/e2e-stack/select_tests.py diff --git a/.github/e2e-stack/assert_tests_ran.py b/.github/e2e-stack/assert_tests_ran.py index 7fd3f7c7c33..c4348c20873 100644 --- a/.github/e2e-stack/assert_tests_ran.py +++ b/.github/e2e-stack/assert_tests_ran.py @@ -20,6 +20,10 @@ def main() -> int: collected: Final = sum(case.get("file") == path for case in cases) skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases) _ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n") + for case in cases: + if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")): + continue + _ = sys.stdout.write(f" failed: {case.get('classname', '')}::{case.get('name', '')}\n") if ( selected and not missing diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py new file mode 100644 index 00000000000..cbe582b5b42 --- /dev/null +++ b/.github/e2e-stack/select_tests.py @@ -0,0 +1,31 @@ +import re +import sys +from typing import Final + +SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$") +OWN_LANE: Final = re.compile( + r"^tests/e2e/(ui|claude_code|load)/|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" +) +HARNESS: Final = re.compile( + r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" + r"|^tests/e2e/gateway/" + r"|^\.github/e2e-stack/" + r"|^\.github/workflows/test-e2e-changed\.yml$" +) + + +def select(changed: tuple[str, ...], canary: tuple[str, ...]) -> tuple[str, ...]: + direct: Final = frozenset(path for path in changed if SELECTABLE.match(path) and not OWN_LANE.match(path)) + harness_changed: Final = any(HARNESS.match(path) for path in changed) + canary_tests: Final = frozenset(path for path in canary if harness_changed and SELECTABLE.match(path)) + return tuple(sorted(direct | canary_tests)) + + +def main() -> int: + changed: Final = tuple(line.strip() for line in sys.stdin if line.strip()) + _ = sys.stdout.write(" ".join(select(changed, tuple(sys.argv[1:]))) + "\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.github/workflows/test-e2e-changed.yml b/.github/workflows/test-e2e-changed.yml index 16d84beab2d..23ab6dfcfe4 100644 --- a/.github/workflows/test-e2e-changed.yml +++ b/.github/workflows/test-e2e-changed.yml @@ -15,11 +15,21 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 permissions: + contents: read pull-requests: read outputs: tests: ${{ steps.changed.outputs.tests }} any: ${{ steps.changed.outputs.any }} steps: + - name: Checkout the selector and the canary suite + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + sparse-checkout: | + .github/e2e-stack + tests/e2e/access_control + persist-credentials: false + ref: ${{ github.sha }} + - name: List the e2e test files this PR added or modified id: changed env: @@ -27,7 +37,6 @@ jobs: REPO: ${{ github.repository }} PR_NUMBER: ${{ github.event.pull_request.number }} HEAD_SHA: ${{ github.event.pull_request.head.sha }} - OWN_LANE: '^tests/e2e/(ui|claude_code|load)/|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$' run: | gh api "repos/${REPO}/pulls/${PR_NUMBER}" \ --jq 'select(.head.sha == env.HEAD_SHA and .changed_files < 3000) | .head.sha' \ @@ -36,9 +45,7 @@ jobs: --jq '.[] | select(.status != "removed") | .filename')" gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}" tests="$(printf '%s\n' "${files}" \ - | grep -E '^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$' \ - | grep -vE "${OWN_LANE}" \ - | sort -u | tr '\n' ' ' | sed 's/ $//')" || true + | python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py)" echo "tests=${tests}" >> "${GITHUB_OUTPUT}" if [ -n "${tests}" ]; then echo "any=true" >> "${GITHUB_OUTPUT}" @@ -170,25 +177,29 @@ jobs: read -r -a test_files <<< "${TESTS}" for pass in 1 2 3; do report="${RUNNER_TEMP}/e2e-pass-${pass}.xml" + log="${RUNNER_TEMP}/e2e-pass-${pass}.log" echo "::group::pass ${pass} of 3" set +e uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \ - -o junit_family=xunit1 --junitxml="${report}" > "${RUNNER_TEMP}/e2e-pass-${pass}.log" 2>&1 + -o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1 status=$? + uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}" + verified=$? set -e + grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1 echo "::endgroup::" if [ "${status}" = "5" ]; then echo "::error::the selected files collected no runnable tests, so nothing was verified" exit 1 fi - if ! uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"; then - echo "::error::pass ${pass} of 3 did not verify every selected file" - exit 1 - fi if [ "${status}" != "0" ]; then echo "::error::pass ${pass} of 3 failed with exit code ${status}" exit "${status}" fi + if [ "${verified}" != "0" ]; then + echo "::error::pass ${pass} of 3 did not verify every selected file" + exit 1 + fi echo "pass ${pass} of 3 passed" done diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 2e95c33cf20..f7c4a5e2527 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -8,6 +8,8 @@ import pytest GATE: Final = Path(__file__).resolve().parents[2] / ".github/e2e-stack/assert_tests_ran.py" SECRETS_TO_ENV: Final = GATE.with_name("secrets_to_env.py") +SELECT_TESTS: Final = GATE.with_name("select_tests.py") +CANARY: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py") SELECTED: Final = ("tests/e2e/access_control/test_a.py", "tests/e2e/access_control/test_b.py") @@ -46,6 +48,29 @@ def test_passing_case_does_not_hide_a_failure_in_the_same_file(tmp_path: Path, o assert result.returncode == 1 +def test_failed_cases_are_named_per_selected_file(tmp_path: Path) -> None: + suite: Final = ET.Element("testsuite") + _ = ET.SubElement(suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_ok") + failed: Final = ET.SubElement( + suite, "testcase", file=SELECTED[0], classname="tests.e2e.access_control.test_a", name="test_boom" + ) + _ = ET.SubElement(failed, "failure", message="secret-bearing message") + errored: Final = ET.SubElement( + suite, "testcase", file=SELECTED[1], classname="tests.e2e.access_control.test_b", name="test_setup" + ) + _ = ET.SubElement(errored, "error") + report: Final = tmp_path / "report.xml" + ET.ElementTree(suite).write(report) + + result: Final = subprocess.run([sys.executable, str(GATE), str(report), *SELECTED], capture_output=True, text=True) + + assert result.returncode == 1 + assert " failed: tests.e2e.access_control.test_a::test_boom\n" in result.stdout + assert " failed: tests.e2e.access_control.test_b::test_setup\n" in result.stdout + assert "test_ok" not in result.stdout + assert "secret-bearing message" not in result.stdout + + @pytest.mark.parametrize("contents", ("", "')) def test_missing_execution_evidence_fails(tmp_path: Path, contents: str) -> None: report: Final = tmp_path / "report.xml" @@ -69,3 +94,63 @@ def test_short_values_are_written_without_masking_every_digit_in_the_log(tmp_pat assert result.returncode == 0, result.stderr assert result.stdout == "::add-mask::sk-0123456789abcdef\n" assert env_path.read_text() == "FLAG='1'\nAPI_KEY='sk-0123456789abcdef'\n" + + +def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: + result: Final = subprocess.run( + [sys.executable, str(SELECT_TESTS), *CANARY], + input="".join(f"{path}\n" for path in changed), + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + return tuple(result.stdout.split()) + + +@pytest.mark.parametrize( + ("changed", "expected"), + ( + (("tests/e2e/logging/test_datadog_e2e.py", "litellm/router.py"), ("tests/e2e/logging/test_datadog_e2e.py",)), + (("tests/e2e/ui/test_keys.py", "tests/e2e/claude_code/test_cli.py", "tests/e2e/load/test_burst.py"), ()), + (("tests/e2e/batches/test_managed_files_enforcement_e2e.py",), ()), + (("tests/e2e/logging/helpers.py", "docs/my-website/docs/index.md", "tests/e2e/CLAUDE.md"), ()), + ( + ("tests/e2e/logging/test_datadog_e2e.py", "tests/e2e/logging/test_datadog_e2e.py"), + ("tests/e2e/logging/test_datadog_e2e.py",), + ), + ), +) +def test_changed_suite_files_are_selected_outside_the_own_lane( + changed: tuple[str, ...], expected: tuple[str, ...] +) -> None: + assert select_tests(changed) == expected + + +@pytest.mark.parametrize( + "harness_file", + ( + "tests/e2e/proxy_client.py", + "tests/e2e/conftest.py", + "tests/e2e/pytest.ini", + "tests/e2e/gateway/stage_mirror_ci_config.yml", + ".github/e2e-stack/up.sh", + ".github/workflows/test-e2e-changed.yml", + ), +) +def test_harness_changes_run_the_canary_suite(harness_file: str) -> None: + assert select_tests((harness_file, "litellm/router.py")) == CANARY + + +def test_a_changed_canary_file_is_selected_once_alongside_a_harness_change() -> None: + assert select_tests((CANARY[1], "tests/e2e/proxy_client.py")) == CANARY + + +def test_the_canary_joins_directly_selected_files_in_sorted_order() -> None: + assert select_tests(("tests/e2e/logging/test_datadog_e2e.py", ".github/e2e-stack/up.sh")) == ( + *CANARY, + "tests/e2e/logging/test_datadog_e2e.py", + ) + + +def test_a_harness_unit_test_change_runs_itself_and_the_canary() -> None: + assert select_tests(("tests/e2e/test_proxy_client.py",)) == (*CANARY, "tests/e2e/test_proxy_client.py") diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 0b8859d9260..be282e2a77e 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -56,15 +56,15 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, harness, configuration, deleted-file, and workflow-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack -Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A failed pass stops the run. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch +Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch Repository admins must require the `e2e-changed-tests` status check for merging and configure the `e2e-changed` environment with required reviewers, self-review disabled, and admin bypass disabled. Each push cancels the previous run; a new run that selects tests needs a fresh approval. Reviewers must inspect the entire executable PR diff, including application code, dependencies, tests, and workflow helpers, before approving the exact revision. Approved code executes with provider credentials, so environment approval is a trust decision about that code Credentials come from the existing AWS Secrets Manager secrets in us-east-1, `litellm-e2e-changed-provider-keys` and `litellm-e2e-changed-license`. The OIDC role must trust only `repo:BerriAI/litellm:environment:e2e-changed` with audience `sts.amazonaws.com` and have read access only to these secrets. The short-lived reader credentials are scoped to the fetch step. Provider credentials must cover the selected suites, including Datadog credentials when logging or MCP tests need them; missing credentials fail the run. `up.sh` refuses to start without `DD_API_KEY`, because the stack's gateway config enables the Datadog callback for every run and a gateway booted without the key fails readiness. Keep provider credentials dedicated to this lane with only the permissions those tests need -Fetched values of eight characters or more are masked before use, while shorter values such as flags stay unmasked because masking a one-character value would blank every matching digit in the log, and credential files and raw output are private to the runner. Public logs contain selected file names, counts, and pass status; raw pytest output, reports, and stack logs are not uploaded or printed. The workflow removes them and the credential files during cleanup. To diagnose a failed pass, reproduce the selected files locally with the appropriate credentials and inspect the local logs +Fetched values of eight characters or more are masked before use, while shorter values such as flags stay unmasked because masking a one-character value would blank every matching digit in the log, and credential files and raw output are private to the runner. Public logs contain selected file names, counts, pytest's summary line, failed test ids, and pass status; raw pytest output, reports, and stack logs are not uploaded or printed. The workflow removes them and the credential files during cleanup. To diagnose a failed pass, reproduce the selected files locally with the appropriate credentials and inspect the local logs To reproduce the CI topology on a dedicated machine, `bash .github/e2e-stack/up.sh` reads `tests/e2e/.env`, writes `stack.env` under `${E2E_STACK_DIR:-/tmp/litellm-e2e-stack}`, and `bash .github/e2e-stack/down.sh` stops it. Keep this directory private and remove its credential files and logs after use From b2e93ba99f3e8ba8e87809d433d2bb9db5cd791a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:58:07 -0700 Subject: [PATCH 38/63] ci(e2e): declare the embedding model the access_control canary calls The first canary run failed pass 1 because the stage-mirror config had no openai-text-embedding-3-small while test_llm_api_routes_group_grants_every_llm_endpoint calls /embeddings with it; the public log named the test, which is the behavior the previous commit added --- tests/e2e/CONTRIBUTING.md | 2 +- tests/e2e/gateway/stage_mirror_ci_config.yml | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index be282e2a77e..106cfa31f0e 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -56,7 +56,7 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 6b3756b5f9f..229e8514dee 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -52,6 +52,10 @@ model_list: litellm_params: model: gemini/gemini-2.5-flash api_key: os.environ/GEMINI_API_KEY + - model_name: openai-text-embedding-3-small + litellm_params: + model: openai/text-embedding-3-small + api_key: os.environ/OPENAI_API_KEY mcp_servers: devin: From b067e836f8df15bb3dbefe345bf503b7d3811b9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:58:11 -0700 Subject: [PATCH 39/63] fix(batches): claim the batch cost spend row in the database before charging The cost callback used to look for an existing `_batch_cost` row before charging a completed batch, which left a window where concurrent retrieves on any instance all charged the key, and it would honor a row any request had written under that id. The spend update writer now inserts the batch cost row itself with `create_many(skip_duplicates=True)` and only the retrieve whose insert lands charges the key, team, and user. An existing row only takes the charge when it is a successful `aretrieve_batch` row, so a client-chosen `x-litellm-call-id` on another endpoint cannot suppress billing. Batch cost rows no longer get their own immediate flush path `batch_cost_is_final` now treats the proxy's normalized `complete` status like `completed`, which the enterprise batch cost poller relies on when it decides whether a completed batch is safe to retire. Tests build that status with `model_copy` since the OpenAI `Batch` model rejects it The `test-quality-ok` markers sit on the `patch(` lines the gate keys on, and the logging tests no longer wrap the priced retrieve in `contextlib.suppress` --- litellm/batches/batch_utils.py | 5 +- litellm/proxy/db/db_spend_update_writer.py | 73 ++++++++-- .../proxy/hooks/proxy_track_cost_callback.py | 68 +++------- litellm/proxy/utils.py | 10 +- .../test_litellm/batches/test_batch_utils.py | 14 +- .../test_litellm_logging.py | 12 +- .../proxy/db/test_db_spend_update_writer.py | 128 +++++++++++++++++- .../hooks/test_proxy_track_cost_callback.py | 116 ++++------------ .../prisma_and_spend/test_spend_functions.py | 15 -- 9 files changed, 249 insertions(+), 192 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index eaac3bf0e9f..959c7498479 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -25,7 +25,8 @@ class BatchCostUsageResult: failed_requests: int -_TERMINAL_BATCH_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) +_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) +_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) def batch_cost_is_final(batch: Batch) -> bool: @@ -39,7 +40,7 @@ def batch_cost_is_final(batch: Batch) -> bool: """ if batch.status not in _TERMINAL_BATCH_STATUSES: return False - if batch.status != "completed" or batch.output_file_id is not None: + if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: return True request_counts: Final = batch.request_counts return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ff48b00dc70..3fad351224b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -82,7 +82,10 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) -IMMEDIATE_FLUSH_CALL_TYPES: Final = RESPONSES_SESSION_CALL_TYPES | frozenset({CallTypes.aretrieve_batch.value}) + + +def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: + return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" class _SpendBatch(Protocol): @@ -216,7 +219,12 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ) -> None: + ) -> bool: + """Record the request's spend, answering whether its cost still needs charging. + + False only for a batch retrieve whose cost row another retrieve already wrote, + so the caller leaves the key, team, and user counters alone (LIT-7048). + """ from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -233,7 +241,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return + return True if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -264,10 +272,8 @@ class DBSpendUpdateWriter: payload["team_id"] = team_id if disable_spend_logs is False: - await self._insert_spend_log_to_db( - payload=payload, - prisma_client=prisma_client, - ) + if not await self._record_spend_log(payload=payload, prisma_client=prisma_client): + return False await self._enqueue_tool_usage_transaction( payload=payload, completion_response=completion_response, @@ -307,6 +313,7 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") + return True except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " @@ -319,7 +326,55 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) - return + return True + + async def _record_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None") -> bool: + if prisma_client is None or not _is_batch_cost_row(payload): + await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + return True + return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + + async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: + """Write the batch's cost row now, or learn that another retrieve already did. + + Every retrieve of one batch shares this row, so the insert that lands first owns + the charge and every later one finds the row and charges nothing (LIT-7048). Only + a row a successful retrieve wrote counts: a failed retrieve, or any request whose + client picked the batch id as its call id, cannot take the charge away. + """ + from litellm.repositories.table_repositories import SpendLogsRepository + + request_id: Final = payload["request_id"] + spend_logs: Final = SpendLogsRepository(prisma_client).table + try: + claimed: Final = await spend_logs.create_many( + data=[prisma_client.jsonify_object(payload)], # mutable-ok: prisma create_many takes a list + skip_duplicates=True, + ) + if claimed == 1: + return True + existing: Final = await spend_logs.find_unique( + where={"request_id": request_id} # mutable-ok: prisma where clause + ) + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log + verbose_proxy_logger.warning( + "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e + ) + await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + return True + if ( + existing is not None + and existing.call_type == CallTypes.aretrieve_batch.value + and existing.status == "success" + ): + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + return False + verbose_proxy_logger.warning( + "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + request_id, + getattr(existing, "call_type", None), + ) + return True async def _enqueue_tool_usage_transaction( self, @@ -940,7 +995,7 @@ class DBSpendUpdateWriter: from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush await enqueue_spend_logs(prisma_client, (payload,)) - if payload.get("call_type") in IMMEDIATE_FLUSH_CALL_TYPES: + if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: request_spend_log_flush() else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 95e61fdd98c..f0c889a8cb4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -34,7 +34,6 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, get_request_model_access_groups, - get_spend_logs_id, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -46,9 +45,7 @@ from litellm.types.utils import ( from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: - from prisma.types import LiteLLM_SpendLogsWhereUniqueInput - - from litellm.proxy.utils import PrismaClient, ProxyLogging + from litellm.proxy.utils import ProxyLogging _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { @@ -229,7 +226,6 @@ class _ProxyDBLogger(CustomLogger): ): from litellm.proxy.proxy_server import ( increment_spend_counters, - prisma_client, proxy_logging_obj, update_cache, ) @@ -257,15 +253,15 @@ class _ProxyDBLogger(CustomLogger): if ( isinstance(completion_response, LiteLLMBatch) and kwargs.get("call_type") == CallTypes.aretrieve_batch.value + and not batch_cost_is_final(completion_response) ): - batch_spend_log_id: Final = get_spend_logs_id( - CallTypes.aretrieve_batch.value, completion_response.model_dump(), kwargs + verbose_proxy_logger.debug( + "Cost tracking deferred for batch %s still in status %s", + completion_response.id, + completion_response.status, ) - if not await _batch_cost_is_trackable_now( - batch=completion_response, spend_log_id=batch_spend_log_id, prisma_client=prisma_client - ): - await _release_budget_reservation(budget_reservation=budget_reservation) - return + await _release_budget_reservation(budget_reservation=budget_reservation) + return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) @@ -307,7 +303,7 @@ class _ProxyDBLogger(CustomLogger): call_type=call_type, ): ## UPDATE DATABASE - await _update_database_and_spend_counters( + charged: Final = await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, increment_spend_counters=increment_spend_counters, user_api_key=user_api_key, @@ -324,6 +320,8 @@ class _ProxyDBLogger(CustomLogger): request_tags=tags, model_access_groups=model_access_groups, ) + if not charged: + return # update cache (fire-and-forget for backward compat: # cached object fields, soft budget alerts, etc.) @@ -509,42 +507,6 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value -async def _batch_cost_is_trackable_now( - batch: LiteLLMBatch, spend_log_id: str | None, prisma_client: "PrismaClient | None" -) -> bool: - """A batch is billed exactly once, from the first retrieve that sees it final. - - Every retrieve of one batch shares a single spend row (its id plus the batch cost - suffix), so a poll that lands before the output exists would write that row at $0 - and pin it there, and every retrieve after the first would add the cost to the - key, team, and user counters again. - """ - if not batch_cost_is_final(batch): - verbose_proxy_logger.debug("Cost tracking deferred for batch %s still in status %s", batch.id, batch.status) - return False - if prisma_client is None or spend_log_id is None: - return True - if not await _spend_log_already_recorded(prisma_client=prisma_client, request_id=spend_log_id): - return True - verbose_proxy_logger.debug( - "Cost tracking skipped for batch %s: spend row %s already recorded", batch.id, spend_log_id - ) - return False - - -async def _spend_log_already_recorded(prisma_client: "PrismaClient", request_id: str) -> bool: - from litellm.proxy.utils import spend_log_is_queued - - if await spend_log_is_queued(prisma_client, request_id): - return True - spend_log_row: Final[LiteLLM_SpendLogsWhereUniqueInput] = {"request_id": request_id} - try: - return await prisma_client.db.litellm_spendlogs.find_unique(where=spend_log_row) is not None - except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreadable DB must not drop the batch's only spend row - verbose_proxy_logger.warning("Could not check for an existing spend row %s, tracking anyway: %s", request_id, e) - return False - - def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse @@ -636,9 +598,9 @@ async def _update_database_and_spend_counters( budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, -) -> None: +) -> bool: try: - await proxy_logging_obj.db_spend_update_writer.update_database( + charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -663,6 +625,9 @@ async def _update_database_and_spend_counters( "Failed to invalidate budget reservation counters after release failed" ) raise + if not charged: + await _release_budget_reservation(budget_reservation=budget_reservation) + return False try: await increment_spend_counters( @@ -688,6 +653,7 @@ async def _update_database_and_spend_counters( finally: budget_reservation["finalized"] = True raise + return True async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f64b51bc6c6..accf7b720fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6251,9 +6251,7 @@ def request_spend_log_flush() -> None: The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. - A batch's cost row is what every other worker checks before charging the same batch - again, so it cannot wait either. Repeated requests coalesce into the monitor's next - pass, so the batching holds. + Repeated requests coalesce into the monitor's next pass, so the batching holds. """ PrismaClient.spend_log_flush_requested.set() @@ -6268,12 +6266,6 @@ async def _wait_for_spend_log_flush_request(interval: float) -> bool: return True -async def spend_log_is_queued(prisma_client: PrismaClient, request_id: str) -> bool: - """Whether a spend log with ``request_id`` is still waiting for the next flush.""" - async with prisma_client._spend_log_transactions_lock: - return any(row.get("request_id") == request_id for row in prisma_client.spend_log_transactions) - - async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 8d4f68164b4..976a96f2db1 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1735,10 +1735,10 @@ def _retrieved_batch( endpoint="/v1/chat/completions", input_file_id="file-in", object="batch", - status=status, + status="validating", output_file_id=output_file_id, request_counts=counts, - ) + ).model_copy(update={"status": status}) class TestBatchCostIsFinal: @@ -1750,8 +1750,9 @@ class TestBatchCostIsFinal: def test_in_flight_batch_is_not_final(self, status): assert bu.batch_cost_is_final(_retrieved_batch(status)) is False - def test_completed_with_output_is_final(self): - assert bu.batch_cost_is_final(_retrieved_batch("completed", output_file_id="file-out")) is True + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_with_output_is_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True def test_completed_without_output_and_unknown_counts_is_not_final(self): assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False @@ -1764,9 +1765,10 @@ class TestBatchCostIsFinal: counts = BatchRequestCounts(total=2, completed=2, failed=0) assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - def test_completed_without_output_and_every_line_failed_is_final(self): + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_without_output_and_every_line_failed_is_final(self, status): counts = BatchRequestCounts(total=2, completed=0, failed=2) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is True + assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) def test_other_terminal_statuses_are_final(self, status): 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 4583429bd11..174efaa4679 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -665,14 +665,14 @@ class TestRetrieveBatchPricesOnlyFinalBatches: endpoint="/v1/chat/completions", input_file_id="file-in", object="batch", - status=status, + status="validating", output_file_id=output_file_id, - ) + ).model_copy(update={"status": status}) @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "output_file_id"), - [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None)], + [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None), ("complete", None)], ) async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None: from litellm.litellm_core_utils import litellm_logging as logging_module @@ -681,8 +681,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches: monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) batch = self._batch(status, output_file_id) - with contextlib.suppress(Exception): - await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) handle_completed_batch.assert_not_awaited() assert "response_cost" not in batch._hidden_params @@ -705,8 +704,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches: monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) batch = self._batch("completed", "file-out") - with contextlib.suppress(Exception): - await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) handle_completed_batch.assert_awaited_once() assert batch._hidden_params["response_cost"] == 8e-06 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index e1b2d151c6d..500a0e7bb06 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -7,6 +7,7 @@ import re from collections.abc import Callable from contextlib import asynccontextmanager from datetime import datetime, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -2934,16 +2935,14 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey @pytest.mark.asyncio @pytest.mark.parametrize( "call_type, expects_flush", - [("aresponses", True), ("responses", True), ("aretrieve_batch", True), ("acompletion", False)], + [("aresponses", True), ("responses", True), ("acompletion", False)], ) async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back( call_type: str, expects_flush: bool ): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a - Responses row cannot sit in this worker's queue until the monitor's next poll. A - batch's cost row is what another worker checks before charging the same batch again - (LIT-7048), so it cannot wait either. + Responses row cannot sit in this worker's queue until the monitor's next poll. """ from litellm.proxy.utils import PrismaClient @@ -2961,6 +2960,127 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker PrismaClient.spend_log_flush_requested.clear() +def _batch_cost_payload() -> dict: + return { + **_minimal_spend_payload(), + "request_id": "batch_abc_batch_cost", + "call_type": "aretrieve_batch", + "status": "success", + } + + +def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock: + prisma = _tool_usage_prisma() + prisma.jsonify_object = lambda data: dict(data) + prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) + prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) + return prisma + + +async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict) -> bool: + with ( + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.disable_spend_logs", False + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.prisma_client", prisma + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), + patch( # test-quality-ok: update_database imports the payload builder inside its body, no seam + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=payload, + ), + ): + charged = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"}, + completion_response=None, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.25, + ) + await asyncio.sleep(0) + return charged + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inserted", "existing", "charged"), + [ + (1, None, True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success"), False), + (0, SimpleNamespace(call_type="aretrieve_batch", status="failure"), True), + (0, SimpleNamespace(call_type="aembedding", status="success"), True), + (0, None, True), + ], + ids=[ + "first_retrieve_owns_the_row", + "another_retrieve_already_charged", + "failed_retrieve_holds_the_row", + "client_chosen_call_id_holds_the_row", + "row_gone_between_insert_and_lookup", + ], +) +async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote_its_row( + inserted: int, existing: object, charged: bool +): + """ + Every retrieve of one batch shares one spend row, so the insert that lands first is + the charge and every later retrieve must leave the counters alone (LIT-7048). A row + written by anything but a successful retrieve, say a request whose client picked the + batch id as its call id, must not be able to take the charge away. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(inserted, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged + + claimed_rows = prisma.db.litellm_spendlogs.create_many.await_args.kwargs + assert claimed_rows["skip_duplicates"] is True + assert [(row["request_id"], row["spend"]) for row in claimed_rows["data"]] == [("batch_abc_batch_cost", 0.25)] + assert prisma.spend_log_transactions == [] + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): + """An unreachable DB must not drop the batch's only spend row, nor its charge.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + assert [row["request_id"] for row in prisma.spend_log_transactions] == ["batch_abc_batch_cost"] + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [{**_batch_cost_payload(), "call_type": "acompletion"}, {**_batch_cost_payload(), "status": "failure"}], + ids=["not_a_batch_retrieve", "failed_batch_retrieve"], +) +async def test_update_database_queues_every_other_spend_row_for_the_next_flush(payload: dict): + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + + assert await _update_database_with(db_writer, prisma, payload) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert prisma.spend_log_transactions == [payload] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize( "injected_deployment, attributed", diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index b7037e8d621..2965f8b4006 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,3 +1,4 @@ +import asyncio from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -6,7 +7,6 @@ import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( - _batch_cost_is_trackable_now, _get_budget_reservation_from_metadata, _ProxyDBLogger, _should_track_cost_callback, @@ -783,110 +783,46 @@ def _retrieved_batch(status: str, output_file_id: str | None): ) -def _prisma_client_with(queued_request_ids: tuple[str, ...], stored_row: object) -> MagicMock: - import asyncio - - prisma_client = MagicMock() - prisma_client._spend_log_transactions_lock = asyncio.Lock() - prisma_client.spend_log_transactions = [{"request_id": request_id} for request_id in queued_request_ids] - prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(return_value=stored_row) - return prisma_client - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("status", "output_file_id", "spend_log_id", "prisma_client", "trackable"), + ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"), [ - ("in_progress", None, "batch_abc_batch_cost", None, False), - ("in_progress", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), - ("completed", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), - ("completed", "file-out", "batch_abc_batch_cost", None, True), - ("completed", "file-out", None, _prisma_client_with((), None), True), - ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with(("batch_abc_batch_cost",), None), False), - ( - "completed", - "file-out", - "batch_abc_batch_cost", - _prisma_client_with((), {"request_id": "batch_abc_batch_cost"}), - False, - ), - ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with((), None), True), - ("failed", None, "batch_abc_batch_cost", _prisma_client_with((), None), True), + ("aretrieve_batch", "in_progress", None, True, False, False), + ("aretrieve_batch", "completed", None, True, False, False), + ("aretrieve_batch", "completed", "file-out", False, True, False), + ("aretrieve_batch", "completed", "file-out", True, True, True), + ("aretrieve_batch", "failed", None, True, True, True), + ("acreate_batch", "validating", None, True, True, True), ], ids=[ - "in_progress_without_db", - "in_progress_never_consults_db", - "completed_without_output_yet", - "final_without_db", - "final_without_spend_log_id", - "final_row_queued_for_flush", - "final_row_already_stored", - "final_first_sighting", - "failed_first_sighting", + "retrieve_before_final", + "retrieve_completed_without_output_yet", + "retrieve_after_another_retrieve_charged", + "retrieve_first_final", + "retrieve_failed_batch", + "create_before_final", ], ) -async def test_batch_cost_is_trackable_now(status, output_file_id, spend_log_id, prisma_client, trackable): - """ - A batch is billed from the first retrieve that sees it final and never again: - a poll before that wrote the shared spend row at $0 and pinned it there, and - every completed retrieve after the first charged the key again (LIT-7048). - """ - assert ( - await _batch_cost_is_trackable_now( - batch=_retrieved_batch(status, output_file_id), spend_log_id=spend_log_id, prisma_client=prisma_client - ) - is trackable - ) - - -@pytest.mark.asyncio -async def test_batch_cost_is_trackable_now_when_the_spend_row_lookup_fails(): - """An unreadable spend log table must not drop the batch's only spend row.""" - prisma_client = _prisma_client_with((), None) - prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(side_effect=RuntimeError("db unreachable")) - - assert ( - await _batch_cost_is_trackable_now( - batch=_retrieved_batch("completed", "file-out"), - spend_log_id="batch_abc_batch_cost", - prisma_client=prisma_client, - ) - is True - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("call_type", "status", "output_file_id", "stored_row", "charged"), - [ - ("aretrieve_batch", "in_progress", None, None, False), - ("aretrieve_batch", "completed", "file-out", {"request_id": "batch_abc_batch_cost"}, False), - ("aretrieve_batch", "completed", "file-out", None, True), - ("acreate_batch", "validating", None, None, True), - ], - ids=["retrieve_before_final", "retrieve_already_recorded", "retrieve_first_final", "create_before_final"], -) -async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs and whether the poll's reservation is handed back is the whole observable contract of the gate - call_type, status, output_file_id, stored_row, charged +async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate + call_type, status, output_file_id, row_claimed, spend_written, charged ): """ - Only retrieves are gated, since creating a batch is its own billable request. - A retrieve that writes nothing hands its budget reservation back instead. + A poll before the batch is final used to pin its shared spend row at $0, and every + completed retrieve after the first charged the key again (LIT-7048). Only retrieves + are gated, since creating a batch is its own billable request, and a retrieve that + charges nothing hands its budget reservation back instead. """ logger = _ProxyDBLogger() budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) with ( - patch( # test-quality-ok: prisma_client is a proxy_server global the callback reads lazily, no seam - "litellm.proxy.proxy_server.prisma_client", _prisma_client_with((), stored_row) - ), patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock - ), + ) as mock_increment_spend_counters, patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock - ), + ) as mock_update_cache, patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam "litellm.proxy.proxy_server.proxy_logging_obj" ) as mock_proxy_logging, @@ -895,7 +831,7 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # ) as mock_release_budget_reservation, ): mock_proxy_logging.failed_tracking_alert = AsyncMock() - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=row_claimed) mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() await logger._PROXY_track_cost_callback( @@ -904,13 +840,15 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # start_time=datetime.now(), end_time=datetime.now(), ) + await asyncio.sleep(0) mock_proxy_logging.failed_tracking_alert.assert_not_called() + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if spend_written else 0) + assert mock_increment_spend_counters.await_count == (1 if charged else 0) + assert mock_update_cache.await_count == (1 if charged else 0) if charged: - mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() mock_release_budget_reservation.assert_not_awaited() else: - mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index fc97d760226..a1eb88a7834 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -6,7 +6,6 @@ Symbols pinned here: - ``update_spend_logs_job`` - ``_monitor_spend_logs_queue`` - ``_raise_failed_update_spend_exception`` - - ``spend_log_is_queued`` """ from __future__ import annotations @@ -23,7 +22,6 @@ from litellm.proxy.utils import ( _monitor_spend_logs_queue, _raise_failed_update_spend_exception, drain_spend_logs_queue, - spend_log_is_queued, update_daily_tag_spend, update_spend, update_spend_logs_job, @@ -631,16 +629,3 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None: with pytest.raises(ValueError, match="specific"): asyncio.run(_runner()) - - -@pytest.mark.asyncio -async def test_spend_log_is_queued_matches_only_rows_awaiting_flush( - mock_prisma_client: Any, make_spend_log_row: Any -) -> None: - mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="batch_abc_batch_cost")] - - assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is True - assert await spend_log_is_queued(mock_prisma_client, "batch_abc") is False - - mock_prisma_client.spend_log_transactions = [] - assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is False From 15372967c6cd5085d5d3d9ebb30c5a158f3f8170 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:06:38 -0700 Subject: [PATCH 40/63] fix(azure_ai): read the azure_ai card for gpt-5 series reasoning effort gates Foundry deployments of gpt-6-astra reached through azure_ai used the bare OpenAI card for the reasoning_effort none gates, so temperature and top_p were refused while the azure_ai card says none is supported. AzureAIStudioConfig now dispatches gpt-5 series params through AzureAIGPT5Config, which looks capabilities up under the azure_ai/ prefix the way the azure route does Also carries the search_context_cost_per_query block azure/gpt-6-astra has, adds a flex service tier cost test that fails at the merge base, and keeps the wildcard test from stripping azure_ai/gpt-6-astra out of the provider set --- litellm/llms/azure_ai/chat/transformation.py | 37 ++++++++++++++++++- ...odel_prices_and_context_window_backup.json | 5 +++ model_prices_and_context_window.json | 5 +++ .../llm_cost_calc/test_llm_cost_calc_utils.py | 15 ++++++++ .../chat/test_azure_ai_transformation.py | 22 +++++++++++ .../proxy/auth/test_model_checks.py | 6 ++- 6 files changed, 87 insertions(+), 3 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index f2d405e9a17..05abd5882c6 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj +from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.llms.openai.common_utils import drop_params_from_unprocessable_entity_error from litellm.llms.openai.openai import OpenAIConfig from litellm.llms.xai.chat.transformation import XAIChatConfig @@ -42,12 +43,25 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( ) +class AzureAIGPT5Config(OpenAIGPT5Config): + @classmethod + def _model_map_lookup_name(cls, model: str) -> str: + return model if model.startswith("azure_ai/") else f"azure_ai/{model}" + + +azureAIGPT5Config: Final = AzureAIGPT5Config() + + class AzureAIStudioConfig(OpenAIConfig): def get_supported_openai_params(self, model: str) -> list: model_supports_tool_choice = True # azure ai supports this by default if not supports_tool_choice(model=f"azure_ai/{model}"): model_supports_tool_choice = False - supported_params = super().get_supported_openai_params(model) + supported_params = ( + azureAIGPT5Config.get_supported_openai_params(model) + if azureAIGPT5Config.is_model_gpt_5_model(model) + else super().get_supported_openai_params(model) + ) if not model_supports_tool_choice: filtered_supported_params: Final = [] for param in supported_params: @@ -61,6 +75,27 @@ class AzureAIStudioConfig(OpenAIConfig): return supported_params + def map_openai_params( + self, + non_default_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature + optional_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: OpenAIConfig.map_openai_params signature + if not azureAIGPT5Config.is_model_gpt_5_model(model): + return super().map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) + return azureAIGPT5Config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) + def _supports_stop_reason(self, model: str) -> bool: """ Check if the model supports stop tokens. diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ac7407c2608..5b4652d38c7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3499,6 +3499,11 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://ai.azure.com/catalog/models/gpt-6-astra", "supported_endpoints": [ "/v1/chat/completions", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ac7407c2608..5b4652d38c7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3499,6 +3499,11 @@ "mode": "chat", "output_cost_per_token": 5e-05, "output_cost_per_token_above_272k_tokens": 7.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://ai.azure.com/catalog/models/gpt-6-astra", "supported_endpoints": [ "/v1/chat/completions", 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 73f1a19d85c..caf97ba791d 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 @@ -2060,6 +2060,21 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( assert completion_cost == pytest.approx(zone_multiplier * output_multiplier * completion_tokens * 5e-5) +def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map): + """Foundry sells gpt-6-astra on Standard Global only, so a flex service_tier bills the standard rate. + The bare OpenAI card the azure_ai route fell back to before this entry existed carries flex prices + at half rate (LIT-7081).""" + usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) + + standard = generic_cost_per_token(model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai") + flex = generic_cost_per_token( + model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai", service_tier="flex" + ) + + assert flex == standard + assert standard == pytest.approx((1000 * 1e-05, 100 * 5e-05)) + + @pytest.mark.parametrize( "model,expected_none,expected_xhigh,expected_minimal", [ 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 33fbb4e8fc7..25eca3b37ad 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 @@ -3,6 +3,8 @@ from unittest.mock import MagicMock, patch import pytest +import litellm +from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map from litellm.llms.azure_ai.azure_model_router.transformation import ( AzureModelRouterConfig, ) @@ -138,6 +140,26 @@ def test_azure_ai_validate_environment_with_azure_ad_token(): assert headers["Content-Type"] == "application/json" +@pytest.fixture +def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url=litellm.model_cost_map_url)) + + +def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none(_local_model_cost_map): + """A Foundry deployment reached through azure_ai reads the azure_ai/ card, where gpt-6-astra supports + reasoning_effort none, so temperature and top_p ride along; the bare OpenAI card says none is + unsupported and the route used to refuse temperature and drop top_p (LIT-7081).""" + optional_params = AzureAIStudioConfig().map_openai_params( + non_default_params={"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9}, + optional_params={}, + model="gpt-6-astra", + drop_params=False, + ) + + assert optional_params == {"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9} + + def test_azure_ai_grok_stop_parameter_handling(): """ Test that Grok models properly handle stop parameter filtering in Azure AI Studio. diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index eb48f70d5da..56dbcca61f3 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -868,12 +868,14 @@ def test_azure_ai_wildcard_lists_the_foundry_gpt_6_astra_entry(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") foundry_key = "azure_ai/gpt-6-astra" local_entry = litellm.get_model_cost_map(url="")[foundry_key] + registered_before = foundry_key in litellm.azure_ai_models try: litellm.add_known_models(model_cost_map={foundry_key: local_entry}) assert foundry_key in get_known_models_from_wildcard("azure_ai/*") finally: - litellm.azure_ai_models.discard(foundry_key) - litellm.add_known_models(model_cost_map={}) + if not registered_before: + litellm.azure_ai_models.discard(foundry_key) + litellm.add_known_models(model_cost_map={}) def test_get_complete_model_list_drops_no_default_models_sentinel(): From a17fcecf7092d0333afed4168d1954a1e9675d18 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:21:34 -0700 Subject: [PATCH 41/63] refactor(azure_ai): type the Foundry param mapping override and drop test docstrings The AzureAIStudioConfig.map_openai_params override now carries dict[str, object] annotations instead of bare dict, and the docstrings added to the new tests go away since the test names already say what they cover. No behavior change --- litellm/llms/azure_ai/chat/transformation.py | 6 +++--- .../llm_cost_calc/test_llm_cost_calc_utils.py | 3 --- .../llms/azure_ai/chat/test_azure_ai_transformation.py | 3 --- tests/test_litellm/proxy/auth/test_model_checks.py | 3 --- .../router_utils/test_reasoning_effort_capability.py | 3 +-- 5 files changed, 4 insertions(+), 14 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 05abd5882c6..7c9a26c3f07 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -77,11 +77,11 @@ class AzureAIStudioConfig(OpenAIConfig): def map_openai_params( self, - non_default_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature - optional_params: dict, # mutable-ok: OpenAIConfig.map_openai_params signature + non_default_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature + optional_params: dict[str, object], # mutable-ok: OpenAIConfig.map_openai_params signature model: str, drop_params: bool, - ) -> dict: # mutable-ok: OpenAIConfig.map_openai_params signature + ) -> dict[str, object]: # mutable-ok: OpenAIConfig.map_openai_params signature if not azureAIGPT5Config.is_model_gpt_5_model(model): return super().map_openai_params( non_default_params=non_default_params, 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 caf97ba791d..40abb5bfca3 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 @@ -2061,9 +2061,6 @@ def test_generic_cost_per_token_azure_gpt_6_astra_foundry_price_sheet( def test_generic_cost_per_token_azure_ai_gpt_6_astra_flex_bills_the_standard_rate(_local_model_cost_map): - """Foundry sells gpt-6-astra on Standard Global only, so a flex service_tier bills the standard rate. - The bare OpenAI card the azure_ai route fell back to before this entry existed carries flex prices - at half rate (LIT-7081).""" usage = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) standard = generic_cost_per_token(model="azure_ai/gpt-6-astra", usage=usage, custom_llm_provider="azure_ai") 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 25eca3b37ad..5ff0b729449 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 @@ -147,9 +147,6 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none(_local_model_cost_map): - """A Foundry deployment reached through azure_ai reads the azure_ai/ card, where gpt-6-astra supports - reasoning_effort none, so temperature and top_p ride along; the bare OpenAI card says none is - unsupported and the route used to refuse temperature and drop top_p (LIT-7081).""" optional_params = AzureAIStudioConfig().map_openai_params( non_default_params={"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9}, optional_params={}, diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index 56dbcca61f3..36bfc4c5dd3 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -859,9 +859,6 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): def test_azure_ai_wildcard_lists_the_foundry_gpt_6_astra_entry(monkeypatch): - """A Foundry (azure_ai) deployment of gpt-6-astra only shows up under an azure_ai/* wildcard - when the cost map carries its own azure_ai/ entry; the azure/ entry from the OpenAI-on-Azure - price sheet never reaches the Foundry provider list (LIT-7081).""" import litellm from litellm.proxy.auth.model_checks import get_known_models_from_wildcard diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index b7499d1c975..3e1f26b6e1c 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -400,8 +400,7 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): """Microsoft Foundry serves the same model but its API accepts reasoning_effort none (verified live: 200 with zero reasoning tokens, and it unlocks temperature), which - OpenAI's rejects, so an Azure deployment offers none on top of low through max, whether - it is reached through the azure route or the azure_ai (Foundry) route.""" + OpenAI's rejects, so an Azure deployment offers none on top of low through max.""" from litellm.utils import _get_model_info_helper model_info = dict(_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)) From 9fd60e4f95228718ed74f75ace470016ee4f42e4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 19:24:00 -0700 Subject: [PATCH 42/63] feat(router): gate heuristic v1 tuning (#39952) --- .../model_management_endpoints.py | 95 ++++++--- litellm/proxy/proxy_server.py | 90 ++++++++- litellm/repositories/model_repository.py | 10 +- .../auto_router_tuning_baseline.py | 157 +++++++++++++++ .../test_model_management_endpoints.py | 186 +++++++++++++++++- .../proxy/proxy_server/test_lifecycle.py | 18 ++ .../test_auto_router_tuning_baseline.py | 185 +++++++++++++++++ 7 files changed, 710 insertions(+), 31 deletions(-) create mode 100644 litellm/router_utils/auto_router_tuning_baseline.py create mode 100644 tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b77108911aa..0f19e9ce149 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( _refresh_cached_team, @@ -109,6 +110,7 @@ from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, validate_strategy_router_model_write, ) +from litellm.router_utils.auto_router_tuning_baseline import is_mutable_tuned_candidate, tuning_quota_violation from litellm.types.proxy.management_endpoints.model_management_endpoints import ( AutoRouterClassifierDefaultPromptResponse, UpdateUsefulLinksRequest, @@ -349,6 +351,36 @@ def _effective_complexity_router_params( ) +def _decrypted_model(stored_model: object) -> str | None: + if not isinstance(stored_model, str): + return None + decrypted: Final = decrypt_value_helper( + value=stored_model, key="model", exception_type="debug", return_original_value=True + ) + return decrypted if isinstance(decrypted, str) else None + + +def _tuning_candidate(effective_params: Mapping[str, object], model_id: str | None) -> Mapping[str, object]: + return MappingProxyType( + { + "litellm_params": effective_params, + "model_info": MappingProxyType({"id": model_id, "db_model": True}), + } + ) + + +def _raise_on_tuning_quota_violation( + *, + candidate: Mapping[str, object], + others: Sequence[Mapping[str, object]], + baselines: Mapping[str, str], + limit: int | None, +) -> None: + violation: Final = tuning_quota_violation(candidate=candidate, others=others, baselines=baselines, limit=limit) + if violation is not None: + raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}") + + @asynccontextmanager async def _auto_router_capability_slot( prisma_client: PrismaClient, *, effective_params: Mapping[str, object], model_id: str | None @@ -366,40 +398,55 @@ async def _auto_router_capability_slot( must wait until the transaction has committed and the lock is released. The transaction writes bypass the repository's publish-on-write, so the config change is published once after commit, the way delete_team_models does. + + A heuristic-v1 router whose tuning has moved off its recorded baseline is judged the same + way under the same lock, against the DB rows plus this proxy's config.yaml routers. """ - from litellm.proxy.proxy_server import _license_check, llm_router + from litellm.proxy.proxy_server import ( + _license_check, # pyright: ignore[reportPrivateUsage] # existing capability slot reads the proxy license singleton + heuristic_v1_tuning_baselines, + llm_router, + ) limit: Final = _license_check.auto_router_capability_limit() capability: Final = gated_capability_of(effective_params) - if limit is None or capability is None: + baselines: Final = heuristic_v1_tuning_baselines + tuning_candidate: Final = _tuning_candidate(effective_params, model_id=model_id) + judges_tuning: Final = baselines is not None and is_mutable_tuned_candidate(tuning_candidate, baselines) + if limit is None or (capability is None and not judges_tuning): yield _proxy_model_table(prisma_client) return async with prisma_client.db.tx() as tx_ctx: tables: Final[_TxModelTables] = tx_ctx await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) - rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw( - _CAPABILITY_DB_ROWS_SQL[capability.key], model_id or "" - ) - db_held: Final = sum( - 1 - for row in rows - for stored_model in (row.get("model"),) - if isinstance(stored_model, str) - and is_complexity_router_model( - decrypt_value_helper( - value=stored_model, - key="model", - exception_type="debug", - return_original_value=True, - ) - ) - ) config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments()) - held: Final = db_held + count_capability_routers(config_rows, capability=capability) - violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit) - if violation is not None: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" + if capability is not None: + rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw( + _CAPABILITY_DB_ROWS_SQL[capability.key], model_id or "" + ) + db_held: Final = sum(1 for row in rows if is_complexity_router_model(_decrypted_model(row.get("model")))) + held: Final = db_held + count_capability_routers(config_rows, capability=capability) + violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit) + if violation is not None: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" + ) + if judges_tuning and baselines is not None: + model_rows: Final = await ModelRepository(WriterPinnedClient(tx_ctx)).find_all_except(model_id or "") + _raise_on_tuning_quota_violation( + candidate=tuning_candidate, + others=tuple( + MappingProxyType( + { + "litellm_params": row.litellm_params, + "model_info": MappingProxyType({"id": row.model_id, "db_model": True}), + } + ) + for row in model_rows + ) + + config_rows, + baselines=baselines, + limit=limit, ) yield tables.litellm_proxymodeltable await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1cd08fe27c0..ba5714fe950 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -125,6 +125,12 @@ from litellm.router_utils.auto_router_model_naming import ( count_capability_routers, validate_complexity_router_config_placement, ) +from litellm.router_utils.auto_router_tuning_baseline import ( + TUNING_BASELINE_PARAM_NAME, + mutable_tuned_identities, + snapshot_tuning_baselines, + tuning_limit_violation, +) from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -892,7 +898,8 @@ def cleanup_router_config_variables(): use_shared_health_check, \ health_check_interval, \ health_check_concurrency, \ - prisma_client + prisma_client, \ + heuristic_v1_tuning_baselines # Set all variables to None master_key = None @@ -911,6 +918,7 @@ def cleanup_router_config_variables(): health_check_interval = None health_check_concurrency = None prisma_client = None + heuristic_v1_tuning_baselines = None async def _flush_spend_logs_queue_on_shutdown() -> None: @@ -2267,6 +2275,7 @@ experimental = False #### GLOBAL VARIABLES #### llm_router: Router | None = None llm_model_list: list | None = None +heuristic_v1_tuning_baselines: Mapping[str, str] | None = None # Serializes every model reconcile (ProxyConfig.add_deployment and clear_cache) so the # read-modify-write of llm_router above is atomic. Without it, two concurrent model # writes each reconcile the router against their OWN db snapshot, and the one holding @@ -9324,6 +9333,77 @@ class ProxyStartupEvent: except Exception as e: verbose_proxy_logger.debug("UI settings sync on startup skipped or failed: %s", e) + @classmethod + async def _load_heuristic_v1_tuning_baselines( + cls, prisma_client: PrismaClient, deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, str] | None: + """Read the recorded tuning baselines, recording current routers on the first boot.""" + from prisma.errors import UniqueViolationError + + try: + config_table: Final = prisma_client.db.litellm_config + row: Final = await config_table.find_unique( + where={"param_name": TUNING_BASELINE_PARAM_NAME} # mutable-ok: Prisma rejects mappingproxy input + ) + if row is not None: + stored: Final = row.param_value + decoded: Final = json.loads(stored) if isinstance(stored, str) else stored + return MappingProxyType( + { + str(identity): str(fingerprint) + for identity, fingerprint in (decoded.items() if isinstance(decoded, Mapping) else ()) + } + ) # mutable-ok: MappingProxyType owns the completed immutable baseline + snapshot: Final = snapshot_tuning_baselines(deployments) + try: + await config_table.create( + data={ # mutable-ok: Prisma rejects mappingproxy input + "param_name": TUNING_BASELINE_PARAM_NAME, + "param_value": json.dumps(dict(snapshot)), # mutable-ok: json only serializes concrete mappings + } + ) + verbose_proxy_logger.info("Recorded heuristic-v1 tuning baseline for %s auto-router(s)", len(snapshot)) + return snapshot + except UniqueViolationError: + competing_row: Final = await config_table.find_unique( + where={"param_name": TUNING_BASELINE_PARAM_NAME} # mutable-ok: Prisma rejects mappingproxy input + ) + competing_value: Final = None if competing_row is None else competing_row.param_value + competing_decoded: Final = ( + json.loads(competing_value) if isinstance(competing_value, str) else competing_value + ) + return MappingProxyType( + { + str(identity): str(fingerprint) + for identity, fingerprint in ( + competing_decoded.items() if isinstance(competing_decoded, Mapping) else () + ) + } + ) # mutable-ok: MappingProxyType owns the completed immutable baseline + except Exception as e: # noqa: BLE001 # enforcement is skipped for this boot; refusing every tuned router on a DB blip is the one outcome the gate forbids + verbose_proxy_logger.warning("Heuristic-v1 tuning baseline unavailable, gate not enforced this boot: %s", e) + return None + + @classmethod + async def enforce_heuristic_v1_tuning_baseline( + cls, prisma_client: PrismaClient, llm_router: Router | None, limit: int | None + ) -> Mapping[str, str] | None: + """Load a complete baseline and reject a startup that exceeds the tuning quota.""" + db_models: Final = await proxy_config._get_models_from_db(prisma_client) + if db_models is None: + verbose_proxy_logger.warning("Heuristic-v1 tuning baseline unavailable, gate not enforced this boot") + return None + config_deployments: Final = () if llm_router is None else tuple(llm_router.config_deployments()) + deployments: Final = (*config_deployments, *proxy_config.decrypt_model_list_from_db(db_models)) + baselines: Final = await cls._load_heuristic_v1_tuning_baselines(prisma_client, deployments) + if baselines is None: + return None + mutable: Final = mutable_tuned_identities(deployments, baselines) + violation: Final = tuning_limit_violation(held=len(mutable), limit=limit) + if violation is not None: + raise ValueError(f"model_list: {violation} {AUTO_ROUTER_LICENSE_REMEDY}") + return baselines + @classmethod async def initialize_scheduled_background_jobs( cls, @@ -9335,7 +9415,7 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, ) -> ProxyWorkerHeartbeat: """Initializes scheduled background jobs""" - global store_model_in_db, scheduler + global heuristic_v1_tuning_baselines, store_model_in_db, scheduler # rebind-ok: startup publishes the one read-only baseline snapshot # MEMORY LEAK FIX: Configure scheduler with optimized settings # Memray analysis showed APScheduler's normalize() and _apply_jitter() causing @@ -9573,6 +9653,12 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + heuristic_v1_tuning_baselines = await cls.enforce_heuristic_v1_tuning_baseline( + prisma_client=prisma_client, + llm_router=llm_router, + limit=_license_check.auto_router_capability_limit(), + ) + await cls._initialize_slack_alerting_jobs( scheduler=scheduler, general_settings=general_settings, diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index acc7c8dcda8..540f197044f 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -3,7 +3,8 @@ Model repository for database operations on LiteLLM_ProxyModelTable. """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol from litellm.models.model import LiteLLM_ProxyModelTable @@ -105,6 +106,13 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): records: Final = await self.table.find_many(where={"blocked": False}) return self._to_model_list(records) + async def find_all_except(self, model_id: str) -> Sequence[LiteLLM_ProxyModelTable]: + """Find every model except the row currently being updated.""" + records: Final = await self.table.find_many( + where=MappingProxyType({"model_id": MappingProxyType({"not": model_id})}) + ) + return tuple(self._to_model_list(records)) + async def find_by_team_id(self, team_id: str) -> list[LiteLLM_ProxyModelTable]: """Find models associated with a specific team. diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py new file mode 100644 index 00000000000..1e6412856a6 --- /dev/null +++ b/litellm/router_utils/auto_router_tuning_baseline.py @@ -0,0 +1,157 @@ +"""Baseline-relative license gate for heuristic-v1 complexity-router tuning.""" + +import hashlib +import json +from collections.abc import Iterable, Mapping +from types import MappingProxyType +from typing import Final + +from pydantic import ValidationError + +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig + +TUNING_BASELINE_PARAM_NAME: Final = "auto_router_tuning_baseline" + +HEURISTIC_V1_TUNING_FIELDS: Final = ( + "tiers", + "tier_model_configs", + "classifier_type", + "tier_boundaries", + "reasoning_override_min_score", + "token_thresholds", + "dimension_weights", + "code_keywords", + "reasoning_keywords", + "technical_keywords", + "custom_technical_keywords", + "simple_keywords", + "escalation_keywords", + "keyword_tier_rules", +) + +_V1_SCORING_CLASSIFIER_TYPES: Final = frozenset({"heuristic", "heuristic_first", "hybrid"}) +_AUTO_ROUTER_COMPLEXITY_PREFIX: Final = "auto_router/complexity_router" +_EMPTY: Final[Mapping[str, object]] = MappingProxyType({}) +_EMPTY_TAGS: Final[tuple[str, ...]] = () + + +def _mapping(value: object) -> Mapping[str, object]: + return value if isinstance(value, Mapping) else _EMPTY + + +def tuning_fingerprint(complexity_router_config: object) -> str | None: + """Digest of normalized heuristic-v1 tuning fields, or None when the config is invalid.""" + try: + validated: Final = ComplexityRouterConfig.model_validate(_mapping(complexity_router_config)) + except ValidationError: + return None + payload: Final = validated.model_dump(mode="json", include=frozenset(HEURISTIC_V1_TUNING_FIELDS)) + return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() + + +DEFAULT_TUNING_FINGERPRINT: Final = tuning_fingerprint(_EMPTY) + + +def uses_heuristic_v1(complexity_router_config: object) -> bool: + """Whether a config's primary classifier path is the heuristic-v1 scorer.""" + return _mapping(complexity_router_config).get("classifier_type", "heuristic") in _V1_SCORING_CLASSIFIER_TYPES + + +def router_identity(deployment: Mapping[str, object]) -> str | None: + """Stable identity for a complexity-router deployment, across tuning edits.""" + model_info: Final = _mapping(deployment.get("model_info")) + model_id: Final = model_info.get("id") + if model_info.get("db_model") is True and isinstance(model_id, str) and model_id: + return f"db:{model_id}" + model_name: Final = deployment.get("model_name") + if not isinstance(model_name, str) or not model_name: + return None + litellm_params: Final = _mapping(deployment.get("litellm_params")) + tags: Final = litellm_params.get("tags") + normalized_tags: Final = ( + tuple(sorted(str(tag) for tag in tags)) + if isinstance(tags, Iterable) and not isinstance(tags, str) + else _EMPTY_TAGS + ) + return f"yaml:{json.dumps((model_name, normalized_tags), separators=(',', ':'))}" + + +def heuristic_v1_router_fingerprint(deployment: Mapping[str, object]) -> tuple[str, str] | None: + """The identity/fingerprint pair for a heuristic-v1 complexity router, else None.""" + litellm_params: Final = _mapping(deployment.get("litellm_params")) + model: Final = litellm_params.get("model") + config: Final = litellm_params.get("complexity_router_config") + if ( + not isinstance(model, str) + or not model.startswith(_AUTO_ROUTER_COMPLEXITY_PREFIX) + or not uses_heuristic_v1(config) + ): + return None + identity: Final = router_identity(deployment) + fingerprint: Final = tuning_fingerprint(config) + if identity is None or fingerprint is None: + return None + return identity, fingerprint + + +def snapshot_tuning_baselines(deployments: Iterable[Mapping[str, object]]) -> Mapping[str, str]: + """One immutable first-observation baseline for every heuristic-v1 complexity router.""" + return MappingProxyType( + { + identity: fingerprint + for deployment in deployments + if (pair := heuristic_v1_router_fingerprint(deployment)) is not None + for identity, fingerprint in (pair,) + } + ) # mutable-ok: MappingProxyType owns the completed immutable snapshot + + +def is_mutable_tuned_candidate(candidate: Mapping[str, object], baselines: Mapping[str, str]) -> bool: + pair: Final = heuristic_v1_router_fingerprint(candidate) + if pair is None: + return False + identity, fingerprint = pair + return fingerprint != baselines.get(identity, DEFAULT_TUNING_FINGERPRINT) + + +def mutable_tuned_identities( + deployments: Iterable[Mapping[str, object]], baselines: Mapping[str, str] +) -> frozenset[str]: + """Heuristic-v1 routers whose current tuning differs from their baseline or shipped default.""" + return frozenset( + identity + for deployment in deployments + if (pair := heuristic_v1_router_fingerprint(deployment)) is not None + for identity, _ in (pair,) + if is_mutable_tuned_candidate(deployment, baselines) + ) + + +def tuning_limit_violation(*, held: int, limit: int | None) -> str | None: + if limit is None or held <= limit: + return None + return ( + f"At most {limit} auto-router(s) with changed heuristic scorer settings or tier models can be modified " + "without an auto-router license. Keep this router on its recorded settings, or revert the other changed " + "router to its baseline, or remove one of them." + ) + + +def tuning_quota_violation( + *, + candidate: Mapping[str, object], + others: Iterable[Mapping[str, object]], + baselines: Mapping[str, str], + limit: int | None, +) -> str | None: + """Why a change to candidate tuning exceeds the baseline-relative free quota.""" + if limit is None: + return None + pair: Final = heuristic_v1_router_fingerprint(candidate) + if pair is None: + return None + identity, _ = pair + if not is_mutable_tuned_candidate(candidate, baselines): + return None + held: Final = mutable_tuned_identities(others, baselines) - frozenset((identity,)) + return tuning_limit_violation(held=len(held) + 1, limit=limit) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 33de2a09626..63c08529b4d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -4440,13 +4440,28 @@ class TestStrategyRouterWriteValidation: class _FakeTx: """Stands in for a prisma transaction: records raw statements and returns encrypted-model candidates.""" - def __init__(self, db_models: list[str]) -> None: + def __init__(self, db_models: list[str], tuning_rows: list[dict[str, object]] | None = None) -> None: self.db_models = db_models + self.tuning_rows = tuning_rows or [] self.raw_calls: list[tuple[str, tuple[object, ...]]] = [] - self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock()) + self.litellm_proxymodeltable = MagicMock( + create=AsyncMock(), + update=AsyncMock(), + find_many=AsyncMock( + return_value=tuple( + LiteLLM_ProxyModelTable.model_validate(row) for row in self.tuning_rows + ) + ), + ) + + @property + def db(self) -> "TestStrategyRouterWriteValidation._FakeTx": + return self async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]: self.raw_calls.append((sql, args)) + if "AS litellm_params" in sql: + return [row for row in self.tuning_rows if row.get("model_id") != args[0]] return [{"model": model} for model in self.db_models] if "AS model" in sql else [] async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx": @@ -4458,9 +4473,11 @@ class TestStrategyRouterWriteValidation: class _FakeDb: """Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity.""" - def __init__(self, db_models: list[str], existing_row: object = None) -> None: + def __init__( + self, db_models: list[str], existing_row: object = None, tuning_rows: list[dict[str, object]] | None = None + ) -> None: self.db = self - self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models) + self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models, tuning_rows=tuning_rows) self.litellm_proxymodeltable = MagicMock( create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row) ) @@ -4617,6 +4634,167 @@ class TestStrategyRouterWriteValidation: assert capability is not None assert capability.sql_config_predicate.split("{config}")[-1].strip() in count_sql + _TUNED_A = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}} + _TUNED_A_EDITED = {**_TUNED_A, "dimension_weights": {"codePresence": 0.9}} + _TUNED_B = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4.1"}} + _TUNED_B_EDITED = {**_TUNED_B, "tiers": {"SIMPLE": "gpt-4o", "MEDIUM": "gpt-4.1"}} + + @staticmethod + def _db_router_row(model_id: str, config: Mapping[str, object]) -> dict[str, object]: + return { + "model_name": f"router-{model_id}", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": dict(config)}, + "model_info": {"id": model_id, "db_model": True}, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "limit,baseline_rows,live_rows,candidate_id,candidate_config,expected", + [ + (1, ["a", "b"], {"a": "_TUNED_A", "b": "_TUNED_B"}, "a", "_TUNED_A", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A", "b": "_TUNED_B"}, "a", "_TUNED_A_EDITED", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "a", "_TUNED_A_EDITED", "allowed"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B_EDITED", "refused"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "c", "_TUNED_B", "refused"), + (1, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B", "allowed"), + (None, ["a", "b"], {"a": "_TUNED_A_EDITED", "b": "_TUNED_B"}, "b", "_TUNED_B_EDITED", "allowed"), + (1, [], {}, "c", "_TUNED_B", "allowed"), + ], + ) + async def test_slot_enforces_baseline_relative_tuning_quota( + self, + limit: int | None, + baseline_rows: list[str], + live_rows: Mapping[str, str], + candidate_id: str, + candidate_config: str, + expected: str, + ) -> None: + """Without the license, one router may move off its recorded tuning baseline and keep being edited; + a change to a second router, or a second new tuned router, is refused. Unchanged baselines and + reverts to baseline are never counted, and a license lifts every check.""" + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + + configs = { + "_TUNED_A": self._TUNED_A, + "_TUNED_A_EDITED": self._TUNED_A_EDITED, + "_TUNED_B": self._TUNED_B, + "_TUNED_B_EDITED": self._TUNED_B_EDITED, + } + baselines = snapshot_tuning_baselines( + [self._db_router_row(row_id, configs["_TUNED_A" if row_id == "a" else "_TUNED_B"]) for row_id in baseline_rows] + ) + effective_params = { + "model": "auto_router/complexity_router", + "complexity_router_config": configs[candidate_config], + } + # The other routers live in the DB, so the slot must read them under its own lock rather than + # trusting this pod's in-memory router: another pod's write is invisible to that list. + fake = self._FakeDb( + [], + tuning_rows=[ + { + "model_id": row_id, + "model_name": f"router-{row_id}", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": configs[name], + }, + } + for row_id, name in live_rows.items() + ], + ) + with ( + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + ): + if expected == "refused": + with pytest.raises(HTTPException) as exc_info: + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id): + pass + assert exc_info.value.status_code == 403 + assert "changed heuristic scorer settings or tier models" in str(exc_info.value.detail) + assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail) + return + async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=candidate_id) as table: + assert hasattr(table, "create") + + @pytest.mark.asyncio + async def test_add_new_model_refuses_a_second_tuned_heuristic_v1_router_without_a_model_id(self) -> None: + """A create request carries no model_info at all, yet the quota still judges it: Deployment mints the + row id before the slot is entered, so a second tuned router is refused before its DB write.""" + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + baselines = snapshot_tuning_baselines([self._db_router_row("a", self._TUNED_A)]) + fake = self._FakeDb( + [], + tuning_rows=[ + { + "model_id": "a", + "model_name": "router-a", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": self._TUNED_A_EDITED, + }, + } + ], + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", baselines), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + patch( # test-quality-ok: prior auth check needs a live DB; only the tuning quota is under test + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + lambda value, new_encryption_key=None: value, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model( + model_params=Deployment( + model_name="second-tuned", + litellm_params=LiteLLM_Params( + model="auto_router/complexity_router", complexity_router_config=self._TUNED_B + ), + ), + user_api_key_dict=admin, + ) + assert exc_info.value.code == "403" + assert "changed heuristic scorer settings or tier models" in str(exc_info.value.message) + fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited() + fake.litellm_proxymodeltable.create.assert_not_awaited() + + @pytest.mark.asyncio + async def test_slot_skips_tuning_quota_when_no_baseline_is_loaded(self) -> None: + """No baseline (DB-less proxy, or the startup read failed) means the gate cannot judge, so it does not.""" + from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot + + fake = self._FakeDb([]) + with ( + patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: the guard reads the proxy license singleton with no injection seam + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: the guard reads the proxy router global with no injection seam + patch("litellm.proxy.proxy_server.heuristic_v1_tuning_baselines", None), # test-quality-ok: baselines are a startup-loaded proxy global with no injection seam + ): + async with _auto_router_capability_slot( + fake, + effective_params={"model": "auto_router/complexity_router", "complexity_router_config": self._TUNED_B}, + model_id="c", + ) as table: + assert hasattr(table, "create") + @pytest.mark.asyncio async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None: """team_model_add needs a second pool connection, so it must run only after the slot transaction diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a06e7142122..a518c9289ae 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -61,6 +61,7 @@ def test_cleanup_router_config_variables_resets_globals(monkeypatch): monkeypatch.setattr(ps, "user_custom_auth", lambda x: x, raising=False) monkeypatch.setattr(ps, "health_check_interval", 42, raising=False) monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False) + monkeypatch.setattr(ps, "heuristic_v1_tuning_baselines", {"router": "baseline"}, raising=False) cleanup_router_config_variables() @@ -70,6 +71,7 @@ def test_cleanup_router_config_variables_resets_globals(monkeypatch): "user_custom_auth": ps.user_custom_auth, "health_check_interval": ps.health_check_interval, "prisma_client": ps.prisma_client, + "heuristic_v1_tuning_baselines": ps.heuristic_v1_tuning_baselines, } assert normalize(observed) == { "master_key": None, @@ -77,6 +79,7 @@ def test_cleanup_router_config_variables_resets_globals(monkeypatch): "user_custom_auth": None, "health_check_interval": None, "prisma_client": None, + "heuristic_v1_tuning_baselines": None, } @@ -818,6 +821,21 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): ) +@pytest.mark.asyncio +async def test_tuning_baseline_waits_for_a_complete_db_model_census(monkeypatch): + prisma_client = MagicMock() + monkeypatch.setattr(ps.proxy_config, "_get_models_from_db", AsyncMock(return_value=None)) + + result = await ProxyStartupEvent.enforce_heuristic_v1_tuning_baseline( + prisma_client=prisma_client, + llm_router=None, + limit=1, + ) + + assert result is None + prisma_client.db.litellm_config.find_unique.assert_not_called() + + # --------------------------------------------------------------------------- # _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809) # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py new file mode 100644 index 00000000000..c3395ab641f --- /dev/null +++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py @@ -0,0 +1,185 @@ +"""Behavior pins for the baseline-relative heuristic-v1 tuning gate.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import pytest + +from litellm.router_utils.auto_router_tuning_baseline import ( + DEFAULT_TUNING_FINGERPRINT, + HEURISTIC_V1_TUNING_FIELDS, + heuristic_v1_router_fingerprint, + mutable_tuned_identities, + router_identity, + snapshot_tuning_baselines, + tuning_fingerprint, + tuning_limit_violation, + tuning_quota_violation, +) + +_TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"} +_ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"} + + +def _router( + name: str, + config: Mapping[str, object] | None, + *, + tags: list[str] | None = None, + db_id: str | None = None, + model: str = "auto_router/complexity_router", +) -> dict[str, object]: + litellm_params: dict[str, object] = {"model": model} + if config is not None: + litellm_params["complexity_router_config"] = dict(config) + if tags is not None: + litellm_params["tags"] = tags + row: dict[str, object] = {"model_name": name, "litellm_params": litellm_params} + if db_id is not None: + row["model_info"] = {"id": db_id, "db_model": True} + return row + + +class TestTuningFingerprint: + def test_normalized_spellings_share_one_fingerprint(self) -> None: + canonical = tuning_fingerprint({"tiers": _TIERS, "dimension_weights": {"codePresence": 0.3}}) + assert canonical == tuning_fingerprint({"dimension_weights": {"codePresence": 0.3}, "tiers": _TIERS}) + assert tuning_fingerprint({"tiers": {"SIMPLE": {"model_name": "x"}}}) == tuning_fingerprint( + {"tiers": {"SIMPLE": "x"}} + ) + + @pytest.mark.parametrize("field", sorted(set(HEURISTIC_V1_TUNING_FIELDS) - {"tier_model_configs"})) + def test_every_tuning_field_changes_the_fingerprint(self, field: str) -> None: + samples: dict[str, object] = { + "tiers": _ALT_TIERS, + "classifier_type": "heuristic_first", + "tier_boundaries": {"simple_medium": 0.2, "medium_complex": 0.4, "complex_reasoning": 0.7}, + "reasoning_override_min_score": 0.05, + "token_thresholds": {"simple": 20, "complex": 500}, + "dimension_weights": {"codePresence": 0.9}, + "code_keywords": ["orionflow"], + "reasoning_keywords": ["deduce"], + "technical_keywords": ["ledgerkit"], + "custom_technical_keywords": ["acmeflow"], + "simple_keywords": ["hey"], + "escalation_keywords": ["ESCALATE"], + "keyword_tier_rules": [{"keywords": ["urgent"], "tier": "COMPLEX"}], + } + config: dict[str, object] = {field: samples[field]} + if field == "classifier_type": + config["heuristic_first_max_tier"] = "MEDIUM" + config["classifier_llm_config"] = {"model": "judge"} + assert tuning_fingerprint(config) != DEFAULT_TUNING_FINGERPRINT + + def test_tier_model_overrides_change_the_fingerprint(self) -> None: + plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}}) + with_override = tuning_fingerprint( + {"tiers": {"SIMPLE": {"model_name": "x", "litellm_params": {"temperature": 0.1}}}} + ) + assert plain != with_override + + def test_non_tuning_fields_do_not_change_the_fingerprint(self) -> None: + assert tuning_fingerprint({"return_raw_model_name": True, "session_affinity": True}) == DEFAULT_TUNING_FINGERPRINT + + def test_invalid_config_has_no_fingerprint(self) -> None: + assert tuning_fingerprint({"tier_boundaries": "not-a-mapping"}) is None + + +class TestRouterIdentity: + def test_db_rows_key_on_model_id_and_yaml_rows_on_name_and_tags(self) -> None: + db_row = _router("renamed", {"tiers": _TIERS}, db_id="row-1") + assert router_identity(db_row) == router_identity(_router("other-name", {"tiers": _TIERS}, db_id="row-1")) + assert router_identity(_router("a", {"tiers": _TIERS}, tags=["x", "y"])) == router_identity( + _router("a", {"tiers": _TIERS}, tags=["y", "x"]) + ) + assert router_identity(_router("a", {"tiers": _TIERS}, tags=["x"])) != router_identity( + _router("a", {"tiers": _TIERS}) + ) + assert router_identity({"litellm_params": {"model": "auto_router/complexity_router"}}) is None + + +class TestHeuristicV1Scope: + @pytest.mark.parametrize( + "config,in_scope", + [ + ({"tiers": _TIERS}, True), + ({"classifier_type": "heuristic", "tiers": _TIERS}, True), + ( + { + "classifier_type": "heuristic_first", + "heuristic_first_max_tier": "MEDIUM", + "classifier_llm_config": {"model": "judge"}, + "tiers": _TIERS, + }, + True, + ), + ( + { + "classifier_type": "hybrid", + "hybrid_boundary_margin": 0.05, + "classifier_llm_config": {"model": "judge"}, + "tiers": _TIERS, + }, + True, + ), + ({"classifier_type": "heuristic_v2", "tiers": _TIERS}, False), + ({"classifier_type": "llm", "classifier_llm_config": {"model": "judge"}, "tiers": _TIERS}, False), + ], + ) + def test_only_v1_scoring_classifiers_are_fingerprinted(self, config: Mapping[str, object], in_scope: bool) -> None: + assert (heuristic_v1_router_fingerprint(_router("r", config)) is not None) is in_scope + + def test_plain_deployments_are_ignored(self) -> None: + assert heuristic_v1_router_fingerprint(_router("gpt", None, model="openai/gpt-4o")) is None + + +class TestQuota: + def test_snapshot_records_every_v1_router_even_at_defaults(self) -> None: + baselines = snapshot_tuning_baselines([_router("a", {"tiers": _TIERS}), _router("b", {})]) + assert set(baselines) == {router_identity(_router("a", {})), router_identity(_router("b", {}))} + assert baselines[router_identity(_router("b", {}))] == DEFAULT_TUNING_FINGERPRINT + + def test_unchanged_snapshot_is_never_mutable(self) -> None: + rows = [_router("a", {"tiers": _TIERS}), _router("b", {"tiers": _ALT_TIERS})] + baselines = snapshot_tuning_baselines(rows) + assert mutable_tuned_identities(rows, baselines) == frozenset() + + def test_router_added_after_snapshot_is_mutable_only_when_tuned(self) -> None: + baselines = snapshot_tuning_baselines([_router("a", {"tiers": _TIERS})]) + assert mutable_tuned_identities([_router("new", {})], baselines) == frozenset() + assert mutable_tuned_identities([_router("new", {"tiers": _TIERS})], baselines) == { + router_identity(_router("new", {})) + } + + def test_quota_matrix(self) -> None: + legacy_a = _router("a", {"tiers": _TIERS}) + legacy_b = _router("b", {"tiers": _ALT_TIERS}) + baselines = snapshot_tuning_baselines([legacy_a, legacy_b]) + edited_a = _router("a", {"tiers": _TIERS, "dimension_weights": {"codePresence": 0.9}}) + edited_b = _router("b", {"tiers": _TIERS}) + new_c = _router("c", {"tiers": _TIERS}) + + assert tuning_quota_violation(candidate=edited_a, others=[legacy_b], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited_a, others=[edited_a, legacy_b], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=legacy_a, others=[edited_b], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited_b, others=[edited_a], baselines=baselines, limit=1) is not None + assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=1) is not None + assert tuning_quota_violation(candidate=new_c, others=[edited_a], baselines=baselines, limit=None) is None + assert tuning_quota_violation(candidate=legacy_a, others=[edited_a, edited_b], baselines=baselines, limit=1) is None + + def test_reverting_to_baseline_frees_the_quota(self) -> None: + legacy_a = _router("a", {"tiers": _TIERS}) + legacy_b = _router("b", {"tiers": _ALT_TIERS}) + baselines = snapshot_tuning_baselines([legacy_a, legacy_b]) + edited_b = _router("b", {"tiers": _TIERS}) + assert tuning_quota_violation(candidate=edited_b, others=[legacy_a], baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited_b, others=[legacy_a, edited_b], baselines=baselines, limit=1) is None + + def test_violation_message_names_the_limit_and_remedy(self) -> None: + message = tuning_limit_violation(held=2, limit=1) + assert message is not None + assert "At most 1 auto-router(s)" in message + assert "revert the other changed router to its baseline" in message + assert tuning_limit_violation(held=1, limit=1) is None + assert tuning_limit_violation(held=5, limit=None) is None From e8f311429ea9afa09195bed21f078bbd50dd791e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:42:15 -0700 Subject: [PATCH 43/63] fix(cost-map): stop advertising reasoning_effort max on azure_ai/gpt-6-astra Foundry rejects reasoning_effort max on the gpt-6-astra deployment with a 400 that names none, low, medium, high, and xhigh as the supported values, so the card no longer lists max. The request path never gated max (only xhigh is opt-in), so this only changes /model_group/info and router capability gating. The azure/ twin stays as is because it was not verified on an Azure OpenAI host --- .../model_prices_and_context_window_backup.json | 2 +- model_prices_and_context_window.json | 2 +- .../test_reasoning_effort_capability.py | 14 +++++++++++++- 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5b4652d38c7..b659c3b65e5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3518,7 +3518,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5b4652d38c7..b659c3b65e5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3518,7 +3518,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index 3e1f26b6e1c..fa3a6dcd95a 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -394,7 +394,6 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: [ ("azure/gpt-6-astra", "azure"), ("azure/us/gpt-6-astra", "azure"), - ("azure_ai/gpt-6-astra", "azure_ai"), ], ) def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): @@ -413,3 +412,16 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: "xhigh", "max", ) + + def test_a_foundry_azure_ai_deployment_advertises_none_but_not_max(self, local_model_cost_map): + from litellm.utils import _get_model_info_helper + + model_info = dict(_get_model_info_helper(model="azure_ai/gpt-6-astra", custom_llm_provider="azure_ai")) + + assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( + "none", + "low", + "medium", + "high", + "xhigh", + ) From 79d47788d9eea16c0038fc7eed10af37b5f8889b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:56:39 -0700 Subject: [PATCH 44/63] fix(anthropic): stream the refusal text on bridged /v1/messages calls Both bridges opened an empty text block on a refused streaming turn and closed it without a single delta, so a client replaying that assistant turn got HTTP 400 "text content blocks must be non-empty" from Anthropic. The safeguard-refusal fallback that motivated withholding the text only runs on the awaited non-streaming response, so nothing needed it withheld Move the refusal readers into the shared messages/utils helpers so the adapters stop reaching into each other's private statics, which is also what put reportPrivateUsage over its budget --- .../adapters/streaming_iterator.py | 16 +++-- .../adapters/transformation.py | 36 ++++------ .../messages/utils.py | 70 ++++++++++++++++++- .../responses_adapters/streaming_iterator.py | 32 ++++++--- .../responses_adapters/transformation.py | 50 ++----------- .../test_streaming_iterator_first_delta.py | 8 +-- ...t_responses_adapters_streaming_iterator.py | 6 +- 7 files changed, 126 insertions(+), 92 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index adf77dc3dd5..c6d744772a5 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -1006,6 +1006,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): delta: Final = cast(Mapping[str, object], processed_chunk["delta"]) # cast-ok: keys checked before use if delta.get("stop_reason") == "max_tokens": return processed_chunk + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + refusal_stop_details, + ) + return cast( # cast-ok: rebuilt dict matches the message_delta TypedDict shape for this branch ContentBlockDelta | MessageBlockDelta, { # mutable-ok: fresh translation payload; never mutated after construction @@ -1013,11 +1017,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): "delta": { # mutable-ok: fresh message_delta payload; never mutated after construction **delta, "stop_reason": "refusal", - "stop_details": { # mutable-ok: fresh stop_details payload; never mutated after construction - "type": "refusal", - "category": None, - "explanation": "".join(self._refusal_text_parts), - }, + "stop_details": refusal_stop_details("".join(self._refusal_text_parts)), }, }, ) @@ -1098,9 +1098,13 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): - Different content types in the response - Specific markers in the content """ + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + openai_chat_refusal_text, + ) + from .transformation import LiteLLMAnthropicMessagesAdapter - refusal_text: Final = LiteLLMAnthropicMessagesAdapter._refusal_text(chunk.choices[0].delta) + refusal_text: Final = openai_chat_refusal_text(chunk.choices[0].delta) if refusal_text is not None: self._refusal_text_parts.append(refusal_text) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 215e3ded908..e3bad54d599 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -117,6 +117,10 @@ from litellm.llms.anthropic.common_utils import ( from litellm.llms.anthropic.experimental_pass_through.context_management import ( PolyfillResult, ) +from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + openai_chat_refusal_text, + refusal_stop_details, +) from litellm.types.llms.anthropic import ( ANTHROPIC_HOSTED_TOOLS, AllAnthropicPassThroughMessageValues, @@ -308,17 +312,6 @@ class LiteLLMAnthropicMessagesAdapter: def __init__(self): pass - @staticmethod - def _refusal_text(message_or_delta: object) -> str | None: - refusal: Final = getattr(message_or_delta, "refusal", None) - if isinstance(refusal, str): - return refusal - provider_specific_fields: Final = getattr(message_or_delta, "provider_specific_fields", None) - if isinstance(provider_specific_fields, Mapping): - provider_refusal: Final = provider_specific_fields.get("refusal") - return provider_refusal if isinstance(provider_refusal, str) else None - return None - ### FOR [BETA] `/v1/messages` endpoint support def _extract_signature_from_tool_call(self, tool_call: object) -> str | None: @@ -1325,7 +1318,7 @@ class LiteLLMAnthropicMessagesAdapter: new_content.append( AnthropicResponseContentBlockText(type="text", text=choice.message.content).model_dump() ) - if (refusal_text := self._refusal_text(choice.message)) is not None: + if (refusal_text := openai_chat_refusal_text(choice.message)) is not None: new_content.append(AnthropicResponseContentBlockText(type="text", text=refusal_text).model_dump()) # Handle tool calls (in parallel to text content) if choice.message.tool_calls is not None and len(choice.message.tool_calls) > 0: @@ -1486,7 +1479,7 @@ class LiteLLMAnthropicMessagesAdapter: tool_name_mapping=tool_name_mapping, ) refusal_text: Final = next( - (text for choice in response.choices if (text := self._refusal_text(choice.message)) is not None), + (text for choice in response.choices if (text := openai_chat_refusal_text(choice.message)) is not None), None, ) @@ -1523,15 +1516,7 @@ class LiteLLMAnthropicMessagesAdapter: usage=anthropic_usage, content=anthropic_content, stop_reason=anthropic_finish_reason, - stop_details=( - { # mutable-ok: fresh refusal stop_details payload built per response - "type": "refusal", - "category": None, - "explanation": refusal_text, - } - if anthropic_finish_reason == "refusal" - else None - ), + stop_details=(refusal_stop_details(refusal_text) if anthropic_finish_reason == "refusal" else None), ) applied_edits: Final = polyfill_result.applied_edits_for_response() if polyfill_result else None @@ -1572,7 +1557,7 @@ class LiteLLMAnthropicMessagesAdapter: "signature": thought_sig, } return "tool_use", cast("ContentBlockContentBlockDict", tool_block) - elif (choice.delta.content is not None and len(choice.delta.content) > 0) or self._refusal_text( + elif (choice.delta.content is not None and len(choice.delta.content) > 0) or openai_chat_refusal_text( choice.delta ) is not None: return "text", TextBlock(type="text", text="") @@ -1646,7 +1631,10 @@ class LiteLLMAnthropicMessagesAdapter: elif reasoning_content: return "thinking_delta", ContentThinkingBlockDelta(type="thinking_delta", thinking=reasoning_content) else: - return "text_delta", ContentTextBlockDelta(type="text_delta", text=text) + refusal_text: Final = "".join( + refusal for choice in choices if (refusal := openai_chat_refusal_text(choice.delta)) is not None + ) + return "text_delta", ContentTextBlockDelta(type="text_delta", text=text + refusal_text) def translate_streaming_openai_response_to_anthropic( self, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 242300c7b6d..7545dff1408 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -1,8 +1,11 @@ -from collections.abc import Mapping +from collections.abc import Iterable, Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints -from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams +from litellm.types.llms.anthropic import ( + AnthropicMessagesRequestOptionalParams, + AnthropicStopDetails, +) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -25,6 +28,69 @@ def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | return stop_details if isinstance(stop_details, dict) else None +def refusal_stop_details(explanation: str | None) -> AnthropicStopDetails: + """The ``stop_details`` object accompanying a translated ``stop_reason: "refusal"``.""" + return AnthropicStopDetails(type="refusal", category=None, explanation=explanation) + + +def _mapping_field(container: object, key: str) -> object | None: + """One key of a raw provider payload, or None when the payload is not a mapping.""" + if not isinstance(container, Mapping): + return None + return cast(Mapping[str, object], container).get(key) # cast-ok: raw payload, callers re-check every value + + +def _mapping_str_field(container: object, key: str) -> str | None: + value: Final = _mapping_field(container, key) + return value if isinstance(value, str) and value else None + + +def openai_chat_refusal_text(message_or_delta: object) -> str | None: + """ + Refusal text carried by an OpenAI Chat Completions message or streaming delta, + read from ``refusal`` or from the ``provider_specific_fields`` LiteLLM parks it + in, or None when the turn is not a refusal. + """ + refusal: Final = getattr(message_or_delta, "refusal", None) + if isinstance(refusal, str) and refusal: + return refusal + return _mapping_str_field(getattr(message_or_delta, "provider_specific_fields", None), "refusal") + + +def _responses_message_refusal_text(item: object) -> str | None: + from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal + + if isinstance(item, ResponseOutputMessage): + return next( + (part.refusal for part in item.content if isinstance(part, ResponseOutputRefusal) and part.refusal), + None, + ) + raw_parts: Final = _mapping_field(item, "content") + if _mapping_str_field(item, "type") != "message" or not isinstance(raw_parts, Sequence): + return None + return next( + ( + refusal + for part in cast(Sequence[object], raw_parts) # cast-ok: members re-validated below + if _mapping_str_field(part, "type") == "refusal" + and isinstance(refusal := _mapping_str_field(part, "refusal"), str) + ), + None, + ) + + +def responses_output_refusal_text(output: Iterable[object]) -> str | None: + """ + Refusal text carried by an OpenAI Responses ``output`` list, in typed + (``ResponseOutputRefusal``) or raw-dictionary shape, or None when none of the + output messages refused. + """ + return next( + (text for item in output if (text := _responses_message_refusal_text(item)) is not None), + None, + ) + + def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError": """The exception a safeguard-refused Anthropic response converts into so the content-policy fallback chain can re-dispatch it.""" diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 06ffcd26d40..5bd662e0f94 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -9,6 +9,10 @@ from typing import TYPE_CHECKING, Any, Final from litellm import verbose_logger from litellm._uuid import uuid +from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + refusal_stop_details, + responses_output_refusal_text, +) from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from .transformation import LiteLLMAnthropicToResponsesAPIAdapter @@ -136,8 +140,20 @@ class AnthropicResponsesStreamWrapper: if event_type == "response.refusal.delta": delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") - if isinstance(delta, str): - self._refusal_text_parts.append(delta) + if not isinstance(delta, str) or not delta: + return + self._refusal_text_parts.append(delta) + item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) + block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index + if block_idx < 0: + block_idx = self._open_block(item_id, {"type": "text", "text": ""}) + self._chunk_queue.append( + { + "type": "content_block_delta", + "index": block_idx, + "delta": {"type": "text_delta", "text": delta}, + } + ) return # ---- text delta ---- @@ -225,9 +241,7 @@ class AnthropicResponsesStreamWrapper: event.get("response") if isinstance(event, dict) else None ) output: Final = (getattr(response_obj, "output", None) or ()) if response_obj is not None else () - refusal_text: Final = LiteLLMAnthropicToResponsesAPIAdapter._refusal_text_from_output(output) or ( - "".join(self._refusal_text_parts) or None - ) + refusal_text: Final = responses_output_refusal_text(output) or ("".join(self._refusal_text_parts) or None) status: Final = getattr(response_obj, "status", None) if response_obj is not None else None has_tool_call: Final = any( getattr(item, "type", None) == "function_call" @@ -255,12 +269,8 @@ class AnthropicResponsesStreamWrapper: "stop_reason": stop_reason, "stop_sequence": None, **( - { # mutable-ok: fresh refusal stop_details payload built per chunk - "stop_details": { # mutable-ok: fresh refusal stop_details payload built per chunk - "type": "refusal", - "category": None, - "explanation": refusal_text, - } + { # mutable-ok: fresh message_delta stop_details entry built per chunk + "stop_details": refusal_stop_details(refusal_text) } if stop_reason == "refusal" else {} # mutable-ok: empty spread placeholder for non-refusal stop diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py index 318065148b8..9fb44127f5b 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py @@ -6,7 +6,7 @@ path used for OpenAI and Azure models. """ import json -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Iterable, Mapping from itertools import groupby from typing import Any, Final, cast @@ -19,6 +19,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking_budget, ) +from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + refusal_stop_details, + responses_output_refusal_text, +) from litellm.llms.anthropic.experimental_pass_through.utils import ( is_reasoning_auto_summary_enabled, prompt_cache_key_from_user_id, @@ -69,38 +73,6 @@ class LiteLLMAnthropicToResponsesAPIAdapter: chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_usage) return LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage(chat_usage) - @staticmethod - def _refusal_text_from_output(output: Iterable[object]) -> str | None: - from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal - - def refusal_text_from_item(item: object) -> str | None: - if isinstance(item, ResponseOutputMessage): - return next( - (part.refusal for part in item.content if isinstance(part, ResponseOutputRefusal)), - None, - ) - if not isinstance(item, Mapping): - return None - item_mapping: Final = cast(Mapping[str, object], item) # cast-ok: keys re-checked before use - raw_parts: Final = item_mapping.get("content") - if item_mapping.get("type") != "message" or not isinstance(raw_parts, Sequence): - return None - for part in cast(Sequence[object], raw_parts): # cast-ok: members re-validated below - if not isinstance(part, Mapping): - continue - part_mapping = cast(Mapping[str, object], part) # cast-ok: keys re-checked before use - if part_mapping.get("type") != "refusal": - continue - refusal = part_mapping.get("refusal") - if isinstance(refusal, str): - return refusal - return None - - return next( - (text for item in output if (text := refusal_text_from_item(item)) is not None), - None, - ) - # ------------------------------------------------------------------ # # Request translation: Anthropic -> Responses API # # ------------------------------------------------------------------ # @@ -656,7 +628,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter: content: Final[list[dict[str, object]]] = [] stop_reason: AnthropicFinishReason = "end_turn" - refusal_text: Final = self._refusal_text_from_output( + refusal_text: Final = responses_output_refusal_text( cast(Iterable[object], response.output) # cast-ok: output items re-validated per item ) @@ -747,13 +719,5 @@ class LiteLLMAnthropicToResponsesAPIAdapter: usage=anthropic_usage, content=content, stop_reason=stop_reason, - stop_details=( - { # mutable-ok: fresh refusal stop_details payload built per response - "type": "refusal", - "category": None, - "explanation": refusal_text, - } - if stop_reason == "refusal" - else None - ), + stop_details=(refusal_stop_details(refusal_text) if stop_reason == "refusal" else None), ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index e7381986ec2..4359ace1870 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -108,7 +108,7 @@ def _text_deltas(events: List[dict]) -> List[str]: ] -def test_streaming_chat_refusal_emits_only_refusal_stop_details(): +def test_streaming_chat_refusal_emits_refusal_text_and_stop_details(): chunks = [ _make_chunk(Delta(content=None, refusal="I cannot fulfill this request.")), _make_chunk(Delta(content=None), finish_reason="stop"), @@ -117,7 +117,7 @@ def test_streaming_chat_refusal_emits_only_refusal_stop_details(): events = _drain_sync(wrapper) - assert _text_deltas(events) == [] + assert _text_deltas(events) == ["I cannot fulfill this request."] message_delta = next(event for event in events if event["type"] == "message_delta") assert message_delta["delta"] == { "stop_reason": "refusal", @@ -130,7 +130,7 @@ def test_streaming_chat_refusal_emits_only_refusal_stop_details(): @pytest.mark.asyncio -async def test_streaming_chat_refusal_emits_only_refusal_stop_details_async(): +async def test_streaming_chat_refusal_emits_refusal_text_and_stop_details_async(): chunks = [ _make_chunk(Delta(content=None, refusal="I cannot fulfill this request.")), _make_chunk(Delta(content=None), finish_reason="stop"), @@ -139,7 +139,7 @@ async def test_streaming_chat_refusal_emits_only_refusal_stop_details_async(): events = await _drain_async(wrapper) - assert _text_deltas(events) == [] + assert _text_deltas(events) == ["I cannot fulfill this request."] message_delta = next(event for event in events if event["type"] == "message_delta") assert message_delta["delta"]["stop_reason"] == "refusal" assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py index d8f16dde3e7..d9df5df426a 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -328,7 +328,7 @@ class TestResponseCompletedUsage: class TestRefusalStreamEvents: - def test_refusal_event_sequence_emits_only_stop_details(self): + def test_refusal_event_sequence_emits_refusal_text_and_stop_details(self): response = SimpleNamespace( status="completed", output=[{"type": "message", "content": [{"type": "refusal", "refusal": "I cannot fulfill this."}]}], @@ -346,11 +346,13 @@ class TestRefusalStreamEvents: assert [chunk["type"] for chunk in chunks] == [ "message_start", "content_block_start", + "content_block_delta", "content_block_stop", "message_delta", "message_stop", ] - assert chunks[3]["delta"] == { + assert chunks[2]["delta"] == {"type": "text_delta", "text": "I cannot fulfill this."} + assert chunks[4]["delta"] == { "stop_reason": "refusal", "stop_sequence": None, "stop_details": { From 116f88b023996a52a46f02b2dd8bffc8c15c9246 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:03:50 -0700 Subject: [PATCH 45/63] fix(e2e-changed): keep the gate off suites the stack cannot run The selector picked up two suites that can never pass in this stack, so editing either one turned the check permanently red: the presidio masking suite calls pytest.fail without an analyzer and anonymizer that up.sh never starts, and the pipecat audio suite skips itself at import time unless the NLTK punkt_tab data is present, which nothing installs. tests/e2e/coverage_registry/test_collector.py had the same problem for a different reason. Its nested pytest.main autoloads pytest-retry from the ci group the workflow installs and dies with "INTERNALERROR: no option named 'filtered_exceptions'", so the collect-only pass now disables that plugin. The plugin's entry point is pytest-retry, not retry, so the same one-word fix lands on mutmut's pytest_add_cli_args, where "-p no:retry" was disabling nothing. Two smaller holes in the harness: a canary argument the shell never expanded used to select nothing and let the gate pass green, and a secret that cannot be represented in both bash and dotenv was rejected without naming the key. --- .github/e2e-stack/secrets_to_env.py | 13 +++-- .github/e2e-stack/select_tests.py | 23 +++++++-- pyproject.toml | 2 +- .../test_e2e_changed_gate.py | 49 ++++++++++++++++++- tests/e2e/CONTRIBUTING.md | 4 +- tests/e2e/coverage_registry/collector.py | 2 + 6 files changed, 80 insertions(+), 13 deletions(-) diff --git a/.github/e2e-stack/secrets_to_env.py b/.github/e2e-stack/secrets_to_env.py index e7870a32e42..7913b25918b 100644 --- a/.github/e2e-stack/secrets_to_env.py +++ b/.github/e2e-stack/secrets_to_env.py @@ -20,10 +20,15 @@ def main() -> int: except (ValidationError, UnicodeError): _ = sys.stderr.write("expected a JSON object containing string environment values\n") return 1 - if any( - ENV_NAME.fullmatch(key) is None or any(char in value for char in "'\n\r\0") for key, value in secrets.items() - ): - _ = sys.stderr.write("environment names or values cannot be represented in both bash and dotenv\n") + unusable: Final = tuple( + key + for key, value in secrets.items() + if ENV_NAME.fullmatch(key) is None or any(char in value for char in "'\n\r\0") + ) + if unusable: + _ = sys.stderr.write( + f"these names or values cannot be represented in both bash and dotenv: {' '.join(sorted(unusable))}\n" + ) return 1 for value in secrets.values(): if len(value) >= MIN_MASKED_LENGTH: diff --git a/.github/e2e-stack/select_tests.py b/.github/e2e-stack/select_tests.py index cbe582b5b42..a62358f81ff 100644 --- a/.github/e2e-stack/select_tests.py +++ b/.github/e2e-stack/select_tests.py @@ -3,8 +3,11 @@ import sys from typing import Final SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$") -OWN_LANE: Final = re.compile( - r"^tests/e2e/(ui|claude_code|load)/|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" +UNSUPPORTED: Final = re.compile( + r"^tests/e2e/(ui|claude_code|load)/" + r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$" + r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$" + r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$" ) HARNESS: Final = re.compile( r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$" @@ -12,18 +15,28 @@ HARNESS: Final = re.compile( r"|^\.github/e2e-stack/" r"|^\.github/workflows/test-e2e-changed\.yml$" ) +UNEXPANDED: Final = re.compile(r"[*?\[]") + + +def is_selectable(path: str) -> bool: + return SELECTABLE.match(path) is not None and UNSUPPORTED.match(path) is None def select(changed: tuple[str, ...], canary: tuple[str, ...]) -> tuple[str, ...]: - direct: Final = frozenset(path for path in changed if SELECTABLE.match(path) and not OWN_LANE.match(path)) + direct: Final = frozenset(path for path in changed if is_selectable(path)) harness_changed: Final = any(HARNESS.match(path) for path in changed) - canary_tests: Final = frozenset(path for path in canary if harness_changed and SELECTABLE.match(path)) + canary_tests: Final = frozenset(path for path in canary if harness_changed and is_selectable(path)) return tuple(sorted(direct | canary_tests)) def main() -> int: + canary: Final = tuple(sys.argv[1:]) + unexpanded: Final = tuple(path for path in canary if UNEXPANDED.search(path)) + if unexpanded: + _ = sys.stderr.write(f"the canary paths reached the selector unexpanded: {' '.join(unexpanded)}\n") + return 1 changed: Final = tuple(line.strip() for line in sys.stdin if line.strip()) - _ = sys.stdout.write(" ".join(select(changed, tuple(sys.argv[1:]))) + "\n") + _ = sys.stdout.write(" ".join(select(changed, canary)) + "\n") return 0 diff --git a/pyproject.toml b/pyproject.toml index b889a3a0e60..f4f238dd4b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -388,7 +388,7 @@ mutate_only_covered_lines = true # registered hash algorithm". Nothing to do with mutation coverage, and one # erroring test is enough to end the stats phase before any mutant runs. pytest_add_cli_args = [ - "-p", "no:retry", + "-p", "no:pytest-retry", "-p", "no:rerunfailures", "-p", "no:xdist", "--ignore=tests/test_litellm/proxy/management_endpoints/test_saml_sso.py", diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index f7c4a5e2527..d2b842364a6 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -113,6 +113,16 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: (("tests/e2e/logging/test_datadog_e2e.py", "litellm/router.py"), ("tests/e2e/logging/test_datadog_e2e.py",)), (("tests/e2e/ui/test_keys.py", "tests/e2e/claude_code/test_cli.py", "tests/e2e/load/test_burst.py"), ()), (("tests/e2e/batches/test_managed_files_enforcement_e2e.py",), ()), + (("tests/e2e/guardrails/test_presidio_masking_e2e.py",), ()), + (("tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py",), ()), + ( + ("tests/e2e/llm_translation/realtime/test_realtime_e2e.py",), + ("tests/e2e/llm_translation/realtime/test_realtime_e2e.py",), + ), + ( + ("tests/e2e/guardrails/test_bedrock_guardrail_e2e.py",), + ("tests/e2e/guardrails/test_bedrock_guardrail_e2e.py",), + ), (("tests/e2e/logging/helpers.py", "docs/my-website/docs/index.md", "tests/e2e/CLAUDE.md"), ()), ( ("tests/e2e/logging/test_datadog_e2e.py", "tests/e2e/logging/test_datadog_e2e.py"), @@ -120,7 +130,7 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: ), ), ) -def test_changed_suite_files_are_selected_outside_the_own_lane( +def test_changed_suite_files_are_selected_unless_the_stack_cannot_run_them( changed: tuple[str, ...], expected: tuple[str, ...] ) -> None: assert select_tests(changed) == expected @@ -154,3 +164,40 @@ def test_the_canary_joins_directly_selected_files_in_sorted_order() -> None: def test_a_harness_unit_test_change_runs_itself_and_the_canary() -> None: assert select_tests(("tests/e2e/test_proxy_client.py",)) == (*CANARY, "tests/e2e/test_proxy_client.py") + + +def test_a_canary_argument_the_shell_never_expanded_fails_the_selector() -> None: + result: Final = subprocess.run( + [sys.executable, str(SELECT_TESTS), "tests/e2e/access_control/test_*.py"], + input="tests/e2e/proxy_client.py\n", + capture_output=True, + text=True, + ) + + assert result.returncode == 1 + assert "tests/e2e/access_control/test_*.py" in result.stderr + assert result.stdout == "" + + +@pytest.mark.parametrize( + ("secrets", "offender", "unprintable"), + ( + ('{"AWS_ACCESS_KEY_ID": "AKIAEXAMPLE", "BAD-NAME": "shibboleth"}', "BAD-NAME", "shibboleth"), + ("""{"AWS_SECRET_ACCESS_KEY": "quote'shibboleth"}""", "AWS_SECRET_ACCESS_KEY", "shibboleth"), + ('{"DD_API_KEY": "line\\nshibboleth"}', "DD_API_KEY", "shibboleth"), + ), +) +def test_an_unusable_secret_is_named_without_printing_its_value( + tmp_path: Path, secrets: str, offender: str, unprintable: str +) -> None: + env_path: Final = tmp_path / ".env" + + result: Final = subprocess.run( + [sys.executable, str(SECRETS_TO_ENV), str(env_path)], input=secrets, capture_output=True, text=True + ) + + assert result.returncode == 1 + assert offender in result.stderr + assert unprintable not in result.stderr + assert result.stdout == "" + assert not env_path.exists() diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 106cfa31f0e..b83300cc8d5 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -56,9 +56,9 @@ A couple of logging destinations are configured on the proxy rather than by the ### The pull request check -Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories and `batches/test_managed_files_enforcement_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack +Every same-repository PR that adds, modifies, or renames a `tests/e2e/**/test_*.py` file runs those changed files three times. A change to the harness itself, meaning a root-level `tests/e2e/*.py` file or `pytest.ini`, `tests/e2e/gateway/`, `.github/e2e-stack/`, or the workflow, also runs the `access_control` suite as a canary, because those files have no test of their own that exercises the stack. `.github/e2e-stack/select_tests.py` applies both rules. The stack config at `tests/e2e/gateway/stage_mirror_ci_config.yml` must declare every model the selected suites use; a missing one shows up as a failed test id in the public log. The suite's own single rerun for network errors and 5xx responses (see `pytest.ini`) applies on every pass, so a transport blip does not fail the check while a race inside a test still does. The stage-mirror stack has a control-plane backend, two gateways behind nginx, Postgres, Jaeger, and TLS cluster-mode Valkey. The stack exports every gateway address in `LITELLM_PROXY_REPLICA_URLS`, so model registration waits until each gateway lists the new model rather than whichever one the load balancer answered from. Documentation, deleted-file, and application-only changes do not start the stack or request environment approval. The `ui/`, `claude_code/`, and `load/` directories, `batches/test_managed_files_enforcement_e2e.py`, `llm_translation/realtime/test_realtime_pipecat_audio_e2e.py`, and `guardrails/test_presidio_masking_e2e.py` remain outside this check because they use separate tooling or need a differently configured stack: the pipecat audio suite skips itself at import time unless the NLTK `punkt_tab` data is installed, and the presidio suite fails without the analyzer and anonymizer services this stack does not start -Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch +Every selected file must execute at least one passing test in each pass, and any test failure, collection error, or entirely skipped or deselected file fails the check. A file whose tests are all marked skip therefore cannot pass this check, so unskip at least one of them, or add the file to `UNSUPPORTED` in `select_tests.py` with the reason, before changing one. A failed pass stops the run. The public log prints pytest's one-line summary for each pass, including the rerun count, and names each failed or errored test as `classname::name`, so a retried network error or a failing test is visible without the raw output. The final `e2e-changed-tests` job succeeds only when no supported test files changed or the approved run completed all three passes. Fork PRs with selected tests fail this gate until a maintainer brings the reviewed change onto a same-repository branch Repository admins must require the `e2e-changed-tests` status check for merging and configure the `e2e-changed` environment with required reviewers, self-review disabled, and admin bypass disabled. Each push cancels the previous run; a new run that selects tests needs a fresh approval. Reviewers must inspect the entire executable PR diff, including application code, dependencies, tests, and workflow helpers, before approving the exact revision. Approved code executes with provider credentials, so environment approval is a trust decision about that code diff --git a/tests/e2e/coverage_registry/collector.py b/tests/e2e/coverage_registry/collector.py index e20f7884f55..c5d09820fe7 100644 --- a/tests/e2e/coverage_registry/collector.py +++ b/tests/e2e/coverage_registry/collector.py @@ -105,6 +105,8 @@ def collect_markers(e2e_dir: Path = E2E_DIR) -> CollectedMarkers: "--continue-on-collection-errors", "-p", "no:cacheprovider", + "-p", + "no:pytest-retry", str(e2e_dir), ], plugins=[sink], From 061c25b5cac15212ded745c51e3299aa4d43f056 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:25:03 -0700 Subject: [PATCH 46/63] fix(spend): let a batch's charge survive an older proxy's $0 poll row A proxy running the old code wrote _batch_cost at $0 every time it polled a batch that was still running, so after an upgrade the claim found that row and read it as proof the batch had already been charged. Only a row that recorded a charge counts now, which leaves those $0 rows, and any row a client planted under the batch id, to be charged over disable_spend_logs skipped the claim entirely, so under that setting every retrieve of a finished batch charged again. The claim now runs either way and writes the one row per batch that makes the charge exactly once, while the per-request logs stay off --- litellm/proxy/db/db_spend_update_writer.py | 42 +++++++------ .../proxy/db/test_db_spend_update_writer.py | 60 ++++++++++++++++--- 2 files changed, 78 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 3fad351224b..48312c025dd 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -271,9 +271,12 @@ class DBSpendUpdateWriter: if team_id is not None and team_id != "": payload["team_id"] = team_id + if not await self._record_spend_log( + payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ): + return False + if disable_spend_logs is False: - if not await self._record_spend_log(payload=payload, prisma_client=prisma_client): - return False await self._enqueue_tool_usage_transaction( payload=payload, completion_response=completion_response, @@ -328,19 +331,23 @@ class DBSpendUpdateWriter: ) return True - async def _record_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None") -> bool: - if prisma_client is None or not _is_batch_cost_row(payload): + async def _record_spend_log( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool + ) -> bool: + if prisma_client is not None and _is_batch_cost_row(payload): + return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + if disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) - return True - return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + return True async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: """Write the batch's cost row now, or learn that another retrieve already did. Every retrieve of one batch shares this row, so the insert that lands first owns the charge and every later one finds the row and charges nothing (LIT-7048). Only - a row a successful retrieve wrote counts: a failed retrieve, or any request whose - client picked the batch id as its call id, cannot take the charge away. + a row that recorded a charge counts: a failed retrieve, a request whose client + picked the batch id as its call id, and the $0 row an older proxy left behind + while the batch was still running all leave the charge to be made. """ from litellm.repositories.table_repositories import SpendLogsRepository @@ -362,17 +369,18 @@ class DBSpendUpdateWriter: ) await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) return True - if ( - existing is not None - and existing.call_type == CallTypes.aretrieve_batch.value - and existing.status == "success" - ): + if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": + verbose_proxy_logger.warning( + "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + request_id, + getattr(existing, "call_type", None), + ) + return True + if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False - verbose_proxy_logger.warning( - "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", - request_id, - getattr(existing, "call_type", None), + verbose_proxy_logger.debug( + "Spend row %s charged nothing for this batch, so this retrieve charges it", request_id ) return True diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 500a0e7bb06..41f2be08545 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2977,10 +2977,12 @@ def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock: return prisma -async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict) -> bool: +async def _update_database_with( + db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict, disable_spend_logs: bool = False +) -> bool: with ( patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam - "litellm.proxy.proxy_server.disable_spend_logs", False + "litellm.proxy.proxy_server.disable_spend_logs", disable_spend_logs ), patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam "litellm.proxy.proxy_server.prisma_client", prisma @@ -3014,14 +3016,16 @@ async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMoc ("inserted", "existing", "charged"), [ (1, None, True), - (0, SimpleNamespace(call_type="aretrieve_batch", status="success"), False), - (0, SimpleNamespace(call_type="aretrieve_batch", status="failure"), True), - (0, SimpleNamespace(call_type="aembedding", status="success"), True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0), True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="failure", spend=0.0), True), + (0, SimpleNamespace(call_type="aembedding", status="success", spend=0.25), True), (0, None, True), ], ids=[ "first_retrieve_owns_the_row", "another_retrieve_already_charged", + "an_older_proxy_left_a_zero_row_while_the_batch_ran", "failed_retrieve_holds_the_row", "client_chosen_call_id_holds_the_row", "row_gone_between_insert_and_lookup", @@ -3033,8 +3037,9 @@ async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote """ Every retrieve of one batch shares one spend row, so the insert that lands first is the charge and every later retrieve must leave the counters alone (LIT-7048). A row - written by anything but a successful retrieve, say a request whose client picked the - batch id as its call id, must not be able to take the charge away. + that recorded no charge must not be able to take the charge away: neither one a + client planted under the batch id, nor the $0 row a pre-upgrade proxy wrote every + time it polled the batch while it was still running. """ db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() @@ -3049,6 +3054,47 @@ async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote assert db_writer._batch_database_updates.await_count == (1 if charged else 0) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inserted", "existing", "charged"), + [ + (1, None, True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25), False), + ], + ids=["first_retrieve_owns_the_row", "another_retrieve_already_charged"], +) +async def test_update_database_charges_a_batch_once_even_with_spend_logs_disabled( + inserted: int, existing: object, charged: bool +): + """ + disable_spend_logs drops the per-request logs, not the batch's charge, so the one row + that makes a batch chargeable exactly once is still written and still read back. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(inserted, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), True) is charged + + assert prisma.db.litellm_spendlogs.create_many.await_count == 1 + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_writes_no_ordinary_spend_row_with_spend_logs_disabled(): + """The batch carve-out above stays a carve-out: every other row still goes unwritten.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + payload = {**_batch_cost_payload(), "call_type": "acompletion"} + + assert await _update_database_with(db_writer, prisma, payload, True) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert prisma.spend_log_transactions == [] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): """An unreachable DB must not drop the batch's only spend row, nor its charge.""" From 82edb9e901c7b87a99a7b5c318b3db333ca6165b Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sun, 6 Sep 2026 00:09:27 -0500 Subject: [PATCH 47/63] fix(vertex): preserve Lyria pricing fallback and audio MIME --- litellm/cost_calculator.py | 14 +++---- litellm/llms/vertex_ai/common_utils.py | 22 ++++++++++- .../text_to_speech/transformation.py | 4 +- .../vertex_passthrough_logging_handler.py | 6 +-- ...test_vertex_passthrough_logging_handler.py | 27 +++++++++++--- .../text_to_speech/test_transformation.py | 17 +++++++++ tests/test_litellm/test_cost_calculator.py | 37 +++++++++++++++++-- 7 files changed, 103 insertions(+), 24 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 5dcbb1d5f37..fa8e1a5264c 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -495,17 +495,15 @@ def cost_per_token( # see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models if call_type == "speech" or call_type == "aspeech": + if custom_llm_provider in ("vertex_ai", "vertex_ai_beta"): + from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_generation_cost + + lyria_generation_cost: Final = get_vertex_ai_lyria_generation_cost(model_without_prefix) + if lyria_generation_cost is not None: + return 0.0, lyria_generation_cost speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) prompt_cost: float = 0.0 completion_cost: float = 0.0 - if not speech_model_info.get("input_cost_per_character") and not speech_model_info.get("input_cost_per_token"): - output_cost_per_generation: Final = speech_model_info.get("output_cost_per_image") - speech_output_cost_per_second: Final = speech_model_info.get("output_cost_per_second") - audio_seconds_per_prediction: Final = speech_model_info.get("audio_seconds_per_prediction") - if output_cost_per_generation is not None: - return prompt_cost, float(output_cost_per_generation) - if speech_output_cost_per_second is not None and audio_seconds_per_prediction is not None: - return prompt_cost, float(speech_output_cost_per_second) * float(audio_seconds_per_prediction) cost_metric: Final = select_cost_metric_for_model(speech_model_info) if cost_metric == "cost_per_character": if prompt_characters is None: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index b649d8dac0e..3895ea25434 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -2,6 +2,7 @@ import re from copy import deepcopy from enum import Enum from functools import lru_cache +from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints import httpx @@ -55,7 +56,26 @@ def _get_bundled_vertex_ai_lyria_model_info(model_key: str) -> VertexAILyriaMode def get_vertex_ai_lyria_model_info(model: str) -> VertexAILyriaModelInfo | None: model_key: Final = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" runtime_model_info: Final = _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key)) - return runtime_model_info or _get_bundled_vertex_ai_lyria_model_info(model_key) + bundled_model_info: Final = _get_bundled_vertex_ai_lyria_model_info(model_key) + if runtime_model_info is None: + return bundled_model_info + if bundled_model_info is None: + return runtime_model_info + return _validate_vertex_ai_lyria_model_info(MappingProxyType({**bundled_model_info, **runtime_model_info})) + + +def get_vertex_ai_lyria_generation_cost(model: str) -> float | None: + model_info: Final = get_vertex_ai_lyria_model_info(model) + if model_info is None: + return None + generation_cost: Final = model_info.get("output_cost_per_image") + if generation_cost is not None: + return generation_cost + cost_per_second: Final = model_info.get("output_cost_per_second") + seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction") + if cost_per_second is None or seconds_per_prediction is None: + return None + return cost_per_second * seconds_per_prediction class VertexAIError(BaseLLMException): diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 2047be2ea37..5a93aca3ac9 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -681,9 +681,11 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): ) # rebind-ok: interactions response supplies its audio MIME type if audio_data is None: raise ValueError(f"No generated audio found in Vertex AI {base_model} response") + decoded_audio: Final = base64.b64decode(audio_data) default_format: Final = model_info["supported_audio_formats"][0] mime_type = ( mime_type + or speech_media_type_from_audio_bytes(decoded_audio) or { # mutable-ok: short-lived lookup selects the default response MIME type; rebind-ok: absent provider MIME type falls back to model metadata "mp3": "audio/mpeg", "wav": "audio/wav", @@ -692,7 +694,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): response: Final = HttpxBinaryResponseContent( httpx.Response( status_code=raw_response.status_code, - content=base64.b64decode(audio_data), + content=decoded_audio, headers={ # mutable-ok: httpx requires a concrete response header dictionary "content-type": mime_type }, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index f7625d71168..afdbaf3b06c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -12,7 +12,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.vertex_ai.common_utils import ( - get_vertex_ai_lyria_model_info, + get_vertex_ai_lyria_generation_cost, get_vertex_location_from_url, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -397,9 +397,7 @@ class VertexPassthroughLoggingHandler: ) if runtime_unit_cost is not None: return runtime_unit_cost - return VertexPassthroughLoggingHandler._audio_prediction_unit_cost_from_model_info( - model_info=get_vertex_ai_lyria_model_info(model=model) - ) + return get_vertex_ai_lyria_generation_cost(model=model) @staticmethod def _audio_prediction_unit_cost_from_model_info(model_info: object) -> float | None: diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index f3d28d58c22..179301a9c81 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -146,13 +146,27 @@ def test_audio_predict_response_supports_bytes_base64_encoded( assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06) -def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_map_omits_model( +@pytest.mark.parametrize( + "missing_fields", + ( + None, + ("output_cost_per_second",), + ("audio_seconds_per_prediction",), + ("output_cost_per_second", "audio_seconds_per_prediction"), + ), +) +def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_metadata_is_incomplete( monkeypatch: pytest.MonkeyPatch, + missing_fields: tuple[str, ...] | None, ) -> None: - stale_runtime_model_cost: Final = { - key: value for key, value in litellm.model_cost.items() if key != "vertex_ai/lyria-002" - } - monkeypatch.setattr(litellm, "model_cost", stale_runtime_model_cost) + if missing_fields is None: + monkeypatch.delitem(litellm.model_cost, "vertex_ai/lyria-002") + else: + monkeypatch.setitem( + litellm.model_cost, + "vertex_ai/lyria-002", + {key: value for key, value in litellm.model_cost["vertex_ai/lyria-002"].items() if key not in missing_fields}, + ) logging_obj = MagicMock() logging_obj.model_call_details = {} response = httpx.Response( @@ -178,7 +192,8 @@ def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_map_omits_mod request_body={"instances": [{"prompt": "ambient piano"}]}, ) - assert "vertex_ai/lyria-002" not in litellm.model_cost + if missing_fields is None: + assert "vertex_ai/lyria-002" not in litellm.model_cost assert result["kwargs"]["model"] == "lyria-002" assert result["kwargs"]["response_cost"] == pytest.approx(0.06) assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06) diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index 457c3f76dbb..4f4df69780a 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -169,6 +169,23 @@ def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled(): class TestVertexAILyriaTextToSpeechConfig: + def test_response_without_mime_type_uses_audio_container(self) -> None: + audio: Final = b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00" + raw_response: Final = httpx.Response( + 200, + json={"outputs": [{"type": "audio", "data": base64.b64encode(audio).decode()}]}, + ) + + response: Final = VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response( + model="lyria-3-pro-preview", + raw_response=raw_response, + logging_obj=MagicMock(), + ) + + assert response.content == audio + assert response.response.headers["content-type"] == "audio/wav" + assert response._hidden_params["audio_mime_type"] == "audio/wav" + @pytest.mark.parametrize( "model", ["lyria-002", "vertex_ai/lyria-3-clip-preview", "lyria-3-pro-preview"], diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index ffc71c76d03..aa630c4916b 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,6 +1,7 @@ import json from pathlib import Path +from typing import Final import pytest @@ -154,14 +155,42 @@ def test_cost_calculator_with_response_cost_in_additional_headers(): ("vertex_ai/lyria-3-pro-preview", 0.08), ], ) -def test_vertex_lyria_speech_cost(model, expected_cost, _local_model_cost_map): - cost = completion_cost( +@pytest.mark.parametrize("runtime_state", ("complete", "missing", "routing_only", "custom_zero", "custom_price")) +@pytest.mark.parametrize("call_type", ("speech", "aspeech")) +def test_vertex_lyria_speech_cost( + model: str, + expected_cost: float, + _local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + runtime_state: str, + call_type: str, +) -> None: + model_info: Final = litellm.model_cost[model] + if runtime_state == "missing": + monkeypatch.delitem(litellm.model_cost, model) + elif runtime_state == "routing_only": + monkeypatch.setitem( + litellm.model_cost, + model, + { + key: value + for key, value in model_info.items() + if key not in ("output_cost_per_image", "output_cost_per_second", "audio_seconds_per_prediction") + }, + ) + elif runtime_state in ("custom_zero", "custom_price"): + cost_key: Final = "output_cost_per_image" if "output_cost_per_image" in model_info else "output_cost_per_second" + multiplier: Final = 0 if runtime_state == "custom_zero" else 2 + monkeypatch.setitem(litellm.model_cost, model, {**model_info, cost_key: model_info[cost_key] * multiplier}) + + cost: Final = completion_cost( model=model, prompt="A bright synth track", - call_type="speech", + call_type=call_type, ) - assert cost == pytest.approx(expected_cost) + expected: Final = 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1) + assert cost == pytest.approx(expected) def test_baseten_model_api_pricing_entries(_local_model_cost_map): From fd24cce2c3d0893e70d88f1e2d6ab3427a09736c Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sun, 6 Sep 2026 00:15:08 -0500 Subject: [PATCH 48/63] test(vertex): isolate Lyria fallback from the remote catalog --- .../llms/vertex_ai/test_vertex_passthrough_logging_handler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index 179301a9c81..2cf6689261e 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -158,6 +158,7 @@ def test_audio_predict_response_supports_bytes_base64_encoded( def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_metadata_is_incomplete( monkeypatch: pytest.MonkeyPatch, missing_fields: tuple[str, ...] | None, + local_model_cost_map: None, ) -> None: if missing_fields is None: monkeypatch.delitem(litellm.model_cost, "vertex_ai/lyria-002") From e79f3ec5205d01323093534a6577905a6dfdb7ac Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:31:32 -0700 Subject: [PATCH 49/63] fix(cost-map): stop advertising reasoning_effort max on the azure gpt-6-astra rows Both Azure routes refuse it. A live call to the same deployment through openai/deployments/gpt-6-astra/chat/completions on api-version 2025-04-01-preview answers reasoning_effort max with a 400 unsupported_value naming none, low, medium, high and xhigh as the values it takes, and xhigh returns 200, so azure/gpt-6-astra and azure/us/gpt-6-astra now match the azure_ai row. --- ...odel_prices_and_context_window_backup.json | 4 +-- .../reasoning_effort_capability.py | 4 +-- model_prices_and_context_window.json | 4 +-- .../test_reasoning_effort_capability.py | 26 ++++++------------- 4 files changed, 14 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b659c3b65e5..48c51f8bbf5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -7237,7 +7237,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -7503,7 +7503,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/litellm/router_utils/reasoning_effort_capability.py b/litellm/router_utils/reasoning_effort_capability.py index 9185d901a28..7b145c15a07 100644 --- a/litellm/router_utils/reasoning_effort_capability.py +++ b/litellm/router_utils/reasoning_effort_capability.py @@ -10,8 +10,8 @@ opt-in. none is opt-out everywhere except the azure gpt-5 family, whose config r 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 +all: outside the gpt-6-astra rows 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. diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b659c3b65e5..48c51f8bbf5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -7237,7 +7237,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, @@ -7503,7 +7503,7 @@ ], "supports_computer_use": true, "supports_function_calling": true, - "supports_max_reasoning_effort": true, + "supports_max_reasoning_effort": false, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, "supports_none_reasoning_effort": true, diff --git a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py index fa3a6dcd95a..ccd6766b13a 100644 --- a/tests/test_litellm/router_utils/test_reasoning_effort_capability.py +++ b/tests/test_litellm/router_utils/test_reasoning_effort_capability.py @@ -394,30 +394,20 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels: [ ("azure/gpt-6-astra", "azure"), ("azure/us/gpt-6-astra", "azure"), + ("azure_ai/gpt-6-astra", "azure_ai"), ], ) - def test_a_foundry_deployment_also_advertises_none(self, local_model_cost_map, model, custom_llm_provider): - """Microsoft Foundry serves the same model but its API accepts reasoning_effort none - (verified live: 200 with zero reasoning tokens, and it unlocks temperature), which - OpenAI's rejects, so an Azure deployment offers none on top of low through max.""" + def test_an_azure_hosted_deployment_advertises_none_but_not_max( + self, local_model_cost_map, model, custom_llm_provider + ): + """Microsoft hosts the same model with a different level set than OpenAI does. Verified live + on both Azure routes: none returns 200 with zero reasoning tokens and unlocks temperature, + which OpenAI's API rejects, while max returns 400 unsupported_value naming none through + xhigh as the levels it does take.""" from litellm.utils import _get_model_info_helper model_info = dict(_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)) - assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( - "none", - "low", - "medium", - "high", - "xhigh", - "max", - ) - - def test_a_foundry_azure_ai_deployment_advertises_none_but_not_max(self, local_model_cost_map): - from litellm.utils import _get_model_info_helper - - model_info = dict(_get_model_info_helper(model="azure_ai/gpt-6-astra", custom_llm_provider="azure_ai")) - assert resolve_supported_reasoning_efforts(model_info, deployment_is_mapped=True) == ( "none", "low", From fa2b64878b6f7be8fed5139ef95961fe26241f99 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:31:33 -0700 Subject: [PATCH 50/63] fix(azure_ai): redirect a gpt-5 capability lookup only when the map has a foundry row gpt-6-astra is the only gpt-5-family name with an azure_ai row. Prefixing the rest cost them every effort flag, since get_llm_provider sends an azure_ai name down the azure provider when a global AZURE_AI_API_BASE points at an openai.azure.com host and azure/ is not a key either, which turned temperature, top_p and logprobs on azure_ai/gpt-5.1-chat-latest from accepted into an UnsupportedParamsError. --- litellm/llms/azure_ai/chat/transformation.py | 14 ++++++++++- .../chat/test_azure_ai_transformation.py | 23 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 7c9a26c3f07..039c462b38a 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -46,7 +46,19 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( class AzureAIGPT5Config(OpenAIGPT5Config): @classmethod def _model_map_lookup_name(cls, model: str) -> str: - return model if model.startswith("azure_ai/") else f"azure_ai/{model}" + """Normalise a Foundry routing name to its cost-map key, when the map has one. + + A Foundry deployment and its OpenAI-hosted namesake are different products with + different capabilities, so ``azure_ai/`` is the entry to read whenever the map + carries it. Most gpt-5-family names have no ``azure_ai/`` row, though, and prefixing + those anyway costs them every flag: ``get_llm_provider`` re-resolves an ``azure_ai/`` + name to the azure provider when a global AZURE_AI_API_BASE points at an + openai.azure.com host, ``azure/`` is not a key either, so the lookup lands + nowhere and every effort answer degrades to False. A missing key defers to the base + resolver instead. + """ + prefixed: Final = model if model.startswith("azure_ai/") else f"azure_ai/{model}" + return prefixed if prefixed in litellm.model_cost else super()._model_map_lookup_name(model) azureAIGPT5Config: Final = AzureAIGPT5Config() 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 5ff0b729449..9924d77eb39 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 @@ -157,6 +157,29 @@ def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none assert optional_params == {"reasoning_effort": "none", "temperature": 0.2, "top_p": 0.9} +def test_a_gpt_5_name_without_a_foundry_row_keeps_reading_its_own_entry( + monkeypatch: pytest.MonkeyPatch, _local_model_cost_map +): + """gpt-6-astra is the only gpt-5-family name with an azure_ai/ row. Reading an azure_ai/ key for + the rest finds nothing, and an openai.azure.com base sends that name down the azure provider, + which has no key for it either, so every effort answer would silently fall back to false and + take temperature, top_p and logprobs down with it.""" + monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.openai.azure.com") + monkeypatch.setenv("AZURE_AI_API_KEY", "placeholder") + + optional_params = litellm.utils.get_optional_params( + model="gpt-5.1-chat-latest", + custom_llm_provider="azure_ai", + temperature=0.2, + top_p=0.9, + logprobs=True, + ) + + assert optional_params["temperature"] == 0.2 + assert optional_params["top_p"] == 0.9 + assert optional_params["logprobs"] is True + + def test_azure_ai_grok_stop_parameter_handling(): """ Test that Grok models properly handle stop parameter filtering in Azure AI Studio. From 0fb3951b2cadca5d9091b7586d0a5a15e4600b42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:32:06 -0700 Subject: [PATCH 51/63] fix(spend): charge a batch once when an older proxy left its cost row at $0 A proxy without this fix wrote the batch's cost row on every poll while the batch was still running, so that row reads $0 and the insert that claims the charge has nowhere to land. The retrieve that charges the batch now writes its own payload over that row under a where clause that still names spend 0.0, so exactly one retrieve takes it over and every later one reads the charge and charges nothing --- litellm/proxy/db/db_spend_update_writer.py | 42 +++++++++- .../proxy/db/test_db_spend_update_writer.py | 76 ++++++++++++++++++- 2 files changed, 112 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 48312c025dd..ae3bc5663eb 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -14,6 +14,7 @@ import time import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload import litellm @@ -379,9 +380,44 @@ class DBSpendUpdateWriter: if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False - verbose_proxy_logger.debug( - "Spend row %s charged nothing for this batch, so this retrieve charges it", request_id - ) + return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client) + + async def _take_over_uncharged_batch_cost_row( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient" + ) -> bool: + """Take the batch's cost row over from the poll that left it charging nothing. + + A pre-upgrade proxy wrote that row every time it polled the batch while it was still + running, so the charge is still to be made and the row still has to end up carrying + it. The row stops matching the moment it carries a charge, so it is one retrieve that + takes it over and charges, and every later one reads the charge and charges nothing. + """ + from litellm.repositories.table_repositories import SpendLogsRepository + + request_id: Final = payload["request_id"] + if payload["spend"] <= 0: + verbose_proxy_logger.debug( + "Cost tracking skipped: this batch costs nothing and spend row %s says so", request_id + ) + return False + try: + taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many( + data=prisma_client.jsonify_object( + MappingProxyType({field: value for field, value in payload.items() if field != "request_id"}) + ), + where={ # mutable-ok: prisma where clause + "request_id": request_id, + "call_type": CallTypes.aretrieve_batch.value, + "status": "success", + "spend": 0.0, + }, + ) + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; a row it cannot take over charges the batch + verbose_proxy_logger.warning("Could not take over spend row %s for a batch's cost: %s", request_id, e) + return True + if taken_over == 0: + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + return False return True async def _enqueue_tool_usage_transaction( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 41f2be08545..33b7af06e1a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2969,16 +2969,21 @@ def _batch_cost_payload() -> dict: } -def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock: +def _spend_logs_prisma(inserted: int, existing: object, taken_over: int = 1) -> MagicMock: prisma = _tool_usage_prisma() prisma.jsonify_object = lambda data: dict(data) prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) + prisma.db.litellm_spendlogs.update_many = AsyncMock(return_value=taken_over) return prisma async def _update_database_with( - db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict, disable_spend_logs: bool = False + db_writer: DBSpendUpdateWriter, + prisma: MagicMock, + payload: dict, + disable_spend_logs: bool = False, + response_cost: float = 0.25, ) -> bool: with ( patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam @@ -3005,7 +3010,7 @@ async def _update_database_with( completion_response=None, start_time=datetime.now(timezone.utc), end_time=datetime.now(timezone.utc), - response_cost=0.25, + response_cost=response_cost, ) await asyncio.sleep(0) return charged @@ -3054,6 +3059,71 @@ async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote assert db_writer._batch_database_updates.await_count == (1 if charged else 0) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("taken_over", "charged"), + [(1, True), (0, False)], + ids=["this_retrieve_takes_it_over", "another_one_got_there_first"], +) +async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left_at_zero( + taken_over: int, charged: bool +): + """ + A proxy without this fix wrote the batch's row at $0 on every poll of a running batch, + and the row outlives the upgrade, so the charge has to land on the row itself. Charging + without writing it there would charge again on every later retrieve (LIT-7048). + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing, taken_over) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged + + taken = prisma.db.litellm_spendlogs.update_many.await_args.kwargs + assert taken["where"] == { + "request_id": "batch_abc_batch_cost", + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.0, + } + assert taken["data"]["spend"] == 0.25 + assert "request_id" not in taken["data"] + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_charges_a_batch_whose_zero_row_it_could_not_take_over(): + """A DB that refuses the takeover must not swallow the batch's cost.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing) + prisma.db.litellm_spendlogs.update_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +async def test_update_database_leaves_a_batch_that_cost_nothing_to_the_retrieve_that_wrote_its_row(): + """ + A batch every line of which failed costs $0, so its row reads $0 for the honest reason + and the retrieve that wrote it is still the one that accounted it. Taking that row over + on every later retrieve would count one batch as many requests. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) + prisma = _spend_logs_prisma(0, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload(), response_cost=0.0) is False + + prisma.db.litellm_spendlogs.update_many.assert_not_called() + assert db_writer._batch_database_updates.await_count == 0 + + @pytest.mark.asyncio @pytest.mark.parametrize( ("inserted", "existing", "charged"), From 6be78fa850d4afc89527488115acb5117a6a34be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:34:31 -0700 Subject: [PATCH 52/63] fix(vertex_ai): bill Lyria per generation, not per audio second Google prices Lyria per generated clip, so every Vertex Lyria entry in the price map now carries a single output_cost_per_image and both the speech and the passthrough cost paths read that one field. The old output_cost_per_second and audio_seconds_per_prediction pair assumed a 30 second clip, which does not match the 32.768 second WAV Vertex returns, and no other model in the map priced audio that way Drops max_audio_length_hours and max_audio_per_prompt from the price map, its schema, the generator, and ModelInfo, since nothing reads them, and drops the audio_mime_type hidden param for the same reason: the response already carries the resolved content type on its own header Folds the per-model bundled catalog lookups into one cached parse of the local cost map, validated with a TypeAdapter over a ReadOnly TypedDict --- ci_cd/generate_model_prices_schema.py | 12 ----- litellm/cost_calculator.py | 16 +++--- litellm/llms/vertex_ai/common_utils.py | 51 ++++++++++--------- .../text_to_speech/transformation.py | 27 +++------- litellm/main.py | 13 ++--- ...odel_prices_and_context_window_backup.json | 7 +-- .../vertex_passthrough_logging_handler.py | 28 +--------- litellm/proxy/proxy_server.py | 11 ++-- litellm/types/llms/openai.py | 3 -- litellm/types/utils.py | 3 -- litellm/utils.py | 3 -- model_prices_and_context_window.json | 7 +-- model_prices_and_context_window.schema.json | 15 ------ .../vertex_ai/test_vertex_ai_common_utils.py | 41 +++++++++++++++ ...test_vertex_passthrough_logging_handler.py | 41 ++++++++------- .../text_to_speech/test_transformation.py | 36 ++++++------- tests/test_litellm/test_cost_calculator.py | 13 +++-- tests/test_litellm/test_utils.py | 6 +-- 18 files changed, 141 insertions(+), 192 deletions(-) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index ee0dad25c81..ab29b70bdd4 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -121,10 +121,6 @@ ARRAY_KEYS: dict[str, JsonSchema] = { } INTEGER_KEYS: dict[str, JsonSchema] = { - "max_audio_per_prompt": { - **NONNEG_INTEGER, - "description": "Maximum number of audio outputs accepted or generated per prompt.", - }, "max_tokens": { **NONNEG_INTEGER, "description": "Legacy field: max output tokens if the provider specifies it, else max input tokens.", @@ -150,14 +146,6 @@ INTEGER_KEYS: dict[str, JsonSchema] = { } NUMBER_KEYS: dict[str, JsonSchema] = { - "audio_seconds_per_prediction": { - **NONNEG_NUMBER, - "description": "Audio duration, in seconds, produced by one prediction.", - }, - "max_audio_length_hours": { - **NONNEG_NUMBER, - "description": "Maximum generated audio duration, expressed in hours.", - }, "regional_processing_uplift_multiplier_eu": { "type": "number", "minimum": 1, diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index fa8e1a5264c..eabf2c6249b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -80,6 +80,7 @@ from litellm.llms.together_ai.cost_calculator import ( get_model_params_and_category, has_together_registry_pricing, ) +from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_generation_cost from litellm.llms.vertex_ai.cost_calculator import ( cost_per_character as google_cost_per_character, ) @@ -495,16 +496,17 @@ def cost_per_token( # see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models if call_type == "speech" or call_type == "aspeech": - if custom_llm_provider in ("vertex_ai", "vertex_ai_beta"): - from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_generation_cost - - lyria_generation_cost: Final = get_vertex_ai_lyria_generation_cost(model_without_prefix) - if lyria_generation_cost is not None: - return 0.0, lyria_generation_cost + lyria_generation_cost: Final = ( + get_vertex_ai_lyria_generation_cost(model=model_without_prefix) + if custom_llm_provider in ("vertex_ai", "vertex_ai_beta") + else None + ) + if lyria_generation_cost is not None: + return 0.0, lyria_generation_cost speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) + cost_metric: Final = select_cost_metric_for_model(speech_model_info) prompt_cost: float = 0.0 completion_cost: float = 0.0 - cost_metric: Final = select_cost_metric_for_model(speech_model_info) if cost_metric == "cost_per_character": if prompt_characters is None: raise ValueError( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 3895ea25434..38cdf334cc3 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1,4 +1,5 @@ import re +from collections.abc import Mapping from copy import deepcopy from enum import Enum from functools import lru_cache @@ -29,8 +30,6 @@ class VertexAILyriaModelInfo(TypedDict): vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"]] supported_audio_formats: ReadOnly[tuple[Literal["mp3", "wav"], ...]] output_cost_per_image: NotRequired[ReadOnly[float]] - output_cost_per_second: NotRequired[ReadOnly[float]] - audio_seconds_per_prediction: NotRequired[ReadOnly[float]] _VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER: Final = TypeAdapter(VertexAILyriaModelInfo) @@ -45,37 +44,41 @@ def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyri return None -@lru_cache(maxsize=32) -def _get_bundled_vertex_ai_lyria_model_info(model_key: str) -> VertexAILyriaModelInfo | None: +@lru_cache(maxsize=1) +def _bundled_vertex_ai_lyria_model_infos() -> Mapping[str, VertexAILyriaModelInfo]: from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap - bundled_model_info: Final = GetModelCostMap.load_local_model_cost_map().get(model_key) - return _validate_vertex_ai_lyria_model_info(bundled_model_info) + return MappingProxyType( + { + model_key: lyria_model_info + for model_key, raw_model_info in GetModelCostMap.load_local_model_cost_map().items() + if (lyria_model_info := _validate_vertex_ai_lyria_model_info(raw_model_info)) is not None + } + ) + + +def _vertex_ai_lyria_model_key(model: str) -> str: + return model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" + + +def _vertex_ai_lyria_generation_cost(model_info: VertexAILyriaModelInfo | None) -> float | None: + return None if model_info is None else model_info.get("output_cost_per_image") def get_vertex_ai_lyria_model_info(model: str) -> VertexAILyriaModelInfo | None: - model_key: Final = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" + model_key: Final = _vertex_ai_lyria_model_key(model) runtime_model_info: Final = _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key)) - bundled_model_info: Final = _get_bundled_vertex_ai_lyria_model_info(model_key) - if runtime_model_info is None: - return bundled_model_info - if bundled_model_info is None: - return runtime_model_info - return _validate_vertex_ai_lyria_model_info(MappingProxyType({**bundled_model_info, **runtime_model_info})) + return runtime_model_info or _bundled_vertex_ai_lyria_model_infos().get(model_key) def get_vertex_ai_lyria_generation_cost(model: str) -> float | None: - model_info: Final = get_vertex_ai_lyria_model_info(model) - if model_info is None: - return None - generation_cost: Final = model_info.get("output_cost_per_image") - if generation_cost is not None: - return generation_cost - cost_per_second: Final = model_info.get("output_cost_per_second") - seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction") - if cost_per_second is None or seconds_per_prediction is None: - return None - return cost_per_second * seconds_per_prediction + model_key: Final = _vertex_ai_lyria_model_key(model) + runtime_cost: Final = _vertex_ai_lyria_generation_cost( + _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key)) + ) + if runtime_cost is not None: + return runtime_cost + return _vertex_ai_lyria_generation_cost(_bundled_vertex_ai_lyria_model_infos().get(model_key)) class VertexAIError(BaseLLMException): diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 5a93aca3ac9..f489366b8d8 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -15,6 +15,7 @@ import httpx import litellm from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.audio_utils.utils import ( + DEFAULT_SPEECH_MEDIA_TYPE, speech_media_type_from_audio_bytes, ) from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -571,13 +572,11 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): VertexAIInteractionsConfig, ) - resolved_project: Final = project - def mint_access_token( _credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, ) -> tuple[str, str]: - return "", project_id or resolved_project + return "", project_id or project return VertexAIInteractionsConfig(mint_access_token=mint_access_token).get_complete_url( api_base=api_base, @@ -681,24 +680,12 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): ) # rebind-ok: interactions response supplies its audio MIME type if audio_data is None: raise ValueError(f"No generated audio found in Vertex AI {base_model} response") - decoded_audio: Final = base64.b64decode(audio_data) - default_format: Final = model_info["supported_audio_formats"][0] - mime_type = ( - mime_type - or speech_media_type_from_audio_bytes(decoded_audio) - or { # mutable-ok: short-lived lookup selects the default response MIME type; rebind-ok: absent provider MIME type falls back to model metadata - "mp3": "audio/mpeg", - "wav": "audio/wav", - }[default_format] - ) - response: Final = HttpxBinaryResponseContent( + binary_data: Final = base64.b64decode(audio_data) + media_type: Final = mime_type or speech_media_type_from_audio_bytes(binary_data) or DEFAULT_SPEECH_MEDIA_TYPE + return HttpxBinaryResponseContent( httpx.Response( status_code=raw_response.status_code, - content=decoded_audio, - headers={ # mutable-ok: httpx requires a concrete response header dictionary - "content-type": mime_type - }, + content=binary_data, + headers=MappingProxyType({"content-type": media_type}), ) ) - response.set_audio_mime_type(mime_type) - return response diff --git a/litellm/main.py b/litellm/main.py index 040af897256..44afce08c0a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8260,14 +8260,11 @@ def speech( # Vertex AI Text-to-Speech (Google Cloud TTS) if text_to_speech_provider_config is None: - if VertexAILyriaTextToSpeechConfig.is_lyria_model(model): - text_to_speech_provider_config = ( - VertexAILyriaTextToSpeechConfig() - ) # rebind-ok: model metadata selects the Lyria provider implementation - else: - text_to_speech_provider_config = ( - VertexAITextToSpeechConfig() - ) # rebind-ok: non-Lyria Vertex models use the standard TTS implementation + text_to_speech_provider_config = ( # rebind-ok: model metadata selects the Vertex TTS implementation + VertexAILyriaTextToSpeechConfig() + if VertexAILyriaTextToSpeechConfig.is_lyria_model(model) + else VertexAITextToSpeechConfig() + ) # Cast to specific Vertex AI config type to access dispatch method vertex_config: Final = cast(VertexAITextToSpeechConfig, text_to_speech_provider_config) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9986d5056d1..0bc6c59a4b7 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45512,12 +45512,9 @@ "supports_tool_choice": true }, "vertex_ai/lyria-002": { - "audio_seconds_per_prediction": 30, "litellm_provider": "vertex_ai", - "max_audio_length_hours": 0.009111111111111111, - "max_audio_per_prompt": 4, - "mode": "audio_speech", - "output_cost_per_second": 0.002, + "mode": "audio_speech", + "output_cost_per_image": 0.06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_audio_formats": [ "wav" diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index afdbaf3b06c..119a53c2411 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -1,6 +1,5 @@ import asyncio import re -from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse @@ -346,9 +345,7 @@ class VertexPassthroughLoggingHandler: prediction_count: Final = VertexPassthroughLoggingHandler._get_audio_prediction_count( json_response=json_response ) - response_cost: Final = ( - VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) or 0.0 - ) * prediction_count + response_cost: Final = (get_vertex_ai_lyria_generation_cost(model=model) or 0.0) * prediction_count logging_obj.model = model # rebind-ok: passthrough attribution records the resolved Vertex model logging_obj.model_call_details[ # rebind-ok: passthrough attribution enriches callback metadata @@ -387,30 +384,9 @@ class VertexPassthroughLoggingHandler: ) -> bool: return ( VertexPassthroughLoggingHandler._get_audio_prediction_count(json_response=json_response) > 0 - and VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) is not None + and get_vertex_ai_lyria_generation_cost(model=model) is not None ) - @staticmethod - def _get_audio_prediction_unit_cost(model: str) -> float | None: - runtime_unit_cost: Final = VertexPassthroughLoggingHandler._audio_prediction_unit_cost_from_model_info( - model_info=litellm.model_cost.get(f"vertex_ai/{model}") - ) - if runtime_unit_cost is not None: - return runtime_unit_cost - return get_vertex_ai_lyria_generation_cost(model=model) - - @staticmethod - def _audio_prediction_unit_cost_from_model_info(model_info: object) -> float | None: - if not isinstance(model_info, Mapping): - return None - output_cost_per_second: Final = model_info.get("output_cost_per_second") - audio_seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction") - if not isinstance(output_cost_per_second, (int, float)) or not isinstance( - audio_seconds_per_prediction, (int, float) - ): - return None - return float(output_cost_per_second * audio_seconds_per_prediction) - @staticmethod def _get_audio_prediction_count( json_response: dict, # mutable-ok: counter inspects the decoded provider response dictionary without mutation diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f11cbc92224..27132c90e05 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11176,14 +11176,9 @@ async def audio_speech( upstream_content_type: Final = ( response.response.headers.get("content-type") if isinstance(response, HttpxBinaryResponseContent) else None ) - hidden_audio_mime_type: Final = hidden_params.get("audio_mime_type") - media_type: Final = ( - hidden_audio_mime_type - if isinstance(hidden_audio_mime_type, str) - else resolve_speech_media_type( - upstream_content_type=upstream_content_type, - response_format=requested_format if isinstance(requested_format, str) else None, - ) + media_type: Final = resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=requested_format if isinstance(requested_format, str) else None, ) return StreamingResponse( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 0b13191d977..32d88da0085 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -119,9 +119,6 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent): return self._hidden_params["response_cost"] = response_cost - def set_audio_mime_type(self, audio_mime_type: str) -> None: - self._hidden_params["audio_mime_type"] = audio_mime_type - class NotGiven: """ diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6c645e70bee..e29944919ba 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -312,9 +312,6 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_video_per_second: float | None # only for vertex ai models output_cost_per_audio_per_second: float | None # only for vertex ai models output_cost_per_second: float | None # for OpenAI Speech models - audio_seconds_per_prediction: ReadOnly[float | None] - max_audio_length_hours: ReadOnly[float | None] - max_audio_per_prompt: ReadOnly[int | None] output_cost_per_second_1080p: ( float | None ) # video_generation tier: key output_cost_per_second_ (e.g. 1080p, 720p) diff --git a/litellm/utils.py b/litellm/utils.py index d6a68b3ce6e..f6b40d33a79 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5880,9 +5880,6 @@ def _get_model_info_helper( "output_cost_per_token_above_512k_tokens", None ), output_cost_per_second=_model_info.get("output_cost_per_second", None), - audio_seconds_per_prediction=_model_info.get("audio_seconds_per_prediction", None), - max_audio_length_hours=_model_info.get("max_audio_length_hours", None), - max_audio_per_prompt=_model_info.get("max_audio_per_prompt", None), output_cost_per_second_1080p=_model_info.get("output_cost_per_second_1080p", None), output_cost_per_second_480p=_model_info.get("output_cost_per_second_480p", None), output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9986d5056d1..0bc6c59a4b7 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45512,12 +45512,9 @@ "supports_tool_choice": true }, "vertex_ai/lyria-002": { - "audio_seconds_per_prediction": 30, "litellm_provider": "vertex_ai", - "max_audio_length_hours": 0.009111111111111111, - "max_audio_per_prompt": 4, - "mode": "audio_speech", - "output_cost_per_second": 0.002, + "mode": "audio_speech", + "output_cost_per_image": 0.06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_audio_formats": [ "wav" diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 58ca91f3977..6faa957dc16 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -53,11 +53,6 @@ "type": "number", "minimum": 0 }, - "audio_seconds_per_prediction": { - "type": "number", - "minimum": 0, - "description": "Audio duration, in seconds, produced by one prediction." - }, "audio_transcription_config": { "type": "string" }, @@ -368,16 +363,6 @@ "type": "string", "description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers." }, - "max_audio_length_hours": { - "type": "number", - "minimum": 0, - "description": "Maximum generated audio duration, expressed in hours." - }, - "max_audio_per_prompt": { - "type": "integer", - "minimum": 0, - "description": "Maximum number of audio outputs accepted or generated per prompt." - }, "max_input_tokens": { "type": "integer", "minimum": 0, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index d1d751989ea..48376872a72 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1696,3 +1696,44 @@ def test_vertex_text_embedding_request_includes_labels_from_metadata(): }, ) assert req.get("labels") == {"project_id": "cost-center-1"} + + +@pytest.mark.parametrize( + ("model", "expected_api"), + [ + ("lyria-002", "lyria_predict"), + ("vertex_ai/lyria-002", "lyria_predict"), + ("lyria-3-clip-preview", "lyria_interactions"), + ("lyria-3-pro-preview", "lyria_interactions"), + ], +) +def test_get_vertex_ai_lyria_model_info_resolves_audio_api(model, expected_api): + from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_model_info + + model_info = get_vertex_ai_lyria_model_info(model=model) + + assert model_info is not None + assert model_info["vertex_ai_audio_api"] == expected_api + + +@pytest.mark.parametrize("model", ["en-US-Studio-O", "gemini-2.5-flash-preview-tts", "chirp-3-hd-charon"]) +def test_get_vertex_ai_lyria_model_info_is_none_for_non_lyria_speech_models(model): + from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_model_info + + assert get_vertex_ai_lyria_model_info(model=model) is None + + +def test_get_vertex_ai_lyria_model_info_falls_back_to_bundled_map(monkeypatch): + import litellm + from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_model_info + + stale_runtime_model_cost = { + key: value for key, value in litellm.model_cost.items() if not key.startswith("vertex_ai/lyria") + } + monkeypatch.setattr(litellm, "model_cost", stale_runtime_model_cost) + + model_info = get_vertex_ai_lyria_model_info(model="lyria-3-pro-preview") + + assert model_info is not None + assert model_info["vertex_ai_audio_api"] == "lyria_interactions" + assert model_info["supported_audio_formats"] == ("mp3", "wav") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index 2cf6689261e..682dadfe854 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -18,8 +18,9 @@ def test_lyria_predict_response_preserves_audio_response_and_logs_cost( litellm.model_cost, "vertex_ai/lyria-002", { - "audio_seconds_per_prediction": 30, - "output_cost_per_second": 0.002, + "vertex_ai_audio_api": "lyria_predict", + "supported_audio_formats": ["wav"], + "output_cost_per_image": 0.06, }, ) logging_obj = MagicMock() @@ -79,8 +80,9 @@ def test_audio_predict_response_uses_model_map_metadata( litellm.model_cost, "vertex_ai/music-audio-preview", { - "audio_seconds_per_prediction": 12, - "output_cost_per_second": 0.5, + "vertex_ai_audio_api": "lyria_predict", + "supported_audio_formats": ["wav"], + "output_cost_per_image": 0.5, }, ) logging_obj = MagicMock() @@ -109,8 +111,8 @@ def test_audio_predict_response_uses_model_map_metadata( ) assert result["kwargs"]["model"] == "music-audio-preview" - assert result["kwargs"]["response_cost"] == pytest.approx(6.0) - assert logging_obj.model_call_details["response_cost"] == pytest.approx(6.0) + assert result["kwargs"]["response_cost"] == pytest.approx(0.5) + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.5) def test_audio_predict_response_supports_bytes_base64_encoded( @@ -120,8 +122,9 @@ def test_audio_predict_response_supports_bytes_base64_encoded( litellm.model_cost, "vertex_ai/lyria-002", { - "audio_seconds_per_prediction": 30, - "output_cost_per_second": 0.002, + "vertex_ai_audio_api": "lyria_predict", + "supported_audio_formats": ["wav"], + "output_cost_per_image": 0.06, }, ) logging_obj = MagicMock() @@ -146,27 +149,23 @@ def test_audio_predict_response_supports_bytes_base64_encoded( assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06) -@pytest.mark.parametrize( - "missing_fields", - ( - None, - ("output_cost_per_second",), - ("audio_seconds_per_prediction",), - ("output_cost_per_second", "audio_seconds_per_prediction"), - ), -) +@pytest.mark.parametrize("runtime_entry_is_missing", (True, False)) def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_metadata_is_incomplete( monkeypatch: pytest.MonkeyPatch, - missing_fields: tuple[str, ...] | None, + runtime_entry_is_missing: bool, local_model_cost_map: None, ) -> None: - if missing_fields is None: + if runtime_entry_is_missing: monkeypatch.delitem(litellm.model_cost, "vertex_ai/lyria-002") else: monkeypatch.setitem( litellm.model_cost, "vertex_ai/lyria-002", - {key: value for key, value in litellm.model_cost["vertex_ai/lyria-002"].items() if key not in missing_fields}, + { + key: value + for key, value in litellm.model_cost["vertex_ai/lyria-002"].items() + if key != "output_cost_per_image" + }, ) logging_obj = MagicMock() logging_obj.model_call_details = {} @@ -193,7 +192,7 @@ def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_metadata_is_i request_body={"instances": [{"prompt": "ambient piano"}]}, ) - if missing_fields is None: + if runtime_entry_is_missing: assert "vertex_ai/lyria-002" not in litellm.model_cost assert result["kwargs"]["model"] == "lyria-002" assert result["kwargs"]["response_cost"] == pytest.approx(0.06) diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index 4f4df69780a..b5eec42b569 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -169,23 +169,6 @@ def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled(): class TestVertexAILyriaTextToSpeechConfig: - def test_response_without_mime_type_uses_audio_container(self) -> None: - audio: Final = b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00" - raw_response: Final = httpx.Response( - 200, - json={"outputs": [{"type": "audio", "data": base64.b64encode(audio).decode()}]}, - ) - - response: Final = VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response( - model="lyria-3-pro-preview", - raw_response=raw_response, - logging_obj=MagicMock(), - ) - - assert response.content == audio - assert response.response.headers["content-type"] == "audio/wav" - assert response._hidden_params["audio_mime_type"] == "audio/wav" - @pytest.mark.parametrize( "model", ["lyria-002", "vertex_ai/lyria-3-clip-preview", "lyria-3-pro-preview"], @@ -385,11 +368,11 @@ class TestVertexAILyriaTextToSpeechConfig: { "predictions": [ { - "bytesBase64Encoded": "bHlyaWEtMi1hdWRpbw==", + "bytesBase64Encoded": "UklGRiQAAABXQVZFZm10IA==", } ] }, - b"lyria-2-audio", + b"RIFF$\x00\x00\x00WAVEfmt ", "audio/wav", ), ( @@ -427,6 +410,19 @@ class TestVertexAILyriaTextToSpeechConfig: b"lyria-3-audio", "audio/mpeg", ), + ( + "lyria-3-pro-preview", + { + "outputs": [ + { + "type": "audio", + "data": "UklGRiQAAABXQVZFZm10IA==", + } + ] + }, + b"RIFF$\x00\x00\x00WAVEfmt ", + "audio/wav", + ), ], ) def test_transform_response( @@ -446,7 +442,7 @@ class TestVertexAILyriaTextToSpeechConfig: ) assert response.content == expected_audio - assert response._hidden_params["audio_mime_type"] == expected_mime_type + assert response.response.headers["content-type"] == expected_mime_type @pytest.mark.parametrize( ("model", "response_format"), diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index aa630c4916b..3236253a10a 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -172,16 +172,15 @@ def test_vertex_lyria_speech_cost( monkeypatch.setitem( litellm.model_cost, model, - { - key: value - for key, value in model_info.items() - if key not in ("output_cost_per_image", "output_cost_per_second", "audio_seconds_per_prediction") - }, + {key: value for key, value in model_info.items() if key != "output_cost_per_image"}, ) elif runtime_state in ("custom_zero", "custom_price"): - cost_key: Final = "output_cost_per_image" if "output_cost_per_image" in model_info else "output_cost_per_second" multiplier: Final = 0 if runtime_state == "custom_zero" else 2 - monkeypatch.setitem(litellm.model_cost, model, {**model_info, cost_key: model_info[cost_key] * multiplier}) + monkeypatch.setitem( + litellm.model_cost, + model, + {**model_info, "output_cost_per_image": model_info["output_cost_per_image"] * multiplier}, + ) cost: Final = completion_cost( model=model, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index a179a82fcb2..f934b2aca1e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -944,14 +944,11 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "code_interpreter_cost_per_session": {"type": "number"}, "inference_geo": {"type": "string"}, "litellm_provider": {"type": "string"}, - "max_audio_length_hours": {"type": "number"}, - "max_audio_per_prompt": {"type": "number"}, "max_input_tokens": {"type": "number"}, "max_output_tokens": {"type": "number"}, "max_tokens": {"type": "number"}, "metadata": {"type": "object"}, "provider_specific_entry": {"type": "object"}, - "audio_seconds_per_prediction": {"type": "number"}, "mode": { "type": "string", "enum": [ @@ -2869,8 +2866,7 @@ def test_vertex_ai_lyria_models_in_cost_map(): assert lyria_2["mode"] == "audio_speech" assert clip["mode"] == "audio_speech" assert pro["mode"] == "audio_speech" - assert lyria_2["audio_seconds_per_prediction"] == 30 - assert lyria_2["output_cost_per_second"] == 0.002 + assert lyria_2["output_cost_per_image"] == 0.06 assert lyria_2["supported_modalities"] == ["text"] assert lyria_2["supported_output_modalities"] == ["audio"] assert lyria_2["supports_audio_output"] is True From b21315fe754ddcb54b6e87759b621e5e17e515ef Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:36:43 -0700 Subject: [PATCH 53/63] refactor(anthropic): type the responses refusal stream iterator state Types the cached sync upstream iterator instead of holding it as Any, so the Responses to Anthropic streaming wrapper carries no untyped state. --- .../responses_adapters/streaming_iterator.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index 5bd662e0f94..bcae732ad42 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -4,7 +4,7 @@ import asyncio import json import traceback from collections import deque -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final from litellm import verbose_logger @@ -55,7 +55,7 @@ class AnthropicResponsesStreamWrapper: self._sent_message_stop = False self._chunk_queue: deque[dict[str, object]] = deque() self._refusal_text_parts: list[str] = [] # mutable-ok: accumulates streamed refusal delta text across chunks - self._sync_responses_iterator: Any = None + self._sync_responses_iterator: Iterator[object] | None = None def _make_message_start(self) -> dict[str, object]: return { @@ -312,11 +312,9 @@ class AnthropicResponsesStreamWrapper: else: if self._sync_responses_iterator is None: self._sync_responses_iterator = iter(self.responses_stream) + sync_iterator: Final = self._sync_responses_iterator missing: Final = object() - while True: - event = await asyncio.to_thread(next, self._sync_responses_iterator, missing) - if event is missing: - break + while (event := await asyncio.to_thread(next, sync_iterator, missing)) is not missing: self._process_event(event) if self._chunk_queue: return self._chunk_queue.popleft() From 24f0be80219cd3403ec28cddbed9be946fad1cb0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:47:41 -0700 Subject: [PATCH 54/63] fix(spend): leave a batch uncharged when the database refuses the takeover The takeover of a $0 row an older proxy left behind used to charge the batch when the update could not reach the database. That leaves the row still reading $0, so every later retrieve finds the same row and charges the batch again, which is the repeat charging this PR exists to stop. The retrieve that does take the row over is the one that charges, and a batch nobody retrieves again after that failure is never charged, the same as one whose proxy died inside the write window. --- litellm/proxy/db/db_spend_update_writer.py | 8 +++++--- .../proxy/db/test_db_spend_update_writer.py | 12 ++++++++---- 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ae3bc5663eb..ee7802a45d3 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -412,9 +412,11 @@ class DBSpendUpdateWriter: "spend": 0.0, }, ) - except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; a row it cannot take over charges the batch - verbose_proxy_logger.warning("Could not take over spend row %s for a batch's cost: %s", request_id, e) - return True + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; the next retrieve takes the row over + verbose_proxy_logger.warning( + "Could not take over spend row %s, leaving this batch's cost to the next retrieve: %s", request_id, e + ) + return False if taken_over == 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 33b7af06e1a..4efb94b60aa 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -3093,17 +3093,21 @@ async def test_update_database_charges_a_batch_whose_row_a_pre_upgrade_poll_left @pytest.mark.asyncio -async def test_update_database_charges_a_batch_whose_zero_row_it_could_not_take_over(): - """A DB that refuses the takeover must not swallow the batch's cost.""" +async def test_update_database_leaves_a_batch_whose_zero_row_it_could_not_take_over_to_the_next_retrieve(): + """ + A DB that refuses the takeover leaves the row reading $0, so charging here would charge + the batch again on every later retrieve. The retrieve that does take the row over is the + one that charges. + """ db_writer = DBSpendUpdateWriter() db_writer._batch_database_updates = AsyncMock() existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.0) prisma = _spend_logs_prisma(0, existing) prisma.db.litellm_spendlogs.update_many = AsyncMock(side_effect=RuntimeError("db unreachable")) - assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False - assert db_writer._batch_database_updates.await_count == 1 + assert db_writer._batch_database_updates.await_count == 0 @pytest.mark.asyncio From 02b44820c4dd5da47abf11e9f1b2dbe2fdce80f5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:59:57 -0700 Subject: [PATCH 55/63] test(vertex_ai): keep imagen predict passthrough off the Lyria audio path The new Lyria passthrough branch runs before the image-generation branch and keys on the same `predictions[0].bytesBase64Encoded` shape imagen returns, so only the cost-map lookup separates them. Cover an imagen predict response end to end so a future change that drops that lookup fails here instead of misbilling images as audio. --- ...test_vertex_passthrough_logging_handler.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index 682dadfe854..98010021bca 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py @@ -9,6 +9,7 @@ import litellm from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) +from litellm.types.utils import PassthroughCallTypes def test_lyria_predict_response_preserves_audio_response_and_logs_cost( @@ -197,3 +198,33 @@ def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_metadata_is_i assert result["kwargs"]["model"] == "lyria-002" assert result["kwargs"]["response_cost"] == pytest.approx(0.06) assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06) + + +def test_image_predict_response_is_not_billed_as_audio( + local_model_cost_map: None, +) -> None: + logging_obj = MagicMock() + logging_obj.model_call_details = {} + response = httpx.Response( + status_code=200, + json={"predictions": [{"bytesBase64Encoded": "frame", "mimeType": "image/png"}]}, + ) + + result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=response, + logging_obj=logging_obj, + url_route=( + "/v1/projects/test/locations/us-central1/publishers/google/models/imagen-4.0-generate-001:predict" + ), + result=response.text, + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"instances": [{"prompt": "a red cube"}]}, + ) + + assert isinstance(result["result"], litellm.ImageResponse) + assert logging_obj.call_type == PassthroughCallTypes.passthrough_image_generation.value + assert result["kwargs"]["response_cost"] == pytest.approx( + litellm.model_cost["vertex_ai/imagen-4.0-generate-001"]["output_cost_per_image"] + ) From 11e45ad95361eb37c00f9cce799bf544fd9990a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:00:01 -0700 Subject: [PATCH 56/63] fix(vertex_ai): mark the Lyria 3 catalog entries text-only `vertex_ai/lyria-3-clip-preview` and `vertex_ai/lyria-3-pro-preview` were registered with `supports_vision`, `supports_image_input`, and an `image` modality, which contradicts their `gemini/lyria-3-*` siblings and makes /model/info advertise image input on text-to-music models. --- .../model_prices_and_context_window_backup.json | 14 +++++--------- model_prices_and_context_window.json | 14 +++++--------- tests/test_litellm/test_utils.py | 11 +++++++---- 3 files changed, 17 insertions(+), 22 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0bc6c59a4b7..5bdf6ed4fb8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -45513,7 +45513,7 @@ }, "vertex_ai/lyria-002": { "litellm_provider": "vertex_ai", - "mode": "audio_speech", + "mode": "audio_speech", "output_cost_per_image": 0.06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_audio_formats": [ @@ -45549,8 +45549,7 @@ "/v1/audio/speech" ], "supported_modalities": [ - "text", - "image" + "text" ], "supported_output_modalities": [ "audio" @@ -45561,11 +45560,10 @@ "supports_audio_input": false, "supports_audio_output": true, "supports_function_calling": false, - "supports_image_input": true, "supports_prompt_caching": false, "supports_response_schema": false, "supports_system_messages": false, - "supports_vision": true, + "supports_vision": false, "supports_web_search": false, "vertex_ai_audio_api": "lyria_interactions" }, @@ -45588,8 +45586,7 @@ "/v1/audio/speech" ], "supported_modalities": [ - "text", - "image" + "text" ], "supported_output_modalities": [ "audio" @@ -45600,11 +45597,10 @@ "supports_audio_input": false, "supports_audio_output": true, "supports_function_calling": false, - "supports_image_input": true, "supports_prompt_caching": false, "supports_response_schema": false, "supports_system_messages": false, - "supports_vision": true, + "supports_vision": false, "supports_web_search": false, "vertex_ai_audio_api": "lyria_interactions" }, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0bc6c59a4b7..5bdf6ed4fb8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -45513,7 +45513,7 @@ }, "vertex_ai/lyria-002": { "litellm_provider": "vertex_ai", - "mode": "audio_speech", + "mode": "audio_speech", "output_cost_per_image": 0.06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria", "supported_audio_formats": [ @@ -45549,8 +45549,7 @@ "/v1/audio/speech" ], "supported_modalities": [ - "text", - "image" + "text" ], "supported_output_modalities": [ "audio" @@ -45561,11 +45560,10 @@ "supports_audio_input": false, "supports_audio_output": true, "supports_function_calling": false, - "supports_image_input": true, "supports_prompt_caching": false, "supports_response_schema": false, "supports_system_messages": false, - "supports_vision": true, + "supports_vision": false, "supports_web_search": false, "vertex_ai_audio_api": "lyria_interactions" }, @@ -45588,8 +45586,7 @@ "/v1/audio/speech" ], "supported_modalities": [ - "text", - "image" + "text" ], "supported_output_modalities": [ "audio" @@ -45600,11 +45597,10 @@ "supports_audio_input": false, "supports_audio_output": true, "supports_function_calling": false, - "supports_image_input": true, "supports_prompt_caching": false, "supports_response_schema": false, "supports_system_messages": false, - "supports_vision": true, + "supports_vision": false, "supports_web_search": false, "vertex_ai_audio_api": "lyria_interactions" }, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index f934b2aca1e..d0916adce9d 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -2887,14 +2887,17 @@ def test_vertex_ai_lyria_models_in_cost_map(): "/v1beta/interactions", "/v1/audio/speech", ] - assert clip["supported_modalities"] == ["text", "image"] - assert pro["supported_modalities"] == ["text", "image"] + assert clip["supported_modalities"] == ["text"] + assert pro["supported_modalities"] == ["text"] + assert clip["supports_vision"] is False + assert pro["supports_vision"] is False + assert "supports_image_input" not in clip + assert "supports_image_input" not in pro assert clip["supported_regions"] == ["global"] assert pro["supported_regions"] == ["global"] assert clip["supports_audio_output"] is True assert pro["supports_audio_output"] is True - assert clip["supports_image_input"] is True - assert pro["supports_image_input"] is True + def test_model_info_for_fireworks_short_form_models(): """ From 3dea1ebb32c96560f9f10171d0597ca30d4bea40 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:18:25 -0700 Subject: [PATCH 57/63] fix(cost-map): keep the prompt cache breakpoint flag on the foundry gpt-6-astra row The openai gpt-6-astra card carries supports_prompt_cache_breakpoint, so a Foundry deployment reported it as true until the azure_ai row took over the lookup. The cache control hook still honours breakpoints for that deployment through the bare name, so /model/info was the only thing that changed, and it now agrees with the hook again. --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 48c51f8bbf5..2f86deffe54 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3524,6 +3524,7 @@ "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 48c51f8bbf5..2f86deffe54 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3524,6 +3524,7 @@ "supports_none_reasoning_effort": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, From fffe0bb0dc94f8b071be85050a0bdaa101c0e2ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:18:25 -0700 Subject: [PATCH 58/63] test(azure_ai): pin the tier the messages bridge sends when astra refuses max The /v1/messages adapter lowers a tier the entry does not accept, so dropping max from the astra rows moves that path from Foundry's 400 to a request at xhigh. Nothing pinned that, and the guard test's docstring named gpt-6-astra as the only gpt-5 name with an azure_ai row, which 11 rows contradict. --- ..._handler_reasoning_effort_normalization.py | 19 +++++++++++++++++++ .../chat/test_azure_ai_transformation.py | 8 ++++---- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py index 56b754c3476..af7befecc33 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -82,3 +82,22 @@ class TestTheNormalizedTierIsTheTierSent: self, local_model_cost_map, model, provider, effort, expected ): assert _reasoning_effort_sent(model, provider, effort) == expected + + @pytest.mark.parametrize( + "model, provider", + [ + ("gpt-6-astra", "azure_ai"), + ("azure_ai/gpt-6-astra", "azure_ai"), + ("gpt-6-astra", "azure"), + ("us/gpt-6-astra", "azure"), + ], + ) + def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts( + self, local_model_cost_map, model, provider + ): + """The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows + say so and the adapter sends the tier below instead of the rejected one.""" + assert _reasoning_effort_sent(model, provider, "max") == "xhigh" + + def test_the_openai_hosted_twin_still_sends_max(self, local_model_cost_map): + assert _reasoning_effort_sent("gpt-6-astra", "openai", "max") == "max" 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 9924d77eb39..f8cc0b5071e 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 @@ -160,10 +160,10 @@ def test_foundry_gpt_6_astra_keeps_sampling_params_when_reasoning_effort_is_none def test_a_gpt_5_name_without_a_foundry_row_keeps_reading_its_own_entry( monkeypatch: pytest.MonkeyPatch, _local_model_cost_map ): - """gpt-6-astra is the only gpt-5-family name with an azure_ai/ row. Reading an azure_ai/ key for - the rest finds nothing, and an openai.azure.com base sends that name down the azure provider, - which has no key for it either, so every effort answer would silently fall back to false and - take temperature, top_p and logprobs down with it.""" + """Most gpt-5-family names have no azure_ai/ row. Reading an azure_ai/ key for those finds + nothing, and an openai.azure.com base sends the name down the azure provider, which has no key + for it either, so every effort answer would silently fall back to false and take temperature, + top_p and logprobs down with it.""" monkeypatch.setenv("AZURE_AI_API_BASE", "https://example-resource.openai.azure.com") monkeypatch.setenv("AZURE_AI_API_KEY", "placeholder") From c09d34fc4b84a619f10eb7e64fe44fea198a8dc7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:20:19 -0700 Subject: [PATCH 59/63] fix(anthropic): stream refusals parked in provider_specific_fields The first-delta guard read `delta.refusal` directly, while the translation three lines later goes through `openai_chat_refusal_text`, which also reads the `provider_specific_fields` LiteLLM parks unrecognized fields in. A provider that sends the refusal that way had its only refusal delta skipped as blank, so the client got `stop_reason: refusal` over an empty content array, which is the symptom this PR set out to fix --- .../adapters/streaming_iterator.py | 5 ++- .../test_streaming_iterator_first_delta.py | 37 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index c6d744772a5..b2513a5c046 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -1064,6 +1064,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): @staticmethod def _is_blank_delta(chunk: "ModelResponseStream") -> bool: from litellm.llms.anthropic.common_utils import is_empty_unsigned_thinking_block + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + openai_chat_refusal_text, + ) choice: Final = chunk.choices[0] if choice.finish_reason is not None: @@ -1073,7 +1076,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): return False if getattr(delta, "content", None): return False - if getattr(delta, "refusal", None): + if openai_chat_refusal_text(delta): return False if getattr(delta, "reasoning_content", None): return False diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 4359ace1870..4c5e2101fbf 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -145,6 +145,43 @@ async def test_streaming_chat_refusal_emits_refusal_text_and_stop_details_async( assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." +def test_streaming_chat_refusal_parked_in_provider_specific_fields_is_emitted(): + """Providers that do not populate ``delta.refusal`` (Azure o-series among + them) hand LiteLLM the refusal as an unrecognized field, which lands in + ``provider_specific_fields``. That first delta still has to stream as text, + otherwise the client gets ``stop_reason: refusal`` over an empty content + array and shows the user nothing. + """ + chunks = [ + _make_chunk(Delta(content=None, provider_specific_fields={"refusal": "I cannot fulfill this request."})), + _make_chunk(Delta(content=None), finish_reason="stop"), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="openai-model") + + events = _drain_sync(wrapper) + + assert _text_deltas(events) == ["I cannot fulfill this request."] + message_delta = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." + + +@pytest.mark.asyncio +async def test_streaming_chat_refusal_parked_in_provider_specific_fields_is_emitted_async(): + chunks = [ + _make_chunk(Delta(content=None, provider_specific_fields={"refusal": "I cannot fulfill this request."})), + _make_chunk(Delta(content=None), finish_reason="stop"), + ] + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="openai-model") + + events = await _drain_async(wrapper) + + assert _text_deltas(events) == ["I cannot fulfill this request."] + message_delta = next(event for event in events if event["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "refusal" + assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." + + def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved(): chunks = [ _make_chunk( From 05cba21763ee40477c2b4b13135f10763bb89957 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:40:20 -0700 Subject: [PATCH 60/63] fix(anthropic): split refusal off a combined finish_reason chunk A fake-streamed provider hands the adapter one chunk carrying both the delta payload and the finish_reason, which is exactly what the combined chunk splitter exists for, but its content check never listed the refusal. The translation short-circuits on finish_reason, so that refusal text was dropped and the client got `stop_reason: refusal` over an empty content array, the symptom this PR set out to fix. Both refusal accumulators also drop their `mutable-ok` lists for a plain string attribute --- .../adapters/streaming_iterator.py | 19 ++++++++++++------- .../responses_adapters/streaming_iterator.py | 6 +++--- .../test_streaming_iterator_first_delta.py | 8 +++++++- 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index b2513a5c046..9158ff4569f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -101,6 +101,10 @@ class _CombinedChunkSplitter: @staticmethod def _is_combined(chunk: "ModelResponseStream") -> bool: """True if ``chunk`` carries response content AND a finish_reason.""" + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + openai_chat_refusal_text, + ) + choices: Final = _optional_attr_sequence(chunk, "choices") if not choices: return False @@ -115,6 +119,7 @@ class _CombinedChunkSplitter: or _optional_attr(delta, "tool_calls") or _optional_attr(delta, "reasoning_content") or _optional_attr(delta, "thinking_blocks") + or openai_chat_refusal_text(delta) ) _PAYLOAD_FIELD_GROUPS: "tuple[tuple[str, ...], ...]" = ( @@ -306,7 +311,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Synthesized compaction block from compact_20260112 polyfill (streaming). self.compaction_block = compaction_block self.iterations_usage = iterations_usage - self._refusal_text_parts: list[str] = [] # mutable-ok: accumulates streamed refusal delta text across chunks + self._refusal_text: str = "" self.sent_compaction_block: bool = False # Per-phase flags so the compaction block's start/delta/stop events # are emitted (and the public state machine is advanced) in @@ -1001,7 +1006,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): self, processed_chunk: ContentBlockDelta | MessageBlockDelta, ) -> ContentBlockDelta | MessageBlockDelta: - if processed_chunk.get("type") != "message_delta" or not self._refusal_text_parts: + if processed_chunk.get("type") != "message_delta" or not self._refusal_text: return processed_chunk delta: Final = cast(Mapping[str, object], processed_chunk["delta"]) # cast-ok: keys checked before use if delta.get("stop_reason") == "max_tokens": @@ -1017,7 +1022,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): "delta": { # mutable-ok: fresh message_delta payload; never mutated after construction **delta, "stop_reason": "refusal", - "stop_details": refusal_stop_details("".join(self._refusal_text_parts)), + "stop_details": refusal_stop_details(self._refusal_text), }, }, ) @@ -1107,13 +1112,13 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): from .transformation import LiteLLMAnthropicMessagesAdapter - refusal_text: Final = openai_chat_refusal_text(chunk.choices[0].delta) - if refusal_text is not None: - self._refusal_text_parts.append(refusal_text) - if chunk.choices[0].finish_reason is not None: return False + refusal_text: Final = openai_chat_refusal_text(chunk.choices[0].delta) + if refusal_text is not None: + self._refusal_text = self._refusal_text + refusal_text + ( block_type, content_block_start, diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py index bcae732ad42..2e0a6a9df8f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/streaming_iterator.py @@ -54,7 +54,7 @@ class AnthropicResponsesStreamWrapper: self._sent_message_start = False self._sent_message_stop = False self._chunk_queue: deque[dict[str, object]] = deque() - self._refusal_text_parts: list[str] = [] # mutable-ok: accumulates streamed refusal delta text across chunks + self._refusal_text: str = "" self._sync_responses_iterator: Iterator[object] | None = None def _make_message_start(self) -> dict[str, object]: @@ -142,7 +142,7 @@ class AnthropicResponsesStreamWrapper: delta = getattr(event, "delta", "") or (event.get("delta", "") if isinstance(event, dict) else "") if not isinstance(delta, str) or not delta: return - self._refusal_text_parts.append(delta) + self._refusal_text = self._refusal_text + delta item_id = getattr(event, "item_id", None) or (event.get("item_id") if isinstance(event, dict) else None) block_idx = self._item_id_to_block_index.get(item_id, -1) if item_id else self._current_block_index if block_idx < 0: @@ -241,7 +241,7 @@ class AnthropicResponsesStreamWrapper: event.get("response") if isinstance(event, dict) else None ) output: Final = (getattr(response_obj, "output", None) or ()) if response_obj is not None else () - refusal_text: Final = responses_output_refusal_text(output) or ("".join(self._refusal_text_parts) or None) + refusal_text: Final = responses_output_refusal_text(output) or (self._refusal_text or None) status: Final = getattr(response_obj, "status", None) if response_obj is not None else None has_tool_call: Final = any( getattr(item, "type", None) == "function_call" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 4c5e2101fbf..fdd08eaa182 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -183,6 +183,10 @@ async def test_streaming_chat_refusal_parked_in_provider_specific_fields_is_emit def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved(): + """Fake-streamed responses arrive as one chunk carrying both the delta and the + finish_reason. The refusal has to be split off and streamed as text, or the + client gets ``stop_reason: refusal`` over an empty content array. + """ chunks = [ _make_chunk( Delta(content=None, refusal="I cannot fulfill this request."), @@ -193,6 +197,7 @@ def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved(): events = _drain_sync(wrapper) + assert _text_deltas(events) == ["I cannot fulfill this request."] message_delta = next(event for event in events if event["type"] == "message_delta") assert message_delta["delta"]["stop_reason"] == "refusal" assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." @@ -202,7 +207,7 @@ def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved(): async def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved_async(): chunks = [ _make_chunk( - Delta(content=None, refusal="I cannot fulfill this request."), + Delta(content=None, provider_specific_fields={"refusal": "I cannot fulfill this request."}), finish_reason="stop", ) ] @@ -210,6 +215,7 @@ async def test_streaming_chat_combined_refusal_and_finish_reason_is_preserved_as events = await _drain_async(wrapper) + assert _text_deltas(events) == ["I cannot fulfill this request."] message_delta = next(event for event in events if event["type"] == "message_delta") assert message_delta["delta"]["stop_reason"] == "refusal" assert message_delta["delta"]["stop_details"]["explanation"] == "I cannot fulfill this request." From fcb6d2267c09389ae1aa9e80e04a35caa5b8b470 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 23:51:44 -0700 Subject: [PATCH 61/63] fix(spend): keep a batch's claim row out of the logs a proxy was told not to write disable_spend_logs has to keep meaning that no request gets logged, and the row that makes a batch chargeable exactly once is the one row it cannot drop, so with logging off that row now carries only what tells the retrieves apart. SPEND_LOGS_URL deployments get their copy back too: the claim writes straight to this table, so the row is queued as well when an external writer is the one that takes the spend logs. --- litellm/proxy/db/db_spend_update_writer.py | 49 +++++++-- .../proxy/db/test_db_spend_update_writer.py | 99 +++++++++++++++++++ 2 files changed, 141 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ee7802a45d3..bdc014d7f13 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -89,6 +89,21 @@ def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" +_BATCH_COST_CLAIM_FIELDS: Final = frozenset({"request_id", "call_type", "spend", "startTime", "endTime", "status"}) + + +def _batch_cost_row_to_write(payload: SpendLogsPayload, disable_spend_logs: bool) -> Mapping[str, object]: + """Reduce a batch's cost row to what tells the retrieves apart when logging is off. + + A proxy run with spend logs disabled still needs one row per batch to charge it once, + so the row is written either way, but it carries no request of its own: no metadata, + no requester IP, no key, model, or token counts (LIT-7048). + """ + if disable_spend_logs is False: + return payload + return MappingProxyType({field: value for field, value in payload.items() if field in _BATCH_COST_CLAIM_FIELDS}) + + class _SpendBatch(Protocol): litellm_usertable: BatchTable litellm_verificationtoken: BatchTable @@ -336,12 +351,16 @@ class DBSpendUpdateWriter: self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool ) -> bool: if prisma_client is not None and _is_batch_cost_row(payload): - return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + return await self._claim_batch_cost_spend_log( + payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ) if disable_spend_logs is False: await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) return True - async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: + async def _claim_batch_cost_spend_log( + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", disable_spend_logs: bool + ) -> bool: """Write the batch's cost row now, or learn that another retrieve already did. Every retrieve of one batch shares this row, so the insert that lands first owns @@ -353,13 +372,17 @@ class DBSpendUpdateWriter: from litellm.repositories.table_repositories import SpendLogsRepository request_id: Final = payload["request_id"] + row: Final = _batch_cost_row_to_write(payload, disable_spend_logs) spend_logs: Final = SpendLogsRepository(prisma_client).table try: claimed: Final = await spend_logs.create_many( - data=[prisma_client.jsonify_object(payload)], # mutable-ok: prisma create_many takes a list + data=[prisma_client.jsonify_object(row)], # mutable-ok: prisma create_many takes a list skip_duplicates=True, ) if claimed == 1: + await self._forward_batch_cost_row( + row=row, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs + ) return True existing: Final = await spend_logs.find_unique( where={"request_id": request_id} # mutable-ok: prisma where clause @@ -368,7 +391,7 @@ class DBSpendUpdateWriter: verbose_proxy_logger.warning( "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e ) - await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) return True if existing is None or existing.call_type != CallTypes.aretrieve_batch.value or existing.status != "success": verbose_proxy_logger.warning( @@ -380,10 +403,22 @@ class DBSpendUpdateWriter: if existing.spend > 0: verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) return False - return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client) + return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) + + async def _forward_batch_cost_row( + self, row: Mapping[str, object], prisma_client: "PrismaClient", disable_spend_logs: bool + ) -> None: + """Queue the claimed row for an external spend log writer, which the claim went around. + + With ``SPEND_LOGS_URL`` set the queue posts every spend log to that writer instead of + inserting it, so a batch's cost row reaches it only by being queued here as well. + """ + if disable_spend_logs is True or os.getenv("SPEND_LOGS_URL") is None: + return + await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) async def _take_over_uncharged_batch_cost_row( - self, payload: SpendLogsPayload, prisma_client: "PrismaClient" + self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] ) -> bool: """Take the batch's cost row over from the poll that left it charging nothing. @@ -403,7 +438,7 @@ class DBSpendUpdateWriter: try: taken_over: Final = await SpendLogsRepository(prisma_client).table.update_many( data=prisma_client.jsonify_object( - MappingProxyType({field: value for field, value in payload.items() if field != "request_id"}) + MappingProxyType({field: value for field, value in row.items() if field != "request_id"}) ), where={ # mutable-ok: prisma where clause "request_id": request_id, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 4efb94b60aa..a8aeaced55f 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import os import re @@ -3169,6 +3170,104 @@ async def test_update_database_writes_no_ordinary_spend_row_with_spend_logs_disa assert db_writer._batch_database_updates.await_count == 1 +_BATCH_CLAIM_FIELDS = {"request_id", "call_type", "status", "spend", "startTime", "endTime"} + + +def _logged_batch_cost_payload() -> dict: + return { + **_batch_cost_payload(), + "api_key": "0e5b0e9e5f", + "model": "gpt-5.6-luna", + "user": "test-user", + "metadata": '{"batch_models": ["gpt-5.6-luna"]}', + "requester_ip_address": "127.0.0.1", + "proxy_server_request": '{"headers": {"user-agent": "litellm-batch-cost-check"}}', + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("disable_spend_logs", "logs_the_request"), + [(False, True), (True, False)], + ids=["spend_logs_on", "spend_logs_off"], +) +async def test_update_database_claims_a_batch_without_logging_the_request_that_polled_it( + disable_spend_logs: bool, logs_the_request: bool +): + """ + disable_spend_logs has to keep meaning that no request gets logged, and the batch's cost + row is the one row it cannot drop, so with logging off that row carries only what tells + the retrieves apart: no metadata, no requester IP, no key, model, or token counts. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + payload = _logged_batch_cost_payload() + + assert await _update_database_with(db_writer, prisma, payload, disable_spend_logs) is True + + claimed = prisma.db.litellm_spendlogs.create_many.await_args.kwargs["data"][0] + assert set(claimed) == (set(payload) if logs_the_request else _BATCH_CLAIM_FIELDS) + assert claimed["spend"] == 0.25 + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend_logs_url", "forwarded"), + [("http://spend-logs.internal", True), (None, False)], + ids=["an_external_writer_takes_the_rows", "rows_are_written_to_this_db"], +) +async def test_update_database_sends_a_claimed_batch_cost_row_on_to_an_external_spend_log_writer( + monkeypatch, spend_logs_url: str | None, forwarded: bool +): + """ + SPEND_LOGS_URL makes the flush post spend logs to that writer instead of inserting them, + and the claim writes straight to this table, so the batch's row reaches the writer only + by being queued as well. Queueing it with no writer configured would insert it twice. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + if spend_logs_url is None: + monkeypatch.delenv("SPEND_LOGS_URL", raising=False) + else: + monkeypatch.setenv("SPEND_LOGS_URL", spend_logs_url) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + queued = [row["request_id"] for row in prisma.spend_log_transactions] + assert queued == (["batch_abc_batch_cost"] if forwarded else []) + + +@pytest.mark.asyncio +async def test_update_database_forwards_no_batch_cost_row_a_later_retrieve_had_already_claimed(monkeypatch): + """The retrieve that lost the claim charges nothing, so it must not post a row either.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25) + prisma = _spend_logs_prisma(0, existing) + monkeypatch.setenv("SPEND_LOGS_URL", "http://spend-logs.internal") + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False + + assert prisma.spend_log_transactions == [] + + +@pytest.mark.asyncio +async def test_update_database_queues_only_the_claim_for_a_batch_it_could_not_write_with_logs_disabled(): + """A refused claim is retried through the queue, so what it queues has to stay unlogged too.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _logged_batch_cost_payload(), True) is True + + assert [set(row) for row in prisma.spend_log_transactions] == [_BATCH_CLAIM_FIELDS] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): """An unreachable DB must not drop the batch's only spend row, nor its charge.""" From b9f5cd60364aae12b654c3e01aa43173e4003e71 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Sat, 5 Sep 2026 23:56:09 -0700 Subject: [PATCH 62/63] ci: run unit tests on Python 3.12 (#39989) --- .github/workflows/_test-unit-base.yml | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 6f6822a975b..75b0f93fd77 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -55,17 +55,14 @@ on: permissions: contents: read +env: + UV_PYTHON: "3.12" + jobs: run: - name: ${{ matrix.python-version == '3.12' && 'Run tests' || format('Run tests (Python {0})', matrix.python-version) }} + name: Run tests runs-on: ubuntu-latest timeout-minutes: ${{ inputs.job-timeout-minutes }} - strategy: - fail-fast: false - matrix: - python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"] - env: - UV_PYTHON: ${{ matrix.python-version }} permissions: contents: read pull-requests: read @@ -88,7 +85,7 @@ jobs: timeout-minutes: 3 uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: - python-version: ${{ matrix.python-version }} + python-version: ${{ env.UV_PYTHON }} - name: Set up uv if: steps.changes.outputs.decision != 'skip' @@ -103,9 +100,9 @@ jobs: uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: path: ${{ env.UV_CACHE_DIR }} - key: ${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }} + key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }} restore-keys: | - ${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}- + ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}- - name: Cache the Rust build if: steps.changes.outputs.decision != 'skip' @@ -139,7 +136,7 @@ jobs: WORKERS: ${{ inputs.workers }} RERUNS: ${{ inputs.reruns }} DIST: ${{ inputs.dist }} - COVERAGE_CORE: ${{ contains(fromJSON('["3.10", "3.11"]'), matrix.python-version) && 'ctrace' || 'sysmon' }} + COVERAGE_CORE: sysmon run: | if [ "${WORKERS}" = "0" ]; then uv run --no-sync pytest ${TEST_PATH:?} \ @@ -166,7 +163,7 @@ jobs: fi - name: Save coverage report - if: always() && matrix.python-version == '3.12' && steps.changes.outputs.decision != 'skip' + if: always() && steps.changes.outputs.decision != 'skip' uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 with: name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }} From defd8661f4e359994bf57a7f9e51ed8c479f17ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:04:24 -0700 Subject: [PATCH 63/63] refactor(spend): stop queueing a batch's claim row for a writer the proxy never builds SPEND_LOGS_URL only diverts spend logs when db_writer_client is set, and nothing in the proxy ever assigns that global, so the queued copy was only ever skipped as a duplicate by the local insert. --- litellm/proxy/db/db_spend_update_writer.py | 15 ------- .../proxy/db/test_db_spend_update_writer.py | 43 ------------------- 2 files changed, 58 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index bdc014d7f13..9230be8055e 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -380,9 +380,6 @@ class DBSpendUpdateWriter: skip_duplicates=True, ) if claimed == 1: - await self._forward_batch_cost_row( - row=row, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs - ) return True existing: Final = await spend_logs.find_unique( where={"request_id": request_id} # mutable-ok: prisma where clause @@ -405,18 +402,6 @@ class DBSpendUpdateWriter: return False return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row) - async def _forward_batch_cost_row( - self, row: Mapping[str, object], prisma_client: "PrismaClient", disable_spend_logs: bool - ) -> None: - """Queue the claimed row for an external spend log writer, which the claim went around. - - With ``SPEND_LOGS_URL`` set the queue posts every spend log to that writer instead of - inserting it, so a batch's cost row reaches it only by being queued here as well. - """ - if disable_spend_logs is True or os.getenv("SPEND_LOGS_URL") is None: - return - await self._insert_spend_log_to_db(payload=prisma_client.jsonify_object(row), prisma_client=prisma_client) - async def _take_over_uncharged_batch_cost_row( self, payload: SpendLogsPayload, prisma_client: "PrismaClient", row: Mapping[str, object] ) -> bool: diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index a8aeaced55f..0bca7c9492c 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,7 +1,6 @@ import asyncio import copy import json -import os import re @@ -3212,48 +3211,6 @@ async def test_update_database_claims_a_batch_without_logging_the_request_that_p assert db_writer._batch_database_updates.await_count == 1 -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("spend_logs_url", "forwarded"), - [("http://spend-logs.internal", True), (None, False)], - ids=["an_external_writer_takes_the_rows", "rows_are_written_to_this_db"], -) -async def test_update_database_sends_a_claimed_batch_cost_row_on_to_an_external_spend_log_writer( - monkeypatch, spend_logs_url: str | None, forwarded: bool -): - """ - SPEND_LOGS_URL makes the flush post spend logs to that writer instead of inserting them, - and the claim writes straight to this table, so the batch's row reaches the writer only - by being queued as well. Queueing it with no writer configured would insert it twice. - """ - db_writer = DBSpendUpdateWriter() - db_writer._batch_database_updates = AsyncMock() - prisma = _spend_logs_prisma(1, None) - if spend_logs_url is None: - monkeypatch.delenv("SPEND_LOGS_URL", raising=False) - else: - monkeypatch.setenv("SPEND_LOGS_URL", spend_logs_url) - - assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True - - queued = [row["request_id"] for row in prisma.spend_log_transactions] - assert queued == (["batch_abc_batch_cost"] if forwarded else []) - - -@pytest.mark.asyncio -async def test_update_database_forwards_no_batch_cost_row_a_later_retrieve_had_already_claimed(monkeypatch): - """The retrieve that lost the claim charges nothing, so it must not post a row either.""" - db_writer = DBSpendUpdateWriter() - db_writer._batch_database_updates = AsyncMock() - existing = SimpleNamespace(call_type="aretrieve_batch", status="success", spend=0.25) - prisma = _spend_logs_prisma(0, existing) - monkeypatch.setenv("SPEND_LOGS_URL", "http://spend-logs.internal") - - assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is False - - assert prisma.spend_log_transactions == [] - - @pytest.mark.asyncio async def test_update_database_queues_only_the_claim_for_a_batch_it_could_not_write_with_logs_disabled(): """A refused claim is retried through the queue, so what it queues has to stay unlogged too."""