chore: merge litellm_internal_staging into fix/batch-retrieve-model-group

This commit is contained in:
mateo-berri 2026-09-06 02:36:39 -07:00
commit f66663cc8b
104 changed files with 5246 additions and 790 deletions

38
.github/e2e-stack/assert_tests_ran.py vendored Normal file
View file

@ -0,0 +1,38 @@
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
def main() -> int:
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")
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
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 file must execute a passing test, with no failures or errors\n")
return 1
if __name__ == "__main__":
sys.exit(main())

17
.github/e2e-stack/down.sh vendored Executable file
View file

@ -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

49
.github/e2e-stack/secrets_to_env.py vendored Normal file
View file

@ -0,0 +1,49 @@
import os
import re
import sys
from pathlib import Path
from typing import Final
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:
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
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:
_ = 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
return 0
if __name__ == "__main__":
sys.exit(main())

44
.github/e2e-stack/select_tests.py vendored Normal file
View file

@ -0,0 +1,44 @@
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$")
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)$"
r"|^tests/e2e/gateway/"
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 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 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, canary)) + "\n")
return 0
if __name__ == "__main__":
sys.exit(main())

207
.github/e2e-stack/up.sh vendored Executable file
View file

@ -0,0 +1,207 @@
#!/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}"
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.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa}"
JAEGER_IMAGE="${E2E_JAEGER_IMAGE:-jaegertracing/jaeger:2.10.0}"
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}"
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}"
chmod 700 "${STACK_DIR}" "${LOGS_DIR}" "${PIDS_DIR}"
chmod 755 "${CERTS_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}"
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 [[ -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}" \
"${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" <<EOF
events {}
http {
map \$http_upgrade \$connection_upgrade {
default upgrade;
'' close;
}
upstream litellm_gateways {
server ${NGINX_UPSTREAM_HOST}:${GATEWAY_PORT_1};
server ${NGINX_UPSTREAM_HOST}:${GATEWAY_PORT_2};
}
server {
listen ${LB_PORT};
client_max_body_size 100m;
location / {
proxy_pass http://litellm_gateways;
proxy_http_version 1.1;
proxy_set_header Host \$host;
proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
proxy_set_header X-Forwarded-Proto \$scheme;
proxy_set_header Upgrade \$http_upgrade;
proxy_set_header Connection \$connection_upgrade;
proxy_buffering off;
proxy_read_timeout 600s;
proxy_send_timeout 600s;
}
}
}
EOF
docker rm -f e2e-nginx >/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" <<EOF
LITELLM_PROXY_URL=http://127.0.0.1:${LB_PORT}
LITELLM_CONTROL_PLANE_URL=http://127.0.0.1:${BACKEND_PORT}
LITELLM_PROXY_REPLICA_URLS=http://127.0.0.1:${GATEWAY_PORT_1},http://127.0.0.1:${GATEWAY_PORT_2}
LITELLM_MASTER_KEY=${MASTER_KEY}
REDIS_HOST=127.0.0.1
REDIS_PORT=${REDIS_PORT}
E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem
DATABASE_URL=postgresql://${DATABASE_USER}:${DATABASE_PASSWORD}@${DATABASE_HOST}:${DATABASE_PORT}/${DATABASE_NAME}
EOF
log "stack is up; pytest env written to ${STACK_DIR}/stack.env"

View file

@ -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 }}

View file

@ -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

240
.github/workflows/test-e2e-changed.yml vendored Normal file
View file

@ -0,0 +1,240 @@
name: e2e-changed-tests
on:
pull_request:
concurrency:
group: e2e-changed-${{ github.event.pull_request.number }}
cancel-in-progress: true
permissions: {}
jobs:
detect:
name: Detect changed e2e tests
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:
GH_TOKEN: ${{ github.token }}
REPO: ${{ github.repository }}
PR_NUMBER: ${{ github.event.pull_request.number }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
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}" \
| 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}"
echo "selected e2e tests: ${tests}"
else
echo "any=false" >> "${GITHUB_OUTPUT}"
echo "no changed e2e test files supported by this stack; 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
ref: ${{ github.sha }}
- 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
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
aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license \
--query SecretString --output text \
| 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
id: boot
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: |
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
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"
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}" > "${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 [ "${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
- 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

View file

@ -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 \

View file

@ -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.",
@ -231,6 +236,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"],
},
}

View file

@ -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 \

View file

@ -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

View file

@ -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 \

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false;

View file

@ -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?

View file

@ -25,6 +25,27 @@ class BatchCostUsageResult:
failed_requests: int
_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:
"""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 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
async def calculate_batch_cost_and_usage(
file_content_dictionary: list[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],

View file

@ -81,6 +81,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,
)
@ -496,6 +497,13 @@ def cost_per_token(
# see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models
if call_type == "speech" or call_type == "aspeech":
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

View file

@ -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

View file

@ -11,6 +11,7 @@ from typing import (
Final,
Literal,
Protocol,
cast, # noqa: TID251 # rebuilt message_delta dict spans the ContentBlockDelta/MessageBlockDelta union
get_args,
)
@ -100,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
@ -114,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, ...], ...]" = (
@ -305,6 +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: 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 +579,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 +814,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 +1002,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:
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":
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
**processed_chunk,
"delta": { # mutable-ok: fresh message_delta payload; never mutated after construction
**delta,
"stop_reason": "refusal",
"stop_details": refusal_stop_details(self._refusal_text),
},
},
)
@staticmethod
def _delta_has_content(processed_chunk: Mapping[str, object]) -> bool:
"""Return True if a translated chunk carries a non-empty
@ -1035,6 +1069,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:
@ -1044,6 +1081,8 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
return False
if getattr(delta, "content", None):
return False
if openai_chat_refusal_text(delta):
return False
if getattr(delta, "reasoning_content", None):
return False
# thinking_blocks whose entries are all empty AND unsigned must not
@ -1067,13 +1106,19 @@ 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
# Example logic - customize based on your needs:
# If chunk indicates a tool call
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,

View file

@ -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,
@ -1314,6 +1318,8 @@ class LiteLLMAnthropicMessagesAdapter:
new_content.append(
AnthropicResponseContentBlockText(type="text", text=choice.message.content).model_dump()
)
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:
for tool_call in choice.message.tool_calls:
@ -1472,14 +1478,23 @@ class LiteLLMAnthropicMessagesAdapter:
choices=response.choices,
tool_name_mapping=tool_name_mapping,
)
refusal_text: Final = next(
(text for choice in response.choices if (text := openai_chat_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)
@ -1501,6 +1516,7 @@ class LiteLLMAnthropicMessagesAdapter:
usage=anthropic_usage,
content=anthropic_content,
stop_reason=anthropic_finish_reason,
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
@ -1541,7 +1557,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 openai_chat_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 []
@ -1613,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,

View file

@ -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."""

View file

@ -1,13 +1,18 @@
# What is this?
## Translates OpenAI call to Anthropic `/v1/messages` format
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
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
@ -49,6 +54,8 @@ class AnthropicResponsesStreamWrapper:
self._sent_message_start = False
self._sent_message_stop = False
self._chunk_queue: deque[dict[str, object]] = deque()
self._refusal_text: str = ""
self._sync_responses_iterator: Iterator[object] | None = None
def _make_message_start(self) -> dict[str, object]:
return {
@ -131,6 +138,24 @@ class AnthropicResponsesStreamWrapper:
)
return
if event_type == "response.refusal.delta":
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 = 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:
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 ----
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)
@ -215,34 +240,47 @@ 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 = 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"
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
message_delta_payload: Final = { # mutable-ok: fresh message_delta payload built per chunk
"stop_reason": stop_reason,
"stop_sequence": None,
**(
{ # 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
),
}
self._chunk_queue.append(
{
"type": "message_delta",
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
"delta": message_delta_payload,
"usage": dict(anthropic_usage),
}
)
@ -266,10 +304,20 @@ 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)
sync_iterator: Final = self._sync_responses_iterator
missing: Final = object()
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()
except StopAsyncIteration:
pass
except Exception as e:

View file

@ -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,
@ -624,6 +628,9 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
content: Final[list[dict[str, object]]] = []
stop_reason: AnthropicFinishReason = "end_turn"
refusal_text: Final = responses_output_refusal_text(
cast(Iterable[object], response.output) # cast-ok: output items re-validated per item
)
for item in response.output:
if isinstance(item, ResponseReasoningItem):
@ -631,10 +638,17 @@ 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":
content.append(
AnthropicResponseContentBlockText(
type="text", text=getattr(part, "refusal", "") or ""
).model_dump()
)
elif isinstance(item, ResponseFunctionToolCall):
try:
@ -654,11 +668,21 @@ 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":
content.append(
AnthropicResponseContentBlockText(
type="text", text=part.get("refusal", "") or ""
).model_dump()
)
elif item_type == "reasoning":
content.extend(
self._thinking_blocks_from_reasoning_item(
@ -679,10 +703,10 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
).model_dump(exclude_none=True)
)
stop_reason = "tool_use"
# 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)
@ -695,4 +719,5 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
usage=anthropic_usage,
content=content,
stop_reason=stop_reason,
stop_details=(refusal_stop_details(refusal_text) if stop_reason == "refusal" else None),
)

View file

@ -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,37 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = (
)
class AzureAIGPT5Config(OpenAIGPT5Config):
@classmethod
def _model_map_lookup_name(cls, model: str) -> str:
"""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/<model>`` 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/<model>`` 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()
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 +87,27 @@ class AzureAIStudioConfig(OpenAIConfig):
return supported_params
def map_openai_params(
self,
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[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,
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.

View file

@ -1,9 +1,14 @@
import re
from collections.abc import Mapping
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
from pydantic import TypeAdapter, ValidationError
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -21,6 +26,61 @@ from litellm.types.utils import TokenCountResponse
from litellm.utils import supports_response_schema, supports_system_messages
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]]
_VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER: Final = 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=1)
def _bundled_vertex_ai_lyria_model_infos() -> Mapping[str, VertexAILyriaModelInfo]:
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
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 = _vertex_ai_lyria_model_key(model)
runtime_model_info: Final = _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key))
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_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):
def __init__(
self,

View file

@ -8,17 +8,25 @@ 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
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
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 (
@ -35,6 +43,10 @@ else:
LiteLLMLoggingObj = Any
HttpxBinaryResponseContent = Any
_LyriaVoice: TypeAlias = (
str | dict | None
) # mutable-ok: inherited interface supports structured provider voice dictionaries
class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
"""
@ -472,3 +484,209 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
# Initialize the HttpxBinaryResponseContent instance
return HttpxBinaryResponseContent(response)
class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
@classmethod
def is_lyria_model(cls, model: str) -> bool:
return get_vertex_ai_lyria_model_info(model=model) is not None
@staticmethod
def _get_model_info(model: str) -> VertexAILyriaModelInfo:
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: # 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, # mutable-ok: inherited provider interface accepts a concrete parameter dictionary
voice: _LyriaVoice = None,
drop_params: bool = False,
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:
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: 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)
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, # mutable-ok: inherited provider interface accepts concrete LiteLLM parameters
) -> str:
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,
)
def mint_access_token(
_credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
) -> tuple[str, str]:
return "", project_id or 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("/")
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,
model: str,
input: str,
voice: str | None,
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),
project_id=self.safe_get_vertex_ai_project(litellm_params),
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: Final = model.removeprefix("vertex_ai/")
model_info: Final = self._get_model_info(model=model)
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
],
"parameters": { # mutable-ok: predict dispatch requires a concrete parameters dictionary
"sample_count": 1
},
}
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
),
}
)
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: 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: Final = response_json.get("predictions") or ()
if predictions:
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 content in content_items:
if content.get("type") == "audio" and content.get("data"):
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")
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=binary_data,
headers=MappingProxyType({"content-type": media_type}),
)
)

View file

@ -8247,6 +8247,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,
)
@ -8271,7 +8272,11 @@ def speech(
# Vertex AI Text-to-Speech (Google Cloud TTS)
if text_to_speech_provider_config is None:
text_to_speech_provider_config = VertexAITextToSpeechConfig()
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)

View file

@ -3485,6 +3485,55 @@
"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,
"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",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": false,
"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_cache_breakpoint": 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,
@ -7189,7 +7238,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,
@ -7455,7 +7504,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,
@ -47141,6 +47190,99 @@
"output_cost_per_token": 4e-07,
"supports_tool_choice": true
},
"vertex_ai/lyria-002": {
"litellm_provider": "vertex_ai",
"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"
],
"supported_endpoints": [
"/v1/audio/speech"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supports_audio_output": true,
"vertex_ai_audio_api": "lyria_predict"
},
"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": "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_audio_formats": [
"mp3"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": false,
"supports_web_search": false,
"vertex_ai_audio_api": "lyria_interactions"
},
"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": "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_audio_formats": [
"mp3",
"wav"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": 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,
"litellm_provider": "vertex_ai-llama_models",

View file

@ -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

View file

@ -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]

View file

@ -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,

View file

@ -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,

View file

@ -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": [
{

View file

@ -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,

View file

@ -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
@ -84,6 +85,25 @@ else:
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
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
@ -215,7 +235,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,
@ -232,7 +257,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:
@ -262,11 +287,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:
await self._insert_spend_log_to_db(
payload=payload,
prisma_client=prisma_client,
)
await self._enqueue_tool_usage_transaction(
payload=payload,
completion_response=completion_response,
@ -306,6 +332,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 "
@ -318,7 +345,102 @@ class DBSpendUpdateWriter:
org_id,
end_user_id,
)
return
return True
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, 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", 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
the charge and every later one finds the row and charges nothing (LIT-7048). Only
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
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(row)], # 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=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(
"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
return await self._take_over_uncharged_batch_cost_row(payload=payload, prisma_client=prisma_client, row=row)
async def _take_over_uncharged_batch_cost_row(
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.
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 row.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; 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
return True
async def _enqueue_tool_usage_transaction(
self,

View file

@ -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 (
@ -37,6 +38,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
from litellm.proxy.utils import ProxyUpdateSpend
from litellm.types.utils import (
CallTypes,
LiteLLMBatch,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@ -248,6 +250,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
and not batch_cost_is_final(completion_response)
):
verbose_proxy_logger.debug(
"Cost tracking deferred for batch %s still in status %s",
completion_response.id,
completion_response.status,
)
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))
@ -285,7 +299,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,
@ -302,6 +316,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.)
@ -578,9 +594,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,
@ -605,6 +621,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(
@ -630,6 +649,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:

View file

@ -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

View file

@ -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")

View file

@ -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(

View file

@ -10,7 +10,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_generation_cost,
get_vertex_location_from_url,
)
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
ModelResponseIterator as VertexModelResponseIterator,
)
@ -44,7 +47,6 @@ else:
PassThroughEndpointLogging = Any
LiteLLMBatch = Any
# Define EndpointType locally to avoid import issues
EndpointType = Any
@ -270,6 +272,16 @@ class VertexPassthroughLoggingHandler:
_json_response: Final[dict[str, object]] = httpx_response.json()
litellm_prediction_response: ModelResponse | EmbeddingResponse | ImageResponse = ModelResponse()
if VertexPassthroughLoggingHandler._is_audio_predict_response(
model=model,
json_response=_json_response,
):
return VertexPassthroughLoggingHandler._handle_audio_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 +335,71 @@ class VertexPassthroughLoggingHandler:
"kwargs": kwargs,
}
@staticmethod
def _handle_audio_predict_response(
json_response: dict, # mutable-ok: passthrough logging receives the decoded provider response dictionary
logging_obj: LiteLLMLoggingObj,
model: str,
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
)
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
"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[ # 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
] = { # mutable-ok: callback contract requires a concrete response dictionary
"response": json_response,
}
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, # mutable-ok: predicate inspects the decoded provider response dictionary without mutation
) -> bool:
return (
VertexPassthroughLoggingHandler._get_audio_prediction_count(json_response=json_response) > 0
and get_vertex_ai_lyria_generation_cost(model=model) is not None
)
@staticmethod
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
return sum(
1
for prediction in predictions
if isinstance(prediction, dict) and (prediction.get("audioContent") or prediction.get("bytesBase64Encoded"))
)
@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."""

View file

@ -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,

View file

@ -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?

View file

@ -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.

View file

@ -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
@ -55,16 +55,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
@ -136,18 +132,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
# ------------

View file

@ -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
@ -26,6 +27,12 @@ class RoutingArgs(LiteLLMPydanticObjectBase):
max_latency_list_size: int = 10
def _average_latency(samples: Sequence[float]) -> float:
if not samples:
return 0.0
return sum(samples) / len(samples)
class LowestLatencyLoggingHandler(CustomLogger):
test_flag: bool = False
logged_success: int = 0
@ -438,23 +445,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
@ -463,7 +460,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
# -------------- #
@ -473,7 +470,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

View file

@ -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)

View file

@ -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.

View file

@ -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: ReadOnly[AnthropicStopDetails]
class ServerToolUsage(TypedDict, total=False):
@ -658,7 +665,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):

View file

@ -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

View file

@ -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

View file

@ -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: 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
@ -335,6 +337,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
"image_generation",
"chat",
"audio_transcription",
"audio_speech",
"responses",
"ocr",
"realtime",

View file

@ -5946,6 +5946,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:
@ -9436,9 +9438,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 (

View file

@ -3485,6 +3485,55 @@
"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,
"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",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_computer_use": true,
"supports_function_calling": true,
"supports_max_reasoning_effort": false,
"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_cache_breakpoint": 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,
@ -7189,7 +7238,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,
@ -7455,7 +7504,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,
@ -47141,6 +47190,99 @@
"output_cost_per_token": 4e-07,
"supports_tool_choice": true
},
"vertex_ai/lyria-002": {
"litellm_provider": "vertex_ai",
"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"
],
"supported_endpoints": [
"/v1/audio/speech"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supports_audio_output": true,
"vertex_ai_audio_api": "lyria_predict"
},
"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": "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_audio_formats": [
"mp3"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": false,
"supports_web_search": false,
"vertex_ai_audio_api": "lyria_interactions"
},
"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": "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_audio_formats": [
"mp3",
"wav"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
],
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": 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,
"litellm_provider": "vertex_ai-llama_models",

View file

@ -623,6 +623,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.",
@ -846,6 +857,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.",

View file

@ -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",

View file

@ -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?

View file

@ -0,0 +1,203 @@
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"
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")
@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
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", ("<testsuite/>", "<testsuite", '<testsuite><testcase name="a"/></testsuite>'))
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
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"
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/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"),
("tests/e2e/logging/test_datadog_e2e.py",),
),
),
)
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
@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")
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()

View file

@ -54,6 +54,20 @@ Some suites need extra services the bare proxy does not start. The `logging/` OT
A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships
### 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, `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 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
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, 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
### 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

View file

@ -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)

View file

@ -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,),
)

View file

@ -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],

View file

@ -33,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)
@ -87,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")

View file

@ -0,0 +1,64 @@
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
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: openai-text-embedding-3-small
litellm_params:
model: openai/text-embedding-3-small
api_key: os.environ/OPENAI_API_KEY
mcp_servers:
devin:
url: "https://mcp.devin.ai/mcp"
auth_type: api_key
auth_value: os.environ/DEVIN_API_KEY

View file

@ -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,
@ -72,6 +75,7 @@ from e2e_config import (
POLL_INTERVAL,
POLL_TIMEOUT,
PROXY_BASE_URL,
PROXY_REPLICA_URLS,
REQUEST_TIMEOUT,
SLOW_PROVIDER_TIMEOUT_SECONDS,
settle_propagation,
@ -80,10 +84,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
@ -109,6 +114,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]],
@ -172,9 +187,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,
@ -185,8 +232,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}"
)
@ -195,6 +242,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
@ -311,12 +359,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",
@ -331,16 +380,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=ModelsListParams(),
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,
@ -352,16 +395,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=ModelsListParams(),
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
@ -548,16 +602,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
@ -574,8 +632,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,
)

View file

@ -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",)

View file

@ -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,57 @@ 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="validating",
output_file_id=output_file_id,
request_counts=counts,
).model_copy(update={"status": status})
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
@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
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
@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(status, 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

View file

@ -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
@ -2051,6 +2060,18 @@ 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):
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",
[

View file

@ -632,6 +632,86 @@ 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="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), ("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
handle_completed_batch = AsyncMock()
monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch)
batch = self._batch(status, output_file_id)
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")
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."""

View file

@ -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(

View file

@ -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"

View file

@ -108,6 +108,136 @@ def _text_deltas(events: List[dict]) -> List[str]:
]
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"),
]
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",
"stop_details": {
"type": "refusal",
"category": None,
"explanation": "I cannot fulfill this request.",
},
}
@pytest.mark.asyncio
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"),
]
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_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():
"""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."),
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_combined_refusal_and_finish_reason_is_preserved_async():
chunks = [
_make_chunk(
Delta(content=None, provider_specific_fields={"refusal": "I cannot fulfill this request."}),
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."
@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"]

View file

@ -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
@ -308,3 +325,60 @@ class TestResponseCompletedUsage:
"cache_creation_input_tokens": 10,
"cache_read_input_tokens": 4004,
}
class TestRefusalStreamEvents:
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."}]}],
usage=None,
)
chunks = _process_all(
[
{"type": "response.created"},
{"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 [chunk["type"] for chunk in chunks] == [
"message_start",
"content_block_start",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
]
assert chunks[2]["delta"] == {"type": "text_delta", "text": "I cannot fulfill this."}
assert chunks[4]["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(
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_incomplete_status_takes_precedence_over_refusal(self):
response = SimpleNamespace(
status="incomplete",
output=[{"type": "message", "content": [{"type": "refusal", "refusal": "Partial refusal"}]}],
usage=None,
)
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"] == "max_tokens"
assert "stop_details" not in message_delta["delta"]

View file

@ -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):
from openai.types.responses import ResponseOutputMessage, ResponseOutputRefusal
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:
"""Build a mock ResponseFunctionToolCall."""
from openai.types.responses import ResponseFunctionToolCall # type: ignore[import]
@ -1280,6 +1290,41 @@ 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"
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 = {
"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"
assert result.get("stop_details", {}).get("explanation") == "Refused by policy"
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 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'."""
response = _make_mock_response(
@ -1338,9 +1383,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."""
@ -1483,9 +1526,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}],
},
]
@ -1632,9 +1673,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}],
},
]
@ -1667,9 +1706,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(
@ -1778,9 +1815,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(
@ -1810,9 +1845,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",
@ -1859,7 +1892,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"] == [

View file

@ -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,46 @@ 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):
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_a_gpt_5_name_without_a_foundry_row_keeps_reading_its_own_entry(
monkeypatch: pytest.MonkeyPatch, _local_model_cost_map
):
"""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")
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.

View file

@ -1738,3 +1738,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")

View file

@ -0,0 +1,230 @@
from datetime import datetime
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
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(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/lyria-002",
{
"vertex_ai_audio_api": "lyria_predict",
"supported_audio_formats": ["wav"],
"output_cost_per_image": 0.06,
},
)
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)
def test_audio_predict_response_uses_model_map_metadata(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/music-audio-preview",
{
"vertex_ai_audio_api": "lyria_predict",
"supported_audio_formats": ["wav"],
"output_cost_per_image": 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(0.5)
assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.5)
def test_audio_predict_response_supports_bytes_base64_encoded(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/lyria-002",
{
"vertex_ai_audio_api": "lyria_predict",
"supported_audio_formats": ["wav"],
"output_cost_per_image": 0.06,
},
)
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)
@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,
runtime_entry_is_missing: bool,
local_model_cost_map: None,
) -> 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 != "output_cost_per_image"
},
)
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"}]},
)
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)
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"]
)

View file

@ -1,14 +1,17 @@
import base64
from typing import Final
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 +44,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 +105,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 +168,391 @@ 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)
@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()
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_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()
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",
},
},
),
],
)
def test_transform_request(
self,
model,
response_format,
expected_body,
):
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,
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": "UklGRiQAAABXQVZFZm10IA==",
}
]
},
b"RIFF$\x00\x00\x00WAVEfmt ",
"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",
),
(
"lyria-3-pro-preview",
{
"outputs": [
{
"type": "audio",
"data": "UklGRiQAAABXQVZFZm10IA==",
}
]
},
b"RIFF$\x00\x00\x00WAVEfmt ",
"audio/wav",
),
],
)
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.response.headers["content-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( # test-quality-ok: litellm.speech has no seam for Vertex token minting
VertexAILyriaTextToSpeechConfig,
"_ensure_access_token",
return_value=("mock-token", "music-project"),
),
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,
):
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 +566,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 +585,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

View file

@ -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),

View file

@ -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

View file

@ -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
)

View file

@ -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()

View file

@ -857,6 +857,24 @@ 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):
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]
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:
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 litellm.proxy.auth.model_checks import get_complete_model_list

View file

@ -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
@ -2936,7 +2937,9 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
"call_type, expects_flush",
[("aresponses", True), ("responses", 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.
@ -2957,6 +2960,303 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c
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, 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,
response_cost: float = 0.25,
) -> 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", 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
),
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=response_cost,
)
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", 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",
],
)
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
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()
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
@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_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 False
assert db_writer._batch_database_updates.await_count == 0
@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"),
[
(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
_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
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."""
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",

View file

@ -1,4 +1,4 @@
import asyncio
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, patch
@ -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,107 @@ 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,
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"),
[
("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=[
"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_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
):
"""
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: 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,
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(return_value=row_claimed)
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(),
)
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_release_budget_reservation.assert_not_awaited()
else:
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 +1175,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 +1762,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 +1841,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.asyncio
@ -1875,9 +1939,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
@ -1923,9 +1985,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:

View file

@ -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

View file

@ -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)
# ---------------------------------------------------------------------------

View file

@ -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}

View file

@ -163,3 +163,31 @@ 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):
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

View file

@ -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

View file

@ -389,14 +389,24 @@ 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):
"""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."""
@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_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="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",
@ -404,5 +414,4 @@ class TestGpt6AstraAdvertisesItsDocumentedLevels:
"medium",
"high",
"xhigh",
"max",
)

View file

@ -1,6 +1,7 @@
import json
from pathlib import Path
from typing import Final
import pytest
@ -146,6 +147,51 @@ 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),
],
)
@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 != "output_cost_per_image"},
)
elif runtime_state in ("custom_zero", "custom_price"):
multiplier: Final = 0 if runtime_state == "custom_zero" else 2
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,
prompt="A bright synth track",
call_type=call_type,
)
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):
expected_pricing = {

View file

@ -1091,6 +1091,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"],
@ -1113,6 +1124,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",
@ -2879,6 +2891,60 @@ 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["mode"] == "audio_speech"
assert clip["mode"] == "audio_speech"
assert pro["mode"] == "audio_speech"
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
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",
]
assert pro["supported_endpoints"] == [
"/v1beta/interactions",
"/v1/audio/speech",
]
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
def test_model_info_for_fireworks_short_form_models():
"""
Test that fireworks_ai short-form model entries (fireworks_ai/<model>)

View file

@ -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 }

View file

@ -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

View file

@ -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<string, unknown>) => null),
const { teamListCall, authorizedSession } = vi.hoisted(() => ({
teamListCall: vi.fn(() => new Promise(() => {})),
authorizedSession: vi.fn(),
}));
vi.mock("@/components/user_dashboard", () => ({
default: (props: Record<string, unknown>) => 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 }) => (
<div>
{headerActions}
<table aria-label="Virtual Keys" />
</div>
),
}));
vi.mock("@/components/organisms/create_key_button", () => ({
default: () => <button type="button">Create Key</button>,
}));
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(<ApiKeysDashboard />);
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(<ApiKeysDashboard />);
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(<ApiKeysDashboard />);
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(<ApiKeysDashboard />);
window.dispatchEvent(new Event("beforeunload"));
expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]');
expect(sessionStorage.getItem("selectedModel")).toBe("gpt-5.5");
});
});

View file

@ -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<Team[] | null>(null);
const [keys, setKeys] = useState<KeyResponse[] | null>([]);
const [createClicked, setCreateClicked] = useState<boolean>(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 (
<UserDashboard
userID={userID}
userRole={userRole}
premiumUser={premiumUser ?? false}
teams={teams}
keys={keys}
setUserRole={setUserRole}
userEmail={userEmail}
setUserEmail={setUserEmail}
setTeams={setTeams}
setKeys={setKeys}
addKey={addKey}
createClicked={createClicked}
autoOpenCreate={autoOpenCreate}
prefillData={prefillData}
/>
<main className="flex h-full flex-col p-8">
<VirtualKeysTable
headerActions={
isViewOnly ? undefined : (
<CreateKey
team={null}
teams={teams}
data={keys}
addKey={addKey}
autoOpenCreate={autoOpenCreate}
prefillData={prefillData}
/>
)
}
/>
</main>
);
}

View file

@ -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 = {

View file

@ -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";

View file

@ -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);
};

View file

@ -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 {

Some files were not shown because too many files have changed in this diff Show more